From 637aa87026c7e820d365fbb477915d983e833435 Mon Sep 17 00:00:00 2001
From: Phillip Munn <25081499+phillipmunn@users.noreply.github.com>
Date: Sun, 23 Aug 2026 20:47:41 +0100
Subject: [PATCH 001/249] feat(server): Add OpenAI /v1/models/{model_id}
 endpoint

---
 python/freetoken/server/openai_api.py | 29 +++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b4becd2631..f88f656ad0 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -143,6 +143,35 @@ async def v1_models():
             default_reasoning_effort=default_effort,
         )])
 
+    @app.get("/v1/models/{model_id:path}")
+    async def v1_model(model_id: str):
+        """Retrieve metadata for the currently served model.
+
+        OpenAI-compatible clients may use this endpoint to validate a model
+        name before sending completion requests.
+        """
+        state = get_state()
+        served_model_id = _served_model_name(state)
+
+        if model_id != served_model_id:
+            return create_error_response(
+                f"The model '{model_id}' does not exist",
+                status_code=404,
+                err_type="invalid_request_error",
+                param="model",
+                code="model_not_found",
+            )
+
+        ctx = _model_context_length(state)
+        efforts, default_effort = await _effort_fields(state)
+        return ModelCard(
+            id=served_model_id,
+            root=state.config.model_path,
+            max_model_len=ctx,
+            context_length=ctx,
+            supported_reasoning_efforts=efforts,
+            default_reasoning_effort=default_effort,
+        )
 
 async def handle_chat_completion(
     req: ChatCompletionRequest,

From 6241178a7550e5a0ffadd85f0eb04af0160f9988 Mon Sep 17 00:00:00 2001
From: ascorb12 <Ascorb@hotmail.com>
Date: Sun, 23 Aug 2026 20:52:10 -0400
Subject: [PATCH 002/249] fix(scheduler): clamp admission to the actual KV pool

A prompt that fits the model's advertised max_seq_len but exceeds the
allocated KV pool passes admission and is queued forever with no error,
no log line, and an idle engine (issue #111). Clamp the admission bound
to num_pages * page_size and reuse the existing too-long rejection path
so such prompts fail loudly with context_length_exceeded instead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 python/freetoken/scheduler/scheduler.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 355411617e..38bf877931 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -487,10 +487,17 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None:
                 )
                 return
             input_len, max_seq_len = len(msg.input_ids), self.engine.max_seq_len
-            max_output_len = max_seq_len - input_len
+            # max_seq_len is the model's advertised context, which can far exceed the
+            # KV pool actually allocated (see issue #111): a prompt that passes this
+            # check but can never be granted enough pages is queued forever with no
+            # error and no log line. Clamp admission to the real pool so oversized
+            # prompts fail loudly with the same error clients already understand.
+            pool_tokens = self.engine.num_pages * self.config.page_size
+            effective_max = min(max_seq_len, pool_tokens)
+            max_output_len = effective_max - input_len
             if max_output_len <= 0:
                 logger.warning_rank0(
-                    f"Input sequence length {input_len} exceeds {max_seq_len}, "
+                    f"Input sequence length {input_len} exceeds {effective_max}, "
                     f"request {msg.uid} is dropped."
                 )
                 # Tell the client instead of dropping silently — otherwise its wait_for_ack
@@ -502,7 +509,7 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None:
                             # "prompt is too long: N tokens > M" is the phrasing Claude Code and
                             # OpenClaw match on; the Anthropic wire has no error code to read.
                             error=(
-                                f"prompt is too long: {input_len} tokens > {max_seq_len} maximum "
+                                f"prompt is too long: {input_len} tokens > {effective_max} maximum "
                                 f"(prompt + generation); shorten the prompt or increase the KV "
                                 f"cache budget"
                             ),

From 102eadc9c2cb1c19a1785394219188aef3c95175 Mon Sep 17 00:00:00 2001
From: Ruslan Rakhimov <risonyo@gmail.com>
Date: Mon, 24 Aug 2026 17:51:44 +0300
Subject: [PATCH 003/249] fix(kernel): clear the latched CUDA error before
 raising in pinned_tensor

TORCH_CHECK reports a failed CUDA call but does not read the status out of
the runtime, so the error stays latched in the calling thread. The errors
raised here are non-sticky and the context remains usable, but the next
unrelated CUDA call picks up the stale status and fails as if it were its
own.

That turns a handled exception into a process-wide fault. host_banks.pin()
catches the RuntimeError from host_register() and re-raises a friendlier
message; any torch call made while handling that failure would report the
stale cudaHostRegister error instead of its own.

Route every CUDA call in the file through an FT_CUDA_CHECK macro that drains
the status with cudaGetLastError() before raising. The two
cudaDeviceGetAttribute calls in host_ptr_identity() were previously
unchecked and are now checked as well. Check messages are unchanged.

Add a test that drives the failure path directly and asserts the context
survives it.
---
 .../freetoken/kernel/csrc/pinned_tensor.cpp   | 57 +++++++++++--------
 tests/kernels/test_pinned_tensor.py           | 17 ++++++
 2 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp
index c3947adfa3..e2678d856e 100644
--- a/python/freetoken/kernel/csrc/pinned_tensor.cpp
+++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp
@@ -4,6 +4,20 @@
 
 namespace {
 
+// A failed runtime call latches its status in the calling thread until someone
+// reads it. TORCH_CHECK reports the error but does not clear it, so the next
+// unrelated CUDA call -- torch's, in practice -- reports *this* failure instead
+// of its own. Drain it before raising: every error below is non-sticky, so the
+// context stays usable for a caller that handles the exception.
+#define FT_CUDA_CHECK(expr, ...)                                               \
+  do {                                                                         \
+    const cudaError_t ft_err_ = (expr);                                        \
+    if (ft_err_ != cudaSuccess) {                                              \
+      cudaGetLastError();                                                      \
+      TORCH_CHECK(false, __VA_ARGS__, cudaGetErrorString(ft_err_));            \
+    }                                                                          \
+  } while (0)
+
 void free_pinned(void *ptr) {
   if (ptr != nullptr) {
     cudaFreeHost(ptr);
@@ -34,9 +48,8 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) {
   const size_t alloc_nbytes = static_cast<size_t>(nbytes == 0 ? 1 : nbytes);
 
   void *data_ptr = nullptr;
-  const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes);
-  TORCH_CHECK(alloc_err == cudaSuccess,
-              "cudaMallocHost failed: ", cudaGetErrorString(alloc_err));
+  FT_CUDA_CHECK(cudaMallocHost(&data_ptr, alloc_nbytes),
+                "cudaMallocHost failed: ");
 
   auto options = input.options().device(torch::kCPU).pinned_memory(true);
 
@@ -58,10 +71,9 @@ torch::Tensor alloc_pinned_tensor(std::vector<int64_t> sizes,
   // Portable + mapped: the offload gather kernel reads these banks straight
   // from host memory (zero-copy), which requires device-mapped pinned pages.
   void *data_ptr = nullptr;
-  const cudaError_t alloc_err = cudaHostAlloc(
-      &data_ptr, alloc_nbytes, cudaHostAllocPortable | cudaHostAllocMapped);
-  TORCH_CHECK(alloc_err == cudaSuccess,
-              "cudaHostAlloc failed: ", cudaGetErrorString(alloc_err));
+  FT_CUDA_CHECK(cudaHostAlloc(&data_ptr, alloc_nbytes,
+                              cudaHostAllocPortable | cudaHostAllocMapped),
+                "cudaHostAlloc failed: ");
 
   auto options = torch::TensorOptions()
                      .dtype(dtype)
@@ -76,37 +88,34 @@ torch::Tensor alloc_pinned_tensor(std::vector<int64_t> sizes,
 // device address). Zero-copy consumers resolve bank base addresses through these.
 bool host_ptr_identity() {
   int device = 0;
-  const cudaError_t err = cudaGetDevice(&device);
-  TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err));
+  FT_CUDA_CHECK(cudaGetDevice(&device), "cudaGetDevice failed: ");
   int uva = 0, reg = 0;
-  cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device);
-  cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device);
+  FT_CUDA_CHECK(cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device),
+                "cudaDeviceGetAttribute(UnifiedAddressing) failed: ");
+  FT_CUDA_CHECK(cudaDeviceGetAttribute(
+                    &reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device),
+                "cudaDeviceGetAttribute(CanUseHostPointerForRegisteredMem) failed: ");
   return uva == 1 && reg == 1;
 }
 
 int64_t host_device_ptr(int64_t host_ptr) {
   void *dev_ptr = nullptr;
-  const cudaError_t err =
-      cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast<void *>(host_ptr), 0);
-  TORCH_CHECK(err == cudaSuccess,
-              "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ",
-              cudaGetErrorString(err));
+  FT_CUDA_CHECK(
+      cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast<void *>(host_ptr), 0),
+      "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ");
   return reinterpret_cast<int64_t>(dev_ptr);
 }
 
 void host_register(int64_t addr, int64_t nbytes) {
-  const cudaError_t err =
-      cudaHostRegister(reinterpret_cast<void *>(addr), static_cast<size_t>(nbytes),
-                       cudaHostRegisterPortable | cudaHostRegisterMapped);
-  TORCH_CHECK(err == cudaSuccess,
-              "cudaHostRegister failed: ", cudaGetErrorString(err));
+  FT_CUDA_CHECK(cudaHostRegister(reinterpret_cast<void *>(addr),
+                                 static_cast<size_t>(nbytes),
+                                 cudaHostRegisterPortable | cudaHostRegisterMapped),
+                "cudaHostRegister failed: ");
 }
 
 int64_t driver_cuda_version() {
   int version = 0;  // stays 0 when no driver is installed
-  const cudaError_t err = cudaDriverGetVersion(&version);
-  TORCH_CHECK(err == cudaSuccess,
-              "cudaDriverGetVersion failed: ", cudaGetErrorString(err));
+  FT_CUDA_CHECK(cudaDriverGetVersion(&version), "cudaDriverGetVersion failed: ");
   return version;
 }
 
diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py
index e61108fd53..589565bd3c 100644
--- a/tests/kernels/test_pinned_tensor.py
+++ b/tests/kernels/test_pinned_tensor.py
@@ -128,6 +128,23 @@ def test_host_device_ptr_is_identity_under_uva():
     assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr()
 
 
+def test_failed_pinned_call_leaves_the_context_usable():
+    if not torch.cuda.is_available():
+        pytest.skip("needs CUDA")
+
+    from freetoken.kernel.pinned import _load_pinned_extension
+
+    torch.cuda.init()
+    ext = _load_pinned_extension()
+    # A rejected registration is a normal outcome -- host_banks.pin() catches it.
+    # The runtime latches the status until someone reads it, so if the extension
+    # raises without draining it, the next unrelated CUDA call reports this
+    # failure as its own and a caller that handled the exception dies anyway.
+    with pytest.raises(RuntimeError, match="cudaHostRegister failed"):
+        ext.host_register(0, 64)
+    torch.randn(4, device="cuda").sum().item()
+
+
 def test_host_bank_pin_registers_and_translates():
     if not torch.cuda.is_available():
         pytest.skip("needs CUDA")

From a1daa7aeea57935c566514a15eae03408831d3d8 Mon Sep 17 00:00:00 2001
From: Ruslan Rakhimov <risonyo@gmail.com>
Date: Mon, 24 Aug 2026 20:54:39 +0300
Subject: [PATCH 004/249] fix(kernel): drain in the deleter too, and keep the
 attribute probe non-fatal

Two corrections to the previous commit.

cudaFreeHost was the one call left uncovered, and it is the worst place to
leak a status: it runs as a from_blob deleter during GC, with no exception
to attribute the failure to. It drains unconditionally and never throws.

host_ptr_identity's two cudaDeviceGetAttribute calls previously discarded
their status with uva/reg pre-initialised to 0 -- the same idiom as
driver_cuda_version's "stays 0 when no driver is installed", i.e. a
deliberate fallback rather than an oversight. Raising there turned an
unqueryable attribute into a fatal error on the offload path, and
_host_ptr_identity is lru_cached, which does not cache exceptions, so it
would have re-raised on every call. Drain without raising and keep the 0
fallback: "no identity" routes device_ptr through host_device_ptr, the real
translation, which is correct everywhere.

Also soften the macro comment: these calls report whatever is latched on
the thread, so "every error below is non-sticky" was a stronger claim than
holds. Draining leaves the context usable for the failures these calls
originate; a sticky error latched elsewhere re-latches on the next call.
---
 .../freetoken/kernel/csrc/pinned_tensor.cpp   | 35 +++++++++++++------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp
index e2678d856e..d35e9b987c 100644
--- a/python/freetoken/kernel/csrc/pinned_tensor.cpp
+++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp
@@ -7,20 +7,27 @@ namespace {
 // A failed runtime call latches its status in the calling thread until someone
 // reads it. TORCH_CHECK reports the error but does not clear it, so the next
 // unrelated CUDA call -- torch's, in practice -- reports *this* failure instead
-// of its own. Drain it before raising: every error below is non-sticky, so the
-// context stays usable for a caller that handles the exception.
+// of its own. Drain it before raising. The failures these calls can originate
+// are non-sticky, so draining leaves the context usable; a sticky error latched
+// elsewhere will simply re-latch on the next call, which is correct.
 #define FT_CUDA_CHECK(expr, ...)                                               \
   do {                                                                         \
     const cudaError_t ft_err_ = (expr);                                        \
     if (ft_err_ != cudaSuccess) {                                              \
-      cudaGetLastError();                                                      \
+      [[maybe_unused]] const cudaError_t ft_drained_ = cudaGetLastError();      \
       TORCH_CHECK(false, __VA_ARGS__, cudaGetErrorString(ft_err_));            \
     }                                                                          \
   } while (0)
 
 void free_pinned(void *ptr) {
-  if (ptr != nullptr) {
-    cudaFreeHost(ptr);
+  if (ptr == nullptr) {
+    return;
+  }
+  // A from_blob deleter runs during GC, with no exception to attribute a failure
+  // to -- so this drains and never throws. Left latched it would be the worst
+  // case of the bug above: a stale status with no visible origin at all.
+  if (cudaFreeHost(ptr) != cudaSuccess) {
+    [[maybe_unused]] const cudaError_t drained = cudaGetLastError();
   }
 }
 
@@ -89,12 +96,20 @@ torch::Tensor alloc_pinned_tensor(std::vector<int64_t> sizes,
 bool host_ptr_identity() {
   int device = 0;
   FT_CUDA_CHECK(cudaGetDevice(&device), "cudaGetDevice failed: ");
+  // An unqueryable attribute stays 0 and answers "no identity", which is the safe
+  // answer: device_ptr() then goes through host_device_ptr(), the real translation,
+  // correct on every platform. Only the latch is new here -- do not raise, or an
+  // attribute query that used to degrade gracefully becomes fatal on the offload path.
   int uva = 0, reg = 0;
-  FT_CUDA_CHECK(cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device),
-                "cudaDeviceGetAttribute(UnifiedAddressing) failed: ");
-  FT_CUDA_CHECK(cudaDeviceGetAttribute(
-                    &reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device),
-                "cudaDeviceGetAttribute(CanUseHostPointerForRegisteredMem) failed: ");
+  if (cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device) != cudaSuccess) {
+    [[maybe_unused]] const cudaError_t drained = cudaGetLastError();
+    uva = 0;
+  }
+  if (cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem,
+                             device) != cudaSuccess) {
+    [[maybe_unused]] const cudaError_t drained = cudaGetLastError();
+    reg = 0;
+  }
   return uva == 1 && reg == 1;
 }
 

From 271420aee138c54c608c0f607be7bf309984185f Mon Sep 17 00:00:00 2001
From: pedro-moser <pedro.niederauer.moser@gmail.com>
Date: Tue, 25 Aug 2026 22:17:54 -0300
Subject: [PATCH 005/249] fix(engine): reserve explicit KV pages during MoE
 auto-sizing

---
 python/freetoken/engine/engine.py |  6 +++++-
 tests/engine/test_cache_budget.py | 29 ++++++++++++++++-------------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index cd6505d2d1..c53acef24f 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -490,7 +490,11 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int
             num_experts=num_experts,
             total_experts=total_experts,
             prefill_overlap=config.moe_prefill_overlap,
-            kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve),
+            kv_reserve_tokens=max(
+                config.kv_reserve_tokens,
+                min_reserve,
+                (config.num_page_override or 0) * page_tokens,
+            ),
             page_size=page_tokens,
             quant_format=banks.quant_format,
         )
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index a164f0b4d9..990ac06604 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -265,7 +265,7 @@ class tp_info:
     assert fixed == 0
 
 
-def test_engine_resolve_auto_moe_cache_size_maps_kwargs():
+def test_engine_resolve_auto_moe_cache_size_maps_kwargs(monkeypatch):
     import torch
 
     from freetoken.engine.engine import Engine
@@ -292,6 +292,7 @@ class StubConfig:
         memory_ratio = 0.9
         moe_prefill_overlap = True
         kv_reserve_tokens = 0
+        num_page_override = 64
         swa_full_tokens_ratio = 0.2
         swa_num_pages_override = None
         model_config = StubModelConfig()
@@ -314,21 +315,23 @@ class StubBanks:
     engine._weights_bytes = 1_000_000
     engine._pool_cls = MHAKVCache  # __init__ skipped -> install the generic pool family
 
-    size, pages, overlap = engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks())
+    captured = {}
 
-    # cross-check against the same pure functions, proving the kwarg mapping is faithful
-    from freetoken.engine.cache_budget import expert_bytes_per_slot, resolve_moe_cache_auto
-    from freetoken.kvcache.mha_pool import MHAKVCache
+    def fake_resolve(**kwargs):
+        captured.update(kwargs)
+        return 8, 64, True
 
-    cache_per_page, fixed, _, _ = MHAKVCache.kv_cost(StubConfig())
-    expected = resolve_moe_cache_auto(
-        baseline_free=10_000_000, weights_bytes=1_000_000, memory_ratio=0.9,
-        cache_per_page=cache_per_page, fixed_cache_size=fixed,
-        per_expert_bytes=expert_bytes_per_slot(StubBanks.sources),
-        num_experts=4, total_experts=8, prefill_overlap=True,
-        kv_reserve_tokens=0, page_size=16, quant_format="bf16",
+    monkeypatch.setattr(
+        "freetoken.engine.cache_budget.resolve_moe_cache_auto", fake_resolve
     )
-    assert (size, pages, overlap) == expected
+    got = engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks())
+
+    assert got == (8, 64, True)
+    assert captured["kv_reserve_tokens"] == 64 * 16
+    assert captured["page_size"] == 16
+    assert captured["num_experts"] == 4
+    assert captured["total_experts"] == 8
+    assert captured["per_expert_bytes"] == 512 + 256
 
 
 # ---------------------------------------------------------------------------

From 8856ae68f6a553a35b3f3b225a838833464ab995 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 26 Aug 2026 23:13:27 +0300
Subject: [PATCH 006/249] server: guarantee abort delivery on stream
 cancellation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

stream_with_cancellation reacted to a client disconnect with an unowned
asyncio.create_task(abort_user(uid)): request teardown could outrun delivery and a
failure inside the task degraded to a never-retrieved-exception warning, so the
scheduler kept decoding for a client that was gone. Await the abort inline behind
asyncio.shield (a second cancellation cannot kill the delivery task), and make
abort_user claim the uid first — exactly one AbortMsg even if cancellation runs
twice, and none at all when the stream already finished normally.

Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 2).
---
 python/freetoken/server/api_server.py    |  13 ++-
 tests/server/test_stream_cancellation.py | 120 +++++++++++++++++++++++
 2 files changed, 128 insertions(+), 5 deletions(-)
 create mode 100644 tests/server/test_stream_cancellation.py

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 3e2acc8542..add3c66050 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -364,15 +364,18 @@ async def stream_with_cancellation(self, generator, request: Request, uid: int):
                     raise asyncio.CancelledError
                 yield chunk
         except asyncio.CancelledError:
-            asyncio.create_task(self.abort_user(uid))
+            try:
+                await asyncio.shield(self.abort_user(uid))
+            except Exception:  # noqa: BLE001
+                logger.exception("Failed to deliver abort for user %s", uid)
             raise
 
     async def abort_user(self, uid: int):
+        claimed = self.event_map.pop(uid, None) is not None
+        self.ack_map.pop(uid, None)
+        if not claimed:
+            return
         await asyncio.sleep(0.1)
-        if uid in self.ack_map:
-            del self.ack_map[uid]
-        if uid in self.event_map:
-            del self.event_map[uid]
         self.stats.on_abort(uid)
         logger.warning("Aborting request for user %s", uid)
         await self.send_one(AbortMsg(uid=uid))
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
new file mode 100644
index 0000000000..120c504564
--- /dev/null
+++ b/tests/server/test_stream_cancellation.py
@@ -0,0 +1,120 @@
+"""Cancellation-path tests for FrontendManager.stream_with_cancellation / abort_user."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from freetoken.message import AbortMsg
+from freetoken.server.api_server import FrontendManager
+
+
+class _Stats:
+    def __init__(self):
+        self.aborts = []
+
+    def on_abort(self, uid):
+        self.aborts.append(uid)
+
+
+def _state(send_impl=None):
+    st = SimpleNamespace(
+        ack_map={7: [object()]},
+        event_map={7: asyncio.Event()},
+        stats=_Stats(),
+    )
+    sent = []
+
+    async def default_send(msg):
+        sent.append(msg)
+
+    st.send_one = send_impl or default_send
+    st.sent = sent
+    return st
+
+
+class _Request:
+    def __init__(self, disconnected=False):
+        self._disconnected = disconnected
+
+    async def is_disconnected(self):
+        return self._disconnected
+
+
+async def _consume(state, request, uid=7):
+    state.abort_user = lambda request_uid: FrontendManager.abort_user(state, request_uid)
+    async for _ in FrontendManager.stream_with_cancellation(state, _never(), request, uid):
+        pass
+
+
+async def _never():
+    await asyncio.sleep(3600)
+    yield b""  # pragma: no cover
+
+
+def test_cancellation_sends_one_abort_and_cleans_maps_inline():
+    async def run():
+        state = _state()
+        task = asyncio.create_task(_consume(state, _Request()))
+        await asyncio.sleep(0.01)
+        task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert asyncio.all_tasks() - {asyncio.current_task()} == set()
+        assert len(state.sent) == 1
+        assert isinstance(state.sent[0], AbortMsg)
+        assert state.sent[0].uid == 7
+        assert state.ack_map == {}
+        assert state.event_map == {}
+        assert state.stats.aborts == [7]
+
+    asyncio.run(run())
+
+
+def test_abort_delivery_failure_preserves_cancellation():
+    async def boom(msg):
+        raise RuntimeError("zmq down")
+
+    async def run():
+        state = _state(send_impl=boom)
+        task = asyncio.create_task(_consume(state, _Request()))
+        await asyncio.sleep(0.01)
+        task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+    asyncio.run(run())
+
+
+def test_abort_user_is_idempotent():
+    async def run():
+        state = _state()
+
+        await FrontendManager.abort_user(state, 7)
+        assert len(state.sent) == 1
+
+        await FrontendManager.abort_user(state, 7)
+        assert len(state.sent) == 1
+        assert state.stats.aborts == [7]
+
+    asyncio.run(run())
+
+
+def test_normal_completion_sends_no_abort():
+    async def run():
+        state = _state()
+
+        async def one():
+            yield b"data: x\n\n"
+
+        async for _ in FrontendManager.stream_with_cancellation(state, one(), _Request(), 7):
+            pass
+        assert state.sent == []
+        assert 7 in state.ack_map  # wait_for_ack owns normal-path cleanup, not the stream
+
+    asyncio.run(run())

From 494aa623a50e238d3e0244f2b859a2df4e96c322 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 26 Aug 2026 23:15:01 +0300
Subject: [PATCH 007/249] server: validate sampling params and rebuild timeout
 at the API layer

resolve_sampling forwarded arbitrary client floats into SamplingParams: negative,
NaN and inf temperature, top_p outside (0, 1], top_k 0 -- all reached the engine
and died (or misbehaved) deep in the sampler instead of failing the request.
Validate the RESOLVED values (checkpoint defaults included) and raise ValueError
in the style of the existing max_tokens check; every adapter already maps that to
a clean 400. CacheRebuildRequest.timeout gets the matching pydantic bound
(0 < t <= 3600) so a NaN/inf/negative timeout cannot reach asyncio.wait_for.

Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defects 3-4); the
rebuild TOCTOU race from that audit is already structurally fixed upstream, so
only the timeout bound is ported.
---
 python/freetoken/server/api_server.py    |  4 +-
 python/freetoken/server/generation.py    | 16 ++++--
 tests/server/test_openai_api.py          | 63 ++++++++++++++++++++++++
 tests/server/test_rebuild_maintenance.py | 15 ++++++
 4 files changed, 93 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 3e2acc8542..94b3e10efd 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -33,7 +33,7 @@
     init_logger,
     load_generation_sampling,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from .args import ServerArgs
 from .anthropic_api import register_anthropic_routes
@@ -494,7 +494,7 @@ class CacheRebuildRequest(BaseModel):
     # is deferred (needs the drain-gate machinery); constraining the Literal makes an
     # unsupported value fail fast with a 422 at the API layer instead of a generic 503.
     mode: Literal["if_idle"] = "if_idle"
-    timeout: float = 300.0
+    timeout: float = Field(default=300.0, gt=0, le=3600)
 
 
 async def dispatch_rebuild(
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index be05d908a0..a60119ecfb 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -15,6 +15,7 @@
 
 import asyncio
 import json
+import math
 import time
 from collections.abc import AsyncIterator
 from dataclasses import dataclass, field
@@ -177,12 +178,21 @@ def pick(value, key, framework):
     # non-positive value is a client error.
     if max_tokens is not None and max_tokens < 1:
         raise ValueError(f"max_tokens must be at least 1, got {max_tokens}")
+    resolved_temperature = pick(temperature, "temperature", 0.0)
+    if not math.isfinite(resolved_temperature) or resolved_temperature < 0:
+        raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}")
+    resolved_top_p = pick(top_p, "top_p", 1.0)
+    if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1:
+        raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}")
+    resolved_top_k = pick(top_k, "top_k", -1)
+    if resolved_top_k != -1 and resolved_top_k < 1:
+        raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}")
     return SamplingParams(
         ignore_eos=ignore_eos,
         max_tokens=DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens,
-        temperature=pick(temperature, "temperature", 0.0),
-        top_k=pick(top_k, "top_k", -1),
-        top_p=pick(top_p, "top_p", 1.0),
+        temperature=resolved_temperature,
+        top_k=resolved_top_k,
+        top_p=resolved_top_p,
         stop_strs=[s for s in stop_list if s],  # drop empty strings (would match everything)
     )
 
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index facd469b35..f6808a7d33 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -676,3 +676,66 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob(
     message = response["choices"][0]["message"]
     assert message["reasoning_content"] == "private thought"
     assert message["content"] == "visible answer"
+
+
+import pytest  # noqa: E402
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("temperature", -0.5),
+        ("top_p", 0.0),
+        ("top_k", 0),
+    ],
+)
+def test_chat_completion_rejects_invalid_sampling(field, value):
+    app = FastAPI()
+    state = FakeState([])
+
+    @app.post("/v1/chat/completions")
+    async def chat_completion(req: ChatCompletionRequest):
+        return await handle_chat_completion(req, request=None, state=state, model_sampling={})
+
+    body = {
+        "model": "unit-model",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream": True,
+        field: value,
+    }
+    response = TestClient(app).post("/v1/chat/completions", json=body)
+
+    assert response.status_code == 400
+    error = response.json()["error"]
+    assert error["type"] == "invalid_request_error"
+    assert field in error["message"]
+    assert state.sent is None
+
+
+@pytest.mark.parametrize(
+    "sampling",
+    [
+        {"temperature": 0.0, "top_p": 1.0, "top_k": -1},
+        {"temperature": 0.0, "top_p": 1.0, "top_k": 1},
+    ],
+)
+def test_chat_completion_accepts_sampling_boundaries(sampling):
+    app = FastAPI()
+    state = FakeState([])
+
+    @app.post("/v1/chat/completions")
+    async def chat_completion(req: ChatCompletionRequest):
+        return await handle_chat_completion(req, request=None, state=state, model_sampling={})
+
+    response = TestClient(app).post(
+        "/v1/chat/completions",
+        json={
+            "model": "unit-model",
+            "messages": [{"role": "user", "content": "hi"}],
+            "stream": True,
+            **sampling,
+        },
+    )
+
+    assert response.status_code == 200
+    assert state.sent is not None
diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py
index a1d10a7e56..0df3d6bfef 100644
--- a/tests/server/test_rebuild_maintenance.py
+++ b/tests/server/test_rebuild_maintenance.py
@@ -327,3 +327,18 @@ def test_cache_rebuild_request_rejects_unknown_mode():
     assert CacheRebuildRequest(mode="if_idle").mode == "if_idle"
     with pytest.raises(ValidationError):
         CacheRebuildRequest(mode="drain")
+
+
+from pydantic import ValidationError  # noqa: E402
+
+from freetoken.server.api_server import CacheRebuildRequest  # noqa: E402
+
+
+def test_rebuild_timeout_bounds():
+    import pytest
+
+    for bad in (0, -5, 3601, float("nan")):
+        with pytest.raises(ValidationError):
+            CacheRebuildRequest(timeout=bad)
+    assert CacheRebuildRequest(timeout=3600).timeout == 3600
+    assert CacheRebuildRequest().timeout == 300.0

From 58a252ed3f2411f4d2a8bf58e3be308657ade17f Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 26 Aug 2026 23:23:29 +0300
Subject: [PATCH 008/249] engine: compute sampled-token logprobs on request

Adds the engine half of OpenAI logprobs support. SamplingParams gains
logprobs/top_logprobs; the sampler computes log_softmax over the PRE-temperature
logits (raw model distribution, so temperature/top-k/top-p do not change reported
values), gathers the chosen token and batch-max top-k, and ships CPU copies
covered by the existing copy_done_event. The scheduler attaches per-request
values (cut to each request's own top_logprobs) to DetokenizeMsg; the detokenizer
builds a neutral entry (token text via single-id decode + UTF-8 bytes, so clients
can reassemble partial-UTF-8 pieces) onto UserReply.logprobs.

Zero cost when off: no row asked -> no mask tensor, no log_softmax, ForwardOutput
carries None. Message fields default to None, so old and new peers interoperate.
Stop-string trimming can hide final visible text; entries still cover every
sampled token.
---
 python/freetoken/core.py                 |   4 +
 python/freetoken/engine/engine.py        |  15 +-
 python/freetoken/engine/sample.py        |  84 +++++++++++-
 python/freetoken/message/frontend.py     |   3 +
 python/freetoken/message/tokenizer.py    |   5 +
 python/freetoken/scheduler/scheduler.py  |  21 ++-
 python/freetoken/tokenizer/detokenize.py |  29 ++++
 python/freetoken/tokenizer/server.py     |  15 +-
 tests/engine/test_sample_logprobs.py     | 167 +++++++++++++++++++++++
 tests/tokenizer/test_logprobs_entry.py   |  59 ++++++++
 10 files changed, 396 insertions(+), 6 deletions(-)
 create mode 100644 tests/engine/test_sample_logprobs.py
 create mode 100644 tests/tokenizer/test_logprobs_entry.py

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index ef0a539cbc..b001fc3e63 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -25,6 +25,10 @@ class SamplingParams:
     # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one
     # appears in the decoded output; the matched substring (and anything after) is trimmed.
     stop_strs: list[str] = field(default_factory=list)
+    # Sampled-token logprobs (OpenAI `logprobs`/`top_logprobs`): when on, the sampler
+    # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives.
+    logprobs: bool = False
+    top_logprobs: int = 0
 
     @property
     def is_greedy(self) -> bool:
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index cd6505d2d1..15121ba605 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -288,6 +288,11 @@ class ForwardOutput(NamedTuple):
     next_tokens_gpu: torch.Tensor
     next_tokens_cpu: torch.Tensor
     copy_done_event: torch.cuda.Event
+    # Sampled-token logprobs (None unless some request in the batch asked): CPU
+    # copies covered by copy_done_event, padded to the batch max top_logprobs.
+    chosen_logprobs_cpu: torch.Tensor | None = None
+    top_ids_cpu: torch.Tensor | None = None
+    top_logprobs_cpu: torch.Tensor | None = None
 
 
 class Engine:
@@ -930,9 +935,17 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         batch_logits = logits[: batch.size]
         next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32)
         next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True)
+        logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args)
         copy_done_event = torch.cuda.Event()
         copy_done_event.record(self.stream)
-        return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
+        if logprobs_out is None:
+            return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
+        chosen_logprobs, top_ids, top_logprobs = logprobs_out
+        return ForwardOutput(
+            next_tokens_gpu, next_tokens_cpu, copy_done_event,
+            chosen_logprobs_cpu=chosen_logprobs, top_ids_cpu=top_ids,
+            top_logprobs_cpu=top_logprobs,
+        )
 
     @torch.inference_mode()
     def _warmup_prefill(self) -> None:
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 01d14b1aab..3d5c27e15f 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -15,6 +15,8 @@ class BatchSamplingArgs:
     temperatures: torch.Tensor | None
     top_k: torch.Tensor | None = None
     top_p: torch.Tensor | None = None
+    logprob_rows: torch.Tensor | None = None
+    max_top_logprobs: int = 0
 
 
 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
@@ -57,8 +59,21 @@ class Sampler:
 
     def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
+        want_logprobs = [p.logprobs for p in params]
+        logprob_rows = (
+            make_device_tensor(want_logprobs, torch.bool, self.device)
+            if any(want_logprobs)
+            else None
+        )
         if all(p.is_greedy for p in params):
-            return BatchSamplingArgs(temperatures=None)
+            max_top_logprobs = max((p.top_logprobs for p in params if p.logprobs), default=0)
+            if max_top_logprobs > self.vocab_size:
+                max_top_logprobs = self.vocab_size
+            return BatchSamplingArgs(
+                temperatures=None,
+                logprob_rows=logprob_rows,
+                max_top_logprobs=max_top_logprobs,
+            )
 
         MIN_P = MIN_T = 1e-6
         ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params]
@@ -70,7 +85,16 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             top_k = make_device_tensor(top_ks, torch.int32, self.device)
         if any(p < 1.0 for p in top_ps):
             top_p = make_device_tensor(top_ps, torch.float32, self.device)
-        return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p)
+        max_top_logprobs = max((p.top_logprobs for p in params if p.logprobs), default=0)
+        if max_top_logprobs > self.vocab_size:
+            max_top_logprobs = self.vocab_size
+        return BatchSamplingArgs(
+            temperatures,
+            top_k=top_k,
+            top_p=top_p,
+            logprob_rows=logprob_rows,
+            max_top_logprobs=max_top_logprobs,
+        )
 
     @nvtx_annotate("Sampler")
     def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
@@ -78,3 +102,59 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
             if args.temperatures is None:  # greedy sampling
                 return torch.argmax(logits, dim=-1)
             return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
+
+    def compute_logprobs(
+        self,
+        logits: torch.Tensor,
+        sampled_tokens: torch.Tensor,
+        args: BatchSamplingArgs,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        if args.logprob_rows is None:
+            return None
+
+        requested_rows = torch.nonzero(args.logprob_rows, as_tuple=False).flatten()
+        if requested_rows.numel() == 0:
+            return None
+
+        request_logits = logits.index_select(0, requested_rows).float()
+        # Reported values are raw model logprobs (pre-temperature log_softmax over logits).
+        request_logprobs = torch.log_softmax(request_logits, dim=-1)
+
+        request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device).index_select(
+            0, requested_rows
+        )
+        request_row_idx = torch.arange(requested_rows.numel(), device=logits.device)
+        request_chosen_logprobs = request_logprobs[request_row_idx, request_tokens]
+
+        chosen_logprobs = torch.full(
+            (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device
+        )
+        chosen_logprobs.index_copy_(0, requested_rows, request_chosen_logprobs)
+
+        if args.max_top_logprobs > 0:
+            request_top_logprobs, request_top_ids = torch.topk(
+                request_logprobs, k=args.max_top_logprobs, dim=-1
+            )
+            top_ids = torch.full(
+                (logits.shape[0], args.max_top_logprobs),
+                -1,
+                dtype=torch.int32,
+                device=logits.device,
+            )
+            top_logprobs = torch.full(
+                (logits.shape[0], args.max_top_logprobs),
+                float("-inf"),
+                dtype=torch.float32,
+                device=logits.device,
+            )
+            top_ids[requested_rows] = request_top_ids.to(torch.int32)
+            top_logprobs[requested_rows] = request_top_logprobs
+        else:
+            top_ids = torch.empty((logits.shape[0], 0), dtype=torch.int32, device=logits.device)
+            top_logprobs = torch.empty((logits.shape[0], 0), dtype=torch.float32, device=logits.device)
+
+        return (
+            chosen_logprobs.to("cpu", non_blocking=True),
+            top_ids.to("cpu", non_blocking=True),
+            top_logprobs.to("cpu", non_blocking=True),
+        )
diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py
index 24725567b9..3da582cf66 100644
--- a/python/freetoken/message/frontend.py
+++ b/python/freetoken/message/frontend.py
@@ -56,6 +56,9 @@ class UserReply(BaseFrontendMsg):
     finish_reason: str | None = None
     # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence').
     matched_stop: str | None = None
+    # Neutral sampled-token logprobs entry for this token (see
+    # tokenizer.detokenize.build_logprobs_entry); None when the request did not ask.
+    logprobs: dict | None = None
 
 
 @dataclass
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index 33b75c785f..f658c9cc44 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -36,6 +36,11 @@ class DetokenizeMsg(BaseTokenizerMsg):
     # The request's stop strings (None when it has none), so the detokenizer can hold back
     # a trailing partial-stop prefix instead of streaming it and then needing to retract.
     stop_strs: list[str] | None = None
+    # Sampled-token logprobs (None unless the request asked): the chosen token's raw
+    # logprob and the top alternatives, already cut to this request's top_logprobs.
+    chosen_logprob: float | None = None
+    top_ids: list[int] | None = None
+    top_logprobs: list[float] | None = None
     # KV page-pool usage snapshot at this step (not-evictable used/total), passed
     # through to the frontend for the shell status bar. 0/0 for owned-KV models.
     kv_used_pages: int = 0
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 48923e3b0a..a5a3c2fe19 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -303,8 +303,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
         if last_data is None:
             return
 
-        batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
-        copy_done.synchronize()
+        batch, outputs = last_data[0].batch, last_data[1]
+        next_tokens_cpu = outputs.next_tokens_cpu
+        outputs.copy_done_event.synchronize()
         reply: List[DetokenizeMsg] = []
         new_finished_reqs: Set[Req] = set()
         with self.cache_manager.lazy_free_region():
@@ -337,6 +338,19 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                 next_token = next_tokens_cpu[i]
                 req.append_host(next_token.unsqueeze(0))
                 next_token = int(next_token.item())
+
+                row_chosen_logprob: float | None = None
+                row_top_ids: list[int] | None = None
+                row_top_logprobs: list[float] | None = None
+                if req.sampling_params.logprobs and outputs.chosen_logprobs_cpu is not None:
+                    row_chosen_logprob = float(outputs.chosen_logprobs_cpu[i].item())
+                    requested_top = req.sampling_params.top_logprobs
+                    if requested_top > 0 and outputs.top_ids_cpu is not None:
+                        row_top_ids = [int(t) for t in outputs.top_ids_cpu[i, :requested_top].tolist()]
+                        row_top_logprobs = outputs.top_logprobs_cpu[i, :requested_top].tolist()
+                    else:
+                        row_top_ids = []
+                        row_top_logprobs = []
                 # EOS / stop-string -> "stop", output budget exhausted -> "length";
                 # EOS and stop strings win over length.
                 hit_length = not req.can_decode
@@ -368,6 +382,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         finish_reason=finish_reason,
                         matched_stop=matched_stop,
                         stop_strs=req.sampling_params.stop_strs or None,
+                        chosen_logprob=row_chosen_logprob,
+                        top_ids=row_top_ids,
+                        top_logprobs=row_top_logprobs,
                     )
                 )
 
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index d138ac1a68..1505a56c59 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -145,3 +145,32 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
                 del self.decode_map[msg.uid]
 
         return incremental_strs
+
+
+def build_logprobs_entry(
+    tokenizer: PreTrainedTokenizerBase,
+    token_id: int,
+    chosen_logprob: float,
+    top_ids: list[int] | None,
+    top_logprobs: list[float] | None,
+) -> dict:
+    """Neutral sampled-token logprobs entry (UserReply.logprobs). Values are the raw
+    pre-temperature logprobs the sampler computed; token text comes from a single-id
+    decode, with UTF-8 bytes alongside so clients can reassemble partial-UTF-8 pieces."""
+
+    def _fields(tid: int) -> tuple[str, list[int]]:
+        text = tokenizer.decode([tid])
+        return text, list(text.encode("utf-8"))
+
+    token, token_bytes = _fields(token_id)
+    top = []
+    for tid, logprob in zip(top_ids or [], top_logprobs or []):
+        text, data = _fields(int(tid))
+        top.append({"token_id": int(tid), "token": text, "bytes": data, "logprob": float(logprob)})
+    return {
+        "token_id": token_id,
+        "token": token,
+        "bytes": token_bytes,
+        "logprob": float(chosen_logprob),
+        "top": top,
+    }
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index 530e862d04..d4f6a03ac1 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -144,7 +144,7 @@ def tokenize_worker(
     tokenizer = load_tokenizer(tokenizer_path)
     logger = init_logger(__name__, f"tokenizer_{tokenizer_id}")
 
-    from .detokenize import DetokenizeManager
+    from .detokenize import DetokenizeManager, build_logprobs_entry
     from .tokenize import TokenizeManager
 
     tokenize_manager = TokenizeManager(tokenizer)
@@ -223,6 +223,19 @@ def tokenize_worker(
                         swa_used_tokens=msg.swa_used_tokens,
                         swa_total_tokens=msg.swa_total_tokens,
                         gpu_mem_bytes=msg.gpu_mem_bytes,
+                        # Stop-string trimming can hide final visible text; keep one logprob
+                        # entry per sampled token.
+                        logprobs=(
+                            build_logprobs_entry(
+                                detokenize_manager.tokenizer,
+                                msg.next_token,
+                                msg.chosen_logprob,
+                                msg.top_ids,
+                                msg.top_logprobs,
+                            )
+                            if msg.chosen_logprob is not None
+                            else None
+                        ),
                     )
                     for msg, reply in zip(detokenize_msg, replies, strict=True)
                 ]
diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
new file mode 100644
index 0000000000..aabbdc5fe7
--- /dev/null
+++ b/tests/engine/test_sample_logprobs.py
@@ -0,0 +1,167 @@
+from types import SimpleNamespace
+
+import torch
+
+from freetoken.engine.sample import BatchSamplingArgs, Sampler
+
+
+def test_compute_logprobs_matches_raw_log_softmax_and_sorted_top() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=4)
+
+    logits = torch.tensor(
+        [
+            [2.0, 0.0, 1.0, -1.0],
+            [0.0, -1.0, 1.0, 3.0],
+            [1.0, 2.0, 3.0, 4.0],
+        ],
+        dtype=torch.float32,
+    )
+    sampled_tokens = torch.tensor([2, 3, 0], dtype=torch.long)
+    args = BatchSamplingArgs(
+        temperatures=torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32),
+        logprob_rows=torch.tensor([True, False, True], dtype=torch.bool),
+        max_top_logprobs=3,
+    )
+
+    result = sampler.compute_logprobs(logits, sampled_tokens, args)
+
+    assert result is not None
+    chosen_logprobs, top_ids, top_logprobs = result
+    expected = torch.log_softmax(logits.float(), dim=-1)
+
+    assert torch.isclose(chosen_logprobs[0], expected[0, sampled_tokens[0]])
+    assert torch.isnan(chosen_logprobs[1])
+    assert torch.isclose(chosen_logprobs[2], expected[2, sampled_tokens[2]])
+
+    expected_top0 = torch.topk(expected[0], k=3)
+    expected_top2 = torch.topk(expected[2], k=3)
+    assert torch.equal(top_ids[0], expected_top0.indices)
+    assert torch.allclose(top_logprobs[0], expected_top0.values)
+    assert torch.equal(top_ids[2], expected_top2.indices)
+    assert torch.allclose(top_logprobs[2], expected_top2.values)
+
+    assert torch.equal(top_ids[1], torch.full((3,), -1, dtype=torch.int32))
+    assert torch.isneginf(top_logprobs[1]).all()
+    assert top_ids.shape == (3, 3)
+    assert top_logprobs.shape == (3, 3)
+
+
+def test_compute_logprobs_ignores_temperatures() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=4)
+
+    logits = torch.tensor(
+        [
+            [0.5, 1.0, 2.0, 3.0],
+            [4.0, 3.0, 2.0, 1.0],
+        ],
+        dtype=torch.float32,
+    )
+    sampled_tokens = torch.tensor([1, 2], dtype=torch.long)
+    logprob_rows = torch.tensor([True, True], dtype=torch.bool)
+
+    cold = BatchSamplingArgs(
+        temperatures=torch.full((2,), 0.3, dtype=torch.float32),
+        logprob_rows=logprob_rows,
+        max_top_logprobs=2,
+    )
+    hot = BatchSamplingArgs(
+        temperatures=torch.full((2,), 2.5, dtype=torch.float32),
+        logprob_rows=logprob_rows,
+        max_top_logprobs=2,
+    )
+
+    cold_out = sampler.compute_logprobs(logits, sampled_tokens, cold)
+    hot_out = sampler.compute_logprobs(logits, sampled_tokens, hot)
+
+    assert cold_out is not None and hot_out is not None
+    cold_chosen, cold_top_ids, cold_top_logprobs = cold_out
+    hot_chosen, hot_top_ids, hot_top_logprobs = hot_out
+
+    assert torch.allclose(cold_chosen, hot_chosen)
+    assert torch.equal(cold_top_ids, hot_top_ids)
+    assert torch.allclose(cold_top_logprobs, hot_top_logprobs)
+
+
+def test_compute_logprobs_returns_none_without_requested_rows() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=4)
+    logits = torch.zeros((2, 4), dtype=torch.float32)
+    sampled_tokens = torch.tensor([0, 1], dtype=torch.long)
+    args = BatchSamplingArgs(
+        temperatures=torch.ones(2),
+        logprob_rows=torch.zeros(2, dtype=torch.bool),
+        max_top_logprobs=2,
+    )
+
+    assert sampler.compute_logprobs(logits, sampled_tokens, args) is None
+
+
+def test_compute_logprobs_handles_zero_top_logprobs() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=4)
+    logits = torch.tensor(
+        [
+            [1.0, 2.0, 3.0, 4.0],
+            [4.0, 3.0, 2.0, 1.0],
+            [0.1, 0.2, 0.3, 0.4],
+        ],
+        dtype=torch.float32,
+    )
+    sampled_tokens = torch.tensor([3, 0, 1], dtype=torch.long)
+    args = BatchSamplingArgs(
+        temperatures=torch.full((3,), 1.0),
+        logprob_rows=torch.tensor([True, False, True], dtype=torch.bool),
+        max_top_logprobs=0,
+    )
+
+    result = sampler.compute_logprobs(logits, sampled_tokens, args)
+
+    assert result is not None
+    chosen_logprobs, top_ids, top_logprobs = result
+    expected = torch.log_softmax(logits, dim=-1)
+    assert torch.isclose(chosen_logprobs[0], expected[0, 3])
+    assert torch.isnan(chosen_logprobs[1])
+    assert torch.isclose(chosen_logprobs[2], expected[2, 1])
+    assert top_ids.shape == (3, 0)
+    assert top_logprobs.shape == (3, 0)
+
+
+def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatch) -> None:
+    # prepare() pins host tensors for the H2D copy; pinning needs a CUDA context,
+    # so stub the transfer helper to keep this test host-agnostic.
+    import freetoken.engine.sample as sample_mod
+
+    monkeypatch.setattr(
+        sample_mod, "make_device_tensor",
+        lambda data, dtype, device: torch.tensor(data, dtype=dtype),
+    )
+    sampler = Sampler(torch.device("cpu"), vocab_size=5)
+
+    batch = SimpleNamespace(
+        reqs=[
+            SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)),
+            SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)),
+        ]
+    )
+    args = sampler.prepare(batch)
+    assert args.max_top_logprobs == sampler.vocab_size
+
+    logits = torch.tensor(
+        [
+            [1.0, 0.0, -1.0, 0.5, 2.0],
+            [2.0, 1.0, 0.0, -1.0, -2.0],
+        ],
+        dtype=torch.float32,
+    )
+    sampled_tokens = torch.tensor([4, 0], dtype=torch.long)
+
+    result = sampler.compute_logprobs(logits, sampled_tokens, args)
+    assert result is not None
+    chosen_logprobs, top_ids, top_logprobs = result
+
+    assert chosen_logprobs.shape == (2,)
+    assert torch.isclose(chosen_logprobs[0], torch.log_softmax(logits[0], dim=-1)[4])
+    assert torch.isnan(chosen_logprobs[1])
+
+    assert top_ids.shape == (2, sampler.vocab_size)
+    assert top_logprobs.shape == (2, sampler.vocab_size)
+    assert torch.equal(top_ids[1], torch.full((sampler.vocab_size,), -1, dtype=torch.int32))
+    assert torch.isneginf(top_logprobs[1]).all()
diff --git a/tests/tokenizer/test_logprobs_entry.py b/tests/tokenizer/test_logprobs_entry.py
new file mode 100644
index 0000000000..89cf35cb5e
--- /dev/null
+++ b/tests/tokenizer/test_logprobs_entry.py
@@ -0,0 +1,59 @@
+from freetoken.tokenizer.detokenize import build_logprobs_entry
+
+
+class _FakeTokenizer:
+    def __init__(self, token_table: dict[int, str]) -> None:
+        self.token_table = token_table
+
+    def decode(self, token_ids: list[int]) -> str:
+        return "".join(self.token_table[token_id] for token_id in token_ids)
+
+
+def test_build_logprobs_entry_matches_contract() -> None:
+    tokenizer = _FakeTokenizer({1: "hel", 2: "lo", 3: "é"})
+    entry = build_logprobs_entry(
+        tokenizer,
+        token_id=3,
+        chosen_logprob=-0.123,
+        top_ids=[3, 2, 1],
+        top_logprobs=[-0.1, -0.9, -1.5],
+    )
+
+    assert entry["token_id"] == 3
+    assert entry["token"] == "é"
+    assert entry["bytes"] == list("é".encode("utf-8"))
+    assert entry["logprob"] == -0.123
+
+    top = entry["top"]
+    assert len(top) == 3
+    assert top[0] == {
+        "token_id": 3,
+        "token": "é",
+        "bytes": list("é".encode("utf-8")),
+        "logprob": -0.1,
+    }
+    assert top[1] == {
+        "token_id": 2,
+        "token": "lo",
+        "bytes": list("lo".encode("utf-8")),
+        "logprob": -0.9,
+    }
+    assert top[2] == {
+        "token_id": 1,
+        "token": "hel",
+        "bytes": list("hel".encode("utf-8")),
+        "logprob": -1.5,
+    }
+
+
+def test_build_logprobs_entry_empty_top() -> None:
+    tokenizer = _FakeTokenizer({7: "x", 8: "y"})
+    entry = build_logprobs_entry(
+        tokenizer,
+        token_id=7,
+        chosen_logprob=-2.0,
+        top_ids=[],
+        top_logprobs=[],
+    )
+
+    assert entry["top"] == []

From ea7ae6d6d718ecd2c648474145aab8cf72f5148a Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 26 Aug 2026 23:25:35 +0300
Subject: [PATCH 009/249] server: OpenAI logprobs for chat and legacy
 completions

API half of logprobs support. Chat: logprobs + top_logprobs (0..20) yield
choice.logprobs.content entries {token, logprob, bytes, top_logprobs[]},
streaming and non-streaming. Completions: the legacy integer field (0..5) yields
{tokens, token_logprobs, top_logprobs, text_offset} with absolute offsets across
the stream; echo+logprobs stays rejected (prompt logprobs need prefill logits
and are out of scope here).

The protocol-neutral event layer carries entries on ContentDelta (a list -- parser
buffering can release several tokens' text in one delta); entries always
accumulate on GenResult for the non-streaming path, and reasoning/tool-call
buffering carries pending entries onto the next content delta. Formatting lives
in server/logprobs.py; wire compatibility follows the OpenAI shapes.
---
 python/freetoken/server/api_models.py |   4 +
 python/freetoken/server/generation.py |  33 +++-
 python/freetoken/server/logprobs.py   |  60 ++++++
 python/freetoken/server/openai_api.py |  86 +++++---
 tests/server/test_logprobs_api.py     | 274 ++++++++++++++++++++++++++
 5 files changed, 423 insertions(+), 34 deletions(-)
 create mode 100644 python/freetoken/server/logprobs.py
 create mode 100644 tests/server/test_logprobs_api.py

diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index ffd7172802..7bd43e642e 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -76,6 +76,10 @@ class ChatCompletionRequest(BaseModel):
     stop: str | list[str] | None = None
     presence_penalty: float = 0.0
     frequency_penalty: float = 0.0
+    # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires
+    # logprobs=true; entries cover generated tokens only (no prompt logprobs).
+    logprobs: bool = False
+    top_logprobs: int | None = None
     chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)
     reasoning_effort: str | None = None
     # DeepSeek-wire thinking toggle ({"type": "enabled"|"disabled"}). Any so a
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index be05d908a0..e4627a06b5 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -70,6 +70,10 @@ class ReasoningDelta:
 @dataclass
 class ContentDelta:
     text: str
+    # Neutral sampled-token logprob entries riding this delta (see UserReply.logprobs);
+    # None when the request did not ask. Parser buffering can attach several entries
+    # to one delta.
+    logprobs: list[dict] | None = None
 
 
 @dataclass
@@ -126,6 +130,9 @@ class GenResult:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    # Neutral sampled-token logprob entries, one per sampled token (empty when the
+    # request did not ask).
+    logprobs: list[dict] = field(default_factory=list)
 
 
 @dataclass
@@ -559,10 +566,17 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte
     completion_tokens = 0
     cached_tokens = 0
     pending = ""
+    pending_logprobs: list[dict] = []
     parse_tools = spec.parse_tools
     reasoning_parser = _make_reasoning_parser(spec, state)
     specials = _leaked_special_tokens(state)
 
+    def _content_delta(text: str) -> ContentDelta:
+        nonlocal pending_logprobs
+        logprobs = pending_logprobs or None
+        pending_logprobs = []
+        return ContentDelta(text, logprobs=logprobs)
+
     tool_parser: FunctionCallParser | None = None
     if parse_tools:
         try:
@@ -616,7 +630,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
                     out.append(done)
                 stripped = strip_special_tokens(payload, specials)
                 if stripped and not (stripped.strip() == "" and suppress_ws):
-                    out.append(ContentDelta(stripped))
+                    out.append(_content_delta(stripped))
                     if stripped.strip():
                         suppress_ws = False
                 continue
@@ -653,6 +667,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
+        if getattr(ack, "logprobs", None) is not None:
+            pending_logprobs.append(ack.logprobs)
         content_delta = ack.incremental_output
         if reasoning_parser is not None and content_delta:
             reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta)
@@ -667,7 +683,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += content_delta
             else:
-                yield ContentDelta(strip_special_tokens(content_delta, specials))
+                yield _content_delta(strip_special_tokens(content_delta, specials))
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
@@ -688,7 +704,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += flush_content
             else:
-                yield ContentDelta(strip_special_tokens(flush_content, specials))
+                yield _content_delta(strip_special_tokens(flush_content, specials))
 
     # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins.
     finish_reason = engine_finish_reason or "stop"
@@ -716,7 +732,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         if residual:
             stripped = strip_special_tokens(residual, specials)
             if stripped and not (stripped.strip() == "" and suppress_ws):
-                yield ContentDelta(stripped)
+                yield _content_delta(stripped)
         if calls_emitted and finish_reason != "length":
             finish_reason = "tool_calls"
     else:
@@ -725,13 +741,14 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             normal_text, tool_calls = parsed
             normal_text = strip_special_tokens(normal_text, specials)
             if normal_text:
-                yield ContentDelta(normal_text)
+                yield _content_delta(normal_text)
             yield ToolCallsDelta(tool_calls)
             if finish_reason != "length":
                 finish_reason = "tool_calls"
         elif parse_tools and pending:
-            yield ContentDelta(strip_special_tokens(pending, specials))
+            yield _content_delta(strip_special_tokens(pending, specials))
 
+    # Entries without a content delta are intentionally dropped in streaming mode.
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
         matched_stop=engine_matched_stop, cached_tokens=cached_tokens,
@@ -742,6 +759,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
     """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse
     tool calls, strip special tokens. The adapters format the GenResult into their wire."""
     full_content = ""
+    logprob_entries: list[dict] = []
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
@@ -754,6 +772,8 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         full_content += ack.incremental_output
+        if getattr(ack, "logprobs", None) is not None:
+            logprob_entries.append(ack.logprobs)
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
@@ -779,4 +799,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens=completion_tokens,
         matched_stop=engine_matched_stop,
         cached_tokens=cached_tokens,
+        logprobs=logprob_entries,
     )
diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py
new file mode 100644
index 0000000000..8618569544
--- /dev/null
+++ b/python/freetoken/server/logprobs.py
@@ -0,0 +1,60 @@
+"""OpenAI logprobs formatting: neutral engine entries -> wire shapes."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def chat_logprobs_error(req: Any) -> str | None:
+    if req.top_logprobs is not None and not req.logprobs:
+        return "top_logprobs requires logprobs=true"
+    if req.top_logprobs is not None and not 0 <= req.top_logprobs <= 20:
+        return "top_logprobs must be between 0 and 20"
+    return None
+
+
+def completion_logprobs_error(req: Any) -> str | None:
+    if req.logprobs is not None and not 0 <= req.logprobs <= 5:
+        return "logprobs must be between 0 and 5"
+    if req.echo and req.logprobs is not None:
+        return "echo with logprobs is not supported"
+    return None
+
+
+def chat_content_entry(entry: dict) -> dict:
+    return {
+        "token": entry["token"],
+        "logprob": entry["logprob"],
+        "bytes": entry["bytes"],
+        "top_logprobs": [
+            {
+                "token": candidate["token"],
+                "logprob": candidate["logprob"],
+                "bytes": candidate["bytes"],
+            }
+            for candidate in entry["top"]
+        ],
+    }
+
+
+def completions_logprobs(entries: list[dict], start_offset: int = 0) -> dict:
+    tokens: list[str] = []
+    token_logprobs: list[float] = []
+    top_logprobs: list[dict[str, float]] = []
+    text_offset: list[int] = []
+    offset = start_offset
+
+    for entry in entries:
+        token = entry["token"]
+        tokens.append(token)
+        token_logprobs.append(entry["logprob"])
+        top_logprobs.append({candidate["token"]: candidate["logprob"] for candidate in entry["top"]})
+        text_offset.append(offset)
+        offset += len(token)
+
+    return {
+        "tokens": tokens,
+        "token_logprobs": token_logprobs,
+        "top_logprobs": top_logprobs,
+        "text_offset": text_offset,
+    }
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b4becd2631..82c32bb78a 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -22,6 +22,12 @@
 )
 from .function_call_parser import ToolCallItem
 from .request_logger import log_request
+from .logprobs import (
+    chat_content_entry,
+    chat_logprobs_error,
+    completion_logprobs_error,
+    completions_logprobs,
+)
 from .generation import (
     ContentDelta,
     GenDone,
@@ -66,17 +72,20 @@ def chat_request_to_genspec(
     thinking_type = _thinking_type(req)
     if req.reasoning_effort or thinking_type:
         ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type)
+    sampling_params = resolve_sampling(
+        temperature=req.temperature,
+        top_k=req.top_k,
+        top_p=req.top_p,
+        max_tokens=req.max_tokens,
+        ignore_eos=req.ignore_eos,
+        model_sampling=model_sampling,
+        stop=req.stop,
+    )
+    sampling_params.logprobs = req.logprobs
+    sampling_params.top_logprobs = req.top_logprobs or 0
     return GenSpec(
         messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]),
-        sampling_params=resolve_sampling(
-            temperature=req.temperature,
-            top_k=req.top_k,
-            top_p=req.top_p,
-            max_tokens=req.max_tokens,
-            ignore_eos=req.ignore_eos,
-            model_sampling=model_sampling,
-            stop=req.stop,
-        ),
+        sampling_params=sampling_params,
         chat_template_kwargs=ctk,
         template_tools=_tools_for_template(req),
         parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None),
@@ -154,6 +163,9 @@ async def handle_chat_completion(
         return create_error_response("function_call is not supported; use tools/tool_choice instead")
     if req.logit_bias is not None:
         return create_error_response("logit_bias is not supported")
+    logprobs_error = chat_logprobs_error(req)
+    if logprobs_error is not None:
+        return create_error_response(logprobs_error, param="top_logprobs")
     if _response_format_unsupported(req.response_format):
         return create_error_response(
             "response_format json_object/json_schema is not supported (no constrained decoding)",
@@ -208,18 +220,20 @@ async def handle_chat_completion(
     if result.tool_calls:
         message["tool_calls"] = _tool_calls_to_openai(result.tool_calls)
 
+    choice: dict[str, Any] = {
+        "index": 0,
+        "message": message,
+        "finish_reason": result.finish_reason,
+    }
+    if req.logprobs:
+        choice["logprobs"] = {"content": [chat_content_entry(e) for e in result.logprobs]}
+
     return {
         "id": f"chatcmpl-{uid}",
         "object": "chat.completion",
         "created": int(time.time()),
         "model": req.model,
-        "choices": [
-            {
-                "index": 0,
-                "message": message,
-                "finish_reason": result.finish_reason,
-            }
-        ],
+        "choices": [choice],
         "usage": _usage(
             result.prompt_tokens,
             result.completion_tokens,
@@ -272,13 +286,10 @@ async def stream_chat_completion_chunks(
                 )
             )
         elif isinstance(ev, ContentDelta):
-            yield _sse(
-                _chat_chunk(
-                    req,
-                    uid,
-                    [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}],
-                )
-            )
+            choice: dict[str, Any] = {"delta": {"content": ev.text}, "index": 0, "finish_reason": None}
+            if ev.logprobs:
+                choice["logprobs"] = {"content": [chat_content_entry(e) for e in ev.logprobs]}
+            yield _sse(_chat_chunk(req, uid, [choice]))
         elif isinstance(ev, ToolCallStart):
             open_tool = {
                 "index": tool_calls_sent,
@@ -412,6 +423,7 @@ async def handle_completion(
         uid = state.new_user()
         await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling)))
         text = ""
+        entries: list[dict] = []
         finish_reason = "stop"
         async for ack in state.wait_for_ack(uid):
             if getattr(ack, "error", None):
@@ -420,10 +432,17 @@ async def handle_completion(
             completion_tokens += ack.completion_tokens_delta
             cached_tokens += ack.cached_tokens
             text += ack.incremental_output
+            if req.logprobs is not None and ack.logprobs is not None:
+                entries.append(ack.logprobs)
             if ack.finished:
                 finish_reason = getattr(ack, "finish_reason", None) or "stop"
                 break
-        choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None})
+        choices.append({
+            "index": index,
+            "text": text,
+            "finish_reason": finish_reason,
+            "logprobs": completions_logprobs(entries) if req.logprobs is not None else None,
+        })
 
     return {
         "id": f"cmpl-{uuid.uuid4().hex}",
@@ -437,6 +456,7 @@ async def handle_completion(
 
 async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]:
     prompt_tokens = 0
+    text_offset = 0
     completion_tokens = 0
     cached_tokens = 0
     finish_reason = "stop"
@@ -460,11 +480,16 @@ async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any)
                             "text": ack.incremental_output,
                             "index": 0,
                             "finish_reason": None,
-                            "logprobs": None,
+                            "logprobs": (
+                                completions_logprobs([ack.logprobs], text_offset)
+                                if req.logprobs is not None and ack.logprobs is not None
+                                else None
+                            ),
                         }
                     ],
                 }
             )
+            text_offset += len(ack.incremental_output)
         if ack.finished:
             finish_reason = getattr(ack, "finish_reason", None) or "stop"
             break
@@ -518,7 +543,7 @@ def _resolve_sampling(
     req: ChatCompletionRequest | CompletionRequest,
     model_sampling: dict[str, Any],
 ) -> SamplingParams:
-    return resolve_sampling(
+    sampling_params = resolve_sampling(
         temperature=req.temperature,
         top_k=req.top_k,
         top_p=req.top_p,
@@ -527,6 +552,10 @@ def _resolve_sampling(
         model_sampling=model_sampling,
         stop=req.stop,
     )
+    if isinstance(req, CompletionRequest) and req.logprobs is not None:
+        sampling_params.logprobs = True
+        sampling_params.top_logprobs = req.logprobs
+    return sampling_params
 
 
 def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None:
@@ -627,8 +656,9 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool
 def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
     if _is_token_prompt(req.prompt):
         return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead"
-    if req.logprobs is not None:
-        return "logprobs is not supported"
+    logprobs_error = completion_logprobs_error(req)
+    if logprobs_error is not None:
+        return logprobs_error
     if req.echo:
         return "echo is not supported"
     if req.suffix is not None:
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
new file mode 100644
index 0000000000..3b4fd5e32a
--- /dev/null
+++ b/tests/server/test_logprobs_api.py
@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from types import SimpleNamespace
+
+from freetoken.message import TokenizeMsg, UserReply
+from freetoken.server.openai_api import (
+    ChatCompletionRequest,
+    CompletionRequest,
+    chat_request_to_genspec,
+    handle_chat_completion,
+    handle_completion,
+    stream_chat_completion_chunks,
+    stream_completion_chunks,
+)
+
+
+def run(coro):
+    return asyncio.run(coro)
+
+
+class FakeState:
+    def __init__(self, replies: list[UserReply], reasoning_parser: str | None = None) -> None:
+        self.config = SimpleNamespace(
+            model_path="/models/unit-model",
+            served_model_name="unit-model",
+            tool_call_parser="llama3",
+            reasoning_parser=reasoning_parser,
+        )
+        self.replies = replies
+        self.sent: TokenizeMsg | None = None
+
+    def new_user(self) -> int:
+        return 42
+
+    async def send_one(self, msg):
+        self.sent = msg
+
+    async def wait_for_ack(self, uid: int):
+        assert uid == 42
+        for reply in self.replies:
+            yield reply
+
+
+def reply(text: str, *, finished: bool = False, logprobs: dict | None = None) -> UserReply:
+    return UserReply(
+        uid=42,
+        incremental_output=text,
+        finished=finished,
+        prompt_tokens_delta=3 if not text else 0,
+        completion_tokens_delta=1 if text else 0,
+        cached_tokens=0,
+        logprobs=logprobs,
+    )
+
+
+def entry(token_id: int, token: str, logprob: float) -> dict:
+    return {
+        "token_id": token_id,
+        "token": token,
+        "bytes": list(token.encode("utf-8")),
+        "logprob": logprob,
+        "top": [
+            {
+                "token_id": token_id,
+                "token": token,
+                "bytes": list(token.encode("utf-8")),
+                "logprob": logprob,
+            },
+            {
+                "token_id": token_id + 1,
+                "token": "x",
+                "bytes": [120],
+                "logprob": logprob - 1,
+            },
+        ],
+    }
+
+
+def chat_request(**kwargs) -> ChatCompletionRequest:
+    payload = {
+        "model": "client-model",
+        "messages": [{"role": "user", "content": "hello"}],
+        "max_tokens": 8,
+    }
+    payload.update(kwargs)
+    return ChatCompletionRequest(**payload)
+
+
+def parse_sse(chunks: list[bytes]) -> list[dict | str]:
+    events: list[dict | str] = []
+    for chunk in chunks:
+        for line in chunk.decode().splitlines():
+            if line.startswith("data: "):
+                data = line.removeprefix("data: ")
+                events.append(data if data == "[DONE]" else json.loads(data))
+    return events
+
+
+def test_chat_non_stream_logprobs() -> None:
+    first = entry(1, "Hello", -0.1)
+    second = entry(2, "!", -0.2)
+    result = run(
+        handle_chat_completion(
+            chat_request(logprobs=True, top_logprobs=2),
+            None,
+            FakeState([reply("Hello", logprobs=first), reply("!", finished=True, logprobs=second)]),
+            {},
+        )
+    )
+
+    content = result["choices"][0]["logprobs"]["content"]
+    assert content == [
+        {
+            "token": "Hello",
+            "logprob": -0.1,
+            "bytes": [72, 101, 108, 108, 111],
+            "top_logprobs": [
+                {"token": "Hello", "logprob": -0.1, "bytes": [72, 101, 108, 108, 111]},
+                {"token": "x", "logprob": -1.1, "bytes": [120]},
+            ],
+        },
+        {
+            "token": "!",
+            "logprob": -0.2,
+            "bytes": [33],
+            "top_logprobs": [
+                {"token": "!", "logprob": -0.2, "bytes": [33]},
+                {"token": "x", "logprob": -1.2, "bytes": [120]},
+            ],
+        },
+    ]
+
+    without_logprobs = run(
+        handle_chat_completion(
+            chat_request(),
+            None,
+            FakeState([reply("Hello", finished=True, logprobs=first)]),
+            {},
+        )
+    )
+    assert without_logprobs["choices"][0].get("logprobs") is None
+
+
+def test_chat_stream_logprobs_follow_content_deltas() -> None:
+    first = entry(1, "Hello", -0.1)
+    state = FakeState(
+        [
+            reply("Hello", logprobs=first),
+            reply(" world", finished=True),
+        ]
+    )
+    req = chat_request(stream=True, logprobs=True, top_logprobs=2)
+    spec = chat_request_to_genspec(req, {})
+    events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec))))
+
+    content_choices = [
+        event["choices"][0]
+        for event in events
+        if isinstance(event, dict)
+        and event["choices"]
+        and event["choices"][0]["delta"].get("content")
+    ]
+    assert content_choices[0]["logprobs"]["content"][0]["token"] == "Hello"
+    assert "logprobs" not in content_choices[1]
+
+
+def test_completion_logprobs_non_stream_and_stream() -> None:
+    first = entry(1, "Hi", -0.1)
+    second = entry(2, "!", -0.2)
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=2, max_tokens=8)
+
+    result = run(
+        handle_completion(
+            req,
+            None,
+            FakeState([reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second)]),
+            {},
+        )
+    )
+    assert result["choices"][0]["logprobs"] == {
+        "tokens": ["Hi", "!"],
+        "token_logprobs": [-0.1, -0.2],
+        "top_logprobs": [
+            {"Hi": -0.1, "x": -1.1},
+            {"!": -0.2, "x": -1.2},
+        ],
+        "text_offset": [0, 2],
+    }
+
+    events = parse_sse(
+        run(
+            _collect(
+                stream_completion_chunks(
+                    42,
+                    CompletionRequest(
+                        model="client-model",
+                        prompt="hello",
+                        logprobs=2,
+                        max_tokens=8,
+                        stream=True,
+                    ),
+                    FakeState([reply("Hi", finished=True, logprobs=first)]),
+                )
+            )
+        )
+    )
+    chunk = next(event for event in events if isinstance(event, dict) and event["choices"][0]["text"])
+    assert chunk["choices"][0]["logprobs"] == {
+        "tokens": ["Hi"],
+        "token_logprobs": [-0.1],
+        "top_logprobs": [{"Hi": -0.1, "x": -1.1}],
+        "text_offset": [0],
+    }
+
+
+def test_logprobs_validation_errors() -> None:
+    chat_top_out_of_range = run(
+        handle_chat_completion(chat_request(logprobs=True, top_logprobs=25), None, FakeState([]), {})
+    )
+    assert chat_top_out_of_range.status_code == 400
+
+    chat_missing_flag = run(
+        handle_chat_completion(chat_request(top_logprobs=1), None, FakeState([]), {})
+    )
+    assert chat_missing_flag.status_code == 400
+
+    completion_out_of_range = run(
+        handle_completion(
+            CompletionRequest(model="client-model", prompt="hello", logprobs=7),
+            None,
+            FakeState([]),
+            {},
+        )
+    )
+    assert completion_out_of_range.status_code == 400
+
+    completion_echo = run(
+        handle_completion(
+            CompletionRequest(model="client-model", prompt="hello", echo=True, logprobs=1),
+            None,
+            FakeState([]),
+            {},
+        )
+    )
+    assert completion_echo.status_code == 400
+
+
+def test_reasoning_logprob_is_carried_to_next_content_delta() -> None:
+    reasoning_entry = entry(1, "thought", -0.1)
+    state = FakeState(
+        [
+            reply("<think>thought</think>", logprobs=reasoning_entry),
+            reply("answer", finished=True),
+        ],
+        reasoning_parser="qwen3",
+    )
+    req = chat_request(stream=True, logprobs=True, top_logprobs=2)
+    spec = chat_request_to_genspec(req, {})
+    events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec))))
+
+    content_choice = next(
+        event["choices"][0]
+        for event in events
+        if isinstance(event, dict)
+        and event["choices"]
+        and event["choices"][0]["delta"].get("content") == "answer"
+    )
+    assert content_choice["logprobs"]["content"][0]["token"] == "thought"
+
+
+async def _collect(iterator):
+    return [chunk async for chunk in iterator]

From 90be981114207b51f7f5364b0293f287c63b64dc Mon Sep 17 00:00:00 2001
From: HaileyStorm <HaileyCollet@gmail.com>
Date: Fri, 28 Aug 2026 07:38:25 -0600
Subject: [PATCH 010/249] fix(server): hide qwen transport markers from
 semantic responses

---
 python/freetoken/server/generation.py         | 367 ++++++++++++++--
 .../test_qwen_semantic_special_tokens.py      | 405 ++++++++++++++++++
 2 files changed, 741 insertions(+), 31 deletions(-)
 create mode 100644 tests/server/test_qwen_semantic_special_tokens.py

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index be05d908a0..cdd32048db 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -42,6 +42,250 @@
 )
 
 
+# Qwen's decoder intentionally keeps protocol markers visible so reasoning and
+# tool-call parsers can consume their (non-special-token) tags.  At the semantic
+# API waist, after those parsers, only transport/internal multimodal markers are
+# suppressed.  Grounding markers (object_ref/box/quad), think tags, tool tags,
+# and FIM/repository markers are deliberately absent from this list.
+_QWEN_SEMANTIC_HIDDEN_TOKENS = (
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_pad|>",
+    "<tts_pad>",
+    "<tts_text_bos>",
+    "<tts_text_eod>",
+    "<tts_text_bos_single>",
+)
+
+
+def _uses_qwen_semantic_protocol(state: Any) -> bool:
+    """Whether this server is configured for a Qwen semantic output dialect.
+
+    Parser selection is the stable frontend-visible family signal available in
+    ``ServerArgs``.  Either half is enough because users may explicitly disable
+    reasoning or omit tools for an individual request.
+    """
+    config = state.config
+    return (
+        getattr(config, "reasoning_parser", None) == "qwen3"
+        or getattr(config, "tool_call_parser", None) in {"qwen", "qwen25", "qwen3_coder"}
+    )
+
+
+class _SemanticSpecialTokenFilter:
+    """Request-local, chunk-safe suppression of semantic protocol markers.
+
+    A proper marker prefix is retained until the next detokenizer delta, so a
+    marker split at any byte boundary cannot leak. A marker inside a quoted/code
+    literal is retained only once the matching closer arrives; an unclosed quote
+    is sanitized at end-of-stream. Structured events can be held inside either
+    ambiguity so their ordering relative to surrounding content stays exact.
+    """
+
+    def __init__(self, tokens: tuple[str, ...]) -> None:
+        self._tokens = tokens
+        self._quote: str | None = None
+        self._escaped = False
+        self._last_char = ""
+        self._candidate_text = ""
+        self._candidate_items: list[Any] = []
+        self._candidate_quoted = False
+        self._literal_items: list[Any] | None = None
+
+    @staticmethod
+    def _append_item(items: list[Any], item: Any) -> None:
+        if item == "":
+            return
+        if isinstance(item, str) and items and isinstance(items[-1], str):
+            items[-1] += item
+        else:
+            items.append(item)
+
+    @classmethod
+    def _strip_items(cls, items: list[Any], tokens: tuple[str, ...]) -> list[Any]:
+        """Strip markers across text segments while preserving event positions."""
+        out: list[Any] = []
+        candidate_text = ""
+        candidate_items: list[Any] = []
+
+        def flush(*, keep_text: bool) -> None:
+            nonlocal candidate_text, candidate_items
+            for pending in candidate_items:
+                if keep_text or not isinstance(pending, str):
+                    cls._append_item(out, pending)
+            candidate_text = ""
+            candidate_items = []
+
+        for item in items:
+            if not isinstance(item, str):
+                if candidate_text:
+                    cls._append_item(candidate_items, item)
+                else:
+                    cls._append_item(out, item)
+                continue
+            index = 0
+            while index < len(item):
+                ch = item[index]
+                if not candidate_text:
+                    if any(token.startswith(ch) for token in tokens):
+                        candidate_text = ch
+                        cls._append_item(candidate_items, ch)
+                    else:
+                        cls._append_item(out, ch)
+                    index += 1
+                    continue
+                extended = candidate_text + ch
+                if extended in tokens:
+                    candidate_text = extended
+                    cls._append_item(candidate_items, ch)
+                    flush(keep_text=False)
+                    index += 1
+                elif any(token.startswith(extended) for token in tokens):
+                    candidate_text = extended
+                    cls._append_item(candidate_items, ch)
+                    index += 1
+                else:
+                    flush(keep_text=True)
+        flush(keep_text=True)
+        return out
+
+    def _start_candidate(self, ch: str, *, quoted: bool) -> None:
+        self._candidate_text = ch
+        self._candidate_items = [ch]
+        self._candidate_quoted = quoted
+
+    def _extend_candidate(self, ch: str) -> tuple[list[Any], bool]:
+        extended = self._candidate_text + ch
+        if extended in self._tokens:
+            self._candidate_text = extended
+            self._append_item(self._candidate_items, ch)
+            if self._candidate_quoted:
+                self._literal_items = self._candidate_items
+                out: list[Any] = []
+            else:
+                out = [item for item in self._candidate_items if not isinstance(item, str)]
+            self._candidate_text = ""
+            self._candidate_items = []
+            self._candidate_quoted = False
+            return out, True
+        if any(token.startswith(extended) for token in self._tokens):
+            self._candidate_text = extended
+            self._append_item(self._candidate_items, ch)
+            return [], True
+        out = self._candidate_items
+        self._last_char = self._candidate_text[-1]
+        self._candidate_text = ""
+        self._candidate_items = []
+        self._candidate_quoted = False
+        return out, False
+
+    def feed_items(self, text: str, *, final: bool = False) -> list[Any]:
+        if not self._tokens:
+            return [text] if text else []
+        out: list[Any] = []
+        index = 0
+        while index < len(text):
+            ch = text[index]
+            if self._literal_items is not None:
+                self._append_item(self._literal_items, ch)
+                if self._escaped:
+                    self._escaped = False
+                elif ch == "\\":
+                    self._escaped = True
+                elif ch == self._quote:
+                    out.extend(self._literal_items)
+                    self._literal_items = None
+                    self._quote = None
+                self._last_char = ch
+                index += 1
+                continue
+
+            if self._candidate_text:
+                emitted, consumed = self._extend_candidate(ch)
+                out.extend(emitted)
+                if consumed:
+                    index += 1
+                continue
+
+            if self._quote is not None:
+                if any(token.startswith(ch) for token in self._tokens):
+                    self._escaped = False
+                    self._start_candidate(ch, quoted=True)
+                    index += 1
+                    continue
+                if self._escaped:
+                    self._escaped = False
+                elif ch == "\\":
+                    self._escaped = True
+                elif ch == self._quote:
+                    self._quote = None
+                self._append_item(out, ch)
+                self._last_char = ch
+                index += 1
+                continue
+
+            opens_quote = ch in {'"', "`", "“", "‘"} or (
+                ch == "'" and not self._last_char.isalnum()
+            )
+            if opens_quote:
+                self._quote = {"“": "”", "‘": "’"}.get(ch, ch)
+                self._append_item(out, ch)
+            elif any(token.startswith(ch) for token in self._tokens):
+                self._start_candidate(ch, quoted=False)
+                index += 1
+                continue
+            else:
+                self._append_item(out, ch)
+            self._last_char = ch
+            index += 1
+
+        if final:
+            if self._literal_items is not None:
+                out.extend(self._strip_items(self._literal_items, self._tokens))
+                self._literal_items = None
+            elif self._candidate_text:
+                out.extend(self._candidate_items)
+            self._candidate_text = ""
+            self._candidate_items = []
+            self._candidate_quoted = False
+            self._quote = None
+            self._escaped = False
+        return out
+
+    def feed(self, text: str, *, final: bool = False) -> str:
+        items = self.feed_items(text, final=final)
+        assert all(isinstance(item, str) for item in items)
+        return "".join(items)
+
+    def boundary(self, item: Any) -> list[Any]:
+        if self._literal_items is not None:
+            self._append_item(self._literal_items, item)
+            return []
+        if self._candidate_text:
+            self._append_item(self._candidate_items, item)
+            return []
+        return [item]
+
+    def finish(self) -> str:
+        return self.feed("", final=True)
+
+    def finish_items(self) -> list[Any]:
+        return self.feed_items("", final=True)
+
+
+def _make_semantic_special_token_filter(state: Any) -> _SemanticSpecialTokenFilter:
+    tokens = _QWEN_SEMANTIC_HIDDEN_TOKENS if _uses_qwen_semantic_protocol(state) else ()
+    return _SemanticSpecialTokenFilter(tokens)
+
+
 class GenerationError(Exception):
     """A request failed before producing output (surfaced via ``UserReply.error`` — e.g. a
     chat template the tokenizer cannot render, or a prompt that exceeds the KV budget). Each
@@ -562,6 +806,14 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte
     parse_tools = spec.parse_tools
     reasoning_parser = _make_reasoning_parser(spec, state)
     specials = _leaked_special_tokens(state)
+    reasoning_special_filter = _make_semantic_special_token_filter(state)
+    content_special_filter = _make_semantic_special_token_filter(state)
+
+    def _filter_reasoning(piece: str) -> str:
+        return reasoning_special_filter.feed(strip_special_tokens(piece, specials))
+
+    def _filter_content(piece: str) -> str:
+        return content_special_filter.feed(strip_special_tokens(piece, specials))
 
     tool_parser: FunctionCallParser | None = None
     if parse_tools:
@@ -583,6 +835,28 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte
     # chunk is a legitimate word gap.
     suppress_ws = False
 
+    def _semantic_content_events(items: list[Any]) -> list[GenEvent]:
+        nonlocal suppress_ws
+        out: list[GenEvent] = []
+        for item in items:
+            if isinstance(item, str):
+                if item and not (item.strip() == "" and suppress_ws):
+                    out.append(ContentDelta(item))
+                    if item.strip():
+                        suppress_ws = False
+            else:
+                out.append(item)
+                if isinstance(item, ToolCallsDelta):
+                    suppress_ws = True
+        return out
+
+    def _filter_tool_text(piece: str) -> list[GenEvent]:
+        text = strip_special_tokens(piece, specials)
+        return _semantic_content_events(content_special_filter.feed_items(text))
+
+    def _filter_tool_boundary(event: GenEvent) -> list[GenEvent]:
+        return _semantic_content_events(content_special_filter.boundary(event))
+
     def _close_open_call() -> ToolCallsDelta | None:
         nonlocal open_call, calls_emitted
         if open_call is None:
@@ -598,8 +872,6 @@ def _close_open_call() -> ToolCallsDelta | None:
         )
         open_call = None
         calls_emitted += 1
-        nonlocal suppress_ws
-        suppress_ws = True
         return ToolCallsDelta([call])
 
     def _route_tool_text(piece: str) -> list[GenEvent]:
@@ -613,12 +885,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
                 # close it first so wire order matches generation order.
                 done = _close_open_call()
                 if done is not None:
-                    out.append(done)
-                stripped = strip_special_tokens(payload, specials)
-                if stripped and not (stripped.strip() == "" and suppress_ws):
-                    out.append(ContentDelta(stripped))
-                    if stripped.strip():
-                        suppress_ws = False
+                    out.extend(_filter_tool_boundary(done))
+                out.extend(_filter_tool_text(payload))
                 continue
             for frag in payload:
                 starts_new = frag.name is not None and (
@@ -627,7 +895,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
                 if starts_new:
                     done = _close_open_call()
                     if done is not None:
-                        out.append(done)
+                        out.extend(_filter_tool_boundary(done))
                 if open_call is None or starts_new:
                     open_call = {
                         "detector_index": frag.tool_index,
@@ -635,14 +903,16 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
                         "params": "",
                         "ordinal": calls_emitted,
                     }
-                    out.append(ToolCallStart(
-                        tool_index=calls_emitted, name=frag.name, args_prefix_stable=frag_stable,
-                    ))
+                    out.extend(_filter_tool_boundary(ToolCallStart(
+                        tool_index=calls_emitted,
+                        name=frag.name,
+                        args_prefix_stable=frag_stable,
+                    )))
                 if frag.parameters:
                     open_call["params"] += frag.parameters
-                    out.append(
+                    out.extend(_filter_tool_boundary(
                         ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters)
-                    )
+                    ))
         return out
 
     engine_finish_reason: str | None = None
@@ -657,7 +927,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         if reasoning_parser is not None and content_delta:
             reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta)
             if reasoning_delta:
-                stripped_reasoning = strip_special_tokens(reasoning_delta, specials)
+                stripped_reasoning = _filter_reasoning(reasoning_delta)
                 if stripped_reasoning:  # a bare special token must not open a thinking block
                     yield ReasoningDelta(stripped_reasoning)
         if content_delta:
@@ -667,7 +937,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += content_delta
             else:
-                yield ContentDelta(strip_special_tokens(content_delta, specials))
+                stripped_content = _filter_content(content_delta)
+                if stripped_content:
+                    yield ContentDelta(stripped_content)
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
@@ -678,7 +950,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
     if reasoning_parser is not None:
         flush_reasoning, flush_content = reasoning_parser.flush()
         if flush_reasoning:
-            stripped_reasoning = strip_special_tokens(flush_reasoning, specials)
+            stripped_reasoning = _filter_reasoning(flush_reasoning)
             if stripped_reasoning:
                 yield ReasoningDelta(stripped_reasoning)
         if flush_content:
@@ -688,7 +960,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += flush_content
             else:
-                yield ContentDelta(strip_special_tokens(flush_content, specials))
+                stripped_content = _filter_content(flush_content)
+                if stripped_content:
+                    yield ContentDelta(stripped_content)
 
     # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins.
     finish_reason = engine_finish_reason or "stop"
@@ -702,35 +976,60 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         for frag in tool_parser.finalize_stream():
             if open_call is not None and frag.parameters:
                 open_call["params"] += frag.parameters
-                yield ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters)
+                for ev in _filter_tool_boundary(
+                    ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters)
+                ):
+                    yield ev
         done = _close_open_call()
         if done is not None:
-            yield done
+            for ev in _filter_tool_boundary(done):
+                yield ev
         for item in tool_parser.recover_truncated_call():
-            yield ToolCallStart(tool_index=calls_emitted, name=item.name)
-            yield ToolCallsDelta(
-                [ToolCallItem(tool_index=calls_emitted, name=item.name, parameters=item.parameters)]
-            )
+            recovered = [
+                ToolCallStart(tool_index=calls_emitted, name=item.name),
+                ToolCallsDelta([
+                    ToolCallItem(
+                        tool_index=calls_emitted,
+                        name=item.name,
+                        parameters=item.parameters,
+                    )
+                ]),
+            ]
+            for event in recovered:
+                for ev in _filter_tool_boundary(event):
+                    yield ev
             calls_emitted += 1
         residual = tool_parser.finish_stream()
         if residual:
-            stripped = strip_special_tokens(residual, specials)
-            if stripped and not (stripped.strip() == "" and suppress_ws):
-                yield ContentDelta(stripped)
+            for ev in _filter_tool_text(residual):
+                yield ev
         if calls_emitted and finish_reason != "length":
             finish_reason = "tool_calls"
     else:
         parsed = _parse_tool_response(pending, spec, state) if parse_tools else None
         if parsed is not None:
             normal_text, tool_calls = parsed
-            normal_text = strip_special_tokens(normal_text, specials)
+            normal_text = _filter_content(normal_text)
             if normal_text:
                 yield ContentDelta(normal_text)
             yield ToolCallsDelta(tool_calls)
             if finish_reason != "length":
                 finish_reason = "tool_calls"
         elif parse_tools and pending:
-            yield ContentDelta(strip_special_tokens(pending, specials))
+            stripped_pending = _filter_content(pending)
+            if stripped_pending:
+                yield ContentDelta(stripped_pending)
+
+    trailing_reasoning = reasoning_special_filter.finish()
+    if trailing_reasoning:
+        yield ReasoningDelta(trailing_reasoning)
+    if tool_parser is not None:
+        for ev in _semantic_content_events(content_special_filter.finish_items()):
+            yield ev
+    else:
+        trailing_content = content_special_filter.finish()
+        if trailing_content:
+            yield ContentDelta(trailing_content)
 
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
@@ -770,9 +1069,15 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
             finish_reason = "tool_calls"
 
     specials = _leaked_special_tokens(state)
+    reasoning_special_filter = _make_semantic_special_token_filter(state)
+    content_special_filter = _make_semantic_special_token_filter(state)
     return GenResult(
-        reasoning=strip_special_tokens(reasoning_text, specials).strip(),
-        content=strip_special_tokens(content_text, specials),
+        reasoning=reasoning_special_filter.feed(
+            strip_special_tokens(reasoning_text, specials), final=True
+        ).strip(),
+        content=content_special_filter.feed(
+            strip_special_tokens(content_text, specials), final=True
+        ),
         tool_calls=tool_calls,
         finish_reason=finish_reason,
         prompt_tokens=prompt_tokens,
diff --git a/tests/server/test_qwen_semantic_special_tokens.py b/tests/server/test_qwen_semantic_special_tokens.py
new file mode 100644
index 0000000000..11e35ccb3e
--- /dev/null
+++ b/tests/server/test_qwen_semantic_special_tokens.py
@@ -0,0 +1,405 @@
+"""Qwen protocol markers are private to semantic chat-style API responses."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+_PY = os.path.join(_ROOT, "python")
+if _PY not in sys.path:
+    sys.path.insert(0, _PY)
+
+from freetoken.llm.llm import LLM, RequestStatus  # noqa: E402
+from freetoken.message import DetokenizeMsg, UserReply  # noqa: E402
+from freetoken.server.api_server import FrontendManager  # noqa: E402
+from freetoken.server.api_models import CompletionRequest  # noqa: E402
+from freetoken.server.generation import (  # noqa: E402
+    _QWEN_SEMANTIC_HIDDEN_TOKENS,
+    ContentDelta,
+    GenDone,
+    GenSpec,
+    ReasoningDelta,
+    ToolCallsDelta,
+    generate_events,
+    generate_full,
+)
+from freetoken.server.openai_api import handle_completion, stream_completion_chunks  # noqa: E402
+
+
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "echo",
+            "parameters": {
+                "type": "object",
+                "properties": {"value": {"type": "string"}},
+            },
+        },
+    }
+]
+
+
+class FakeState:
+    def __init__(
+        self,
+        chunks: list[str],
+        *,
+        completion_tokens: list[int] | None = None,
+        reasoning_parser: str | None = "qwen3",
+        tool_call_parser: str = "qwen3_coder",
+    ) -> None:
+        self.config = SimpleNamespace(
+            reasoning_parser=reasoning_parser,
+            tool_call_parser=tool_call_parser,
+            served_model_name="qwen-unit",
+        )
+        self._chunks = chunks
+        self._completion_tokens = completion_tokens or [1] * len(chunks)
+        self.sent = None
+
+    def new_user(self) -> int:
+        return 42
+
+    async def send_one(self, msg) -> None:
+        self.sent = msg
+
+    async def wait_for_ack(self, uid: int):
+        assert uid == 42
+        for index, (chunk, token_delta) in enumerate(zip(self._chunks, self._completion_tokens)):
+            last = index == len(self._chunks) - 1
+            yield UserReply(
+                uid=uid,
+                incremental_output=chunk,
+                finished=last,
+                finish_reason="stop" if last else None,
+                prompt_tokens_delta=5 if index == 0 else 0,
+                completion_tokens_delta=token_delta,
+            )
+
+
+def _spec(*, tools=None) -> GenSpec:
+    return GenSpec(
+        messages=[],
+        sampling_params=SimpleNamespace(),
+        chat_template_kwargs={"enable_thinking": False},
+        template_tools=tools,
+        parser_tools=tools,
+    )
+
+
+async def _events(chunks: list[str], *, tools=None):
+    return [event async for event in generate_events(42, _spec(tools=tools), FakeState(chunks))]
+
+
+def _visible(events) -> str:
+    return "".join(event.text for event in events if isinstance(event, ContentDelta))
+
+
+@pytest.mark.parametrize("marker", _QWEN_SEMANTIC_HIDDEN_TOKENS)
+def test_every_hidden_marker_is_chunk_boundary_safe(marker: str):
+    text = f"before{marker}after"
+    events = asyncio.run(_events(list(text)))
+
+    assert _visible(events) == "beforeafter"
+    assert isinstance(events[-1], GenDone)
+    assert events[-1].completion_tokens == len(text)
+
+
+def test_stream_and_nonstream_match_while_preserving_semantic_markers_and_literals():
+    grounding = (
+        "<|object_ref_start|>cat<|object_ref_end|>"
+        "<|box_start|>(1,2),(3,4)<|box_end|>"
+        "<|quad_start|>(1,2),(3,4),(5,6),(7,8)<|quad_end|>"
+    )
+    text = (
+        "<|im_start|><think>plan</think>answer"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{grounding} literal \"<|im_start|>\" and `<|image_pad|>`<|im_end|>"
+    )
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text], completion_tokens=[len(text)])))
+
+    reasoning = "".join(event.text for event in streamed if isinstance(event, ReasoningDelta))
+    expected_content = f"answer{grounding} literal \"<|im_start|>\" and `<|image_pad|>`"
+    assert reasoning == full.reasoning == "plan"
+    assert _visible(streamed) == full.content == expected_content
+    assert full.completion_tokens == streamed[-1].completion_tokens == len(text)
+
+
+def test_tool_and_think_markers_reach_their_parsers_before_filtering():
+    block = (
+        "<tool_call><function=echo><parameter=value>"
+        'literal <|im_start|>'
+        "</parameter></function></tool_call>"
+    )
+    text = "<think>choose tool</think><|im_start|>" + block + "<|im_end|>"
+    streamed = asyncio.run(_events(list(text), tools=TOOLS))
+    full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState([text])))
+
+    stream_calls = [
+        call for event in streamed if isinstance(event, ToolCallsDelta) for call in event.calls
+    ]
+    assert "".join(
+        event.text for event in streamed if isinstance(event, ReasoningDelta)
+    ) == full.reasoning == "choose tool"
+    assert _visible(streamed) == full.content == ""
+    assert len(stream_calls) == len(full.tool_calls) == 1
+    assert stream_calls[0].name == full.tool_calls[0].name == "echo"
+    assert json.loads(stream_calls[0].parameters) == json.loads(full.tool_calls[0].parameters) == {
+        "value": "literal <|im_start|>"
+    }
+
+
+def test_filter_state_is_isolated_per_concurrent_request():
+    async def run_pair():
+        return await asyncio.gather(
+            _events(["left<|im_", "start|>right"]),
+            _events(['quoted "<|im_start|>"']),
+        )
+
+    first, second = asyncio.run(run_pair())
+    assert _visible(first) == "leftright"
+    assert _visible(second) == 'quoted "<|im_start|>"'
+
+
+@pytest.mark.parametrize("quote", ['"', "`", "'"])
+def test_unterminated_literal_cannot_disable_filtering(quote: str):
+    text = f"answer {quote}oops <|im_end|>"
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == f"answer {quote}oops "
+
+
+@pytest.mark.parametrize("slashes", ["\\", "\\\\"])
+def test_escaped_marker_opener_does_not_bypass_unterminated_quote_filter(slashes: str):
+    text = f'answer "oops {slashes}<|im_end|>'
+    expected = f'answer "oops {slashes}'
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == expected
+
+
+@pytest.mark.parametrize(("opening", "closing"), [("“", "”"), ("‘", "’")])
+def test_typographic_quoted_marker_is_preserved(opening: str, closing: str):
+    text = f"literal {opening}<|im_start|>{closing}"
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == text
+
+
+def test_suppressed_marker_does_not_turn_a_possessive_into_an_open_quote():
+    text = "foo<|im_end|>'s <|vision_pad|>done"
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == "foo's done"
+
+
+@pytest.mark.parametrize(
+    ("text", "expected"),
+    [
+        ("<<|im_end|>", "<"),
+        ("<|im_<|im_end|>", "<|im_"),
+        ("x<<|vision_pad|>y", "x<y"),
+    ],
+)
+def test_overlapping_candidate_start_reprocesses_the_mismatch(text: str, expected: str):
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == expected
+
+
+def test_partial_candidate_before_quote_closer_does_not_extend_literal_scope():
+    text = 'literal "<|im_" after <|im_end|> "done"'
+    expected = 'literal "<|im_" after  "done"'
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    assert _visible(streamed) == full.content == expected
+
+
+def test_tool_event_inside_marker_fragments_keeps_stream_nonstream_parity():
+    block = (
+        "<tool_call><function=echo><parameter=value>ok</parameter>"
+        "</function></tool_call>"
+    )
+    chunks = ["before<|im_", block, "start|>after"]
+    streamed = asyncio.run(_events(chunks, tools=TOOLS))
+    full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)])))
+
+    assert _visible(streamed) == full.content == "beforeafter"
+    assert len([event for event in streamed if isinstance(event, ToolCallsDelta)]) == 1
+    assert len(full.tool_calls) == 1
+
+
+def test_disproved_marker_prefix_stays_before_intervening_tool_event():
+    block = (
+        "<tool_call><function=echo><parameter=value>ok</parameter>"
+        "</function></tool_call>"
+    )
+    streamed = asyncio.run(_events(["pre<|im_", block, "Xpost"], tools=TOOLS))
+
+    kinds = [type(event).__name__ for event in streamed]
+    call_start_index = kinds.index("ToolCallStart")
+    call_end_index = kinds.index("ToolCallsDelta")
+    before = "".join(
+        event.text
+        for event in streamed[:call_start_index]
+        if isinstance(event, ContentDelta)
+    )
+    after = "".join(
+        event.text
+        for event in streamed[call_end_index + 1 :]
+        if isinstance(event, ContentDelta)
+    )
+
+    assert before == "pre<|im_"
+    assert after == "Xpost"
+
+
+def test_tool_event_flushes_prior_quoted_prose_before_call():
+    block = (
+        "<tool_call><function=echo><parameter=value>ok</parameter>"
+        "</function></tool_call>"
+    )
+    chunks = [
+        'before "<|im_start|> quoted',
+        block,
+        ' tail <|vision_pad|>" after<|im_end|>',
+    ]
+    streamed = asyncio.run(_events(chunks, tools=TOOLS))
+    full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)])))
+
+    kinds = [type(event).__name__ for event in streamed]
+    call_start_index = kinds.index("ToolCallStart")
+    call_end_index = kinds.index("ToolCallsDelta")
+    prefix = "".join(
+        event.text
+        for event in streamed[:call_start_index]
+        if isinstance(event, ContentDelta)
+    )
+    suffix = "".join(
+        event.text
+        for event in streamed[call_end_index + 1 :]
+        if isinstance(event, ContentDelta)
+    )
+    assert prefix == 'before "<|im_start|> quoted'
+    assert suffix == ' tail <|vision_pad|>" after'
+    assert call_start_index < call_end_index
+    assert _visible(streamed) == full.content == (
+        'before "<|im_start|> quoted tail <|vision_pad|>" after'
+    )
+
+
+def test_reasoning_filter_hides_transport_marker_and_preserves_fim_tokens():
+    text = (
+        "<think>plan<|vision_pad|></think>"
+        "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|><|im_end|>"
+    )
+    streamed = asyncio.run(_events(list(text)))
+    full = asyncio.run(generate_full(42, _spec(), FakeState([text])))
+
+    reasoning = "".join(
+        event.text for event in streamed if isinstance(event, ReasoningDelta)
+    )
+    expected_content = "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|>"
+    assert reasoning == full.reasoning == "plan"
+    assert _visible(streamed) == full.content == expected_content
+
+
+def test_non_qwen_semantic_dialect_is_unchanged():
+    text = '<|im_start|>literal "<|vision_pad|>"<|im_end|>'
+    state = FakeState(
+        [text],
+        reasoning_parser=None,
+        tool_call_parser="llama3",
+    )
+
+    result = asyncio.run(generate_full(42, _spec(), state))
+
+    assert result.content == text
+
+
+def test_legacy_qwen_parser_alias_enables_semantic_filter_without_reasoning():
+    text = "<|im_start|>visible<|vision_pad|><|im_end|>"
+    state = FakeState([text], reasoning_parser=None, tool_call_parser="qwen")
+    result = asyncio.run(generate_full(42, _spec(), state))
+
+    assert result.content == "visible"
+
+
+def test_legacy_completions_paths_remain_raw_for_qwen():
+    raw = '<|im_start|>raw "<|vision_pad|>"<|im_end|>'
+
+    full = asyncio.run(
+        handle_completion(
+            CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8),
+            request=None,
+            state=FakeState([raw]),
+            model_sampling={},
+        )
+    )
+    assert full["choices"][0]["text"] == raw
+
+    async def stream_text() -> str:
+        chunks = [
+            chunk async for chunk in stream_completion_chunks(
+                42,
+                CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8, stream=True),
+                FakeState([raw]),
+            )
+        ]
+        frames = []
+        for chunk in chunks:
+            for line in chunk.decode().splitlines():
+                if line.startswith("data: ") and line != "data: [DONE]":
+                    frames.append(json.loads(line[6:]))
+        return "".join(frame["choices"][0]["text"] for frame in frames)
+
+    assert asyncio.run(stream_text()) == raw
+
+
+def test_raw_generate_stream_remains_unfiltered_for_qwen():
+    raw = "<|im_start|>raw<|vision_pad|><|im_end|>"
+
+    async def stream_text() -> str:
+        chunks = [
+            chunk async for chunk in FrontendManager.stream_generate(FakeState([raw]), 42)
+        ]
+        frames = []
+        for chunk in chunks:
+            for line in chunk.decode().splitlines():
+                if line.startswith("data: ") and line != "data: [DONE]":
+                    frames.append(json.loads(line[6:]))
+        return "".join(frame["text"] for frame in frames)
+
+    assert asyncio.run(stream_text()) == raw
+
+
+def test_offline_generation_keeps_exact_output_token_ids():
+    offline = SimpleNamespace(
+        status_map={7: RequestStatus(uid=7, input_ids=[], output_ids=[])},
+        eos_token_ids=set(),
+    )
+    marker_token_ids = [248045, 248044, 248046]
+
+    LLM.offline_send_result(
+        offline,
+        [
+            DetokenizeMsg(uid=7, next_token=token_id, finished=False)
+            for token_id in marker_token_ids
+        ],
+    )
+
+    assert offline.status_map[7].output_ids == marker_token_ids

From f139235a7599e33f923c3ef962cc860306aae071 Mon Sep 17 00:00:00 2001
From: sime2408 <simun.sunjic@gmail.com>
Date: Sat, 29 Aug 2026 09:40:12 +0200
Subject: [PATCH 011/249] fix(server): report why a tokenizer/detokenizer
 worker died during load
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_run_scheduler wraps startup and pushes an ("error", reason) ack before dying,
so the supervisor can report the real cause. _run_tokenize_worker did not, so
every tokenizer-side startup failure reached the parent only as a dead process
and surfaced as the generic "backend worker freetoken-detokenizer-0 exited
during load" — the actionable reason (unreadable tokenizer.json, missing
checkpoint file, ZMQ bind clash) was lost.

Wrap tokenize_worker in the same try/report/re-raise the scheduler already uses.
The traceback still prints and the process still exits non-zero; only the ack is
added. No ack_queue (direct call) stays a plain propagate.

Adds tests/server/test_launch_worker_errors.py, which fails before this change.
---
 python/freetoken/server/launch.py         | 13 ++++-
 tests/server/test_launch_worker_errors.py | 69 +++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 tests/server/test_launch_worker_errors.py

diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py
index acef55175a..a2d6512812 100644
--- a/python/freetoken/server/launch.py
+++ b/python/freetoken/server/launch.py
@@ -52,7 +52,18 @@ def _run_tokenize_worker(detach: bool, **kwargs) -> None:
         _detach_process_group()
     from freetoken.tokenizer import tokenize_worker
 
-    tokenize_worker(**kwargs)
+    ack_queue = kwargs.get("ack_queue")
+    try:
+        tokenize_worker(**kwargs)
+    except Exception as exc:  # noqa: BLE001 -- surface the reason, then let it propagate
+        # Same contract as _run_scheduler: without this, a tokenizer/detokenizer that dies
+        # before its ready ack reaches the parent only as a dead process, so the supervisor
+        # reports the generic "backend worker freetoken-detokenizer-0 exited during load" and
+        # the real cause (an unreadable tokenizer.json, a missing checkpoint file, a ZMQ bind
+        # clash) is lost. The traceback still prints and the process still exits non-zero.
+        if ack_queue is not None:
+            _report_startup_error(ack_queue, exc)
+        raise
 
 
 def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None:
diff --git a/tests/server/test_launch_worker_errors.py b/tests/server/test_launch_worker_errors.py
new file mode 100644
index 0000000000..8779016bdd
--- /dev/null
+++ b/tests/server/test_launch_worker_errors.py
@@ -0,0 +1,69 @@
+"""A tokenizer/detokenizer worker that dies during load must report WHY.
+
+Regression test for the asymmetry that made every tokenizer-side startup failure surface as
+the generic "backend worker freetoken-detokenizer-0 exited during load": ``_run_scheduler``
+pushed an ("error", reason) ack before dying, ``_run_tokenize_worker`` did not.
+"""
+
+from __future__ import annotations
+
+import queue
+
+import pytest
+
+from freetoken.server.launch import _run_tokenize_worker
+
+
+def _drain(q: "queue.Queue") -> list:
+    out = []
+    while True:
+        try:
+            out.append(q.get_nowait())
+        except queue.Empty:
+            return out
+
+
+def test_tokenize_worker_reports_the_real_reason_before_dying(monkeypatch):
+    """A raising tokenize_worker pushes ("error", reason) so the supervisor can surface it."""
+    import freetoken.tokenizer as tok_mod
+
+    def boom(**_kwargs):
+        raise FileNotFoundError("tokenizer.json is missing from the checkpoint")
+
+    monkeypatch.setattr(tok_mod, "tokenize_worker", boom)
+    q: "queue.Queue" = queue.Queue()
+
+    with pytest.raises(FileNotFoundError):
+        _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope")
+
+    acks = _drain(q)
+    assert acks, "the dying worker pushed no ack at all"
+    kind, reason = acks[0]
+    assert kind == "error"
+    assert "FileNotFoundError" in reason
+    assert "tokenizer.json is missing" in reason
+
+
+def test_tokenize_worker_without_ack_queue_still_propagates(monkeypatch):
+    """No ack_queue (direct call / test harness) must not turn the failure into an AttributeError."""
+    import freetoken.tokenizer as tok_mod
+
+    def boom(**_kwargs):
+        raise ValueError("bad tokenizer config")
+
+    monkeypatch.setattr(tok_mod, "tokenize_worker", boom)
+
+    with pytest.raises(ValueError, match="bad tokenizer config"):
+        _run_tokenize_worker(detach=False, tokenizer_path="/nope")
+
+
+def test_tokenize_worker_success_pushes_no_error_ack(monkeypatch):
+    """The happy path must stay clean — no spurious ("error", …) on the ack queue."""
+    import freetoken.tokenizer as tok_mod
+
+    monkeypatch.setattr(tok_mod, "tokenize_worker", lambda **_k: None)
+    q: "queue.Queue" = queue.Queue()
+
+    _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope")
+
+    assert not [a for a in _drain(q) if isinstance(a, tuple) and a and a[0] == "error"]

From 855650d60465d66ca3a5feb1e403fb8eba33e623 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Tue, 1 Sep 2026 21:08:01 +0300
Subject: [PATCH 012/249] server: fail-close chat logprobs under semantic
 parsing

Review finding on #224 (HaileyStorm): streaming chat carried logprob entries
of hidden reasoning tokens onto the next visible content delta -- pending
entries accumulated before the reasoning parser classified the text, and
_content_delta() drained the whole list into whatever visible chunk came
next. That broke token<->content alignment and leaked hidden reasoning token
strings through logprobs entries; non-streaming meanwhile kept every sampled
entry and end-of-stream dropped undrained ones -- three behaviors for one
surface.

Close it at both layers:

- /v1/chat/completions rejects logprobs=true with a 400 when the server runs
  a reasoning parser or the request enables tool parsing: the semantic layer
  can hide or reclassify tokens, so entries cannot be aligned 1:1 with
  visible content tokens. tool_choice="none" keeps logprobs available; the
  raw /v1/completions path (no semantic layer) is unchanged.

- _generate_events_impl now collects entries only on the passthrough path
  (no reasoning parser, no tool parsing), so a hidden-token entry can never
  ride a later visible delta even for a caller that skips API validation.

Replaces test_reasoning_logprob_is_carried_to_next_content_delta (which
enshrined the carry) with fail-close coverage for both conflict cases, the
tool_choice="none" pass-through, and a no-carry guard test.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 python/freetoken/server/api_models.py |  2 ++
 python/freetoken/server/generation.py | 10 ++++++-
 python/freetoken/server/openai_api.py | 29 ++++++++++++++++++
 tests/server/test_logprobs_api.py     | 42 ++++++++++++++++++++++++---
 4 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 7bd43e642e..307c275755 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -78,6 +78,8 @@ class ChatCompletionRequest(BaseModel):
     frequency_penalty: float = 0.0
     # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires
     # logprobs=true; entries cover generated tokens only (no prompt logprobs).
+    # Fail-closed (400) when a reasoning parser or tool parsing is active: entries
+    # must align 1:1 with visible content tokens (see _chat_logprobs_conflict).
     logprobs: bool = False
     top_logprobs: int | None = None
     chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index e4627a06b5..1ee44f11c4 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -587,6 +587,14 @@ def _content_delta(text: str) -> ContentDelta:
             tool_parser = candidate
     frag_stable = tool_parser.args_fragments_prefix_stable() if tool_parser else True
 
+    # Logprob entries are collected only on the passthrough path: with a reasoning
+    # parser or tool parsing active, text can be hidden, held back, or reclassified,
+    # so an entry could come to describe a token that never reaches visible content.
+    # The API layer fail-closes those request combinations; this guard is the
+    # structural half of the contract — a hidden-token entry must never ride a
+    # later visible delta, whatever the caller.
+    collect_logprobs = reasoning_parser is None and tool_parser is None and not parse_tools
+
     # Streaming tool-call assembly: detectors emit fragments (name first, then
     # argument diffs); they accumulate here and the call is emitted complete when
     # the next call starts, trailing text arrives, or the stream ends.
@@ -667,7 +675,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
-        if getattr(ack, "logprobs", None) is not None:
+        if collect_logprobs and getattr(ack, "logprobs", None) is not None:
             pending_logprobs.append(ack.logprobs)
         content_delta = ack.incremental_output
         if reasoning_parser is not None and content_delta:
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 82c32bb78a..b3e71ed99d 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -153,6 +153,32 @@ async def v1_models():
         )])
 
 
+def _chat_logprobs_conflict(req: ChatCompletionRequest, state: Any) -> str | None:
+    """Fail-close chat logprobs when a semantic parsing layer can hide or reclassify
+    generated tokens. A reasoning parser routes tokens out of visible content (and its
+    stream buffering re-chunks text), and tool parsing consumes tokens into tool_calls,
+    so logprob entries could not be aligned 1:1 with visible content tokens — and hidden
+    reasoning-token strings must never surface through logprobs entries. The raw
+    /v1/completions path has no semantic layer and keeps full logprobs support."""
+    if not req.logprobs:
+        return None
+    if getattr(state.config, "reasoning_parser", None):
+        return (
+            "logprobs on /v1/chat/completions are not supported when the server runs a "
+            "reasoning parser: reasoning tokens are hidden from message content, so "
+            "logprob entries cannot be aligned with it. Use /v1/completions for raw "
+            "token logprobs."
+        )
+    if _should_parse_tools(req):
+        return (
+            "logprobs with tool parsing on /v1/chat/completions are not supported: "
+            "tool-call tokens are consumed into tool_calls, so logprob entries cannot "
+            "be aligned with message content. Send tool_choice='none' or use "
+            "/v1/completions for raw token logprobs."
+        )
+    return None
+
+
 async def handle_chat_completion(
     req: ChatCompletionRequest,
     request: Request | None,
@@ -166,6 +192,9 @@ async def handle_chat_completion(
     logprobs_error = chat_logprobs_error(req)
     if logprobs_error is not None:
         return create_error_response(logprobs_error, param="top_logprobs")
+    logprobs_conflict = _chat_logprobs_conflict(req, state)
+    if logprobs_conflict is not None:
+        return create_error_response(logprobs_conflict, param="logprobs")
     if _response_format_unsupported(req.response_format):
         return create_error_response(
             "response_format json_object/json_schema is not supported (no constrained decoding)",
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 3b4fd5e32a..8cf0387c60 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -247,11 +247,45 @@ def test_logprobs_validation_errors() -> None:
     assert completion_echo.status_code == 400
 
 
-def test_reasoning_logprob_is_carried_to_next_content_delta() -> None:
-    reasoning_entry = entry(1, "thought", -0.1)
+def test_chat_logprobs_fail_closed_under_semantic_parsing() -> None:
+    # A server-side reasoning parser hides <think> tokens from message content, so
+    # logprob entries cannot be aligned with it: reject up front, stream and
+    # non-stream alike, before any engine work is submitted.
+    with_parser = FakeState([], reasoning_parser="qwen3")
+    for stream in (False, True):
+        resp = run(
+            handle_chat_completion(chat_request(logprobs=True, stream=stream), None, with_parser, {})
+        )
+        assert resp.status_code == 400
+        assert json.loads(resp.body)["error"]["param"] == "logprobs"
+
+    # Tool parsing consumes tokens into tool_calls -- same conflict.
+    tool = {"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}
+    tools_resp = run(
+        handle_chat_completion(chat_request(logprobs=True, tools=[tool]), None, FakeState([]), {})
+    )
+    assert tools_resp.status_code == 400
+    assert json.loads(tools_resp.body)["error"]["param"] == "logprobs"
+
+    # tool_choice="none" disables parsing, so logprobs stay available.
+    ok = run(
+        handle_chat_completion(
+            chat_request(logprobs=True, tools=[tool], tool_choice="none"),
+            None,
+            FakeState([reply("Hi", finished=True, logprobs=entry(1, "Hi", -0.1))]),
+            {},
+        )
+    )
+    assert ok["choices"][0]["logprobs"]["content"][0]["token"] == "Hi"
+
+
+def test_reasoning_logprob_is_not_carried_to_content_delta() -> None:
+    # The generation-layer guard behind the API fail-close: even when a caller
+    # bypasses request validation, a hidden reasoning token's entry is dropped,
+    # never attached to a later visible delta (where its token string would leak).
     state = FakeState(
         [
-            reply("<think>thought</think>", logprobs=reasoning_entry),
+            reply("<think>thought</think>", logprobs=entry(1, "thought", -0.1)),
             reply("answer", finished=True),
         ],
         reasoning_parser="qwen3",
@@ -267,7 +301,7 @@ def test_reasoning_logprob_is_carried_to_next_content_delta() -> None:
         and event["choices"]
         and event["choices"][0]["delta"].get("content") == "answer"
     )
-    assert content_choice["logprobs"]["content"][0]["token"] == "thought"
+    assert "logprobs" not in content_choice
 
 
 async def _collect(iterator):

From 274a2ce5212969a2d8b3a149bcbedf04d63c0814 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Tue, 1 Sep 2026 21:11:57 +0300
Subject: [PATCH 013/249] server: abort abandoned non-streaming requests on
 client disconnect

Follow-up to the review evidence on #222 (benwilson): the non-streaming path
kept generating after the client was gone -- with --max-running-requests 1 an
abandoned request is a full outage for its remaining max_tokens (measured
repro: ~70 s of dead decode, the next client's first token 61 s late).

stream_with_cancellation was the only place a disconnect was observed. Give
the plain handlers its non-streaming twin: _await_watching_disconnect() runs
the generation drain as a task and polls request.is_disconnected() once a
second; when the client goes away it delivers the same shielded abort_user
(claim + AbortMsg first, so the drain task's own cleanup cannot swallow the
claim), then winds the drain down and answers 499 (client closed request --
for the access log; the wire is dead). Handler cancellation (server shutdown)
delivers the abort too, mirroring the streaming path.

Covers /v1/chat/completions and each prompt of a non-streaming
/v1/completions batch. Requests with request=None (adapter-internal callers)
are unaffected.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 python/freetoken/server/openai_api.py    | 104 ++++++++++++++++++++---
 tests/server/test_stream_cancellation.py |  97 +++++++++++++++++++++
 2 files changed, 187 insertions(+), 14 deletions(-)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b4becd2631..f053054841 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+import logging
 import time
 import uuid
 from collections.abc import AsyncIterator, Callable
@@ -39,6 +40,8 @@
     submit_generation,
 )
 
+logger = logging.getLogger(__name__)
+
 #: The wire superset plus "off", DeepSeek's disable synonym that
 #: effort_toggle_kwargs has always honored.
 _ACCEPTED_EFFORTS = (*KNOWN_REASONING_EFFORTS, "off")
@@ -144,6 +147,65 @@ async def v1_models():
         )])
 
 
+# How often a non-streaming handler looks at the transport while the engine
+# generates. The streaming path checks per chunk; one second bounds an abandoned
+# request's extra decode work without measurable polling overhead.
+_DISCONNECT_POLL_SECONDS = 1.0
+
+
+async def _await_watching_disconnect(awaitable, request: Request | None, state: Any, uid: int):
+    """Run an in-flight generation awaitable while watching the transport — the
+    non-streaming twin of stream_with_cancellation. Returns the awaitable's result,
+    or None when the client disconnected first: the same shielded abort_user is
+    delivered so an abandoned request stops burning decode slots instead of running
+    to max_tokens (with --max-running-requests 1 that is a full outage for its
+    remaining budget)."""
+    gen = asyncio.ensure_future(awaitable)
+    if request is None:
+        return await gen
+    try:
+        while True:
+            done, _ = await asyncio.wait({gen}, timeout=_DISCONNECT_POLL_SECONDS)
+            if done:
+                return gen.result()
+            if await request.is_disconnected():
+                break
+    except asyncio.CancelledError:
+        # The handler itself was cancelled (server shutdown): deliver the abort,
+        # then let the cancellation propagate — mirrors stream_with_cancellation.
+        try:
+            await asyncio.shield(state.abort_user(uid))
+        except Exception:  # noqa: BLE001
+            logger.exception("Failed to deliver abort for user %s", uid)
+        gen.cancel()
+        raise
+    # Client gone. Claim + AbortMsg first (abort_user is a no-op once the ack
+    # loop's own cleanup has run), then wind down the drain task; its result is
+    # undeliverable either way.
+    try:
+        await asyncio.shield(state.abort_user(uid))
+    except Exception:  # noqa: BLE001
+        logger.exception("Failed to deliver abort for user %s", uid)
+    gen.cancel()
+    try:
+        await gen
+    except asyncio.CancelledError:
+        pass
+    except Exception:  # noqa: BLE001 — the response is undeliverable, only cleanup matters
+        pass
+    return None
+
+
+def _client_disconnected_response() -> JSONResponse:
+    # 499 — nginx's "client closed request". The client is gone; this status is
+    # for the access log, not the wire.
+    return create_error_response(
+        "client disconnected before the response was ready",
+        status_code=499,
+        err_type="client_disconnected",
+    )
+
+
 async def handle_chat_completion(
     req: ChatCompletionRequest,
     request: Request | None,
@@ -199,9 +261,13 @@ async def handle_chat_completion(
         return StreamingResponse(chunks, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/chat/completions")
+        result = await _await_watching_disconnect(
+            generate_full(uid, spec, state, source="/v1/chat/completions"), request, state, uid
+        )
     except GenerationError as exc:
         return create_error_response(str(exc), code=exc.code)
+    if result is None:
+        return _client_disconnected_response()
     message: dict[str, Any] = {"role": "assistant", "content": result.content}
     if result.reasoning:
         message["reasoning_content"] = result.reasoning
@@ -411,19 +477,29 @@ async def handle_completion(
     for index, prompt in enumerate(prompts):
         uid = state.new_user()
         await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling)))
-        text = ""
-        finish_reason = "stop"
-        async for ack in state.wait_for_ack(uid):
-            if getattr(ack, "error", None):
-                return create_error_response(ack.error)
-            prompt_tokens += ack.prompt_tokens_delta
-            completion_tokens += ack.completion_tokens_delta
-            cached_tokens += ack.cached_tokens
-            text += ack.incremental_output
-            if ack.finished:
-                finish_reason = getattr(ack, "finish_reason", None) or "stop"
-                break
-        choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None})
+
+        async def _drain_one(uid: int = uid, index: int = index):
+            nonlocal prompt_tokens, completion_tokens, cached_tokens
+            text = ""
+            finish_reason = "stop"
+            async for ack in state.wait_for_ack(uid):
+                if getattr(ack, "error", None):
+                    return create_error_response(ack.error)
+                prompt_tokens += ack.prompt_tokens_delta
+                completion_tokens += ack.completion_tokens_delta
+                cached_tokens += ack.cached_tokens
+                text += ack.incremental_output
+                if ack.finished:
+                    finish_reason = getattr(ack, "finish_reason", None) or "stop"
+                    break
+            return {"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}
+
+        choice = await _await_watching_disconnect(_drain_one(), request, state, uid)
+        if choice is None:
+            return _client_disconnected_response()
+        if isinstance(choice, JSONResponse):
+            return choice
+        choices.append(choice)
 
     return {
         "id": f"cmpl-{uuid.uuid4().hex}",
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index 120c504564..ded9178bf9 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -118,3 +118,100 @@ async def one():
         assert 7 in state.ack_map  # wait_for_ack owns normal-path cleanup, not the stream
 
     asyncio.run(run())
+
+
+# --------------------------------------------------------------------------- #
+# Non-streaming handlers: a disconnected client must abort the engine request
+# (same shielded abort_user as the streaming path) instead of letting it run to
+# max_tokens while later requests queue behind it.
+# --------------------------------------------------------------------------- #
+
+from freetoken.server import openai_api
+from freetoken.server.api_models import ChatCompletionRequest, CompletionRequest
+
+
+def _ack(text, finished=False):
+    return SimpleNamespace(
+        error=None,
+        incremental_output=text,
+        finished=finished,
+        prompt_tokens_delta=0,
+        completion_tokens_delta=1,
+        cached_tokens=0,
+        finish_reason="stop" if finished else None,
+        matched_stop=None,
+        logprobs=None,
+    )
+
+
+class _ApiState:
+    """State fake for the openai_api handlers: hangs forever when given no acks."""
+
+    def __init__(self, acks=None):
+        self.config = SimpleNamespace(
+            model_path="/models/unit-model",
+            served_model_name="unit-model",
+            tool_call_parser="llama3",
+            reasoning_parser=None,
+        )
+        self.acks = acks
+        self.sent = []
+        self.aborted = []
+
+    def new_user(self):
+        return 7
+
+    async def send_one(self, msg):
+        self.sent.append(msg)
+
+    async def wait_for_ack(self, uid):
+        if self.acks is None:
+            await asyncio.sleep(3600)
+        for ack in self.acks or []:
+            yield ack
+
+    async def abort_user(self, uid):
+        self.aborted.append(uid)
+
+
+def _chat_req(**kwargs):
+    payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8}
+    payload.update(kwargs)
+    return ChatCompletionRequest(**payload)
+
+
+def test_chat_non_stream_disconnect_delivers_abort(monkeypatch):
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+    state = _ApiState(acks=None)  # generation never finishes on its own
+
+    resp = asyncio.run(
+        openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=True), state, {})
+    )
+
+    assert resp.status_code == 499
+    assert state.aborted == [7]
+
+
+def test_completion_non_stream_disconnect_delivers_abort(monkeypatch):
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+    state = _ApiState(acks=None)
+    req = CompletionRequest(model="m", prompt="hello", max_tokens=8)
+
+    resp = asyncio.run(
+        openai_api.handle_completion(req, _Request(disconnected=True), state, {})
+    )
+
+    assert resp.status_code == 499
+    assert state.aborted == [7]
+
+
+def test_non_stream_connected_client_gets_result_without_abort(monkeypatch):
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+    state = _ApiState(acks=[_ack("Hi", finished=True)])
+
+    result = asyncio.run(
+        openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=False), state, {})
+    )
+
+    assert result["choices"][0]["message"]["content"] == "Hi"
+    assert state.aborted == []

From d51a8e2f9146d0a3d2ad0e6806273ddbc5e55b00 Mon Sep 17 00:00:00 2001
From: dejay2 <218806300+dejay2@users.noreply.github.com>
Date: Wed, 2 Sep 2026 05:04:04 +0100
Subject: [PATCH 014/249] fix(engine): reserve the pool's dummy page in the
 --moe-cache-auto KV floor

plan_cache_budget works in USABLE pages, while create_kv_pool allocates
num_pages + 1: every pool family keeps page 0 as an unreachable
dummy/sentinel. The plan prices the KV half as num_pages * cache_per_page
and hands experts everything the reserve does not protect (the fill is
greedy and deliberately leaves no headroom), so it under-counted the KV
pool by exactly one page's worth of bytes, and that page was then
allocated out of memory the split had already given to expert slots. On
a model with a large cache_per_page that page is not small, and the
failure mode is the one --moe-cache-auto exists to prevent: an
arithmetic plan that says it fits, followed by a CUDA OOM in the KV
allocation at boot.

The same off-by-one made --kv-reserve-tokens quietly under-deliver: it
is documented as a KV-cache token floor, but ceil(N / page_size) usable
pages minus the dummy is less than N tokens of reachable capacity.
Reserve the sentinel page on top of the user-visible floor and say so in
the --kv-reserve-tokens help text.

Tested (Windows, RTX 5090):

  python -m pytest -q tests/engine tests/server

  before: 643 passed, 2 failed   after: 644 passed, 2 failed

The two failures are pre-existing on this box and unrelated: both
tests/engine/test_cache_budget.py::test_adjust_config_* need flashinfer,
which is not installed here.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK
---
 python/freetoken/engine/cache_budget.py |  4 +++-
 python/freetoken/server/args.py         |  5 ++++-
 tests/engine/test_cache_budget.py       | 20 ++++++++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py
index ab7c0a9f1d..8f85cf44dc 100644
--- a/python/freetoken/engine/cache_budget.py
+++ b/python/freetoken/engine/cache_budget.py
@@ -118,7 +118,9 @@ def resolve_moe_cache_auto(
     """
     budget_bytes = net_cache_budget_bytes(memory_ratio, baseline_free, weights_bytes, fixed_cache_size)
     max_slots = 992 if quant_format == "nvfp4_marlin" else total_experts
-    kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size)
+    # Every pool keeps page 0 as an unreachable dummy/sentinel. The CLI floor is expressed in
+    # usable tokens, so reserve that internal page in addition to the user-visible capacity.
+    kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + 1
     return plan_cache_budget(
         budget_bytes=budget_bytes,
         per_expert_bytes=per_expert_bytes,
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 5b4db587d1..ab0e0cd630 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -537,7 +537,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "--kv-reserve-tokens",
         type=int,
         default=ServerArgs.kv_reserve_tokens,
-        help="KV-cache token floor reserved before --moe-cache-auto fills experts.",
+        help=(
+            "Usable KV-cache token floor reserved before --moe-cache-auto fills experts "
+            "(the internal dummy page is additional)."
+        ),
     )
 
     parser.add_argument(
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 9ac2a4f4cc..40e1055b2c 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -112,6 +112,26 @@ def test_resolve_auto_applies_ratio_once_and_marlin_cap():
     assert size == 8 and pages == 40 and overlap is True
 
 
+def test_resolve_auto_reserves_usable_tokens_beyond_the_dummy_page():
+    size, pages, overlap = resolve_moe_cache_auto(
+        baseline_free=940,
+        weights_bytes=0,
+        memory_ratio=1.0,
+        cache_per_page=10,
+        fixed_cache_size=0,
+        per_expert_bytes=100,
+        num_experts=2,
+        total_experts=50,
+        prefill_overlap=False,
+        kv_reserve_tokens=256,
+        page_size=64,
+        quant_format="bf16",
+    )
+    assert overlap is False
+    assert (pages - 1) * 64 >= 256
+    assert size == 8 and pages == 14
+
+
 def test_resolve_auto_marlin_caps_slots():
     size, _, _ = resolve_moe_cache_auto(
         baseline_free=10_000_000, weights_bytes=0, memory_ratio=1.0,

From b6e799507621b5c0f96ceb4104365ddcaef559b3 Mon Sep 17 00:00:00 2001
From: tspeaks <speaks_tyler@yahoo.com>
Date: Tue, 8 Sep 2026 00:10:17 -0500
Subject: [PATCH 015/249] fix(kernel): order the batch-memcpy probe against the
 current stream

_probe allocates its destination with torch.zeros, which enqueues the fill on
the current stream, then enqueues the verification copy on a fresh probe
stream. Nothing joins the two. When the current stream has a backlog the copy
completes first on the independent probe stream and the fill lands on top of
it, so the probe reads back zeros and load_batch_memcpy raises
"cudaMemcpyBatchAsync probe copied wrong bytes" on a GPU that supports the API.

OffloadMoeCache catches that and falls back to full-layer copies, so
--moe-prefill-hit-d2d silently does nothing whenever the current stream is busy
as the probe runs -- in practice during prefill warmup, which is exactly when
the flag is first exercised.

A cold process hides the bug: the first torch.zeros pays a cudaMalloc and the
first torch.cuda.Stream() populates the per-device stream pool, and each of
those synchronizes the device, draining the backlog before the copy is
enqueued. That is why the probe passes when run standalone and fails inside a
warmed-up server.

Join the probe stream to the current stream before the copy.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 python/freetoken/kernel/batch_memcpy.py |  6 +++
 tests/kernels/test_batch_memcpy.py      | 71 +++++++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 tests/kernels/test_batch_memcpy.py

diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py
index b39e5cde24..67e2bf9870 100644
--- a/python/freetoken/kernel/batch_memcpy.py
+++ b/python/freetoken/kernel/batch_memcpy.py
@@ -28,6 +28,12 @@ def _probe(fn) -> None:
     src = torch.arange(16, dtype=torch.uint8).pin_memory()
     dst = torch.zeros(16, dtype=torch.uint8, device="cuda")
     stream = torch.cuda.Stream()
+    # torch.zeros filled dst on the CURRENT stream; the copy below runs on this one.
+    # Without the join the fill can land after the copy and clobber it, leaving the
+    # probe to read back zeros -- a false "copied wrong bytes" on hardware that
+    # supports the API. A cold process hides this (the first torch.zeros and
+    # torch.cuda.Stream() each synchronize the device); a warmed-up server does not.
+    stream.wait_stream(torch.cuda.current_stream())
     fn(
         torch.tensor([dst.data_ptr()]),
         torch.tensor([src.data_ptr()]),
diff --git a/tests/kernels/test_batch_memcpy.py b/tests/kernels/test_batch_memcpy.py
new file mode 100644
index 0000000000..bc20c13c3c
--- /dev/null
+++ b/tests/kernels/test_batch_memcpy.py
@@ -0,0 +1,71 @@
+"""The cudaMemcpyBatchAsync load probe must be ordered against the current stream.
+
+_probe allocates its destination with torch.zeros (a fill enqueued on the CURRENT
+stream) but enqueues the copy on a fresh probe stream. With nothing joining the two,
+the fill can land *after* the copy and clobber it, so the probe reads back zeros and
+load_batch_memcpy raises "probe copied wrong bytes" on a GPU that supports the API
+perfectly well. It only reproduces once the caching allocator and the stream pool are
+warm -- in a cold process torch.zeros and torch.cuda.Stream() each synchronize the
+device and hide the race, which is why the probe passes standalone but fails inside a
+server that is mid-warmup.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+JIT = pytest.mark.skipif(
+    os.getenv("FREETOKEN_DISABLE_JIT", "").strip().lower() in {"1", "true", "yes", "on"},
+    reason="batch_memcpy has no AOT prebuild; needs runtime JIT",
+)
+
+
+def _cuda_at_least(major: int, minor: int) -> bool:
+    cuda = torch.version.cuda
+    if cuda is None:
+        return False
+    return tuple(int(x) for x in cuda.split(".")[:2]) >= (major, minor)
+
+
+BATCH_API = pytest.mark.skipif(
+    not _cuda_at_least(13, 0), reason="the cudaMemcpyBatchAsync binding needs CUDA >= 13.0"
+)
+
+# ~1s of spin on the current stream. The probe's own copy is microseconds, so any
+# ordering bug shows up as the whole backlog landing on top of the copied bytes.
+_BACKLOG_CYCLES = 1_500_000_000
+
+
+def _warm_device_state() -> None:
+    """Reach the state a server is in at prefill warmup: the caching allocator holds a
+    free block of the probe's 16-byte size and the per-device stream pool is populated,
+    so neither torch.zeros nor torch.cuda.Stream() synchronizes on the probe's behalf."""
+    for _ in range(4):
+        block = torch.zeros(16, dtype=torch.uint8, device="cuda")
+        del block
+        pinned = torch.arange(16, dtype=torch.uint8).pin_memory()
+        del pinned
+    pool = [torch.cuda.Stream() for _ in range(40)]
+    torch.cuda.synchronize()
+    del pool
+
+
+@CUDA
+@JIT
+@BATCH_API
+def test_probe_survives_a_busy_current_stream():
+    from freetoken.kernel.batch_memcpy import load_batch_memcpy
+
+    _warm_device_state()
+    try:
+        torch.cuda._sleep(_BACKLOG_CYCLES)
+        # Guard against a vacuous test: if something drained the queue the race is
+        # not being exercised at all and a pass would mean nothing.
+        assert not torch.cuda.current_stream().query(), "backlog drained before the probe ran"
+        load_batch_memcpy()
+    finally:
+        torch.cuda.synchronize()

From 6c1adb65865c88c7a35b0d82771939a6c9180c68 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 9 Sep 2026 15:05:54 +0300
Subject: [PATCH 016/249] server: key abort claims on aborted uids, deliver
 from the stream's finally

On a disconnect the CancelledError unwinds from the response task's
innermost await -- wait_for_ack -- whose finally pops ack_map and
event_map on the way out. stream_with_cancellation then called
abort_user, whose claim was `event_map.pop(uid) is not None`: always
False by that point, so the AbortMsg was never sent and the scheduler
kept decoding to max_tokens, with no log line either (the cancellation
interrupts the async for, not the is_disconnected branch). A body
closed after the response started arrives as GeneratorExit, which the
`except CancelledError` never saw either. Diagnosed with a live repro
by gdevenyi on the PR thread.

Key the claim on an insertion-ordered aborted_uids dict, FIFO-bounded
like the scheduler's abort tombstones (a duplicate AbortMsg is safe:
the scheduler acks aborts for uids it no longer has), and move delivery
into the stream wrapper's finally, gated on the stream not finishing.
One path now covers server-side cancellation, late closes, and
generator errors; the accounting drain's abort barrier also stops
missing requests whose maps were already emptied.

Assisted-by: Claude
---
 python/freetoken/server/api_server.py    | 35 ++++++++---
 python/freetoken/server/openai_api.py    |  5 +-
 tests/server/test_stream_cancellation.py | 79 ++++++++++++++++++++++++
 3 files changed, 107 insertions(+), 12 deletions(-)

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index add3c66050..9cdb46ee12 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -137,6 +137,13 @@ class FrontendManager:
     initialized: bool = False
     ack_map: Dict[int, List[UserReply]] = field(default_factory=dict)
     event_map: Dict[int, asyncio.Event] = field(default_factory=dict)
+    # Abort claims live HERE, not in the maps above: a disconnect cancels the response
+    # task at its innermost await (inside wait_for_ack, whose finally empties both maps
+    # on the way out), so by the time any abort path runs, a map-based claim would no-op
+    # exactly when the AbortMsg matters most. Insertion-ordered, FIFO-bounded like the
+    # scheduler's abort tombstones; a duplicate AbortMsg for an evicted or finished uid
+    # is safe (the scheduler acks aborts for uids it no longer has).
+    aborted_uids: Dict[int, None] = field(default_factory=dict)
     # Stable identity for this serve process. Generated before the backend is ready so every
     # /health state (loading/ok/error) and /v1/stats can identify the same engine generation.
     instance_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -356,6 +363,7 @@ async def stream_generate(self, uid: int):
         logger.debug("Finished streaming response for user %s", uid)
 
     async def stream_with_cancellation(self, generator, request: Request, uid: int):
+        finished = False
         try:
             async for chunk in generator:
                 # detect if the client has disconnected
@@ -363,18 +371,27 @@ async def stream_with_cancellation(self, generator, request: Request, uid: int):
                     logger.info("Client disconnected for user %s", uid)
                     raise asyncio.CancelledError
                 yield chunk
-        except asyncio.CancelledError:
-            try:
-                await asyncio.shield(self.abort_user(uid))
-            except Exception:  # noqa: BLE001
-                logger.exception("Failed to deliver abort for user %s", uid)
-            raise
+            finished = True
+        finally:
+            # finally, not `except CancelledError`: the abort must go out however the
+            # stream dies -- the server cancelling the response task (the CancelledError
+            # lands on wait_for_ack's await, never on the poll above), a late close
+            # delivering GeneratorExit, or an exception out of the generator. Only a
+            # stream that ran to completion leaves the engine with nothing to stop.
+            if not finished:
+                try:
+                    await asyncio.shield(self.abort_user(uid))
+                except Exception:  # noqa: BLE001
+                    logger.exception("Failed to deliver abort for user %s", uid)
 
     async def abort_user(self, uid: int):
-        claimed = self.event_map.pop(uid, None) is not None
-        self.ack_map.pop(uid, None)
-        if not claimed:
+        if uid in self.aborted_uids:
             return
+        self.aborted_uids[uid] = None
+        while len(self.aborted_uids) > 65_536:
+            self.aborted_uids.pop(next(iter(self.aborted_uids)))
+        self.event_map.pop(uid, None)
+        self.ack_map.pop(uid, None)
         await asyncio.sleep(0.1)
         self.stats.on_abort(uid)
         logger.warning("Aborting request for user %s", uid)
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index f053054841..81fdb3f75e 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -179,9 +179,8 @@ async def _await_watching_disconnect(awaitable, request: Request | None, state:
             logger.exception("Failed to deliver abort for user %s", uid)
         gen.cancel()
         raise
-    # Client gone. Claim + AbortMsg first (abort_user is a no-op once the ack
-    # loop's own cleanup has run), then wind down the drain task; its result is
-    # undeliverable either way.
+    # Client gone. Deliver the AbortMsg first (the decode stops sooner), then
+    # wind down the drain task; its result is undeliverable either way.
     try:
         await asyncio.shield(state.abort_user(uid))
     except Exception:  # noqa: BLE001
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index ded9178bf9..e38fa56bfb 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -23,6 +23,7 @@ def _state(send_impl=None):
     st = SimpleNamespace(
         ack_map={7: [object()]},
         event_map={7: asyncio.Event()},
+        aborted_uids={},
         stats=_Stats(),
     )
     sent = []
@@ -105,6 +106,84 @@ async def run():
     asyncio.run(run())
 
 
+def test_cancellation_after_ack_loop_cleanup_still_aborts():
+    """The production disconnect path: the CancelledError lands on wait_for_ack's
+    innermost await, and its finally empties both maps BEFORE the stream wrapper
+    gets to run. A claim keyed on the maps no-ops exactly then; the AbortMsg must
+    still reach the scheduler (which acks aborts for uids it no longer has)."""
+
+    async def run():
+        state = _state()
+
+        async def drain_like_wait_for_ack():
+            try:
+                await asyncio.sleep(3600)
+                yield b""  # pragma: no cover
+            finally:
+                state.ack_map.pop(7, None)
+                state.event_map.pop(7, None)
+
+        async def consume():
+            state.abort_user = lambda uid: FrontendManager.abort_user(state, uid)
+            async for _ in FrontendManager.stream_with_cancellation(
+                state, drain_like_wait_for_ack(), _Request(), 7
+            ):
+                pass
+
+        task = asyncio.create_task(consume())
+        await asyncio.sleep(0.01)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert state.ack_map == {} and state.event_map == {}
+        assert [type(m) for m in state.sent] == [AbortMsg]
+        assert state.sent[0].uid == 7
+        assert state.stats.aborts == [7]
+
+    asyncio.run(run())
+
+
+def test_late_close_delivers_abort():
+    """A body closed after the response started surfaces as aclose() on the stream
+    wrapper: GeneratorExit, which `except CancelledError` never sees. The finally
+    must deliver the abort for any stream that did not run to completion."""
+
+    async def run():
+        state = _state()
+        state.abort_user = lambda uid: FrontendManager.abort_user(state, uid)
+
+        async def chunks():
+            while True:
+                yield b"data: x\n\n"
+
+        agen = FrontendManager.stream_with_cancellation(state, chunks(), _Request(), 7)
+        assert await agen.__anext__() == b"data: x\n\n"
+        await agen.aclose()
+
+        assert [type(m) for m in state.sent] == [AbortMsg]
+        assert state.sent[0].uid == 7
+        assert state.stats.aborts == [7]
+
+    asyncio.run(run())
+
+
+def test_concurrent_aborts_send_one_message():
+    """The accounting drain gathers abort_user over every in-flight uid and can race
+    a disconnect abort for the same request; the claim set must collapse them."""
+
+    async def run():
+        state = _state()
+        await asyncio.gather(
+            FrontendManager.abort_user(state, 7),
+            FrontendManager.abort_user(state, 7),
+        )
+        assert len(state.sent) == 1
+        assert state.stats.aborts == [7]
+
+    asyncio.run(run())
+
+
 def test_normal_completion_sends_no_abort():
     async def run():
         state = _state()

From 7e72c6153ae827de7b0485ddb0a83b3475ea2e38 Mon Sep 17 00:00:00 2001
From: mac <av_samarin@mail.ru>
Date: Wed, 9 Sep 2026 15:07:45 +0300
Subject: [PATCH 017/249] tests: feed _process_last_data a ForwardOutput-shaped
 fake, not a tuple

The logprobs change moved the drain from positional unpacking to named
ForwardOutput attributes (next_tokens_cpu, copy_done_event, the *_cpu
logprob columns), which the bare-tuple fake in the in-flight-abort
scheduler tests predates: three of them died with 'tuple' object has no
attribute 'next_tokens_cpu' on a GPU run. Caught by gdevenyi on the PR
thread; a macOS run cannot see it (the file already fails on the missing
kernel packages there, identically on both sides of the series).

A namespace mirroring the drain-facing fields keeps the test off the
engine import chain, whose kernel packages a scheduler unit test should
not require.

Assisted-by: Claude
---
 tests/scheduler/test_abort_inflight_prefill.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py
index 1c9a9fba51..d9900eb2f8 100644
--- a/tests/scheduler/test_abort_inflight_prefill.py
+++ b/tests/scheduler/test_abort_inflight_prefill.py
@@ -94,11 +94,19 @@ def _launch_req(pool, cm, tm, prompt, *, cls=Req, track_seqlen=None):
 
 
 def _as_last_data(batch):
-    return (
-        SimpleNamespace(batch=batch),
-        (None, torch.tensor([42], dtype=torch.int32),
-         SimpleNamespace(synchronize=lambda: None)),
+    # Mirrors engine.ForwardOutput's drain-facing fields. _process_last_data reads
+    # attributes (next_tokens_cpu, copy_done_event, the *_cpu logprob columns), so a
+    # bare tuple no longer works; a namespace keeps this test off the engine import
+    # chain, whose kernel packages a scheduler unit test should not require.
+    outputs = SimpleNamespace(
+        next_tokens_gpu=None,
+        next_tokens_cpu=torch.tensor([42], dtype=torch.int32),
+        copy_done_event=SimpleNamespace(synchronize=lambda: None),
+        chosen_logprobs_cpu=None,
+        top_ids_cpu=None,
+        top_logprobs_cpu=None,
     )
+    return (SimpleNamespace(batch=batch), outputs)
 
 
 def test_abort_inflight_final_chunk_marks_then_drains():

From c3d0f975bba332f364c695f4ac146ef46a2fa072 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 14:19:37 -0400
Subject: [PATCH 018/249] feat(qwen4_exp): tensor parallelism for
 Qwen3.8-Flash-Next (offload backend)

Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as
its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert
gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab
rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds
half the experts and each MoE layer needs one all-reduce (routed + gate * shared
are combined before the reduce). Router, QSA indexer, norms, hyper-connections
and PLE stay replicated so all ranks select the same blocks and n-gram rows.

Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and
distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes
apart behind a 100+ GiB load).

Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense
checkpoints raise under TP.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/engine/config.py             |  2 +-
 python/freetoken/layers/linear.py             |  8 +-
 .../layers/quantization/moe/nvfp4.py          | 14 +--
 python/freetoken/models/nvfp4_banks.py        | 44 ++++++++-
 .../freetoken/models/qwen4_exp/attention.py   | 42 ++++++---
 python/freetoken/models/qwen4_exp/gdn.py      | 63 +++++++++----
 python/freetoken/models/qwen4_exp/moe.py      | 31 ++++++-
 python/freetoken/models/qwen4_exp/weight.py   | 69 ++++++++++++--
 tests/models/qwen4_exp/test_tp_shard.py       | 92 +++++++++++++++++++
 tests/models/test_nvfp4_banks_tp.py           | 76 +++++++++++++++
 10 files changed, 392 insertions(+), 49 deletions(-)
 create mode 100644 tests/models/qwen4_exp/test_tp_shard.py
 create mode 100644 tests/models/test_nvfp4_banks_tp.py

diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 7ab792f9e9..e14fd3118d 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -79,7 +79,7 @@ class EngineConfig:
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
-    distributed_timeout: float = 60.0
+    distributed_timeout: float = 1800.0  # ranks reach the first collective minutes apart on a 100+ GiB offload load
     use_dummy_weight: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None
diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py
index d8a3e8619f..40ebe934c0 100644
--- a/python/freetoken/layers/linear.py
+++ b/python/freetoken/layers/linear.py
@@ -83,13 +83,17 @@ def __init__(
         input_size: int,
         output_sizes: List[int],
         has_bias: bool,
+        local_output_sizes: List[int] | None = None,
         *,
         quant_config: QuantConfig | None = None,
         prefix: str = "",
     ):
-        # check that all output sizes are divisible by tp_size
+        # check that all output sizes are divisible by tp_size (a caller that replicates
+        # GQA kv heads across ranks passes the per-rank sizes explicitly)
         tp_info = get_tp_info()
-        tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        if local_output_sizes is None:
+            local_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        tp_output_sizes = local_output_sizes
         output_size = sum(output_sizes)
         tp_output_size = sum(tp_output_sizes)
         super().__init__(
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index bb341cdfb7..dfe35604b3 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -48,7 +48,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
         return f"triton nvfp4 MoE kernel: {reason}" if reason else None
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((2 * i, h // 2), torch.uint8),
             "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8),
@@ -61,7 +61,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
     def pack(self, pieces, cfg: MoEConfig, out):
         out["gate_up"].copy_(fused_piece(pieces, "gate_up"))
         out["gate_up_scale"].copy_(fused_piece(pieces, "gate_up_scale"))
-        out["gate_up_global"].copy_(fused_global(pieces, cfg.intermediate))
+        out["gate_up_global"].copy_(fused_global(pieces, cfg.local_intermediate))
         out["down"].copy_(pieces["down"])
         out["down_scale"].copy_(pieces["down_scale"])
         out["down_global"].copy_(global_rows(pieces["down_global"], cfg.hidden))
@@ -241,7 +241,7 @@ def worth_it(self, cfg: MoEConfig) -> bool:
         return (8, 0) <= backend.device_capability() < (10, 0)
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32),
             "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8),
@@ -252,7 +252,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         }
 
     def pack(self, pieces, cfg: MoEConfig, out):
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         device = torch.device("cuda")
         gu, gus, gug = fused_piece(pieces, "gate_up"), fused_piece(pieces, "gate_up_scale"), fused_global(pieces, i)
         dn, dns, dng = pieces["down"], pieces["down_scale"], global_rows(pieces["down_global"], h)
@@ -503,7 +503,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
 
     def worth_it(self, cfg: MoEConfig) -> bool:
         # NOTE: never auto-selected. flashinfer's cute launcher indexes each bank with int32 element offsets, so the GPU slot cache is capped at (2^31 - 1) / elements-per-slot (about 1000 slots for GLM-5.3-Flash's 12960 experts); until flashinfer lifts that, b12x stays behind triton in the table; the selection does not weigh the cap, cache-auto clamps to slot_limit() and OffloadMoeCache refuses a larger cache.
-        return cfg.intermediate >= B12X_MIN_INTERMEDIATE
+        return cfg.local_intermediate >= B12X_MIN_INTERMEDIATE
 
     def slot_limit(self, cfg: MoEConfig) -> int | None:
         # the cute launcher indexes each bank with int32 element offsets
@@ -512,7 +512,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None:
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         # flashinfer's prepared tiles: the Marlin shapes, byte-identical to the native rows
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32),
             "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8),
@@ -525,7 +525,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
     def pack(self, pieces, cfg: MoEConfig, out):
         from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_w4a16_prepare import prepare_w4a16_packed_weights
 
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         device = torch.device("cuda")
         gu = fused_piece(pieces, "gate_up").to(device)
         gug = fused_global(pieces, i).to(device).float()
diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py
index 1d1b42859c..0feddaf2a7 100644
--- a/python/freetoken/models/nvfp4_banks.py
+++ b/python/freetoken/models/nvfp4_banks.py
@@ -60,6 +60,36 @@ def _bank_layer(spec: Nvfp4ExpertSourceSpec, layer: int, config) -> int | None:
     return bank_layer
 
 
+def _tp_slice() -> tuple[int, int] | None:
+    """``(i_local, i_lo)``: this rank's slice of the expert intermediate axis, or None at TP=1.
+
+    TP shards every routed expert along I (gate/up output rows, down input columns), so each
+    rank caches only its own half and the MoE layer all-reduces the partial sum. The expert
+    kernel sizes its banks from ``MoEConfig.local_intermediate``, which is the same split.
+    """
+    from freetoken.distributed import try_get_tp_info
+
+    tp = try_get_tp_info()
+    if tp is None or tp.size == 1:
+        return None
+    return tp.size, tp.rank
+
+
+def _tp_shard(role: str, tensor: torch.Tensor, inter: int, tp_size: int, tp_rank: int):
+    """This rank's I-slice of one checkpoint expert tensor. ``_global`` is a per-tensor scalar
+    and has no I axis; gate/up carry I on the row axis, down on the column axis (halved for the
+    packed FP4 codes, sixteenthed for the fp8 block scales)."""
+    if role.endswith("_global"):
+        return tensor
+    n = inter // tp_size
+    assert n % 16 == 0, f"NVFP4 TP shard {n} must cover whole 16-wide scale blocks"
+    lo = tp_rank * n
+    if role.startswith("down"):
+        d = 2 if role == "down" else 16
+        return tensor[:, lo // d : (lo + n) // d]
+    return tensor[lo : lo + n]
+
+
 def _kind_suffix(kind: str) -> str:
     return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind]
 
@@ -133,7 +163,19 @@ def _parallel():
                 tensor = _ingest_global(spec, tensor)
             yield name, tensor
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    stream = _parallel() if parallel else _serial()
+    shard = _tp_slice()
+    if shard is not None:
+        tp_size, tp_rank = shard
+        inter = config.moe_intermediate_size
+        base = stream
+
+        def _sharded():
+            for name, tensor in base:
+                yield name, _tp_shard(wanted[name][2], tensor, inter, tp_size, tp_rank)
+
+        stream = _sharded()
+    return per_expert_pieces(stream, wanted.get, tensors_per_expert=9)
 
 
 __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]
diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py
index 66211c7c29..651da87d93 100644
--- a/python/freetoken/models/qwen4_exp/attention.py
+++ b/python/freetoken/models/qwen4_exp/attention.py
@@ -19,9 +19,16 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.distributed import get_tp_info
+from freetoken.layers import (
+    BaseOP,
+    GemmaPlusOneRMSNorm,
+    LinearColParallelMerged,
+    LinearOProj,
+    LinearReplicated,
+)
 from freetoken.layers.rotary import get_rope
-from freetoken.utils import nvtx_annotate
+from freetoken.utils import div_even, nvtx_annotate
 
 if TYPE_CHECKING:
     from freetoken.core import Batch
@@ -123,12 +130,25 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
         self.head_dim = config.head_dim
         self.qo_attn_dim = self.num_q * self.head_dim
         self.kv_attn_dim = self.num_kv * self.head_dim
-        self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
+        # TP: q heads split across ranks, kv heads split or (num_kv < tp) replicated; the
+        # indexer stays replicated so every rank selects the same blocks.
+        tp = get_tp_info()
+        self._local_num_q = div_even(self.num_q, tp.size)
+        self._local_num_kv = div_even(self.num_kv, tp.size, allow_replicate=True)
+        self._local_qo_dim = self._local_num_q * self.head_dim
+        self._local_kv_dim = self._local_num_kv * self.head_dim
+        self._qkv_split = [self._local_qo_dim * 2, self._local_kv_dim, self._local_kv_dim]
         self.qkv_proj = LinearColParallelMerged(
-            config.hidden_size, self._qkv_split, has_bias=False,
+            config.hidden_size,
+            [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim],
+            has_bias=False,
+            local_output_sizes=self._qkv_split,
             quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
         )
-        self.o_proj = LinearReplicated(
+        # row-parallel, NOT LinearReplicated: under TP>1 qkv_proj is column-parallel, so this
+        # rank's attention output is its own head slice and o_proj must consume the sharded
+        # input dim and all-reduce the partial sums. At TP=1 the two are equivalent.
+        self.o_proj = LinearOProj(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )
@@ -147,21 +167,21 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
     @nvtx_annotate("QSA")
     def forward(self, x: torch.Tensor, batch: Batch) -> torch.Tensor:
         qg, k, v = self.qkv_proj.forward(x).split(self._qkv_split, dim=-1)
-        qg = qg.view(-1, self.num_q, self.head_dim * 2)
+        qg = qg.view(-1, self._local_num_q, self.head_dim * 2)
         q = qg[..., : self.head_dim].contiguous()
-        gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim)
-        k = k.contiguous().view(-1, self.num_kv, self.head_dim)
+        gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_dim)
+        k = k.contiguous().view(-1, self._local_num_kv, self.head_dim)
         v = v.contiguous()
         self.q_norm.forward_inplace(q)
         self.k_norm.forward_inplace(k)
         q, k = self.rotary.forward(
-            batch.positions, q.view(-1, self.qo_attn_dim), k.view(-1, self.kv_attn_dim)
+            batch.positions, q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim)
         )
         index = self.indexer.forward(x)
         o = get_global_ctx().attn_backend.qsa_forward(
-            q.view(-1, self.num_q, self.head_dim), k, v, index, self.layer_id, batch
+            q.view(-1, self._local_num_q, self.head_dim), k, v, index, self.layer_id, batch
         )
-        gated = o.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate)
+        gated = o.reshape(-1, self._local_qo_dim) * torch.sigmoid(gate)
         return self.o_proj.forward(gated)
 
 
diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py
index ed49dd42f0..3b2e029340 100644
--- a/python/freetoken/models/qwen4_exp/gdn.py
+++ b/python/freetoken/models/qwen4_exp/gdn.py
@@ -3,9 +3,12 @@
 import torch
 import torch.nn.functional as F
 from freetoken.core import get_global_ctx
+from freetoken.distributed import get_tp_info
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
-from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.distributed import get_tp_info
+from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.quantization import QuantConfig
+from freetoken.utils import div_even
 from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
 
 
@@ -50,12 +53,29 @@ def __init__(
         self.value_dim = num_v_heads * head_v_dim
         self.conv_dim = 2 * self.key_dim + self.value_dim
         self.conv_kernel_size = conv_kernel_size
+        # TP-local head counts: k heads and their v-head groups split evenly across ranks
+        # (the fla kernels take the GQA ratio from the shapes); the state pool is sharded the
+        # same way (kvcache.linear_state_pool._linear_local_dims).
+        tp = get_tp_info()
+        self._local_num_k_heads = div_even(num_k_heads, tp.size, allow_replicate=True)
+        self._local_num_v_heads = div_even(num_v_heads, tp.size, allow_replicate=True)
+        self._local_key_dim = self._local_num_k_heads * head_k_dim
+        self._local_value_dim = self._local_num_v_heads * head_v_dim
+        self._local_conv_dim = 2 * self._local_key_dim + self._local_value_dim
         # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM)
         self._split_in_proj = (
             quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None
         )
 
-        self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads]
+        self._in_proj_split = [
+            self._local_conv_dim, self._local_value_dim,
+            self._local_num_v_heads, self._local_num_v_heads,
+        ]
+        # The split (quantized) branch has no TP-aware variant: its two linears are built from
+        # the checkpoint's own scheme and are not sharded here.
+        assert not (tp.size > 1 and self._split_in_proj), (
+            "qwen4_exp TP shards bf16 GDN projections only"
+        )
         if self._split_in_proj:
             self.in_proj_qkvz = LinearColParallelMerged(
                 hidden_size, [self.conv_dim, self.value_dim], has_bias=False,
@@ -68,18 +88,23 @@ def __init__(
         else:
             # Fused input projection (one GEMM instead of four): qkv | z | b | a.
             self.in_proj = LinearColParallelMerged(
-                hidden_size, self._in_proj_split, has_bias=False,
+                hidden_size, [self.conv_dim, self.value_dim, num_v_heads, num_v_heads],
+                has_bias=False, local_output_sizes=self._in_proj_split,
                 quant_config=quant_config, prefix=f"{prefix}.in_proj",
             )
-        self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size)
+        self.conv1d = _DepthwiseConv1d(self._local_conv_dim, conv_kernel_size)
         # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive,
         # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a
         # per-call .float() upcast in the decode wrapper. The weight loader exempts
         # *.A_log / *.dt_bias from the model-dtype downcast.
-        self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32)
-        self.A_log = torch.empty(num_v_heads, dtype=torch.float32)
+        self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32)
+        self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32)
         self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps, activation=output_gate)
-        self.out_proj = LinearReplicated(
+        # Row-parallel over the local v heads, all-reduce inside. At TP=1 this is exactly
+        # LinearReplicated (div_even(x, 1) == x and the reduction is skipped), so one class
+        # covers both -- and under TP>1 a replicated o_proj would silently take the wrong
+        # input width and skip the reduction.
+        self.out_proj = LinearOProj(
             self.value_dim, hidden_size, has_bias=False,
             quant_config=quant_config, prefix=f"{prefix}.out_proj",
         )
@@ -148,7 +173,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         else:
             proj = self.in_proj.forward(hidden_states)
             conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1)
-        z = z.reshape(total, self.num_v_heads, self.head_v_dim)
+        nk, nv = self._local_num_k_heads, self._local_num_v_heads
+        kd, vd = self._local_key_dim, self._local_value_dim
+        z = z.reshape(total, nv, self.head_v_dim)
         li = pool.local_index(self.layer_id)
 
         if batch.is_decode:
@@ -157,10 +184,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA).
             mixed = self._conv_decode(conv_in, fla.cache_indices, pool)  # [B, conv_dim]
             B = mixed.shape[0]
-            qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-            q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
-            k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
-            v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype)
+            qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1)
+            q = qf.reshape(1, B, nk, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, B, nk, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, B, nv, self.head_v_dim).to(dtype)
             core_out = gdn_decode_fla(
                 q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias,
                 state_source=pool.recurrent_states[li], indices=fla.cache_indices,
@@ -170,13 +197,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             mixed = self._conv_prefill(
                 conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
-            qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-            q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)
-            k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)
-            v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype)
+            qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1)
+            q = qf.reshape(1, total, nk, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, total, nk, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, total, nv, self.head_v_dim).to(dtype)
             g, beta = self._gate_params(a, b)
-            g = g.reshape(1, total, self.num_v_heads)
-            beta = beta.float().reshape(1, total, self.num_v_heads)
+            g = g.reshape(1, total, nv)
+            beta = beta.float().reshape(1, total, nv)
             # The chunk kernel reads + writes back initial_state[cache_indices] in place;
             # fresh sequences (cached_len==0) must start from a zeroed slot.
             if fla.fresh_state_indices is not None:
diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py
index 3687f28035..7d9f6adf7a 100644
--- a/python/freetoken/models/qwen4_exp/moe.py
+++ b/python/freetoken/models/qwen4_exp/moe.py
@@ -3,7 +3,12 @@
 from typing import TYPE_CHECKING
 
 import torch
+import torch.nn.functional as F
+from freetoken.core import get_global_ctx
+from freetoken.distributed import DistributedCommunicator, get_tp_info
 from freetoken.kernel.triton.moe_shared_gate import shared_gate_mul_add, shared_gate_sigmoid
+from freetoken.layers import LinearRowParallel, silu_and_mul
+from freetoken.layers.moe import OffloadMoELayer
 from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE
 
 if TYPE_CHECKING:
@@ -14,15 +19,37 @@ class Qwen4ExpMoE(Qwen3_5MoE):
     """Qwen3_5MoE with the shared-expert gate on triton instead of gemv + sigmoid + mul + add.
 
     Same weights, same state dict. The gate reduction stays ahead of the routed experts, which may write into ``hidden_states`` in place.
+
+    TP: the offload experts are sharded along the intermediate axis and the bf16 shared expert is
+    row-parallel, so both produce partial sums; ``routed + gate * shared`` is linear in them and
+    is reduced once (one all-reduce per MoE layer instead of two).
     """
 
+    def __init__(self, config: ModelConfig, layer_id: int | None = None) -> None:
+        super().__init__(config, layer_id=layer_id)
+        self._comm = DistributedCommunicator()
+        self._tp_size = get_tp_info().size
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         router_logits = self.gate.forward(hidden_states)
-        shared = self.shared_expert.forward(hidden_states)
         gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1))
-        routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits)
+        se, ex = self.shared_expert, self.experts
+        if (
+            self._tp_size > 1
+            and isinstance(se.down_proj, LinearRowParallel)
+            and isinstance(ex, OffloadMoELayer)
+        ):
+            shared = F.linear(silu_and_mul(se.gate_up_proj.forward(hidden_states)), se.down_proj.weight)
+            if get_global_ctx().batch.is_prefill:
+                routed = ex.prefill_forward(hidden_states, router_logits)
+            else:
+                routed = ex.decode_forward(hidden_states, router_logits)
+            out = self._comm.all_reduce(shared_gate_mul_add(routed, shared, gate))
+            return out.view(num_tokens, hidden_dim)
+        shared = se.forward(hidden_states)
+        routed = ex.forward(hidden_states=hidden_states, router_logits=router_logits)
         return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim)
 
 
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 2ce7462b82..97198f408f 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -21,12 +21,12 @@
 import safetensors
 import torch
 from freetoken.distributed import get_tp_info
-from freetoken.models.loader import drop_page_cache, iter_weight_files
+from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor
 from freetoken.models.nvfp4_banks import (
     Nvfp4ExpertSourceSpec,
 )
 from freetoken.moe.host_banks import HostBank, read_range_into
-from freetoken.utils import download_hf_weight
+from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight
 from freetoken.utils.progress import byte_bar
 from tqdm import tqdm
 
@@ -137,6 +137,58 @@ def _try_fuse(
     return None
 
 
+def _shard_rows(
+    t: torch.Tensor, parts: list[tuple[int, int]], rank: int, world: int
+) -> torch.Tensor:
+    """Column-parallel slice (dim 0) of a ``[part0 | part1 | ...]`` fusion; ``parts`` gives each
+    part as ``(heads, rows_per_head)``. Heads split evenly across ranks; a part with fewer heads
+    than ranks (GQA kv) replicates head ``rank * heads // world``, the ``div_even(...,
+    allow_replicate=True)`` convention of the TP-aware layers."""
+    out, off = [], 0
+    for heads, rows in parts:
+        local = div_even(heads, world, allow_replicate=True)
+        first = rank * heads // world
+        out.append(t[off + first * rows : off + (first + local) * rows])
+        off += heads * rows
+    assert off == t.shape[0], f"fusion parts {parts} cover {off} rows, tensor has {t.shape[0]}"
+    return torch.cat(out, dim=0)
+
+
+def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.Tensor:
+    """TP shard of one state-dict tensor (fused projections included); identity at TP=1.
+
+    Column-parallel (dim 0, by head): attention ``qkv_proj`` [q|gate per head | k | v], GDN
+    ``in_proj`` [q | k | v | z | b | a] and the matching ``conv1d`` channels, ``A_log`` /
+    ``dt_bias``, shared-expert ``gate_up_proj``. Row-parallel (dim 1): ``o_proj``,
+    ``out_proj``, shared-expert ``down_proj``. Vocab rows: ``embed_tokens`` / ``lm_head``.
+    Everything else (router, indexer, norms, HC, PLE, shared-expert gate) is replicated.
+    """
+    if world == 1:
+        return t
+    if name.endswith(".self_attn.qkv_proj.weight"):
+        q = (config.num_qo_heads, 2 * config.head_dim)
+        kv = (config.num_kv_heads, config.head_dim)
+        return _shard_rows(t, [q, kv, kv], rank, world)
+    if ".linear_attn." in name:
+        g = config.linear_attention_group()
+        k = (g.num_key_heads, g.key_head_dim)
+        v = (g.num_value_heads, g.value_head_dim)
+        if name.endswith(".in_proj.weight"):
+            return _shard_rows(t, [k, k, v, v, (v[0], 1), (v[0], 1)], rank, world)
+        if name.endswith(".conv1d.weight"):
+            return _shard_rows(t, [k, k, v], rank, world)
+        if name.endswith((".A_log", ".dt_bias")):
+            return _shard_rows(t, [(v[0], 1)], rank, world)
+        if name.endswith(".out_proj.weight"):
+            return t.chunk(world, dim=1)[rank].clone()
+        return t
+    if name.endswith(".shared_expert.gate_up_proj.weight"):
+        half = t.shape[0] // 2
+        return _shard_rows(t, [(half, 1), (half, 1)], rank, world)
+    # o_proj / down_proj: dim 1; embed_tokens / lm_head: vocab rows; others unchanged.
+    return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None)
+
+
 def iter_weights(
     model_path: str,
     device: torch.device,
@@ -158,16 +210,18 @@ def iter_weights(
     ``include_moe_experts`` is accepted for the loader contract but never yields anything: the
     routed experts are NVFP4 and always come from the offload cache's expert reader.
     """
-    if get_tp_info().size > 1:
-        raise NotImplementedError("qwen4_exp weight loading supports TP=1 only")
     if not include_non_moe:
         return
 
+    from .config import parse_config
+
+    tp = get_tp_info()
+    config = parse_config(cached_load_hf_config(model_path))
     fuse_buf: dict[str, dict[int, torch.Tensor]] = {}
     for file in tqdm(
         iter_weight_files(model_path),
         desc="Loading weights",
-        disable=not get_tp_info().is_primary(),
+        disable=not tp.is_primary(),
     ):
         with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
             for raw_name in f.keys():
@@ -178,9 +232,10 @@ def iter_weights(
                 fused = _try_fuse(name, tensor, fuse_buf)
                 if fused is not None:
                     if fused != ():  # () means buffered, not yet complete
-                        yield fused
+                        name, tensor = fused
+                        yield name, _shard(name, tensor, config, tp.rank, tp.size)
                     continue
-                yield name, tensor
+                yield name, _shard(name, tensor, config, tp.rank, tp.size)
 
     assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}"
 
diff --git a/tests/models/qwen4_exp/test_tp_shard.py b/tests/models/qwen4_exp/test_tp_shard.py
new file mode 100644
index 0000000000..87a305ef0c
--- /dev/null
+++ b/tests/models/qwen4_exp/test_tp_shard.py
@@ -0,0 +1,92 @@
+"""TP sharding of the qwen4_exp dense weights (pure tensor math, no TP runtime needed)."""
+
+from types import SimpleNamespace
+
+import torch
+from freetoken.models.qwen4_exp.weight import _shard, _shard_rows
+
+
+def _cfg():
+    g = SimpleNamespace(
+        num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8
+    )
+    return SimpleNamespace(
+        num_qo_heads=6, num_kv_heads=2, head_dim=8, linear_attention_group=lambda: g
+    )
+
+
+def _gather(name, t, cfg, world, dim=0):
+    return torch.cat([_shard(name, t, cfg, r, world) for r in range(world)], dim=dim)
+
+
+def test_shard_rows_splits_each_part_by_head():
+    t = torch.arange(4 * 4 + 2 * 4).reshape(
+        -1, 1
+    )  # part A: 4 heads x 4 rows, part B: 2 heads x 4
+    s = [_shard_rows(t, [(4, 4), (2, 4)], r, 4) for r in range(4)]
+    assert all(x.shape[0] == 8 for x in s)
+    assert s[0][:4].flatten().tolist() == [0, 1, 2, 3]
+    assert s[3][:4].flatten().tolist() == [12, 13, 14, 15]
+    # 2 kv heads over 4 ranks replicate: rank * heads // world -> heads 0, 0, 1, 1
+    assert s[0][4:].equal(s[1][4:]) and s[2][4:].flatten().tolist() == [20, 21, 22, 23]
+
+
+def test_shard_round_trips_the_fused_projections():
+    cfg, world = _cfg(), 2
+    qkv = torch.randn(6 * 16 + 2 * 8 + 2 * 8, 5)
+    assert (
+        _gather("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, world).shape
+        == qkv.shape
+    )
+    q0 = _shard("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, 0, world)
+    assert q0.shape[0] == 3 * 16 + 8 + 8
+    torch.testing.assert_close(q0[:48], qkv[:48])  # q heads 0-2 (16 rows each: q|gate)
+    torch.testing.assert_close(q0[48:56], qkv[96:104])  # kv head 0 of k
+    torch.testing.assert_close(q0[56:64], qkv[112:120])  # kv head 0 of v
+    kd, vd, nv = 4 * 8, 12 * 8, 12
+    in_proj = torch.randn(kd + kd + vd + vd + nv + nv, 5)
+    s = _shard("model.layers.1.linear_attn.in_proj.weight", in_proj, cfg, 1, world)
+    assert s.shape[0] == (kd + kd + vd + vd + nv + nv) // 2
+    torch.testing.assert_close(s[:16], in_proj[16:32])  # q: k heads 2,3
+    torch.testing.assert_close(s[-6:], in_proj[-6:])  # a: v heads 6-11
+    conv = torch.randn(kd + kd + vd, 1, 4)
+    assert _shard(
+        "model.layers.1.linear_attn.conv1d.weight", conv, cfg, 0, world
+    ).shape == ((kd + kd + vd) // 2, 1, 4)
+    a_log = torch.randn(nv)
+    torch.testing.assert_close(
+        _shard("model.layers.1.linear_attn.A_log", a_log, cfg, 1, world), a_log[6:]
+    )
+
+
+def test_shard_row_parallel_and_vocab_and_replicated():
+    cfg, world = _cfg(), 2
+    for name in (
+        "model.layers.0.self_attn.o_proj.weight",
+        "model.layers.1.linear_attn.out_proj.weight",
+        "model.layers.0.mlp.shared_expert.down_proj.weight",
+    ):
+        t = torch.randn(3, 8)
+        torch.testing.assert_close(_gather(name, t, cfg, world, dim=1), t)
+    gate_up = torch.randn(2 * 6, 3)
+    s1 = _shard(
+        "model.layers.0.mlp.shared_expert.gate_up_proj.weight", gate_up, cfg, 1, world
+    )
+    torch.testing.assert_close(s1, torch.cat([gate_up[3:6], gate_up[9:12]]))
+    emb = torch.randn(10, 3)
+    torch.testing.assert_close(
+        _gather("model.embed_tokens.weight", emb, cfg, world), emb
+    )
+    torch.testing.assert_close(_gather("lm_head.weight", emb, cfg, world), emb)
+    for name in (
+        "model.layers.0.mlp.gate.weight",
+        "model.layers.0.self_attn.indexer.index_qk_proj.weight",
+        "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight",
+        "model.layers.1.ple.value_proj.weight",
+        "model.layers.0.mlp.shared_expert_gate.weight",
+    ):
+        t = torch.randn(4, 6)
+        assert _shard(name, t, cfg, 1, world).equal(t)
+    assert _shard("model.layers.0.self_attn.qkv_proj.weight", gate_up, cfg, 0, 1).equal(
+        gate_up
+    )
diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py
new file mode 100644
index 0000000000..13bf2a2c72
--- /dev/null
+++ b/tests/models/test_nvfp4_banks_tp.py
@@ -0,0 +1,76 @@
+"""TP sharding of the NVFP4 expert source banks: the two ranks' banks concatenated along the
+intermediate axis must equal the unsharded placement."""
+
+import torch
+from freetoken.distributed.info import DistributedInfo
+from freetoken.models import nvfp4_banks
+from freetoken.models.nvfp4_banks import _alloc_nvfp4_host_banks, _Placer
+
+E, H, INTER = 2, 64, 32
+
+
+def _place(monkeypatch, rank, world):
+    monkeypatch.setattr(
+        nvfp4_banks, "get_tp_info", lambda: DistributedInfo(rank, world)
+    )
+    hb = _alloc_nvfp4_host_banks(1, E, H, INTER // world)
+    banks = {name: [b.tensor for b in layers] for name, layers in hb.items()}
+    placer = _Placer(banks, INTER)
+    torch.manual_seed(0)
+    for e in range(E):
+        for role in ("gate", "up"):
+            placer.put(
+                0,
+                e,
+                role,
+                "weight",
+                torch.randint(0, 255, (INTER, H // 2), dtype=torch.uint8),
+            )
+            placer.put(
+                0,
+                e,
+                role,
+                "weight_scale",
+                torch.randn(INTER, H // 16).to(torch.float8_e4m3fn),
+                torch.tensor(0.5 + e, dtype=torch.float16),
+            )
+        placer.put(
+            0,
+            e,
+            "down",
+            "weight",
+            torch.randint(0, 255, (H, INTER // 2), dtype=torch.uint8),
+        )
+        placer.put(
+            0,
+            e,
+            "down",
+            "weight_scale",
+            torch.randn(H, INTER // 16).to(torch.float8_e4m3fn),
+            torch.tensor(2.0 + e, dtype=torch.float16),
+        )
+    return {k: v[0] for k, v in banks.items()}
+
+
+def test_rank_banks_concatenate_to_the_full_placement(monkeypatch):
+    full = _place(monkeypatch, 0, 1)
+    r0, r1 = _place(monkeypatch, 0, 2), _place(monkeypatch, 1, 2)
+    n = INTER // 2
+    for name in (
+        "gate_up_packed",
+        "gate_up_scale",
+        "gate_up_global",
+    ):  # rows: [gate I | up I]
+        gate = torch.cat([r0[name][:, :n], r1[name][:, :n]], dim=1)
+        up = torch.cat([r0[name][:, n:], r1[name][:, n:]], dim=1)
+        assert torch.equal(
+            torch.cat([gate, up], dim=1).view(torch.uint8), full[name].view(torch.uint8)
+        ), name
+    for name in ("down_packed", "down_scale"):  # columns
+        assert torch.equal(
+            torch.cat([r0[name], r1[name]], dim=2).view(torch.uint8),
+            full[name].view(torch.uint8),
+        ), name
+    assert torch.equal(r0["down_global"], full["down_global"]) and torch.equal(
+        r1["down_global"], full["down_global"]
+    )

From 96e190c3685b5a5fc1f3ac6c75983cb1c5921598 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 15:09:33 -0400
Subject: [PATCH 019/249] fix(qwen4_exp): load the HF config for sharding only
 when TP > 1

tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose
config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/weight.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 97198f408f..905d32d89a 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -216,7 +216,9 @@ def iter_weights(
     from .config import parse_config
 
     tp = get_tp_info()
-    config = parse_config(cached_load_hf_config(model_path))
+    # The sharding geometry needs the HF config; TP=1 never shards, so keep the plain path free
+    # of a config load (synthetic test checkpoints carry no model_type).
+    config = parse_config(cached_load_hf_config(model_path)) if tp.size > 1 else None
     fuse_buf: dict[str, dict[int, torch.Tensor]] = {}
     for file in tqdm(
         iter_weight_files(model_path),

From a99433c683a3ed214d5baabc1d9b041799d17295 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 14:55:22 -0400
Subject: [PATCH 020/249] test(qwen4_exp): TP shard test against the new
 expert-piece stream

---
 tests/models/test_nvfp4_banks_tp.py | 124 +++++++++++++---------------
 1 file changed, 56 insertions(+), 68 deletions(-)

diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py
index 13bf2a2c72..b943cff1c0 100644
--- a/tests/models/test_nvfp4_banks_tp.py
+++ b/tests/models/test_nvfp4_banks_tp.py
@@ -1,76 +1,64 @@
-"""TP sharding of the NVFP4 expert source banks: the two ranks' banks concatenated along the
-intermediate axis must equal the unsharded placement."""
+"""TP sharding of the NVFP4 routed-expert stream.
 
+Every rank reads the same checkpoint tensors and keeps only its own slice of the expert
+intermediate axis, so the ranks' slices must partition the original exactly -- the MoE layer
+all-reduces the partial sums, which is only correct if nothing is dropped or double-counted.
+The expert kernel sizes its banks from ``MoEConfig.local_intermediate``, the same split.
+"""
+
+import pytest
 import torch
-from freetoken.distributed.info import DistributedInfo
-from freetoken.models import nvfp4_banks
-from freetoken.models.nvfp4_banks import _alloc_nvfp4_host_banks, _Placer
+from freetoken.layers.quantization.moe.base import MoEConfig
+from freetoken.models.nvfp4_banks import _tp_shard
 
-E, H, INTER = 2, 64, 32
+INTER, HIDDEN = 640, 2560
 
 
-def _place(monkeypatch, rank, world):
-    monkeypatch.setattr(
-        nvfp4_banks, "get_tp_info", lambda: DistributedInfo(rank, world)
-    )
-    hb = _alloc_nvfp4_host_banks(1, E, H, INTER // world)
-    banks = {name: [b.tensor for b in layers] for name, layers in hb.items()}
-    placer = _Placer(banks, INTER)
-    torch.manual_seed(0)
-    for e in range(E):
-        for role in ("gate", "up"):
-            placer.put(
-                0,
-                e,
-                role,
-                "weight",
-                torch.randint(0, 255, (INTER, H // 2), dtype=torch.uint8),
-            )
-            placer.put(
-                0,
-                e,
-                role,
-                "weight_scale",
-                torch.randn(INTER, H // 16).to(torch.float8_e4m3fn),
-                torch.tensor(0.5 + e, dtype=torch.float16),
-            )
-        placer.put(
-            0,
-            e,
-            "down",
-            "weight",
-            torch.randint(0, 255, (H, INTER // 2), dtype=torch.uint8),
-        )
-        placer.put(
-            0,
-            e,
-            "down",
-            "weight_scale",
-            torch.randn(H, INTER // 16).to(torch.float8_e4m3fn),
-            torch.tensor(2.0 + e, dtype=torch.float16),
+def _roles():
+    """One expert's checkpoint tensors, in the modelopt layout."""
+    return {
+        "gate": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2),
+        "gate_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16),
+        "gate_global": torch.tensor([2.5]),
+        "up": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2),
+        "up_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16),
+        "up_global": torch.tensor([2.5]),
+        "down": torch.arange(HIDDEN * (INTER // 2), dtype=torch.uint8).reshape(HIDDEN, INTER // 2),
+        "down_scale": torch.arange(HIDDEN * (INTER // 16), dtype=torch.uint8).reshape(HIDDEN, INTER // 16),
+        "down_global": torch.tensor([3.5]),
+    }
+
+
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+def test_rank_slices_partition_the_original_tensor(tp_size):
+    full = _roles()
+    for role, tensor in full.items():
+        parts = [_tp_shard(role, tensor, INTER, tp_size, r) for r in range(tp_size)]
+        if role.endswith("_global"):
+            # per-tensor scalar: no I axis, every rank keeps it whole
+            assert all(torch.equal(p, tensor) for p in parts)
+            continue
+        axis = 1 if role.startswith("down") else 0
+        assert torch.equal(torch.cat(parts, dim=axis), tensor), role
+        assert all(p.shape[axis] == tensor.shape[axis] // tp_size for p in parts), role
+
+
+def test_tp1_is_the_identity():
+    for role, tensor in _roles().items():
+        assert _tp_shard(role, tensor, INTER, 1, 0) is tensor or torch.equal(
+            _tp_shard(role, tensor, INTER, 1, 0), tensor
         )
-    return {k: v[0] for k, v in banks.items()}
 
 
-def test_rank_banks_concatenate_to_the_full_placement(monkeypatch):
-    full = _place(monkeypatch, 0, 1)
-    r0, r1 = _place(monkeypatch, 0, 2), _place(monkeypatch, 1, 2)
-    n = INTER // 2
-    for name in (
-        "gate_up_packed",
-        "gate_up_scale",
-        "gate_up_global",
-    ):  # rows: [gate I | up I]
-        gate = torch.cat([r0[name][:, :n], r1[name][:, :n]], dim=1)
-        up = torch.cat([r0[name][:, n:], r1[name][:, n:]], dim=1)
-        assert torch.equal(
-            torch.cat([gate, up], dim=1).view(torch.uint8), full[name].view(torch.uint8)
-        ), name
-    for name in ("down_packed", "down_scale"):  # columns
-        assert torch.equal(
-            torch.cat([r0[name], r1[name]], dim=2).view(torch.uint8),
-            full[name].view(torch.uint8),
-        ), name
-    assert torch.equal(r0["down_global"], full["down_global"]) and torch.equal(
-        r1["down_global"], full["down_global"]
-    )
+def test_a_shard_that_splits_a_scale_block_is_rejected():
+    # scales cover 16 values; a rank slice that is not a whole number of blocks cannot be
+    # sliced consistently across the weight and its scale
+    with pytest.raises(AssertionError, match="16-wide scale blocks"):
+        _tp_shard("gate", torch.zeros(24, 8, dtype=torch.uint8), 24, 3, 0)
+
+
+def test_moe_config_local_intermediate_matches_the_shard():
+    for tp_size in (1, 2, 4):
+        cfg = MoEConfig(num_experts=8, hidden=HIDDEN, intermediate=INTER, top_k=2, tp_size=tp_size)
+        assert cfg.local_intermediate == INTER // tp_size
+        assert _tp_shard("gate", _roles()["gate"], INTER, tp_size, 0).shape[0] == cfg.local_intermediate

From be41deae6bff788f3392e45477452e450720c834 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 18:55:34 -0400
Subject: [PATCH 021/249] feat(qwen4_exp): load-time per-tensor FP8 dense
 projections (W8A8 via _scaled_mm)

Opt-in with FREETOKEN_FP8_DENSE=1 on a bf16-dense checkpoint (e.g. the
RadixArk NVFP4 build): the weight reader quantizes qkv_proj / o_proj, GDN
in_proj (q|k|v|z; the b|a gate rows stay bf16 as in_proj_ba) and out_proj to
per-tensor e4m3 after TP sharding, and layers/fp8_dynamic.py runs them as
cuBLASLt W8A8 GEMMs with a dynamic per-tensor activation scale (one fused
Triton launch at decode sizes; no host sync, CUDA-graph safe). Column-merged
and row-parallel variants, so it works at TP>1.

Why: on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) these projections are
2.67 GB of the ~4 GB a TP=2 rank reads per token; bf16 cuBLAS takes
3.2-3.4 ms per step per rank, raw _scaled_mm 1.9 ms, while the existing
Triton FP8 kernels are slower than bf16 there (measured, weights rotated
past the L2).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/layers/fp8_dynamic.py        | 154 ++++++++++++++++++
 python/freetoken/models/config.py             |   7 +
 .../freetoken/models/qwen4_exp/attention.py   |  37 +++--
 python/freetoken/models/qwen4_exp/config.py   |   8 +
 python/freetoken/models/qwen4_exp/gdn.py      |  48 ++++--
 python/freetoken/models/qwen4_exp/model.py    |   1 +
 python/freetoken/models/qwen4_exp/weight.py   |  67 +++++++-
 tests/models/qwen4_exp/test_fp8_dense.py      | 121 ++++++++++++++
 8 files changed, 409 insertions(+), 34 deletions(-)
 create mode 100644 python/freetoken/layers/fp8_dynamic.py
 create mode 100644 tests/models/qwen4_exp/test_fp8_dense.py

diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py
new file mode 100644
index 0000000000..1a956fbabf
--- /dev/null
+++ b/python/freetoken/layers/fp8_dynamic.py
@@ -0,0 +1,154 @@
+"""Load-time per-tensor FP8 dense linear (W8A8 through cuBLASLt ``torch._scaled_mm``).
+
+The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``),
+produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding).
+The activation is quantized per call with a dynamic per-tensor scale: one fused Triton launch
+(amax pass, then cast) at decode sizes, a torch reduction plus a cast kernel above that. No
+host sync anywhere (the scales stay on the device), so the decode path is CUDA-graph safe;
+the branch between the two paths is on the tensor *shape*, never on its values.
+
+Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per
+decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16.
+Requires sm_89+ (``_scaled_mm``'s floor; Ampere has no FP8 tensor cores).
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+from freetoken.distributed import DistributedCommunicator, get_tp_info
+from freetoken.utils import div_even
+
+from .base import BaseOP
+
+FP8 = torch.float8_e4m3fn
+E4M3_MAX = 448.0
+_FUSED_MAX_ELEMENTS = (
+    65536  # one program handles the whole tensor below this (decode sizes)
+)
+_SCALE_FLOOR = 1e-12  # an all-zero activation (graph warmup buffers) must not give 1/0
+
+
+@triton.jit
+def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
+    """One program: max|x| over the tensor, then the cast. Decode-sized inputs only."""
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for start in range(0, n, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32)
+        acc = tl.maximum(acc, tl.abs(v))
+    amax = tl.maximum(tl.max(acc, axis=0), 1e-12)
+    tl.store(scale_ptr, amax / 448.0)
+    inv = 448.0 / amax
+    for start in range(0, n, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        mask = offs < n
+        v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
+        v = tl.minimum(tl.maximum(v, -448.0), 448.0)
+        tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
+
+
+@triton.jit
+def _quant_cast_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
+    """Cast under a scale already on the device (prefill sizes; the amax is a torch reduction)."""
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    inv = 1.0 / tl.load(scale_ptr)
+    v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
+    v = tl.minimum(tl.maximum(v, -448.0), 448.0)
+    tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
+
+
+def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``."""
+    n = x.numel()
+    out = torch.empty_like(x, dtype=FP8)
+    if n <= _FUSED_MAX_ELEMENTS:
+        scale = torch.empty((), dtype=torch.float32, device=x.device)
+        _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=4096, num_warps=8)
+        return out, scale
+    amax = torch.linalg.vector_norm(x, ord=float("inf")).float()
+    scale = amax.clamp_min_(_SCALE_FLOOR).div_(E4M3_MAX)
+    _quant_cast_kernel[(triton.cdiv(n, 4096),)](
+        x, out, scale, n, BLOCK=4096, num_warps=4
+    )
+    return out, scale
+
+
+def fp8_dynamic_linear(
+    x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor
+) -> torch.Tensor:
+    """``x @ (weight * weight_scale)^T`` in W8A8; ``weight`` [N, K] e4m3 row-major, whose ``.t()``
+    is the column-major operand cuBLASLt wants (a stride change, never a copy)."""
+    *lead, k = x.shape
+    x2 = x.reshape(-1, k).contiguous()
+    x8, scale = quant_per_tensor(x2)
+    y = torch._scaled_mm(
+        x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype
+    )
+    return y.reshape(*lead, weight.shape[0])
+
+
+class Fp8DynamicLinear(BaseOP):
+    """Per-tensor FP8 linear over the local shard; ``all_reduce`` adds the TP sum (row-parallel)."""
+
+    def __init__(self, local_isize: int, local_osize: int, *, all_reduce: bool = False):
+        assert local_isize % 16 == 0 and local_osize % 16 == 0, (
+            local_isize,
+            local_osize,
+        )
+        self.local_input_size = local_isize
+        self.local_output_size = local_osize
+        self.weight = torch.empty(local_osize, local_isize, dtype=FP8)
+        self.weight_scale = torch.empty((), dtype=torch.float32)
+        self._comm = (
+            DistributedCommunicator() if all_reduce and get_tp_info().size > 1 else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = fp8_dynamic_linear(x, self.weight, self.weight_scale)
+        if self._comm is not None:
+            y = self._comm.all_reduce(y)
+        return y
+
+
+class Fp8DynamicColMerged(Fp8DynamicLinear):
+    """Drop-in for ``LinearColParallelMerged``: one weight concatenating several projections
+    along the output dim; the caller splits the output by the local sizes as before."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_sizes: List[int],
+        local_output_sizes: List[int] | None = None,
+    ):
+        tp = get_tp_info()
+        if local_output_sizes is None:
+            local_output_sizes = [div_even(size, tp.size) for size in output_sizes]
+        self.output_sizes = list(output_sizes)
+        self.local_output_sizes = list(local_output_sizes)
+        super().__init__(input_size, sum(local_output_sizes))
+
+
+class Fp8DynamicRowParallel(Fp8DynamicLinear):
+    """Drop-in for ``LinearOProj`` / ``LinearRowParallel``: the input dim is sharded, the
+    all-reduce runs after the local GEMM (each rank scales its own shard)."""
+
+    def __init__(self, input_size: int, output_size: int):
+        super().__init__(
+            div_even(input_size, get_tp_info().size), output_size, all_reduce=True
+        )
+
+
+__all__ = [
+    "FP8",
+    "E4M3_MAX",
+    "Fp8DynamicColMerged",
+    "Fp8DynamicLinear",
+    "Fp8DynamicRowParallel",
+    "fp8_dynamic_linear",
+    "quant_per_tensor",
+]
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index 07a16df5c7..0ad73b752b 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -18,6 +18,13 @@ def vision_load_enabled() -> bool:
     return os.getenv("FREETOKEN_LOAD_VISION", "0").strip().lower() in _VISION_TRUE
 
 
+def fp8_dense_enabled() -> bool:
+    """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF):
+    per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs
+    (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``."""
+    return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE
+
+
 def detect_expert_quant(hf_config: Any) -> str:
     """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for
     a ModelOpt FP4 build (``quant_algo: NVFP4``) OR an llm-compressor NVFP4 export
diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py
index 651da87d93..5b7e903971 100644
--- a/python/freetoken/models/qwen4_exp/attention.py
+++ b/python/freetoken/models/qwen4_exp/attention.py
@@ -27,6 +27,7 @@
     LinearOProj,
     LinearReplicated,
 )
+from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel
 from freetoken.layers.rotary import get_rope
 from freetoken.utils import div_even, nvtx_annotate
 
@@ -138,20 +139,30 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
         self._local_qo_dim = self._local_num_q * self.head_dim
         self._local_kv_dim = self._local_num_kv * self.head_dim
         self._qkv_split = [self._local_qo_dim * 2, self._local_kv_dim, self._local_kv_dim]
-        self.qkv_proj = LinearColParallelMerged(
-            config.hidden_size,
-            [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim],
-            has_bias=False,
-            local_output_sizes=self._qkv_split,
-            quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
-        )
-        # row-parallel, NOT LinearReplicated: under TP>1 qkv_proj is column-parallel, so this
-        # rank's attention output is its own head slice and o_proj must consume the sharded
-        # input dim and all-reduce the partial sums. At TP=1 the two are equivalent.
-        self.o_proj = LinearOProj(
-            self.qo_attn_dim, config.hidden_size, has_bias=False,
-            quant_config=config.quant, prefix=f"{prefix}.o_proj",
+        qkv_sizes = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
+        # Load-time per-tensor FP8 applies only where the checkpoint declares no scheme of
+        # its own; a quantized checkpoint always wins, and its modules go through QuantConfig.
+        fp8_dyn = config.attn_quant == "fp8_dynamic" and (
+            config.quant is None or config.quant.scheme_for(f"{prefix}.qkv_proj") is None
         )
+        if fp8_dyn:  # W8A8 through _scaled_mm, weights quantized by the loader
+            self.qkv_proj = Fp8DynamicColMerged(
+                config.hidden_size, qkv_sizes, local_output_sizes=self._qkv_split
+            )
+            self.o_proj = Fp8DynamicRowParallel(self.qo_attn_dim, config.hidden_size)
+        else:
+            self.qkv_proj = LinearColParallelMerged(
+                config.hidden_size, qkv_sizes, has_bias=False,
+                local_output_sizes=self._qkv_split,
+                quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
+            )
+            # row-parallel, NOT LinearReplicated: under TP>1 qkv_proj is column-parallel, so
+            # this rank's attention output is its own head slice and o_proj must consume the
+            # sharded input dim and all-reduce. At TP=1 the two are equivalent.
+            self.o_proj = LinearOProj(
+                self.qo_attn_dim, config.hidden_size, has_bias=False,
+                quant_config=config.quant, prefix=f"{prefix}.o_proj",
+            )
         self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps)
         rotary = config.rotary_config
diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index 4bcd28de30..e8b62766af 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -7,6 +7,7 @@
 
 from freetoken.layers.quantization import QuantConfig
 from freetoken.models.config import (
+    fp8_dense_enabled,
     FullAttentionGroupConfig,
     LinearGatedDeltaGroupConfig,
     ModelConfig,
@@ -138,6 +139,12 @@ def parse_config(hf_config: Any) -> ModelConfig:
         else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))}
     )
 
+    # FREETOKEN_FP8_DENSE=1: the bf16 attention / GDN projections are quantized at LOAD to
+    # per-tensor e4m3 and served W8A8 through cuBLASLt (layers/fp8_dynamic.py). This is a
+    # synthetic scheme for a checkpoint that ships those modules unquantized -- distinct from
+    # the QuantConfig path, which serves what the checkpoint declares. The module builders
+    # apply it only where the checkpoint has no scheme of its own, so the two never collide.
+    attn_quant = "fp8_dynamic" if fp8_dense_enabled() else "none"
     layer_types = _layer_types(text)
     full_ids = tuple(i for i, t in enumerate(layer_types) if t == "full_attention")
     linear_ids = tuple(i for i, t in enumerate(layer_types) if t == "linear_attention")
@@ -245,6 +252,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         image_token_id=getattr(hf_config, "image_token_id", None),
         attention_groups=groups,
         expert_quant=expert_quant,
+        attn_quant=attn_quant,
         qwen4_args=qwen4_args,
         slot_states=ple_slot_states(qwen4_args),
     )
diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py
index 3b2e029340..f0d95c3d74 100644
--- a/python/freetoken/models/qwen4_exp/gdn.py
+++ b/python/freetoken/models/qwen4_exp/gdn.py
@@ -9,6 +9,7 @@
 from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.quantization import QuantConfig
 from freetoken.utils import div_even
+from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel
 from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
 
 
@@ -36,6 +37,7 @@ def __init__(
         self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim,
         conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "sigmoid",
         *, quant_config: QuantConfig | None = None, prefix: str = "",
+        attn_quant: str = "none",
     ):
         self.layer_id = layer_id
         # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as
@@ -66,17 +68,29 @@ def __init__(
         self._split_in_proj = (
             quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None
         )
+        # Load-time per-tensor FP8 splits the fusion the same way, but it is synthetic: it
+        # applies only where the checkpoint declares no scheme of its own.
+        self._dynamic_fp8 = attn_quant == "fp8_dynamic" and not self._split_in_proj
 
         self._in_proj_split = [
             self._local_conv_dim, self._local_value_dim,
             self._local_num_v_heads, self._local_num_v_heads,
         ]
-        # The split (quantized) branch has no TP-aware variant: its two linears are built from
-        # the checkpoint's own scheme and are not sharded here.
+        # A checkpoint-declared scheme has no TP-aware variant here: those two linears are
+        # built from the checkpoint's own scheme and are not sharded. Load-time fp8 is.
         assert not (tp.size > 1 and self._split_in_proj), (
-            "qwen4_exp TP shards bf16 GDN projections only"
+            "qwen4_exp TP shards bf16 or load-time fp8 GDN projections only"
         )
-        if self._split_in_proj:
+        if self._dynamic_fp8:
+            self.in_proj_qkvz = Fp8DynamicColMerged(
+                hidden_size, [self.conv_dim, self.value_dim],
+                local_output_sizes=[self._local_conv_dim, self._local_value_dim],
+            )
+            self.in_proj_ba = LinearColParallelMerged(
+                hidden_size, [num_v_heads, num_v_heads], has_bias=False,
+                local_output_sizes=[self._local_num_v_heads, self._local_num_v_heads],
+            )
+        elif self._split_in_proj:
             self.in_proj_qkvz = LinearColParallelMerged(
                 hidden_size, [self.conv_dim, self.value_dim], has_bias=False,
                 quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz",
@@ -100,14 +114,17 @@ def __init__(
         self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32)
         self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32)
         self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps, activation=output_gate)
-        # Row-parallel over the local v heads, all-reduce inside. At TP=1 this is exactly
-        # LinearReplicated (div_even(x, 1) == x and the reduction is skipped), so one class
-        # covers both -- and under TP>1 a replicated o_proj would silently take the wrong
-        # input width and skip the reduction.
-        self.out_proj = LinearOProj(
-            self.value_dim, hidden_size, has_bias=False,
-            quant_config=quant_config, prefix=f"{prefix}.out_proj",
-        )
+        # Row-parallel over the local v heads, all-reduce inside. At TP=1 LinearOProj is
+        # exactly LinearReplicated (div_even(x, 1) == x and the reduction is skipped), so one
+        # class covers both -- and under TP>1 a replicated o_proj would silently take the
+        # wrong input width and skip the reduction.
+        if self._dynamic_fp8:
+            self.out_proj = Fp8DynamicRowParallel(self.value_dim, hidden_size)
+        else:
+            self.out_proj = LinearOProj(
+                self.value_dim, hidden_size, has_bias=False,
+                quant_config=quant_config, prefix=f"{prefix}.out_proj",
+            )
 
     def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
         beta = b.sigmoid()
@@ -165,15 +182,14 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             fla = build_fla_metadata(batch, hidden_states.device)
             batch.fla_metadata = fla
 
-        if self._split_in_proj:
+        if self._split_in_proj or self._dynamic_fp8:
             qkvz = self.in_proj_qkvz.forward(hidden_states)
-            conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1)
+            conv_in, z = torch.split(qkvz, self._in_proj_split[:2], dim=-1)
             ba = self.in_proj_ba.forward(hidden_states)
-            b, a = torch.split(ba, [self.num_v_heads, self.num_v_heads], dim=-1)
+            b, a = torch.split(ba, [nv, nv], dim=-1)
         else:
             proj = self.in_proj.forward(hidden_states)
             conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1)
-        nk, nv = self._local_num_k_heads, self._local_num_v_heads
         kd, vd = self._local_key_dim, self._local_value_dim
         z = z.reshape(total, nv, self.head_v_dim)
         li = pool.local_index(self.layer_id)
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index a4b63362de..c9cf9c131e 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -50,6 +50,7 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseO
         output_gate=g.output_gate,
         quant_config=config.quant,
         prefix=prefix,
+        attn_quant=config.attn_quant,
     )
 
 
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 905d32d89a..71aad61540 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -189,6 +189,48 @@ def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.T
     return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None)
 
 
+# Load-time per-tensor FP8 (attn_quant == "fp8_dynamic", layers/fp8_dynamic.py): these keep
+# their name and gain a sibling ``weight_scale``; GDN ``in_proj`` splits into the fp8
+# ``in_proj_qkvz`` and the bf16 ``in_proj_ba`` (the gate projections stay bf16, as in the
+# block-fp8 checkpoints and in sglang / vLLM).
+_FP8_DENSE_SUFFIXES = (
+    ".self_attn.qkv_proj.weight",
+    ".self_attn.o_proj.weight",
+    ".linear_attn.out_proj.weight",
+)
+_E4M3_MAX = 448.0
+
+
+def _quantize_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(e4m3 weight, fp32 scale ())`` with ``w ~= weight * scale``."""
+    w = w.float()
+    scale = (w.abs().amax() / _E4M3_MAX).clamp_min(1e-12)
+    return (w / scale).clamp_(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn), scale.reshape(())
+
+
+def _fp8_dense(
+    name: str, t: torch.Tensor, config, world: int
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The (already sharded) dense tensor as the fp8_dynamic model expects it."""
+    if name.endswith(".linear_attn.in_proj.weight"):
+        g = config.linear_attention_group()
+        nk = div_even(g.num_key_heads, world, allow_replicate=True)
+        nv = div_even(g.num_value_heads, world, allow_replicate=True)
+        qkvz = 2 * nk * g.key_head_dim + 2 * nv * g.value_head_dim  # [q | k | v | z] local rows
+        assert t.shape[0] == qkvz + 2 * nv, (name, t.shape, qkvz, nv)
+        base = name[: -len("in_proj.weight")]
+        w8, scale = _quantize_per_tensor(t[:qkvz])
+        yield base + "in_proj_qkvz.weight", w8
+        yield base + "in_proj_qkvz.weight_scale", scale
+        yield base + "in_proj_ba.weight", t[qkvz:].contiguous()
+    elif name.endswith(_FP8_DENSE_SUFFIXES):
+        w8, scale = _quantize_per_tensor(t)
+        yield name, w8
+        yield name[: -len("weight")] + "weight_scale", scale
+    else:
+        yield name, t
+
+
 def iter_weights(
     model_path: str,
     device: torch.device,
@@ -213,12 +255,27 @@ def iter_weights(
     if not include_non_moe:
         return
 
+    from freetoken.models.config import fp8_dense_enabled
+
     from .config import parse_config
 
     tp = get_tp_info()
-    # The sharding geometry needs the HF config; TP=1 never shards, so keep the plain path free
-    # of a config load (synthetic test checkpoints carry no model_type).
-    config = parse_config(cached_load_hf_config(model_path)) if tp.size > 1 else None
+    # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so
+    # keep that path free of a config load (synthetic test checkpoints carry no model_type).
+    config = (
+        parse_config(cached_load_hf_config(model_path))
+        if tp.size > 1 or fp8_dense_enabled()
+        else None
+    )
+    fp8 = config is not None and config.attn_quant == "fp8_dynamic"
+
+    def emit(name: str, tensor: torch.Tensor):
+        tensor = _shard(name, tensor, config, tp.rank, tp.size)
+        if fp8:
+            yield from _fp8_dense(name, tensor, config, tp.size)
+        else:
+            yield name, tensor
+
     fuse_buf: dict[str, dict[int, torch.Tensor]] = {}
     for file in tqdm(
         iter_weight_files(model_path),
@@ -235,9 +292,9 @@ def iter_weights(
                 if fused is not None:
                     if fused != ():  # () means buffered, not yet complete
                         name, tensor = fused
-                        yield name, _shard(name, tensor, config, tp.rank, tp.size)
+                        yield from emit(name, tensor)
                     continue
-                yield name, _shard(name, tensor, config, tp.rank, tp.size)
+                yield from emit(name, tensor)
 
     assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}"
 
diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
new file mode 100644
index 0000000000..058562e5cf
--- /dev/null
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -0,0 +1,121 @@
+"""Load-time per-tensor FP8 for the qwen4_exp dense projections (FREETOKEN_FP8_DENSE=1)."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from freetoken.models.qwen4_exp.weight import _fp8_dense, _quantize_per_tensor
+
+from .common import requires_cuda
+
+E4M3_MAX = 448.0
+
+
+def _cfg():
+    g = SimpleNamespace(
+        num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8
+    )
+    return SimpleNamespace(linear_attention_group=lambda: g)
+
+
+def _assert_e4m3_close(
+    deq: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor
+) -> None:
+    # e4m3 keeps 3 mantissa bits (rel 2^-4); below scale * 2^-6 it is subnormal (abs 2^-9 steps)
+    tol = ref.abs() * 2**-4 + float(scale) * 2**-9 + 1e-7
+    assert ((deq - ref).abs() <= tol).all()
+
+
+def test_quantize_per_tensor_round_trips_within_e4m3():
+    w = torch.randn(64, 32) * 0.02
+    w8, scale = _quantize_per_tensor(w)
+    assert (
+        w8.dtype == torch.float8_e4m3fn
+        and scale.dtype == torch.float32
+        and scale.shape == ()
+    )
+    assert torch.isclose(scale * E4M3_MAX, w.abs().max())
+    _assert_e4m3_close(w8.float() * scale, w, scale)
+
+
+def test_in_proj_splits_into_fp8_qkvz_and_bf16_ba_per_rank():
+    cfg, world = (
+        _cfg(),
+        2,
+    )  # local: 2 k heads, 6 v heads -> qkvz = 2*2*8 + 2*6*8 = 128, ba = 12
+    t = torch.randn(128 + 12, 16, dtype=torch.bfloat16)
+    out = dict(_fp8_dense("model.layers.3.linear_attn.in_proj.weight", t, cfg, world))
+    assert sorted(out) == [
+        "model.layers.3.linear_attn.in_proj_ba.weight",
+        "model.layers.3.linear_attn.in_proj_qkvz.weight",
+        "model.layers.3.linear_attn.in_proj_qkvz.weight_scale",
+    ]
+    w8 = out["model.layers.3.linear_attn.in_proj_qkvz.weight"]
+    scale = out["model.layers.3.linear_attn.in_proj_qkvz.weight_scale"]
+    assert w8.shape == (128, 16) and w8.dtype == torch.float8_e4m3fn
+    _assert_e4m3_close(w8.float() * scale, t[:128].float(), scale)
+    ba = out["model.layers.3.linear_attn.in_proj_ba.weight"]
+    assert ba.dtype == torch.bfloat16 and torch.equal(ba, t[128:])
+
+
+def test_other_projections_gain_a_scale_and_the_rest_pass_through():
+    cfg = _cfg()
+    t = torch.randn(32, 16, dtype=torch.bfloat16)
+    for name in (
+        "x.self_attn.qkv_proj.weight",
+        "x.self_attn.o_proj.weight",
+        "x.linear_attn.out_proj.weight",
+    ):
+        out = dict(_fp8_dense(name, t, cfg, 1))
+        assert sorted(out) == [name, name[: -len("weight")] + "weight_scale"]
+        assert out[name].dtype == torch.float8_e4m3fn
+    out = dict(_fp8_dense("x.mlp.shared_expert.gate_up_proj.weight", t, cfg, 1))
+    assert (
+        list(out) == ["x.mlp.shared_expert.gate_up_proj.weight"]
+        and out.popitem()[1] is t
+    )
+
+
+def test_ops_declare_fp8_weight_and_scalar_scale():
+    from freetoken.distributed import set_tp_info, try_get_tp_info
+    from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    col = Fp8DynamicColMerged(32, [64, 16, 16])
+    row = Fp8DynamicRowParallel(64, 32)
+    for op in (col, row):
+        sd = op.state_dict()
+        assert set(sd) == {"weight", "weight_scale"}
+        assert (
+            sd["weight"].dtype == torch.float8_e4m3fn and sd["weight_scale"].shape == ()
+        )
+    assert col.weight.shape == (96, 32) and row.weight.shape == (32, 64)
+
+
+@requires_cuda
+@pytest.mark.parametrize("rows", [1, 16, 300], ids=["decode-1", "decode-16", "prefill"])
+def test_fp8_linear_matches_bf16_on_the_dequantized_weight(rows: int):
+    from freetoken.layers.fp8_dynamic import fp8_dynamic_linear
+
+    torch.manual_seed(0)
+    w = (torch.randn(256, 128, device="cuda") * 0.02).to(torch.bfloat16)
+    w8, scale = _quantize_per_tensor(w)
+    x = torch.randn(rows, 128, device="cuda", dtype=torch.bfloat16)
+    ref = torch.nn.functional.linear(x, (w8.float() * scale).to(torch.bfloat16))
+    got = fp8_dynamic_linear(x, w8, scale.to("cuda"))
+    assert got.shape == ref.shape and got.dtype == torch.bfloat16
+    # the per-tensor activation cast is the only extra rounding: 2^-4 relative on the inputs
+    torch.testing.assert_close(
+        got.float(), ref.float(), rtol=0.1, atol=0.08 * ref.abs().max().item()
+    )
+
+
+@requires_cuda
+def test_quant_per_tensor_zero_input_is_finite():
+    from freetoken.layers.fp8_dynamic import quant_per_tensor
+
+    x8, scale = quant_per_tensor(
+        torch.zeros(16, 128, device="cuda", dtype=torch.bfloat16)
+    )
+    assert torch.isfinite(scale).item() and (x8.float() == 0).all()

From 0968c24011915d10682c20326f5cabec20261c77 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 18:57:14 -0400
Subject: [PATCH 022/249] chore(qwen4_exp): log the load-time FP8 dense mode
 once at weight load

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/weight.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 71aad61540..e7c67658f4 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -26,10 +26,12 @@
     Nvfp4ExpertSourceSpec,
 )
 from freetoken.moe.host_banks import HostBank, read_range_into
-from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight
+from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger
 from freetoken.utils.progress import byte_bar
 from tqdm import tqdm
 
+logger = init_logger(__name__)
+
 # Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW
 # weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's
 # stacked ``mtp.layers.N.mlp.experts.*`` tensors.
@@ -268,6 +270,11 @@ def iter_weights(
         else None
     )
     fp8 = config is not None and config.attn_quant == "fp8_dynamic"
+    if fp8:
+        logger.info(
+            "qwen4_exp dense projections: load-time per-tensor FP8 (W8A8 via _scaled_mm), "
+            "FREETOKEN_FP8_DENSE=1"
+        )
 
     def emit(name: str, tensor: torch.Tensor):
         tensor = _shard(name, tensor, config, tp.rank, tp.size)

From 8d37231d3793499bec36c53d6cc939641f1b61f9 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 19:09:48 -0400
Subject: [PATCH 023/249] fix(qwen4_exp): return the FP8 quantization slack to
 the allocator before the cache planner runs

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/weight.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index e7c67658f4..7f347cec7c 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -304,6 +304,11 @@ def emit(name: str, tensor: torch.Tensor):
                 yield from emit(name, tensor)
 
     assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}"
+    if fp8 and device.type == "cuda":
+        # The bf16 originals and fp32 temporaries of the quantization sit in the caching
+        # allocator; hand them back so the expert-cache planner (free VRAM after load) sees
+        # the halved dense footprint instead of the slack.
+        torch.cuda.empty_cache()
 
 
 # ======================================================================================

From 1a6f504c05d061a6af577c6ac23103897136e92c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 19:23:02 -0400
Subject: [PATCH 024/249] fix(qwen4_exp): give in_proj_ba its own storage under
 FP8 dense

t[qkvz:].contiguous() on a contiguous row slice returns a view, so every GDN
layer's bf16 gate rows kept the whole sharded bf16 in_proj resident next to
the fp8 copy: 36 x 42 MB = 1.5 GiB per TP=2 rank, which is why the expert
cache planner saw no saving (22,594 -> 22,458 slots) after the FP8 switch.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/weight.py | 4 +++-
 tests/models/qwen4_exp/test_fp8_dense.py    | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 7f347cec7c..57f9d5f4fd 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -224,7 +224,9 @@ def _fp8_dense(
         w8, scale = _quantize_per_tensor(t[:qkvz])
         yield base + "in_proj_qkvz.weight", w8
         yield base + "in_proj_qkvz.weight_scale", scale
-        yield base + "in_proj_ba.weight", t[qkvz:].contiguous()
+        # clone, not contiguous(): a contiguous row slice IS contiguous, so .contiguous() would
+        # hand back a view that keeps the whole bf16 in_proj (36 x 42 MB per rank) resident
+        yield base + "in_proj_ba.weight", t[qkvz:].clone()
     elif name.endswith(_FP8_DENSE_SUFFIXES):
         w8, scale = _quantize_per_tensor(t)
         yield name, w8
diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
index 058562e5cf..fac51bdb5a 100644
--- a/tests/models/qwen4_exp/test_fp8_dense.py
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -56,6 +56,8 @@ def test_in_proj_splits_into_fp8_qkvz_and_bf16_ba_per_rank():
     _assert_e4m3_close(w8.float() * scale, t[:128].float(), scale)
     ba = out["model.layers.3.linear_attn.in_proj_ba.weight"]
     assert ba.dtype == torch.bfloat16 and torch.equal(ba, t[128:])
+    # its own storage: a view would keep the whole bf16 in_proj alive next to the fp8 copy
+    assert ba.untyped_storage().data_ptr() != t.untyped_storage().data_ptr()
 
 
 def test_other_projections_gain_a_scale_and_the_rest_pass_through():

From 74ccbea7a3cc60bdb90f0e28109d6af862913bc8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sat, 5 Sep 2026 17:30:28 -0400
Subject: [PATCH 025/249] perf(fp8): quantize activations with a block-parallel
 amax

quant_per_tensor ran a single Triton program over the whole tensor, so its cost
grew linearly with the input: 9.8 us for an [8, 10240] activation and 68.6 us at
[64, 10240], against ~1.2 us of useful work. Above 16384 elements, split it into
a block-parallel partial amax and a reduce+cast, which flattens the cost to
~2.4 us. The arithmetic is the one the single-program kernel already used, so a
given tensor quantizes bit for bit as before; the old three-launch torch
reduction path above 65536 elements goes away with it.

The partial count is a runtime argument and the pass-1 grid is strided and capped
at _MAX_PARTS, so both kernels have exactly ONE compiled Triton variant. Letting
the partial count reach the kernel as a constexpr instead costs a fresh
compilation for every distinct input length -- unbounded variant growth in a
server that sees arbitrary prompt lengths, and a compile stall mid-generation.

Measured on an RTX 6000 Ada (sm_89, torch 2.11) under CUDA-graph capture,
one-program -> split: 20480 elts 2.59 -> 2.22 us, 81920 elts 9.83 -> 2.44 us,
655360 elts 68.60 -> 3.06 us. Below the threshold the single program still wins
(2560 elts: 1.25 vs 2.07 us) and is kept.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
(cherry picked from commit d254729340b4ac87d02d6956b57069904cb0e73d)
---
 python/freetoken/layers/fp8_dynamic.py   | 79 +++++++++++++++++-------
 tests/models/qwen4_exp/test_fp8_dense.py | 17 +++++
 2 files changed, 75 insertions(+), 21 deletions(-)

diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py
index 1a956fbabf..5e84a929c3 100644
--- a/python/freetoken/layers/fp8_dynamic.py
+++ b/python/freetoken/layers/fp8_dynamic.py
@@ -2,10 +2,12 @@
 
 The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``),
 produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding).
-The activation is quantized per call with a dynamic per-tensor scale: one fused Triton launch
-(amax pass, then cast) at decode sizes, a torch reduction plus a cast kernel above that. No
-host sync anywhere (the scales stay on the device), so the decode path is CUDA-graph safe;
-the branch between the two paths is on the tensor *shape*, never on its values.
+The activation is quantized per call with a dynamic per-tensor scale: one Triton launch (a
+single program walking the tensor) for tiny inputs, two above ``_SPLIT_MIN_ELEMENTS`` (a
+block-parallel amax, then a reduce+cast). Both compute the same arithmetic, so the split is
+a pure speed choice. No host sync anywhere (the scales stay on the device), so the decode
+path is CUDA-graph safe; the branch between the two paths is on the tensor *shape*, never
+on its values.
 
 Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per
 decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16.
@@ -26,15 +28,21 @@
 
 FP8 = torch.float8_e4m3fn
 E4M3_MAX = 448.0
-_FUSED_MAX_ELEMENTS = (
-    65536  # one program handles the whole tensor below this (decode sizes)
-)
-_SCALE_FLOOR = 1e-12  # an all-zero activation (graph warmup buffers) must not give 1/0
+_BLOCK = 4096
+# One program per block above this, two launches (partial amax, then reduce+cast); below it
+# a single program walks the whole tensor. The one-program kernel is serial over the tensor,
+# so it must not be given the [T, hc_count*hidden] hyper-connection activations.
+_SPLIT_MIN_ELEMENTS = 16384
+# Partial-amax programs, and so the constexpr width of the fold in pass 2. Capping it (rather
+# than letting it follow the input) keeps BOTH split kernels to ONE compiled Triton variant:
+# the partial count is a runtime argument, so a server that sees a new prompt length does not
+# compile a new kernel mid-generation.
+_MAX_PARTS = 512
 
 
 @triton.jit
 def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
-    """One program: max|x| over the tensor, then the cast. Decode-sized inputs only."""
+    """One program: max|x| over the tensor, then the cast. Tiny inputs only (see SPLIT_MIN)."""
     acc = tl.zeros([BLOCK], dtype=tl.float32)
     for start in range(0, n, BLOCK):
         offs = start + tl.arange(0, BLOCK)
@@ -52,12 +60,38 @@ def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
 
 
 @triton.jit
-def _quant_cast_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
-    """Cast under a scale already on the device (prefill sizes; the amax is a torch reduction)."""
-    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+def _amax_partial_kernel(x_ptr, part_ptr, n, nprog, BLOCK: tl.constexpr):
+    """max|x| -> one fp32 partial per program (pass 1). Grid-strided, so ``nprog`` bounds the
+    partial count however large the tensor is."""
+    pid = tl.program_id(0)
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for start in range(pid * BLOCK, n, nprog * BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32)
+        acc = tl.maximum(acc, tl.abs(v))
+    tl.store(part_ptr + pid, tl.max(acc, axis=0))
+
+
+@triton.jit
+def _reduce_cast_kernel(
+    x_ptr, out_ptr, part_ptr, scale_ptr, n, nprog,
+    MAX_PARTS: tl.constexpr, BLOCK: tl.constexpr,
+):
+    """Pass 2: fold the partials to the tensor amax, then cast this program's block.
+
+    Every program repeats the (nprog-element, L2-resident) fold instead of paying a third
+    launch for it; program 0 also publishes the scale. The arithmetic is the one
+    ``_quant_fused_kernel`` uses, so both paths quantize a given tensor identically.
+    """
+    poffs = tl.arange(0, MAX_PARTS)
+    amax = tl.max(tl.load(part_ptr + poffs, mask=poffs < nprog, other=0.0), axis=0)
+    amax = tl.maximum(amax, 1e-12)
+    pid = tl.program_id(0)
+    if pid == 0:
+        tl.store(scale_ptr, amax / 448.0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
     mask = offs < n
-    inv = 1.0 / tl.load(scale_ptr)
-    v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
+    v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * (448.0 / amax)
     v = tl.minimum(tl.maximum(v, -448.0), 448.0)
     tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
 
@@ -66,14 +100,17 @@ def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``."""
     n = x.numel()
     out = torch.empty_like(x, dtype=FP8)
-    if n <= _FUSED_MAX_ELEMENTS:
-        scale = torch.empty((), dtype=torch.float32, device=x.device)
-        _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=4096, num_warps=8)
+    scale = torch.empty((), dtype=torch.float32, device=x.device)
+    if n <= _SPLIT_MIN_ELEMENTS:
+        # below this, the second launch costs more than the parallelism buys
+        _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=_BLOCK, num_warps=8)
         return out, scale
-    amax = torch.linalg.vector_norm(x, ord=float("inf")).float()
-    scale = amax.clamp_min_(_SCALE_FLOOR).div_(E4M3_MAX)
-    _quant_cast_kernel[(triton.cdiv(n, 4096),)](
-        x, out, scale, n, BLOCK=4096, num_warps=4
+    nblock = triton.cdiv(n, _BLOCK)
+    nprog = min(nblock, _MAX_PARTS)
+    part = torch.empty(nprog, dtype=torch.float32, device=x.device)
+    _amax_partial_kernel[(nprog,)](x, part, n, nprog, BLOCK=_BLOCK, num_warps=8)
+    _reduce_cast_kernel[(nblock,)](
+        x, out, part, scale, n, nprog, MAX_PARTS=_MAX_PARTS, BLOCK=_BLOCK, num_warps=8,
     )
     return out, scale
 
diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
index fac51bdb5a..4b12598a13 100644
--- a/tests/models/qwen4_exp/test_fp8_dense.py
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -121,3 +121,20 @@ def test_quant_per_tensor_zero_input_is_finite():
         torch.zeros(16, 128, device="cuda", dtype=torch.bfloat16)
     )
     assert torch.isfinite(scale).item() and (x8.float() == 0).all()
+
+
+@requires_cuda
+@pytest.mark.parametrize("n", [4096, 16384, 20480, 81920, 200000])
+def test_split_quant_path_agrees_with_the_single_program_one(n: int):
+    """The two-launch path above _SPLIT_MIN_ELEMENTS must quantize bit for bit like the
+    one-program kernel: same amax (max is exact under any grouping), same arithmetic."""
+    import freetoken.layers.fp8_dynamic as m
+
+    torch.manual_seed(0)
+    x = (torch.randn(n, device="cuda") * 3.0).to(torch.bfloat16)
+    got8, got_scale = m.quant_per_tensor(x)
+    ref8 = torch.empty_like(x, dtype=m.FP8)
+    ref_scale = torch.empty((), dtype=torch.float32, device="cuda")
+    m._quant_fused_kernel[(1,)](x, ref8, ref_scale, n, BLOCK=m._BLOCK, num_warps=8)
+    assert torch.equal(got_scale, ref_scale)
+    assert torch.equal(got8.view(torch.uint8), ref8.view(torch.uint8))

From 85c4b6042a55211d6983a839bcbc2dc2b9d457b0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sat, 5 Sep 2026 17:10:21 -0400
Subject: [PATCH 026/249] feat(layers): optional per-tensor FP8 lm_head
 (FREETOKEN_FP8_LMHEAD=1)

The vocab-parallel head is the last large bf16 read on the decode path: 0.64 GB
per step per rank at the Qwen3.8-Flash-Next geometry ([124160, 2560] after the
TP=2 vocab split), 0.80 ms of the 9.9 ms step in an nsys trace of production.

Its own flag, not FREETOKEN_FP8_DENSE, because this one moves the logits: every
other quantized module feeds a norm or a sigmoid downstream, while a per-tensor
e4m3 vocab matrix changes each sampled token's score directly, so it carries its
own quality gate rather than riding along with the pure-throughput changes.

ParallelLMHead.forward grows a _logits() seam (the local vocab-shard GEMM);
Fp8ParallelLMHead overrides only that, leaving the all_gather of the logits
above it untouched. Untied embeddings only -- a tied head shares the bf16
embedding table, which the lookup side still reads as bf16.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
(cherry picked from commit 35e80b30e30052fbac360baa167f0f12a791ed34)
(cherry picked from commit 0693735a60efd17fded6f9b686f3b6d298c5bea0)
(cherry picked from commit 1797ab56e8ee9cd7138111770896bca54b7d9ce7)
---
 python/freetoken/layers/embedding.py        | 16 +++++++++----
 python/freetoken/layers/fp8_dynamic.py      | 20 ++++++++++++++++
 python/freetoken/models/config.py           |  8 +++++++
 python/freetoken/models/qwen4_exp/model.py  | 26 ++++++++++++++-------
 python/freetoken/models/qwen4_exp/weight.py | 24 ++++++++++++-------
 tests/models/qwen4_exp/test_fp8_dense.py    | 13 +++++++++++
 6 files changed, 87 insertions(+), 20 deletions(-)

diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index a22216548a..6939156cd6 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -121,6 +121,17 @@ def state_dict(
             return super().state_dict(prefix=prefix, result=result)
         return {} if result is None else result
 
+    def _logits(self, x: torch.Tensor) -> torch.Tensor:
+        """The local vocab-shard GEMM; the seam a quantized head overrides."""
+        module = self.tied_embedding or self
+        return F.linear(x, module.weight, self.bias)
+
+    def _logits(self, x: torch.Tensor) -> torch.Tensor:
+        """The local vocab-shard GEMM; the seam a quantized head overrides."""
+        if self.tied_embedding is not None:
+            return F.linear(x, self.tied_embedding.weight, self.bias)
+        return self.quant_method.apply(self, x)
+
     @nvtx_annotate("LMHead")
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         ctx = get_global_ctx()
@@ -131,10 +142,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = x[indices].contiguous()
             del indices
 
-        if self.tied_embedding is not None:
-            logits = F.linear(x, self.tied_embedding.weight, self.bias)
-        else:
-            logits = self.quant_method.apply(self, x)
+        logits = self._logits(x)
         if self.tp_size == 1:
             return logits
         input_shape = logits.shape
diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py
index 5e84a929c3..caeebb3221 100644
--- a/python/freetoken/layers/fp8_dynamic.py
+++ b/python/freetoken/layers/fp8_dynamic.py
@@ -25,6 +25,7 @@
 from freetoken.utils import div_even
 
 from .base import BaseOP
+from .embedding import ParallelLMHead
 
 FP8 = torch.float8_e4m3fn
 E4M3_MAX = 448.0
@@ -180,9 +181,28 @@ def __init__(self, input_size: int, output_size: int):
         )
 
 
+class Fp8ParallelLMHead(ParallelLMHead):
+    """``ParallelLMHead`` whose vocab shard is a per-tensor e4m3 weight (FREETOKEN_FP8_LMHEAD=1).
+
+    Only the GEMM changes; the vocab-parallel all_gather of the logits above it is untouched.
+    Untied embeddings only -- a tied head shares the bf16 embedding table, which the lookup
+    side still reads as bf16.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int):
+        super().__init__(num_embeddings, embedding_dim)
+        self.weight = torch.empty(self.num_embeddings_tp, embedding_dim, dtype=FP8)
+        self.weight_scale = torch.empty((), dtype=torch.float32)
+
+    def _logits(self, x: torch.Tensor) -> torch.Tensor:
+        y = fp8_dynamic_linear(x, self.weight, self.weight_scale)
+        return y if self.bias is None else y + self.bias
+
+
 __all__ = [
     "FP8",
     "E4M3_MAX",
+    "Fp8ParallelLMHead",
     "Fp8DynamicColMerged",
     "Fp8DynamicLinear",
     "Fp8DynamicRowParallel",
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index 0ad73b752b..62340ab754 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -25,6 +25,14 @@ def fp8_dense_enabled() -> bool:
     return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE
 
 
+def fp8_lmhead_enabled() -> bool:
+    """Load-time FP8 for the lm_head as well (opt-in, default OFF, needs FREETOKEN_FP8_DENSE=1
+    for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits:
+    a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own
+    quality gate. ``FREETOKEN_FP8_LMHEAD=1``."""
+    return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE
+
+
 def detect_expert_quant(hf_config: Any) -> str:
     """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for
     a ModelOpt FP4 build (``quant_algo: NVFP4``) OR an llm-compressor NVFP4 export
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index c9cf9c131e..848b4d645d 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -21,6 +21,7 @@
 from freetoken.core import get_global_ctx
 from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding
 from freetoken.models.blocks import BaseLLMModel
+from freetoken.models.config import fp8_lmhead_enabled
 from freetoken.utils import nvtx_annotate
 
 from .attention import Qwen4ExpAttention
@@ -129,14 +130,23 @@ class Qwen4ExpForCausalLM(BaseLLMModel):
     def __init__(self, config: ModelConfig) -> None:
         self._config = config
         self.model = Qwen4ExpModel(config)
-        self.lm_head = ParallelLMHead(
-            num_embeddings=config.vocab_size,
-            embedding_dim=config.hidden_size,
-            tie_word_embeddings=config.tie_word_embeddings,
-            tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None,
-            quant_config=config.quant,
-            prefix="lm_head",
-        )
+        if fp8_lmhead_enabled() and not config.tie_word_embeddings and config.quant is None:
+            # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized.
+            # A checkpoint that declares its own scheme goes through QuantConfig instead.
+            from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead
+
+            self.lm_head = Fp8ParallelLMHead(
+                num_embeddings=config.vocab_size, embedding_dim=config.hidden_size
+            )
+        else:
+            self.lm_head = ParallelLMHead(
+                num_embeddings=config.vocab_size,
+                embedding_dim=config.hidden_size,
+                tie_word_embeddings=config.tie_word_embeddings,
+                tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None,
+                quant_config=config.quant,
+                prefix="lm_head",
+            )
         super().__init__()
 
     def load_host_tables(self, engine_config) -> int:
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 57f9d5f4fd..125cea2c0e 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -200,6 +200,8 @@ def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.T
     ".self_attn.o_proj.weight",
     ".linear_attn.out_proj.weight",
 )
+# Behind its own flag: this one moves the logits (see models.config.fp8_lmhead_enabled).
+_FP8_LMHEAD_SUFFIX = "lm_head.weight"
 _E4M3_MAX = 448.0
 
 
@@ -211,10 +213,11 @@ def _quantize_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
 
 def _fp8_dense(
-    name: str, t: torch.Tensor, config, world: int
+    name: str, t: torch.Tensor, config, world: int, *,
+    dense: bool = True, lm_head: bool = False,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """The (already sharded) dense tensor as the fp8_dynamic model expects it."""
-    if name.endswith(".linear_attn.in_proj.weight"):
+    if dense and name.endswith(".linear_attn.in_proj.weight"):
         g = config.linear_attention_group()
         nk = div_even(g.num_key_heads, world, allow_replicate=True)
         nv = div_even(g.num_value_heads, world, allow_replicate=True)
@@ -227,7 +230,9 @@ def _fp8_dense(
         # clone, not contiguous(): a contiguous row slice IS contiguous, so .contiguous() would
         # hand back a view that keeps the whole bf16 in_proj (36 x 42 MB per rank) resident
         yield base + "in_proj_ba.weight", t[qkvz:].clone()
-    elif name.endswith(_FP8_DENSE_SUFFIXES):
+    elif (dense and name.endswith(_FP8_DENSE_SUFFIXES)) or (
+        lm_head and name.endswith(_FP8_LMHEAD_SUFFIX)
+    ):
         w8, scale = _quantize_per_tensor(t)
         yield name, w8
         yield name[: -len("weight")] + "weight_scale", scale
@@ -259,16 +264,17 @@ def iter_weights(
     if not include_non_moe:
         return
 
-    from freetoken.models.config import fp8_dense_enabled
+    from freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled
 
     from .config import parse_config
 
     tp = get_tp_info()
     # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so
     # keep that path free of a config load (synthetic test checkpoints carry no model_type).
+    lmhead = fp8_lmhead_enabled()
     config = (
         parse_config(cached_load_hf_config(model_path))
-        if tp.size > 1 or fp8_dense_enabled()
+        if tp.size > 1 or fp8_dense_enabled() or lmhead
         else None
     )
     fp8 = config is not None and config.attn_quant == "fp8_dynamic"
@@ -277,11 +283,13 @@ def iter_weights(
             "qwen4_exp dense projections: load-time per-tensor FP8 (W8A8 via _scaled_mm), "
             "FREETOKEN_FP8_DENSE=1"
         )
+    if lmhead:
+        logger.info("qwen4_exp lm_head: load-time per-tensor FP8, FREETOKEN_FP8_LMHEAD=1")
 
     def emit(name: str, tensor: torch.Tensor):
         tensor = _shard(name, tensor, config, tp.rank, tp.size)
-        if fp8:
-            yield from _fp8_dense(name, tensor, config, tp.size)
+        if fp8 or lmhead:
+            yield from _fp8_dense(name, tensor, config, tp.size, dense=fp8, lm_head=lmhead)
         else:
             yield name, tensor
 
@@ -306,7 +314,7 @@ def emit(name: str, tensor: torch.Tensor):
                 yield from emit(name, tensor)
 
     assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}"
-    if fp8 and device.type == "cuda":
+    if (fp8 or lmhead) and device.type == "cuda":
         # The bf16 originals and fp32 temporaries of the quantization sit in the caching
         # allocator; hand them back so the expert-cache planner (free VRAM after load) sees
         # the halved dense footprint instead of the slack.
diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
index 4b12598a13..058515b183 100644
--- a/tests/models/qwen4_exp/test_fp8_dense.py
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -138,3 +138,16 @@ def test_split_quant_path_agrees_with_the_single_program_one(n: int):
     m._quant_fused_kernel[(1,)](x, ref8, ref_scale, n, BLOCK=m._BLOCK, num_warps=8)
     assert torch.equal(got_scale, ref_scale)
     assert torch.equal(got8.view(torch.uint8), ref8.view(torch.uint8))
+
+
+def test_reader_gates_lm_head_behind_its_own_flag():
+    cfg = _cfg()
+    t = torch.randn(64, 32, dtype=torch.bfloat16)
+    assert list(dict(_fp8_dense("lm_head.weight", t, cfg, 1))) == ["lm_head.weight"]
+    out = dict(_fp8_dense("lm_head.weight", t, cfg, 1, dense=False, lm_head=True))
+    assert sorted(out) == ["lm_head.weight", "lm_head.weight_scale"]
+    assert out["lm_head.weight"].dtype == torch.float8_e4m3fn
+    # the lm_head flag alone must not pull in the dense rewrites
+    passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1,
+                                  dense=False, lm_head=True))
+    assert list(passthrough) == ["x.self_attn.qkv_proj.weight"]

From 62a9a8c1031fa779c674d18b03a86d37fdeeed5f Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 18:02:07 -0400
Subject: [PATCH 027/249] fix(qwen4_exp): forward prefix through Qwen4ExpMoE,
 and test the CPU build

Qwen4ExpDecoderLayer builds its MoE as `Qwen4ExpMoE(config, layer_id, prefix=...)`,
but this PR's override of __init__ (added to hold the TP communicator) took only
(config, layer_id), so a server boot died with

    TypeError: Qwen4ExpMoE.__init__() got an unexpected keyword argument 'prefix'

The whole CPU test suite was green with that bug in place, because every test that
builds a decoder layer is behind requires_cuda -- nothing without a GPU ever
constructed the model. tests/models/qwen4_exp/test_build_cpu.py closes that: it
builds the full model on the meta device (no GPU, no memory) and asserts the state
dict has both layer families, an lm_head, and MoE weights on more than one layer,
so a dropped or shared prefix fails too.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/moe.py |  6 ++--
 tests/models/qwen4_exp/test_build_cpu.py | 43 ++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 tests/models/qwen4_exp/test_build_cpu.py

diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py
index 7d9f6adf7a..c10407861f 100644
--- a/python/freetoken/models/qwen4_exp/moe.py
+++ b/python/freetoken/models/qwen4_exp/moe.py
@@ -25,8 +25,10 @@ class Qwen4ExpMoE(Qwen3_5MoE):
     is reduced once (one all-reduce per MoE layer instead of two).
     """
 
-    def __init__(self, config: ModelConfig, layer_id: int | None = None) -> None:
-        super().__init__(config, layer_id=layer_id)
+    def __init__(
+        self, config: ModelConfig, layer_id: int | None = None, *, prefix: str = ""
+    ) -> None:
+        super().__init__(config, layer_id=layer_id, prefix=prefix)
         self._comm = DistributedCommunicator()
         self._tp_size = get_tp_info().size
 
diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py
new file mode 100644
index 0000000000..0cf2027d14
--- /dev/null
+++ b/tests/models/qwen4_exp/test_build_cpu.py
@@ -0,0 +1,43 @@
+"""The whole model must CONSTRUCT without a GPU.
+
+Every other model test that builds a decoder layer is behind requires_cuda, so a constructor
+signature mismatch between Qwen4ExpDecoderLayer and the ops it builds is invisible to a
+CPU-only run -- it only shows up when a server boots. That happened: Qwen4ExpMoE did not
+accept the `prefix` kwarg model.py passes, and a full CPU suite still went green.
+
+Building on the meta device costs no memory and no GPU, so there is no reason not to.
+"""
+
+import torch
+from freetoken.distributed import set_tp_info, try_get_tp_info
+from freetoken.layers import set_rope_device
+from freetoken.models.qwen4_exp.config import parse_config
+from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM
+
+from .common import toy_hf_config
+
+
+def _build():
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    set_rope_device(torch.device("cpu"))
+    with torch.device("meta"):
+        return Qwen4ExpForCausalLM(parse_config(toy_hf_config()))
+
+
+def test_model_constructs_on_meta_without_cuda():
+    model = _build()
+    sd = model.state_dict()
+    assert sd, "state dict is empty"
+    # both layer families and the head must be present
+    assert any(".self_attn." in k for k in sd), "no full-attention layer built"
+    assert any(".linear_attn." in k for k in sd), "no GDN layer built"
+    assert any(k.startswith("lm_head") for k in sd), "no lm_head built"
+
+
+def test_every_layer_gets_its_own_prefixed_weights():
+    # a prefix that is dropped or shared silently collapses layers onto one another
+    sd = _build().state_dict()
+    mlp_keys = {k for k in sd if ".mlp." in k}
+    layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k}
+    assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}"

From d2ad39b8d55550938741f7788183557fd2524591 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 18:40:25 -0400
Subject: [PATCH 028/249] feat(quant): let the triton NVFP4 expert kernel serve
 a tensor-parallel shard

With the expert piece stream sliced per rank and the banks sized from
MoEConfig.local_intermediate, a rank holds exactly its half of every expert, so this
kernel can serve TP>1 -- the routed output is a partial sum and the MoE layer already
reduces it (_maybe_all_reduce, or the single combined all-reduce in qwen4_exp's block).

Without this the whole selection table is empty under TP=2 on sm_89 and the server
refuses to start:

    KernelSelectionError: no usable kernel in table;
      triton: TP > 1 is not supported for this expert format;
      marlin: vLLM is not installed;
      b12x: b12x requires sm_120+, got sm_89

marlin and b12x keep tp_ok=False deliberately: their pack() repacks the native rows and
neither has been verified against a per-rank bank.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/layers/quantization/moe/nvfp4.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index dfe35604b3..e2dc1478c2 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -41,7 +41,13 @@ class TritonNvfp4MoEKernel(MoEKernel):
     cpu_format = "nvfp4"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+        # tp_ok: the source stream is sliced along the intermediate axis per rank
+        # (nvfp4_banks._tp_shard) and this kernel sizes its banks from
+        # cfg.local_intermediate, so a rank holds and reads only its own half. The routed
+        # output is then a partial sum, which the MoE layer reduces (_maybe_all_reduce, or
+        # one combined all-reduce in qwen4_exp's block). marlin and b12x stay off: their
+        # pack() repacks the rows and neither has been checked against a sharded bank.
+        reason = self._common_reject(cfg, resident_ok=False, tp_ok=True, cpu_ok=True, plain_silu_only=False)
         if reason:
             return reason
         reason = gated_epilogue_reason(cfg)

From 1da3c6c7353b9ee304a8c797bee42b8cd54629fa Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 19:00:04 -0400
Subject: [PATCH 029/249] fix(qwen4_exp): one predicate decides the FP8
 lm_head, for builder and loader alike

The model builder and the weight reader tested this separately and drifted, so on the
shipping checkpoint the reader emitted lm_head.weight_scale while the builder made a
plain ParallelLMHead, and startup died with

    RuntimeError: Unexpected keys in state_dict: ['lm_head.weight_scale']

The builder's extra condition was `config.quant is None`, which is wrong twice over:
what matters is whether the checkpoint declares a scheme for **lm_head**, not whether it
has a QuantConfig at all. This model ships NVFP4 routed experts (so quant is not None)
with lm_head in the modelopt ignore list (so it has no scheme and the synthetic FP8 head
is exactly what is wanted).

Both sides now call config.use_fp8_lmhead(), which owns the flag, the tied-embedding
exclusion and the scheme test in one place.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/config.py | 17 +++++++++++++++++
 python/freetoken/models/qwen4_exp/model.py  |  4 ++--
 python/freetoken/models/qwen4_exp/weight.py |  7 ++++---
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index e8b62766af..6f1939eb32 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -110,6 +110,23 @@ def _layer_types(text: Any) -> list[str]:
     ]
 
 
+def use_fp8_lmhead(config: ModelConfig) -> bool:
+    """Whether to build/load the lm_head as load-time per-tensor FP8.
+
+    The MODEL BUILDER and the WEIGHT READER must agree exactly: if one says yes and the
+    other no, the state dict gains or loses ``lm_head.weight_scale`` and load fails. They
+    used to test this separately and drifted. Conditions: the operator asked
+    (FREETOKEN_FP8_LMHEAD=1), the head is untied (a tied head shares the bf16 embedding
+    table), and the checkpoint declares no scheme OF ITS OWN for lm_head -- a quantized head
+    goes through QuantConfig instead. Note it is the lm_head scheme that matters, not whether
+    the checkpoint has a QuantConfig at all: this model ships NVFP4 experts with lm_head in
+    the modelopt ignore list."""
+    if not fp8_lmhead_enabled() or config.tie_word_embeddings:
+        return False
+    quant = getattr(config, "quant", None)
+    return quant is None or quant.scheme_for("lm_head") is None
+
+
 def parse_config(hf_config: Any) -> ModelConfig:
     text = getattr(hf_config, "text_config", hf_config)
 
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index 848b4d645d..3238459a6a 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -21,7 +21,7 @@
 from freetoken.core import get_global_ctx
 from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding
 from freetoken.models.blocks import BaseLLMModel
-from freetoken.models.config import fp8_lmhead_enabled
+from .config import use_fp8_lmhead
 from freetoken.utils import nvtx_annotate
 
 from .attention import Qwen4ExpAttention
@@ -130,7 +130,7 @@ class Qwen4ExpForCausalLM(BaseLLMModel):
     def __init__(self, config: ModelConfig) -> None:
         self._config = config
         self.model = Qwen4ExpModel(config)
-        if fp8_lmhead_enabled() and not config.tie_word_embeddings and config.quant is None:
+        if use_fp8_lmhead(config):
             # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized.
             # A checkpoint that declares its own scheme goes through QuantConfig instead.
             from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 125cea2c0e..10f9608522 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -266,17 +266,18 @@ def iter_weights(
 
     from freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled
 
-    from .config import parse_config
+    from .config import parse_config, use_fp8_lmhead
 
     tp = get_tp_info()
     # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so
     # keep that path free of a config load (synthetic test checkpoints carry no model_type).
-    lmhead = fp8_lmhead_enabled()
     config = (
         parse_config(cached_load_hf_config(model_path))
-        if tp.size > 1 or fp8_dense_enabled() or lmhead
+        if tp.size > 1 or fp8_dense_enabled() or fp8_lmhead_enabled()
         else None
     )
+    # the SAME predicate the model builder uses, so the emitted keys always match its slots
+    lmhead = config is not None and use_fp8_lmhead(config)
     fp8 = config is not None and config.attn_quant == "fp8_dynamic"
     if fp8:
         logger.info(

From d377a20c09fe90b148c03f66d72aad22e4d81e15 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 19:12:27 -0400
Subject: [PATCH 030/249] fix(qwen4_exp): import fp8_lmhead_enabled, which
 use_fp8_lmhead calls

The predicate moved into config.py but its import did not come with it, so building the
model raised NameError on the very first line of use_fp8_lmhead -- under every flag
setting, including the default one that wants no FP8 head at all.

Found by the CPU key check (build the model on meta, diff its slots against the names the
loader emits) before it could cost a GPU window. No behaviour change beyond the module now
importing.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/config.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index 6f1939eb32..bb99669876 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -8,6 +8,7 @@
 from freetoken.layers.quantization import QuantConfig
 from freetoken.models.config import (
     fp8_dense_enabled,
+    fp8_lmhead_enabled,
     FullAttentionGroupConfig,
     LinearGatedDeltaGroupConfig,
     ModelConfig,

From bd106c9aa5963f3d19c543dc557d7649f2fb61a6 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 10 Sep 2026 19:53:05 -0400
Subject: [PATCH 031/249] fix(qwen4_exp): restore the local head counts the GDN
 forward reads

The rebase onto the quantization refactor dropped

    nk, nv = self._local_num_k_heads, self._local_num_v_heads

from the top of GatedDeltaNet.forward while keeping all eleven uses of nk and nv below
it. Every decode died at the first CUDA-graph capture:

    File "python/freetoken/models/qwen4_exp/gdn.py", line 189, in forward
      b, a = torch.split(ba, [nv, nv], dim=-1)
    NameError: name 'nv' is not defined

Also drops a duplicate `from freetoken.distributed import get_tp_info` two lines under
the first, left by the same merge.

`ruff check --select F821` reports both, and would have reported the missing
fp8_lmhead_enabled import a commit earlier. Lint the tree before asking for the GPU, not
after: this file's eleven undefined names cost a window that a sub-second check would
have saved. The eight F821 hits outside qwen4_exp are pre-existing on main.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
---
 python/freetoken/models/qwen4_exp/gdn.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py
index f0d95c3d74..763cb78ca4 100644
--- a/python/freetoken/models/qwen4_exp/gdn.py
+++ b/python/freetoken/models/qwen4_exp/gdn.py
@@ -5,7 +5,6 @@
 from freetoken.core import get_global_ctx
 from freetoken.distributed import get_tp_info
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
-from freetoken.distributed import get_tp_info
 from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.quantization import QuantConfig
 from freetoken.utils import div_even
@@ -182,6 +181,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             fla = build_fla_metadata(batch, hidden_states.device)
             batch.fla_metadata = fla
 
+        nk, nv = self._local_num_k_heads, self._local_num_v_heads
         if self._split_in_proj or self._dynamic_fp8:
             qkvz = self.in_proj_qkvz.forward(hidden_states)
             conv_in, z = torch.split(qkvz, self._in_proj_split[:2], dim=-1)

From b482b86221c7a3de88f7ed0347e11366d3d46882 Mon Sep 17 00:00:00 2001
From: Shikha Kathrani <99456505+shikhakath@users.noreply.github.com>
Date: Sun, 13 Sep 2026 15:56:34 -0700
Subject: [PATCH 032/249] fix(server): loopback-bind the TP rendezvous store,
 add --dist-port

---
 docs/cli.md                            |  1 +
 python/freetoken/engine/config.py      |  3 +-
 python/freetoken/engine/engine.py      | 49 +++++++++++++------
 python/freetoken/server/args.py        | 28 +++++++++--
 tests/engine/test_distributed_store.py | 66 ++++++++++++++++++++++++++
 tests/server/test_dist_port_arg.py     | 41 ++++++++++++++++
 6 files changed, 168 insertions(+), 20 deletions(-)
 create mode 100644 tests/engine/test_distributed_store.py
 create mode 100644 tests/server/test_dist_port_arg.py

diff --git a/docs/cli.md b/docs/cli.md
index e1eaa8aa3e..144597bb06 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) |
 | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 7ab792f9e9..0f16eeeca1 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -80,6 +80,7 @@ class EngineConfig:
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
     distributed_timeout: float = 60.0
+    distributed_port: int = 2333
     use_dummy_weight: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None
@@ -121,4 +122,4 @@ def max_forward_len(self) -> int:
 
     @property
     def distributed_addr(self) -> str:
-        return "tcp://127.0.0.1:2333"
+        return f"tcp://127.0.0.1:{self.distributed_port}"
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 006c19089e..fd0d1c084a 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -3,6 +3,7 @@
 import gc
 import math
 import os
+import socket
 from datetime import timedelta
 from typing import Any, Dict, Iterable, NamedTuple, Tuple
 
@@ -434,15 +435,40 @@ def __init__(self, config: EngineConfig):
             # Prefill runs on the first comma part; warm its autotune cache.
             self._warmup_prefill()
 
-    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
-        if config.tp_info.size == 1 or config.use_pynccl:
-            torch.distributed.init_process_group(
-                backend="gloo",
-                rank=config.tp_info.rank,
-                world_size=config.tp_info.size,
-                timeout=timedelta(seconds=config.distributed_timeout),
-                init_method=config.distributed_addr,
+    def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Store:
+        """The rendezvous store's C10d server ignores the host it's given and always listens
+        on every interface, so an unauthenticated TCPStore ends up reachable off-box even
+        when distributed_addr says 127.0.0.1. Rank 0 pre-binds the listening socket to
+        loopback itself and hands the fd to TCPStore (master_listen_fd) to force that."""
+        timeout = timedelta(seconds=config.distributed_timeout)
+        if not config.tp_info.is_primary():
+            return torch.distributed.TCPStore(
+                "127.0.0.1", config.distributed_port, config.tp_info.size,
+                is_master=False, timeout=timeout, multi_tenant=True,
             )
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind(("127.0.0.1", config.distributed_port))
+        sock.listen(1024)
+        # TCPStore uses this fd directly rather than duplicating it, so it must outlive the
+        # store; kept alive on self, closed only when the process (and the store) exits.
+        self._distributed_listen_socket = sock
+        return torch.distributed.TCPStore(
+            "127.0.0.1", config.distributed_port, config.tp_info.size,
+            is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.fileno(),
+        )
+
+    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
+        store = self._make_distributed_store(config)
+        use_gloo = config.tp_info.size == 1 or config.use_pynccl
+        torch.distributed.init_process_group(
+            backend="gloo" if use_gloo else "nccl",
+            rank=config.tp_info.rank,
+            world_size=config.tp_info.size,
+            timeout=timedelta(seconds=config.distributed_timeout),
+            store=store,
+        )
+        if use_gloo:
             tp_cpu_group = torch.distributed.group.WORLD
             assert tp_cpu_group is not None
             max_bytes = (
@@ -450,13 +476,6 @@ def _init_communication(self, config: EngineConfig) -> torch.distributed.Process
             )
             enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes)
         else:
-            torch.distributed.init_process_group(
-                backend="nccl",
-                rank=config.tp_info.rank,
-                world_size=config.tp_info.size,
-                timeout=timedelta(seconds=config.distributed_timeout),
-                init_method=config.distributed_addr,
-            )
             tp_cpu_group = torch.distributed.new_group(backend="gloo")
             assert tp_cpu_group is not None
         return tp_cpu_group
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 32e8d1266d..c773367dce 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -93,10 +93,6 @@ def backend_create_detokenizer_link(self) -> bool:
     def frontend_create_tokenizer_link(self) -> bool:
         return not self.share_tokenizer
 
-    @property
-    def distributed_addr(self) -> str:
-        return f"tcp://127.0.0.1:{self.server_port + 1}"
-
 
 def parse_args(
     args: List[str],
@@ -143,6 +139,15 @@ def _positive_int(value: str) -> int:
             raise argparse.ArgumentTypeError("must be >= 1")
         return n
 
+    def _valid_port(value: str) -> int:
+        try:
+            n = int(value)
+        except ValueError as exc:
+            raise argparse.ArgumentTypeError("must be an integer") from exc
+        if not 0 <= n <= 65535:
+            raise argparse.ArgumentTypeError("must be between 0 and 65535")
+        return n
+
     def _lazy_gpu_arg(value: str) -> tuple[str, ...]:
         from freetoken.gpu_select import gpu_arg
 
@@ -338,6 +343,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="The port number for the server to listen on.",
     )
 
+    parser.add_argument(
+        "--dist-port",
+        type=_valid_port,
+        dest="distributed_port",
+        default=None,
+        help=(
+            "Port for the internal TP rendezvous store, loopback-only regardless of --host. "
+            "Defaults to --port + 1; override when that collides with another instance or "
+            "service on the same host."
+        ),
+    )
+
     parser.add_argument(
         "--cuda-graph-max-bs",
         "--graph",
@@ -721,6 +738,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         if entry:
             kwargs["quant_backend"] = entry
 
+    if kwargs["distributed_port"] is None:
+        kwargs["distributed_port"] = kwargs["server_port"] + 1
+
     if kwargs["model_path"].startswith("~"):
         kwargs["model_path"] = os.path.expanduser(kwargs["model_path"])
 
diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py
new file mode 100644
index 0000000000..f70be12668
--- /dev/null
+++ b/tests/engine/test_distributed_store.py
@@ -0,0 +1,66 @@
+"""Regression test for #301: the TP rendezvous store must not fall back to a wildcard
+bind. torch's tcp:// rendezvous handler ignores the host it's given when binding the
+master's listening socket and always listens on every interface, so
+init_method="tcp://127.0.0.1:PORT" looks loopback-only but isn't -- see
+Engine._make_distributed_store. This drives the real store-construction path across two
+spawned local ranks (CPU/gloo) and checks the actual bind address, not just the URL.
+"""
+
+from __future__ import annotations
+
+import multiprocessing as mp
+from datetime import timedelta
+from types import SimpleNamespace
+
+import torch
+import torch.distributed as dist
+
+from freetoken.distributed import DistributedInfo
+from freetoken.engine.engine import Engine
+
+WORLD_SIZE = 2
+TIMEOUT = timedelta(seconds=15)
+
+
+def _config(rank: int, port: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        tp_info=DistributedInfo(rank, WORLD_SIZE),
+        distributed_timeout=TIMEOUT.total_seconds(),
+        distributed_port=port,
+    )
+
+
+def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None:
+    engine = Engine.__new__(Engine)  # only _make_distributed_store is under test
+    store = engine._make_distributed_store(_config(rank, port))
+    dist.init_process_group(
+        backend="gloo", rank=rank, world_size=WORLD_SIZE, timeout=TIMEOUT, store=store
+    )
+    total = torch.tensor([float(rank)])
+    dist.all_reduce(total)
+    bind_host = engine._distributed_listen_socket.getsockname()[0] if rank == 0 else None
+    result_q.put((rank, total.item(), bind_host))
+    dist.destroy_process_group()
+
+
+def test_two_ranks_rendezvous_with_loopback_only_store():
+    ctx = mp.get_context("spawn")
+    port = 29511
+    result_q = ctx.Queue()
+    procs = [ctx.Process(target=_run_rank, args=(rank, port, result_q)) for rank in range(WORLD_SIZE)]
+    for p in procs:
+        p.start()
+
+    results = {}
+    for _ in procs:
+        rank, total, bind_host = result_q.get(timeout=30)
+        results[rank] = (total, bind_host)
+    for p in procs:
+        p.join(timeout=30)
+        assert p.exitcode == 0
+
+    # both ranks completed the same all_reduce over the store -- rendezvous worked
+    assert results[0][0] == results[1][0] == 1.0  # sum(0, 1)
+    # the master's listening socket was pre-bound to loopback, not the wildcard
+    # address the C10d TCPStore server binds to on its own
+    assert results[0][1] == "127.0.0.1"
diff --git a/tests/server/test_dist_port_arg.py b/tests/server/test_dist_port_arg.py
new file mode 100644
index 0000000000..a49def7b65
--- /dev/null
+++ b/tests/server/test_dist_port_arg.py
@@ -0,0 +1,41 @@
+"""Regression test for #301: the TP rendezvous port must be independently configurable
+from --port, so two instances on adjacent ports don't have one's API port collide with
+the other's rendezvous port.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from freetoken.server.args import parse_args
+
+ANON_PATH = "/models/anon"
+
+
+class _Config:
+    def to_dict(self) -> dict:
+        return {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"}
+
+
+def _parse(argv: list[str]):
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()):
+        return parse_args(["--model", ANON_PATH, *argv])
+
+
+def test_dist_port_defaults_to_server_port_plus_one():
+    args, _ = _parse(["--port", "8081"])
+    assert args.distributed_port == 8082
+    assert args.distributed_addr == "tcp://127.0.0.1:8082"
+
+
+def test_dist_port_override_is_independent_of_server_port():
+    args, _ = _parse(["--port", "8082", "--dist-port", "9000"])
+    assert args.distributed_port == 9000
+    assert args.distributed_addr == "tcp://127.0.0.1:9000"
+
+
+def test_dist_port_rejects_out_of_range_value():
+    with pytest.raises(SystemExit):
+        _parse(["--dist-port", "70000"])

From 3882ef689c672588eae43c8365f1991878d0f8d5 Mon Sep 17 00:00:00 2001
From: taking-lying-flat <1615405@qq.com>
Date: Mon, 14 Sep 2026 09:06:55 +0800
Subject: [PATCH 033/249] fix(scheduler): match stop strings with incremental
 decoding

---
 python/freetoken/core.py                 |  2 ++
 python/freetoken/scheduler/scheduler.py  | 27 ++++++++++++++-----
 python/freetoken/tokenizer/detokenize.py | 33 ++++++++++++++++--------
 3 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 269b94ffe9..9171354e4d 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -12,6 +12,7 @@
     from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool
     from freetoken.kvcache.linear_state_pool import LinearStatePool
     from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.tokenizer.detokenize import DecodeStatus
 
 
 @dataclass
@@ -64,6 +65,7 @@ class Req:
     # handler must not free resources under an in-flight forward; it sets this flag and
     # _process_last_data frees the request when the batch drains (after copy_done.synchronize).
     aborted: bool = False
+    stop_decode_status: DecodeStatus | None = None
 
     def __post_init__(self) -> None:
         assert self.input_ids.is_cpu
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 28520e13df..ce28665b76 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -425,16 +425,31 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
         self.send_result(reply)
 
     def _match_stop_str(self, req: Req) -> str | None:
-        """First stop string present in this request's generated tail, else None. Decodes
-        only a short suffix (bounded by the longest stop string's char length, so a stop of
-        N chars spans at most N tokens) to keep the per-step cost small."""
+        """Match stops against the same incrementally decoded text as the frontend."""
+        from freetoken.tokenizer.detokenize import DecodeStatus
+
         stop_strs = req.sampling_params.stop_strs
         prompt_len = req.max_device_len - req.output_len
-        if len(req.input_ids) <= prompt_len:
+        end = len(req.input_ids)
+        # The frontend omits a terminal EOS, including at the output limit with ignore_eos.
+        if (
+            end > prompt_len
+            and not req.can_decode
+            and int(req.input_ids[-1]) in self.eos_token_ids
+        ):
+            end -= 1
+        if end <= prompt_len:
             return None
+        if req.stop_decode_status is None:
+            req.stop_decode_status = DecodeStatus(
+                decoded_ids=[], decoded_str="", read_offset=0, surr_offset=0, sent_offset=0
+            )
+        state = req.stop_decode_status
+        state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist())
+        tail = state.decode(self.tokenizer)
         max_chars = max(len(s) for s in stop_strs)
-        tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1))
-        tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist())
+        # Bound the text used for matching, retaining token context for incomplete characters.
+        state.decoded_str = state.decoded_str[-max_chars:]
         for s in stop_strs:
             if s in tail:
                 return s
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index d138ac1a68..26f2f4f9eb 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -72,6 +72,21 @@ class DecodeStatus:
     surr_offset: int  # length of surr ids
     sent_offset: int  # length of sent out string
 
+    def update(self, read_str: str, surr_str: str) -> str:
+        new_text = read_str[len(surr_str) :]
+        if len(new_text) > 0 and not new_text.endswith("�"):
+            self.decoded_str += new_text
+            self.surr_offset = self.read_offset
+            self.read_offset = len(self.decoded_ids)
+            return self.decoded_str
+        return self.decoded_str + find_printable_text(new_text)
+
+    def decode(self, tokenizer: PreTrainedTokenizerBase) -> str:
+        return self.update(
+            tokenizer.decode(self.decoded_ids[self.surr_offset :]),
+            tokenizer.decode(self.decoded_ids[self.surr_offset : self.read_offset]),
+        )
+
 
 class DetokenizeManager:
     def __init__(
@@ -91,6 +106,9 @@ def discard(self, uid: int) -> None:
         self.decode_map.pop(uid, None)
 
     def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
+        # Each message must advance its request's decode state before the next one.
+        if len({msg.uid for msg in msgs}) != len(msgs):
+            return [self.detokenize([msg])[0] for msg in msgs]
         read_ids: List[List[int]] = []
         surr_ids: List[List[int]] = []
         for msg in msgs:
@@ -103,7 +121,9 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
                     sent_offset=0,
                 )
             s = self.decode_map[msg.uid]
-            if not (msg.finished and msg.next_token in self.eos_token_ids):
+            if not (
+                msg.finished and not msg.matched_stop and msg.next_token in self.eos_token_ids
+            ):
                 s.decoded_ids.append(msg.next_token)
             read_ids.append(s.decoded_ids[s.surr_offset :])
             surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset])
@@ -114,16 +134,7 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
         incremental_strs: List[str] = []
         for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True):
             s = self.decode_map[msg.uid]
-            new_text = read_str[len(surr_str) :]
-            # Streaming chunk: update the decode status
-            if len(new_text) > 0 and not new_text.endswith("�"):
-                output_str = s.decoded_str + new_text
-                s.decoded_str = output_str
-                s.surr_offset = s.read_offset
-                s.read_offset = len(s.decoded_ids)
-            else:
-                new_text = find_printable_text(new_text)
-                output_str = s.decoded_str + new_text
+            output_str = s.update(read_str, surr_str)
 
             prev_sent = s.sent_offset
             if msg.finished:

From 40816fce85ad5148d419b61eeebae5d2d21ae9ea Mon Sep 17 00:00:00 2001
From: Cyber-Marty <15370813896@163.com>
Date: Mon, 14 Sep 2026 15:52:49 +0800
Subject: [PATCH 034/249] fix(utils): tolerate coalesced msgpack frames in zmq
 pull queues

ZmqPullQueue/ZmqSubQueue/ZmqAsyncPullQueue decoded each received frame
with msgpack.unpackb, which raises ExtraData when a frame carries more
than one packed object. On Windows with offload MoE this surfaced in
production as two mid-batch worker deaths in 81 requests: the tokenizer
and scheduler processes crashed on the same helper
(utils/mp.py), taking the whole API server down.

Route every received frame through one buffered msgpack.Unpacker per
queue: the first object is returned by get() and any remainder is held
for subsequent calls instead of crashing. empty() accounts for the
buffer, and get_raw()/decode() stay unbuffered (the multi-rank
broadcast path pairs them with a rank-wide count that a buffered
remainder would desynchronize) -- mixing them with buffered get() now
raises instead of silently skipping messages.

Regression tests cover the coalesced frame on sync pull, sub queue,
decode buffer, and empty() accounting; the async variant is skipped on
Windows (zmq.asyncio needs a Selector loop; prod runs on Linux CI).

Fixes #452
---
 python/freetoken/utils/mp.py     |  69 +++++++++++++++---
 tests/utils/__init__.py          |   0
 tests/utils/test_mp_coalesced.py | 116 +++++++++++++++++++++++++++++++
 3 files changed, 175 insertions(+), 10 deletions(-)
 create mode 100644 tests/utils/__init__.py
 create mode 100644 tests/utils/test_mp_coalesced.py

diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py
index 181baf1a07..f91498fbb2 100644
--- a/python/freetoken/utils/mp.py
+++ b/python/freetoken/utils/mp.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Callable, Dict, Generic, TypeVar
+from typing import Any, Callable, Dict, Generic, TypeVar
 
 import msgpack
 import zmq
@@ -9,6 +9,36 @@
 T = TypeVar("T")
 
 
+class _CoalescedUnpacker:
+    """
+    Buffered msgpack decoder tolerant of coalesced frames.
+
+    A single ZMQ frame should carry exactly one msgpack object, but on some platforms
+    (issue #452: Windows, offload MoE) a read can surface a frame containing two packed
+    objects back to back. ``msgpack.unpackb`` raises ``ExtraData`` in that case and kills
+    the worker process. Feeding every received frame through one ``Unpacker`` instead
+    decodes the first object and buffers the remainder, so the next ``get()`` returns it
+    in order instead of crashing.
+    """
+
+    def __init__(self) -> None:
+        self._unpacker = msgpack.Unpacker(raw=False)
+        self._pending: list[Any] = []
+
+    def feed(self, frame: bytes) -> None:
+        """Queue up every msgpack object contained in ``frame`` (usually exactly one)."""
+        self._unpacker.feed(frame)
+        for obj in self._unpacker:
+            self._pending.append(obj)
+
+    def take(self) -> Any:
+        """Pop the oldest undelivered object; raises StopIteration when the buffer is empty."""
+        return self._pending.pop(0)
+
+    def __len__(self) -> int:
+        return len(self._pending)
+
+
 class ZmqPushQueue(Generic[T]):
     def __init__(
         self,
@@ -62,19 +92,30 @@ def __init__(
         self.socket = self.context.socket(zmq.PULL)
         self.socket.bind(addr) if create else self.socket.connect(addr)
         self.decoder = decoder
+        self._unpacker = _CoalescedUnpacker()
 
     def get(self) -> T:
-        event = self.socket.recv()
-        return self.decoder(msgpack.unpackb(event, raw=False))
+        if not len(self._unpacker):
+            event = self.socket.recv()
+            self._unpacker.feed(event)
+        obj = self._unpacker.take()
+        return self.decoder(obj)
 
     def get_raw(self) -> bytes:
+        # Raw path (multi-rank broadcast) must stay unbuffered: callers pair empty()/get_raw()
+        # with a rank-wide count, so a buffered remainder would desynchronize the loop.
+        if len(self._unpacker):
+            raise RuntimeError(
+                "get_raw() cannot be mixed with buffered get() on the same queue: "
+                "the unpacker holds undelivered objects that would be skipped."
+            )
         return self.socket.recv()
 
     def decode(self, raw: bytes) -> T:
         return self.decoder(msgpack.unpackb(raw, raw=False))
 
     def empty(self) -> bool:
-        return self.socket.poll(timeout=0) == 0
+        return len(self._unpacker) == 0 and self.socket.poll(timeout=0) == 0
 
     def stop(self):
         self.socket.close()
@@ -92,10 +133,14 @@ def __init__(
         self.socket = self.context.socket(zmq.PULL)
         self.socket.bind(addr) if create else self.socket.connect(addr)
         self.decoder = decoder
+        self._unpacker = _CoalescedUnpacker()
 
     async def get(self) -> T:
-        event = await self.socket.recv()
-        return self.decoder(msgpack.unpackb(event, raw=False))
+        if not len(self._unpacker):
+            event = await self.socket.recv()
+            self._unpacker.feed(event)
+        obj = self._unpacker.take()
+        return self.decoder(obj)
 
     def stop(self):
         self.socket.close()
@@ -138,14 +183,18 @@ def __init__(
         self.socket.bind(addr) if create else self.socket.connect(addr)
         self.socket.setsockopt_string(zmq.SUBSCRIBE, "")
         self.decoder = decoder
+        self._unpacker = _CoalescedUnpacker()
 
     def get(self) -> T:
-        event = self.socket.recv()
-        return self.decoder(msgpack.unpackb(event, raw=False))
+        if not len(self._unpacker):
+            event = self.socket.recv()
+            self._unpacker.feed(event)
+        obj = self._unpacker.take()
+        return self.decoder(obj)
 
     def empty(self) -> bool:
-        return self.socket.poll(timeout=0) == 0
+        return len(self._unpacker) == 0 and self.socket.poll(timeout=0) == 0
 
     def stop(self):
         self.socket.close()
-        self.context.term()
+        self.context.term()
\ No newline at end of file
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py
new file mode 100644
index 0000000000..e599db7265
--- /dev/null
+++ b/tests/utils/test_mp_coalesced.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for coalesced msgpack frames on ZMQ queues (issue #452).
+
+A single ZMQ frame should carry exactly one msgpack object, but a read can surface
+a frame containing two packed objects back to back. ``msgpack.unpackb`` raises
+``ExtraData`` there and the worker process dies. The buffered unpacker decodes the
+first object and holds the remainder for the next get().
+"""
+
+import asyncio
+
+import msgpack
+import pytest
+
+from freetoken.utils.mp import (
+    ZmqAsyncPullQueue,
+    ZmqAsyncPushQueue,
+    ZmqPubQueue,
+    ZmqPullQueue,
+    ZmqPushQueue,
+    ZmqSubQueue,
+)
+
+
+def _identity_decoder(obj):
+    return obj
+
+
+def _addr(port):
+    return f"tcp://127.0.0.1:{port}"
+
+
+class TestCoalescedFrames:
+    """A frame carrying two packed objects must yield two get() calls, not a crash."""
+
+    def test_sync_pull_coalesced_frame(self):
+        push = ZmqPushQueue(_addr(5601), create=True, encoder=lambda o: o)
+        pull = ZmqPullQueue(_addr(5601), create=False, decoder=_identity_decoder)
+        try:
+            # Simulate the coalesced frame: two packb outputs concatenated.
+            frame = msgpack.packb({"i": 1}, use_bin_type=True) + msgpack.packb({"i": 2}, use_bin_type=True)
+            pull._unpacker.feed(frame)
+
+            assert pull.get() == {"i": 1}
+            assert pull.get() == {"i": 2}
+            assert len(pull._unpacker) == 0
+        finally:
+            push.stop()
+            pull.stop()
+
+    def test_decode_coalesced_raw(self):
+        from freetoken.utils.mp import _CoalescedUnpacker
+
+        unpacker = _CoalescedUnpacker()
+        frame = msgpack.packb("first", use_bin_type=True) + msgpack.packb("second", use_bin_type=True)
+        unpacker.feed(frame)
+        assert unpacker.take() == "first"
+        assert unpacker.take() == "second"
+
+    def test_empty_still_true_with_buffered_only(self):
+        push = ZmqPushQueue(_addr(5602), create=True, encoder=lambda o: o)
+        pull = ZmqPullQueue(_addr(5602), create=False, decoder=_identity_decoder)
+        try:
+            # nothing on socket and no buffer -> empty
+            assert pull.empty()
+            # a coalesced frame arrives and is fully buffered -> socket empty but buffer has 1 left
+            push.put({"x": 1})
+            pull._unpacker.feed(pull.get_raw())
+            assert pull.get() == {"x": 1}
+            assert pull.empty()
+        finally:
+            push.stop()
+            pull.stop()
+
+    def test_sub_queue_coalesced(self):
+        import msgpack
+
+        pub = ZmqPubQueue(_addr(5603), create=True, encoder=lambda o: o)
+        sub = ZmqSubQueue(_addr(5603), create=False, decoder=_identity_decoder)
+        try:
+            import time
+
+            time.sleep(0.1)  # SUB subscription propagation
+            pub.put({"a": 1})
+            pub.put({"a": 2})
+            time.sleep(0.1)
+            # Deliver both objects even if they arrive in one frame.
+            assert sub.get() == {"a": 1}
+            assert sub.get() == {"a": 2}
+        finally:
+            pub.stop()
+            sub.stop()
+
+
+async def test_async_pull_coalesced():
+    import sys
+
+    from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull
+    from freetoken.utils.mp import ZmqAsyncPushQueue as _AsyncPush
+
+    # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI.
+    if sys.platform == "win32":
+        pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)")
+    push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o)
+    pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder)
+    try:
+        await push.put({"n": 1})
+        await push.put({"n": 2})
+        await asyncio.sleep(0.1)
+
+        # Even if both messages arrive in a single frame, two gets must succeed.
+        assert await pull.get() == {"n": 1}
+        assert await pull.get() == {"n": 2}
+    finally:
+        push.stop()
+        pull.stop()
\ No newline at end of file

From 7a15be2248989b253f1f188cbf41bec6399324d2 Mon Sep 17 00:00:00 2001
From: Rafael Rincon <rafaeldrincon@gmail.com>
Date: Mon, 14 Sep 2026 15:30:28 -0400
Subject: [PATCH 035/249] feat(server): add --default-thinking-mode flag for
 server-wide thinking default
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

OpenAI-compatible clients that never send chat_template_kwargs (Vercel AI
SDK openai-compatible provider, llama-swap, LiteLLM proxies) currently get
the model's template default, which for reasoning models like Qwen3.6 means
thinking is always on — the token budget goes to reasoning_content and
content comes back empty.

--default-thinking-mode {auto,chat,thinking} lets the operator pick the
server-wide default:

- auto (default): current behavior, template decides
- chat: fills enable_thinking=False into every request that does not set
  any explicit thinking control
- thinking: fills enable_thinking=True the same way

An explicit per-request value (enable_thinking / thinking / thinking_mode
in chat_template_kwargs, or the DeepSeek thinking.type wire toggle) always
wins over the server default.

Applied to all three frontends: OpenAI chat completions, Anthropic
messages, and Responses API.

Closes #472
---
 python/freetoken/server/anthropic_api.py   |  5 +++
 python/freetoken/server/args.py            | 20 +++++++++
 python/freetoken/server/openai_api.py      | 44 +++++++++++++++++--
 python/freetoken/server/responses_api.py   |  8 +++-
 tests/server/test_default_thinking_mode.py | 51 ++++++++++++++++++++++
 5 files changed, 124 insertions(+), 4 deletions(-)
 create mode 100644 tests/server/test_default_thinking_mode.py

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 0d60940c97..68b006a40b 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -120,6 +120,7 @@ async def handle_anthropic_messages(
             default_max_tokens=(
                 getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS
             ),
+            default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"),
         )
         uid = await submit_generation(spec, state)
     except ValueError as exc:
@@ -308,10 +309,14 @@ def convert_anthropic_to_genspec(
     model_sampling: dict[str, Any],
     reasoning_parser: str | None = None,
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    default_thinking_mode: str | None = None,
 ) -> GenSpec:
+    from .openai_api import apply_default_thinking_mode
+
     messages, template_tools, parser_tools, ctk = convert_anthropic_prompt(
         req, reasoning_parser=reasoning_parser
     )
+    ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
     return GenSpec(
         messages=messages,
         sampling_params=resolve_sampling(
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..ca640fa1da 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -49,6 +49,12 @@ class ServerArgs(SchedulerConfig):
     # Reasoning parser that splits <think> reasoning from content for OpenAI
     # responses. None disables it (default for models without a reasoning protocol).
     reasoning_parser: str | None = None
+    # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the
+    # current per-request behavior; "chat" forces enable_thinking=False for every request
+    # that does not explicitly set it in chat_template_kwargs (for OpenAI-compatible
+    # clients that never send template kwargs, like Vercel AI SDK or llama-swap);
+    # "thinking" forces thinking on the same way.
+    default_thinking_mode: str = "auto"
     # "model": fill unspecified request sampling params from generation_config.json
     # (temperature/top_k/top_p), like sglang. "none": use framework defaults only.
     sampling_defaults: str = "model"
@@ -583,6 +589,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--default-thinking-mode",
+        type=str,
+        default=ServerArgs.default_thinking_mode,
+        choices=["auto", "chat", "thinking"],
+        help=(
+            "Server-wide default thinking mode for reasoning-capable models. 'auto' keeps "
+            "the current per-request behavior; 'chat' forces enable_thinking=False for "
+            "every request that does not explicitly set it in chat_template_kwargs (for "
+            "OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK "
+            "or llama-swap); 'thinking' forces thinking on the same way."
+        ),
+    )
+
     parser.add_argument(
         "--moe-strategy",
         default=ServerArgs.moe_strategy,
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index dc2f73a979..2b2973e5b9 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -56,15 +56,44 @@ def _thinking_type(req: Any) -> str | None:
 
 
 
+def apply_default_thinking_mode(
+    ctk: dict[str, Any] | None,
+    default_mode: str | None,
+) -> dict[str, Any] | None:
+    """Merge the server's --default-thinking-mode into a request's template kwargs.
+
+    An explicit per-request value always wins; the server default only fills
+    in what the request left unset. "auto" (the default) is a no-op.
+    """
+    if not default_mode or default_mode == "auto":
+        return ctk
+    if ctk is None:
+        ctk = {}
+    else:
+        ctk = dict(ctk)
+    has_explicit = (
+        "enable_thinking" in ctk
+        or "thinking" in ctk
+        or "thinking_mode" in ctk
+    )
+    if not has_explicit:
+        if default_mode == "chat":
+            ctk["enable_thinking"] = False
+        elif default_mode == "thinking":
+            ctk["enable_thinking"] = True
+    return ctk
+
+
 def chat_request_to_genspec(
     req: ChatCompletionRequest,
     model_sampling: dict[str, Any],
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    default_thinking_mode: str | None = None,
 ) -> GenSpec:
     """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params')."""
     from .model_meta import effort_toggle_kwargs
 
-    ctk = req.chat_template_kwargs
+    ctk = apply_default_thinking_mode(req.chat_template_kwargs, default_thinking_mode)
     thinking_type = _thinking_type(req)
     if req.reasoning_effort or thinking_type:
         ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type)
@@ -185,7 +214,12 @@ async def handle_chat_completion(
         default_max_tokens = (
             getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS
         )
-        spec = chat_request_to_genspec(req, model_sampling, default_max_tokens=default_max_tokens)
+        spec = chat_request_to_genspec(
+            req,
+            model_sampling,
+            default_max_tokens=default_max_tokens,
+            default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"),
+        )
     except ValueError as exc:
         return create_error_response(str(exc))
 
@@ -245,7 +279,11 @@ async def stream_chat_completion_chunks(
 ) -> AsyncIterator[bytes]:
     """Format generate_events() into the OpenAI chat.completion.chunk SSE stream."""
     if spec is None:
-        spec = chat_request_to_genspec(req, {})
+        spec = chat_request_to_genspec(
+            req,
+            {},
+            default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"),
+        )
     yield _sse(
         _chat_chunk(
             req,
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index 71d6eed873..6fecdbdc16 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -155,6 +155,7 @@ async def handle_responses(
         spec = convert_responses_to_genspec(
             req, model_sampling, default_max_tokens=default_max,
             reasoning_parser=getattr(state.config, "reasoning_parser", None),
+            default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"),
         )
         uid = await submit_generation(spec, state)
     except GenerationError as exc:
@@ -190,6 +191,7 @@ def convert_responses_to_genspec(
     model_sampling: dict[str, Any],
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
     reasoning_parser: str | None = None,
+    default_thinking_mode: str | None = None,
 ) -> GenSpec:
     # Collect every system/developer text — the top-level `instructions` PLUS any
     # system/developer-role input items (codex sends both: a system prompt as `instructions`
@@ -234,7 +236,11 @@ def convert_responses_to_genspec(
 
     from .model_meta import effort_toggle_kwargs
 
-    ctk = dict(getattr(req, "chat_template_kwargs", None) or {})
+    from .openai_api import apply_default_thinking_mode
+
+    ctk = apply_default_thinking_mode(
+        dict(getattr(req, "chat_template_kwargs", None) or {}), default_thinking_mode
+    )
     if req.reasoning:
         ctk = effort_toggle_kwargs(req.reasoning.get("effort"), ctk)
 
diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py
new file mode 100644
index 0000000000..ac979e625d
--- /dev/null
+++ b/tests/server/test_default_thinking_mode.py
@@ -0,0 +1,51 @@
+"""Unit tests for the --default-thinking-mode server flag merge logic."""
+from __future__ import annotations
+
+from freetoken.server.openai_api import apply_default_thinking_mode
+
+
+def test_auto_is_noop():
+    assert apply_default_thinking_mode(None, "auto") is None
+    assert apply_default_thinking_mode({"enable_thinking": True}, "auto") == {
+        "enable_thinking": True
+    }
+    assert apply_default_thinking_mode(None, None) is None
+
+
+def test_chat_fills_unset_kwargs():
+    assert apply_default_thinking_mode(None, "chat") == {"enable_thinking": False}
+    assert apply_default_thinking_mode({}, "chat") == {"enable_thinking": False}
+    assert apply_default_thinking_mode({"other": 1}, "chat") == {
+        "other": 1,
+        "enable_thinking": False,
+    }
+
+
+def test_thinking_fills_unset_kwargs():
+    assert apply_default_thinking_mode(None, "thinking") == {"enable_thinking": True}
+    assert apply_default_thinking_mode({}, "thinking") == {"enable_thinking": True}
+
+
+def test_explicit_request_value_wins():
+    # Explicit per-request values override the server default, both directions.
+    assert apply_default_thinking_mode({"enable_thinking": True}, "chat") == {
+        "enable_thinking": True
+    }
+    assert apply_default_thinking_mode({"enable_thinking": False}, "thinking") == {
+        "enable_thinking": False
+    }
+    assert apply_default_thinking_mode({"thinking": False}, "chat") == {"thinking": False}
+    assert apply_default_thinking_mode({"thinking_mode": "thinking"}, "chat") == {
+        "thinking_mode": "thinking"
+    }
+
+
+def test_original_dict_not_mutated():
+    original = {"other": 1}
+    result = apply_default_thinking_mode(original, "chat")
+    assert result == {"other": 1, "enable_thinking": False}
+    assert original == {"other": 1}
+
+
+def test_unknown_mode_is_noop():
+    assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1}
\ No newline at end of file

From 6a47d2d3026fccf8af0f7ebaf88c6fa9e6fdca25 Mon Sep 17 00:00:00 2001
From: cue <gxf.alpha@gmail.com>
Date: Mon, 14 Sep 2026 22:48:37 -0700
Subject: [PATCH 036/249] feat(scheduler): interleave decode steps into long
 prefill runs

_schedule_next_batch gives prefill unconditional priority (or-short-circuit), so a
long prompt prefilled in chunks occupies that many consecutive scheduler steps
during which in-flight decodes are never scheduled. Measured on the live service:
four cold prefills of 680/797/506/250 chunks over 267/317/190/93 s, a
prefill-burst share of 70% in a 15-minute window, at 0.388 s per 384-token chunk.
An already-decoding request therefore waited up to ~5 minutes for its next token.

Reversing the order outright would starve prefill, and a request that never
prefills never starts. The scheduler already carried the marker for this: a TODO
reading "support other policies: e.g. DECODE first".

Add --decode-interleave-every N: prefill keeps priority, but after N consecutive
prefill steps one decode step is taken if a decode is runnable. At N=8 that bounds
the stall to ~3 s instead of ~300 s, for about 1% of prefill wall time. Unset
keeps the historical order bit-for-bit, and a Scheduler constructed without the
policy (as the accounting tests build it) still works.

Assisted-by: Claude Code
---
 python/freetoken/engine/config.py        |   5 +
 python/freetoken/scheduler/interleave.py |  73 ++++++++
 python/freetoken/scheduler/scheduler.py  |  36 +++-
 python/freetoken/server/args.py          |  14 ++
 tests/scheduler/test_interleave.py       | 207 +++++++++++++++++++++++
 5 files changed, 330 insertions(+), 5 deletions(-)
 create mode 100644 python/freetoken/scheduler/interleave.py
 create mode 100644 tests/scheduler/test_interleave.py

diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 44c5a38e73..d90c02ecdd 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -77,6 +77,11 @@ class EngineConfig:
     # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime
     # rebuild can pin an absolute window instead).
     swa_full_tokens_ratio: float = 0.2
+    # Force one decode step after this many consecutive prefill steps, so a long chunked
+    # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first
+    # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE
+    # first"). At 8 the cost is ~1% of prefill wall time.
+    decode_interleave_every: int | None = None
     # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
diff --git a/python/freetoken/scheduler/interleave.py b/python/freetoken/scheduler/interleave.py
new file mode 100644
index 0000000000..76f293c8f5
--- /dev/null
+++ b/python/freetoken/scheduler/interleave.py
@@ -0,0 +1,73 @@
+"""Decode-interleave scheduling policy.
+
+Background (measured on production, 2026-09-15)
+-----------------------------------------------
+``Scheduler._schedule_next_batch`` picks::
+
+    batch = (prefill_manager.schedule_next_batch(budget)
+             or decode_manager.schedule_next_batch())
+
+Prefill wins unconditionally.  A long prompt is chunked at the window pool's
+budget (384 tokens in production), so a 300k-token context is 797 consecutive
+prefill steps -- and while they run, *decoding requests already in flight are
+never scheduled*.  Measured on the live service:
+
+    four long cold prefills: 680/797/506/250 chunks, 267/317/190/93 s
+    prefill-burst share of a 15-minute window: 70%
+
+so an in-flight request waits up to ~5 minutes for its next token.  The
+scheduler already carries the marker for this::
+
+    # TODO: support other policies: e.g. DECODE first
+
+Reversing the order outright is wrong -- it starves prefill, and a request that
+never prefills never starts.  What is needed is *interleaving*: let prefill keep
+priority, but force a decode step every N prefill steps so in-flight requests
+make progress.  The cost is small by construction: measured chunk cost is
+0.388 s, and one decode step for a handful of requests is a fraction of that,
+so every 8th step costs ~1% of prefill wall time while bounding the stall at
+8 chunks (~3 s) instead of the whole burst (~300 s).
+
+This module holds the *policy* only -- a counter and a decision -- so it can be
+unit-tested without a Scheduler, an engine, or a GPU.
+"""
+
+from __future__ import annotations
+
+
+class DecodeInterleavePolicy:
+    """Decide whether to force a decode step between prefill steps.
+
+    ``every`` is the number of consecutive prefill steps allowed before a decode
+    step is forced.  ``None`` or <= 0 disables the policy, reproducing the
+    historical prefill-first behaviour bit-for-bit.
+
+    The counter only advances on prefill steps that were actually scheduled, so
+    a scheduler with nothing to prefill never accumulates credit, and the first
+    decode of a burst is not delayed.
+    """
+
+    def __init__(self, every: int | None = None) -> None:
+        self.every = int(every) if every else None
+        self._prefill_streak = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.every is not None and self.every > 0
+
+    def note_prefill(self) -> None:
+        """Record that a prefill step was scheduled in this slot."""
+        self._prefill_streak += 1
+
+    def note_decode(self) -> None:
+        """Record that a decode step was scheduled in this slot (streak resets)."""
+        self._prefill_streak = 0
+
+    def wants_decode(self) -> bool:
+        """True when the streak has reached the threshold and a decode is due."""
+        if not self.enabled:
+            return False
+        return self._prefill_streak >= self.every
+
+    def reset(self) -> None:
+        self._prefill_streak = 0
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index ac1bf322e3..d605e66dce 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -31,6 +31,7 @@
 from .cache import CacheManager
 from .config import SchedulerConfig
 from .decode import DecodeManager
+from .interleave import DecodeInterleavePolicy
 from .io import SchedulerIOMixin
 from .mm import cut_image_spans, plan_mm_batch
 from .prefill import ChunkedReq, PrefillManager
@@ -91,6 +92,14 @@ def __init__(self, config: SchedulerConfig):
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
         )
         self.decode_manager = DecodeManager(config.page_size)
+        # Bound how long a run of prefill steps may starve in-flight decodes. A long
+        # prompt chunked at the window budget is hundreds of consecutive prefill steps;
+        # measured on an 8xRTX4090 DSV4 deployment, that left an already-decoding
+        # request unscheduled for 267-317 s. Unset reproduces the historical
+        # prefill-first order exactly.
+        self._interleave = DecodeInterleavePolicy(
+            getattr(config, "decode_interleave_every", None)
+        )
         self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups)
         self.prefill_manager = PrefillManager(
             self.cache_manager,
@@ -867,11 +876,28 @@ def _gather_multimodal(self, batch: Batch) -> None:
             )
 
     def _schedule_next_batch(self) -> ForwardInput | None:
-        # TODO: support other policies: e.g. DECODE first
-        batch = (
-            self.prefill_manager.schedule_next_batch(self.prefill_budget)
-            or self.decode_manager.schedule_next_batch()
-        )
+        # Prefill keeps priority -- a request that never prefills never starts -- but a
+        # long run of prefill chunks must not starve in-flight decodes. When the policy
+        # says a decode is due, take one if a decode is actually runnable; otherwise stay
+        # with prefill (never spend a slot idle just to keep a promise).
+        # getattr, not attribute access: the accounting tests drive this method on a
+        # partially built Scheduler, and interleaving must be an opt-in that a stripped
+        # object simply does not have rather than something that breaks it. A missing
+        # policy is the historical prefill-first order.
+        policy = getattr(self, "_interleave", None)
+        batch = None
+        if policy is not None and policy.wants_decode():
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None:
+                policy.note_decode()
+        if batch is None:
+            batch = self.prefill_manager.schedule_next_batch(self.prefill_budget)
+            if batch is not None and policy is not None:
+                policy.note_prefill()
+        if batch is None:
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None and policy is not None:
+                policy.note_decode()
         if batch is None:
             return None
         forward_input = self._prepare_batch(batch)
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..3166d66d72 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -407,6 +407,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--decode-interleave-every",
+        type=int,
+        default=ServerArgs.decode_interleave_every,
+        help=(
+            "Force one decode step after this many consecutive prefill steps. A long "
+            "prompt is prefilled in chunks, so it occupies that many consecutive "
+            "scheduler steps; measured on an 8xRTX4090 DSV4 deployment, a 300k-token "
+            "context left an already-decoding request unscheduled for 267-317 s. "
+            "Unset keeps the historical prefill-first order (the scheduler's own TODO "
+            "names this policy)."
+        ),
+    )
+
     parser.add_argument(
         "--page-size",
         type=int,
diff --git a/tests/scheduler/test_interleave.py b/tests/scheduler/test_interleave.py
new file mode 100644
index 0000000000..61eed3802e
--- /dev/null
+++ b/tests/scheduler/test_interleave.py
@@ -0,0 +1,207 @@
+"""Tests for DecodeInterleavePolicy and its wiring into Scheduler._schedule_next_batch.
+
+The production symptom this exists for (measured 2026-09-15): a 300k-token prompt
+is chunked at the window pool budget (384 tok), so it becomes 797 consecutive
+prefill steps, and an already-decoding request is not scheduled once during the
+267-317 s that takes.  Prefill must keep priority (a request that never prefills
+never starts) but must yield to decode periodically.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from freetoken.scheduler.interleave import DecodeInterleavePolicy
+
+
+# --------------------------------------------------------------------------- #
+# policy, in isolation
+# --------------------------------------------------------------------------- #
+def test_disabled_by_default_and_never_forces_decode():
+    """Unset reproduces historical prefill-first behaviour exactly."""
+    p = DecodeInterleavePolicy()
+    assert not p.enabled
+    for _ in range(10_000):
+        p.note_prefill()
+    assert not p.wants_decode()
+
+
+@pytest.mark.parametrize("every", [1, 2, 8, 16])
+def test_forces_decode_after_exactly_every_prefills(every):
+    p = DecodeInterleavePolicy(every)
+    assert p.enabled
+    for i in range(1, every):
+        p.note_prefill()
+        assert not p.wants_decode(), f"fired early at {i} < {every}"
+    p.note_prefill()
+    assert p.wants_decode(), f"did not fire at {every}"
+
+
+def test_decode_resets_the_streak():
+    p = DecodeInterleavePolicy(4)
+    for _ in range(4):
+        p.note_prefill()
+    assert p.wants_decode()
+    p.note_decode()
+    assert not p.wants_decode(), "streak must reset after a decode step"
+    for _ in range(3):
+        p.note_prefill()
+    assert not p.wants_decode()
+    p.note_prefill()
+    assert p.wants_decode()
+
+
+def test_every_one_interleaves_strictly():
+    """every=1 is the extreme: decode after every single prefill step."""
+    p = DecodeInterleavePolicy(1)
+    p.note_prefill()
+    assert p.wants_decode()
+
+
+@pytest.mark.parametrize("bad", [0, -1, -100])
+def test_non_positive_disables(bad):
+    p = DecodeInterleavePolicy(bad)
+    assert not p.enabled
+    for _ in range(50):
+        p.note_prefill()
+    assert not p.wants_decode()
+
+
+def test_reset_clears_credit():
+    p = DecodeInterleavePolicy(8)
+    for _ in range(7):
+        p.note_prefill()
+    p.reset()
+    assert not p.wants_decode()
+    p.note_prefill()
+    assert not p.wants_decode()
+
+
+# --------------------------------------------------------------------------- #
+# the decision the scheduler actually makes
+# --------------------------------------------------------------------------- #
+def _decide(policy, has_prefill, has_decode, every=8):
+    """Mirror of the scheduler's new decision, as a pure function.
+
+    prefill wins unless the policy says decode is due AND a decode is available.
+    """
+    if has_prefill and not (policy.wants_decode() and has_decode):
+        policy.note_prefill()
+        return "prefill"
+    if has_decode:
+        policy.note_decode()
+        return "decode"
+    if has_prefill:
+        policy.note_prefill()
+        return "prefill"
+    return None
+
+
+def test_burst_is_punctuated_by_decode_every_n():
+    """797 prefill steps in a row -- the production shape -- must contain decode steps."""
+    policy = DecodeInterleavePolicy(8)
+    seq = []
+    for _ in range(797):
+        seq.append(_decide(policy, has_prefill=True, has_decode=True))
+    n_dec = seq.count("decode")
+    # A decode step resets the streak, so a cycle is every+1 slots (8 prefill + 1 decode):
+    # 797 slots hold 88 full cycles, i.e. 88 decode steps.
+    assert n_dec == 797 // 9, f"expected {797 // 9} decode steps, got {n_dec}"
+    # and no run of prefills exceeds the threshold
+    run = 0
+    worst = 0
+    for s in seq:
+        run = run + 1 if s == "prefill" else 0
+        worst = max(worst, run)
+    assert worst == 8, f"longest prefill run {worst} exceeds the threshold"
+
+
+def test_disabled_policy_yields_the_historical_sequence():
+    policy = DecodeInterleavePolicy()
+    seq = [_decide(policy, True, True) for _ in range(100)]
+    assert seq.count("decode") == 0, "disabled must never take a decode slot"
+
+
+def test_decode_is_not_forced_when_none_is_runnable():
+    """Nothing to decode -> prefill keeps the slot (no idle step)."""
+    policy = DecodeInterleavePolicy(2)
+    seq = [_decide(policy, has_prefill=True, has_decode=False) for _ in range(10)]
+    assert seq == ["prefill"] * 10
+
+
+def test_prefill_keeps_priority_below_the_threshold():
+    """The first N-1 slots of a burst must still go to prefill."""
+    policy = DecodeInterleavePolicy(8)
+    seq = [_decide(policy, True, True) for _ in range(7)]
+    assert seq.count("decode") == 0
+
+
+def test_pure_decode_phase_unaffected():
+    """With nothing to prefill, every slot is decode (no policy interference)."""
+    policy = DecodeInterleavePolicy(8)
+    seq = [_decide(policy, has_prefill=False, has_decode=True) for _ in range(50)]
+    assert seq == ["decode"] * 50
+
+
+def test_scheduler_without_policy_keeps_prefill_first():
+    """A Scheduler built without __init__ (as the accounting tests do) must not break.
+
+    _schedule_next_batch reads the policy through getattr because the accounting tests
+    build a Scheduler this way; a stripped object keeps the historical prefill-first
+    order instead of raising.
+    """
+    from types import SimpleNamespace
+
+    from freetoken.scheduler.scheduler import Scheduler
+
+    prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[])
+    decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[])
+    s = Scheduler.__new__(Scheduler)
+    s.prefill_budget = 384
+    s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch)
+    s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch)
+    s._prepare_batch = lambda value: value
+    s.send_result = lambda messages: None
+    # no _interleave attribute at all
+    assert not hasattr(s, "_interleave")
+    assert Scheduler._schedule_next_batch(s) is prefill_batch, "must stay prefill-first"
+
+
+def test_scheduler_with_policy_and_no_decode_stays_prefill():
+    """When no decode is runnable the policy must not steal a slot."""
+    from types import SimpleNamespace
+
+    from freetoken.scheduler.scheduler import Scheduler
+    from freetoken.scheduler.interleave import DecodeInterleavePolicy
+
+    prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[])
+    s = Scheduler.__new__(Scheduler)
+    s.prefill_budget = 384
+    s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch)
+    s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None)
+    s._prepare_batch = lambda value: value
+    s.send_result = lambda messages: None
+    s._interleave = DecodeInterleavePolicy(1)      # fires every step
+    for _ in range(5):
+        assert Scheduler._schedule_next_batch(s) is prefill_batch
+
+
+def test_scheduler_with_policy_takes_decode_when_due():
+    """Once the streak reaches the threshold, a runnable decode gets the slot."""
+    from types import SimpleNamespace
+
+    from freetoken.scheduler.scheduler import Scheduler
+    from freetoken.scheduler.interleave import DecodeInterleavePolicy
+
+    prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[])
+    decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[])
+    s = Scheduler.__new__(Scheduler)
+    s.prefill_budget = 384
+    s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch)
+    s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch)
+    s._prepare_batch = lambda value: value
+    s.send_result = lambda messages: None
+    s._interleave = DecodeInterleavePolicy(3)
+    got = [Scheduler._schedule_next_batch(s) for _ in range(4)]
+    assert got[0] is prefill_batch and got[1] is prefill_batch and got[2] is prefill_batch
+    assert got[3] is decode_batch, "4th slot must be the forced decode"

From e96388f748b64e7bc6f65f1118ae3f0b0df1092d Mon Sep 17 00:00:00 2001
From: cue <gxf.alpha@gmail.com>
Date: Mon, 14 Sep 2026 22:48:41 -0700
Subject: [PATCH 037/249] test(scheduler): exercise the real
 _schedule_next_batch, not a model of it

Review of the previous commit returned "merge with 2 test nits" and no
correctness defect, and named the gap precisely: several tests drove _decide, a
pure re-implementation of the decision written inside the test file. That is a
model of the code, not the code, so it can drift and the tests keep passing while
the real _schedule_next_batch diverges.

Adds coverage against the real method via Scheduler.__new__ with stub managers: a
truth table over policy=None/1/3/8 asserting the exact slot sequence, strict
alternation at every=1, the latch case (decode due but unavailable, prefills
continue, then decode fires exactly once), the pure-decode and no-work paths, a
policy that exists but is disabled, decode-manager call counts, and the
config -> Scheduler.__init__ -> CLI wiring.

Two expectations written from reasoning were wrong and the real scheduler
corrected them (every=1 starts with prefill, because the streak begins at 0),
which is the value this coverage adds.

Assisted-by: Claude Code
---
 tests/scheduler/test_interleave.py | 139 +++++++++++++++++++++++++++++
 1 file changed, 139 insertions(+)

diff --git a/tests/scheduler/test_interleave.py b/tests/scheduler/test_interleave.py
index 61eed3802e..42a3f7221e 100644
--- a/tests/scheduler/test_interleave.py
+++ b/tests/scheduler/test_interleave.py
@@ -205,3 +205,142 @@ def test_scheduler_with_policy_takes_decode_when_due():
     got = [Scheduler._schedule_next_batch(s) for _ in range(4)]
     assert got[0] is prefill_batch and got[1] is prefill_batch and got[2] is prefill_batch
     assert got[3] is decode_batch, "4th slot must be the forced decode"
+
+
+# --------------------------------------------------------------------------- #
+# Truth table against the REAL scheduler (codex review: _decide is a model of
+# the decision, not the decision -- it can drift from the code it describes).
+# --------------------------------------------------------------------------- #
+def _real_scheduler(every=None, prefill=True, decode=True):
+    """Build a Scheduler whose two managers are stubs with call counters."""
+    from types import SimpleNamespace
+
+    from freetoken.scheduler.scheduler import Scheduler
+    from freetoken.scheduler.interleave import DecodeInterleavePolicy
+
+    calls = {"prefill": 0, "decode": 0}
+    pf = SimpleNamespace(is_prefill=True, prompt_admissions=[])
+    dc = SimpleNamespace(is_prefill=False, prompt_admissions=[])
+
+    def _pf(budget):
+        calls["prefill"] += 1
+        return pf if prefill else None
+
+    def _dc():
+        calls["decode"] += 1
+        return dc if decode else None
+
+    s = Scheduler.__new__(Scheduler)
+    s.prefill_budget = 384
+    s.prefill_manager = SimpleNamespace(schedule_next_batch=_pf)
+    s.decode_manager = SimpleNamespace(schedule_next_batch=_dc)
+    s._prepare_batch = lambda value: value
+    s.send_result = lambda messages: None
+    if every is not None:
+        s._interleave = DecodeInterleavePolicy(every)
+    return s, calls, pf, dc
+
+
+@pytest.mark.parametrize("every,expected", [
+    # (policy N, expected sequence of slot winners over 9 slots)
+    (None, ["p"] * 9),                       # disabled -> always prefill
+    # streak starts at 0, so slot 0 is prefill; then strict alternation
+    (1, ["p", "d", "p", "d", "p", "d", "p", "d", "p"]),
+    (3, ["p", "p", "p", "d", "p", "p", "p", "d", "p"]),
+    (8, ["p"] * 8 + ["d"]),
+])
+def test_truth_table_against_real_scheduler(every, expected):
+    """The real _schedule_next_batch must produce exactly this slot sequence."""
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=every)
+    got = []
+    for _ in range(9):
+        b = Scheduler._schedule_next_batch(s)
+        got.append("p" if b is pf else ("d" if b is dc else "?"))
+    assert got == expected, f"every={every}: got {got}, want {expected}"
+
+
+def test_every_one_alternates_while_both_stay_runnable():
+    """every=1 is the degenerate case: one decode between every prefill."""
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=1)
+    seq = []
+    for _ in range(6):
+        b = Scheduler._schedule_next_batch(s)
+        seq.append("p" if b is pf else "d")
+    # slot 0 prefill (streak starts at 0), then alternation
+    assert seq == ["p", "d", "p", "d", "p", "d"]
+
+
+def test_latch_releases_when_decode_becomes_available():
+    """Due decode unavailable -> keep prefilling -> decode becomes available -> it fires.
+
+    Exercises the subtle path: the policy must not lose its 'decode is due' state
+    while waiting, and must not fire twice once it does.
+    """
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=2)
+    # phase 1: no decode runnable at all
+    s.decode_manager.schedule_next_batch = lambda: None
+    for _ in range(5):
+        assert Scheduler._schedule_next_batch(s) is pf
+    # phase 2: decode becomes available -> next slot must be decode
+    s.decode_manager.schedule_next_batch = lambda: dc
+    assert Scheduler._schedule_next_batch(s) is dc, "latch did not release"
+
+
+def test_no_work_returns_none():
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, _, _, _ = _real_scheduler(every=8, prefill=False, decode=False)
+    assert Scheduler._schedule_next_batch(s) is None
+
+
+def test_pure_decode_fallback_when_nothing_to_prefill():
+    """With no prefill available every slot is decode (no policy interference)."""
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=2, prefill=False, decode=True)
+    got = [Scheduler._schedule_next_batch(s) for _ in range(5)]
+    assert all(b is dc for b in got)
+
+
+def test_disabled_policy_object_not_just_missing_attribute():
+    """codex: cover a policy that EXISTS but is disabled, not only an absent one."""
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=None)
+    s._interleave = __import__(
+        "freetoken.scheduler.interleave", fromlist=["DecodeInterleavePolicy"]
+    ).DecodeInterleavePolicy(0)          # present but disabled
+    got = [Scheduler._schedule_next_batch(s) for _ in range(6)]
+    assert all(b is pf for b in got), "a disabled policy must never take a slot"
+
+
+def test_decode_manager_not_probed_when_policy_disabled():
+    """When disabled, the decode manager must not be consulted at all (no extra call)."""
+    from freetoken.scheduler.scheduler import Scheduler
+
+    s, calls, pf, dc = _real_scheduler(every=None)
+    Scheduler._schedule_next_batch(s)
+    assert calls["decode"] == 0, "disabled policy should not probe decode"
+
+
+def test_config_to_scheduler_wiring():
+    """codex: cover CLI/config -> Scheduler.__init__ wiring."""
+    import inspect
+
+    from freetoken.engine.config import EngineConfig
+    from freetoken.scheduler import scheduler as sched_mod
+
+    assert "decode_interleave_every" in {f.name for f in __import__("dataclasses").fields(EngineConfig)}
+    src = inspect.getsource(sched_mod.Scheduler.__init__)
+    assert "_interleave" in src and "decode_interleave_every" in src, (
+        "Scheduler.__init__ must build the policy from the config field"
+    )
+    # and the CLI exposes it
+    from freetoken.server.args import ServerArgs
+    assert hasattr(ServerArgs, "decode_interleave_every")

From b37eb800c89fa167f798567ddc72203d501e5a1a Mon Sep 17 00:00:00 2001
From: Rafael Rincon <rafaeldrincon@gmail.com>
Date: Tue, 15 Sep 2026 11:40:39 -0400
Subject: [PATCH 038/249] fix(server): hoist system messages to front for
 templates requiring system-first

Some chat templates (e.g. Qwen3.6) raise 'System message must be at the
beginning' when the system message is not at index 0. render_messages()
now reorders system messages to the front before passing to the template.

This is a no-op when system is already first or absent.
---
 python/freetoken/server/generation.py | 12 ++++++++--
 tests/server/test_mm_input.py         | 33 +++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 91056e8d1f..d2956a44fd 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -194,8 +194,16 @@ def pick(value, key, framework):
 def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
     """Normalize OpenAI-shaped message dicts for the chat template: flatten text
     content parts to a string and decode tool-call arguments from JSON. Raises
-    ValueError on a non-text content part (text-only server). Shared by all adapters."""
-    return [_render_message(m) for m in messages]
+    ValueError on a non-text content part (text-only server). Shared by all adapters.
+
+    Some chat templates (e.g. Qwen3.6) require the system message at index 0;
+    hoist system messages to the front to satisfy that constraint."""
+    rendered = [_render_message(m) for m in messages]
+    system_msgs = [m for m in rendered if m.get("role") == "system"]
+    if system_msgs and rendered[0].get("role") != "system":
+        non_system = [m for m in rendered if m.get("role") != "system"]
+        rendered = system_msgs + non_system
+    return rendered
 
 
 def _render_message(message: dict[str, Any]) -> dict[str, Any]:
diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py
index 19d9f86529..5ac5834bf5 100644
--- a/tests/server/test_mm_input.py
+++ b/tests/server/test_mm_input.py
@@ -140,3 +140,36 @@ def test_image_token_budget_flags_land_in_the_multimodal_config():
         assert parse_args(["--model", "/models/anon"])[0].mm.processor_kwargs == {}
         with pytest.raises(SystemExit):  # argparse reports the bad pair and exits
             parse_args(["--model", "/models/anon", "--image-min-tokens", "2048", "--image-max-tokens", "1024"])
+
+
+def test_render_messages_hoists_system_to_front():
+    """render_messages moves system messages to index 0 for templates that
+    require it (e.g. Qwen3.6 'System message must be at the beginning')."""
+    msgs = render_messages([
+        {"role": "user", "content": "What is 2+2?"},
+        {"role": "system", "content": "You are a math tutor."},
+        {"role": "user", "content": "Answer briefly."},
+    ])
+    assert msgs[0]["role"] == "system"
+    assert msgs[1]["role"] == "user"
+    assert msgs[2]["role"] == "user"
+
+
+def test_render_messages_preserves_system_first():
+    """When system is already first, no reordering happens."""
+    msgs = render_messages([
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Hi"},
+    ])
+    assert msgs[0]["role"] == "system"
+    assert msgs[1]["role"] == "user"
+
+
+def test_render_messages_no_system_unchanged():
+    """No system message → order preserved."""
+    msgs = render_messages([
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there"},
+        {"role": "user", "content": "Bye"},
+    ])
+    assert [m["role"] for m in msgs] == ["user", "assistant", "user"]

From 464125213e26ebcb7e957bec375e389c33d2a8ec Mon Sep 17 00:00:00 2001
From: Rafael Rincon <rafaeldrincon@gmail.com>
Date: Tue, 15 Sep 2026 23:05:37 -0400
Subject: [PATCH 039/249] fix: merge multiple system messages into one

Qwen3.6 chat template raises 'System message must be at the beginning'
when a second system message appears after loop.first. Merging all
system messages into a single one at index 0 fixes this.

Tested with: system+user+system, user+system+user (hoisted+merged).
---
 python/freetoken/server/generation.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index d2956a44fd..8dd9142867 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -197,12 +197,23 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
     ValueError on a non-text content part (text-only server). Shared by all adapters.
 
     Some chat templates (e.g. Qwen3.6) require the system message at index 0;
-    hoist system messages to the front to satisfy that constraint."""
+    hoist system messages to the front and merge multiples into one to satisfy
+    that constraint."""
     rendered = [_render_message(m) for m in messages]
     system_msgs = [m for m in rendered if m.get("role") == "system"]
     if system_msgs and rendered[0].get("role") != "system":
         non_system = [m for m in rendered if m.get("role") != "system"]
         rendered = system_msgs + non_system
+    # Merge multiple system messages into one (Qwen3.6 template rejects
+    # system messages that are not at loop.first, i.e. after the first).
+    system_msgs = [m for m in rendered if m.get("role") == "system"]
+    if len(system_msgs) > 1:
+        sep = chr(10) + chr(10)
+        merged_content = sep.join(
+            str(m.get("content") or "") for m in system_msgs if m.get("content")
+        )
+        non_system = [m for m in rendered if m.get("role") != "system"]
+        rendered = [{"role": "system", "content": merged_content}] + non_system
     return rendered
 
 

From 849e072bae02ba254968b3afa0e8360adbdab247 Mon Sep 17 00:00:00 2001
From: AlReis <258541404+alvarorsouza-arch@users.noreply.github.com>
Date: Thu, 17 Sep 2026 06:03:02 +0000
Subject: [PATCH 040/249] feat(kvcache): host-RAM KV tier for sparse-attention
 models (--kv-host-pages)

---
 python/freetoken/attention/qsa_sparse.py      |  82 +++-
 python/freetoken/core.py                      |   4 +
 python/freetoken/engine/config.py             |   5 +
 python/freetoken/engine/engine.py             |  58 ++-
 python/freetoken/kernel/triton/qsa/offload.py | 235 +++++++++++
 python/freetoken/kvcache/__init__.py          |  11 +-
 python/freetoken/kvcache/kv_host_offload.py   | 368 ++++++++++++++++++
 python/freetoken/kvcache/qsa_pool.py          |   7 +-
 python/freetoken/server/args.py               |  11 +
 9 files changed, 775 insertions(+), 6 deletions(-)
 create mode 100644 python/freetoken/kernel/triton/qsa/offload.py
 create mode 100644 python/freetoken/kvcache/kv_host_offload.py

diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index b07c952323..de48768c8c 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -92,6 +92,10 @@ class QSASparseMetadata(BaseAttnMetadata):
     rope_rows:        torch.Tensor | None = None  # [T] int32 arange
     q_rope_cache:     torch.Tensor | None = None  # [T, rotary_dim] float32
     k_rope_cache:     torch.Tensor | None = None  # [T, rotary_dim] float32
+    # KV host offload (--kv-host-pages): this forward's write pages (re-included in every
+    # fetch query so no mid-forward eviction can take them) and the physical out_loc.
+    write_pages:      torch.Tensor | None = None
+    out_loc_gpu:      torch.Tensor | None = None
     # fmt: on
 
     def get_last_indices(self, bs: int) -> torch.Tensor:
@@ -142,6 +146,12 @@ def __init__(self, config: ModelConfig) -> None:
             ).to(self.device)
 
         self._block_topk_kernel = _resolve_block_topk()
+        # KV host offload (None unless --kv-host-pages > 0): the GPU K/V pool is an LRU cache
+        # over the pinned host mirror; reads/writes translate logical -> physical pages.
+        self.kv_offloader = getattr(get_global_ctx(), "kv_offloader", None)
+        if self.kv_offloader is not None:
+            assert self.kv_offloader.pool is self.kvcache
+            self.kv_offloader.set_select_width(self.select_width, self.block_topk)
         # decode staging (static buffers under CUDA graphs; eager decode snapshots per step)
         self._graph: dict[str, torch.Tensor] = {}
         self.capture_bs: List[int] = []
@@ -288,7 +298,15 @@ def qsa_forward(
         md = batch.attn_metadata
         assert isinstance(md, QSASparseMetadata)
         slot = self._idx_slot[layer_id]
-        self.kvcache.store_kv(k, v, batch.out_loc, layer_id)
+        off = self.kv_offloader
+        if off is not None and slot == 0:
+            # Bind this forward's write pages before the first store; they ride every fetch
+            # query below, so no mid-forward eviction can take them (and the write-through
+            # mirror makes any other eviction lossless).
+            md.write_pages = off.ensure_write_pages(batch.out_loc, md.is_decode,
+                                                    md.token_to_req)
+            md.out_loc_gpu = off.translate_slots(batch.out_loc, md.is_decode)
+        self.kvcache.store_kv(k, v, batch.out_loc if off is None else md.out_loc_gpu, layer_id)
         if md.block_table is None:
             self._snapshot_decode(md, batch)
         if slot == 0 or md.cmp_rows is None:
@@ -298,7 +316,12 @@ def qsa_forward(
             self._plan_index_writes(md, batch)
 
         self._update_index_cache(index, md, slot)
+        if off is not None:
+            # Write-through: the pinned host mirror (the backing store) sees every token.
+            off.mirror_store(k, v, batch.out_loc, layer_id)
         indices = self._select(index, md, slot)
+        if off is not None:
+            return self._attend_offloaded(qsa_sparse_paged_attention, q, indices, md, layer_id)
         return qsa_sparse_paged_attention(
             q,
             self.kvcache.k_cache(layer_id),
@@ -309,6 +332,61 @@ def qsa_forward(
             torch.empty_like(q),
         )
 
+    def _attend_offloaded(self, attend, q, indices, md, layer_id) -> torch.Tensor:
+        """Attend through the host-mirrored page space: fetch the selected pages into
+        physical slots, translate the block table, call the attend kernel unmodified.
+        Prefill attends one row at a time (a row's scattered selection nearly fills the
+        physical pool) over a once-per-layer compacted page list -- no host syncs."""
+        off = self.kv_offloader
+        k_cache = self.kvcache.k_cache(layer_id)
+        v_cache = self.kvcache.v_cache(layer_id)
+        # fp8 pool (PR #354): the attend kernel dequantizes with these; None on a bf16 pool.
+        # Scale accessors only exist on a quantized KV pool (fp8, e.g. PR #354); the
+        # plain bf16 pool on main has none. Probe instead of assuming, so the offload
+        # composes with either.
+        _k_scale = getattr(self.kvcache, "k_scale", None)
+        _v_scale = getattr(self.kvcache, "v_scale", None)
+        ks = _k_scale(layer_id) if _k_scale is not None else None
+        vs = _v_scale(layer_id) if _v_scale is not None else None
+        # The attend kernel on an unquantized pool takes NO scale arguments at all, so the
+        # kwargs must be omitted entirely (passing k_scale=None is still unexpected).
+        scales: dict = {}
+        if ks is not None:
+            scales["k_scale"] = ks
+        if vs is not None:
+            scales["v_scale"] = vs
+        if md.is_decode:
+            eff = off.fetch_for_attend(indices, md)
+            return attend(q, k_cache, v_cache, indices, eff, md.token_to_req,
+                          torch.empty_like(q), **scales)
+        rows = q.shape[0]
+        if rows == 0:
+            return attend(q, k_cache, v_cache, indices, md.block_table, md.token_to_req,
+                          torch.empty_like(q), **scales)
+        # Row-group loop: one compact for the whole layer (grid = rows, no sync), then the
+        # offloader packs rows into residency-feasible groups CPU-side with REAL page unions
+        # (one D2H per layer) -- consecutive rows select mostly the same pages, so groups
+        # pack several rows. Per group: one tight ensure + one translate + one attend.
+        sel_all = off.compact_all(indices, md.token_to_req, md.block_table)
+        out = torch.empty_like(q)
+        for (start, end), eff in off.iter_prefill_groups(sel_all, md.write_pages, md.block_table):
+            attend(q[start:end], k_cache, v_cache, indices[start:end], eff,
+                   md.token_to_req[start:end], out[start:end], **scales)
+        dropped = off.trunc_count()  # one sync per layer
+        if dropped:
+            msg = (
+                f"KV host offload: {dropped} selected page(s) dropped this layer because a "
+                f"query row spans more than {off.max_sel_pages} distinct pages; those rows "
+                "would attend over unrelated physical slots. Raise --kv-host-pages or "
+                "--memory-ratio, or run with --max-running-requests 1."
+            )
+            if md.is_decode:
+                # Inside a CUDA graph we cannot raise without tearing down the capture.
+                logger.warning(msg)
+            else:
+                raise RuntimeError(msg)
+        return out
+
     def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None:
         """Per-token slab row and ring row for this forward; the other QSA layers reuse it
         (it is layer-invariant). Pure device arithmetic: no host sync, graph-capturable."""
@@ -518,6 +596,8 @@ def empty(*shape: int, dtype: torch.dtype) -> torch.Tensor:
         }
         if topk_scratch:
             self._graph["topk_scratch"] = empty(chunk, topk_scratch, dtype=torch.int32)
+        if self.kv_offloader is not None:
+            self.kv_offloader.init_graph_buffers(max_bs, pages, self.select_width)
 
     def prepare_for_capture(self, batch: Batch) -> None:
         self.prepare_metadata(batch)
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 269b94ffe9..f2d6485116 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -10,6 +10,7 @@
     from freetoken.attention import BaseAttnBackend, BaseAttnMetadata
     from freetoken.attention.linear import FLAMetadata
     from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool
+    from freetoken.kvcache.kv_host_offload import KVHostOffloader
     from freetoken.kvcache.linear_state_pool import LinearStatePool
     from freetoken.moe.offload_cache import OffloadMoeCache
 
@@ -181,6 +182,9 @@ class Context:
     attn_backend: BaseAttnBackend = field(init=False)
     moe_offload_cache: OffloadMoeCache | None = None
     kv_cache: BaseKVCachePool = field(init=False)
+    # KV host offload (--kv-host-pages): the QSA backend's logical->physical page bridge.
+    # Set by the engine before create_attention_backend; None when disabled.
+    kv_offloader: KVHostOffloader | None = None
     # Per-request recurrent state for GatedDeltaNet layers; set by the engine for
     # hybrid linear-attention models, otherwise None.
     linear_state_pool: LinearStatePool | None = None
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 44c5a38e73..5df70492bb 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -89,6 +89,11 @@ class EngineConfig:
     # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size
     # is final. Mutually exclusive with num_page_override.
     num_token_override: int | None = None
+    # Extra KV pages mirrored to pinned host RAM (QSA paged pools, e.g. Qwen3.8-Flash-Next):
+    # the GPU pool becomes an LRU cache over the logical page space (scheduler, page table and
+    # radix tree all see num_pages + kv_host_pages pages), extending context past VRAM
+    # capacity. 0 (default) = off.
+    kv_host_pages: int = 0
     # Runtime knobs of the multimodal path; the architecture side (vision_config, mrope) lives in ModelConfig.
     mm: MultimodalConfig = field(default_factory=MultimodalConfig)
 
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 2f0c89fbb1..47db7b4e91 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -46,6 +46,12 @@ def _require_offload_cache_size(cache_size: int, num_experts: int) -> None:
             f"cache-sizing flags)."
         )
 
+# Headroom the KV host offload keeps out of the GPU pool solve: the bigger logical index slab
+# is priced exactly, but prefill transients (GDN chunk workspaces, QSA score logits) also grow
+# with the logical width and need this slack (a 0.8 GiB post-capture margin OOMed a 4096-token
+# prefill; the pre-offload runs sat at ~1.05 GiB).
+_KV_OFFLOAD_SERVING_MARGIN = 512 << 20
+
 
 def _flashinfer_available() -> bool:
     from freetoken.kernel.backend import is_flashinfer_installed
@@ -389,10 +395,35 @@ def __init__(self, config: EngineConfig):
         # off it; the KV pool family owns every geometry-specific formula behind the rest.
         available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free)
         available_memory -= state_pool_bytes(config)
-        self.num_pages = self._pool_cls.solve_num_pages(config, available_memory)
+        num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory)
+        if config.kv_host_pages > 0:
+            # KV host offload: the scheduler, page table and radix tree run in a LOGICAL page
+            # space of num_gpu_pages + kv_host_pages; the GPU K/V buffer stays a
+            # num_gpu_pages LRU cache over a pinned host mirror (kvcache/kv_host_offload.py).
+            # The QSA index slab is per LOGICAL page, so reserve its growth before solving
+            # the GPU page count.
+            from freetoken.kvcache.qsa_pool import QSAKVCache
+
+            if not issubclass(self._pool_cls, QSAKVCache):
+                raise ValueError(
+                    f"--kv-host-pages requires the QSA paged KV pool, got {self._pool_cls.__name__}"
+                )
+            spec = next(s for s in config.model_config.kv_cache_group_specs() if s.num_layers > 0)
+            slab_per_page = (
+                (config.page_size // spec.index_ratio)
+                * spec.num_index_layers * spec.index_head_dim * 2
+            )
+            # The two-pass solve prices the logical-space slab growth AND hands back the
+            # serving-margin the growth would otherwise eat (prefill transients need it).
+            reserve = config.kv_host_pages * slab_per_page + _KV_OFFLOAD_SERVING_MARGIN
+            num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory - reserve)
+            self.num_pages = num_gpu_pages + config.kv_host_pages
+        else:
+            self.num_pages = num_gpu_pages
         num_tokens = self.num_pages * config.page_size
         self.ctx.kv_cache = self.kv_cache = create_kv_pool(
-            config, self.num_pages, device=self.device, dtype=self.dtype
+            config, num_gpu_pages, device=self.device, dtype=self.dtype,
+            num_index_pages=self.num_pages if config.kv_host_pages > 0 else None,
         )
 
         # ======================= Linear (GatedDeltaNet) state initialization ========================
@@ -426,6 +457,21 @@ def __init__(self, config: EngineConfig):
         # through them belongs to the attention backend, built later in init_capture_graph.
         self.kv_cache.attach_page_table(self.page_table)
 
+        # ======================= KV host offload ========================
+        # Attached before create_attention_backend: the QSA backend reads ctx.kv_offloader
+        # at construction.
+        self.kv_offloader = None
+        if config.kv_host_pages > 0:
+            from freetoken.kvcache.kv_host_offload import KVHostOffloader
+
+            self.kv_offloader = KVHostOffloader(self.kv_cache, self.num_pages, self.device)
+            self.ctx.kv_offloader = self.kv_offloader
+            logger.info_rank0(
+                f"Allocating {num_tokens} tokens for KV cache (logical): "
+                f"{num_gpu_pages} GPU-resident pages + {config.kv_host_pages} host pages, "
+                f"pinned mirror {mem_GB(self.kv_offloader.mirror_bytes)}"
+            )
+
         # ======================= Attention backend initialization ========================
         self.ctx.attn_backend = self.attn_backend = create_attention_backend(
             config.attention_backend, config.model_config
@@ -865,6 +911,14 @@ def rebuild_runtime_cache(
                 raise CacheRebuildRejected(str(e)) from e
         if num_pages is not None and num_pages <= 0:
             raise CacheRebuildRejected(f"num_pages must be positive, got {num_pages}")
+        if num_pages is not None and self.kv_offloader is not None:
+            # The logical/physical split (host mirror, LRU maps, QSA index slab sized to the
+            # logical space) is built once at startup; a runtime resize would have to rebuild
+            # all of it. Reject recoverably instead.
+            raise CacheRebuildRejected(
+                "num_pages rebuild is not supported with --kv-host-pages (restart with a "
+                "different --kv-host-pages / --memory-ratio instead)"
+            )
         if num_mamba_slots is not None:
             if self.linear_state_pool is None:
                 raise CacheRebuildRejected(
diff --git a/python/freetoken/kernel/triton/qsa/offload.py b/python/freetoken/kernel/triton/qsa/offload.py
new file mode 100644
index 0000000000..9b669bf5a7
--- /dev/null
+++ b/python/freetoken/kernel/triton/qsa/offload.py
@@ -0,0 +1,235 @@
+"""Triton kernels for the QSA KV host offload (``kvcache/kv_host_offload.py``).
+
+The page table, radix tree and compressed-index slab all run in LOGICAL page space; the GPU
+KV buffer is an LRU cache of physical slots over the pinned host mirror. These kernels are
+the fixed-shape, CUDA-graph-safe glue between the two spaces:
+
+* :func:`write_pages` -- logical page id of every written token slot (``out_loc // page_size``).
+* :func:`compact_selected_pages` -- per query row, the distinct logical pages the sparse
+  selection touches (``block_table[req, indices//page_size]``), first-element padded.
+* :func:`translate_table` -- logical block table -> physical slot table (post-ensure).
+* :func:`mirror_store` -- write-through of this forward's K/V rows into the pinned host
+  mirror (UVA stores), so a physical slot is always drop-clean.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _write_pages_kernel(out_loc, write_pages, page_size: tl.constexpr):
+    i = tl.program_id(0)
+    tl.store(write_pages + i, tl.load(out_loc + i) // page_size)
+
+
+def write_pages(out_loc: torch.Tensor, out: torch.Tensor, page_size: int) -> None:
+    """``out[i] = out_loc[i] // page_size`` for i < out.numel() (out_loc may be wider)."""
+    n = out.numel()
+    if n:
+        _write_pages_kernel[(n,)](out_loc, out, page_size)
+
+
+@triton.jit
+def _compact_sel_kernel(
+    indices,        # [rows, SEL] int32: selected LOGICAL token ids, -1 padded
+    token_to_req,   # [rows] int32
+    block_table,    # [reqs, W] int32: logical page ids
+    out,            # [rows, MAXP] int32: distinct logical pages per row
+    trunc,          # [1] int32: rows whose selection overflowed MAXP (0 = none dropped)
+    counts,         # [rows] int32: distinct pages stored per row (<= MAXP)
+    sel_width: tl.constexpr,
+    width,          # W: block_table row stride
+    page_size: tl.constexpr,
+    maxp: tl.constexpr,      # rows of `out` per row written (<= MAXV)
+    maxv: tl.constexpr,      # next_pow2(maxp)
+    dummy_page,              # logical id used when a row selects nothing (never happens in practice)
+):
+    r = tl.program_id(0)
+    req = tl.load(token_to_req + r).to(tl.int64)
+    lanes = tl.arange(0, maxv)
+    seen = tl.full([maxv], -1, tl.int32)
+    first = tl.full((), dummy_page, tl.int32)
+    count = tl.zeros((), tl.int32)
+    ndrop = tl.zeros((), tl.int32)
+    for j in range(sel_width):
+        tok = tl.load(indices + r * sel_width + j)
+        col = tok // page_size
+        pg = tl.load(block_table + req * width + col, mask=tok >= 0, other=-1)
+        new = (tok >= 0) & (tl.sum((seen == pg).to(tl.int32)) == 0)
+        take = new & (count < maxp)
+        ndrop += (new & (count >= maxp)).to(tl.int32)
+        seen = tl.where((lanes == count) & take, pg, seen)
+        first = tl.where((count == 0) & take, pg, first)
+        count += take.to(tl.int32)
+    row = tl.where(lanes < count, seen, first)
+    tl.store(out + r * maxp + lanes, row, mask=lanes < maxp)
+    tl.store(counts + r, count)
+    tl.atomic_add(trunc, ndrop)
+
+
+def compact_selected_pages(
+    indices: torch.Tensor,
+    token_to_req: torch.Tensor,
+    block_table: torch.Tensor,
+    out: torch.Tensor,
+    page_size: int,
+    dummy_page: int,
+    trunc: torch.Tensor | None = None,
+    counts: torch.Tensor | None = None,
+) -> None:
+    """Per query row, the distinct logical pages its selected tokens live on.
+
+    ``out`` is ``[rows, maxp]`` int32; rows with fewer hits are padded with the row's first
+    page (duplicates collapse in lru_ensure), or ``dummy_page`` when the row selected nothing.
+    ``trunc`` (optional [1] int32, device) accumulates the TOTAL number of pages dropped
+    across all rows because they exceeded ``maxp`` -- 0 means no truncation happened. ``counts`` (optional
+    [rows] int32, device) receives each row's stored distinct-page count.
+    """
+    rows, sel = indices.shape
+    maxp = out.shape[1]
+    if trunc is None:
+        trunc = getattr(compact_selected_pages, "_dummy", None)
+        if trunc is None or trunc.device != out.device:
+            trunc = compact_selected_pages._dummy = torch.zeros(1, dtype=torch.int32, device=out.device)
+    if counts is None:
+        c = getattr(compact_selected_pages, "_dummy_counts", None)
+        if c is None or c.numel() < rows or c.device != out.device:
+            c = compact_selected_pages._dummy_counts = torch.empty(
+                max(rows, 4096), dtype=torch.int32, device=out.device
+            )
+        counts = c
+    if rows:
+        _compact_sel_kernel[(rows,)](
+            indices, token_to_req, block_table, out, trunc, counts,
+            sel, block_table.stride(0), page_size, maxp, triton.next_power_of_2(maxp),
+            dummy_page,
+        )
+
+
+@triton.jit
+def _translate_kernel(table, phys_of, out, total, BLOCK: tl.constexpr):
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = i < total
+    pg = tl.load(table + i, mask=mask, other=0)
+    ph = tl.load(phys_of + pg, mask=mask, other=0)
+    # A -1 entry means the page was not ensured (only possible for a truncated selection,
+    # reported via the trunc counter): point it at slot 0 rather than indexing garbage.
+    tl.store(out + i, tl.where(ph < 0, 0, ph), mask=mask)
+
+
+def translate_table(table: torch.Tensor, phys_of: torch.Tensor, out: torch.Tensor) -> None:
+    """``out = phys_of[table]`` elementwise over ``out.numel()`` entries (all int32)."""
+    n = out.numel()
+    if n:
+        _translate_kernel[(triton.cdiv(n, 1024),)](table, phys_of, out, n, 1024)
+
+
+@triton.jit
+def _translate_slots_kernel(out_loc, phys_of, out, page_size: tl.constexpr):
+    i = tl.program_id(0)
+    slot = tl.load(out_loc + i)
+    tl.store(out + i, tl.load(phys_of + slot // page_size) * page_size + slot % page_size)
+
+
+def translate_slots(out_loc: torch.Tensor, phys_of: torch.Tensor, out: torch.Tensor, page_size: int) -> None:
+    """Logical token slot -> physical token slot (page translate, offset preserved)."""
+    n = out.numel()
+    if n:
+        _translate_slots_kernel[(n,)](out_loc, phys_of, out, page_size)
+
+
+@triton.jit
+def _mirror_store_kernel(
+    k_bytes,      # [T, ROW] uint8 view of this layer's K rows
+    v_bytes,      # [T, ROW] uint8 view of this layer's V rows
+    out_loc,      # [T] int32 logical token slots
+    host_ptrs,    # [2*L] int64: per dense layer the K/V host bank base addresses
+    layer,
+    row_bytes: tl.constexpr,
+    block: tl.constexpr,
+):
+    t = tl.program_id(0)
+    slot = tl.load(out_loc + t).to(tl.int64)
+    offs = tl.arange(0, block)
+    mask = offs < row_bytes
+    kb = tl.load(host_ptrs + 2 * layer).to(tl.pointer_type(tl.uint8))
+    vb = tl.load(host_ptrs + 2 * layer + 1).to(tl.pointer_type(tl.uint8))
+    tl.store(kb + slot * row_bytes + offs, tl.load(k_bytes + t * row_bytes + offs, mask=mask), mask=mask)
+    tl.store(vb + slot * row_bytes + offs, tl.load(v_bytes + t * row_bytes + offs, mask=mask), mask=mask)
+
+
+def mirror_store(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out_loc: torch.Tensor,
+    host_ptrs: torch.Tensor,
+    dense_layer: int,
+) -> None:
+    """Write-through: this layer's K/V rows of every token in ``out_loc`` into the host mirror."""
+    k = k.reshape(k.shape[0], -1)
+    v = v.reshape(v.shape[0], -1)
+    if not k.is_contiguous():
+        k = k.contiguous()
+    if not v.is_contiguous():
+        v = v.contiguous()
+    assert k.shape == v.shape and k.dtype == v.dtype
+    t = k.shape[0]
+    row_bytes = k.shape[1] * k.element_size()
+    if t:
+        _mirror_store_kernel[(t,)](
+            k.view(torch.uint8), v.view(torch.uint8), out_loc, host_ptrs, dense_layer,
+            row_bytes, triton.next_power_of_2(row_bytes),
+        )
+
+
+
+
+@triton.jit
+def _write_pages_stride_kernel(out_loc, out, page_size: tl.constexpr):
+    """One page per page_size tokens: covers every page of a run but (possibly) its last."""
+    i = tl.program_id(0)
+    tl.store(out + i, tl.load(out_loc + i * page_size) // page_size)
+
+
+@triton.jit
+def _write_pages_pad_kernel(out, nstride):
+    """Pad the boundary slots with out[0]: a real logical page, so a run with fewer
+    boundaries than `extra` leaves duplicates behind instead of garbage. Duplicates
+    collapse inside lru_ensure."""
+    i = tl.program_id(0)
+    tl.store(out + nstride + i, tl.load(out + 0))
+
+
+@triton.jit
+def _write_pages_runs_kernel(out_loc, token_to_req, out, counter, n, nstride, extra):
+    """The last token of every request run (a prefill batch concatenates runs): its page is
+    the one the stride can miss. Bounded by `extra`, so the buffer stays fixed-shape."""
+    i = tl.program_id(0)
+    cur = tl.load(token_to_req + i)
+    nxt = tl.load(token_to_req + i + 1, mask=(i + 1 < n), other=-1)
+    if (i + 1 >= n) or (nxt != cur):
+        slot = tl.atomic_add(counter, 1)
+        if slot < extra:
+            tl.store(out + nstride + slot, tl.load(out_loc + i) // page_size)
+
+
+def write_pages_prefill(out_loc, token_to_req, out, counter, nstride, extra, page_size):
+    """Prefill write-page set: stride samples + the last token of each request run.
+
+    Fully device-side (no host sync): a sync here stalls the MoE prefetch pipeline that
+    `--moe-prefill-overlap` depends on, which costs more than the sampling itself.
+    """
+    n = out_loc.numel()
+    if not n:
+        return
+    _write_pages_stride_kernel[(nstride,)](out_loc, out, page_size)
+    if extra:
+        _write_pages_pad_kernel[(extra,)](out, nstride)
+        counter.zero_()
+        _write_pages_runs_kernel[(n,)](out_loc, token_to_req, out, counter, n, nstride, extra)
+
+
+__all__ = ["write_pages", "compact_selected_pages", "translate_table", "translate_slots", "mirror_store"]
diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py
index 1792da2b3c..7eea0fc7b3 100644
--- a/python/freetoken/kvcache/__init__.py
+++ b/python/freetoken/kvcache/__init__.py
@@ -76,10 +76,14 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]:
     return MHAKVCache
 
 
-def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dtype):
+def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dtype,
+                   num_index_pages: int | None = None):
     """Build the engine's KV pool for ``num_pages`` USABLE pages (the dummy page and every
     secondary tier -- window pool, index slab, state rings -- are derived here or inside
-    the pool). Single factory entry for all pool families, DSV4 included."""
+    the pool). Single factory entry for all pool families, DSV4 included.
+
+    ``num_index_pages`` (KV host offload only) sizes the QSA compressed-index slab in the
+    LOGICAL page space while the K/V slabs stay at ``num_pages`` physical pages."""
     from .dsv4_cost_model import _dsv4_pool_sizes
     from .hybrid_swa_pool import _naive_swa_num_tokens, _swa_paged_num_tokens
     from .dsv4_paged_pool import DSV4PagedKVCache
@@ -117,6 +121,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt
         device=device,
         dtype=dtype,
         num_req_slots=config.max_running_req + 1,  # + 1 for the dummy request row
+        num_index_pages=None if num_index_pages is None else num_index_pages + 1,
     )
 
 
@@ -128,6 +133,7 @@ def create_kvcache_pool(
     device: torch.device,
     num_swa_tokens: int | None = None,
     num_req_slots: int | None = None,
+    num_index_pages: int | None = None,
 ) -> BaseKVCachePool:
     if model_config.has_swa_attention:
         from .hybrid_swa_pool import HybridSWAKVCache
@@ -205,6 +211,7 @@ def create_kvcache_pool(
             num_req_slots=num_req_slots,
             layer_ids=spec.layer_ids,
             mrope=model_config.model_is_mrope,
+            num_index_pages=num_index_pages,
         )
 
     if len(kv_specs) == 1 and kv_specs[0].mla:
diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py
new file mode 100644
index 0000000000..40e5940cad
--- /dev/null
+++ b/python/freetoken/kvcache/kv_host_offload.py
@@ -0,0 +1,368 @@
+"""KV host offload for the QSA paged pool (Qwen3.8-Flash-Next, ``--kv-host-pages``).
+
+The engine keeps TWO page counts: ``num_pages`` logical pages (what the scheduler, the page
+table and the radix tree see) and a smaller GPU-resident physical pool. This module is the
+bridge:
+
+* the pinned host **mirror** holds one K/V page per logical page (the backing store);
+* the GPU K/V buffer is an **LRU cache** of physical slots over the mirror, driven by
+  flashlib's ``lru_ensure`` (device-side, fixed-shape, CUDA-graph safe -- the same machinery
+  the MoE expert offload uses);
+* every ``store_kv`` is **write-through**: a Triton kernel mirrors the token's K/V rows into
+  the host bank over UVA, so a physical slot is always clean and eviction is a pure
+  rebinding (no write-back, no dirty tracking);
+* reads translate: the QSA backend selects pages in logical space (the compressed index slab
+  is per logical page and always resident), ensures the selected pages are resident
+  (``lru_ensure`` + one fused H2D multi-bank copy), rewrites the block table to physical
+  slots, and calls the attend kernel unmodified.
+
+The dummy page is an ordinary cacheable id (``num_logical_pages``); its mirror page stays
+zeroed. Stale bindings of freed/reallocated pages are harmless: a new tenant's pages are
+fully written (or position-masked) before they are read, and the mirror catches up through
+the same write-through.
+"""
+
+from __future__ import annotations
+
+import torch
+from flashlib.kernels.slot_cache import lru_ensure
+
+from freetoken.moe.host_banks import HostBank
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class KVHostOffloader:
+    """GPU-slot LRU cache over the pinned host mirror of the logical KV page space."""
+
+    def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None:
+        buf = pool._kv_buffer  # [2, L, P_gpu + 1, page_size, kv_heads, head_dim]
+        assert buf.shape[0] == 2 and buf.is_contiguous()
+        _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape
+        self.pool = pool
+        self.page_size = page_size
+        self.num_layers = num_layers
+        self.num_slots = num_slots  # buffer pages, dummy GPU page included
+        self.num_logical = num_logical_pages  # scheduler-visible pages; dummy id == this
+        self.dummy_page = num_logical_pages
+        self.device = device
+        # The selection touches whole pages: select_width tokens can span at most
+        # select_width // page_size + 2 distinct pages (boundary straddle + open group).
+        self.max_sel_pages = 0  # set by set_select_width (backend init)
+        # Static decode buffers, created by init_graph_buffers (capture) or lazily.
+        self._g: dict[str, torch.Tensor] = {}
+
+        self.mirror = HostBank(
+            (2, num_layers, num_logical_pages + 1, page_size, kv_heads, head_dim), buf.dtype
+        )
+        self.mirror.tensor.zero_()  # fresh/dummy pages read back as zeros, never stale bytes
+        self.mirror.pin()
+
+        # LRU slot maps (lru_ensure's id space is logical pages, slot space is buffer pages).
+        self.phys_of = torch.full((num_logical_pages + 1,), -1, dtype=torch.int32, device=device)
+        self.logical_of = torch.full((num_slots,), -1, dtype=torch.int32, device=device)
+        self.usage = torch.zeros((num_slots,), dtype=torch.int64, device=device)
+        self.step = torch.zeros((), dtype=torch.int64, device=device)
+        self.src_ids = torch.empty((num_slots,), dtype=torch.int32, device=device)
+        self.dst_slots = torch.empty((num_slots,), dtype=torch.int32, device=device)
+        self.num_copy = torch.zeros((1,), dtype=torch.int64, device=device)
+        self._out_scratch = torch.empty((num_slots,), dtype=torch.int32, device=device)
+        # Eager-prefill scratch (never captured): per-layer compacted selections (+ their
+        # pinned host copy for the CPU group packer), the per-group tight query (pinned
+        # staging + device buffer) and the translated block table. Lazy, grow-only.
+        self._sel_all: torch.Tensor | None = None
+        self._counts_dev: torch.Tensor | None = None
+        self._sel_host = None
+        self._counts_host = None
+        self._query_host = None
+        self._query_buf: torch.Tensor | None = None
+        self._eff_buf: torch.Tensor | None = None
+        self._trunc_eager = torch.zeros(1, dtype=torch.int32, device=device)
+        # Prefill write-page scratch (grow-only) + counter for the run-aware sampler.
+        self.max_bs = 1
+        self._wp_buf: torch.Tensor | None = None
+        self._wp_count = torch.zeros((), dtype=torch.int32, device=device)
+
+        from freetoken.kernel.pinned import device_ptr
+
+        # Fused-copy descriptor: 2 * num_layers banks (K/V of each layer), one page row each.
+        row_bytes = page_size * kv_heads * head_dim * buf.dtype.itemsize
+        assert row_bytes % 16 == 0
+        mirror = self.mirror.tensor
+        dst_ptrs, src_ptrs, host_ptrs = [], [], []
+        for kv in range(2):
+            for layer in range(num_layers):
+                assert buf[kv, layer].is_contiguous() and mirror[kv, layer].is_contiguous()
+                dst_ptrs.append(buf[kv, layer].data_ptr())
+                src_ptrs.append(device_ptr(mirror[kv, layer]))
+        for layer in range(num_layers):  # K/V interleaved, the mirror-store kernel's layout
+            host_ptrs.append(device_ptr(mirror[0, layer]))
+            host_ptrs.append(device_ptr(mirror[1, layer]))
+        self._dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=device)
+        self._src_ptrs = torch.tensor(src_ptrs, dtype=torch.int64, device=device)
+        self._feats = torch.full((2 * num_layers,), row_bytes, dtype=torch.int64, device=device)
+        self._host_ptrs = torch.tensor(host_ptrs, dtype=torch.int64, device=device)
+        self.mirror_bytes = self.mirror.nbytes
+
+    # ----- geometry ---------------------------------------------------------
+    def set_select_width(self, select_width: int, block_topk: int = 0,
+                         max_bs: int = 1) -> None:
+        # The sparse selection is block_topk short runs of consecutive tokens SCATTERED
+        # across the sequence -- each run can touch its own page (two if it straddles a
+        # boundary). Contiguity (select_width // page_size) badly underestimates this.
+        # Hard cap: the decode query is `max_bs * (1 + max_sel_pages)` wide and must fit
+        # the physical pool (lru_ensure requires |distinct| <= num_cached), so the
+        # per-row width is the pool split across the concurrent requests, not the whole
+        # pool. Concurrency and per-row selection width trade off here by construction;
+        # exceeding the resulting width now fails loudly (see `trunc_count`).
+        worst = 2 * (block_topk or (select_width // 4)) + 8
+        per_row = max(16, (self.num_slots - 73) // max(1, max_bs))
+        self.max_sel_pages = min(worst, per_row)
+
+    def prefill_fetch_rows(self, n_write_pages: int) -> int:
+        """Query rows per prefill attend sub-chunk: the write pages plus this many rows' worth
+        of selected pages must always fit the physical pool (lru_ensure's distinct bound)."""
+        margin = 8
+        budget = self.num_slots - margin - n_write_pages
+        if budget < self.max_sel_pages:
+            raise RuntimeError(
+                f"KV host offload: {self.num_slots} physical pages cannot hold {n_write_pages} "
+                f"write pages + one row's selection ({self.max_sel_pages}); the GPU pool is too small"
+            )
+        return max(1, budget // self.max_sel_pages)
+
+    # ----- graph buffers ------------------------------------------------------
+    def init_graph_buffers(self, max_bs: int, width_pages: int, select_width: int | None = None) -> None:
+        """Static decode buffers (capture-safe addresses); sliced to the live bs per step."""
+        self.max_bs = max(1, max_bs)
+        if select_width is not None:
+            self.set_select_width(select_width, max_bs=max_bs)
+        assert self.max_sel_pages > 0, "set_select_width must run before init_graph_buffers"
+        d = self.device
+        need = max(self.num_slots, max_bs * (1 + self.max_sel_pages))
+        if self._out_scratch.numel() < need:
+            self._out_scratch = torch.empty((need,), dtype=torch.int32, device=d)
+        self._g = {
+            "write_pages": torch.zeros(max_bs, dtype=torch.int32, device=d),
+            "out_loc_gpu": torch.zeros(max_bs, dtype=torch.int32, device=d),
+            "sel_pages": torch.zeros((max_bs, self.max_sel_pages), dtype=torch.int32, device=d),
+            "query": torch.zeros(max_bs * (1 + self.max_sel_pages), dtype=torch.int32, device=d),
+            "eff_table": torch.zeros((max_bs, width_pages), dtype=torch.int32, device=d),
+            "trunc": torch.zeros(1, dtype=torch.int32, device=d),
+        }
+
+    def trunc_count(self) -> int:
+        """Dropped selection pages accumulated since the last read (0 = no truncation)."""
+        n = int(self._trunc_eager.item())
+        if n:
+            self._trunc_eager.zero_()
+        t = self._g.get("trunc")
+        if t is not None:
+            m = int(t.item())
+            if m:
+                t.zero_()
+            n += m
+        return n
+
+    def _graph_buffers(self, bs: int, width_pages: int) -> None:
+        """Lazy (eager-only) static-buffer init; capture paths get them from init_graph_buffers."""
+        g = self._g
+        if not g or g["write_pages"].shape[0] < bs or g["eff_table"].shape[1] != width_pages:
+            self.init_graph_buffers(max(bs, g["write_pages"].shape[0] if g else 0), width_pages)
+
+    # ----- residency engine ---------------------------------------------------
+    # lru_ensure's phase-1 register tile is O(K^2): the kernel is tuned for K<=512
+    # (MoE's whole-layer materialize). Chunk bigger queries into calls inside that envelope.
+    _ENSURE_CHUNK = 512
+
+    def _ensure(self, query: torch.Tensor) -> None:
+        """Bind every logical page in ``query`` to a physical slot and fill the misses."""
+        from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
+
+        n = query.numel()
+        # The chunks are disjoint slices, so distinct(query) <= n; keeping n <= num_slots
+        # makes the kernel's |distinct| <= num_cached requirement hold by construction.
+        assert n <= self.num_slots, (
+            f"ensure query of {n} pages exceeds the {self.num_slots} physical slots"
+        )
+        if n > self._out_scratch.numel():
+            # eager-only path (prefill); capture-time queries always fit the pre-sized buffer
+            self._out_scratch = torch.empty((n,), dtype=torch.int32, device=self.device)
+        for off in range(0, n, self._ENSURE_CHUNK):
+            q = query[off : off + self._ENSURE_CHUNK]
+            k = q.numel()
+            lru_ensure(
+                q,
+                self.phys_of,
+                self.logical_of,
+                self.usage,
+                self.step,
+                self._out_scratch[:k],
+                self.src_ids,
+                self.dst_slots,
+                self.num_copy,
+            )
+            fast_index_copy_multi_jit(
+                self._dst_ptrs, self._src_ptrs, self._feats,
+                self.dst_slots, self.src_ids, self.num_copy,
+            )
+
+    # ----- write path ---------------------------------------------------------
+    def ensure_write_pages(self, out_loc: torch.Tensor, is_decode: bool,
+                           token_to_req: torch.Tensor | None = None) -> torch.Tensor:
+        """Make every page ``out_loc`` writes to resident; returns the write-page query
+        (the fetch path re-includes it so a write page can never be evicted mid-forward)."""
+        from freetoken.kernel.triton.qsa.offload import write_pages, write_pages_prefill
+
+        if is_decode:
+            n = out_loc.numel()
+            wp = self._g["write_pages"][:n]
+            write_pages(out_loc, wp, self.page_size)
+            self._ensure(wp)
+            return wp
+        # A prefill batch concatenates several requests, and a plain every-page_size-th
+        # stride lands at an arbitrary offset inside each later request: its LAST page is
+        # skipped whenever `len % page_size <= offset`, and store_kv would then write to a
+        # physical slot that was never ensured. The stride provably covers every page but
+        # the last of each run, so add the final token of each run (one page per
+        # concurrent request). All device-side: no host sync.
+        n = out_loc.numel()
+        if not n:
+            return out_loc
+        nstride = (n + self.page_size - 1) // self.page_size
+        extra = max(1, self.max_bs)
+        if self._wp_buf is None or self._wp_buf.numel() < nstride + extra:
+            self._wp_buf = torch.empty(max(nstride + extra, 1024), dtype=torch.int32,
+                                       device=self.device)
+        pages = self._wp_buf[: nstride + extra]
+        write_pages_prefill(out_loc, token_to_req, pages, self._wp_count,
+                            nstride, extra, self.page_size)
+        self._ensure(pages)
+        return pages
+
+    def translate_slots(self, out_loc: torch.Tensor, is_decode: bool) -> torch.Tensor:
+        """Logical token slots -> physical (post-ensure; the write pages are pinned-hot)."""
+        from freetoken.kernel.triton.qsa.offload import translate_slots
+
+        if is_decode:
+            n = out_loc.numel()
+            out = self._g["out_loc_gpu"][:n]
+        else:
+            out = torch.empty_like(out_loc)
+        translate_slots(out_loc, self.phys_of, out, self.page_size)
+        return out
+
+    def mirror_store(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int) -> None:
+        from freetoken.kernel.triton.qsa.offload import mirror_store
+
+        mirror_store(k, v, out_loc, self._host_ptrs, self.pool._dense(layer_id))
+
+    # ----- read path ------------------------------------------------------------
+    def compact_all(self, indices: torch.Tensor, token_to_req: torch.Tensor,
+                    block_table: torch.Tensor) -> torch.Tensor:
+        """One launch for the whole layer: per query row, the distinct logical pages its
+        selection touches ([rows, max_sel_pages], first-page padded) plus per-row distinct
+        counts. No host sync."""
+        from freetoken.kernel.triton.qsa.offload import compact_selected_pages
+
+        rows = indices.shape[0]
+        buf = self._sel_all
+        if buf is None or buf.shape[0] < rows:
+            n = max(rows, 4096)
+            buf = self._sel_all = torch.empty((n, self.max_sel_pages), dtype=torch.int32, device=self.device)
+            self._counts_dev = torch.empty((n,), dtype=torch.int32, device=self.device)
+        sel = buf[:rows]
+        compact_selected_pages(
+            indices, token_to_req, block_table, sel, self.page_size, self.dummy_page,
+            self._trunc_eager, counts=self._counts_dev[:rows],
+        )
+        return sel
+
+    def iter_prefill_groups(self, sel_all: torch.Tensor, write_pages: torch.Tensor,
+                            block_table: torch.Tensor):
+        """Pack the layer's query rows into residency-feasible groups, then per group upload
+        the tight query (write set + the group's REAL union), ensure, translate; yields
+        ``((row_start, row_end), eff_table)`` and the caller attends between iterations.
+
+        One D2H per layer (the compacted page lists); the CPU packer then works with exact
+        sets, so the per-group distinct bound (|write| + |union| + slack <= num_slots) is
+        VERIFIED, not estimated -- lru_ensure's capacity requirement holds by construction.
+        Consecutive rows select mostly the same pages, so groups pack several rows each.
+        Buffers are reused across groups; single-stream ordering (attend of group i is
+        enqueued before group i+1's H2D/ensure/translate) makes that safe.
+        """
+        import numpy as np
+
+        rows = sel_all.shape[0]
+        host = self._sel_host
+        if host is None or host.shape[0] < rows:
+            n = max(rows, 4096)
+            host = self._sel_host = torch.empty((n, self.max_sel_pages), dtype=torch.int32, pin_memory=True)
+            self._counts_host = torch.empty((n,), dtype=torch.int32, pin_memory=True)
+        sel_h, cnt_h = host[:rows], self._counts_host[:rows]
+        sel_h.copy_(sel_all, non_blocking=True)
+        cnt_h.copy_(self._counts_dev[:rows], non_blocking=True)
+        torch.cuda.synchronize(self.device)  # the one sync per layer
+        sel_np, cnt_np = sel_h.numpy(), cnt_h.numpy()
+        write_np = write_pages.cpu().numpy().astype(np.int64)
+        # |write| + |union of the group's pages| + slack <= physical slots
+        cap = self.num_slots - write_np.size - 9
+        n_ids = self.num_logical + 1
+
+        gmask = np.zeros(n_ids, dtype=bool)
+        gmask[write_np] = True
+        gstart, gcost = 0, 0
+        groups: list[tuple[int, int, np.ndarray]] = []
+        for r in range(rows):
+            pages = sel_np[r, : cnt_np[r]].astype(np.int64)
+            extra = int((~gmask[pages]).sum())
+            if gcost + extra > cap and r > gstart:
+                groups.append((gstart, r, np.flatnonzero(gmask)))
+                gmask[:] = False
+                gmask[write_np] = True
+                gstart, gcost = r, 0
+                extra = int((~gmask[pages]).sum())
+            gmask[pages] = True
+            gcost += extra
+        groups.append((gstart, rows, np.flatnonzero(gmask)))
+
+        from freetoken.kernel.triton.qsa.offload import translate_table
+
+        for gstart, gend, pages_np in groups:
+            n = pages_np.size
+            qh = self._query_host
+            if qh is None or qh.numel() < n:
+                qh = self._query_host = torch.empty((max(n, 2048),), dtype=torch.int32, pin_memory=True)
+                self._query_buf = torch.empty_like(qh, device=self.device)
+            qh[:n] = torch.from_numpy(pages_np.astype(np.int32))
+            qd = self._query_buf[:n]
+            qd.copy_(qh[:n], non_blocking=True)
+            self._ensure(qd)
+            eff = self._eff_buf
+            if eff is None or eff.shape != block_table.shape:
+                eff = self._eff_buf = torch.empty_like(block_table)
+            translate_table(block_table, self.phys_of, eff)
+            yield (gstart, gend), eff
+
+    def fetch_for_attend(self, indices: torch.Tensor, md) -> torch.Tensor:
+        """Decode (graph-safe): ensure the selected pages resident, return the PHYSICAL
+        block table for the attend kernel. Fixed shapes throughout."""
+        from freetoken.kernel.triton.qsa.offload import compact_selected_pages, translate_table
+
+        rows = indices.shape[0]
+        self._graph_buffers(rows, md.block_table.shape[1])
+        g = self._g
+        sel = g["sel_pages"][:rows]
+        compact_selected_pages(indices, md.token_to_req, md.block_table, sel, self.page_size, self.dummy_page, g["trunc"])
+        per_row = 1 + self.max_sel_pages
+        query = g["query"][: rows * per_row].view(rows, per_row)
+        query[:, 0].copy_(md.write_pages)
+        query[:, 1:].copy_(sel)
+        self._ensure(query.view(-1))
+        eff = g["eff_table"][:rows]
+        translate_table(md.block_table, self.phys_of, eff)
+        return eff
+
+
+__all__ = ["KVHostOffloader"]
diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py
index 7de70977d5..31ec60beb1 100644
--- a/python/freetoken/kvcache/qsa_pool.py
+++ b/python/freetoken/kvcache/qsa_pool.py
@@ -66,6 +66,7 @@ def __init__(
         ring_capacity: int | None = None,
         layer_ids: Sequence[int] | None = None,
         mrope: bool = False,
+        num_index_pages: int | None = None,
     ) -> None:
         if index_ratio < 1 or page_size % index_ratio != 0:
             # slot // index_ratio only names one group when a group never straddles a page.
@@ -93,6 +94,10 @@ def __init__(
         self._index_dtype = dtype
         self._page_size = page_size
         self._mrope = mrope
+        # KV host offload: the K/V slabs stay at num_pages PHYSICAL pages while the
+        # compressed index slab covers the (larger) logical page space, so selection keeps
+        # working for host-resident pages. None -> the slab follows the K/V pages.
+        self._num_index_pages = num_pages if num_index_pages is None else num_index_pages
         super().__init__(
             num_kv_heads=num_kv_heads,
             num_layers=num_layers,
@@ -104,7 +109,7 @@ def __init__(
             layer_ids=layer_ids,
         )
         self._zero_kv_slabs()
-        self._alloc_index_tiers(num_pages)
+        self._alloc_index_tiers(self._num_index_pages)
 
     def _zero_kv_slabs(self) -> None:
         # Defense-in-depth: the attend kernels pos-mask every K/V load (the real fix for
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..d8908210cc 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -677,6 +677,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="KV-cache token floor reserved before --moe-cache-auto fills experts.",
     )
 
+    parser.add_argument(
+        "--kv-host-pages",
+        type=int,
+        default=ServerArgs.kv_host_pages,
+        help=(
+            "Extra KV pages mirrored to pinned host RAM (QSA models): the GPU pool becomes "
+            "an LRU cache over the logical page space, extending context past VRAM capacity. "
+            "0 = off."
+        ),
+    )
+
     parser.add_argument(
         "--moe-cache-policy",
         default=ServerArgs.moe_cache_policy,

From 0f26fd7482f6c18591df49c4f7a4c40759998d4c Mon Sep 17 00:00:00 2001
From: AlReis <258541404+alvarorsouza-arch@users.noreply.github.com>
Date: Thu, 17 Sep 2026 06:03:03 +0000
Subject: [PATCH 041/249] perf(moe): single-launch prefill buffer invalidation
 without hidden sync

---
 .../freetoken/kernel/triton/moe/__init__.py   |  5 ++
 .../freetoken/kernel/triton/moe/invalidate.py | 60 +++++++++++++++++++
 python/freetoken/moe/offload_cache.py         | 17 +++---
 3 files changed, 75 insertions(+), 7 deletions(-)
 create mode 100644 python/freetoken/kernel/triton/moe/__init__.py
 create mode 100644 python/freetoken/kernel/triton/moe/invalidate.py

diff --git a/python/freetoken/kernel/triton/moe/__init__.py b/python/freetoken/kernel/triton/moe/__init__.py
new file mode 100644
index 0000000000..3f620fec6e
--- /dev/null
+++ b/python/freetoken/kernel/triton/moe/__init__.py
@@ -0,0 +1,5 @@
+"""Triton kernels for the MoE offload cache (``moe/offload_cache.py``)."""
+
+from .invalidate import invalidate_prefill_slots
+
+__all__ = ["invalidate_prefill_slots"]
diff --git a/python/freetoken/kernel/triton/moe/invalidate.py b/python/freetoken/kernel/triton/moe/invalidate.py
new file mode 100644
index 0000000000..4fa01db394
--- /dev/null
+++ b/python/freetoken/kernel/triton/moe/invalidate.py
@@ -0,0 +1,60 @@
+"""Single-launch invalidation of the MoE prefill overlap buffers.
+
+``OffloadMoeCache._invalidate_prefill_buffer`` runs twice per prefill chunk (once per
+double-buffer slot) and used to clear the buffer's slot map with a boolean-mask index:
+
+    slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
+
+A boolean index produces a data-dependent shape, so every call hid a device-to-host
+synchronization. With two buffer reuses per chunk x 48 layers, the host waited -- per
+layer -- for ALL enqueued GPU work (the cached-context attention included) before
+prefetching the next layer, which serialized the prefill-overlap pipeline. On long
+cached contexts that turned a ~3 s chunk into a 20-150 s turn (py-spy: 92% of the
+scheduler inside this invalidation + ``synchronize``, GPU 0-3% utilized).
+
+This kernel is fixed-shape: one launch, no synchronization, identical result. Zeroing
+``usage`` here (instead of a separate ``zero_``) keeps the buffer's slots as the
+argmin(usage) eviction victims, exactly like the code it replaces.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _invalidate_slots_kernel(
+    id_of_slot_ptr,   # (num_slots,) int32 -- slot -> expert id, -1 if empty
+    slot_for_id_ptr,  # (num_layers * num_experts) int32 -- expert id -> slot, -1 if not resident
+    usage_ptr,        # (num_slots,) int64 -- last-used step; zeroed so these slots are
+                      # the argmin(usage) victims in ensure_experts' eviction
+    slot_start,
+    n,
+    BLOCK: tl.constexpr,
+):
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    m = i < n
+    e = tl.load(id_of_slot_ptr + slot_start + i, mask=m, other=-1)
+    # clear the expert's residency flag; empty slots (-1) are masked out
+    tl.store(slot_for_id_ptr + e, -1, mask=m & (e >= 0))
+    tl.store(id_of_slot_ptr + slot_start + i, -1, mask=m)
+    tl.store(usage_ptr + slot_start + i, 0, mask=m)
+
+
+def invalidate_prefill_slots(
+    id_of_slot: torch.Tensor,
+    slot_for_id: torch.Tensor,
+    usage: torch.Tensor,
+    slot_start: int,
+    num_experts: int,
+) -> None:
+    """Free the ``num_experts`` slots starting at ``slot_start`` (fixed shape, no sync).
+
+    ``slot_for_id`` may be the 2-D ``(num_layers, num_experts)`` map: ids are flat
+    (``layer * num_experts + expert``) and the tensor is contiguous, so the kernel's
+    flat addressing matches the ``view(-1)`` the eager path used.
+    """
+    grid = ((num_experts + 255) // 256,)
+    _invalidate_slots_kernel[grid](
+        id_of_slot, slot_for_id, usage, slot_start, num_experts, BLOCK=256
+    )
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 7abb9c8003..6e0c56a916 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -634,13 +634,16 @@ def _init_prefill_overlap_buffers(self) -> None:
 
     def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
         slot_start = buffer_id * self.num_experts
-        slot_end = slot_start + self.num_experts
-        old_ids = self.id_of_slot[slot_start:slot_end]
-        self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
-        old_ids.fill_(-1)
-        # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in
-        # ensure_experts evicts them first.
-        self.usage[slot_start:slot_end].zero_()
+        # One fixed-shape launch. The previous boolean-mask index
+        # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every
+        # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48
+        # layers the host waited, per layer, for ALL enqueued GPU work (the cached-context
+        # attention included), serializing the prefill-overlap pipeline.
+        from freetoken.kernel.triton.moe import invalidate_prefill_slots
+
+        invalidate_prefill_slots(
+            self.id_of_slot, self.slot_for_id, self.usage, slot_start, self.num_experts
+        )
 
     def begin_prefill(self) -> None:
         if not self.prefill_overlap:

From a76d9ecc1f5043aa3c8eebd5880eb0334ddbe80b Mon Sep 17 00:00:00 2001
From: yajun <leamon@stu.xidian.edu.cn>
Date: Thu, 17 Sep 2026 16:42:16 +0800
Subject: [PATCH 042/249] feat(server): serve per-request inference metrics
 under --enable-metrics-report

- the scheduler measures the prefill span itself (admission to the token sampled off the request's last prefill chunk) and ships it as a duration on that token's DetokenizeMsg: it is the one boundary a client timing SSE frames cannot see, and a duration keeps the wire free of any cross-process clock assumption
- the remaining timings come off the raw ack stream, not the emitted GenEvents: a reasoning or tool-call parser holds text back, sometimes to the end of the stream, so an event-timed TTFT runs long, and acks are the one path the streaming and buffered generators share
- cached_prompt_tokens reports the real prefix-cache hit whether or not --enable-cache-report is set; that flag governs the billing fields in usage, and gating the hit here would leave prefill_tokens_per_second computed over tokens that were never forwarded
- decode_tokens_per_second divides by decode_tokens - 1, as vLLM's and sglang's serving benchmarks do: the first token falls out of prefill and spans no decode interval
- served on /v1/chat/completions, /v1/messages and /v1/responses; /v1/completions does not go through the shared generation core and is left out

Closes #503
---
 docs/cli.md                                  |  56 ++++++++
 python/freetoken/message/frontend.py         |   3 +
 python/freetoken/message/tokenizer.py        |   5 +
 python/freetoken/scheduler/scheduler.py      |  22 ++++
 python/freetoken/server/anthropic_api.py     |  26 +++-
 python/freetoken/server/anthropic_models.py  |   6 +
 python/freetoken/server/args.py              |  20 +++
 python/freetoken/server/generation.py        | 128 +++++++++++++++++--
 python/freetoken/server/openai_api.py        |  45 +++++--
 python/freetoken/server/responses_api.py     |  37 +++++-
 python/freetoken/tokenizer/server.py         |   1 +
 tests/scheduler/test_cost_accounting_core.py |  17 +++
 tests/server/test_anthropic_api.py           |  56 +++++++-
 tests/server/test_generation_accounting.py   | 113 +++++++++++++++-
 tests/server/test_openai_api.py              |  62 +++++++++
 tests/server/test_responses_api.py           |  41 +++++-
 16 files changed, 599 insertions(+), 39 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 0c33943987..ab4429e0d3 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -98,6 +98,62 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--tool-call-parser` | auto | Tool-call format; auto-inferred from the model family |
 | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables |
 | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block |
+| `--enable-metrics-report` | off | Serve a per-request `metrics` object next to usage ([below](#per-request-performance-metrics)) |
+
+### Per-request performance metrics
+
+`--enable-metrics-report` adds a `metrics` object to each response on `/v1/chat/completions`,
+`/v1/messages` and `/v1/responses`. It is not part of any of those protocols, which is why it is
+off by default; clients that do not read it are unaffected.
+
+```json
+{
+  "usage": { "prompt_tokens": 8192, "completion_tokens": 256, "total_tokens": 8448 },
+  "metrics": {
+    "ttft_ms": 320.5,
+    "prefill_tokens": 2048,
+    "cached_prompt_tokens": 6144,
+    "prefill_time_ms": 410.0,
+    "prefill_tokens_per_second": 4995.12,
+    "decode_tokens": 256,
+    "decode_time_ms": 5120.8,
+    "decode_tokens_per_second": 49.8,
+    "total_time_ms": 5480.2
+  }
+}
+```
+
+| Field | Meaning |
+| --- | --- |
+| `ttft_ms` | Time to the first generated token, from the point the request entered the generation path |
+| `prefill_tokens` | Prompt tokens actually forwarded: `prompt_tokens - cached_prompt_tokens` |
+| `cached_prompt_tokens` | Prompt tokens served from the prefix cache instead of recomputed |
+| `prefill_time_ms` | The scheduler's own prefill span: admission to the token sampled off the last prefill chunk |
+| `prefill_tokens_per_second` | `prefill_tokens / prefill_time_ms` |
+| `decode_tokens` | Generated tokens (same as `usage.completion_tokens`) |
+| `decode_time_ms` | First generated token to the last |
+| `decode_tokens_per_second` | `(decode_tokens - 1) / decode_time_ms`; the first token comes out of prefill, so it spans no decode interval |
+| `total_time_ms` | Request latency up to the terminal engine reply, excluding HTTP framing |
+
+Reading the numbers:
+
+- `prefill_time_ms` is measured inside the scheduler, which is the point of the flag: a client
+  timing SSE frames cannot see when prefill starts or ends. The others are measured at the API
+  layer and carry the same IPC hops a client's own timestamps would.
+- `ttft_ms` is larger than `prefill_time_ms` by the time the request spent queued before
+  admission. Under load that gap is the queue, not the model.
+- Every span is **this request's share of shared work**. Its prefill chunk is co-scheduled with
+  other prompts and its decode steps are batched with other requests, so the throughputs
+  describe this request under that load, not isolated engine benchmarks. For those,
+  send one request at a time.
+- `cached_prompt_tokens` here always reports the real prefix-cache hit, whether or not
+  `--enable-cache-report` is set; that flag governs the billing fields in `usage`. Without it
+  `prefill_tokens_per_second` would be computed over tokens that were never forwarded.
+
+Streaming responses carry `metrics` on the same final message as usage, so the request has to ask
+for usage too: `stream_options: {"include_usage": true}` on `/v1/chat/completions`. On
+`/v1/messages` it rides `message_delta`; on `/v1/responses`, `response.completed`.
+`/v1/completions` does not serve `metrics`: it does not go through the shared generation core.
 
 ### Image input
 
diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py
index 24725567b9..392faea10e 100644
--- a/python/freetoken/message/frontend.py
+++ b/python/freetoken/message/frontend.py
@@ -45,6 +45,9 @@ class UserReply(BaseFrontendMsg):
     swa_total_tokens: int = 0
     # Bytes the engine process holds on the GPU (torch reserved pool). 0 when not reported.
     gpu_mem_bytes: int = 0
+    # Scheduler-measured prefill span (see DetokenizeMsg.prefill_ms). Arrives once, on the
+    # reply carrying the request's first generated token; 0.0 on every other reply.
+    prefill_ms: float = 0.0
     # Set (with finished=True) when a request failed before producing output — e.g. a chat
     # template that the tokenizer cannot render, or a prompt that exceeds the KV budget the
     # scheduler can serve. Carries a human-readable reason. Without this, such a request would
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index 9442ee059e..f0336f73da 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -48,6 +48,11 @@ class DetokenizeMsg(BaseTokenizerMsg):
     swa_total_tokens: int = 0
     # Bytes this engine process holds on the GPU (torch reserved pool). 0 on CPU.
     gpu_mem_bytes: int = 0
+    # Scheduler-measured prefill span for this request: admission (the first prefill batch
+    # was prepared) to the token sampled off its last prefill chunk. Set on that first token
+    # only, 0.0 on every later one -- it is the one span a client cannot derive from HTTP
+    # timestamps, which is why it rides the wire instead of being re-estimated frontend-side.
+    prefill_ms: float = 0.0
 
 
 @dataclass
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index ac1bf322e3..88b4558744 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import time
 
 from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias
 
@@ -111,6 +112,11 @@ def __init__(self, config: SchedulerConfig):
         # tombstone so an abort-before-admission request can never be resurrected after its
         # terminal accounting acknowledgement has already been published.
         self._abort_tombstones: dict[int, None] = {}
+        # uid -> monotonic clock at admission, for the per-request prefill span shipped on the
+        # first sampled token. An entry existing IS the "this is the first token" test, so it is
+        # popped there; _free_req_resources pops too, so a request aborted mid-prefill (which
+        # never samples) cannot leave one behind.
+        self._prefill_start: dict[int, float] = {}
         self._forward_iter = 0  # global forward counter; drives the SWA proactive-eviction cadence
         # The launched-but-not-yet-drained batch (overlap): set at the top of each overlap_loop
         # iteration so the abort handler can tell whether a request's forward is still in flight
@@ -370,6 +376,10 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                     and not finished
                 ):
                     req.toolcall_anchor_len = req.input_ids.numel()
+                # Present only until this request's first token is sampled, i.e. off its last
+                # prefill chunk -- so popping it both ends the prefill span and marks this as
+                # the one reply that carries it.
+                started_at = self._prefill_start.pop(req.uid, None)
                 reply.append(
                     DetokenizeMsg(
                         uid=req.uid,
@@ -378,6 +388,10 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         finish_reason=finish_reason,
                         matched_stop=matched_stop,
                         stop_strs=req.sampling_params.stop_strs or None,
+                        prefill_ms=(
+                            0.0 if started_at is None
+                            else (time.monotonic() - started_at) * 1000.0
+                        ),
                     )
                 )
 
@@ -631,6 +645,9 @@ def _free_req_resources(self, req: Req) -> None:
         # slots to two later requests. table_idx == -1 marks an already-freed request.
         if req.table_idx == -1:
             return
+        # A request freed without ever sampling (aborted mid-prefill) would otherwise keep its
+        # admission timestamp forever; a normally-finished one popped it at its first token.
+        self._prefill_start.pop(req.uid, None)
         # Polymorphic free: the DSV4 manager returns the request's window pages + cmp/idx blocks
         # to their tier free-lists; the generic manager frees its KV pages (it reads
         # page_table[req.table_idx], so free the table entry after).
@@ -886,6 +903,11 @@ def _report_prompt_admissions(self, batch: Batch) -> None:
         """
         if not batch.is_prefill or not batch.prompt_admissions:
             return
+        # Before the forward, which is what makes this the prefill START: _schedule_next_batch
+        # has prepared the batch but not run it.
+        now = time.monotonic()
+        for uid, _, _ in batch.prompt_admissions:
+            self._prefill_start[uid] = now
         self.send_result(
             [
                 PromptAdmittedMsg(uid=uid, prompt_tokens=prompt_tokens, cached_tokens=cached_tokens)
diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 0d60940c97..883c285bf3 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -44,9 +44,11 @@
     ToolCallArgsDelta,
     ToolCallsDelta,
     ToolCallStart,
+    build_metrics,
     count_prompt_tokens,
     generate_events,
     generate_full,
+    metrics_enabled,
     render_messages,
     resolve_sampling,
     split_tool_lists,
@@ -126,10 +128,11 @@ async def handle_anthropic_messages(
         return _anthropic_error_response(400, "invalid_request_error", str(exc))
 
     cache_report = getattr(state.config, "enable_cache_report", False)
+    metrics = metrics_enabled(state)
     if req.stream:
         events = anthropic_event_stream(
             generate_events(uid, spec, state, source="/v1/messages"),
-            req.model, uid, cache_report=cache_report,
+            req.model, uid, cache_report=cache_report, metrics=metrics,
         )
         if request is not None:
             events = state.stream_with_cancellation(events, request, uid)
@@ -139,7 +142,9 @@ async def handle_anthropic_messages(
         result = await generate_full(uid, spec, state, source="/v1/messages")
     except GenerationError as exc:
         return _anthropic_error_response(400, "invalid_request_error", str(exc))
-    response = anthropic_full_response(result, req.model, uid, cache_report=cache_report)
+    response = anthropic_full_response(
+        result, req.model, uid, cache_report=cache_report, metrics=metrics
+    )
     return JSONResponse(content=response.model_dump(exclude_none=True))
 
 
@@ -375,7 +380,7 @@ def _tool_result_parts(content) -> tuple[str, list[dict[str, Any]]]:
 # Output formatting: GenResult / GenEvent -> Anthropic response / events
 # --------------------------------------------------------------------------- #
 def anthropic_full_response(
-    result: GenResult, model: str, uid: int, cache_report: bool = False
+    result: GenResult, model: str, uid: int, cache_report: bool = False, metrics: bool = False
 ) -> AnthropicMessagesResponse:
     content: list[AnthropicContentBlock] = []
     if result.reasoning:
@@ -402,6 +407,12 @@ def anthropic_full_response(
         usage=_anthropic_usage(
             result.prompt_tokens, result.completion_tokens, result.cached_tokens, cache_report
         ),
+        metrics=build_metrics(
+            prompt_tokens=result.prompt_tokens,
+            completion_tokens=result.completion_tokens,
+            cached_tokens=result.cached_tokens,
+            timings=result.timings,
+        ) if metrics else None,
     )
 
 
@@ -420,7 +431,8 @@ def _anthropic_usage(
 
 
 async def anthropic_event_stream(
-    events: AsyncIterator[Any], model: str, uid: int, cache_report: bool = False
+    events: AsyncIterator[Any], model: str, uid: int, cache_report: bool = False,
+    metrics: bool = False,
 ) -> AsyncIterator[str]:
     """Format the protocol-neutral GenEvent stream into Anthropic SSE events.
 
@@ -576,6 +588,12 @@ def _tool_args_delta(fragment: str) -> str:
                     usage=_anthropic_usage(
                         ev.prompt_tokens, ev.completion_tokens, ev.cached_tokens, cache_report
                     ),
+                    metrics=build_metrics(
+                        prompt_tokens=ev.prompt_tokens,
+                        completion_tokens=ev.completion_tokens,
+                        cached_tokens=ev.cached_tokens,
+                        timings=ev.timings,
+                    ) if metrics else None,
                 ))
                 yield _event(AnthropicStreamEvent(type="message_stop"))
                 # Anthropic streams terminate on message_stop — no OpenAI-style
diff --git a/python/freetoken/server/anthropic_models.py b/python/freetoken/server/anthropic_models.py
index 42ab545bec..2a6232ee6f 100644
--- a/python/freetoken/server/anthropic_models.py
+++ b/python/freetoken/server/anthropic_models.py
@@ -153,6 +153,10 @@ class AnthropicMessagesResponse(BaseModel):
     ) = None
     stop_sequence: str | None = None
     usage: AnthropicUsage | None = None
+    # FreeToken extension, served only under --enable-metrics-report: per-request inference
+    # timings. Anthropic's wire has no such field, so it stays absent by default rather than
+    # being emitted as null -- both response paths dump with exclude_none.
+    metrics: dict[str, Any] | None = None
 
     def model_post_init(self, __context):
         if not self.id:
@@ -176,3 +180,5 @@ class AnthropicStreamEvent(BaseModel):
     index: int | None = None
     error: AnthropicError | None = None
     usage: AnthropicUsage | None = None
+    # Rides message_delta next to usage, under --enable-metrics-report only.
+    metrics: dict[str, Any] | None = None
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..5b37b6b55b 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -59,6 +59,10 @@ class ServerArgs(SchedulerConfig):
     # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses
     # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report.
     enable_cache_report: bool = False
+    # Serve a per-request `metrics` object (TTFT, prefill/decode times and throughputs,
+    # prefix-cache hit) alongside usage. Off by default: it is a non-standard field on
+    # every protocol we speak.
+    enable_metrics_report: bool = False
     # Comma-separated hostname allowlist for client-supplied image URLs; empty admits any domain.
     allowed_media_domains: str = ""
     # Directory file:// image refs may be read from; empty rejects local files.
@@ -524,6 +528,22 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--enable-metrics-report",
+        action="store_true",
+        default=ServerArgs.enable_metrics_report,
+        help=(
+            "Serve a per-request `metrics` object next to usage on /v1/chat/completions, "
+            "/v1/completions, /v1/messages and /v1/responses: ttft_ms, prefill_time_ms and "
+            "prefill_tokens_per_second (the prefill span measured by the scheduler itself), "
+            "decode_time_ms and decode_tokens_per_second, cached_prompt_tokens and "
+            "total_time_ms. Streaming responses carry it on the same final chunk as usage, so "
+            "the request must also ask for usage (OpenAI stream_options.include_usage). "
+            "Non-standard on every protocol, hence opt-in. Under concurrency the spans are "
+            "this request's share of shared batches, not isolated engine throughput."
+        ),
+    )
+
     parser.add_argument(
         "--sampling-defaults",
         type=str,
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 91056e8d1f..108b666b61 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -107,6 +107,26 @@ class ToolCallsDelta:
     calls: list[ToolCallItem]
 
 
+@dataclass
+class GenTimings:
+    """Per-request inference timings, in milliseconds.
+
+    ``prefill_ms`` is the scheduler's own span (admission to the token sampled off the
+    request's last prefill chunk), shipped over the wire because it is the one boundary a
+    client cannot see; the rest are measured here, around the engine ack loop, so they carry
+    the same IPC hops the client's own HTTP timestamps would.
+
+    Under concurrency every span is this request's SHARE of shared work -- its prefill chunk
+    is co-scheduled with other prompts, its decode steps batched with other requests -- so
+    the derived rates describe this request under that load, not isolated engine throughput.
+    """
+
+    ttft_ms: float = 0.0
+    prefill_ms: float = 0.0
+    decode_ms: float = 0.0
+    total_ms: float = 0.0
+
+
 @dataclass
 class GenDone:
     finish_reason: str
@@ -114,6 +134,7 @@ class GenDone:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    timings: GenTimings = field(default_factory=GenTimings)
 
 
 GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone
@@ -129,6 +150,7 @@ class GenResult:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    timings: GenTimings = field(default_factory=GenTimings)
 
 
 @dataclass
@@ -491,6 +513,83 @@ async def with_keepalive(events: AsyncIterator[GenEvent], interval: float):
             task.cancel()
 
 
+class _Timer:
+    """Times one generation off the raw ack stream.
+
+    Deliberately fed from acks rather than from emitted GenEvents: a reasoning or tool-call
+    parser holds text back, sometimes to the end of the stream, so the first EVENT can lag the
+    first generated token by a lot. Acks are also the one path both the streaming and the
+    buffered generator share, so TTFT means the same thing on either.
+    """
+
+    def __init__(self) -> None:
+        self.start = time.monotonic()
+        self.first_token_at: float | None = None
+        self.end: float | None = None
+        self.prefill_ms = 0.0
+
+    def observe(self, ack: Any) -> None:
+        if ack.prefill_ms:
+            self.prefill_ms = ack.prefill_ms
+        # completion_tokens_delta, not incremental_output: the detokenizer holds back a
+        # trailing partial-stop prefix, so a real generated token can arrive with empty text.
+        if self.first_token_at is None and ack.completion_tokens_delta:
+            self.first_token_at = time.monotonic()
+        if ack.finished:
+            # Stamped here, not in finish(): the callers run the reasoning split and the
+            # tool-call drain between the terminal ack and building their result, and that
+            # post-processing is not inference time.
+            self.end = time.monotonic()
+
+    def finish(self) -> GenTimings:
+        # A stream cut short (client disconnect) never saw a terminal ack; time it to here.
+        end = self.end if self.end is not None else time.monotonic()
+        first = self.first_token_at
+        return GenTimings(
+            ttft_ms=0.0 if first is None else (first - self.start) * 1000.0,
+            prefill_ms=self.prefill_ms,
+            decode_ms=0.0 if first is None else (end - first) * 1000.0,
+            total_ms=(end - self.start) * 1000.0,
+        )
+
+
+def build_metrics(
+    *, prompt_tokens: int, completion_tokens: int, cached_tokens: int, timings: GenTimings
+) -> dict[str, Any]:
+    """The `metrics` object served under --enable-metrics-report, identical on every protocol.
+
+    ``cached_prompt_tokens`` is the real prefix-cache hit whatever --enable-cache-report says:
+    that flag governs BILLING fields in `usage`, and gating the hit here instead would leave
+    prefill_tokens_per_second silently computed over tokens that were never forwarded.
+
+    Decode throughput divides by completion_tokens - 1, not completion_tokens: the first token
+    falls out of prefill, so decode_time_ms spans only the intervals after it (the convention
+    vLLM's and sglang's serving benchmarks use).
+    """
+    prefill_tokens = max(prompt_tokens - cached_tokens, 0)
+    prefill_s = timings.prefill_ms / 1000.0
+    decode_s = timings.decode_ms / 1000.0
+    return {
+        "ttft_ms": round(timings.ttft_ms, 3),
+        "prefill_tokens": prefill_tokens,
+        "cached_prompt_tokens": cached_tokens,
+        "prefill_time_ms": round(timings.prefill_ms, 3),
+        "prefill_tokens_per_second": round(prefill_tokens / prefill_s, 2) if prefill_s > 0 else 0.0,
+        "decode_tokens": completion_tokens,
+        "decode_time_ms": round(timings.decode_ms, 3),
+        "decode_tokens_per_second": (
+            round((completion_tokens - 1) / decode_s, 2)
+            if decode_s > 0 and completion_tokens > 1
+            else 0.0
+        ),
+        "total_time_ms": round(timings.total_ms, 3),
+    }
+
+
+def metrics_enabled(state: Any) -> bool:
+    return bool(getattr(state.config, "enable_metrics_report", False))
+
+
 def _record_generation(
     *,
     source: str | None,
@@ -532,27 +631,24 @@ async def generate_events(
     """Wraps `_generate_events_impl` to log the request with its totals, read off the terminal
     `GenDone`. The `finally` still records the row on a mid-stream disconnect — but with 0 tokens
     if the drop lands before `GenDone`, the only event carrying the totals."""
-    start = time.monotonic()
+    timer = _Timer()
     prompt_tokens = 0
     completion_tokens = 0
-    first_token_at: float | None = None
     error: str | None = None
     try:
-        async for ev in _generate_events_impl(uid, spec, state):
+        async for ev in _generate_events_impl(uid, spec, state, timer):
             if isinstance(ev, GenDone):
                 prompt_tokens = ev.prompt_tokens
                 completion_tokens = ev.completion_tokens
-            elif first_token_at is None:
-                first_token_at = time.monotonic()
             yield ev
     except GenerationError as exc:
         error = str(exc)
         raise
     finally:
         _record_generation(
-            source=source, stream=True, start=start,
+            source=source, stream=True, start=timer.start,
             prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, error=error,
-            first_token_at=first_token_at,
+            first_token_at=timer.first_token_at,
         )
 
 
@@ -561,25 +657,29 @@ async def generate_full(
 ) -> GenResult:
     """Wraps `_generate_full_impl` to log the request with its totals; the `finally` also records
     a `GenerationError` as a failed row."""
-    start = time.monotonic()
+    timer = _Timer()
     result: GenResult | None = None
     error: str | None = None
     try:
-        result = await _generate_full_impl(uid, spec, state)
+        result = await _generate_full_impl(uid, spec, state, timer)
         return result
     except GenerationError as exc:
         error = str(exc)
         raise
     finally:
         _record_generation(
-            source=source, stream=False, start=start,
+            source=source, stream=False, start=timer.start,
             prompt_tokens=result.prompt_tokens if result else 0,
             completion_tokens=result.completion_tokens if result else 0,
             error=error,
+            # No first_token_at: non-streaming rows deliberately carry no TTFT, because
+            # requests_ttft_mean_ms averages only the rows that have one. `metrics` still does.
         )
 
 
-async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIterator[GenEvent]:
+async def _generate_events_impl(
+    uid: int, spec: GenSpec, state: Any, timer: _Timer
+) -> AsyncIterator[GenEvent]:
     """Protocol-neutral streaming generation. Yields semantic events (reasoning /
     content / tool-call deltas) terminated by exactly one GenDone. Produces no wire
     format — the OpenAI/Anthropic/Responses streamers format these into their own.
@@ -686,6 +786,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
     async for ack in state.wait_for_ack(uid):
         if getattr(ack, "error", None):
             raise GenerationError(ack.error, getattr(ack, "error_code", None))
+        timer.observe(ack)
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
@@ -771,10 +872,11 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
         matched_stop=engine_matched_stop, cached_tokens=cached_tokens,
+        timings=timer.finish(),
     )
 
 
-async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
+async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer) -> GenResult:
     """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse
     tool calls, strip special tokens. The adapters format the GenResult into their wire."""
     full_content = ""
@@ -786,6 +888,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
     async for ack in state.wait_for_ack(uid):
         if getattr(ack, "error", None):
             raise GenerationError(ack.error, getattr(ack, "error_code", None))
+        timer.observe(ack)
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
@@ -815,4 +918,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens=completion_tokens,
         matched_stop=engine_matched_stop,
         cached_tokens=cached_tokens,
+        timings=timer.finish(),
     )
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index dc2f73a979..4b08bc83a3 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -28,12 +28,15 @@
     GenDone,
     GenerationError,
     GenSpec,
+    GenTimings,
     ReasoningDelta,
     ToolCallArgsDelta,
     ToolCallsDelta,
     ToolCallStart,
+    build_metrics,
     generate_events,
     generate_full,
+    metrics_enabled,
     prerender_error,
     render_messages,
     resolve_sampling,
@@ -217,7 +220,7 @@ async def handle_chat_completion(
     if result.tool_calls:
         message["tool_calls"] = _tool_calls_to_openai(result.tool_calls)
 
-    return {
+    response: dict[str, Any] = {
         "id": f"chatcmpl-{uid}",
         "object": "chat.completion",
         "created": int(time.time()),
@@ -235,6 +238,14 @@ async def handle_chat_completion(
             _reported_cached(state, result.cached_tokens),
         ),
     }
+    if metrics_enabled(state):
+        response["metrics"] = build_metrics(
+            prompt_tokens=result.prompt_tokens,
+            completion_tokens=result.completion_tokens,
+            cached_tokens=result.cached_tokens,
+            timings=result.timings,
+        )
+    return response
 
 
 async def stream_chat_completion_chunks(
@@ -257,6 +268,7 @@ async def stream_chat_completion_chunks(
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
+    timings = GenTimings()
     tool_calls_sent = 0
     open_tool: dict[str, Any] | None = None
     events = generate_events(uid, spec, state, source="/v1/chat/completions")
@@ -366,21 +378,28 @@ async def stream_chat_completion_chunks(
             prompt_tokens = ev.prompt_tokens
             completion_tokens = ev.completion_tokens
             cached_tokens = ev.cached_tokens
+            timings = ev.timings
             yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}]))
 
     if req.stream_options and req.stream_options.include_usage:
-        yield _sse(
-            {
-                "id": f"chatcmpl-{uid}",
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": req.model,
-                "choices": [],
-                "usage": _usage(
-                    prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)
-                ),
-            }
-        )
+        final: dict[str, Any] = {
+            "id": f"chatcmpl-{uid}",
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": req.model,
+            "choices": [],
+            "usage": _usage(
+                prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)
+            ),
+        }
+        if metrics_enabled(state):
+            final["metrics"] = build_metrics(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cached_tokens=cached_tokens,
+                timings=timings,
+            )
+        yield _sse(final)
 
     yield b"data: [DONE]\n\n"
 
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index 71d6eed873..e333fb70e9 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -69,8 +69,10 @@
     ToolCallArgsDelta,
     ToolCallsDelta,
     ToolCallStart,
+    build_metrics,
     generate_events,
     generate_full,
+    metrics_enabled,
     render_messages,
     resolve_sampling,
     split_tool_lists,
@@ -163,10 +165,11 @@ async def handle_responses(
         return _error_response(400, str(exc))
 
     cache_report = getattr(state.config, "enable_cache_report", False)
+    metrics = metrics_enabled(state)
     if req.stream:
         events = responses_stream_generator(
             generate_events(uid, spec, state, source="/v1/responses"), req, response_id, created,
-            cache_report=cache_report,
+            cache_report=cache_report, metrics=metrics,
         )
         if request is not None:
             events = state.stream_with_cancellation(events, request, uid)
@@ -176,7 +179,9 @@ async def handle_responses(
         result = await generate_full(uid, spec, state, source="/v1/responses")
     except GenerationError as exc:
         return _error_response(400, str(exc), exc.code)
-    response = build_responses_response(result, req, response_id, created, cache_report=cache_report)
+    response = build_responses_response(
+        result, req, response_id, created, cache_report=cache_report, metrics=metrics
+    )
     return JSONResponse(content=response.model_dump(mode="json"))
 
 
@@ -410,6 +415,7 @@ def build_responses_response(
     response_id: str,
     created: int,
     cache_report: bool = False,
+    metrics: bool = False,
 ) -> Response:
     truncated = result.finish_reason == "length"
     item_status = "incomplete" if truncated else "completed"
@@ -454,15 +460,21 @@ def build_responses_response(
             result.cached_tokens if cache_report else 0,
         ),
         incomplete_reason="max_output_tokens" if truncated else None,
+        metrics=build_metrics(
+            prompt_tokens=result.prompt_tokens,
+            completion_tokens=result.completion_tokens,
+            cached_tokens=result.cached_tokens,
+            timings=result.timings,
+        ) if metrics else None,
     )
 
 
 def _response_obj(
     response_id: str, created: int, model: str, output: list[Any],
     *, status: str, usage: ResponseUsage | None, error: ResponseError | None = None,
-    incomplete_reason: str | None = None,
+    incomplete_reason: str | None = None, metrics: dict[str, Any] | None = None,
 ) -> Response:
-    return Response(
+    response = Response(
         id=response_id,
         created_at=created,
         model=model,
@@ -476,6 +488,11 @@ def _response_obj(
         tool_choice="auto",
         tools=[],
     )
+    if metrics is not None:
+        # The SDK model allows extras, so the FreeToken-only `metrics` object rides the real
+        # Response type rather than forcing a parallel dict-shaped response path.
+        response.metrics = metrics
+    return response
 
 
 def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> ResponseUsage:
@@ -499,13 +516,15 @@ async def responses_stream_generator(
     response_id: str,
     created: int,
     cache_report: bool = False,
+    metrics: bool = False,
 ) -> AsyncIterator[str]:
     seq = _Seq()
 
-    def snapshot(status, output, usage=None, incomplete_reason=None):
+    def snapshot(status, output, usage=None, incomplete_reason=None, request_metrics=None):
         return _response_obj(
             response_id, created, req.model, output,
             status=status, usage=usage, incomplete_reason=incomplete_reason,
+            metrics=request_metrics,
         )
 
     yield _sse(ResponseCreatedEvent(
@@ -707,6 +726,12 @@ def args_delta_frame(fragment: str) -> str:
                 finish_reason = ev.finish_reason
                 usage_pt, usage_ct = ev.prompt_tokens, ev.completion_tokens
                 usage_cached = ev.cached_tokens if cache_report else 0
+                request_metrics = build_metrics(
+                    prompt_tokens=ev.prompt_tokens,
+                    completion_tokens=ev.completion_tokens,
+                    cached_tokens=ev.cached_tokens,
+                    timings=ev.timings,
+                ) if metrics else None
                 for f in close_current():
                     yield f
                 if finish_reason == "length":
@@ -716,6 +741,7 @@ def args_delta_frame(fragment: str) -> str:
                             "incomplete", output_items,
                             usage=_usage(usage_pt, usage_ct, usage_cached),
                             incomplete_reason="max_output_tokens",
+                            request_metrics=request_metrics,
                         ),
                     ))
                 else:
@@ -724,6 +750,7 @@ def args_delta_frame(fragment: str) -> str:
                         response=snapshot(
                             "completed", output_items,
                             usage=_usage(usage_pt, usage_ct, usage_cached),
+                            request_metrics=request_metrics,
                         ),
                     ))
     except GenerationError as exc:
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index d3365bfecd..8182c4551a 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -225,6 +225,7 @@ def tokenize_worker(
                         swa_used_tokens=msg.swa_used_tokens,
                         swa_total_tokens=msg.swa_total_tokens,
                         gpu_mem_bytes=msg.gpu_mem_bytes,
+                        prefill_ms=msg.prefill_ms,
                     )
                     for msg, reply in zip(detokenize_msg, replies, strict=True)
                 ]
diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py
index 045010c6b4..37df2b79a1 100644
--- a/tests/scheduler/test_cost_accounting_core.py
+++ b/tests/scheduler/test_cost_accounting_core.py
@@ -94,6 +94,7 @@ def test_schedule_reports_admission_only_after_prepare_succeeds():
     scheduler.prefill_budget = 99
     scheduler.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: batch)
     scheduler.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None)
+    scheduler._prefill_start = {}
     events = []
 
     def prepare(value):
@@ -110,6 +111,8 @@ def send(messages):
     assert events[0] == ("prepared", batch)
     sent = events[1][1]
     assert [(m.uid, m.prompt_tokens, m.cached_tokens) for m in sent] == [(1, 12, 4), (2, 34, 0)]
+    # Admission also opens each request's prefill span, closed at its first sampled token.
+    assert sorted(scheduler._prefill_start) == [1, 2]
 
 
 def test_prepare_failure_emits_no_prompt_admission():
@@ -130,6 +133,20 @@ def fail_prepare(_batch):
     assert sent == []
 
 
+def test_freeing_a_request_that_never_sampled_closes_its_prefill_span():
+    """The first sampled token is what normally pops the span. A request aborted mid-prefill
+    never reaches that point, so the free path has to pop it -- otherwise a long-lived server
+    leaks one entry per aborted request."""
+    scheduler = Scheduler.__new__(Scheduler)
+    scheduler._prefill_start = {7: 100.0, 9: 200.0}
+    scheduler.cache_manager = SimpleNamespace(cache_req=lambda req, finished: None)
+    scheduler.table_manager = SimpleNamespace(free=lambda table_idx: None)
+
+    Scheduler._free_req_resources(scheduler, SimpleNamespace(uid=7, table_idx=3))
+
+    assert scheduler._prefill_start == {9: 200.0}  # only the freed request's span is dropped
+
+
 def test_scheduler_rejection_emits_error_but_no_admission():
     scheduler = Scheduler.__new__(Scheduler)
     scheduler.engine = SimpleNamespace(max_seq_len=4)
diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py
index 327bfcd87d..8f7ef15443 100644
--- a/tests/server/test_anthropic_api.py
+++ b/tests/server/test_anthropic_api.py
@@ -33,6 +33,7 @@
     ContentDelta,
     GenDone,
     GenResult,
+    GenTimings,
     ReasoningDelta,
     ToolCallsDelta,
 )
@@ -43,12 +44,14 @@ async def _aiter(items):
         yield it
 
 
-def _collect_events(events, model="claude-x", uid=1, cache_report=False):
+def _collect_events(events, model="claude-x", uid=1, cache_report=False, metrics=False):
     """Run the Anthropic event stream over neutral GenEvents; return [(type, data), ...]."""
 
     async def run():
         out = []
-        async for frame in A.anthropic_event_stream(_aiter(events), model, uid, cache_report=cache_report):
+        async for frame in A.anthropic_event_stream(
+            _aiter(events), model, uid, cache_report=cache_report, metrics=metrics
+        ):
             etype = None
             data = None
             for line in frame.split("\n"):
@@ -976,3 +979,52 @@ def test_image_only_tool_result_keeps_an_empty_tool_message():
         "role": "user",
         "content": [{"type": "image", "freetoken_ref": {"kind": "b64", "data": "aGk="}}],
     }
+
+
+# --------------------------------------------------------------------------- #
+# Per-request metrics (--enable-metrics-report)
+# --------------------------------------------------------------------------- #
+def _timed_result() -> GenResult:
+    return GenResult(
+        reasoning="", content="hi", tool_calls=[], finish_reason="stop",
+        prompt_tokens=11, completion_tokens=5, cached_tokens=8,
+        timings=GenTimings(ttft_ms=30.0, prefill_ms=12.0, decode_ms=40.0, total_ms=70.0),
+    )
+
+
+def test_full_response_metrics_absent_unless_asked_for():
+    """Anthropic's wire has no `metrics` field, so /v1/messages must stay byte-identical for
+    every client that did not opt in -- absent, not null."""
+    resp = A.anthropic_full_response(_timed_result(), "claude-x", uid=9)
+    assert "metrics" not in resp.model_dump(exclude_none=True)
+
+
+def test_full_response_metrics_report_the_untouched_prompt_split():
+    """cache_report rewrites usage.input_tokens to exclude the cached prefix; metrics must
+    still describe the real prefill, which is prompt_tokens - cached_tokens."""
+    resp = A.anthropic_full_response(
+        _timed_result(), "claude-x", uid=9, cache_report=True, metrics=True
+    )
+    body = resp.model_dump(exclude_none=True)
+    assert body["usage"]["input_tokens"] == 3
+    assert body["metrics"]["prefill_tokens"] == 3
+    assert body["metrics"]["cached_prompt_tokens"] == 8
+    assert body["metrics"]["prefill_time_ms"] == 12.0
+    assert body["metrics"]["ttft_ms"] == 30.0
+
+
+def test_stream_message_delta_carries_metrics():
+    events = [
+        ContentDelta("hi"),
+        GenDone("stop", 4, 2, cached_tokens=1,
+                timings=GenTimings(ttft_ms=5.0, prefill_ms=2.0, decode_ms=9.0, total_ms=14.0)),
+    ]
+    collected = _collect_events(events, metrics=True)
+    md = next(e[1] for e in collected if e[0] == "message_delta")
+    assert md["metrics"]["prefill_time_ms"] == 2.0
+    assert md["metrics"]["prefill_tokens"] == 3
+    # Only the terminal event carries them.
+    assert sum(1 for _, data in collected if isinstance(data, dict) and "metrics" in data) == 1
+
+    off = _collect_events(events)
+    assert all("metrics" not in data for _, data in off if isinstance(data, dict))
diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py
index 8276163b0c..5615460ecd 100644
--- a/tests/server/test_generation_accounting.py
+++ b/tests/server/test_generation_accounting.py
@@ -22,6 +22,8 @@
 from freetoken.server.generation import (  # noqa: E402
     GenDone,
     GenSpec,
+    GenTimings,
+    build_metrics,
     generate_events,
     generate_full,
 )
@@ -63,13 +65,22 @@ async def wait_for_ack(self, uid: int):
             yield reply
 
 
-def _ack(prompt: int = 0, completion: int = 0, out: str = "", finished: bool = False) -> UserReply:
+def _ack(
+    prompt: int = 0,
+    completion: int = 0,
+    out: str = "",
+    finished: bool = False,
+    cached: int = 0,
+    prefill_ms: float = 0.0,
+) -> UserReply:
     return UserReply(
         uid=42,
         incremental_output=out,
         finished=finished,
         prompt_tokens_delta=prompt,
         completion_tokens_delta=completion,
+        cached_tokens=cached,
+        prefill_ms=prefill_ms,
         finish_reason="stop" if finished else None,
     )
 
@@ -193,3 +204,103 @@ def test_ttft_mean_is_zero_without_samples():
     request_ring.reset()
     request_ring.record_request(_row(ttft_ms=None))
     assert request_ring.requests_ttft_mean_ms() == 0
+
+
+# ------------------------------------------------------- per-request metrics
+# The timings ride the same ack stream the token totals do, so they are covered here rather
+# than in each adapter's file; the adapter tests only check that the object reaches the wire.
+def test_prefill_ms_from_the_engine_reaches_the_result_unchanged():
+    """The scheduler measures the prefill span; the generation layer must pass it through
+    rather than re-deriving it from ack arrival times (which batch together)."""
+    st = FakeState([
+        _ack(prompt=100, cached=40),
+        _ack(completion=1, out="a", prefill_ms=12.5),
+        _ack(completion=1, out="b", finished=True),
+    ])
+    result = asyncio.run(generate_full(42, _spec(), st))
+    assert result.timings.prefill_ms == 12.5
+    assert result.cached_tokens == 40
+
+
+# A real gap between acks, so a TTFT taken at the wrong ack is distinguishable from one taken
+# at the right one; FakeState yields its whole list within a single event-loop tick.
+_GAP_S = 0.02
+_GAP_MS = _GAP_S * 1000
+
+
+class PacedState(FakeState):
+    async def wait_for_ack(self, uid: int):
+        assert uid == 42
+        for index, reply in enumerate(self._replies):
+            if index:
+                await asyncio.sleep(_GAP_S)
+            yield reply
+
+
+def test_ttft_is_measured_at_the_first_generated_token_not_the_admission_ack():
+    """The admission ack carries prompt_tokens and no token; timing TTFT from it would report
+    roughly the queue time for every request."""
+    st = PacedState([
+        _ack(prompt=5),
+        _ack(completion=1, out="hi"),
+        _ack(completion=1, out="!", finished=True),
+    ])
+    result = asyncio.run(generate_full(42, _spec(), st))
+    # One gap to the first token (not zero, as an admission-timed TTFT would give), two to the end.
+    assert _GAP_MS * 0.5 < result.timings.ttft_ms < _GAP_MS * 1.9
+    assert result.timings.total_ms > result.timings.ttft_ms + _GAP_MS * 0.5
+    assert result.timings.decode_ms > _GAP_MS * 0.5
+
+
+def test_ttft_counts_a_token_the_detokenizer_held_back():
+    """A token whose text is withheld (a trailing partial-stop prefix) is still a generated
+    token: TTFT keys off completion_tokens_delta, not off non-empty output -- otherwise it
+    would skip to the ack after it."""
+    st = PacedState([
+        _ack(prompt=5),
+        _ack(completion=1, out=""),
+        _ack(completion=1, out="done", finished=True),
+    ])
+    result = asyncio.run(generate_full(42, _spec(), st))
+    assert _GAP_MS * 0.5 < result.timings.ttft_ms < _GAP_MS * 1.9
+
+
+def test_stream_gendone_carries_the_same_timings():
+    st = FakeState([
+        _ack(prompt=8, cached=2),
+        _ack(completion=1, out="x", prefill_ms=7.0),
+        _ack(completion=1, out="y", finished=True),
+    ])
+
+    async def drain():
+        return [ev async for ev in generate_events(42, _spec(), st) if isinstance(ev, GenDone)]
+
+    (done,) = asyncio.run(drain())
+    assert done.timings.prefill_ms == 7.0
+    assert done.timings.ttft_ms > 0.0
+
+
+def test_build_metrics_splits_prompt_tokens_into_cached_and_prefilled():
+    metrics = build_metrics(
+        prompt_tokens=8192, completion_tokens=256, cached_tokens=6144,
+        timings=GenTimings(ttft_ms=320.5, prefill_ms=410.0, decode_ms=5120.0, total_ms=5480.2),
+    )
+    # The identity the issue asks for: prompt_tokens == cached_prompt_tokens + prefill_tokens.
+    assert metrics["cached_prompt_tokens"] + metrics["prefill_tokens"] == 8192
+    assert metrics["prefill_tokens"] == 2048
+    assert metrics["prefill_tokens_per_second"] == round(2048 / 0.410, 2)
+    # Decode divides by completion_tokens - 1: the first token came out of prefill.
+    assert metrics["decode_tokens_per_second"] == round(255 / 5.120, 2)
+    assert metrics["decode_tokens"] == 256
+    assert metrics["ttft_ms"] == 320.5
+    assert metrics["total_time_ms"] == 5480.2
+
+
+def test_build_metrics_rates_are_zero_rather_than_dividing_by_zero():
+    metrics = build_metrics(
+        prompt_tokens=10, completion_tokens=1, cached_tokens=0, timings=GenTimings()
+    )
+    assert metrics["prefill_tokens_per_second"] == 0.0
+    # One generated token spans no decode interval, so there is no rate to report.
+    assert metrics["decode_tokens_per_second"] == 0.0
+    assert metrics["prefill_tokens"] == 10
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index e33018e195..ce675a2508 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -687,3 +687,65 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob(
     message = response["choices"][0]["message"]
     assert message["reasoning_content"] == "private thought"
     assert message["content"] == "visible answer"
+
+
+# ----------------------------------------------------- per-request metrics
+def _metrics_replies() -> list[UserReply]:
+    return [
+        UserReply(uid=42, incremental_output="", finished=False, prompt_tokens_delta=10, cached_tokens=4),
+        UserReply(uid=42, incremental_output="hi", finished=True, completion_tokens_delta=2, prefill_ms=8.0),
+    ]
+
+
+def test_non_stream_chat_metrics_only_with_flag():
+    state = FakeState(_metrics_replies())
+    state.config.enable_metrics_report = True
+    response = run(handle_chat_completion(chat_request(tools=None), request=None, state=state, model_sampling={}))
+    metrics = response["metrics"]
+    assert metrics["prefill_time_ms"] == 8.0
+    assert metrics["prefill_tokens"] == 6 and metrics["cached_prompt_tokens"] == 4
+    assert metrics["decode_tokens"] == 2
+    assert metrics["total_time_ms"] > 0
+
+    off = run(handle_chat_completion(chat_request(tools=None), request=None, state=FakeState(_metrics_replies()), model_sampling={}))
+    assert "metrics" not in off
+
+
+def test_non_stream_chat_metrics_report_the_cache_hit_without_enable_cache_report():
+    """--enable-cache-report governs the billing fields in `usage`. Gating the hit inside
+    `metrics` on it too would leave prefill_tokens_per_second computed over tokens that were
+    never forwarded."""
+    state = FakeState(_metrics_replies())
+    state.config.enable_metrics_report = True
+    response = run(handle_chat_completion(chat_request(tools=None), request=None, state=state, model_sampling={}))
+    assert "prompt_tokens_details" not in response["usage"]
+    assert response["metrics"]["cached_prompt_tokens"] == 4
+
+
+def test_stream_chat_metrics_ride_the_usage_chunk():
+    state = FakeState(_metrics_replies())
+    state.config.enable_metrics_report = True
+    req = chat_request(tools=None, stream_options={"include_usage": True})
+
+    async def collect():
+        return [chunk async for chunk in stream_chat_completion_chunks(42, req, state)]
+
+    events = parse_sse(run(collect()))
+    final = next(e for e in reversed(events) if isinstance(e, dict) and e.get("usage"))
+    assert final["metrics"]["prefill_time_ms"] == 8.0
+    assert final["metrics"]["decode_tokens"] == 2
+    # No metrics on the content chunks -- exactly one carries them.
+    assert sum(1 for e in events if isinstance(e, dict) and "metrics" in e) == 1
+
+
+def test_stream_chat_without_include_usage_has_no_metrics_chunk():
+    """Metrics ride the usage chunk, so a client that opted out of usage keeps the plain
+    OpenAI stream rather than getting an extra trailing chunk it never asked for."""
+    state = FakeState(_metrics_replies())
+    state.config.enable_metrics_report = True
+
+    async def collect():
+        return [chunk async for chunk in stream_chat_completion_chunks(42, chat_request(tools=None), state)]
+
+    events = parse_sse(run(collect()))
+    assert not any(isinstance(e, dict) and "metrics" in e for e in events)
diff --git a/tests/server/test_responses_api.py b/tests/server/test_responses_api.py
index 887e78e573..6e7721b891 100644
--- a/tests/server/test_responses_api.py
+++ b/tests/server/test_responses_api.py
@@ -338,10 +338,11 @@ def test_stream_tool_call_events():
 # Route smoke tests
 # --------------------------------------------------------------------------- #
 class FakeState:
-    def __init__(self, outputs, finish_reason=None, cached_tokens=0):
+    def __init__(self, outputs, finish_reason=None, cached_tokens=0, prefill_ms=0.0):
         self._outputs = outputs
         self._finish_reason = finish_reason  # stamped on the terminal ack
         self._cached_tokens = cached_tokens  # stamped on the first ack (admission reply)
+        self._prefill_ms = prefill_ms  # the scheduler stamps it on the first generated token
         self.maintenance_state = "serving"
         self.config = SimpleNamespace(
             mm=SimpleNamespace(text_model_only=False, disabled_encoders=frozenset()),
@@ -364,7 +365,8 @@ async def wait_for_ack(self, uid):
             yield UserReply(uid=uid, incremental_output=text, finished=finished,
                             finish_reason=self._finish_reason if finished else None,
                             prompt_tokens_delta=pt, completion_tokens_delta=ct,
-                            cached_tokens=self._cached_tokens if i == 0 else 0)
+                            cached_tokens=self._cached_tokens if i == 0 else 0,
+                            prefill_ms=self._prefill_ms if i == 0 else 0.0)
 
     async def stream_with_cancellation(self, gen, request, uid):
         async for chunk in gen:
@@ -967,3 +969,38 @@ def test_convert_function_call_output_text_list_stays_a_plain_tool_message():
     spec = RP.convert_responses_to_genspec(req, {})
     assert [m["role"] for m in spec.messages] == ["user", "assistant", "tool"]
     assert spec.messages[2]["content"] == "ab"
+
+
+# --------------------------------------------------------------------------- #
+# Per-request metrics (--enable-metrics-report)
+# --------------------------------------------------------------------------- #
+def test_route_nonstream_metrics_only_with_flag():
+    fake = FakeState([("Hello world", True, 5, 2)], cached_tokens=3, prefill_ms=6.0)
+    fake.config.enable_metrics_report = True
+    body = _client(fake).post("/v1/responses", json={"model": "gpt-x", "input": "hi"}).json()
+    assert body["metrics"]["prefill_time_ms"] == 6.0
+    assert body["metrics"]["prefill_tokens"] == 2 and body["metrics"]["cached_prompt_tokens"] == 3
+    assert body["metrics"]["total_time_ms"] > 0
+
+    off = _client(FakeState([("Hello world", True, 5, 2)])).post(
+        "/v1/responses", json={"model": "gpt-x", "input": "hi"}
+    ).json()
+    assert "metrics" not in off
+
+
+def test_route_stream_completed_event_carries_metrics():
+    fake = FakeState([("Hello world", True, 5, 2)], cached_tokens=3, prefill_ms=6.0)
+    fake.config.enable_metrics_report = True
+    r = _client(fake).post(
+        "/v1/responses", json={"model": "gpt-x", "input": "hi", "stream": True}
+    )
+    completed = None
+    for block in r.text.split("\n\n"):
+        for line in block.split("\n"):
+            if line.startswith("data:"):
+                payload = json.loads(line[len("data:"):].strip())
+                if payload.get("type") == "response.completed":
+                    completed = payload
+    assert completed is not None
+    assert completed["response"]["metrics"]["prefill_time_ms"] == 6.0
+    assert completed["response"]["metrics"]["cached_prompt_tokens"] == 3

From ae8b3cfaef74bb3b687dce8e761f6355171c6137 Mon Sep 17 00:00:00 2001
From: Ran Bao <worksev@gmail.com>
Date: Fri, 18 Sep 2026 01:23:21 +1200
Subject: [PATCH 043/249] fix(scheduler): preserve pending hybrid checkpoint
 across prefill chunks

---
 python/freetoken/scheduler/prefill.py        |  5 ++
 tests/scheduler/test_hybrid_cache_manager.py | 83 +++++++++++++++++++-
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index bd8129aae8..98272c3248 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -142,6 +142,7 @@ def _add_one_req(
         next_track_idx: int = 0,
         restore_src: int | None = None,
         swa_evicted_seqlen: int = 0,
+        last_track_seqlen: int | None = None,
     ) -> Req | None:
         remain_len = pending_req.input_len - cached_len
         chunk_size = min(self.token_budget, remain_len)
@@ -223,6 +224,9 @@ def _add_one_req(
         req.linear_slot_idx = linear_slot_idx
         req.mamba_ping_pong = ping_pong
         req.mamba_next_track_idx = next_track_idx
+        # A short continuation cannot make a new snapshot; retain the pending one
+        # together with its ping-pong cursor until the final prefill commit.
+        req.mamba_last_track_seqlen = last_track_seqlen
         req.mamba_restore_src = restore_src
         req.swa_evicted_seqlen = swa_evicted_seqlen  # carry the extend-free watermark across chunks
         return req
@@ -240,6 +244,7 @@ def try_add_one(self, pending_req: PendingReq) -> Req | None:
                 linear_slot_idx=chunked_req.linear_slot_idx,
                 ping_pong=chunked_req.mamba_ping_pong,
                 next_track_idx=chunked_req.mamba_next_track_idx,
+                last_track_seqlen=chunked_req.mamba_last_track_seqlen,
                 restore_src=None,  # continuation chunk already has live state
                 swa_evicted_seqlen=chunked_req.swa_evicted_seqlen,  # extend-free watermark so far
             )
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index 098f27660c..282ab2ca43 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -6,6 +6,7 @@
 from types import SimpleNamespace
 
 import torch
+import pytest
 
 from freetoken.core import Req, SamplingParams
 from freetoken.kvcache.linear_state_pool import LinearStatePool
@@ -153,6 +154,83 @@ def test_naive_cache_does_not_align_prefill_chunks():
     assert adder.try_add_one(pending).extend_len == 100
 
 
+@pytest.mark.parametrize("full_chunks", [2, 3])
+@pytest.mark.parametrize("tail", [1, 59, 64, 65])
+@pytest.mark.parametrize("finish_early", [False, True])
+def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, tail, finish_early):
+    """A short final extend must retain the prior state; a longer one must replace it."""
+    import freetoken.core as core
+    from freetoken.attention.linear import build_fla_metadata
+    from freetoken.scheduler.decode import DecodeManager
+    from freetoken.scheduler.prefill import ChunkedReq, PrefillManager
+    from freetoken.scheduler.scheduler import Scheduler
+    from freetoken.scheduler.table import TableManager
+    from freetoken.scheduler.utils import PendingReq
+
+    pool = _pool()
+    pt = torch.zeros(2, 1024, dtype=torch.int32)
+    cm = CacheManager(32, 64, pt, "hybrid_radix", linear_state_pool=pool)
+    tm = TableManager(max_running_reqs=1, page_table=pt)
+    pm = PrefillManager(cm, tm, DecodeManager(page_size=64))
+    monkeypatch.setattr(core, "_GLOBAL_CTX", core.Context(page_size=64, linear_state_pool=pool))
+    prompt = torch.arange(full_chunks * 128 + tail, dtype=torch.int32)
+    pm.pending_list = [PendingReq(0, prompt, SamplingParams(max_tokens=8))]
+    last_batch = None
+    states = {}
+    final = None
+
+    while pm.runnable or last_batch is not None:
+        # Match overlap ordering: prepare the continuation before draining the prior chunk.
+        batch = pm.schedule_next_batch(128)
+        if batch is not None:
+            batch.padded_reqs = batch.reqs
+            cm.allocate_paged(batch.reqs)
+            req = batch.reqs[0]
+            meta = build_fla_metadata(batch, torch.device("cpu"))
+            if meta.track_dst is not None:
+                slot = meta.track_dst.item()
+                value = len(states) + 1
+                states[req.mamba_last_track_seqlen] = value
+                pool.recurrent_states[:, slot].fill_(value)
+                pool.conv_states[:, slot].fill_(value)
+            pool.recurrent_states[:, req.linear_slot_idx].fill_(-9)
+            pool.conv_states[:, req.linear_slot_idx].fill_(-9)
+            req.complete_one()
+        if last_batch is not None:
+            previous = last_batch.reqs[0]
+            if not isinstance(previous, ChunkedReq):
+                final = previous
+                if not finish_early:
+                    cm.cache_req(final, finished=False)
+        last_batch = batch
+
+    # Before the finish path can donate an aligned live state, require the frozen checkpoint.
+    expected = full_chunks * 128 - 64 if tail <= 64 else full_chunks * 128 + 64
+    if not finish_early:
+        match = cm.match_req(_pend(prompt.tolist() + [999]))
+        assert match.cuda_handle.cached_len == expected
+        assert torch.all(pool.recurrent_states[:, match.mamba_value] == states[expected])
+        assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected])
+
+    # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path.
+    stub = SimpleNamespace(cache_manager=cm, table_manager=tm)
+    Scheduler._free_req_resources(stub, final)
+    Scheduler._free_req_resources(stub, final)
+    match = cm.match_req(_pend(prompt.tolist() + [999]))
+    expected_finish = len(prompt) if tail == 64 else expected
+    assert match.cuda_handle.cached_len == expected_finish
+    value = -9 if tail == 64 else states[expected]
+    assert torch.all(pool.recurrent_states[:, match.mamba_value] == value)
+    assert torch.all(pool.conv_states[:, match.mamba_value] == value)
+    assert tm.available_size == 1
+    cm.check_integrity()
+    assert pool.num_free_slots + cm.prefix_cache.mamba_evictable_size == pool.num_slots - 1
+    cm.ensure_mamba_slots(pool.num_slots - 1)
+    cm.check_integrity()
+    assert pool.num_free_slots == pool.num_slots - 1
+    assert len(cm.free_slots) == cm.num_pages
+
+
 def test_pool_sizing_covers_4mr_floor():
     """C6: pool must reserve the 4-slot-per-request non-evictable floor even at a tiny ratio."""
     from types import SimpleNamespace
@@ -164,7 +242,4 @@ def test_pool_sizing_covers_4mr_floor():
 
 
 if __name__ == "__main__":
-    for name, fn in list(globals().items()):
-        if name.startswith("test_") and callable(fn):
-            fn()
-            print(f"{name}: PASS")
+    raise SystemExit(pytest.main([__file__]))

From 56ed5700337a9fdcb00ad9821b3592794cc13a22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= <Avicennasis@gmail.com>
Date: Mon, 31 Aug 2026 03:15:12 -0400
Subject: [PATCH 044/249] fix(checkpoint): validate FTW v1 indexes eagerly

Malformed or truncated FTW indexes used to surface as mmap/O_DIRECT failures
deep inside the load. Parse and check the index once at open (integer ranges,
required fields, duplicate keys, entry names, shard coverage, shard size), raise
FTWFormatError, and reject nothing the writer produces. Shards are stat'd
through symlinks because Hugging Face hub snapshots are symlink farms.
---
 python/freetoken/checkpoint/__init__.py |   3 +-
 python/freetoken/checkpoint/ftw.py      | 434 +++++++++++++++-
 tests/checkpoint/__init__.py            |   1 +
 tests/checkpoint/test_ftw_validation.py | 640 ++++++++++++++++++++++++
 4 files changed, 1054 insertions(+), 24 deletions(-)
 create mode 100644 tests/checkpoint/__init__.py
 create mode 100644 tests/checkpoint/test_ftw_validation.py

diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py
index f3ce70fcdc..406fdf997a 100644
--- a/python/freetoken/checkpoint/__init__.py
+++ b/python/freetoken/checkpoint/__init__.py
@@ -5,6 +5,7 @@
 """
 
 from .ftw import (
+    FTWFormatError,
     FTWReader,
     FTWWriter,
     is_ftw_checkpoint,
@@ -14,6 +15,6 @@
 from .convert import convert_checkpoint
 
 __all__ = [
-    "FTWReader", "FTWWriter", "is_ftw_checkpoint",
+    "FTWFormatError", "FTWReader", "FTWWriter", "is_ftw_checkpoint",
     "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint",
 ]
diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index e2bd0de24e..76729b79e6 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -37,6 +37,7 @@
 import mmap
 import os
 import re
+import stat
 import threading
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
@@ -59,6 +60,14 @@
 # Per-layer expert-bank entry name (converter streaming path, see checkpoint/convert.py):
 # each layer of a bank is its own FTW tensor instead of one flat [num_layers*E, ...] region.
 _LAYER_ENTRY_RE = re.compile(r"^(?P<base>.+)#L(?P<layer>\d{5})$")
+_MAX_INDEX_INT = (1 << 63) - 1
+_INDEX_FIELDS = {
+    "format", "version", "align", "shard_limit", "total_bytes", "tensors", "shards",
+}
+
+
+class FTWFormatError(ValueError):
+    """The FTW v1 index or one of its declared shards is structurally invalid."""
 
 
 def layer_bank_entry_name(bank_name: str, layer_id: int) -> str:
@@ -102,6 +111,321 @@ def _elsize(dt: torch.dtype) -> int:
     return torch.empty((), dtype=dt).element_size()
 
 
+def _index_int(value, field: str, *, minimum: int = 0) -> int:
+    """Return a JSON integer that is safe to pass to file/tensor APIs.
+
+    ``bool`` is deliberately excluded even though it subclasses ``int`` in Python.
+    File offsets, byte counts, and dimensions ultimately cross signed 64-bit APIs, so
+    accepting larger arbitrary-precision JSON numbers only postpones a less useful
+    overflow failure until allocation or I/O.
+    """
+    if type(value) is not int or not minimum <= value <= _MAX_INDEX_INT:
+        raise FTWFormatError(
+            f"{field} must be an integer in [{minimum}, {_MAX_INDEX_INT}], got {value!r}"
+        )
+    return value
+
+
+def _required(obj: dict, key: str, owner: str):
+    if key not in obj:
+        raise FTWFormatError(f"{owner} is missing required field {key!r}")
+    return obj[key]
+
+
+def _json_object(pairs: list[tuple[str, object]]) -> dict:
+    """Build one JSON object while rejecting ambiguous duplicate keys."""
+    obj = {}
+    for key, value in pairs:
+        if key in obj:
+            raise FTWFormatError(f"FTW index contains duplicate JSON key {key!r}")
+        obj[key] = value
+    return obj
+
+
+def _invalid_json_constant(value: str):
+    raise FTWFormatError(f"FTW index contains non-standard JSON constant {value}")
+
+
+def _json_float(value: str) -> float:
+    parsed = float(value)
+    if not math.isfinite(parsed):
+        raise FTWFormatError(f"FTW index floating-point value is out of range: {value}")
+    return parsed
+
+
+def _validate_metadata(value, field: str = "metadata") -> None:
+    """Reject metadata that cannot round-trip through strict JSON unambiguously."""
+    if isinstance(value, dict):
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise ValueError(f"{field} keys must be strings, got {key!r}")
+            _validate_metadata(item, f"{field}.{key}")
+    elif isinstance(value, (list, tuple)):
+        for i, item in enumerate(value):
+            _validate_metadata(item, f"{field}[{i}]")
+    elif isinstance(value, float):
+        if not math.isfinite(value):
+            raise ValueError(f"{field} must not contain NaN or infinity")
+    elif value is not None and not isinstance(value, (str, int, bool)):
+        raise ValueError(
+            f"{field} must contain only JSON-compatible values, got {type(value).__name__}"
+        )
+
+
+def _safe_shard_basename(value, field: str) -> str:
+    if (not isinstance(value, str) or not value or value in {".", ".."}
+            or os.path.basename(value) != value or "/" in value or "\\" in value
+            or "\x00" in value):
+        raise FTWFormatError(f"{field} must be a safe shard basename, got {value!r}")
+    return value
+
+
+def _validate_ftw_v1_index(path: str, index) -> tuple[list[dict], list[dict]]:
+    """Eagerly validate all v1 structure used by :class:`FTWReader`.
+
+    This is intentionally a structural check, not a content-authenticity check. It
+    prevents malformed JSON from turning into path traversal, integer overflow,
+    overlapping tensor views, or delayed mmap/O_DIRECT failures. A future format can
+    add hashes without weakening the valid v1 contract checked here.
+    """
+    if not isinstance(index, dict):
+        raise FTWFormatError(f"FTW index root must be an object, got {type(index).__name__}")
+
+    fmt = _required(index, "format", "FTW index")
+    if fmt != FORMAT_TAG:
+        raise FTWFormatError(f"not a {FORMAT_TAG} checkpoint: {path} (format={fmt!r})")
+
+    version = _index_int(_required(index, "version", "FTW index"), "version")
+    if version != FORMAT_VERSION:
+        raise FTWFormatError(
+            f"unsupported FTW version {version}; this reader supports {FORMAT_VERSION}"
+        )
+
+    align = _index_int(_required(index, "align", "FTW index"), "align", minimum=1)
+    if align != ALIGN:
+        raise FTWFormatError(f"unsupported FTW alignment {align}; expected {ALIGN}")
+
+    shard_limit = _index_int(
+        _required(index, "shard_limit", "FTW index"), "shard_limit", minimum=ALIGN
+    )
+    if shard_limit % ALIGN:
+        raise FTWFormatError(f"shard_limit must be a multiple of {ALIGN}, got {shard_limit}")
+
+    total_bytes = _index_int(
+        _required(index, "total_bytes", "FTW index"), "total_bytes"
+    )
+    if total_bytes % ALIGN:
+        raise FTWFormatError(f"total_bytes must be {ALIGN}-aligned, got {total_bytes}")
+
+    raw_shards = _required(index, "shards", "FTW index")
+    if not isinstance(raw_shards, list):
+        raise FTWFormatError("shards must be an array")
+
+    shard_names: set[str] = set()
+    shards: list[dict] = []
+    for i, shard in enumerate(raw_shards):
+        owner = f"shards[{i}]"
+        if not isinstance(shard, dict):
+            raise FTWFormatError(f"{owner} must be an object")
+        file = _safe_shard_basename(_required(shard, "file", owner), f"{owner}.file")
+        if file in shard_names:
+            raise FTWFormatError(f"duplicate shard file {file!r}")
+        shard_names.add(file)
+
+        global_off = _index_int(
+            _required(shard, "global_off", owner), f"{owner}.global_off"
+        )
+        nbytes = _index_int(_required(shard, "nbytes", owner), f"{owner}.nbytes")
+        if global_off % ALIGN:
+            raise FTWFormatError(
+                f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}"
+            )
+        if nbytes % ALIGN:
+            raise FTWFormatError(f"{owner}.nbytes must be {ALIGN}-aligned, got {nbytes}")
+        if nbytes > shard_limit:
+            raise FTWFormatError(
+                f"{owner}.nbytes {nbytes} exceeds shard_limit {shard_limit}"
+            )
+        if global_off > _MAX_INDEX_INT - nbytes:
+            raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range")
+
+        shard_path = os.path.join(path, file)
+        try:
+            # Follow symlinks: Hugging Face hub snapshots are symlink farms, so a shard
+            # entry routinely resolves into blobs/. The regular-file and size checks
+            # below therefore apply to the link target.
+            shard_stat = os.stat(shard_path)
+        except (OSError, UnicodeError) as exc:
+            # UnicodeError: a JSON-escaped lone surrogate is a valid str but not a
+            # filesystem name; keep it inside the format-error contract.
+            raise FTWFormatError(f"cannot stat FTW shard {file!r}: {exc}") from exc
+        if not stat.S_ISREG(shard_stat.st_mode):
+            raise FTWFormatError(f"FTW shard {file!r} is not a regular file")
+        if shard_stat.st_size != nbytes:
+            raise FTWFormatError(
+                f"FTW shard {file!r} size mismatch: index declares {nbytes}, "
+                f"file has {shard_stat.st_size}"
+            )
+        shards.append(shard)
+
+    shards.sort(key=lambda shard: shard["global_off"])
+    if any(shard["nbytes"] == 0 for shard in shards):
+        # FTWWriter can produce one empty shard for a checkpoint containing only
+        # zero-sized tensors. It cannot produce empty shards in a non-empty stream.
+        if total_bytes != 0 or len(shards) != 1:
+            raise FTWFormatError("zero-length shards are only valid as one empty FTW shard")
+
+    cursor = 0
+    for shard in shards:
+        if shard["global_off"] != cursor:
+            relation = "an overlap" if shard["global_off"] < cursor else "a gap"
+            raise FTWFormatError(
+                f"FTW shard coverage has {relation}: expected offset {cursor}, "
+                f"got {shard['global_off']} for {shard['file']!r}"
+            )
+        cursor += shard["nbytes"]
+    if cursor != total_bytes:
+        raise FTWFormatError(
+            f"shards cover {cursor} logical bytes but total_bytes is {total_bytes}"
+        )
+
+    raw_tensors = _required(index, "tensors", "FTW index")
+    if not isinstance(raw_tensors, list):
+        raise FTWFormatError("tensors must be an array")
+
+    tensor_names: set[str] = set()
+    tensors: list[dict] = []
+    allocations: list[tuple[int, int, str]] = []
+    zero_allocations: list[tuple[int, str]] = []
+    dtype_sizes: dict[torch.dtype, int] = {}
+    for i, tensor in enumerate(raw_tensors):
+        owner = f"tensors[{i}]"
+        if not isinstance(tensor, dict):
+            raise FTWFormatError(f"{owner} must be an object")
+
+        name = _required(tensor, "name", owner)
+        if not isinstance(name, str):
+            raise FTWFormatError(f"{owner}.name must be a string")
+        if name in tensor_names:
+            raise FTWFormatError(f"duplicate tensor name {name!r}")
+        tensor_names.add(name)
+
+        kind = _required(tensor, "kind", owner)
+        if not isinstance(kind, str):
+            raise FTWFormatError(f"{owner}.kind must be a string")
+
+        dtype_name = _required(tensor, "dtype", owner)
+        if not isinstance(dtype_name, str):
+            raise FTWFormatError(f"{owner}.dtype must be a string")
+        dtype = getattr(torch, dtype_name, None)
+        if not isinstance(dtype, torch.dtype):
+            raise FTWFormatError(f"{owner}.dtype is not a known torch dtype: {dtype_name!r}")
+        element_size = dtype_sizes.get(dtype)
+        if element_size is None:
+            try:
+                element_size = _elsize(dtype)
+                # Match the actual reconstruction primitive in iter_ftw_weights rather
+                # than accepting every object torch happens to classify as a dtype.
+                torch.frombuffer(bytearray(element_size), dtype=dtype, count=1)
+            except (RuntimeError, TypeError, ValueError) as exc:
+                raise FTWFormatError(
+                    f"{owner}.dtype is not readable FTW storage: {dtype_name!r}"
+                ) from exc
+            dtype_sizes[dtype] = element_size
+
+        shape = _required(tensor, "shape", owner)
+        if not isinstance(shape, list):
+            raise FTWFormatError(f"{owner}.shape must be an array")
+        dims = [_index_int(dim, f"{owner}.shape[{j}]") for j, dim in enumerate(shape)]
+        numel = 1
+        for dim in dims:
+            if dim == 0:
+                numel = 0
+            elif numel and numel > _MAX_INDEX_INT // dim:
+                raise FTWFormatError(f"{owner}.shape element count exceeds signed 64-bit range")
+            else:
+                numel *= dim
+        if numel == 0:
+            try:
+                # Zero-numel shapes allocate no storage, so checking PyTorch's stride
+                # arithmetic here is cheap and catches dimensions that cannot be rebuilt.
+                torch.empty(tuple(dims), dtype=dtype)
+            except (RuntimeError, TypeError, ValueError) as exc:
+                raise FTWFormatError(
+                    f"{owner}.shape cannot be reconstructed by torch: {shape!r}"
+                ) from exc
+        if numel and numel > _MAX_INDEX_INT // element_size:
+            raise FTWFormatError(f"{owner} byte size exceeds signed 64-bit range")
+        expected_nbytes = numel * element_size
+
+        global_off = _index_int(
+            _required(tensor, "global_off", owner), f"{owner}.global_off"
+        )
+        nbytes = _index_int(_required(tensor, "nbytes", owner), f"{owner}.nbytes")
+        if nbytes != expected_nbytes:
+            raise FTWFormatError(
+                f"{owner}.nbytes is {nbytes}, but shape {shape!r} and dtype "
+                f"{dtype_name!r} require {expected_nbytes}"
+            )
+        if global_off % ALIGN:
+            raise FTWFormatError(
+                f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}"
+            )
+        if global_off > _MAX_INDEX_INT - nbytes:
+            raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range")
+        end = global_off + nbytes
+        padded_end = _align_up(end)
+        if end > total_bytes or padded_end > total_bytes:
+            raise FTWFormatError(
+                f"{owner} range [{global_off}, {end}) (padded to {padded_end}) "
+                f"exceeds total_bytes {total_bytes}"
+            )
+        if nbytes:
+            allocations.append((global_off, padded_end, name))
+        else:
+            zero_allocations.append((global_off, name))
+        tensors.append(tensor)
+
+    allocations.sort(key=lambda item: (item[0], item[1], item[2]))
+    allocation_boundaries = {0}
+    cursor = 0
+    previous_name: str | None = None
+    for start, end, name in allocations:
+        if start < cursor:
+            raise FTWFormatError(
+                f"tensor {name!r} overlaps padded allocation for {previous_name!r}"
+            )
+        if start > cursor:
+            raise FTWFormatError(
+                f"FTW tensor coverage has a gap: expected offset {cursor}, "
+                f"got {start} for {name!r}"
+            )
+        allocation_boundaries.add(start)
+        allocation_boundaries.add(end)
+        cursor = end
+        previous_name = name
+    if cursor != total_bytes:
+        raise FTWFormatError(
+            f"tensors cover {cursor} padded bytes but total_bytes is {total_bytes}"
+        )
+
+    for offset, name in zero_allocations:
+        if offset not in allocation_boundaries:
+            raise FTWFormatError(
+                f"zero-sized tensor {name!r} at offset {offset} is not on a tensor boundary"
+            )
+
+    # Preserve both zero-byte forms emitted by FTWWriter: finalize-without-add_tensor
+    # has no shards, while one or more zero-sized tensors produce one empty shard.
+    if total_bytes == 0:
+        if tensors and not shards:
+            raise FTWFormatError("a zero-sized tensor checkpoint must contain one empty shard")
+        if not tensors and shards:
+            raise FTWFormatError("an empty checkpoint must not contain shard entries")
+
+    return shards, tensors
+
+
 def is_ftw_checkpoint(path: str) -> bool:
     """True if ``path`` is a directory holding a FreeToken Weight (FTW) index."""
     return os.path.isfile(os.path.join(path, INDEX_NAME))
@@ -133,11 +457,17 @@ class FTWWriter:
     """
 
     def __init__(self, out_dir: str, *, shard_limit: int = DEFAULT_SHARD_LIMIT):
-        assert shard_limit % ALIGN == 0, "shard_limit must be a multiple of ALIGN"
+        if (type(shard_limit) is not int or shard_limit < ALIGN
+                or shard_limit > _MAX_INDEX_INT or shard_limit % ALIGN):
+            raise ValueError(
+                f"shard_limit must be an integer multiple of {ALIGN} in "
+                f"[{ALIGN}, {_MAX_INDEX_INT}], got {shard_limit!r}"
+            )
         os.makedirs(out_dir, exist_ok=True)
         self.out_dir = out_dir
         self.shard_limit = shard_limit
         self._tensors: list[dict] = []
+        self._tensor_names: set[str] = set()
         self._shards: list[dict] = []
         self._global = 0  # running FTW offset (incl. padding)
         self._f = None  # current shard file handle
@@ -171,6 +501,10 @@ def _write_raw(self, data: memoryview) -> None:
             self._global += take
 
     def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> None:
+        if not isinstance(name, str) or not isinstance(kind, str):
+            raise ValueError("FTW tensor name and kind must be strings")
+        if name in self._tensor_names:
+            raise ValueError(f"duplicate FTW tensor name {name!r}")
         t = tensor.detach().cpu().contiguous()
         raw = t.reshape(-1).view(torch.uint8)
         nbytes = int(raw.numel())
@@ -179,16 +513,24 @@ def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> N
                                and self._cur + nbytes > self.shard_limit):
             self._roll()
         global_off = self._global
-        assert global_off % ALIGN == 0, "tensor start must be aligned (invariant)"
+        if global_off % ALIGN:
+            raise RuntimeError("FTW writer invariant failed: tensor start is not aligned")
         self._write_raw(memoryview(raw.numpy()))
         self._tensors.append({"name": name, "kind": kind, "dtype": _dtype_str(t.dtype),
                               "shape": list(t.shape), "global_off": global_off, "nbytes": nbytes})
+        self._tensor_names.add(name)
         # pad to ALIGN so the next tensor starts aligned
         pad = _align_up(self._global) - self._global
         if pad:
             self._write_raw(memoryview(bytes(pad)))
 
     def finalize(self, meta: dict) -> dict:
+        if not isinstance(meta, dict):
+            raise ValueError(f"FTW metadata must be an object, got {type(meta).__name__}")
+        collisions = _INDEX_FIELDS & meta.keys()
+        if collisions:
+            raise ValueError(f"FTW metadata cannot override index fields: {sorted(collisions)}")
+        _validate_metadata(meta)
         if self._f is not None:
             self._shards.append({"file": _SHARD_FMT.format(self._shard_idx),
                                  "global_off": self._shard_start, "nbytes": self._cur})
@@ -199,7 +541,7 @@ def finalize(self, meta: dict) -> dict:
                  "tensors": self._tensors, "shards": self._shards, **meta}
         tmp = os.path.join(self.out_dir, INDEX_NAME + ".tmp")
         with open(tmp, "w") as f:
-            json.dump(index, f)
+            json.dump(index, f, allow_nan=False)
         os.replace(tmp, os.path.join(self.out_dir, INDEX_NAME))
         return index
 
@@ -215,12 +557,22 @@ class FTWReader:
     (the rounding reads into the region's padding, which is discarded by the tensor view)."""
 
     def __init__(self, path: str):
-        with open(os.path.join(path, INDEX_NAME)) as f:
-            self.index = json.load(f)
-        assert self.index.get("format") == FORMAT_TAG, f"not a {FORMAT_TAG}: {path}"
+        index_path = os.path.join(path, INDEX_NAME)
+        try:
+            with open(index_path, encoding="utf-8") as f:
+                self.index = json.load(
+                    f,
+                    object_pairs_hook=_json_object,
+                    parse_constant=_invalid_json_constant,
+                    parse_float=_json_float,
+                )
+        except FTWFormatError:
+            raise
+        except (ValueError, RecursionError) as exc:
+            raise FTWFormatError(f"malformed FTW index {index_path!r}: {exc}") from exc
         self.dir = path
-        self.shards = sorted(self.index["shards"], key=lambda s: s["global_off"])
-        self.tensors = {t["name"]: t for t in self.index["tensors"]}
+        self.shards, tensors = _validate_ftw_v1_index(path, self.index)
+        self.tensors = {tensor["name"]: tensor for tensor in tensors}
         self._fds: dict[str, int] = {}
         self._maps: dict[str, tuple[mmap.mmap, memoryview]] = {}
         # O_DIRECT (DMA straight from disk, bypassing the page cache) is the fast path but a
@@ -359,7 +711,9 @@ def rd(job):
 
 
 def _transient_buffer(nbytes: int) -> mmap.mmap:
-    return mmap.mmap(-1, _align_up(nbytes))
+    # mmap rejects a zero-length mapping, but FTWWriter legitimately emits zero-sized
+    # tensors. Give those tensors one aligned backing page while exposing count=0 below.
+    return mmap.mmap(-1, _align_up(nbytes) or ALIGN)
 
 
 def iter_ftw_weights(path: str, *, kinds=("weight",), keep: Callable[[str], bool] | None = None,
@@ -402,9 +756,14 @@ def _producer():
                 buf = _transient_buffer(e["nbytes"])
                 reader.read_into(memoryview(buf), e, workers=workers, chunk=chunk)
                 dt = _dtype_of(e["dtype"])
-                t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt))
-                # a 0-d entry (a per-tensor scale) comes back 0-d, not [1]
-                if not _put((e["name"], t.view(*e["shape"]) if e["shape"] else t.view(()), buf, e["nbytes"])):
+                if e["nbytes"]:
+                    t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt))
+                    # a 0-d entry (a per-tensor scale) comes back 0-d, not [1]
+                    t = t.view(tuple(e["shape"])) if e["shape"] else t.view(())
+                else:
+                    # torch.frombuffer rejects count=0 even with a non-empty backing map.
+                    t = torch.empty(tuple(e["shape"]), dtype=dt)
+                if not _put((e["name"], t, buf, e["nbytes"])):
                     return
         except BaseException as ex:  # surface to consumer
             err.append(ex)
@@ -465,13 +824,20 @@ def load_ftw_banks(
     vectors, unaffected by the row split (fixed GPU residency; see
     ``cache_budget.expert_bytes_per_slot``).
     """
+    if type(num_layers) is not int or num_layers < 1:
+        raise ValueError(f"num_layers must be a positive integer, got {num_layers!r}")
+    if layer_residency is not None and len(layer_residency) != num_layers:
+        raise ValueError(
+            f"layer_residency has {len(layer_residency)} entries; expected {num_layers}"
+        )
+
     from freetoken.moe.host_banks import (
         HostBank, HostResidency, PinPipeline, alloc_banks, born_pinned_default,
     )
     from freetoken.utils.progress import byte_bar
 
-    residency = layer_residency or [HostResidency.PINNED.value] * num_layers
-    assert len(residency) == num_layers, (len(residency), num_layers)
+    residency = ([HostResidency.PINNED.value] * num_layers
+                 if layer_residency is None else layer_residency)
 
     # PINNED layers are born-pinned (cudaHostAlloc) where that wins (see born_pinned_default); LOCKED/PAGEABLE layers stay lazy mmaps
     born = born_pinned_default()
@@ -514,7 +880,11 @@ def _backing(layer_id: int) -> str:
         per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e
 
     mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys()
-    assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+    if mixed:
+        reader.close()
+        raise FTWFormatError(
+            f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+        )
 
     # Row banks: one padded-window HostBank per (name, layer_id) for the flat layout, plus
     # how to carve the real [num_experts, *row_shape] tensor out of its head; ``None`` marks
@@ -526,13 +896,25 @@ def _backing(layer_id: int) -> str:
 
     for e in flat_entries:
         name = e["name"]
+        if not e["shape"]:
+            reader.close()
+            raise FTWFormatError(f"FTW bank {name!r} must have at least one dimension")
         total, *row_shape = e["shape"]
-        assert total % num_layers == 0, (name, total, num_layers)
+        if total % num_layers:
+            reader.close()
+            raise FTWFormatError(
+                f"FTW bank {name!r} has {total} rows, not divisible by "
+                f"num_layers={num_layers}"
+            )
         num_experts = total // num_layers
         dtype = _dtype_of(e["dtype"])
         row_bytes = (math.prod(row_shape) if row_shape else 1) * _elsize(dtype)
         layer_bytes = num_experts * row_bytes
-        assert layer_bytes * num_layers == e["nbytes"], (name, layer_bytes, num_layers, e["nbytes"])
+        if layer_bytes * num_layers != e["nbytes"]:
+            reader.close()
+            raise FTWFormatError(
+                f"FTW bank {name!r} byte geometry does not match its layer layout"
+            )
         row_hb[name] = []
         row_view_args[name] = []
         for layer_id in range(num_layers):
@@ -546,15 +928,21 @@ def _backing(layer_id: int) -> str:
             row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
 
     for base, by_layer in per_layer_groups.items():
-        assert sorted(by_layer) == list(range(num_layers)), (
-            f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
-            f"expected exactly range({num_layers})"
-        )
+        if sorted(by_layer) != list(range(num_layers)):
+            reader.close()
+            raise FTWFormatError(
+                f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
+                f"expected exactly range({num_layers})"
+            )
         row_hb[base] = []
         row_view_args[base] = []
         for layer_id in range(num_layers):
             e = by_layer[layer_id]
-            assert e["global_off"] % ALIGN == 0, (base, layer_id, e["global_off"])  # writer invariant
+            if e["global_off"] % ALIGN:
+                reader.close()
+                raise FTWFormatError(
+                    f"FTW bank {base!r} layer {layer_id} is not {ALIGN}-aligned"
+                )
             bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id))
             row_hb[base].append(bank)
             row_view_args[base].append(None)
@@ -661,6 +1049,6 @@ def _read_layer(job):
 
 __all__ = [
     "INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT",
-    "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader",
+    "FTWFormatError", "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader",
     "iter_ftw_weights", "load_ftw_banks", "layer_bank_entry_name",
 ]
diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py
new file mode 100644
index 0000000000..5b8db6dafd
--- /dev/null
+++ b/tests/checkpoint/__init__.py
@@ -0,0 +1 @@
+"""Checkpoint-format tests."""
diff --git a/tests/checkpoint/test_ftw_validation.py b/tests/checkpoint/test_ftw_validation.py
new file mode 100644
index 0000000000..0f9acb4463
--- /dev/null
+++ b/tests/checkpoint/test_ftw_validation.py
@@ -0,0 +1,640 @@
+from __future__ import annotations
+
+import copy
+import json
+import mmap
+import os
+from pathlib import Path
+import subprocess
+import sys
+import types
+
+import pytest
+import torch
+
+from freetoken.checkpoint.ftw import (
+    ALIGN,
+    FORMAT_TAG,
+    FORMAT_VERSION,
+    INDEX_NAME,
+    FTWFormatError,
+    FTWReader,
+    FTWWriter,
+    iter_ftw_weights,
+)
+
+MAX_INDEX_INT = (1 << 63) - 1
+
+
+def _write_index(path, index) -> None:
+    with open(path / INDEX_NAME, "w") as f:
+        json.dump(index, f)
+
+
+def _write_raw_index(path, data: bytes) -> None:
+    with open(path / INDEX_NAME, "wb") as f:
+        f.write(data)
+
+
+def _two_tensor_checkpoint(tmp_path):
+    path = tmp_path / "checkpoint"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("first", torch.arange(1000, dtype=torch.float32))  # 4000 + padding
+    writer.add_tensor("second", torch.arange(4, dtype=torch.int16))  # 8 + padding
+    return path, writer.finalize({"model_type": "test"})
+
+
+def _three_tensor_checkpoint(tmp_path):
+    path = tmp_path / "checkpoint"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    for name in ("first", "middle", "last"):
+        writer.add_tensor(name, torch.arange(8, dtype=torch.uint8))
+    return path, writer.finalize({})
+
+
+def _expect_invalid(path, index, match: str) -> None:
+    _write_index(path, index)
+    with pytest.raises(FTWFormatError, match=match):
+        FTWReader(str(path))
+
+
+def test_valid_writer_output_reads_tensor_spanning_shards(tmp_path):
+    path = tmp_path / "checkpoint"
+    expected = (torch.arange(5000, dtype=torch.int64) % 251).to(torch.uint8)
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("spanning", expected)
+    index = writer.finalize({})
+
+    assert len(index["shards"]) == 2
+    reader = FTWReader(str(path))
+    reader._direct = 0  # make this portable across tmpfs/overlayfs and CI kernels
+    entry = reader.tensors["spanning"]
+    pieces = list(reader._pieces(entry["global_off"], entry["nbytes"]))
+    assert [piece[3] for piece in pieces] == [ALIGN, 5000 - ALIGN]
+
+    buf = mmap.mmap(-1, 2 * ALIGN)
+    dest = memoryview(buf)
+    try:
+        reader.read_into(dest, entry, workers=1, chunk=ALIGN)
+        actual = torch.frombuffer(buf, dtype=torch.uint8, count=5000).clone()
+        assert torch.equal(actual, expected)
+    finally:
+        dest.release()
+        buf.close()
+        reader.close()
+
+
+def test_valid_empty_and_zero_tensor_writer_outputs(tmp_path):
+    empty_path = tmp_path / "empty"
+    empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({})
+    assert empty_index["shards"] == []
+    empty_reader = FTWReader(str(empty_path))
+    assert empty_reader.shards == []
+    empty_reader.close()
+
+    zero_path = tmp_path / "zero"
+    zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN)
+    zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32))
+    zero_index = zero_writer.finalize({})
+    assert zero_index["shards"][0]["nbytes"] == 0
+    zero_reader = FTWReader(str(zero_path))
+    assert zero_reader.tensors["zero"]["nbytes"] == 0
+    zero_reader.close()
+
+
+def test_valid_zero_tensors_at_writer_allocation_boundaries(tmp_path):
+    path = tmp_path / "mixed-zero"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("zero-before", torch.empty(0, dtype=torch.float32))
+    writer.add_tensor("first", torch.arange(1000, dtype=torch.float32))
+    writer.add_tensor("zero-between", torch.empty((2, 0, 3), dtype=torch.int16))
+    writer.add_tensor("second", torch.arange(4, dtype=torch.int16))
+    writer.add_tensor("zero-after", torch.empty(0, dtype=torch.uint8))
+    index = writer.finalize({})
+
+    reader = FTWReader(str(path))
+    assert [reader.tensors[name]["global_off"] for name in (
+        "zero-before", "zero-between", "zero-after",
+    )] == [0, ALIGN, 2 * ALIGN]
+    assert index["total_bytes"] == 2 * ALIGN
+    reader.close()
+
+
+def test_zero_tensor_can_be_iterated(tmp_path, monkeypatch):
+    path = tmp_path / "zero-iteration"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("zero", torch.empty((2, 0, 3), dtype=torch.float32))
+    writer.finalize({})
+
+    class _Progress:
+        def update(self, _nbytes):
+            pass
+
+        def close(self):
+            pass
+
+    progress = types.ModuleType("freetoken.utils.progress")
+    progress.byte_bar = lambda *_args, **_kwargs: _Progress()
+    monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress)
+
+    loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1))
+    assert len(loaded) == 1
+    assert loaded[0][0] == "zero"
+    assert loaded[0][1].shape == (2, 0, 3)
+
+
+def test_scalar_tensor_round_trips_with_scalar_shape(tmp_path, monkeypatch):
+    path = tmp_path / "scalar-iteration"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("scalar", torch.tensor(17, dtype=torch.int32))
+    writer.finalize({})
+
+    class _Progress:
+        def update(self, _nbytes):
+            pass
+
+        def close(self):
+            pass
+
+    progress = types.ModuleType("freetoken.utils.progress")
+    progress.byte_bar = lambda *_args, **_kwargs: _Progress()
+    monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress)
+
+    loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1))
+    assert loaded[0][0] == "scalar"
+    assert loaded[0][1].shape == ()
+    assert loaded[0][1].item() == 17
+
+
+def test_shards_may_be_declared_out_of_order(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["shards"].reverse()
+    _write_index(path, index)
+
+    reader = FTWReader(str(path))
+    assert [shard["global_off"] for shard in reader.shards] == [0, ALIGN]
+    reader.close()
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "match"),
+    [
+        ("format", "other", "not a freetoken_weight"),
+        ("version", FORMAT_VERSION + 1, "unsupported FTW version"),
+        ("version", True, "version must be an integer"),
+        ("align", ALIGN * 2, "unsupported FTW alignment"),
+        ("shard_limit", ALIGN + 1, "shard_limit must be a multiple"),
+        ("shard_limit", 0, "shard_limit must be an integer"),
+        ("total_bytes", ALIGN + 1, "total_bytes must be 4096-aligned"),
+        ("total_bytes", -1, "total_bytes must be an integer"),
+        ("total_bytes", MAX_INDEX_INT + 1, "total_bytes must be an integer"),
+    ],
+)
+def test_rejects_invalid_header_fields(tmp_path, field, value, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index[field] = value
+    _expect_invalid(path, index, match)
+
+
+def test_rejects_non_object_root_and_missing_fields(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    _expect_invalid(path, [], "index root must be an object")
+
+    index.pop("tensors")
+    _expect_invalid(path, index, "missing required field 'tensors'")
+
+
+@pytest.mark.parametrize("raw", [b"{", b"\xff"])
+def test_rejects_malformed_json_and_encoding(tmp_path, raw):
+    path, _index = _two_tensor_checkpoint(tmp_path)
+    _write_raw_index(path, raw)
+    with pytest.raises(FTWFormatError, match="malformed FTW index"):
+        FTWReader(str(path))
+
+
+def test_rejects_duplicate_json_object_keys(tmp_path):
+    path, _index = _two_tensor_checkpoint(tmp_path)
+    _write_raw_index(
+        path,
+        b'{"format":"freetoken_weight","version":1,"version":1}',
+    )
+    with pytest.raises(FTWFormatError, match="duplicate JSON key 'version'"):
+        FTWReader(str(path))
+
+
+@pytest.mark.parametrize("constant", ["NaN", "Infinity", "-Infinity"])
+def test_rejects_non_standard_json_constants_even_in_metadata(tmp_path, constant):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    raw = json.dumps(index).removesuffix("}") + f', "metadata_value": {constant}}}'
+    _write_raw_index(path, raw.encode())
+    with pytest.raises(FTWFormatError, match="non-standard JSON constant"):
+        FTWReader(str(path))
+
+
+def test_rejects_json_float_that_overflows_python_float(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    raw = json.dumps(index).removesuffix("}") + ', "metadata_value": 1e999}'
+    _write_raw_index(path, raw.encode())
+    with pytest.raises(FTWFormatError, match="floating-point value is out of range"):
+        FTWReader(str(path))
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "match"),
+    [
+        ("shards", {}, "shards must be an array"),
+        ("tensors", {}, "tensors must be an array"),
+    ],
+)
+def test_rejects_non_array_index_collections(tmp_path, field, value, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index[field] = value
+    _expect_invalid(path, index, match)
+
+
+@pytest.mark.parametrize(
+    "file",
+    ["../outside.ftw", "/tmp/outside.ftw", "sub/file.ftw", "sub\\file.ftw", "..", "bad\x00.ftw"],
+)
+def test_rejects_unsafe_shard_names_before_filesystem_access(tmp_path, file):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["shards"][0]["file"] = file
+    _expect_invalid(path, index, "safe shard basename")
+
+
+def test_rejects_duplicate_shard_names(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["shards"][1]["file"] = index["shards"][0]["file"]
+    _expect_invalid(path, index, "duplicate shard file")
+
+
+@pytest.mark.parametrize(
+    ("entry", "match"),
+    [
+        (1, r"shards\[0\] must be an object"),
+        ({"global_off": 0, "nbytes": ALIGN}, "missing required field 'file'"),
+        ({"file": 1, "global_off": 0, "nbytes": ALIGN}, "safe shard basename"),
+    ],
+)
+def test_rejects_malformed_shard_entries(tmp_path, entry, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["shards"][0] = entry
+    _expect_invalid(path, index, match)
+
+
+def test_rejects_missing_truncated_extended_and_non_regular_shards_eagerly(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    first = path / index["shards"][0]["file"]
+    first.unlink()
+    _expect_invalid(path, index, "cannot stat FTW shard")
+
+    path, index = _two_tensor_checkpoint(tmp_path / "truncated")
+    first = path / index["shards"][0]["file"]
+    with open(first, "r+b") as f:
+        f.truncate(ALIGN - 1)
+    _expect_invalid(path, index, "size mismatch")
+
+    path, index = _two_tensor_checkpoint(tmp_path / "extended")
+    first = path / index["shards"][0]["file"]
+    with open(first, "ab") as f:
+        f.write(b"\x00")
+    _expect_invalid(path, index, "size mismatch")
+
+    path, index = _two_tensor_checkpoint(tmp_path / "directory")
+    first = path / index["shards"][0]["file"]
+    first.unlink()
+    first.mkdir()
+    _expect_invalid(path, index, "not a regular file")
+
+
+def test_accepts_symlinked_shard_whose_target_matches(tmp_path):
+    # Hugging Face hub snapshots are symlink farms (every file links into blobs/), so a
+    # shard entry that resolves through a symlink to a matching regular file must load.
+    path, index = _two_tensor_checkpoint(tmp_path)
+    blobs = tmp_path / "blobs"
+    blobs.mkdir()
+    first = path / index["shards"][0]["file"]
+    target = blobs / "first-blob"
+    target.write_bytes(first.read_bytes())
+    first.unlink()
+    first.symlink_to(target)
+
+    reader = FTWReader(str(path))
+    reader._direct = 0
+    buf = mmap.mmap(-1, ALIGN)
+    dest = memoryview(buf)
+    try:
+        reader.read_into(dest, reader.tensors["first"], workers=1, chunk=ALIGN)
+        actual = torch.frombuffer(buf, dtype=torch.float32, count=1000).clone()
+        assert torch.equal(actual, torch.arange(1000, dtype=torch.float32))
+    finally:
+        dest.release()
+        buf.close()
+        reader.close()
+
+
+def test_symlinked_shard_is_checked_against_its_target(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    external = tmp_path / "outside.ftw"
+    external.write_bytes(b"\x00" * (ALIGN + 1))
+    first = path / index["shards"][0]["file"]
+    first.unlink()
+    first.symlink_to(external)
+
+    _expect_invalid(path, index, "size mismatch")
+
+
+def test_rejects_shard_name_that_is_not_a_filesystem_name(tmp_path):
+    # A JSON-escaped lone surrogate decodes to a valid str that os.stat cannot encode;
+    # it must surface as FTWFormatError, not UnicodeEncodeError.
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["shards"][0]["file"] = "\ud800.ftw"
+
+    _expect_invalid(path, index, "cannot stat FTW shard")
+
+
+@pytest.mark.parametrize(
+    ("mutation", "match"),
+    [
+        (lambda index: index["shards"][1].update(global_off=3 * ALIGN), "coverage has a gap"),
+        (lambda index: index["shards"][1].update(global_off=0), "coverage has an overlap"),
+        (lambda index: index.update(total_bytes=3 * ALIGN), "shards cover .* total_bytes"),
+        (lambda index: index["shards"][0].update(global_off=1), "global_off must be 4096-aligned"),
+        (lambda index: index["shards"][0].update(global_off=MAX_INDEX_INT + 1), "global_off must be an integer"),
+        (lambda index: index["shards"][0].update(
+            global_off=(MAX_INDEX_INT // ALIGN) * ALIGN,
+        ), "byte range exceeds signed 64-bit"),
+        (lambda index: index["shards"][0].update(nbytes=2 * ALIGN), "exceeds shard_limit"),
+        (lambda index: index["shards"][0].update(nbytes=ALIGN - 1), "nbytes must be 4096-aligned"),
+        (lambda index: index["shards"][0].update(nbytes=True), "nbytes must be an integer"),
+    ],
+)
+def test_rejects_invalid_shard_geometry(tmp_path, mutation, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    mutation(index)
+    _expect_invalid(path, index, match)
+
+
+def test_rejects_zero_length_shard_in_nonempty_stream(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    empty_name = "empty.ftw"
+    (path / empty_name).touch()
+    index["shards"].insert(0, {"file": empty_name, "global_off": 0, "nbytes": 0})
+    _expect_invalid(path, index, "zero-length shards")
+
+
+def test_rejects_non_writer_zero_byte_shard_forms(tmp_path):
+    empty_path = tmp_path / "empty"
+    empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({})
+    (empty_path / "empty.ftw").touch()
+    empty_index["shards"] = [{"file": "empty.ftw", "global_off": 0, "nbytes": 0}]
+    _expect_invalid(empty_path, empty_index, "empty checkpoint must not contain shard")
+
+    zero_path = tmp_path / "zero"
+    zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN)
+    zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32))
+    zero_index = zero_writer.finalize({})
+    zero_index["shards"] = []
+    _expect_invalid(zero_path, zero_index, "zero-sized tensor checkpoint must contain")
+
+
+@pytest.mark.parametrize(
+    ("mutation", "match"),
+    [
+        (lambda index: index["tensors"][1].update(name="first"), "duplicate tensor name"),
+        (lambda index: index["tensors"][0].update(name=1), "name must be a string"),
+        (lambda index: index["tensors"][0].update(kind=1), "kind must be a string"),
+        (lambda index: index["tensors"][0].update(dtype="not_a_dtype"), "not a known torch dtype"),
+        (lambda index: index["tensors"][0].update(dtype="strided"), "not a known torch dtype"),
+        (lambda index: index["tensors"][0].update(dtype=1), "dtype must be a string"),
+        (lambda index: index["tensors"][0].update(shape="1000"), "shape must be an array"),
+        (lambda index: index["tensors"][0].update(shape=[True]), r"shape\[0\] must be an integer"),
+        (lambda index: index["tensors"][0].update(shape=[1.0]), r"shape\[0\] must be an integer"),
+        (lambda index: index["tensors"][0].update(shape=[-1]), r"shape\[0\] must be an integer"),
+        (lambda index: index["tensors"][0].update(shape=[(1 << 63) - 1, 2]), "element count exceeds"),
+        (lambda index: index["tensors"][0].update(
+            shape=[0, MAX_INDEX_INT, MAX_INDEX_INT], nbytes=0,
+        ), "shape cannot be reconstructed by torch"),
+        (lambda index: index["tensors"][0].update(shape=[MAX_INDEX_INT], dtype="float32"), "byte size exceeds"),
+        (lambda index: index["tensors"][0].update(nbytes=4), "shape .* require 4000"),
+        (lambda index: index["tensors"][0].update(global_off=1), "global_off must be 4096-aligned"),
+        (lambda index: index["tensors"][0].update(
+            dtype="uint8", shape=[ALIGN], nbytes=ALIGN,
+            global_off=(MAX_INDEX_INT // ALIGN) * ALIGN,
+        ), "byte range exceeds signed 64-bit"),
+        (lambda index: index["tensors"][1].update(global_off=2 * ALIGN), "exceeds total_bytes"),
+        (lambda index: index["tensors"][0].update(nbytes=True), "nbytes must be an integer"),
+    ],
+)
+def test_rejects_invalid_tensor_entries(tmp_path, mutation, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    mutation(index)
+    _expect_invalid(path, index, match)
+
+
+@pytest.mark.parametrize(
+    ("entry", "match"),
+    [
+        (1, r"tensors\[0\] must be an object"),
+        ({"kind": "weight", "dtype": "float32", "shape": [1],
+          "global_off": 0, "nbytes": 4}, "missing required field 'name'"),
+    ],
+)
+def test_rejects_malformed_tensor_entries(tmp_path, entry, match):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["tensors"][0] = entry
+    _expect_invalid(path, index, match)
+
+
+def test_rejects_dtype_that_torch_cannot_reconstruct(tmp_path, monkeypatch):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    original = torch.frombuffer
+
+    def reject_float32(buffer, *, dtype, count):
+        if dtype == torch.float32:
+            raise RuntimeError("unsupported storage dtype")
+        return original(buffer, dtype=dtype, count=count)
+
+    monkeypatch.setattr(torch, "frombuffer", reject_float32)
+    _expect_invalid(path, index, "dtype is not readable FTW storage")
+
+
+def test_rejects_tensor_overlap_including_padded_allocation(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["tensors"][1]["global_off"] = 0
+    _expect_invalid(path, index, "overlaps padded allocation")
+
+
+@pytest.mark.parametrize(
+    ("remove", "match"),
+    [
+        (0, "tensor coverage has a gap"),
+        (1, "tensor coverage has a gap"),
+        (2, "tensors cover .* total_bytes"),
+    ],
+)
+def test_rejects_leading_internal_and_trailing_tensor_gaps(tmp_path, remove, match):
+    path, index = _three_tensor_checkpoint(tmp_path)
+    index["tensors"].pop(remove)
+    _expect_invalid(path, index, match)
+
+
+def test_rejects_zero_sized_tensor_inside_nonzero_allocation(tmp_path):
+    path = tmp_path / "checkpoint"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("spanning", torch.arange(5000, dtype=torch.uint8))
+    index = writer.finalize({})
+    index["tensors"].append({
+        "name": "zero-inside",
+        "kind": "weight",
+        "dtype": "float32",
+        "shape": [0],
+        "global_off": ALIGN,
+        "nbytes": 0,
+    })
+    _expect_invalid(path, index, "zero-sized tensor .* is not on a tensor boundary")
+
+
+def test_index_validation_does_not_mutate_tensor_order(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    index["tensors"].reverse()
+    expected = copy.deepcopy(index["tensors"])
+    _write_index(path, index)
+
+    reader = FTWReader(str(path))
+    assert reader.index["tensors"] == expected
+    assert [entry["name"] for entry in reader.entries()] == ["second", "first"]
+    reader.close()
+
+
+def test_writer_emits_declared_v1_header(tmp_path):
+    path, index = _two_tensor_checkpoint(tmp_path)
+    assert index["format"] == FORMAT_TAG
+    assert index["version"] == FORMAT_VERSION
+    assert index["align"] == ALIGN
+    assert sum(os.path.getsize(path / shard["file"]) for shard in index["shards"]) == index["total_bytes"]
+
+
+@pytest.mark.parametrize("shard_limit", [0, ALIGN - 1, ALIGN + 1, True, MAX_INDEX_INT + 1])
+def test_writer_rejects_invalid_shard_limit_without_assertions(tmp_path, shard_limit):
+    with pytest.raises(ValueError, match="shard_limit must be an integer multiple"):
+        FTWWriter(str(tmp_path / "checkpoint"), shard_limit=shard_limit)
+
+
+def test_writer_does_not_emit_indexes_its_reader_would_reject(tmp_path):
+    path = tmp_path / "checkpoint"
+    writer = FTWWriter(str(path), shard_limit=ALIGN)
+    writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32))
+
+    with pytest.raises(ValueError, match="duplicate FTW tensor name"):
+        writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32))
+    with pytest.raises(ValueError, match="cannot override index fields"):
+        writer.finalize({"version": FORMAT_VERSION})
+    with pytest.raises(ValueError, match="must not contain NaN or infinity"):
+        writer.finalize({"metadata_value": float("nan")})
+    with pytest.raises(ValueError, match="keys must be strings"):
+        writer.finalize({"nested": {1: "ambiguous with a JSON string key"}})
+    with pytest.raises(ValueError, match="JSON-compatible values"):
+        writer.finalize({"metadata_value": object()})
+
+    writer.finalize({"metadata_value": 1.5})
+    FTWReader(str(path)).close()
+
+
+def test_validation_remains_active_under_optimized_python(tmp_path):
+    repo_root = Path(__file__).resolve().parents[2]
+    script = r'''
+import json
+import logging
+from pathlib import Path
+import sys
+import types
+
+import torch
+
+utils = types.ModuleType("freetoken.utils")
+utils.init_logger = logging.getLogger
+sys.modules["freetoken.utils"] = utils
+
+from freetoken.checkpoint.ftw import FTWFormatError, FTWReader, FTWWriter, INDEX_NAME, load_ftw_banks
+
+moe = types.ModuleType("freetoken.moe")
+moe.__path__ = []
+host_banks = types.ModuleType("freetoken.moe.host_banks")
+host_banks.HostBank = object
+host_banks.HostResidency = types.SimpleNamespace(PINNED=types.SimpleNamespace(value="pinned"))
+host_banks.PinPipeline = object
+host_banks.alloc_banks = lambda _specs: {}
+host_banks.born_pinned_default = lambda: False
+progress = types.ModuleType("freetoken.utils.progress")
+progress.byte_bar = lambda *_args, **_kwargs: None
+sys.modules["freetoken.moe"] = moe
+sys.modules["freetoken.moe.host_banks"] = host_banks
+sys.modules["freetoken.utils.progress"] = progress
+
+path = Path(sys.argv[1]) / "optimized"
+path.mkdir()
+(path / INDEX_NAME).write_text(json.dumps({"format": "wrong"}))
+
+try:
+    FTWReader(str(path))
+except FTWFormatError:
+    pass
+else:
+    raise SystemExit("malformed index was accepted under -O")
+
+for action in (
+    lambda: FTWWriter(str(path / "writer"), shard_limit=0),
+    lambda: load_ftw_banks(str(path), num_layers=0),
+    lambda: load_ftw_banks(str(path), num_layers=1, layer_residency=[]),
+):
+    try:
+        action()
+    except ValueError:
+        pass
+    else:
+        raise SystemExit("invalid writer/reader argument was accepted under -O")
+
+schema_cases = []
+
+mixed = path / "mixed"
+writer = FTWWriter(str(mixed), shard_limit=4096)
+writer.add_tensor("bank", torch.ones(1), kind="experts_bank")
+writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((mixed, 1))
+
+nondivisible = path / "nondivisible"
+writer = FTWWriter(str(nondivisible), shard_limit=4096)
+writer.add_tensor("bank", torch.ones(3), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((nondivisible, 2))
+
+missing_layer = path / "missing-layer"
+writer = FTWWriter(str(missing_layer), shard_limit=4096)
+writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((missing_layer, 2))
+
+for checkpoint, layers in schema_cases:
+    try:
+        load_ftw_banks(str(checkpoint), num_layers=layers)
+    except FTWFormatError:
+        pass
+    else:
+        raise SystemExit("malformed expert-bank schema was accepted under -O")
+'''
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.pathsep.join(filter(None, [
+        str(repo_root / "python"), env.get("PYTHONPATH", ""),
+    ]))
+    result = subprocess.run(
+        [sys.executable, "-O", "-c", script, str(tmp_path)],
+        cwd=repo_root,
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=30,
+        check=False,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr

From 20f3b237063c4cf20592c935a55c691df6f50bb4 Mon Sep 17 00:00:00 2001
From: alvarorsouza-arch <alvaro.r.souza@gmail.com>
Date: Sat, 19 Sep 2026 07:54:09 +0000
Subject: [PATCH 045/249] fix(kvcache): fp8-aware write-through mirror +
 batch-safe selection cap

- mirror_store_quant (new Triton kernel, graph-safe UVA): the write-through
  mirror now copies e4m3 codes + fp32 scales read from the pool AFTER
  quantization (the pool is the source of truth). The bf16-era mirror
  corrupted fp8 pages silently on rehydrate.
- mirror_scales host bank + fused-copy descriptor extended with the scale
  banks (per-bank row sizes).
- set_select_width takes max_bs: the whole batch selection must fit the
  physical pool by construction (fixes assert-crash on multi-row decode with
  small pools).
- measured: 299,648 logical tokens with kv-cache-dtype fp8 + kv-host-pages,
  needle at 253k tokens answered correctly (~1026 t/s), quality suite
  identical to baseline.
---
 python/freetoken/attention/qsa_sparse.py      |   2 +-
 python/freetoken/kernel/triton/qsa/offload.py | 103 +++++++------
 python/freetoken/kvcache/kv_host_offload.py   | 141 +++++++++++++-----
 3 files changed, 161 insertions(+), 85 deletions(-)

diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index de48768c8c..ae36fc73fb 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -318,7 +318,7 @@ def qsa_forward(
         self._update_index_cache(index, md, slot)
         if off is not None:
             # Write-through: the pinned host mirror (the backing store) sees every token.
-            off.mirror_store(k, v, batch.out_loc, layer_id)
+            off.mirror_store(k, v, batch.out_loc, layer_id, getattr(md, "out_loc_gpu", None))
         indices = self._select(index, md, slot)
         if off is not None:
             return self._attend_offloaded(qsa_sparse_paged_attention, q, indices, md, layer_id)
diff --git a/python/freetoken/kernel/triton/qsa/offload.py b/python/freetoken/kernel/triton/qsa/offload.py
index 9b669bf5a7..12223f48b8 100644
--- a/python/freetoken/kernel/triton/qsa/offload.py
+++ b/python/freetoken/kernel/triton/qsa/offload.py
@@ -67,7 +67,7 @@ def _compact_sel_kernel(
     row = tl.where(lanes < count, seen, first)
     tl.store(out + r * maxp + lanes, row, mask=lanes < maxp)
     tl.store(counts + r, count)
-    tl.atomic_add(trunc, ndrop)
+    tl.atomic_max(trunc, ndrop)
 
 
 def compact_selected_pages(
@@ -84,8 +84,8 @@ def compact_selected_pages(
 
     ``out`` is ``[rows, maxp]`` int32; rows with fewer hits are padded with the row's first
     page (duplicates collapse in lru_ensure), or ``dummy_page`` when the row selected nothing.
-    ``trunc`` (optional [1] int32, device) accumulates the TOTAL number of pages dropped
-    across all rows because they exceeded ``maxp`` -- 0 means no truncation happened. ``counts`` (optional
+    ``trunc`` (optional [1] int32, device) accumulates how many pages were dropped per row
+    because they exceeded ``maxp`` -- 0 means no truncation happened. ``counts`` (optional
     [rows] int32, device) receives each row's stored distinct-page count.
     """
     rows, sel = indices.shape
@@ -185,51 +185,60 @@ def mirror_store(
         )
 
 
-
-
-@triton.jit
-def _write_pages_stride_kernel(out_loc, out, page_size: tl.constexpr):
-    """One page per page_size tokens: covers every page of a run but (possibly) its last."""
-    i = tl.program_id(0)
-    tl.store(out + i, tl.load(out_loc + i * page_size) // page_size)
-
-
-@triton.jit
-def _write_pages_pad_kernel(out, nstride):
-    """Pad the boundary slots with out[0]: a real logical page, so a run with fewer
-    boundaries than `extra` leaves duplicates behind instead of garbage. Duplicates
-    collapse inside lru_ensure."""
-    i = tl.program_id(0)
-    tl.store(out + nstride + i, tl.load(out + 0))
-
-
 @triton.jit
-def _write_pages_runs_kernel(out_loc, token_to_req, out, counter, n, nstride, extra):
-    """The last token of every request run (a prefill batch concatenates runs): its page is
-    the one the stride can miss. Bounded by `extra`, so the buffer stays fixed-shape."""
-    i = tl.program_id(0)
-    cur = tl.load(token_to_req + i)
-    nxt = tl.load(token_to_req + i + 1, mask=(i + 1 < n), other=-1)
-    if (i + 1 >= n) or (nxt != cur):
-        slot = tl.atomic_add(counter, 1)
-        if slot < extra:
-            tl.store(out + nstride + slot, tl.load(out_loc + i) // page_size)
-
-
-def write_pages_prefill(out_loc, token_to_req, out, counter, nstride, extra, page_size):
-    """Prefill write-page set: stride samples + the last token of each request run.
-
-    Fully device-side (no host sync): a sync here stalls the MoE prefetch pipeline that
-    `--moe-prefill-overlap` depends on, which costs more than the sampling itself.
-    """
-    n = out_loc.numel()
-    if not n:
+def _mirror_store_quant_kernel(
+    src_ptrs,     # [4L] int64 GPU: por camada, bases pool (kcode, kscale, vcode, vscale)
+    host_ptrs,    # [4L] int64 GPU: por camada, bases host UVA (kcode, kscale, vcode, vscale)
+    out_phys,     # [T] int32: slots físicos (pool) de cada token
+    out_log,      # [T] int32: slots lógicos (espelho) de cada token
+    layer,
+    code_bytes: tl.constexpr, scale_bytes: tl.constexpr,
+    BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr,
+):
+    """fp8 write-through: lê códigos+scales do pool (pós-quantização) e espelha no host."""
+    t = tl.program_id(0)
+    phys = tl.load(out_phys + t).to(tl.int64)
+    logi = tl.load(out_log + t).to(tl.int64)
+    oc = tl.arange(0, BLOCK_C)
+    mc = oc < code_bytes
+    os_ = tl.arange(0, BLOCK_S)
+    ms = os_ < scale_bytes
+    kc_h = tl.load(host_ptrs + 4 * layer).to(tl.pointer_type(tl.uint8))
+    ks_h = tl.load(host_ptrs + 4 * layer + 1).to(tl.pointer_type(tl.uint8))
+    vc_h = tl.load(host_ptrs + 4 * layer + 2).to(tl.pointer_type(tl.uint8))
+    vs_h = tl.load(host_ptrs + 4 * layer + 3).to(tl.pointer_type(tl.uint8))
+    kc_s = tl.load(src_ptrs + 4 * layer).to(tl.pointer_type(tl.uint8))
+    ks_s = tl.load(src_ptrs + 4 * layer + 1).to(tl.pointer_type(tl.uint8))
+    vc_s = tl.load(src_ptrs + 4 * layer + 2).to(tl.pointer_type(tl.uint8))
+    vs_s = tl.load(src_ptrs + 4 * layer + 3).to(tl.pointer_type(tl.uint8))
+    tl.store(kc_h + logi * code_bytes + oc, tl.load(kc_s + phys * code_bytes + oc, mask=mc), mask=mc)
+    tl.store(vc_h + logi * code_bytes + oc, tl.load(vc_s + phys * code_bytes + oc, mask=mc), mask=mc)
+    tl.store(ks_h + logi * scale_bytes + os_, tl.load(ks_s + phys * scale_bytes + os_, mask=ms), mask=ms)
+    tl.store(vs_h + logi * scale_bytes + os_, tl.load(vs_s + phys * scale_bytes + os_, mask=ms), mask=ms)
+
+
+def mirror_store_quant(
+    kc_pool: torch.Tensor,   # [P*page, H*D] e4m3 (view)
+    ks_pool: torch.Tensor,   # [P*page, H] fp32
+    vc_pool: torch.Tensor,
+    vs_pool: torch.Tensor,
+    out_phys: torch.Tensor,  # [T] int32 slots físicos
+    out_log: torch.Tensor,   # [T] int32 slots lógicos
+    src_ptrs: torch.Tensor,  # [4L] int64 GPU
+    host_ptrs: torch.Tensor, # [4L] int64 GPU
+    dense_layer: int,
+) -> None:
+    """Espelha códigos+scales de uma camada fp8 pro host (graph-safe: UVA stores)."""
+    t = out_phys.numel()
+    if not t:
         return
-    _write_pages_stride_kernel[(nstride,)](out_loc, out, page_size)
-    if extra:
-        _write_pages_pad_kernel[(extra,)](out, nstride)
-        counter.zero_()
-        _write_pages_runs_kernel[(n,)](out_loc, token_to_req, out, counter, n, nstride, extra)
+    code_bytes = kc_pool.shape[1] * kc_pool.element_size()
+    scale_bytes = ks_pool.shape[1] * ks_pool.element_size()
+    _mirror_store_quant_kernel[(t,)](
+        src_ptrs, host_ptrs, out_phys, out_log, dense_layer,
+        code_bytes, scale_bytes,
+        triton.next_power_of_2(code_bytes), triton.next_power_of_2(scale_bytes),
+    )
 
 
-__all__ = ["write_pages", "compact_selected_pages", "translate_table", "translate_slots", "mirror_store"]
+__all__ = ["write_pages", "compact_selected_pages", "translate_table", "translate_slots", "mirror_store", "mirror_store_quant"]
diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py
index 40e5940cad..dbb49f4814 100644
--- a/python/freetoken/kvcache/kv_host_offload.py
+++ b/python/freetoken/kvcache/kv_host_offload.py
@@ -59,6 +59,21 @@ def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None:
         self.mirror.tensor.zero_()  # fresh/dummy pages read back as zeros, never stale bytes
         self.mirror.pin()
 
+        # fp8 pool (kv_quant="fp8", PR#354): the pool stores e4m3 CODES + fp32 per-token
+        # scales. The mirror must carry BOTH or rehydrated pages decode garbage (the read
+        # path below restores codes+scales into the pool; a bf16-era mirror of raw k/v
+        # would silently corrupt). The raw-k/v Triton mirror_store only runs on bf16 pools.
+        self._scale_buf = getattr(pool, "_scale_buffer", None)  # [2, L, P*page, H] | None
+        self.mirror_scales: HostBank | None = None
+        if self._scale_buf is not None:
+            scale_heads = self._scale_buf.shape[-1]
+            self.mirror_scales = HostBank(
+                (2, num_layers, (num_logical_pages + 1) * page_size, scale_heads),
+                self._scale_buf.dtype,
+            )
+            self.mirror_scales.tensor.zero_()
+            self.mirror_scales.pin()
+
         # LRU slot maps (lru_ensure's id space is logical pages, slot space is buffer pages).
         self.phys_of = torch.full((num_logical_pages + 1,), -1, dtype=torch.int32, device=device)
         self.logical_of = torch.full((num_slots,), -1, dtype=torch.int32, device=device)
@@ -79,10 +94,6 @@ def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None:
         self._query_buf: torch.Tensor | None = None
         self._eff_buf: torch.Tensor | None = None
         self._trunc_eager = torch.zeros(1, dtype=torch.int32, device=device)
-        # Prefill write-page scratch (grow-only) + counter for the run-aware sampler.
-        self.max_bs = 1
-        self._wp_buf: torch.Tensor | None = None
-        self._wp_count = torch.zeros((), dtype=torch.int32, device=device)
 
         from freetoken.kernel.pinned import device_ptr
 
@@ -103,7 +114,42 @@ def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None:
         self._src_ptrs = torch.tensor(src_ptrs, dtype=torch.int64, device=device)
         self._feats = torch.full((2 * num_layers,), row_bytes, dtype=torch.int64, device=device)
         self._host_ptrs = torch.tensor(host_ptrs, dtype=torch.int64, device=device)
-        self.mirror_bytes = self.mirror.nbytes
+        # fp8: estende o descritor com os bancos de scales (mesma lista de páginas —
+        # o fused copy roda uma linha por banco por página; feats heterogêneos ok).
+        # E prepara os ponteiros do mirror_store quantizado (graph-safe, UVA).
+        self._quant_src_ptrs = None
+        self._quant_host_ptrs = None
+        if self.mirror_scales is not None:
+            sc = self._scale_buf  # [2, L, P*page, H] -> por página: [page, H]
+            scale_row = page_size * sc.shape[-1] * sc.element_size()
+            assert scale_row % 16 == 0
+            qsrc, qhost = [], []
+            for kv in range(2):
+                for layer in range(num_layers):
+                    pool_v = sc[kv, layer].view(-1, page_size, sc.shape[-1])
+                    mir_v = self.mirror_scales.tensor[kv, layer].view(
+                        -1, page_size, sc.shape[-1])
+                    assert pool_v.is_contiguous() and mir_v.is_contiguous()
+                    dst_ptrs.append(pool_v.data_ptr())
+                    src_ptrs.append(device_ptr(mir_v))
+            for layer in range(num_layers):
+                # ordem por camada: (kcode, kscale, vcode, vscale) — casa com o kernel
+                qsrc += [buf[0, layer].data_ptr(), sc[0, layer].data_ptr(),
+                         buf[1, layer].data_ptr(), sc[1, layer].data_ptr()]
+                qhost += [device_ptr(mirror[0, layer]),
+                          device_ptr(self.mirror_scales.tensor[0, layer]),
+                          device_ptr(mirror[1, layer]),
+                          device_ptr(self.mirror_scales.tensor[1, layer])]
+            self._quant_src_ptrs = torch.tensor(qsrc, dtype=torch.int64, device=device)
+            self._quant_host_ptrs = torch.tensor(qhost, dtype=torch.int64, device=device)
+            self._dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=device)
+            self._src_ptrs = torch.tensor(src_ptrs, dtype=torch.int64, device=device)
+            self._feats = torch.cat([
+                self._feats,
+                torch.full((2 * num_layers,), scale_row, dtype=torch.int64, device=device),
+            ])
+        self.mirror_bytes = self.mirror.nbytes + (
+            self.mirror_scales.nbytes if self.mirror_scales is not None else 0)
 
     # ----- geometry ---------------------------------------------------------
     def set_select_width(self, select_width: int, block_topk: int = 0,
@@ -111,14 +157,11 @@ def set_select_width(self, select_width: int, block_topk: int = 0,
         # The sparse selection is block_topk short runs of consecutive tokens SCATTERED
         # across the sequence -- each run can touch its own page (two if it straddles a
         # boundary). Contiguity (select_width // page_size) badly underestimates this.
-        # Hard cap: the decode query is `max_bs * (1 + max_sel_pages)` wide and must fit
-        # the physical pool (lru_ensure requires |distinct| <= num_cached), so the
-        # per-row width is the pool split across the concurrent requests, not the whole
-        # pool. Concurrency and per-row selection width trade off here by construction;
-        # exceeding the resulting width now fails loudly (see `trunc_count`).
+        # Hard cap: the WHOLE decode batch's query (max_bs x (1 + sel)) must fit the pool,
+        # or mid-ensure evictions corrupt the attend. Divide by the batch, not the row.
         worst = 2 * (block_topk or (select_width // 4)) + 8
-        per_row = max(16, (self.num_slots - 73) // max(1, max_bs))
-        self.max_sel_pages = min(worst, per_row)
+        per_batch = max(16, (self.num_slots - 73) // max(1, max_bs))
+        self.max_sel_pages = min(worst, per_batch)
 
     def prefill_fetch_rows(self, n_write_pages: int) -> int:
         """Query rows per prefill attend sub-chunk: the write pages plus this many rows' worth
@@ -135,7 +178,6 @@ def prefill_fetch_rows(self, n_write_pages: int) -> int:
     # ----- graph buffers ------------------------------------------------------
     def init_graph_buffers(self, max_bs: int, width_pages: int, select_width: int | None = None) -> None:
         """Static decode buffers (capture-safe addresses); sliced to the live bs per step."""
-        self.max_bs = max(1, max_bs)
         if select_width is not None:
             self.set_select_width(select_width, max_bs=max_bs)
         assert self.max_sel_pages > 0, "set_select_width must run before init_graph_buffers"
@@ -152,6 +194,22 @@ def init_graph_buffers(self, max_bs: int, width_pages: int, select_width: int |
             "trunc": torch.zeros(1, dtype=torch.int32, device=d),
         }
 
+    def invalidate_pages(self, logical_pages: torch.Tensor) -> None:
+        """Marca páginas lógicas como frias (phys_of=-1), desfazendo bindings defasados.
+        Necessário quando o espelho foi reescrito por fora (reidratação host->espelho):
+        sem isso, um binding antigo apontaria pra um slot físico com conteúdo de outro
+        inquilino e o fetch seria pulado. Só desfaz se o slot ainda aponta pra esta página."""
+        lp = logical_pages.long().to(self.device)
+        old = self.phys_of[lp]
+        mask = old >= 0
+        if not bool(mask.any()):
+            return
+        old_slots = old[mask].long()
+        cur = self.logical_of[old_slots]
+        still = cur == lp[mask].long()
+        self.logical_of[old_slots[still]] = -1
+        self.phys_of[lp[mask]] = -1
+
     def trunc_count(self) -> int:
         """Dropped selection pages accumulated since the last read (0 = no truncation)."""
         n = int(self._trunc_eager.item())
@@ -209,11 +267,10 @@ def _ensure(self, query: torch.Tensor) -> None:
             )
 
     # ----- write path ---------------------------------------------------------
-    def ensure_write_pages(self, out_loc: torch.Tensor, is_decode: bool,
-                           token_to_req: torch.Tensor | None = None) -> torch.Tensor:
+    def ensure_write_pages(self, out_loc: torch.Tensor, is_decode: bool) -> torch.Tensor:
         """Make every page ``out_loc`` writes to resident; returns the write-page query
         (the fetch path re-includes it so a write page can never be evicted mid-forward)."""
-        from freetoken.kernel.triton.qsa.offload import write_pages, write_pages_prefill
+        from freetoken.kernel.triton.qsa.offload import write_pages
 
         if is_decode:
             n = out_loc.numel()
@@ -221,23 +278,9 @@ def ensure_write_pages(self, out_loc: torch.Tensor, is_decode: bool,
             write_pages(out_loc, wp, self.page_size)
             self._ensure(wp)
             return wp
-        # A prefill batch concatenates several requests, and a plain every-page_size-th
-        # stride lands at an arbitrary offset inside each later request: its LAST page is
-        # skipped whenever `len % page_size <= offset`, and store_kv would then write to a
-        # physical slot that was never ensured. The stride provably covers every page but
-        # the last of each run, so add the final token of each run (one page per
-        # concurrent request). All device-side: no host sync.
-        n = out_loc.numel()
-        if not n:
-            return out_loc
-        nstride = (n + self.page_size - 1) // self.page_size
-        extra = max(1, self.max_bs)
-        if self._wp_buf is None or self._wp_buf.numel() < nstride + extra:
-            self._wp_buf = torch.empty(max(nstride + extra, 1024), dtype=torch.int32,
-                                       device=self.device)
-        pages = self._wp_buf[: nstride + extra]
-        write_pages_prefill(out_loc, token_to_req, pages, self._wp_count,
-                            nstride, extra, self.page_size)
+        # The chunk's tokens are position-contiguous, so every 64th slot names a written
+        # page exactly once -- no torch.unique (a GPU sync) needed.
+        pages = (out_loc[:: self.page_size] // self.page_size).contiguous()
         self._ensure(pages)
         return pages
 
@@ -253,10 +296,34 @@ def translate_slots(self, out_loc: torch.Tensor, is_decode: bool) -> torch.Tenso
         translate_slots(out_loc, self.phys_of, out, self.page_size)
         return out
 
-    def mirror_store(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int) -> None:
-        from freetoken.kernel.triton.qsa.offload import mirror_store
-
-        mirror_store(k, v, out_loc, self._host_ptrs, self.pool._dense(layer_id))
+    def mirror_store(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor,
+                     layer_id: int, out_loc_gpu: torch.Tensor | None = None) -> None:
+        if self.mirror_scales is None:
+            from freetoken.kernel.triton.qsa.offload import mirror_store
+
+            mirror_store(k, v, out_loc, self._host_ptrs, self.pool._dense(layer_id))
+            return
+        # fp8 pool: o pool JÁ quantizou (store_kv rodou antes). Espelhamos codes+scales
+        # LIDOS DO POOL (a única fonte da verdade pós-quantização) — nunca re-quantizar.
+        assert out_loc_gpu is not None, "fp8 mirror_store precisa dos slots físicos"
+        self._mirror_store_fp8(out_loc, out_loc_gpu, self.pool._dense(layer_id))
+
+    def _mirror_store_fp8(self, out_loc_logical: torch.Tensor, out_loc_phys: torch.Tensor,
+                          dense_layer: int) -> None:
+        """codes+scales do pool (slots físicos, pós-quantização) -> espelho (slots lógicos).
+        Graph-safe: kernel Triton com stores UVA, nada de .cpu() no caminho."""
+        from freetoken.kernel.triton.qsa.offload import mirror_store_quant
+
+        H, D = self.mirror.tensor.shape[4], self.mirror.tensor.shape[5]
+        scH = self._scale_buf.shape[-1]
+        mirror_store_quant(
+            self.pool._kv_buffer[0, dense_layer].view(-1, H * D),
+            self._scale_buf[0, dense_layer],
+            self.pool._kv_buffer[1, dense_layer].view(-1, H * D),
+            self._scale_buf[1, dense_layer],
+            out_loc_phys, out_loc_logical,
+            self._quant_src_ptrs, self._quant_host_ptrs, dense_layer,
+        )
 
     # ----- read path ------------------------------------------------------------
     def compact_all(self, indices: torch.Tensor, token_to_req: torch.Tensor,

From 5046e4a00f94e5dea3407072241839a258a0510d Mon Sep 17 00:00:00 2001
From: alvarorsouza-arch <alvaro.r.souza@gmail.com>
Date: Sat, 19 Sep 2026 07:55:45 +0000
Subject: [PATCH 046/249] feat(kvcache): host-resident prefix tier for hybrid
 GDN models

Prefix reuse on hybrid GDN models needs KV pages AND a GDN state snapshot
alive at the match point; the small global snapshot pool gets drained by
interleaved auxiliary calls (agent title generators) and every match
truncates to zero. Measured on real WorkBuddy agent turns: TTFT 20-150s
recompute per message.

This adds a host-resident tier: evicted GDN snapshots and KV prefix pages
move to pinned RAM instead of dying. Matches rehydrate from host in ~4s
(24k-token prefix) instead of a full recompute.

- HybridRadixCache hooks: on_evict_node / on_match_dead_snapshot; evicted
  leaves become tombstones instead of cascade-freeing KV
- HostPrefixTier: pinned GDN snapshot bank + cold-prefix index with its own
  KV slot space (logical ids get reused by the scheduler, so content is
  copied D2H on eviction, never referenced by id)
- Admission rehydration; with KVHostOffloader active the copy is
  host-to-host into its mirror (no bulk H2D; lazy fetch covers residency)
- Eviction copies run in 128-page slices (staging the whole prefix OOMed
  exactly under memory pressure, which is when eviction runs)

Measured: churn test cached-token 0 -> 23,936 (16.5s -> 3.8s); total pool
flush rehydrates at 24,000 cached tokens in 4.4s; quality suite identical
(39/40 QA, 30/30 cloze); decode steady 33-37 t/s.

All off by default (FT_GDN_HOST_TIER / FT_PREFIX_HOST envs).
---
 python/freetoken/engine/config.py             |   9 +-
 python/freetoken/kvcache/host_prefix_tier.py  | 318 ++++++++++++++++++
 .../freetoken/kvcache/hybrid_radix_cache.py   |  43 ++-
 python/freetoken/scheduler/cache.py           | 131 ++++++++
 tests/kvcache/test_host_prefix_tier.py        |  60 ++++
 tests/kvcache/test_tier3.py                   |  85 +++++
 tests/kvcache/test_tree_hooks.py              |  77 +++++
 7 files changed, 720 insertions(+), 3 deletions(-)
 create mode 100644 python/freetoken/kvcache/host_prefix_tier.py
 create mode 100644 tests/kvcache/test_host_prefix_tier.py
 create mode 100644 tests/kvcache/test_tier3.py
 create mode 100644 tests/kvcache/test_tree_hooks.py

diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 5df70492bb..5f0a0314d2 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import copy
+import os
 from dataclasses import dataclass, field, replace
 from functools import cached_property
 from typing import TYPE_CHECKING, List
@@ -66,11 +67,17 @@ class EngineConfig:
     cuda_graph_bs: List[int] | None = None
     cuda_graph_max_bs: int | None = None
     page_size: int = 1
+    # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3
+    # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens
+    # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by
+    # _adjust_config, which also refuses it on a pool family or attention backend that
+    # cannot read the scales.
+    kv_quant: str = "none"
     memory_ratio: float = 0.9
     # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse);
     # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as
     # ceil(ratio * max_running_req) extra slots.
-    linear_state_cache_ratio: float = 2.0
+    linear_state_cache_ratio: float = float(os.environ.get("FT_LINEAR_STATE_CACHE_RATIO", "2.0"))
     # Window/full ratio for the SWA radix cache (`--cache-type radix` on SWA models) and the DSV4
     # window tier: the DEFAULT window-pool size = max(working-set floor, ratio x full-pool tokens).
     # < 1.0 trades retained window-prefix capacity for memory savings; must be in (0, 1]. It is the
diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
new file mode 100644
index 0000000000..0265b126ef
--- /dev/null
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -0,0 +1,318 @@
+"""Host-side cold-prefix tier (Nível 3 — Degrau 1: snapshots GDN; Degrau 3: KV de prefixo).
+
+Dois artefatos em RAM pinned:
+1. um banco de snapshots GDN (write na evicção do tree, read no match/admissão);
+2. um índice de prefixos frios: token ids + slots KV host + slot GDN host.
+
+Degrau 3: o KV vive num **slot space host separado** (banco pinned [2, L, S, page, H, D] de
+códigos + scales, mesma geometria do pool GPU). Ids lógicos de página são reutilizados pelo
+scheduler, então o conteúdo é COPIADO (D2H) na hora do evict do tree — nunca referenciado por
+id lógico depois. A reidratação aloca páginas novas no pool e copia H2D de volta.
+
+Ver IMPLEMENTACAO-NIVEL3.md para o design completo.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import torch
+
+from freetoken.moe.host_banks import HostBank
+
+
+class HostMatch:
+    __slots__ = ("host_len", "kv_pages", "gdn_slot")
+
+    def __init__(self, host_len: int, kv_pages: torch.Tensor | None, gdn_slot: int | None) -> None:
+        self.host_len = host_len            # tokens cobertos (page-aligned)
+        self.kv_pages = kv_pages            # int32 CPU: slots KV do tier (não ids lógicos!)
+        self.gdn_slot = gdn_slot            # slot host do snapshot GDN (None = sem snapshot)
+
+
+@dataclass
+class _Entry:
+    key_tokens: torch.Tensor                      # CPU, o prefixo (page-aligned)
+    kv_slots: torch.Tensor | None = None          # int32 CPU: um slot host por página
+    gdn_slot: int | None = None                   # snapshot GDN host do boundary
+
+
+class HostPrefixTier:
+    """Índice + storage RAM pra prefixos evictados do radix tree GPU.
+
+    O índice é LRU por entrada (hit re-marca); quando a RAM do tier enche (KV) ou os slots
+    GDN acabam, drop-oldest libera a entrada mais fria inteira (KV slots + slot GDN).
+    As páginas KV são copiadas para slots host próprios na evicção — o conteúdo sobrevive à
+    reutilização dos ids lógicos pelo scheduler.
+
+    Geometria GDN: os tensores do LinearStatePool são [layers, slots, ...]; cada banco host
+    é [host_slots, layers, ...] — a cópia de um slot é [layers, ...] <-> [layers, ...].
+    Geometria KV: pool._kv_buffer [2, L, P, page, H, D] (códigos e4m3 no pool fp8) +
+    pool._scale_buffer [2, L, P*page, H] fp32; os banks host espelham isso com S host slots.
+    """
+
+    def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
+                 kv_budget_bytes: int = 0, device: torch.device | str = "cuda") -> None:
+        self.device = torch.device(device)
+        self.pool = state_pool
+        tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states}
+        tensors.update(state_pool.slot_states)
+        self._order = list(tensors)
+        self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {}
+        for name, t in tensors.items():
+            # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest]
+            bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype)
+            if self.device.type == "cuda" and torch.cuda.is_available():
+                bank.pin()
+            self._banks[name] = (t, bank)
+        self.num_slots = gdn_slots_host
+        self._free = list(range(gdn_slots_host))
+        self._in_use: set[int] = set()
+        self._index: OrderedDict[bytes, _Entry] = OrderedDict()
+        # key: bytes dos token ids page-aligned -> _Entry
+
+        # ---- KV page store (Degrau 3) ----
+        self.kv_pool = kv_pool
+        self.kv_page_bytes = 0
+        self._kv_codes: HostBank | None = None
+        self._kv_scales: HostBank | None = None
+        self._kv_free: list[int] = []
+        if kv_pool is not None and kv_budget_bytes > 0:
+            codes = kv_pool._kv_buffer            # [2, L, P, page, H, D]
+            scales = kv_pool._scale_buffer        # [2, L, P*page, H] | None (fp8)
+            _, layers, _, page_size, heads, dim = codes.shape
+            page_bytes = 2 * layers * page_size * heads * dim * codes.element_size()
+            if scales is not None:
+                page_bytes += 2 * layers * page_size * heads * scales.element_size()
+            n = max(1, kv_budget_bytes // page_bytes)
+            self._kv_codes = HostBank((2, layers, n, page_size, heads, dim), codes.dtype)
+            if scales is not None:
+                self._kv_scales = HostBank((2, layers, n, page_size, heads), scales.dtype)
+            if self.device.type == "cuda" and torch.cuda.is_available():
+                self._kv_codes.pin()
+                if self._kv_scales is not None:
+                    self._kv_scales.pin()
+            self._kv_free = list(range(n))
+            self.kv_page_bytes = page_bytes
+        self.copy_stream = (
+            torch.cuda.Stream(self.device) if self.device.type == "cuda" else None)
+
+    # ---------------- GDN snapshot storage ----------------
+    def write_gdn(self, gpu_slot: int) -> int:
+        """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot.
+        Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do
+        engine garante leitura antes de qualquer reescrita do slot (ver spec)."""
+        if not self._free:
+            raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)")
+        hs = self._free.pop()
+        self._in_use.add(hs)
+        if self.copy_stream is not None:
+            with torch.cuda.stream(self.copy_stream):
+                for name in self._order:
+                    t, bank = self._banks[name]
+                    bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True)
+            self.copy_stream.synchronize()
+        else:
+            for name in self._order:
+                t, bank = self._banks[name]
+                bank.tensor[hs].copy_(t[:, gpu_slot])
+        return hs
+
+    def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None:
+        """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller)."""
+        if host_slot not in self._in_use:
+            raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso")
+        if self.copy_stream is not None:
+            with torch.cuda.stream(self.copy_stream):
+                for name in self._order:
+                    t, bank = self._banks[name]
+                    t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True)
+            self.copy_stream.synchronize()
+        else:
+            for name in self._order:
+                t, bank = self._banks[name]
+                t[:, dst_gpu_slot].copy_(bank.tensor[host_slot])
+
+    def free_gdn(self, host_slot: int) -> None:
+        self._in_use.discard(host_slot)
+        self._free.append(host_slot)
+
+    # ---------------- KV page storage (Degrau 3) ----------------
+    @property
+    def has_kv(self) -> bool:
+        return self._kv_codes is not None
+
+    def write_kv(self, kv_indices: torch.Tensor) -> torch.Tensor | None:
+        """Copia as páginas (token slots do pool, page-aligned) para slots host livres.
+        Retorna os slots host (int32 CPU), ou None se o tier não tem KV store.
+        Deve rodar DURANTE o evict do tree, com os bytes ainda válidos no buffer GPU."""
+        if not self.has_kv:
+            return None
+        page_size = self._kv_codes.tensor.shape[3]
+        pages = (kv_indices[::page_size] // page_size).long().to(self.device)
+        n = pages.numel()
+        if n == 0:
+            return None
+        codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer
+        # Defesa: um slot fora do pool indicaria um bug de lifetime no tree. Pular a escrita
+        # (o prefixo simplesmente não vai pro tier) em vez de derrubar o scheduler.
+        max_page = int(pages.max().item())
+        if max_page >= codes.shape[2]:
+            from freetoken.utils import init_logger
+            init_logger(__name__).warning(
+                f"HostPrefixTier.write_kv: page {max_page} fora do pool "
+                f"({codes.shape[2]} páginas): n={n}, kv_slots min/max = "
+                f"{int(kv_indices.min())}/{int(kv_indices.max())}, prefixo de {kv_indices.numel()} tokens -- "
+                "escrita ignorada"
+            )
+            return None
+        while len(self._kv_free) < n:
+            self._drop_oldest()  # libera os slots da entrada mais fria (inclui GDN)
+        slots = torch.tensor([self._kv_free.pop() for _ in range(n)], dtype=torch.int64)
+        # NB: sem staging do prefixo inteiro — cada fatia calcula seus token-ids.
+
+        def _go():
+            # Chunked: NUNCA estagiar o prefixo inteiro (o evict acontece sob pressão de
+            # memória — um staging GB-sized explode a VRAM e derruba o scheduler).
+            CH = 128  # páginas por fatia (~100 MB de temp por fatia em fp8)
+            for i in range(0, n, CH):
+                p = pages[i : i + CH]
+                s = slots[i : i + CH]
+                staged = codes[:, :, p].contiguous()
+                self._kv_codes.tensor[:, :, s] = staged.cpu() if staged.is_cuda else staged
+                del staged
+            if scales is not None:
+                for i in range(0, n, CH):
+                    p = pages[i : i + CH]
+                    s = slots[i : i + CH]
+                    tk = (p.unsqueeze(1) * page_size
+                          + torch.arange(page_size, device=self.device)).flatten()
+                    staged_s = scales[:, :, tk].view(
+                        2, scales.shape[1], p.numel(), page_size, scales.shape[3]).contiguous()
+                    self._kv_scales.tensor[:, :, s] = (
+                        staged_s.cpu() if staged_s.is_cuda else staged_s)
+                    del staged_s
+
+        if self.copy_stream is not None:
+            with torch.cuda.stream(self.copy_stream):
+                _go()
+            self.copy_stream.synchronize()
+        else:
+            _go()
+        return slots.to(torch.int32).cpu()
+
+    def read_kv(self, host_slots: torch.Tensor, dst_pages: torch.Tensor) -> None:
+        """H2D dos slots host para as páginas dst (ids de página do pool, device)."""
+        assert self.has_kv
+        slots = host_slots.long()
+        pages = dst_pages.long().to(self.device)
+        codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer
+        page_size = self._kv_codes.tensor.shape[3]
+        n = slots.numel()
+
+        def _go():
+            CH = 128
+            for i in range(0, n, CH):
+                s = slots[i : i + CH]
+                p = pages[i : i + CH]
+                staged = self._kv_codes.tensor[:, :, s].to(self.device, non_blocking=True)
+                codes[:, :, p] = staged
+                del staged
+            if scales is not None:
+                for i in range(0, n, CH):
+                    s = slots[i : i + CH]
+                    p = pages[i : i + CH]
+                    src = self._kv_scales.tensor[:, :, s].reshape(
+                        2, scales.shape[1], p.numel() * page_size, scales.shape[3])
+                    tk = (p.unsqueeze(1) * page_size
+                          + torch.arange(page_size, device=self.device)).flatten()
+                    staged_s = src.to(self.device, non_blocking=True)
+                    scales[:, :, tk] = staged_s
+                    del staged_s
+
+        if self.copy_stream is not None:
+            with torch.cuda.stream(self.copy_stream):
+                _go()
+            self.copy_stream.synchronize()
+        else:
+            _go()
+
+    def read_kv_to_mirror(self, host_slots: torch.Tensor, mirror_codes,
+                          mirror_scales, logical_pages: torch.Tensor) -> None:
+        """H2H (pinned memcpy): slots KV do tier -> espelho do offloader nas páginas
+        lógicas dadas. Usado quando o offloader está ativo: reidratar NÃO sobe bytes pra
+        GPU — o fetch lazy do offloader cuida da residência na atenção."""
+        assert self.has_kv
+        s = host_slots.long()
+        lp = logical_pages.long().cpu()
+        page = self._kv_codes.tensor.shape[3]
+        mirror_codes[:, :, lp] = self._kv_codes.tensor[:, :, s]
+        if self._kv_scales is not None and mirror_scales is not None:
+            t_s = self._kv_scales.tensor.view(2, -1, self._kv_codes.tensor.shape[2], page,
+                                              self._kv_scales.tensor.shape[-1])
+            m_s = mirror_scales.view(2, -1, mirror_codes.shape[2], page,
+                                     self._kv_scales.tensor.shape[-1])
+            m_s[:, :, lp] = t_s[:, :, s]
+
+    def _drop_oldest(self) -> None:
+        if not self._index:
+            raise RuntimeError("HostPrefixTier: KV host store cheio e índice vazio")
+        _, ent = self._index.popitem(last=False)  # o mais frio
+        if ent.kv_slots is not None:
+            self._kv_free.extend(ent.kv_slots.tolist())
+            ent.kv_slots = None
+        if ent.gdn_slot is not None:
+            self.free_gdn(ent.gdn_slot)
+            ent.gdn_slot = None
+
+    # ---------------- índice de prefixos ----------------
+    @staticmethod
+    def _key_of(tokens: torch.Tensor) -> bytes:
+        return tokens.to(torch.int64).cpu().numpy().tobytes()
+
+    def put_prefix(self, key_tokens: torch.Tensor, kv_pages: torch.Tensor | None = None,
+                   host_gdn_slot: int | None = None) -> None:
+        """Registra (ou faz merge numa entrada existente) um prefixo frio. key_tokens deve
+        ser page-aligned (o tree já alinha). kv_pages = slots host do KV store."""
+        key = self._key_of(key_tokens)
+        ent = self._index.get(key)
+        if ent is None:
+            ent = _Entry(key_tokens.cpu())
+            self._index[key] = ent
+        else:
+            self._index.move_to_end(key)  # refresh LRU
+        if kv_pages is not None:
+            if ent.kv_slots is not None and ent.kv_slots is not kv_pages:
+                self._kv_free.extend(ent.kv_slots.tolist())  # substituição: libera os velhos
+            ent.kv_slots = kv_pages.to(torch.int32).cpu()
+        if host_gdn_slot is not None:
+            if ent.gdn_slot is not None and ent.gdn_slot != host_gdn_slot:
+                self.free_gdn(ent.gdn_slot)  # a cópia nova substitui a anterior
+            ent.gdn_slot = host_gdn_slot
+
+    def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | None:
+        """Maior prefixo page-aligned presente no índice (python/CPU, fora do hot path)."""
+        toks = input_ids.to(torch.int64).cpu()
+        n = (toks.numel() // page_size) * page_size
+        while n > 0:
+            key = self._key_of(toks[:n])
+            hit = self._index.get(key)
+            if hit is not None:
+                self._index.move_to_end(key)
+                return HostMatch(n, hit.kv_slots, hit.gdn_slot)
+            n -= page_size
+        return None
+
+    def evict_host(self, key_tokens: torch.Tensor) -> None:
+        """Remove uma entrada do índice e libera os slots host (KV + GDN)."""
+        ent = self._index.pop(self._key_of(key_tokens), None)
+        if ent is None:
+            return
+        if ent.kv_slots is not None:
+            self._kv_free.extend(ent.kv_slots.tolist())
+        if ent.gdn_slot is not None:
+            self.free_gdn(ent.gdn_slot)
+
+
+__all__ = ["HostPrefixTier", "HostMatch"]
diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py
index 1bd9e0b9cd..cb4a7f90b2 100644
--- a/python/freetoken/kvcache/hybrid_radix_cache.py
+++ b/python/freetoken/kvcache/hybrid_radix_cache.py
@@ -71,6 +71,16 @@ def __init__(self, device: torch.device, page_size: int) -> None:
         self.full_protected = 0
         self.mamba_evictable = 0     # number of live, unlocked snapshots
         self.mamba_protected = 0
+        # Host tier (FT_GDN_HOST_TIER): wired by CacheManager. ``on_evict_node`` fires with
+        # the node IMMEDIATELY before its snapshot is freed (the slot is still valid --
+        # enqueue the D2H before returning). ``on_match_dead_snapshot`` may rehydrate a
+        # tombstone during match: return a fresh GPU slot to attach, None to keep walking.
+        self.on_evict_node = None
+        self.on_match_dead_snapshot = None
+        # KV host tier (FT_PREFIX_HOST): fires for every node whose KV pages LEAVE the tree
+        # (evict_full's leaf + cascade-exposed tombstone parents) -- BEFORE the free, while
+        # the bytes are still valid in the GPU buffer (the D2H must happen inside the call).
+        self.on_evict_kv = None
 
     # ---------------------------------------------------------------- match / insert
     def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch:
@@ -83,6 +93,14 @@ def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch:
         while not cur.is_root():
             if cur.mamba_value is not None:
                 return HybridMatch(self._collect_kv(cur), end_len, cur.mamba_value, cur)
+            if self.on_match_dead_snapshot is not None:
+                slot = self.on_match_dead_snapshot(cur)
+                if slot is not None:
+                    # revive the tombstone: the rehydrated snapshot rides the tree again
+                    cur.mamba_value = slot
+                    if cur.mamba_ref_count == 0:
+                        self.mamba_evictable += 1
+                    return HybridMatch(self._collect_kv(cur), end_len, slot, cur)
             end_len -= cur.length
             cur = cur.parent
         return HybridMatch(self.empty, 0, None, self.root)
@@ -156,8 +174,12 @@ def evict_full(self, num_tokens: int) -> EvictResult:
             if node.ref_count != 0 or not node.is_leaf() or node.is_root():
                 continue
             freed += node.length
+            if self.on_evict_kv is not None:
+                self.on_evict_kv(node)
             kv.append(node.value)
             self.full_evictable -= node.length
+            if node.mamba_value is not None and self.on_evict_node is not None:
+                self.on_evict_node(node)
             self._free_node_mamba(node, mamba)
             parent, casc = self._cascade_tombstone_leaves(self._unlink(node), kv)
             freed += casc
@@ -177,14 +199,19 @@ def evict_mamba(self, num: int) -> EvictResult:
             node = heapq.heappop(cands)
             if node.mamba_value is None or node.mamba_ref_count != 0 or node.is_root():
                 continue
-            if node.is_leaf() and node.ref_count == 0:
+            if self.on_evict_node is not None:
+                self.on_evict_node(node)
+            if node.is_leaf() and node.ref_count == 0 and self.on_evict_node is None:
                 kv.append(node.value)
                 self.full_evictable -= node.length
                 self._free_node_mamba(node, mamba)
                 freed += 1
                 self._cascade_tombstone_leaves(self._unlink(node), kv)
             else:
-                self._free_node_mamba(node, mamba)  # tombstone internal (or locked-KV) node
+                # internal (or locked-KV) node -> tombstone; with a host tier wired, LEAF too:
+                # the snapshot is safely on the host, and keeping the leaf's KV in the tree
+                # is what lets a later match rehydrate the boundary instead of re-prefilling.
+                self._free_node_mamba(node, mamba)
                 freed += 1
         return EvictResult(torch.cat(kv) if kv else self.empty, mamba)
 
@@ -230,6 +257,8 @@ def _cascade_tombstone_leaves(self, parent: RadixTreeNode, kv_out: List[torch.Te
         freed = 0
         while (parent.mamba_value is None and parent.is_leaf()
                and parent.ref_count == 0 and not parent.is_root()):
+            if self.on_evict_kv is not None:
+                self.on_evict_kv(parent)
             kv_out.append(parent.value)
             self.full_evictable -= parent.length
             freed += parent.length
@@ -252,6 +281,16 @@ def _collect_kv(self, node: RadixTreeNode) -> torch.Tensor:
         vals.reverse()
         return torch.cat(vals) if vals else self.empty
 
+    def _collect_key(self, node: RadixTreeNode) -> torch.Tensor:
+        """The full token prefix ending at ``node`` (segment keys concatenated root->node)."""
+        keys: List[torch.Tensor] = []
+        n = node
+        while not n.is_root():
+            keys.append(n._key)
+            n = n.parent
+        keys.reverse()
+        return torch.cat(keys) if keys else torch.empty(0, dtype=torch.int32)
+
     def _leaves(self) -> List[RadixTreeNode]:
         out, stack = [], [self.root]
         while stack:
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 8f77be9330..3f833eaf18 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -56,6 +56,31 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ
         if swa_pool is not None:
             self.prefill_chunk_budget = getattr(swa_pool, "prefill_chunk_budget", None)
         self.prefix_cache = self._make_prefix_cache(device, page_size, type)
+        # Host prefix tier: FT_GDN_HOST_TIER=1 guarda snapshots GDN na RAM (Degrau 2);
+        # FT_PREFIX_HOST=1 + FT_PREFIX_HOST_GB adiciona o KV de prefixo num slot space host
+        # separado (Degrau 3 — o radix nunca perde o mapeamento). Hybrid path only.
+        self.host_tier = None
+        gdn_on = os.environ.get("FT_GDN_HOST_TIER", "0") == "1"
+        kv_gb = (float(os.environ.get("FT_PREFIX_HOST_GB", "16"))
+                 if os.environ.get("FT_PREFIX_HOST", "0") == "1" else 0)
+        if self.is_hybrid and (gdn_on or kv_gb > 0):
+            from freetoken.core import get_global_ctx
+            from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+            from freetoken.utils import init_logger
+
+            slots = int(os.environ.get("FT_GDN_HOST_SLOTS", "32"))
+            kv_pool = get_global_ctx().kv_cache if kv_gb > 0 else None
+            self.host_tier = HostPrefixTier(
+                linear_state_pool, kv_pool, gdn_slots_host=slots,
+                kv_budget_bytes=int(kv_gb * (1 << 30)), device=device)
+            self.prefix_cache.on_evict_node = self._on_evict_node_to_host
+            self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot
+            if kv_pool is not None:
+                self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host
+            init_logger(__name__).info(
+                f"host prefix tier ON: {slots} GDN slots"
+                + (f" + {kv_gb:.0f} GiB KV store" if kv_pool is not None else "")
+            )
         self.device = device
         self.num_pages = num_pages
         self.page_table = page_table
@@ -101,6 +126,16 @@ def match_req(self, req: PendingReq) -> MatchResult:
         if self.is_hybrid:
             from freetoken.kvcache.hybrid_radix_cache import HybridCacheHandle
             m = self.prefix_cache.match_prefix(ids)
+            # Host prefix tier (FT_PREFIX_HOST): a prefix evicted from the tree lives in RAM
+            # (KV bytes + GDN snapshot). If the tree match fell short, rehydrate the deeper
+            # host entry into fresh pages, re-insert it into the tree and re-match.
+            tier = self.host_tier
+            if tier is not None and tier.has_kv:
+                hit = tier.lookup(ids, self.page_size)
+                if (hit is not None and hit.kv_pages is not None and hit.gdn_slot is not None
+                        and hit.host_len > m.cached_len):
+                    self._rehydrate_host_prefix(ids, hit)
+                    m = self.prefix_cache.match_prefix(ids)
             return MatchResult(
                 HybridCacheHandle(m.cached_len, m.node, m.kv_indices), mamba_value=m.mamba_value)
         return self.prefix_cache.match_prefix(ids)
@@ -548,6 +583,96 @@ def check_integrity(self) -> None:
         if self.page_size > 1:
             assert torch.all(self.free_slots % self.page_size == 0)
 
+    # ----- host GDN-snapshot tier (FT_GDN_HOST_TIER) -------------------------------
+    def _on_evict_node_to_host(self, node) -> None:
+        """Tree hook: the node's GDN snapshot is about to be freed. Copy it to the host tier
+        and index the node's prefix (the GPU slot is freed only after we return)."""
+        tier = self.host_tier
+        prefix = self.prefix_cache._collect_key(node)
+        try:
+            hs = tier.write_gdn(node.mamba_value)
+        except RuntimeError:
+            # host tier full: drop the OLDEST indexed prefix and retry once
+            if not tier._index:
+                return
+            oldest = next(iter(tier._index))
+            tier.evict_host(tier._index[oldest].key_tokens)
+            try:
+                hs = tier.write_gdn(node.mamba_value)
+            except RuntimeError:
+                return
+        tier.put_prefix(prefix, None, hs)
+
+    def _on_evict_kv_to_host(self, node) -> None:
+        """Tree hook (FT_PREFIX_HOST): the node's KV pages are about to be freed. Copy the
+        WHOLE prefix's pages to the host store while the bytes are still valid (eviction is
+        bottom-up, so the ancestors' pages are still live), and index the prefix."""
+        tier = self.host_tier
+        prefix = self.prefix_cache._collect_key(node)
+        kv = self.prefix_cache._collect_kv(node)
+        slots = tier.write_kv(kv)
+        if slots is not None:
+            tier.put_prefix(prefix, slots, None)
+
+    def _rehydrate_snapshot(self, node) -> int | None:
+        """Tree hook (match path): the node is a snapshot tombstone. If the host tier has its
+        prefix, copy the snapshot back into a fresh pool slot and return the slot; None =
+        keep walking (the match truncates as before)."""
+        tier = self.host_tier
+        hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size)
+        if hit is None or hit.gdn_slot is None:
+            return None
+        pool = self.linear_state_pool
+        if pool.num_free_slots < 1:
+            # make room: evicting another snapshot writes IT to the host tier first
+            self.ensure_mamba_slots(1)
+            if pool.num_free_slots < 1:
+                return None
+        slot = pool.alloc(1)[0]
+        tier.read_gdn(hit.gdn_slot, slot)
+        return slot
+
+    def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None:
+        """FT_PREFIX_HOST: sobe um prefixo frio do tier pro tree+GPU -- aloca páginas novas
+        (pode evictar: a vítima é escrita no host primeiro), copia os bytes KV H2D, restaura
+        o snapshot GDN num slot novo e re-insere o prefixo no tree (o match seguinte acha).
+        """
+        tier = self.host_tier
+        n_pages = hit.host_len // self.page_size
+        kv_new = self._page_to_token(self._allocate(n_pages))
+        pool = self.linear_state_pool
+        if hit.gdn_slot is not None:
+            self.ensure_mamba_slots(1)
+            if pool.num_free_slots < 1:
+                self._free(kv_new)  # sem slot GDN não dá pra reidratar; devolve as páginas
+                return
+            slot = pool.alloc(1)[0]
+            tier.read_gdn(hit.gdn_slot, slot)
+        else:
+            self._free(kv_new)
+            return  # o modelo híbrido não resume sem snapshot GDN
+        # Com o KVHostOffloader ativo, as páginas recém-alocadas são LÓGICAS: subir bytes
+        # pra GPU em bulk estoura o pool físico quando o prefixo > pool (o fetch lazy já
+        # resolve residência na atenção). Copiamos host->host pro espelho e invalidamos
+        # bindings defasados, em vez de H2D.
+        from freetoken.core import get_global_ctx
+        off = getattr(get_global_ctx(), "kv_offloader", None)
+        new_pages = kv_new[:: self.page_size] // self.page_size
+        if off is not None:
+            tier.read_kv_to_mirror(hit.kv_pages, off.mirror.tensor,
+                                   off.mirror_scales.tensor if off.mirror_scales is not None else None,
+                                   new_pages)
+            off.invalidate_pages(new_pages)
+        else:
+            tier.read_kv(hit.kv_pages, new_pages)
+        prefix_len, mamba_exist = self.prefix_cache.insert(ids[: hit.host_len], kv_new, slot)
+        if prefix_len > 0:
+            # o tree já tinha [0, prefix_len) com as páginas antigas (dedup do insert) --
+            # as páginas frescas que aloquei pra essa parte voltam pro free list
+            self._free(kv_new[:prefix_len])
+        if mamba_exist:
+            pool.free([slot])  # o boundary já tinha snapshot vivo (dedup do insert)
+
     def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None:
         """Re-point the page table and reset page accounting + prefix cache IN PLACE.
 
@@ -561,6 +686,12 @@ def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None:
         self.page_table = page_table
         self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * self.page_size
         self.prefix_cache = self._make_prefix_cache(device, self.page_size, self.cache_type)
+        if self.host_tier is not None:
+            # the rebuilt tree gets the same hooks (the host index survives a rebuild)
+            self.prefix_cache.on_evict_node = self._on_evict_node_to_host
+            self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot
+            if self.host_tier.has_kv:
+                self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host
         # The discarded hybrid tree owned donated GDN-snapshot slots; rebuild is idle-only, so
         # reclaim the whole LinearStatePool free-list (else those slots leak -> admission hangs).
         if self.is_hybrid:
diff --git a/tests/kvcache/test_host_prefix_tier.py b/tests/kvcache/test_host_prefix_tier.py
new file mode 100644
index 0000000000..fdb552e8e0
--- /dev/null
+++ b/tests/kvcache/test_host_prefix_tier.py
@@ -0,0 +1,60 @@
+"""Degrau 1 — teste standalone do HostPrefixTier (roda sem servidor; GPU tiny)."""
+import torch
+from types import SimpleNamespace
+
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+
+def main():
+    import os
+    dev = "cuda" if (os.environ.get("TIER_TEST_CPU") is None and torch.cuda.is_available()) else "cpu"
+    # pool fake: mesmas geometrias, mini
+    L, S = 4, 3   # 4 "layers", 3 slots GPU
+    pool = SimpleNamespace(
+        conv_states=torch.randn(L, S, 8, 3, dtype=torch.bfloat16, device=dev),
+        recurrent_states=torch.randn(L, S, 2, 4, 4, dtype=torch.float32, device=dev),
+        slot_states={"extra": torch.randn(L, S, 5, dtype=torch.bfloat16, device=dev)},
+    )
+    tier = HostPrefixTier(pool, gdn_slots_host=4, device=dev)
+
+    # 1) round-trip D2H->H2D preserva bytes
+    original = [t.clone() for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"])]
+    hs = tier.write_gdn(1)                    # snapshot do slot 1 -> host
+    # suja o slot GPU pra provar que o restore vem do host
+    for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"]):
+        t[:, 1].zero_()
+    tier.read_gdn(hs, 2)                      # host -> slot 2
+    ok = all(torch.equal(t[:, 1].cpu(), u[:, 2].cpu())
+             for t, u in zip(original, (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"])))
+    assert ok, "snapshot divergiu no round-trip!"
+    print("PASS 1: snapshot GDN round-trip D2H->H2D bit-exato")
+
+    # 2) índice: put + lookup longest-prefix page-aligned
+    toks = torch.arange(64 * 5, dtype=torch.int64)      # 5 páginas
+    pages = torch.tensor([10, 11, 12, 13, 14], dtype=torch.int32)
+    tier.put_prefix(toks[: 64 * 4], pages[:4], hs)       # registro 4 páginas
+    m = tier.lookup(torch.arange(64 * 5, dtype=torch.int64))  # query 5 páginas
+    assert m is not None and m.host_len == 64 * 4 and m.gdn_slot == hs
+    assert torch.equal(m.kv_pages, pages[:4])
+    print("PASS 2: lookup acha o maior prefixo page-aligned")
+
+    # 3) miss quando nada bate
+    assert tier.lookup(torch.full((64 * 3,), 999999, dtype=torch.int64)) is None
+    print("PASS 3: lookup retorna None em miss total")
+
+    # 4) evict_host libera o slot GDN host
+    tier.evict_host(toks[: 64 * 4])
+    assert hs in tier._free
+    print("PASS 4: evict_host libera slot host")
+
+    # 5) proteção: tier novo com 2 slots, 3ª escrita levanta erro claro
+    tier2 = HostPrefixTier(pool, gdn_slots_host=2, device=dev)
+    tier2.write_gdn(0); tier2.write_gdn(0)
+    try:
+        tier2.write_gdn(0)
+        raise SystemExit("ERA PRA TER FALHADO")
+    except RuntimeError:
+        print("PASS 5: capacidade cheia levanta RuntimeError claro")
+
+    print("\nTODOS OS TESTES PASSARAM ✓")
+
+main()
diff --git a/tests/kvcache/test_tier3.py b/tests/kvcache/test_tier3.py
new file mode 100644
index 0000000000..335a0c7094
--- /dev/null
+++ b/tests/kvcache/test_tier3.py
@@ -0,0 +1,85 @@
+"""Degrau 3 — teste standalone do KV store host do HostPrefixTier (CPU).
+
+Cobre: round-trip D2H->H2D de páginas KV (códigos + scales), a serpente dos ids lógicos
+(a página lógica reescrita depois do evict NÃO contamina a cópia host), merge de entrada
+(gdn depois kv), e drop-oldest quando o budget enche.
+"""
+import torch
+from types import SimpleNamespace
+
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+
+
+def main():
+    # pool fake: [2, L=2, P=8, page=4, H=2, D=3] códigos bf16 + scales [2, L, P*4, 2] fp32
+    L, P, PAGE, H, D = 2, 8, 4, 2, 3
+    kv_pool = SimpleNamespace(
+        _kv_buffer=torch.randn(2, L, P, PAGE, H, D, dtype=torch.bfloat16),
+        _scale_buffer=torch.randn(2, L, P * PAGE, H, dtype=torch.float32),
+    )
+    state_pool = SimpleNamespace(
+        conv_states=torch.randn(L, 3, 8, 3, dtype=torch.bfloat16),
+        recurrent_states=torch.randn(L, 3, 2, 4, 4, dtype=torch.float32),
+        slot_states={},
+    )
+    page_bytes = 2 * L * PAGE * H * D * 2 + 2 * L * PAGE * H * 4
+    # budget pra 3 páginas (força drop-oldest cedo)
+    tier = HostPrefixTier(state_pool, kv_pool, gdn_slots_host=2,
+                          kv_budget_bytes=3 * page_bytes, device="cpu")
+    assert tier.has_kv and len(tier._kv_free) == 3
+
+    # 1) write/read round-trip: prefixo de 2 páginas (páginas lógicas 3 e 5)
+    kv_idx = torch.cat([torch.arange(3 * PAGE, 4 * PAGE), torch.arange(5 * PAGE, 6 * PAGE)]).to(torch.int32)
+    orig_codes = kv_pool._kv_buffer[:, :, 3].clone(), kv_pool._kv_buffer[:, :, 5].clone()
+    orig_scales = (kv_pool._scale_buffer[:, :, 3 * PAGE:4 * PAGE].clone(),
+                   kv_pool._scale_buffer[:, :, 5 * PAGE:6 * PAGE].clone())
+    hs = tier.write_kv(kv_idx)
+    assert hs is not None and hs.numel() == 2
+
+    # 2) A SERPENTE: sobrescreve as páginas lógicas 3 e 5 (novo inquilino)
+    kv_pool._kv_buffer[:, :, 3].fill_(99)
+    kv_pool._kv_buffer[:, :, 5].fill_(99)
+    kv_pool._scale_buffer[:, :, 3 * PAGE:6 * PAGE].fill_(99)
+    # lê de volta pra páginas 0 e 7
+    tier.read_kv(hs, torch.tensor([0, 7], dtype=torch.int32))
+    assert torch.equal(kv_pool._kv_buffer[:, :, 0], orig_codes[0])
+    assert torch.equal(kv_pool._kv_buffer[:, :, 7], orig_codes[1])
+    assert torch.equal(kv_pool._scale_buffer[:, :, 0 * PAGE:1 * PAGE], orig_scales[0])
+    assert torch.equal(kv_pool._scale_buffer[:, :, 7 * PAGE:8 * PAGE], orig_scales[1])
+    print("PASS 1: round-trip KV host + ids lógicos reescritos não contaminam")
+
+    # 3) merge: entrada com gdn primeiro, kv depois
+    toks = torch.arange(64 * 2, dtype=torch.int64)
+    g = tier.write_gdn(1)
+    tier.put_prefix(toks, None, g)
+    tier.put_prefix(toks, hs, None)  # merge: adota kv, mantém gdn
+    m = tier.lookup(toks)
+    assert m is not None and m.gdn_slot == g and torch.equal(m.kv_pages, hs)
+    print("PASS 2: merge gdn+kv na mesma entrada")
+
+    # 4) drop-oldest quando o budget enche: escreve 2 páginas (resta 1), depois 2 de novo
+    #    -> a entrada mais velha (toks) é dropada INTEIRA (KV slots + slot GDN)
+    kv2 = torch.cat([torch.arange(0 * PAGE, 1 * PAGE), torch.arange(2 * PAGE, 3 * PAGE)]).to(torch.int32)
+    toks2 = torch.arange(64 * 2, 64 * 3, dtype=torch.int64)
+    s2 = tier.write_kv(kv2)   # precisa de 2 slots, tem 1 -> drop de toks (inteiro)
+    tier.put_prefix(toks2, s2, None)
+    assert tier.lookup(toks) is None, "o mais velho devia ter sido dropado inteiro"
+    assert g in tier._free, "o slot GDN do dropado devia voltar pro free-list"
+    kv3 = torch.arange(1 * PAGE, 3 * PAGE).to(torch.int32)
+    toks3 = torch.arange(64 * 3, 64 * 4, dtype=torch.int64)
+    s3 = tier.write_kv(kv3)   # precisa de 2, tem 1 -> drop de toks2
+    tier.put_prefix(toks3, s3, None)
+    assert tier.lookup(toks2) is None
+    m3 = tier.lookup(toks3)
+    assert m3 is not None and torch.equal(m3.kv_pages, s3)
+    print("PASS 3: drop-oldest remove a entrada mais fria inteira (KV + GDN)")
+
+    # 5) evict_host libera tudo
+    tier.evict_host(toks3)
+    assert tier.lookup(toks3) is None and len(tier._kv_free) == 3
+    print("PASS 4: evict_host remove e libera")
+
+    print("\nTODOS OS TESTES PASSARAM ✓")
+
+
+main()
diff --git a/tests/kvcache/test_tree_hooks.py b/tests/kvcache/test_tree_hooks.py
new file mode 100644
index 0000000000..840629310f
--- /dev/null
+++ b/tests/kvcache/test_tree_hooks.py
@@ -0,0 +1,77 @@
+"""Degrau 2 — teste standalone dos hooks do HybridRadixCache (CPU, sem servidor)."""
+import torch
+from freetoken.kvcache.hybrid_radix_cache import HybridRadixCache
+
+
+def main():
+    tree = HybridRadixCache(torch.device("cpu"), 64)
+    evicted = []       # (prefix_len, mamba_slot) capturados pelo hook de evicção
+    rehydrated = []    # prefix_len reidratado no match
+    host_index = set()
+
+    tree.on_evict_node = lambda node: evicted.append(
+        (len(tree._collect_key(node)), node.mamba_value))
+    next_slot = [100]
+
+    def fake_rehydrate(node):
+        n = len(tree._collect_key(node))
+        if n in host_index:
+            rehydrated.append(n)
+            s = next_slot[0]
+            next_slot[0] += 1
+            return s
+        return None
+
+    tree.on_match_dead_snapshot = fake_rehydrate
+
+    # A (4096 toks, snapshot 7) + filho B (mais 2048, snapshot 8) -> A é nó INTERNO
+    idsA = torch.arange(4096, dtype=torch.int32)
+    kvA = torch.arange(4096, dtype=torch.int32) * 10
+    assert tree.insert(idsA, kvA, mamba_value=7) == (0, False)
+    idsB = torch.cat([idsA, torch.full((2048,), 555, dtype=torch.int32)])
+    kvB = torch.arange(6144, dtype=torch.int32) * 10
+    assert tree.insert(idsB, kvB, mamba_value=8) == (4096, False)
+
+    # match enquanto tudo vivo: ganha o snapshot mais profundo (B, 6144)
+    m = tree.match_prefix(idsB)
+    assert m.cached_len == 6144 and m.mamba_value == 8
+
+    # evicta os 2 snapshots: A (interno -> tombstone) e B (folha -> com tier, tombstone
+    # in-place: o KV da folha FICA no tree, não é liberado)
+    er = tree.evict_mamba(2)
+    assert sorted(er.mamba_slots) == [7, 8], er
+    assert sorted(evicted) == [(4096, 7), (6144, 8)], evicted
+    assert er.kv_indices.numel() == 0, "com o tier, nenhuma página KV é liberada"
+    print("PASS 1: on_evict_node dispara; folha vira tombstone e mantém o KV:", evicted)
+
+    # sem cópia host: match trunca pra 0 (A e B viraram tombstones)
+    m = tree.match_prefix(idsA)
+    assert m.cached_len == 0 and m.mamba_value is None
+    print("PASS 2: match trunca sem tier host")
+
+    # host tem A: match(idsA) reidrata A
+    host_index.add(4096)
+    m = tree.match_prefix(idsA)
+    assert m.cached_len == 4096 and m.mamba_value == 100, (m.cached_len, m.mamba_value)
+    assert rehydrated == [4096]
+    print("PASS 3: match reidrata tombstone interno via tier host (slot 100)")
+
+    # host tem B também: match(idsB) reidrata a FOLHA tombstone (mais profunda)
+    host_index.add(6144)
+    m = tree.match_prefix(idsB)
+    assert m.cached_len == 6144 and m.mamba_value == 101, (m.cached_len, m.mamba_value)
+    print("PASS 4: match reidrata folha-tombstone (KV preservado) -> 6144")
+
+    # nós reidratados ficam vivos: segundo match não chama o hook
+    m2 = tree.match_prefix(idsB)
+    assert m2.mamba_value == 101 and rehydrated == [4096, 6144]
+    print("PASS 5: nós reidratados ficam vivos (sem H2D repetido)")
+
+    # integridade de contagem: os dois snapshots revividos voltam a contar como evictable
+    assert tree.mamba_evictable == 2, tree.mamba_evictable
+    print("PASS 6: contagem mamba_evictable consistente")
+
+    print("\nTODOS OS TESTES PASSARAM ✓")
+
+
+main()

From 6c316e2c93719eb1acb6a6874869120da64d1a30 Mon Sep 17 00:00:00 2001
From: alvarorsouza-arch <alvaro.r.souza@gmail.com>
Date: Sat, 19 Sep 2026 10:58:20 +0000
Subject: [PATCH 047/249] fix(qsa): 16-bit compute dtype guard for the indexer
 + fixed-shape write-pages

From the Sep 15-16 debugging sessions (boolean-indexing anti-pattern family):
- the indexer buffers must use the pool COMPUTE dtype (bf16), never the e4m3
  store dtype of an fp8 KV pool, or CUDA-graph capture fails on tl.dot
- ensure_write_pages takes token_to_req for the fixed-shape kernel path
---
 python/freetoken/attention/qsa_sparse.py | 53 +++++++++++-------------
 1 file changed, 24 insertions(+), 29 deletions(-)

diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index ae36fc73fb..e889515f12 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -116,7 +116,17 @@ def __init__(self, config: ModelConfig) -> None:
             f"qsa_sparse backend needs a QSA pool, got {type(self.kvcache).__name__}"
         )
         self.device = self.kvcache.device
+        # The pool's COMPUTE dtype, never its store dtype (the contract lives in
+        # kvcache/base.py). These buffers feed the indexer -- qsa_index_norm_rope and
+        # qsa_mqa_paged -- whose tl.dot has no fp8 path, so an e4m3 q_index does not
+        # fail here, it fails at CUDA-graph capture with "Unsupported rhs dtype
+        # fp8e4nv". --kv-cache-dtype fp8 quantizes only the KV tiers; the index tiers
+        # stay 16-bit by design (kvcache/qsa_pool.py).
         self.dtype = self.kvcache.dtype
+        assert self.dtype.itemsize == 2, (
+            f"QSA block selection needs a 16-bit compute dtype, got {self.dtype} -- "
+            "the KV pool must report its compute dtype, not e4m3 codes"
+        )
         self.index_head_dim = self.kvcache.index_head_dim
         self.ratio = self.kvcache.index_ratio
         self.ring_capacity = self.kvcache.ring_capacity
@@ -303,8 +313,7 @@ def qsa_forward(
             # Bind this forward's write pages before the first store; they ride every fetch
             # query below, so no mid-forward eviction can take them (and the write-through
             # mirror makes any other eviction lossless).
-            md.write_pages = off.ensure_write_pages(batch.out_loc, md.is_decode,
-                                                    md.token_to_req)
+            md.write_pages = off.ensure_write_pages(batch.out_loc, md.is_decode)
             md.out_loc_gpu = off.translate_slots(batch.out_loc, md.is_decode)
         self.kvcache.store_kv(k, v, batch.out_loc if off is None else md.out_loc_gpu, layer_id)
         if md.block_table is None:
@@ -320,6 +329,8 @@ def qsa_forward(
             # Write-through: the pinned host mirror (the backing store) sees every token.
             off.mirror_store(k, v, batch.out_loc, layer_id, getattr(md, "out_loc_gpu", None))
         indices = self._select(index, md, slot)
+        # Scale tensors only exist on an fp8 pool (k_scale returns None otherwise); the
+        # index tier stays bf16 either way, so _select above is quantization-agnostic.
         if off is not None:
             return self._attend_offloaded(qsa_sparse_paged_attention, q, indices, md, layer_id)
         return qsa_sparse_paged_attention(
@@ -330,6 +341,8 @@ def qsa_forward(
             md.block_table,
             md.token_to_req,
             torch.empty_like(q),
+            k_scale=self.kvcache.k_scale(layer_id),
+            v_scale=self.kvcache.v_scale(layer_id),
         )
 
     def _attend_offloaded(self, attend, q, indices, md, layer_id) -> torch.Tensor:
@@ -341,28 +354,16 @@ def _attend_offloaded(self, attend, q, indices, md, layer_id) -> torch.Tensor:
         k_cache = self.kvcache.k_cache(layer_id)
         v_cache = self.kvcache.v_cache(layer_id)
         # fp8 pool (PR #354): the attend kernel dequantizes with these; None on a bf16 pool.
-        # Scale accessors only exist on a quantized KV pool (fp8, e.g. PR #354); the
-        # plain bf16 pool on main has none. Probe instead of assuming, so the offload
-        # composes with either.
-        _k_scale = getattr(self.kvcache, "k_scale", None)
-        _v_scale = getattr(self.kvcache, "v_scale", None)
-        ks = _k_scale(layer_id) if _k_scale is not None else None
-        vs = _v_scale(layer_id) if _v_scale is not None else None
-        # The attend kernel on an unquantized pool takes NO scale arguments at all, so the
-        # kwargs must be omitted entirely (passing k_scale=None is still unexpected).
-        scales: dict = {}
-        if ks is not None:
-            scales["k_scale"] = ks
-        if vs is not None:
-            scales["v_scale"] = vs
+        ks = self.kvcache.k_scale(layer_id)
+        vs = self.kvcache.v_scale(layer_id)
         if md.is_decode:
             eff = off.fetch_for_attend(indices, md)
             return attend(q, k_cache, v_cache, indices, eff, md.token_to_req,
-                          torch.empty_like(q), **scales)
+                          torch.empty_like(q), k_scale=ks, v_scale=vs)
         rows = q.shape[0]
         if rows == 0:
             return attend(q, k_cache, v_cache, indices, md.block_table, md.token_to_req,
-                          torch.empty_like(q), **scales)
+                          torch.empty_like(q), k_scale=ks, v_scale=vs)
         # Row-group loop: one compact for the whole layer (grid = rows, no sync), then the
         # offloader packs rows into residency-feasible groups CPU-side with REAL page unions
         # (one D2H per layer) -- consecutive rows select mostly the same pages, so groups
@@ -371,20 +372,14 @@ def _attend_offloaded(self, attend, q, indices, md, layer_id) -> torch.Tensor:
         out = torch.empty_like(q)
         for (start, end), eff in off.iter_prefill_groups(sel_all, md.write_pages, md.block_table):
             attend(q[start:end], k_cache, v_cache, indices[start:end], eff,
-                   md.token_to_req[start:end], out[start:end], **scales)
+                   md.token_to_req[start:end], out[start:end], k_scale=ks, v_scale=vs)
         dropped = off.trunc_count()  # one sync per layer
         if dropped:
-            msg = (
-                f"KV host offload: {dropped} selected page(s) dropped this layer because a "
-                f"query row spans more than {off.max_sel_pages} distinct pages; those rows "
-                "would attend over unrelated physical slots. Raise --kv-host-pages or "
-                "--memory-ratio, or run with --max-running-requests 1."
+            logger.warning(
+                f"KV host offload: {dropped} selected pages dropped this layer "
+                f"(selection spans > {off.max_sel_pages} pages); the affected rows attended "
+                "over fallback slots"
             )
-            if md.is_decode:
-                # Inside a CUDA graph we cannot raise without tearing down the capture.
-                logger.warning(msg)
-            else:
-                raise RuntimeError(msg)
         return out
 
     def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None:

From b9b8f1b330640b95f6ea04ec1e9dd91bbeb468e3 Mon Sep 17 00:00:00 2001
From: Oscar Rojas <116426609+oscar-investmatic@users.noreply.github.com>
Date: Sun, 20 Sep 2026 16:03:39 +0200
Subject: [PATCH 048/249] fix(server): clamp /v1/models context_length to the
 allocated KV pool

/v1/models published `config.max_seq_len` -- the model's positional ceiling -- while the
engine clamps its own `max_seq_len` to `min(config.max_seq_len, num_pages * page_size)` at
startup and again in `_refresh_seq_state` after a rebuild, and the scheduler admits against
that clamped value. The frontend process holds an unclamped ServerArgs copy, so the two
disagreed about one quantity: on a 12 GiB card, 262144 advertised against a 178176-token pool.

`ft launch` reads this number to size each agent's context window (opencode `limit.context`,
codex `context_window`/`max_context_window`, CLAUDE_CODE_MAX_CONTEXT_TOKENS), so agents set
their compaction threshold past what the server can hold and hit `context_length_exceeded`
instead of compacting in time.

Split the pool resolution out of `cache_geometry` into `kv_pool_geometry` so /v1/models and
/v1/cache/status cannot drift apart, and clamp the advertised value with it. Falls back to the
model ceiling while `num_pages` is still 0 (loading, or an engine build that sends no meta ack).

The existing route test passed either way -- FakeState carries no pool geometry, so it fell
through to the ceiling. It now asserts that fallback explicitly, alongside new cases for the
clamp, a pool larger than the ceiling, and a rebuild superseding the load-time allocation.

Fixes #448

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 python/freetoken/server/api_server.py | 27 +++++++++++++-
 python/freetoken/server/openai_api.py | 27 ++++++++++++--
 tests/server/test_openai_api.py       | 54 +++++++++++++++++++++++++--
 3 files changed, 100 insertions(+), 8 deletions(-)

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 5592a2b0d5..7267b1a9c3 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -710,6 +710,30 @@ def _reasoning_geometry(state: Any) -> dict | None:
     return {"gears": list(gears), "default": default, "kwargs": kwargs}
 
 
+def kv_pool_geometry(state: Any) -> tuple[int, int]:
+    """The allocated KV pool as ``(num_pages, page_size)``, most-recent-truth first: the last
+    rebuild's result, else the running UserReply snapshot, else the load-time ("meta", …) ack.
+    ``num_pages`` is 0 until one of those lands (still loading / older engine build).
+
+    Split out of cache_geometry so /v1/models can clamp the advertised context to the pool
+    without recomputing the whole panel payload -- both callers must agree on the resolution
+    order or they would report different capacities for the same server.
+    """
+    stats = getattr(state, "stats", None)
+    last = getattr(state, "last_rebuild", None) or {}
+    pools = getattr(state, "cache_pools", None) or {}
+    num_pages = int(
+        last.get("num_pages")
+        or getattr(stats, "kv_total_pages", 0)
+        or pools.get("num_pages", 0)
+        or 0
+    )
+    page_size = int(
+        pools.get("page_size", 0) or getattr(getattr(state, "config", None), "page_size", 1) or 1
+    )
+    return num_pages, page_size
+
+
 def cache_geometry(state: Any) -> dict:
     """Current cache geometry for the desktop cache panel. Each pool size resolves
     most-recent-truth first: the last rebuild's result, else the running UserReply snapshot
@@ -725,11 +749,10 @@ def cache_geometry(state: Any) -> dict:
     config = state.config
     last = getattr(state, "last_rebuild", None) or {}
     pools = getattr(state, "cache_pools", None) or {}
-    num_pages = int(last.get("num_pages") or tr.kv_total_pages or pools.get("num_pages", 0) or 0)
+    num_pages, page_size = kv_pool_geometry(state)
     num_mamba_slots = int(
         last.get("mamba_slots") or tr.mamba_total_slots or pools.get("num_mamba_slots", 0) or 0
     )
-    page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1)
     moe_cache_size = last.get("moe_cache_size")
     if moe_cache_size is None:
         moe_cache_size = int(pools.get("moe_cache_size", 0) or 0) or configured_moe_cache_size(
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index dc2f73a979..087b2a4e48 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -701,10 +701,31 @@ def _served_model_name(state: Any) -> str:
 
 
 def _model_context_length(state: Any) -> int | None:
-    """The model ceiling, not `min(ceiling, KV budget)`: a rebuild moves the latter, and agents
-    read this once at startup."""
+    """The usable context: ``min(model ceiling, allocated KV tokens)``.
+
+    This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__,
+    and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler
+    admits against. The frontend process holds an unclamped ServerArgs copy, so publishing
+    ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144
+    against a 178176-token pool, so the engine and this route disagreed about one quantity.
+
+    `ft launch` reads this to size each agent's context window (opencode's ``limit.context``,
+    codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops
+    those agents compacting before the KV pool runs out.
+    """
     try:  # never 500 a metadata route: max_seq_len walks into the HF config on some builds
         value = int(state.config.max_seq_len)
     except Exception:  # noqa: BLE001
         return None
-    return value if value > 0 else None
+    if value <= 0:
+        return None
+    try:
+        # Local import: api_server imports this module, so a module-level one would cycle.
+        from .api_server import kv_pool_geometry
+
+        num_pages, page_size = kv_pool_geometry(state)
+        kv_tokens = num_pages * page_size
+    except Exception:  # noqa: BLE001
+        kv_tokens = 0
+    # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then.
+    return min(value, kv_tokens) if kv_tokens > 0 else value
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index e33018e195..d9d565ff64 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -442,19 +442,67 @@ def test_models_route_returns_served_model_name():
     assert card["max_model_len"] is None and card["context_length"] is None
 
 
+def _models_card(state):
+    app = FastAPI()
+    register_openai_routes(app, lambda: state, lambda: {})
+    return TestClient(app).get("/v1/models").json()["data"][0]
+
+
 def test_models_route_publishes_the_model_context_length():
     """`ft launch` reads this to size each agent's context window."""
     state = FakeState([])
     state.config.max_seq_len = 262144
-    app = FastAPI()
-    register_openai_routes(app, lambda: state, lambda: {})
 
-    card = TestClient(app).get("/v1/models").json()["data"][0]
+    card = _models_card(state)
 
+    # No pool geometry resolved yet (still loading, or pre-("meta", …) ack): the model
+    # ceiling is the only answer available.
     assert card["max_model_len"] == 262144
     assert card["context_length"] == 262144
 
 
+def test_models_route_clamps_the_context_length_to_the_kv_pool():
+    """The engine clamps its own max_seq_len to the allocated pool and the scheduler admits
+    against that, so publishing the unclamped ceiling here made `ft launch` size each agent's
+    compaction window past what the server can actually hold (#448)."""
+    state = FakeState([])
+    state.config.max_seq_len = 262144
+    state.config.page_size = 1
+    state.cache_pools = {"num_pages": 178176, "page_size": 1}
+
+    card = _models_card(state)
+
+    assert card["max_model_len"] == 178176
+    assert card["context_length"] == 178176
+
+
+def test_models_route_keeps_the_ceiling_when_the_pool_exceeds_it():
+    """A pool larger than the model's positional ceiling does not extend the context."""
+    state = FakeState([])
+    state.config.max_seq_len = 32768
+    state.config.page_size = 1
+    state.cache_pools = {"num_pages": 178176, "page_size": 1}
+
+    card = _models_card(state)
+
+    assert card["max_model_len"] == 32768
+    assert card["context_length"] == 32768
+
+
+def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool():
+    """A rebuild moves the pool; /v1/models must follow it rather than freeze the load-time
+    allocation (same most-recent-truth order /v1/cache/status reports)."""
+    state = FakeState([])
+    state.config.max_seq_len = 262144
+    state.config.page_size = 1
+    state.cache_pools = {"num_pages": 178176, "page_size": 1}
+    state.last_rebuild = {"num_pages": 40000}
+
+    card = _models_card(state)
+
+    assert card["max_model_len"] == 40000
+
+
 async def _collect(generator):
     return [chunk async for chunk in generator]
 

From f9947db772b9e6217540bb8e7a6ae9db54f17ce2 Mon Sep 17 00:00:00 2001
From: chrisqian <chrisqian@example.com>
Date: Fri, 18 Sep 2026 12:57:32 +0800
Subject: [PATCH 049/249] fix(moe): keep eager fallback and a bounds guard in
 the prefill invalidation helper

CPU-resident test caches reach this helper through prefetch_prefill_layer and
cannot run the Triton kernel (triton is a Linux-only dependency). The eager
path clamped via Python slices; raw pointer arithmetic cannot, so refuse a
buffer that does not fit instead of corrupting the heap.

Assisted-by: pi
---
 python/freetoken/kernel/triton/moe/invalidate.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/python/freetoken/kernel/triton/moe/invalidate.py b/python/freetoken/kernel/triton/moe/invalidate.py
index 4fa01db394..24c5fbf3f9 100644
--- a/python/freetoken/kernel/triton/moe/invalidate.py
+++ b/python/freetoken/kernel/triton/moe/invalidate.py
@@ -54,6 +54,21 @@ def invalidate_prefill_slots(
     (``layer * num_experts + expert``) and the tensor is contiguous, so the kernel's
     flat addressing matches the ``view(-1)`` the eager path used.
     """
+    # The raw-pointer kernel cannot clamp like the eager slices did; refuse a buffer
+    # that does not fit instead of reading past the cache.
+    if slot_start + num_experts > id_of_slot.numel():
+        raise ValueError(
+            f"prefill buffer [{slot_start}, {slot_start + num_experts}) exceeds "
+            f"id_of_slot size {id_of_slot.numel()}"
+        )
+    # The kernel needs a CUDA context; CPU test caches keep the old eager path, where
+    # the hidden sync is harmless (no enqueued GPU work to drain).
+    if id_of_slot.device.type != "cuda":
+        old_ids = id_of_slot[slot_start : slot_start + num_experts]
+        slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
+        old_ids.fill_(-1)
+        usage[slot_start : slot_start + num_experts].zero_()
+        return
     grid = ((num_experts + 255) // 256,)
     _invalidate_slots_kernel[grid](
         id_of_slot, slot_for_id, usage, slot_start, num_experts, BLOCK=256

From dc9216c11313fce88997b750ac6af68f72e5e7b1 Mon Sep 17 00:00:00 2001
From: chrisqian <chrisqian@example.com>
Date: Fri, 18 Sep 2026 12:57:32 +0800
Subject: [PATCH 050/249] test(moe): check the prefill invalidation kernel
 against the eager reference on cuda

Element-wise compare over three configs x 20 random trials, covering duplicate
expert ids, empty slots and a nonzero slot_start, plus the bounds guard.

Assisted-by: pi
---
 tests/moe/test_offload.py | 44 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index cd5f5dc773..aaf3a4c097 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -277,6 +277,50 @@ def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots():
     assert torch.equal(cache.bank_caches["down"][:num_experts], down_source[0])
 
 
+def test_invalidate_prefill_slots_kernel_matches_eager_reference():
+    from freetoken.kernel.triton.moe import invalidate_prefill_slots
+
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the invalidation kernel")
+
+    for num_layers, num_experts, cache_size, slot_start in [
+        (3, 4, 8, 0),
+        (3, 4, 8, 4),
+        (2, 256, 1024, 512),
+    ]:
+        for trial in range(20):
+            torch.manual_seed(1234 + trial)
+            ids = torch.full((cache_size,), -1, dtype=torch.int32)
+            k = torch.randint(0, num_experts + 1, (1,)).item()
+            slots = torch.randperm(cache_size)[:k]
+            ids[slots] = torch.randint(0, num_layers * num_experts, (k,), dtype=torch.int32)
+            slot_for_id = torch.full((num_layers, num_experts), -1, dtype=torch.int32)
+            for slot, expert in enumerate(ids.tolist()):
+                if expert >= 0:
+                    slot_for_id.view(-1)[expert] = slot
+            usage = torch.randint(0, 999, (cache_size,), dtype=torch.int64)
+
+            ids_c, sfi_c, usage_c = ids.cuda(), slot_for_id.cuda(), usage.cuda()
+
+            ref_sfi = slot_for_id.clone()
+            ref_usage = usage.clone()
+            ref_ids = ids.clone()
+            buf = ref_ids[slot_start : slot_start + num_experts]
+            ref_sfi.view(-1)[buf[buf >= 0].long()] = -1
+            buf.fill_(-1)
+            ref_usage[slot_start : slot_start + num_experts].zero_()
+
+            invalidate_prefill_slots(ids_c, sfi_c, usage_c, slot_start, num_experts)
+            torch.cuda.synchronize()
+
+            assert torch.equal(ref_sfi, sfi_c.cpu())
+            assert torch.equal(ref_ids, ids_c.cpu())
+            assert torch.equal(ref_usage, usage_c.cpu())
+
+    with pytest.raises(ValueError):
+        invalidate_prefill_slots(ids.new_full((1,), -1), slot_for_id, usage.new_zeros(1), 0, 2)
+
+
 def test_prefill_overlap_waits_for_previous_prefill_release_after_begin(monkeypatch):
     from freetoken.moe.offload_cache import OffloadMoeCache
 

From cee23bf44a86e48d70f8fda3ca3b6d4d9890072e Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 23 Sep 2026 22:22:48 -0700
Subject: [PATCH 051/249] fix(runtime): correct overlap token limits and cache
 safety (#546)

* fix(scheduler): count delivered tokens under overlap

* fix(moe): initialize slots used by inactive hybrid routes

* fix(kvcache): report usable window capacity floors

* refactor(moe): inline the zeroed slot cache allocation

* test: drop the new regression tests

---------

Co-authored-by: andy-yang-1 <andy_yang@berkeley.edu>
---
 python/freetoken/kvcache/cache_status.py | 11 ++++-------
 python/freetoken/moe/offload_cache.py    |  4 ++--
 python/freetoken/scheduler/scheduler.py  |  3 ++-
 tests/kvcache/test_cache_unit_bytes.py   |  2 +-
 4 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py
index 10169d6515..6aa13e7f75 100644
--- a/python/freetoken/kvcache/cache_status.py
+++ b/python/freetoken/kvcache/cache_status.py
@@ -120,21 +120,18 @@ def _mamba() -> int:
         return int(_linear_pool_min_slots(config) - 1)
 
     def _swa() -> int:
-        # Window-pool floor in tokens (matches the pool's own page unit x count):
-        #   DSV4      -- the working-set floor pages _dsv4_pool_sizes clamps n_win_pages to, x P.
-        #   radix-SWA -- the concurrent working set (+1 slot-0 sentinel) _swa_paged_num_tokens
-        #                floors to.
-        # 0 for models without a window pool.
+        # Report usable tokens, matching num_swa_pages in rebuild requests.
+        # DSV4's physical floor includes a dummy page; the SWA floor already excludes slot 0.
         from .dsv4_cost_model import _dsv4_window_floor_pages
         from .hybrid_swa_pool import _swa_pool_floor
 
         mc = config.model_config
         if mc.dsv4_args is not None:
             P = mc.dsv4_args.window_size
-            return int(_dsv4_window_floor_pages(config, P) * P)
+            return int((_dsv4_window_floor_pages(config, P) - 1) * P)
         if not (mc.has_swa_attention and config.cache_type == "swa_radix"):
             return 0
-        return int(_swa_pool_floor(config) + 1)
+        return int(_swa_pool_floor(config))
 
     for key, fn in (("kv_tokens", _kv), ("moe_experts", _moe), ("mamba_slots", _mamba),
                     ("swa_tokens", _swa)):
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 7abb9c8003..32e0f3cc1e 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -354,7 +354,7 @@ def set_bank_sources(
                     name, layer_id, source.shape, source.dtype,
                 )
             self.bank_sources[name] = list(per_layer)
-            self.bank_caches[name] = torch.empty(
+            self.bank_caches[name] = torch.zeros(
                 (self.cache_size, *head.shape[1:]),
                 dtype=head.dtype,
                 device=self.device,
@@ -492,7 +492,7 @@ def rebuild(self, cache_size: int) -> None:
         # 3. Reallocate the slot cache from the retained host sources.
         for name in self.bank_schema:
             head = self.bank_sources[name][0]
-            self.bank_caches[name] = torch.empty(
+            self.bank_caches[name] = torch.zeros(
                 (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device
             )
         self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema]
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index ac1bf322e3..fa0d96a714 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -349,7 +349,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                 next_token = int(next_token.item())
                 # EOS / stop-string -> "stop", output budget exhausted -> "length";
                 # EOS and stop strings win over length.
-                hit_length = not req.can_decode
+                # Overlap can advance device_len ahead of the token delivered to the host.
+                hit_length = req.input_ids.numel() >= req.max_device_len
                 hit_eos = (
                     not req.sampling_params.ignore_eos and next_token in self.eos_token_ids
                 )
diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py
index c0eec34eff..496622f985 100644
--- a/tests/kvcache/test_cache_unit_bytes.py
+++ b/tests/kvcache/test_cache_unit_bytes.py
@@ -213,7 +213,7 @@ def test_floors_dsv4_reports_real_window_floor():
     assert floors["kv_tokens"] > 0
     # The window (swa) floor is reported too, in tokens (window pages x P) -- not the bogus 0 the
     # radix-only branch used to give for DSV4.
-    assert floors["swa_tokens"] == _dsv4_window_floor_pages(cfg, P) * P
+    assert floors["swa_tokens"] == (_dsv4_window_floor_pages(cfg, P) - 1) * P
 
 
 def test_floors_missing_config_all_zero():

From f5b9700cc10312d52169d9195af043c040f66d70 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 24 Sep 2026 20:16:03 -0700
Subject: [PATCH 052/249] refactor(ple): extract a shared disk row store (#548)

---
 .../row_store_ext.cpp}                        | 170 ++++++++++++------
 python/freetoken/kernel/row_store.py          |  54 ++++++
 python/freetoken/models/qwen4_exp/ple_disk.py |  37 +---
 setup.py                                      |   6 +-
 tests/kernels/test_row_store.py               |  87 +++++++++
 tests/models/qwen4_exp/test_ple_disk.py       |   8 +-
 6 files changed, 270 insertions(+), 92 deletions(-)
 rename python/freetoken/kernel/csrc/{ple_store/ple_store_ext.cpp => row_store/row_store_ext.cpp} (83%)
 create mode 100644 python/freetoken/kernel/row_store.py
 create mode 100644 tests/kernels/test_row_store.py

diff --git a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp
similarity index 83%
rename from python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp
rename to python/freetoken/kernel/csrc/row_store/row_store_ext.cpp
index 2d54d692d6..e33b18a8e0 100644
--- a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp
+++ b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp
@@ -1,7 +1,12 @@
-// Disk-backed PLE row store: rows read straight from the checkpoint's fp8 shard tensors
-// through an extent table (PleRowSource in ple_ssd.py). Engine-thread only, no locks.
-// Duplicate rows in one fill dedup into ONE batched read round; no RAM cache and no
-// per-sequence state. Hash reference: tests/models/qwen4_exp/test_ple_disk.py.
+// Disk-backed row store: fixed-width rows read straight from checkpoint shard tensors
+// through an extent table (row i of extent e at extent_base[e] + i * row_stride).
+// Engine-thread only, no locks. Duplicate rows in one fill dedup into ONE batched read
+// round; no RAM cache and no per-sequence state.
+//
+//   RowStore  -- hash-agnostic: stage_rows(row_ids) + flush(); callers supply row ids.
+//   PleStore  -- RowStore plus the Qwen3.8-Flash-Next PLE n-gram hash: stage(tokens).
+//                Hash reference: tests/models/qwen4_exp/test_ple_disk.py.
+//
 // Platform seams: TableFile (O_DIRECT+pread; Win: NO_BUFFERING), BatchReader (io_uring,
 // pread-pool fallback = the portable shape), cumemop_* (dlopen libcuda; Win: nvcuda).
 
@@ -183,7 +188,9 @@ class ThreadPoolBatchReader final : public BatchReader {
   ThreadPoolBatchReader() {
     // these threads block on I/O, not compute, so the core count is only a default
     unsigned n = std::max(1u, std::min(kReaderThreads, std::thread::hardware_concurrency()));
-    if (const char *env = std::getenv("FREETOKEN_PLE_READER_THREADS")) {
+    const char *env = std::getenv("FREETOKEN_ROW_STORE_READER_THREADS");
+    if (env == nullptr) env = std::getenv("FREETOKEN_PLE_READER_THREADS");  // the PLE-era name
+    if (env != nullptr) {
       // kBatchEntries is the submit depth, so threads past it never get a read
       const int v = std::atoi(env);
       if (v > 0) n = std::min((unsigned)v, kBatchEntries);
@@ -432,29 +439,24 @@ int64_t pos_mod(int64_t v, int64_t m) {
   return r < 0 ? r + m : r;
 }
 
-class PleStore {
+class RowStore {
   struct Extent {
     const TableFile *file;
     int64_t base;
   };
 
  public:
-  PleStore(std::vector<std::string> paths, std::vector<int64_t> extent_file,
+  RowStore(std::vector<std::string> paths, std::vector<int64_t> extent_file,
            std::vector<int64_t> extent_base, int64_t rows_per_extent, int64_t row_bytes,
-           int64_t row_stride, std::vector<int64_t> multipliers, std::vector<int64_t> head_vocab_sizes,
-           std::vector<int64_t> head_offsets, int64_t eos_token_id, bool use_io_uring)
-      : row_bytes_(row_bytes),
-        row_stride_(row_stride),
-        rows_per_extent_(rows_per_extent),
-        mult_(std::move(multipliers)),
-        sizes_(std::move(head_vocab_sizes)),
-        offsets_(std::move(head_offsets)),
-        eos_(eos_token_id) {
-    if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty())
-      throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables");
+           int64_t row_stride, bool use_io_uring)
+      : row_bytes_(row_bytes), row_stride_(row_stride), rows_per_extent_(rows_per_extent) {
+    if (row_bytes_ <= 0 || row_stride_ < row_bytes_ || rows_per_extent_ <= 0)
+      throw std::runtime_error("row store geometry: want row_bytes > 0, row_stride >= row_bytes, rows_per_extent > 0");
     if (row_bytes_ > kPage)
-      throw std::runtime_error("PLE row_bytes " + std::to_string(row_bytes_) +
+      throw std::runtime_error("row_bytes " + std::to_string(row_bytes_) +
                                " exceeds a page; bounce slots assume one-page rows");
+    if (extent_file.size() != extent_base.size() || extent_file.empty())
+      throw std::runtime_error("row store geometry: extent_file and extent_base must be equal-length and non-empty");
     for (const std::string &p : paths)
       files_.push_back(std::make_unique<TableFile>(p));
     const int64_t extent_bytes = (rows_per_extent_ - 1) * row_stride_ + row_bytes_;
@@ -472,36 +474,33 @@ class PleStore {
   }
 
   // reader first: a still-running read must not land in freed bounce memory
-  ~PleStore() {
+  virtual ~RowStore() {
     reader_.reset();
     free(bounce_);
   }
 
-  PleStore(const PleStore &) = delete;
-  PleStore &operator=(const PleStore &) = delete;
+  RowStore(const RowStore &) = delete;
+  RowStore &operator=(const RowStore &) = delete;
 
-  // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier.
-  void hash_rows(const int64_t *w, int64_t *rows) {
-    const int64_t prev1 = w[1];
-    const int64_t prev2 = prev1 == eos_ ? eos_ : w[0];
-    const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]);
-    const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]);
-    const size_t half = sizes_.size() / 2;
-    for (size_t h = 0; h < sizes_.size(); h++)
-      rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h];
-  }
+  int64_t row_bytes() const { return row_bytes_; }
+  int64_t total_rows() const { return (int64_t)extents_.size() * rows_per_extent_; }
 
-  // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush().
-  void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) {
-    const int64_t *tokens = reinterpret_cast<const int64_t *>(tokens_addr);
+  // Queue n row ids (int64 at rows_addr); row i lands at staging_addr + i * dst_stride
+  // (dst_stride = 0 packs rows back to back). Out-of-range ids throw before any I/O is
+  // issued. No I/O until flush().
+  void stage_rows(uintptr_t rows_addr, int64_t n, uintptr_t staging_addr, int64_t dst_stride) {
+    const int64_t *rows = reinterpret_cast<const int64_t *>(rows_addr);
     uint8_t *staging = reinterpret_cast<uint8_t *>(staging_addr);
-    const size_t heads = sizes_.size();
-    std::vector<int64_t> rows(heads);
+    if (dst_stride == 0) dst_stride = row_bytes_;
+    if (dst_stride < row_bytes_)
+      throw std::runtime_error("stage_rows: dst_stride below row_bytes would overlap rows");
+    const int64_t total = total_rows();
     for (int64_t i = 0; i < n; i++) {
-      hash_rows(tokens + i, rows.data());
-      for (size_t h = 0; h < heads; h++)
-        request_row(rows[h], staging + ((size_t)i * heads + h) * row_bytes_);
+      if (rows[i] < 0 || rows[i] >= total)
+        throw std::out_of_range("stage_rows: row id " + std::to_string(rows[i]) + " outside [0, " +
+                                std::to_string(total) + ")");
     }
+    for (int64_t i = 0; i < n; i++) request_row(rows[i], staging + (size_t)i * (size_t)dst_stride);
   }
 
   // One batched disk round for everything staged; signals even when nothing was.
@@ -519,15 +518,7 @@ class PleStore {
            std::to_string(files_.size()) + " files";
   }
 
- private:
-  struct Pending {
-    const TableFile *file;
-    int64_t read_off;
-    int64_t read_len;
-    int64_t row_off;  // row payload start inside the read buffer
-    std::vector<uint8_t *> dsts;
-  };
-
+ protected:
   // Queue dst on this fill's pending batch; duplicate rows fan out from one read.
   void request_row(int64_t row_id, uint8_t *dst) {
     auto pit = pending_index_.find(row_id);
@@ -548,11 +539,20 @@ class PleStore {
     pending_.push_back(std::move(p));
   }
 
+ private:
+  struct Pending {
+    const TableFile *file;
+    int64_t read_off;
+    int64_t read_len;
+    int64_t row_off;  // row payload start inside the read buffer
+    std::vector<uint8_t *> dsts;
+  };
+
   // Read every pending row in reader-capacity batches and fan out the copies.
   void flush_pending() {
     if (pending_.empty()) return;
     struct Cleanup {
-      PleStore *s;
+      RowStore *s;
       ~Cleanup() {
         s->pending_.clear();
         s->pending_index_.clear();
@@ -586,8 +586,6 @@ class PleStore {
   }
 
   int64_t row_bytes_, row_stride_, rows_per_extent_;
-  std::vector<int64_t> mult_, sizes_, offsets_;
-  int64_t eos_;
   std::vector<std::unique_ptr<TableFile>> files_;
   std::vector<Extent> extents_;
   std::unique_ptr<BatchReader> reader_;
@@ -597,10 +595,71 @@ class PleStore {
   std::unordered_map<int64_t, size_t> pending_index_;
 };
 
+// The Qwen3.8-Flash-Next PLE table: RowStore plus the checkpoint's n-gram hash, so a fill
+// takes raw token runs and never crosses the Python boundary per row.
+class PleStore final : public RowStore {
+ public:
+  PleStore(std::vector<std::string> paths, std::vector<int64_t> extent_file,
+           std::vector<int64_t> extent_base, int64_t rows_per_extent, int64_t row_bytes,
+           int64_t row_stride, std::vector<int64_t> multipliers, std::vector<int64_t> head_vocab_sizes,
+           std::vector<int64_t> head_offsets, int64_t eos_token_id, bool use_io_uring)
+      : RowStore(std::move(paths), std::move(extent_file), std::move(extent_base), rows_per_extent,
+                 row_bytes, row_stride, use_io_uring),
+        mult_(std::move(multipliers)),
+        sizes_(std::move(head_vocab_sizes)),
+        offsets_(std::move(head_offsets)),
+        eos_(eos_token_id) {
+    if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty())
+      throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables");
+  }
+
+  // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier.
+  void hash_rows(const int64_t *w, int64_t *rows) {
+    const int64_t prev1 = w[1];
+    const int64_t prev2 = prev1 == eos_ ? eos_ : w[0];
+    const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]);
+    const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]);
+    const size_t half = sizes_.size() / 2;
+    for (size_t h = 0; h < sizes_.size(); h++)
+      rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h];
+  }
+
+  // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush().
+  void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) {
+    const int64_t *tokens = reinterpret_cast<const int64_t *>(tokens_addr);
+    uint8_t *staging = reinterpret_cast<uint8_t *>(staging_addr);
+    const size_t heads = sizes_.size();
+    std::vector<int64_t> rows(heads);
+    for (int64_t i = 0; i < n; i++) {
+      hash_rows(tokens + i, rows.data());
+      for (size_t h = 0; h < heads; h++)
+        request_row(rows[h], staging + ((size_t)i * heads + h) * (size_t)row_bytes());
+    }
+  }
+
+ private:
+  std::vector<int64_t> mult_, sizes_, offsets_;
+  int64_t eos_;
+};
+
 }  // namespace
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  py::class_<PleStore>(m, "PleStore")
+  py::class_<RowStore>(m, "RowStore")
+      .def(py::init<std::vector<std::string>, std::vector<int64_t>, std::vector<int64_t>,
+                    int64_t, int64_t, int64_t, bool>(),
+           py::arg("paths"), py::arg("extent_file"), py::arg("extent_base"),
+           py::arg("rows_per_extent"), py::arg("row_bytes"), py::arg("row_stride"),
+           py::arg("use_io_uring") = true)
+      .def("stage_rows", &RowStore::stage_rows, py::arg("rows_addr"), py::arg("n"),
+           py::arg("staging_addr"), py::arg("dst_stride") = 0,
+           py::call_guard<py::gil_scoped_release>())
+      .def("flush", &RowStore::flush, py::arg("signal_addr") = 0,
+           py::call_guard<py::gil_scoped_release>())
+      .def("io_backend", &RowStore::io_backend)
+      .def_property_readonly("row_bytes", &RowStore::row_bytes)
+      .def_property_readonly("total_rows", &RowStore::total_rows);
+  py::class_<PleStore, RowStore>(m, "PleStore")
       .def(py::init<std::vector<std::string>, std::vector<int64_t>, std::vector<int64_t>,
                     int64_t, int64_t, int64_t, std::vector<int64_t>,
                     std::vector<int64_t>, std::vector<int64_t>, int64_t, bool>(),
@@ -610,10 +669,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            py::arg("head_vocab_sizes"), py::arg("head_offsets"), py::arg("eos_token_id"),
            py::arg("use_io_uring") = true)
       .def("stage", &PleStore::stage, py::arg("tokens_addr"), py::arg("n"),
-           py::arg("staging_addr"), py::call_guard<py::gil_scoped_release>())
-      .def("flush", &PleStore::flush, py::arg("signal_addr") = 0,
-           py::call_guard<py::gil_scoped_release>())
-      .def("io_backend", &PleStore::io_backend);
+           py::arg("staging_addr"), py::call_guard<py::gil_scoped_release>());
   m.def("memop_write", &memop_write, py::arg("stream"), py::arg("addr"), py::arg("value"));
   m.def("memop_wait_geq", &memop_wait_geq, py::arg("stream"), py::arg("addr"), py::arg("value"));
   m.def("memop_wait_reset", &memop_wait_reset, py::arg("stream"), py::arg("flag_addr"));
diff --git a/python/freetoken/kernel/row_store.py b/python/freetoken/kernel/row_store.py
new file mode 100644
index 0000000000..5b42479ebc
--- /dev/null
+++ b/python/freetoken/kernel/row_store.py
@@ -0,0 +1,54 @@
+"""Python face of the disk row store extension (``kernel/csrc/row_store``) and its CUDA-graph sync.
+
+``RowStore`` reads fixed-width rows straight out of checkpoint shard files (io_uring or a pread pool,
+O_DIRECT) into pinned staging by row id; ``PleStore`` adds PLE's n-gram hashing.
+
+Graph sync: a captured graph cannot wait for host I/O, so a captured ``lookup`` WAITs on a pinned
+flag through a stream memop (``wait_reset``) that the host ``signal``s once the fill landed; where
+the driver rejects memops in capture, the caller falls back to launch-gating (fill before replay).
+``probe_wait_sync`` decides which, once, at startup.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from freetoken.kernel import _row_store
+from freetoken.kernel.pinned import alloc_pinned_tensor
+
+RowStore = _row_store.RowStore
+PleStore = _row_store.PleStore
+
+
+def probe_wait_sync(mode: str, device: torch.device) -> bool:
+    """``mode``: ``auto`` (probe), ``wait`` (require memops) or ``gate`` (never use them)."""
+    if mode not in ("auto", "wait", "gate"):
+        raise ValueError(f"unknown row-store sync mode {mode!r}; expected auto, wait or gate")
+    if mode == "gate":
+        return False
+    scratch = alloc_pinned_tensor(1, dtype=torch.int64)
+    scratch.zero_()
+    stream = torch.cuda.current_stream(device)
+    ok = (
+        _row_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0
+        and _row_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0
+    )
+    if ok:
+        stream.synchronize()
+        ok = int(scratch[0]) == 7
+    if mode == "wait" and not ok:
+        raise RuntimeError("wait-sync requested but stream memops are unavailable")
+    return ok
+
+
+def wait_reset(stream: torch.cuda.Stream, flag: torch.Tensor) -> None:
+    """Enqueue WAIT(flag >= 1) then flag := 0 on ``stream`` (capturable)."""
+    _row_store.memop_wait_reset(stream.cuda_stream, flag.data_ptr())
+
+
+def signal(flag: torch.Tensor) -> None:
+    """Host side: release a pending WAIT (a release-store of 1)."""
+    _row_store.signal_flag(flag.data_ptr())
+
+
+__all__ = ["RowStore", "PleStore", "probe_wait_sync", "wait_reset", "signal"]
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index cea6269357..ec7da80d56 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -112,7 +112,7 @@ def __init__(
         max_extend_tokens: int = 8192,
         dtype: torch.dtype = torch.bfloat16,
     ) -> None:
-        from freetoken.kernel import _ple_store
+        from freetoken.kernel.row_store import PleStore
 
         self.num_rows = source.total_rows
         self.head_dim = source.row_bytes  # fp8: one byte per element
@@ -128,7 +128,7 @@ def __init__(
             raise ValueError(
                 f"PLE row source holds {source.total_rows} rows but the hash addresses {need}; incomplete checkpoint?"
             )
-        self._store = _ple_store.PleStore(
+        self._store = PleStore(
             paths=list(source.paths),
             extent_file=list(source.extent_file),
             extent_base=list(source.extent_base),
@@ -155,7 +155,9 @@ def __init__(
         self._eager_pinned.zero_()  # the warmup prefill stages nothing and reads whatever sits here
         self._eager_dev = torch.empty(eager_bytes, dtype=torch.uint8, device=self._device)
         # probe picks flag-sync (graph WAITs at the consume, host fills then signals) or launch-gating
-        self._wait_sync = self._probe_wait_sync(os.getenv(_SYNC_ENV, "auto"))
+        from freetoken.kernel.row_store import probe_wait_sync
+
+        self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device)
         # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap
         self._flag = alloc_pinned_tensor(1, dtype=torch.int64)
         self._flag.zero_()
@@ -164,25 +166,6 @@ def __init__(
         sync = "wait-sync" if self._wait_sync else "launch-gating"
         logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}")
 
-    def _probe_wait_sync(self, mode: str) -> bool:
-        from freetoken.kernel import _ple_store
-
-        if mode == "gate":
-            return False
-        scratch = alloc_pinned_tensor(1, dtype=torch.int64)
-        scratch.zero_()
-        stream = torch.cuda.current_stream(self._device)
-        ok = (
-            _ple_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0
-            and _ple_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0
-        )
-        if ok:
-            stream.synchronize()
-            ok = int(scratch[0]) == 7
-        if mode == "wait" and not ok:
-            raise RuntimeError("FREETOKEN_PLE_SYNC=wait but stream memops are unavailable")
-        return ok
-
     # ---------------- host side (engine thread, before the forward launches) ----------------
 
     def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
@@ -222,10 +205,10 @@ def _complete() -> None:
                                 for r, t in zip(reqs, tokens)]
                         self.fill(runs, graph=True)
                     except BaseException:
-                        from freetoken.kernel import _ple_store
+                        from freetoken.kernel.row_store import signal
 
                         # unblock the stream before surfacing; the step's output is discarded
-                        _ple_store.signal_flag(self._flag.data_ptr())
+                        signal(self._flag)
                         raise
 
                 return _complete
@@ -260,11 +243,9 @@ def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torc
         rows = row_ids.shape[0]
         capturing = torch.cuda.is_current_stream_capturing()
         if capturing and self._wait_sync:
-            from freetoken.kernel import _ple_store
+            from freetoken.kernel.row_store import wait_reset
 
-            _ple_store.memop_wait_reset(
-                torch.cuda.current_stream(self._device).cuda_stream, self._flag.data_ptr()
-            )
+            wait_reset(torch.cuda.current_stream(self._device), self._flag)
         pinned, dev = (
             (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned, self._eager_dev)
         )
diff --git a/setup.py b/setup.py
index d8e705ab69..7132eea088 100644
--- a/setup.py
+++ b/setup.py
@@ -64,12 +64,12 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
             libraries=["cudart"],
             extra_compile_args=["-O3", "-std=c++17", "-pthread"],
         ),
-        # --ple-backend disk row store; Linux-only until the TableFile/BatchReader seams grow Windows bodies
+        # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies
         *([
             CppExtension(
-                name="freetoken.kernel._ple_store",
+                name="freetoken.kernel._row_store",
                 sources=[
-                    "python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp",
+                    "python/freetoken/kernel/csrc/row_store/row_store_ext.cpp",
                 ],
                 extra_compile_args=["-O3", "-std=c++17"],
             )
diff --git a/tests/kernels/test_row_store.py b/tests/kernels/test_row_store.py
new file mode 100644
index 0000000000..d4a195026a
--- /dev/null
+++ b/tests/kernels/test_row_store.py
@@ -0,0 +1,87 @@
+"""The hash-agnostic ``RowStore`` (``kernel/csrc/row_store``): byte-exact staged rows across extents and files,
+dedup of repeated ids, a destination stride, and bounds checks before any I/O. The PLE hash wrapper on top of
+it is covered by ``tests/models/qwen4_exp/test_ple_disk.py``."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+_row_store = pytest.importorskip("freetoken.kernel._row_store")
+
+ROW_BYTES = 264  # deliberately unaligned rows exercise reads across page boundaries
+
+
+def _table(rows: int, seed: int) -> torch.Tensor:
+    gen = torch.Generator().manual_seed(seed)
+    return torch.randint(0, 256, (rows, ROW_BYTES), dtype=torch.uint8, generator=gen)
+
+
+def _stage(store, ids: torch.Tensor, dst_stride: int = 0) -> torch.Tensor:
+    stride = dst_stride or store.row_bytes
+    staging = torch.zeros(ids.numel() * stride, dtype=torch.uint8)
+    store.stage_rows(ids.data_ptr(), ids.numel(), staging.data_ptr(), dst_stride)
+    store.flush(0)
+    return staging.view(ids.numel(), stride)
+
+
+@pytest.mark.parametrize("use_io_uring", [True, False])
+def test_rows_come_back_bitwise_across_extents_and_files(tmp_path, use_io_uring):
+    # two files; file 0 carries extents 0 and 2 behind an unrelated header, file 1 carries extent 1
+    per_extent = 1000
+    table = _table(3 * per_extent, seed=1)
+    header = b"safetensors-like header, not a row" + bytes(4096 - 34)
+    f0 = tmp_path / "shard0.bin"
+    f1 = tmp_path / "shard1.bin"
+    f0.write_bytes(header + table[:per_extent].numpy().tobytes() + table[2 * per_extent :].numpy().tobytes())
+    f1.write_bytes(table[per_extent : 2 * per_extent].numpy().tobytes())
+    store = _row_store.RowStore(
+        paths=[str(f0), str(f1)],
+        extent_file=[0, 1, 0],
+        extent_base=[len(header), 0, len(header) + per_extent * ROW_BYTES],
+        rows_per_extent=per_extent,
+        row_bytes=ROW_BYTES,
+        row_stride=ROW_BYTES,
+        use_io_uring=use_io_uring,
+    )
+    assert store.total_rows == 3 * per_extent and store.row_bytes == ROW_BYTES
+
+    ids = torch.tensor([0, 999, 1000, 1999, 2000, 2999, 1234, 1234, 42, 0], dtype=torch.int64)
+    got = _stage(store, ids)
+    assert torch.equal(got, table[ids])
+
+    # a wider destination stride leaves the gap untouched
+    got = _stage(store, ids, dst_stride=ROW_BYTES + 40)
+    assert torch.equal(got[:, :ROW_BYTES], table[ids]) and int(got[:, ROW_BYTES:].sum()) == 0
+
+    # a fill larger than the reader's in-flight window pipelines through
+    many = torch.randint(0, 3 * per_extent, (5000,), dtype=torch.int64, generator=torch.Generator().manual_seed(2))
+    assert torch.equal(_stage(store, many), table[many])
+
+    # flushing nothing still signals the flag
+    flag = torch.zeros(1, dtype=torch.int64)
+    store.flush(flag.data_ptr())
+    assert int(flag[0]) == 1
+
+
+def test_geometry_and_bounds_are_checked_up_front(tmp_path):
+    table = _table(64, seed=3)
+    path = tmp_path / "t.bin"
+    path.write_bytes(table.numpy().tobytes())
+    make = lambda **kw: _row_store.RowStore(
+        paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=64,
+        row_bytes=ROW_BYTES, row_stride=ROW_BYTES, use_io_uring=False, **kw,
+    )
+    store = make()
+    ids = torch.tensor([3, 64], dtype=torch.int64)
+    staging = torch.zeros(2 * ROW_BYTES, dtype=torch.uint8)
+    with pytest.raises(IndexError):
+        store.stage_rows(ids.data_ptr(), 2, staging.data_ptr(), 0)
+    with pytest.raises(RuntimeError, match="overlap"):
+        store.stage_rows(ids[:1].data_ptr(), 1, staging.data_ptr(), ROW_BYTES - 1)
+    with pytest.raises(RuntimeError, match="extent needs"):
+        _row_store.RowStore(paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=65,
+                            row_bytes=ROW_BYTES, row_stride=ROW_BYTES, use_io_uring=False)
+    with pytest.raises(RuntimeError, match="exceeds a page"):
+        _row_store.RowStore(paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=1,
+                            row_bytes=4097, row_stride=4097, use_io_uring=False)
diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py
index fa11807390..1112210b84 100644
--- a/tests/models/qwen4_exp/test_ple_disk.py
+++ b/tests/models/qwen4_exp/test_ple_disk.py
@@ -13,7 +13,7 @@
 from .common import EOS, hash_constants, requires_cuda, toy_hf_config
 from .test_ple import _meta
 
-_ple_store = pytest.importorskip("freetoken.kernel._ple_store")
+_row_store = pytest.importorskip("freetoken.kernel._row_store")
 
 _KEY_PREFIX = "model.layers.1.ple.ple_embedding.ngram_embedding"
 
@@ -43,7 +43,7 @@ def _make_store(tmp_path, *, write=True, use_io_uring=True):
     path = tmp_path / "ple-table.bin"
     if write:
         path.write_bytes(table.numpy().tobytes())
-    store = _ple_store.PleStore(
+    store = _row_store.PleStore(
         paths=[str(path)],
         extent_file=[0],
         extent_base=[0],
@@ -175,10 +175,10 @@ def test_layouts_readers_and_errors(tmp_path):
         multipliers=[3, 5, 7], head_vocab_sizes=sizes, head_offsets=offsets,
         eos_token_id=eos,
     )
-    ref = _ple_store.PleStore(
+    ref = _row_store.PleStore(
         paths=[str(flat)], extent_file=[0, 0, 0, 0], extent_base=[0, nb, 2 * nb, 3 * nb], **kwargs
     )
-    multi = _ple_store.PleStore(
+    multi = _row_store.PleStore(
         paths=[str(fa), str(fb)], extent_file=[0, 1, 0, 1],
         extent_base=[1231, 0, 1231 + nb + 77, nb + 4095], **kwargs,
     )

From 0d652e73a452d014ac5441a15baa75348e9fcb0a Mon Sep 17 00:00:00 2001
From: ZihaoMu <zmu@amd.com>
Date: Sat, 26 Sep 2026 08:49:45 +0800
Subject: [PATCH 053/249] feat(rocm): add RDNA3 and RDNA4 runtime foundation
 (#132)

* feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103)

- Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents
- Update pinned_tensor.cpp to compile under both nvcc and hipcc
- Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family)
- Guard NVIDIA arch checks to return None on ROCm
- Skip nvcc version check in _toolchain.py when on ROCm
- Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch)
- Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm
- Add is_rocm() and driver_hip_version() in backend.py
- Add rocm-smi fallback in __main__.py for clangd generation
- Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds,
  Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache
- Add AMD ROCm classifier in pyproject.toml

* fix(rocm): fall back for single-bank fast index copy

* test(rocm): cover fast index copy fallback

* test(rocm): use explicit fast copy module import

* fix(rocm): complete RDNA3 runtime path

* fix(rocm): preserve CPU MoE implementation

* build(rocm): support RDNA4 discovery and modular SDKs

* fix(rocm): complete modular SDK runtime support

* test(rocm): harden foundation platform contracts

* fix(rocm): pin jit target and split install guide

* perf(kernel): restore ROCm non-temporal index copy

---------

Co-authored-by: bouclem <bouclement2012@gmail.com>
Co-authored-by: Yuu <206304251+nekomario28@users.noreply.github.com>
Co-authored-by: LZ-QWQ <lz1830837889@163.com>
---
 docs/install.md                               |   2 +
 docs/install_amd.md                           |  39 +++++
 pyproject.toml                                |   6 +-
 python/freetoken/kernel/__main__.py           |  35 ++--
 python/freetoken/kernel/_toolchain.py         |   9 +-
 python/freetoken/kernel/backend.py            |  18 ++
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       |   2 +-
 .../csrc/include/freetoken/hip_compat.h       | 162 ++++++++++++++++++
 .../kernel/csrc/include/freetoken/utils.cuh   |   6 +
 .../kernel/csrc/jit/fast_index_copy.cuh       |  55 ++++--
 .../freetoken/kernel/csrc/pinned_tensor.cpp   |   2 +-
 python/freetoken/kernel/pynccl.py             |   1 +
 python/freetoken/kernel/triton/activation.py  |  10 +-
 python/freetoken/kernel/triton/e4m3_compat.py |   4 +
 python/freetoken/kernel/triton/norm.py        |   8 +-
 python/freetoken/kernel/utils.py              | 145 ++++++++++++++--
 python/freetoken/utils/__init__.py            |   4 +
 python/freetoken/utils/arch.py                |  52 ++++++
 setup.py                                      |  77 +++++++--
 tests/kernels/test_e4m3_compat.py             |  23 +++
 tests/kernels/test_pinned_tensor.py           |   7 +-
 tests/kernels/test_rocm_launch_kwargs.py      |  51 ++++++
 tests/utils/test_rocm_arch.py                 | 127 ++++++++++++++
 23 files changed, 780 insertions(+), 65 deletions(-)
 create mode 100644 docs/install_amd.md
 create mode 100644 python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
 create mode 100644 tests/kernels/test_rocm_launch_kwargs.py
 create mode 100644 tests/utils/test_rocm_arch.py

diff --git a/docs/install.md b/docs/install.md
index 9438c21a67..f474216bf0 100644
--- a/docs/install.md
+++ b/docs/install.md
@@ -1,5 +1,7 @@
 # Install
 
+AMD users: see the [AMD ROCm installation guide (WIP)](install_amd.md).
+
 ## Requirements
 
 - Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13)
diff --git a/docs/install_amd.md b/docs/install_amd.md
new file mode 100644
index 0000000000..d8b0c1bcb1
--- /dev/null
+++ b/docs/install_amd.md
@@ -0,0 +1,39 @@
+# AMD ROCm installation (WIP)
+
+AMD support is experimental and remains a work in progress. See the
+[AMD Support roadmap](https://github.com/FlashML-org/FreeToken/issues/541) for
+the current integration and qualification status.
+
+## Requirements
+
+- Linux x86_64
+- AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`)
+- ROCm 7.14
+- Python >= 3.10
+
+## Install from source
+
+Use an official ROCm PyTorch image whose PyTorch version satisfies the project's
+`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is:
+
+```bash
+VIDEO_GID="$(getent group video | cut -d: -f3)"
+RENDER_GID="$(getent group render | cut -d: -f3)"
+docker run --rm -it \
+  --device=/dev/kfd --device=/dev/dri \
+  --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \
+  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+  -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \
+  -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \
+  rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash
+```
+
+Inside the container, preserve the ROCm-enabled PyTorch already supplied by the
+image and disable build isolation so it is also used to compile the extensions:
+
+```bash
+python -m pip install --no-build-isolation -e .
+```
+
+Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the
+actual target reported by `rocminfo`.
diff --git a/pyproject.toml b/pyproject.toml
index 25cb204bb4..8d00d2bbc6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,6 +23,7 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Operating System :: POSIX :: Linux",
     "Environment :: GPU :: NVIDIA CUDA",
+    "Environment :: GPU :: AMD ROCm",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -59,7 +60,10 @@ dependencies = [
     "torchvision>=0.26,<0.27",
     "tqdm>=4.66,<5",
     "transformers>=5.16,<5.17",
-    "triton==3.6.0; platform_system == 'Linux'",
+    # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its
+    # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
+    # runtime-specific wheel selected by the PyTorch distribution.
+    "triton>=3.6,<3.8; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 
diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py
index 7be541a67c..5b66d484df 100644
--- a/python/freetoken/kernel/__main__.py
+++ b/python/freetoken/kernel/__main__.py
@@ -6,27 +6,32 @@ def generate_clangd():
     import subprocess
 
     from freetoken.kernel.utils import DEFAULT_INCLUDE
-    from freetoken.utils import init_logger
+    from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm
     from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
     logger = init_logger(__name__)
     logger.info("Generating .clangd file...")
     include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
-    status = subprocess.run(
-        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
-        capture_output=True,
-        check=True,
-    )
-    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
-    major, minor = compute_cap.split(".")
+
+    # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries.
+    if is_rocm():
+        arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"]
+    else:
+        try:
+            status = subprocess.run(
+                args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+                capture_output=True,
+                check=True,
+            )
+            compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+            major, minor = compute_cap.split(".")
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            import torch
+
+            major, minor = torch.cuda.get_device_capability()
+        arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"]
     compile_flags = ",\n    ".join(
-        [
-            "-xcuda",
-            f"--cuda-gpu-arch=sm_{major}{minor}",
-            "-std=c++20",
-            "-Wall",
-            "-Wextra",
-        ]
+        arch_flags + ["-std=c++20", "-Wall", "-Wextra"]
         + [f"-isystem{path}" for path in include_paths]
     )
     clangd_content = f"""
diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py
index b49cebbb43..93d0322151 100644
--- a/python/freetoken/kernel/_toolchain.py
+++ b/python/freetoken/kernel/_toolchain.py
@@ -1,4 +1,4 @@
-"""CUDA toolchain/torch consistency checks.
+"""CUDA/HIP toolchain/torch consistency checks.
 
 Standalone on purpose: setup.py and the kernel-cache build backend load this
 file by path, so it must not import the freetoken package.
@@ -16,6 +16,11 @@
 _TRUE_VALUES = {"1", "true", "yes", "on"}
 
 
+def _is_rocm() -> bool:
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
 def _nvcc_path() -> str | None:
     from torch.utils.cpp_extension import CUDA_HOME
 
@@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None:
     nvcc-built binaries link libcudart.so.<nvcc major>; at runtime only the
     torch wheel's own CUDA runtime is guaranteed to be loadable.
     """
+    if _is_rocm():
+        return  # ROCm uses hipcc, not nvcc
     if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES:
         return
     torch_major = torch_cuda_major()
diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py
index 2a886cf77d..6d1b4a9f8e 100644
--- a/python/freetoken/kernel/backend.py
+++ b/python/freetoken/kernel/backend.py
@@ -47,6 +47,24 @@ def device_capability() -> tuple[int, int]:
     return (int(major), int(minor))
 
 
+@functools.cache
+def is_rocm() -> bool:
+    """True when torch is built for ROCm (AMD GPU)."""
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
+@functools.cache
+def driver_hip_version() -> int | None:
+    """ROCm driver version, or None if undetermined."""
+    # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used.
+    try:
+        from freetoken.kernel.pinned import _load_pinned_extension
+        return int(_load_pinned_extension().driver_cuda_version()) or None
+    except Exception:
+        return None
+
+
 @functools.cache
 def driver_cuda_version() -> int | None:
     """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0),
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index 8349123bd6..e2cb0d9c33 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -29,7 +29,7 @@
 #include <thread>
 #include <vector>
 
-#include <cuda_runtime_api.h>
+#include <freetoken/hip_compat.h>
 #include <torch/extension.h>
 
 #if defined(__linux__)
diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
new file mode 100644
index 0000000000..1898b41189
--- /dev/null
+++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
@@ -0,0 +1,162 @@
+#pragma once
+
+// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so
+// the same C++ source compiles under both nvcc and hipcc.  Include this instead
+// of <cuda_runtime_api.h> directly when the file needs the runtime API.
+//
+// On NVIDIA platforms the CUDA headers are included as-is and every macro below
+// resolves to the original CUDA symbol, so there is zero overhead.
+//
+// Supported ROCm targets:
+//   gfx1100 — RX 7900 XTX / XT
+//   gfx1101 — RX 7900 GRE
+//   gfx1102 — RX 7700 / XT
+//   gfx1103 — RX 7600 / XT
+//   gfx1200 — RX 9060 family
+//   gfx1201 — RX 9070 family / Radeon AI PRO R9700
+
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+
+#define FREETOKEN_USE_ROCM 1
+
+// --- HIP runtime headers ---
+#include <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+// --- API name mapping (CUDA -> HIP) ---
+// HIP already defines most cuda* names as macros that expand to hip* equivalents
+// via hip_runtime.h, but a few are missing or differ in signature.  Define them
+// here so call-sites stay unchanged.
+
+#ifndef cudaSuccess
+#define cudaSuccess hipSuccess
+#endif
+
+#ifndef cudaError_t
+#define cudaError_t hipError_t
+#endif
+
+#ifndef cudaGetErrorString
+#define cudaGetErrorString hipGetErrorString
+#endif
+
+#ifndef cudaGetLastError
+#define cudaGetLastError hipGetLastError
+#endif
+
+#ifndef cudaMallocHost
+#define cudaMallocHost hipHostMalloc
+#endif
+
+#ifndef cudaFreeHost
+#define cudaFreeHost hipHostFree
+#endif
+
+#ifndef cudaHostAlloc
+#define cudaHostAlloc hipHostMalloc
+#endif
+
+#ifndef cudaHostRegister
+#define cudaHostRegister hipHostRegister
+#endif
+
+#ifndef cudaHostRegisterPortable
+#define cudaHostRegisterPortable hipHostRegisterPortable
+#endif
+
+#ifndef cudaHostRegisterMapped
+#define cudaHostRegisterMapped hipHostRegisterMapped
+#endif
+
+#ifndef cudaHostAllocPortable
+#define cudaHostAllocPortable hipHostMallocPortable
+#endif
+
+#ifndef cudaHostAllocMapped
+#define cudaHostAllocMapped hipHostMallocMapped
+#endif
+
+#ifndef cudaHostGetDevicePointer
+#define cudaHostGetDevicePointer hipHostGetDevicePointer
+#endif
+
+#ifndef cudaGetDevice
+#define cudaGetDevice hipGetDevice
+#endif
+
+#ifndef cudaDriverGetVersion
+#define cudaDriverGetVersion hipDriverGetVersion
+#endif
+
+#ifndef cudaDeviceGetAttribute
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#endif
+
+#ifndef cudaDevAttrUnifiedAddressing
+#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing
+#endif
+
+#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem
+// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux).
+// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute.
+#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing
+#endif
+
+#ifndef cudaFuncSetAttribute
+#define cudaFuncSetAttribute hipFuncSetAttribute
+#endif
+
+#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
+#endif
+
+#ifndef cudaLaunchKernelEx
+// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP.
+#define cudaLaunchKernelEx hipLaunchKernelEx
+#endif
+
+#ifndef cudaLaunchConfig_t
+#define cudaLaunchConfig_t hipLaunchConfig_t
+#endif
+
+#ifndef cudaLaunchAttribute
+#define cudaLaunchAttribute hipLaunchAttribute
+#endif
+
+#ifndef cudaLaunchAttributeProgrammaticStreamSerialization
+// PDL (Programmatic Dependent Launch) is NVIDIA-specific.
+// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency.
+#define cudaLaunchAttributeProgrammaticStreamSerialization 0
+#endif
+
+#ifndef cudaStream_t
+#define cudaStream_t hipStream_t
+#endif
+
+#ifndef cudaStreamSynchronize
+#define cudaStreamSynchronize hipStreamSynchronize
+#endif
+
+#ifndef cudaLaunchHostFunc
+#define cudaLaunchHostFunc hipLaunchHostFunc
+#endif
+
+#ifndef CUDART_CB
+#define CUDART_CB
+#endif
+
+#ifndef __grid_constant__
+#define __grid_constant__
+#endif
+
+#ifndef dim3
+// HIP already provides dim3; this is a no-op guard.
+#endif
+
+#else  // NVIDIA CUDA path
+
+#define FREETOKEN_USE_ROCM 0
+
+#include <cuda_runtime_api.h>
+
+#endif
diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
index 8e917832c1..72a4959698 100644
--- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
+++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <freetoken/hip_compat.h>
 #include <freetoken/utils.h>
 
 #include <dlpack/dlpack.h>
@@ -115,6 +116,10 @@ public:
   }
 
   auto with_attr(bool use_pdl) -> LaunchKernel & {
+#if FREETOKEN_USE_ROCM
+    (void)use_pdl;
+    m_config.numAttrs = 0;
+#else
     if (use_pdl) {
       m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization;
       m_attr_cache.val.programmaticStreamSerializationAllowed = 1;
@@ -123,6 +128,7 @@ public:
     } else {
       m_config.numAttrs = 0;
     }
+#endif
     return *this;
   }
 
diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
index bb83c23ed2..43f3fc1ba5 100644
--- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
+++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
@@ -20,6 +20,13 @@ struct device_vec {
 
 namespace details {
 
+#if FREETOKEN_USE_ROCM
+// Use native Clang vector types so each __builtin_nontemporal_load/store
+// maps to one vectorized 8B/16B global memory instruction.
+using native_uint2 = uint32_t __attribute__((ext_vector_type(2)));
+using native_uint4 = uint32_t __attribute__((ext_vector_type(4)));
+#endif
+
 template <std::size_t kUnit>
 inline constexpr auto get_mem_package() {
     if constexpr (kUnit == 16) {
@@ -34,40 +41,68 @@ inline constexpr auto get_mem_package() {
 }
 
 __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+#if FREETOKEN_USE_ROCM
+    return uint1{__builtin_nontemporal_load(&src->x)};
+#else
     uint32_t tmp;
     asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
     return uint1{tmp};
+#endif
 }
 
 __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
+#if FREETOKEN_USE_ROCM
+    const auto value = __builtin_nontemporal_load(reinterpret_cast<const native_uint2*>(src));
+    return __builtin_bit_cast(uint2, value);
+#else
     uint32_t tmp0, tmp1;
     asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
     return uint2{tmp0, tmp1};
+#endif
 }
 
 __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
+#if FREETOKEN_USE_ROCM
+    const auto value = __builtin_nontemporal_load(reinterpret_cast<const native_uint4*>(src));
+    return __builtin_bit_cast(uint4, value);
+#else
     uint32_t tmp0, tmp1, tmp2, tmp3;
     asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
     return uint4{tmp0, tmp1, tmp2, tmp3};
+#endif
 }
 
 __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
+#if FREETOKEN_USE_ROCM
+    __builtin_nontemporal_store(value.x, &dst->x);
+#else
     uint32_t tmp = value.x;
     asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
+#endif
 }
 
 __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
+#if FREETOKEN_USE_ROCM
+    __builtin_nontemporal_store(
+        __builtin_bit_cast(native_uint2, value), reinterpret_cast<native_uint2*>(dst));
+#else
     uint32_t tmp0 = value.x;
     uint32_t tmp1 = value.y;
     asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
+#endif
 }
 
 __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
+#if FREETOKEN_USE_ROCM
+    __builtin_nontemporal_store(
+        __builtin_bit_cast(native_uint4, value), reinterpret_cast<native_uint4*>(dst));
+#else
     uint32_t tmp0 = value.x;
     uint32_t tmp1 = value.y;
     uint32_t tmp2 = value.z;
     uint32_t tmp3 = value.w;
     asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
+#endif
 }
 
 __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) {
@@ -75,7 +110,7 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag
     auto* flag = reinterpret_cast<int*>(const_cast<int32_t*>(flag_ptr));
     uint32_t sleep_ns = 128;
     while (atomicAdd(flag, 0) > 0) {
-#if __CUDA_ARCH__ >= 700
+#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700
         __nanosleep(sleep_ns);
 #endif
         sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048;
@@ -147,7 +182,7 @@ inline bool host_ptr_identity() {
 }
 
 inline void* device_alias(void* ptr, DLDevice dev) {
-    if (dev.device_type == kDLCUDA || host_ptr_identity()) {
+    if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) {
         return ptr;
     }
     void* mapped = nullptr;
@@ -269,7 +304,7 @@ inline auto get_sync_flag_ptr(
     auto flag_dtype = host::SymbolicDType{};
     host::TensorMatcher({1})
         .with_dtype<int32_t>(flag_dtype)
-        .with_device<kDLCUDA>(device)
+        .with_device<kDLCUDA, kDLROCM>(device)
         .verify(sync_flag);
     return static_cast<int32_t*>(sync_flag.data_ptr());
 }
@@ -344,17 +379,17 @@ struct FastIndexCopyKernel {
 
         TensorMatcher({-1, D})
         .with_dtype(data_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLCUDA, kDLROCM, kDLCUDAHost, kDLROCMHost, kDLCPU>()
         .verify(src);
 
         TensorMatcher({-1, D})
         .with_dtype(data_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLCUDA, kDLROCM, kDLCUDAHost, kDLROCMHost, kDLCPU>()
         .verify(dst);
 
         TensorMatcher({L})
         .with_dtype<int32_t, int64_t>(indices_dtype)
-        .with_device<kDLCUDA>(device)
+        .with_device<kDLCUDA, kDLROCM>(device)
         .verify(src_indices)
         .verify(dst_indices);
 
@@ -363,7 +398,7 @@ struct FastIndexCopyKernel {
             const auto num_indices_tensor = num_indices.value();
             TensorMatcher({1})
                 .with_dtype<int64_t>(num_indices_dtype)
-                .with_device<kDLCUDA>(device)
+                .with_device<kDLCUDA, kDLROCM>(device)
                 .verify(num_indices_tensor);
 
             num_indices_data_ptr = static_cast<const int64_t*>(num_indices_tensor.data_ptr());
@@ -529,14 +564,14 @@ struct MultiIndexCopyKernel {
         auto indices_dtype = SymbolicDType{};
         auto num_indices_dtype = SymbolicDType{};
 
-        TensorMatcher({B}).with_dtype<int64_t>(ptr_dtype).with_device<kDLCUDA>(device)
+        TensorMatcher({B}).with_dtype<int64_t>(ptr_dtype).with_device<kDLCUDA, kDLROCM>(device)
             .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes);
-        TensorMatcher({L}).with_dtype<int32_t, int64_t>(indices_dtype).with_device<kDLCUDA>(device)
+        TensorMatcher({L}).with_dtype<int32_t, int64_t>(indices_dtype).with_device<kDLCUDA, kDLROCM>(device)
             .verify(dst_indices).verify(src_indices);
 
         const int64_t* valid_length = nullptr;
         if (num_indices.has_value()) {
-            TensorMatcher({1}).with_dtype<int64_t>(num_indices_dtype).with_device<kDLCUDA>(device)
+            TensorMatcher({1}).with_dtype<int64_t>(num_indices_dtype).with_device<kDLCUDA, kDLROCM>(device)
                 .verify(num_indices.value());
             valid_length = static_cast<const int64_t*>(num_indices.value().data_ptr());
         }
diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp
index c3947adfa3..9355f57aa5 100644
--- a/python/freetoken/kernel/csrc/pinned_tensor.cpp
+++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp
@@ -1,5 +1,5 @@
 #include <cstdint>
-#include <cuda_runtime_api.h>
+#include <freetoken/hip_compat.h>
 #include <torch/extension.h>
 
 namespace {
diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py
index 23ea57351d..71bde7342a 100644
--- a/python/freetoken/kernel/pynccl.py
+++ b/python/freetoken/kernel/pynccl.py
@@ -27,6 +27,7 @@ def get_buffer(self) -> int: ...
 
 @functools.cache
 def _load_nccl_module() -> Module:
+    # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD.
     return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"])
 
 
diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py
index 0e4d28c8fc..9645960d11 100644
--- a/python/freetoken/kernel/triton/activation.py
+++ b/python/freetoken/kernel/triton/activation.py
@@ -20,8 +20,9 @@
 import triton.language as tl
 from triton.language.extra import libdevice
 from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents
+from triton.language import target_info
 
-from freetoken.utils.arch import is_sm90_supported
+from freetoken.utils.arch import is_rocm, is_sm90_supported
 
 SILU = 0
 GELU = 1
@@ -50,6 +51,8 @@ def _pdl_supported() -> bool:
 
 @triton.jit
 def _fast_tanh(x):
+    if target_info.is_hip():
+        return libdevice.tanh(x)
     # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh.
     return tl.inline_asm_elementwise(
         "tanh.approx.f32 $0, $1;", "=f,f", [x],
@@ -59,6 +62,8 @@ def _fast_tanh(x):
 
 @triton.jit
 def _fast_ex2(x):
+    if target_info.is_hip():
+        return libdevice.exp2(x)
     # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu.
     return tl.inline_asm_elementwise(
         "ex2.approx.f32 $0, $1;", "=f,f", [x],
@@ -141,8 +146,9 @@ def _act_and_mul(
     block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512)
     num_stages = 2 if block_d == 1024 else 3
     _act_and_mul_kernel[grid](
-        o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl,
+        o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl,
         BLOCK_D=block_d, num_warps=4, num_stages=num_stages,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
     )
     return out
 
diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py
index 1d9f744ce1..6ea0c7537d 100644
--- a/python/freetoken/kernel/triton/e4m3_compat.py
+++ b/python/freetoken/kernel/triton/e4m3_compat.py
@@ -59,6 +59,10 @@ def e4m3_native() -> bool:
     if _native is None:
         if FORCE_EMU:
             _native = False
+        elif torch.version.hip is not None:
+            # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA
+            # compute capability and must not select the native fp8e4nv path.
+            _native = False
         else:
             from freetoken.gpu_select import assigned_visible_gpu
 
diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py
index 3f95c29f0d..9071e1dfff 100644
--- a/python/freetoken/kernel/triton/norm.py
+++ b/python/freetoken/kernel/triton/norm.py
@@ -30,7 +30,7 @@
 import triton.language as tl
 from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
 
-from freetoken.utils.arch import is_sm90_supported
+from freetoken.utils.arch import is_rocm, is_sm90_supported
 
 _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])}
 
@@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool):
     pdl = contig and is_sm90_supported()
     _rmsnorm_kernel[(A, B)](
         out, input, weight, eps, H, sxa, sxb, soa, sob,
-        CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma,
+        CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
         num_warps=_num_warps(A * B), num_stages=1,
     )
     return out
@@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool):
     pdl = contig and is_sm90_supported()
     _fused_add_rmsnorm_kernel[(A, B)](
         input, residual, weight, eps, H, sxa, sxb, sra, srb,
-        CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma,
+        CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
         num_warps=_num_warps(A * B), num_stages=1,
     )
 
diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py
index cdd765c7d6..9af61526ae 100644
--- a/python/freetoken/kernel/utils.py
+++ b/python/freetoken/kernel/utils.py
@@ -5,6 +5,7 @@
 import os
 import pathlib
 import re
+from functools import cache
 from typing import TYPE_CHECKING, Iterator, List, NamedTuple, Tuple, TypeAlias, Union
 
 if TYPE_CHECKING:
@@ -20,15 +21,23 @@
 DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")]
 DEFAULT_CFLAGS = ["-std=c++20", "-O3"]
 DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"]
+DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"]
 DEFAULT_LDFLAGS = []
 
 
-ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST"
+CUDA_ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST"
+ROCM_ARCH_LIST_ENV = "TVM_FFI_ROCM_ARCH_LIST"
+
+
+def _is_rocm() -> bool:
+    import torch
+
+    return getattr(torch.version, "hip", None) is not None
 
 
 def _cuda_arch_list() -> List[str]:
     """Archs a CUDA build targets: the AOT build's TVM_FFI_CUDA_ARCH_LIST, else the GPU this process is bound to."""
-    arch_list = os.getenv(ARCH_LIST_ENV, "").split()
+    arch_list = os.getenv(CUDA_ARCH_LIST_ENV, "").split()
     if arch_list:
         return arch_list
     import torch
@@ -39,17 +48,34 @@ def _cuda_arch_list() -> List[str]:
     return [f"{major}.{minor}"]
 
 
+def _rocm_arch_list() -> List[str]:
+    """The single ROCm GPU this process is bound to, or an explicit cross-compile target."""
+    from freetoken.utils.arch import get_rocm_gfx_arch
+
+    arch = get_rocm_gfx_arch()
+    if arch is None:
+        raise RuntimeError(
+            "Could not determine the bound ROCm GPU architecture; select a visible GPU "
+            "or set FREETOKEN_ROCM_ARCH for cross-compilation"
+        )
+    return [arch]
+
+
 @contextlib.contextmanager
-def _pin_tvm_ffi_arch_ctx(arch_list: List[str]) -> Iterator[None]:
-    """Hand tvm-ffi the arch list through its env var for one build. Left unset, tvm-ffi asks nvidia-smi and takes the first GPU listed, which under --gpu or CUDA_VISIBLE_DEVICES on a mixed box is not the bound one."""
-    if os.getenv(ARCH_LIST_ENV) or not arch_list:
+def _pin_tvm_ffi_arch_ctx(arch_list: List[str], env_var: str) -> Iterator[None]:
+    """Pin tvm-ffi to the resolved targets for one build, then restore the caller's environment."""
+    if not arch_list:
         yield
         return
-    os.environ[ARCH_LIST_ENV] = " ".join(arch_list)
+    old_arch_list = os.environ.get(env_var)
+    os.environ[env_var] = " ".join(arch_list)
     try:
         yield
     finally:
-        os.environ.pop(ARCH_LIST_ENV, None)
+        if old_arch_list is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_arch_list
 
 
 def _cuda_cflags(extra: List[str], arch_list: List[str]) -> List[str]:
@@ -63,6 +89,61 @@ def _rank(a: str) -> int:
         cc = max(arch_list, key=_rank).rstrip("a").replace(".", "")
         flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"]
     return flags
+
+
+def _hip_cflags(extra: List[str], arch_list: List[str]) -> List[str]:
+    """HIP flags for a kernel build on ROCm."""
+    # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size).
+    return DEFAULT_HIP_CFLAGS + extra + [f"--offload-arch={arch}" for arch in arch_list]
+
+
+@cache
+def _rocm_link_flags() -> List[str]:
+    """Make ROCm's runtime library discoverable to JIT link commands.
+
+    Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``.
+    ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI
+    still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via
+    an explicit linker search path without modifying the Python environment.
+    """
+    candidates: list[pathlib.Path] = []
+    if os.getenv("ROCM_HOME"):
+        candidates.append(pathlib.Path(os.environ["ROCM_HOME"]))
+    try:
+        from torch.utils.cpp_extension import ROCM_HOME
+
+        if ROCM_HOME:
+            candidates.append(pathlib.Path(ROCM_HOME))
+    except ImportError:
+        pass
+    spec = importlib.util.find_spec("_rocm_sdk_core")
+    if spec and spec.submodule_search_locations:
+        candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations))))
+    candidates.append(pathlib.Path("/opt/rocm"))
+
+    for rocm_home in dict.fromkeys(candidates):
+        library_dir = rocm_home / "lib"
+        unversioned = library_dir / "libamdhip64.so"
+        link_dir = library_dir
+        if not unversioned.exists():
+            versioned = sorted(library_dir.glob("libamdhip64.so.*"))
+            if not versioned:
+                continue
+            link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib"
+            link_dir.mkdir(parents=True, exist_ok=True)
+            compat_link = link_dir / "libamdhip64.so"
+            if not compat_link.exists() and not compat_link.is_symlink():
+                try:
+                    compat_link.symlink_to(versioned[-1])
+                except FileExistsError:
+                    # Multiple tensor-parallel ranks may prepare the same cache.
+                    pass
+
+        return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"]
+
+    raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking")
+
+
 CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool]
 
 
@@ -223,12 +304,16 @@ def load_aot(
     if prebuilt is not None:
         return prebuilt
 
+    is_rocm = _is_rocm()
     arch_list: List[str] = []
     if cuda_files:
-        from freetoken.kernel._toolchain import check_nvcc_matches_torch
+        if is_rocm:
+            arch_list = _rocm_arch_list()
+        else:
+            from freetoken.kernel._toolchain import check_nvcc_matches_torch
 
-        check_nvcc_matches_torch()
-        arch_list = _cuda_arch_list()
+            check_nvcc_matches_torch()
+            arch_list = _cuda_arch_list()
 
     from tvm_ffi.cpp import load
 
@@ -242,14 +327,23 @@ def load_aot(
     cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files]
     cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files]
 
-    with _pin_tvm_ffi_arch_ctx(arch_list):
+    if is_rocm:
+        cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list)
+        runtime_ldflags = _rocm_link_flags()
+        arch_list_env = ROCM_ARCH_LIST_ENV
+    else:
+        cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list)
+        runtime_ldflags = []
+        arch_list_env = CUDA_ARCH_LIST_ENV
+
+    with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env):
         return load(
             name,
             cpp_files=cpp_files,
             cuda_files=cuda_files,
             extra_cflags=DEFAULT_CFLAGS + extra_cflags,
-            extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list),
-            extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
+            extra_cuda_cflags=cuda_cflags,
+            extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags,
             extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
             build_directory=build_directory,
         )
@@ -272,12 +366,16 @@ def load_jit(
     if prebuilt is not None:
         return prebuilt
 
+    is_rocm = _is_rocm()
     arch_list: List[str] = []
     if cuda_files or cuda_wrappers:
-        from freetoken.kernel._toolchain import check_nvcc_matches_torch
+        if is_rocm:
+            arch_list = _rocm_arch_list()
+        else:
+            from freetoken.kernel._toolchain import check_nvcc_matches_torch
 
-        check_nvcc_matches_torch()
-        arch_list = _cuda_arch_list()
+            check_nvcc_matches_torch()
+            arch_list = _cuda_arch_list()
 
     from tvm_ffi.cpp import load_inline
 
@@ -300,14 +398,23 @@ def load_jit(
     cuda_sources = [f'#include "{path}"' for path in cuda_paths]
     cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers]
 
-    with _pin_tvm_ffi_arch_ctx(arch_list):
+    if is_rocm:
+        cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list)
+        runtime_ldflags = _rocm_link_flags()
+        arch_list_env = ROCM_ARCH_LIST_ENV
+    else:
+        cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list)
+        runtime_ldflags = []
+        arch_list_env = CUDA_ARCH_LIST_ENV
+
+    with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env):
         return load_inline(
             name,
             cpp_sources=cpp_sources,
             cuda_sources=cuda_sources,
             extra_cflags=DEFAULT_CFLAGS + extra_cflags,
-            extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list),
-            extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
+            extra_cuda_cflags=cuda_cflags,
+            extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags,
             extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
             build_directory=build_directory,
         )
diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py
index 2e4ad15f2f..af68bf5816 100644
--- a/python/freetoken/utils/__init__.py
+++ b/python/freetoken/utils/__init__.py
@@ -1,5 +1,7 @@
 from .arch import (
     is_arch_supported,
+    is_rocm,
+    get_rocm_gfx_arch,
     is_sm90_family,
     is_sm90_supported,
     is_sm100_family,
@@ -35,6 +37,8 @@
     "load_toolcall_anchor_id",
     "init_logger",
     "is_arch_supported",
+    "is_rocm",
+    "get_rocm_gfx_arch",
     "is_sm90_family",
     "is_sm90_supported",
     "is_sm100_family",
diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py
index 8c1c6c3d56..304505966c 100644
--- a/python/freetoken/utils/arch.py
+++ b/python/freetoken/utils/arch.py
@@ -1,14 +1,66 @@
 from __future__ import annotations
 
 import functools
+import os
+import re
 from typing import Tuple
 
 
+_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?")
+
+
+def _gfx_arch_from(value: object) -> str | None:
+    match = _GFX_ARCH_RE.search(str(value).lower())
+    return match.group(0) if match else None
+
+
+@functools.cache
+def is_rocm() -> bool:
+    """True when torch is built for ROCm (AMD GPU) instead of CUDA."""
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
+@functools.cache
+def get_rocm_gfx_arch() -> str | None:
+    """Return the current AMD GPU target (for example ``gfx1201``).
+
+    Prefer the runtime device because build variables may contain multiple
+    semicolon-separated targets. Environment variables remain useful for
+    cross-compilation and systems where no GPU is currently visible; in that
+    fallback mode the first target is returned. The result is process-cached
+    for FreeToken's one-process-per-GPU execution model, so callers must select
+    the intended device before the first call.
+    """
+    if not is_rocm():
+        return None
+
+    import torch
+
+    if torch.cuda.is_available():
+        try:
+            props = torch.cuda.get_device_properties(torch.cuda.current_device())
+            for attr in ("gcnArchName", "arch"):
+                arch = _gfx_arch_from(getattr(props, attr, ""))
+                if arch:
+                    return arch
+        except (AttributeError, RuntimeError):
+            pass
+
+    for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"):
+        arch = _gfx_arch_from(os.getenv(env_var, ""))
+        if arch:
+            return arch
+    return None
+
+
 @functools.cache
 def _get_torch_cuda_version() -> Tuple[int, int] | None:
     import torch
     import torch.version
 
+    if is_rocm():
+        return None
     if not torch.cuda.is_available() or not torch.version.cuda:
         return None
     return torch.cuda.get_device_capability()
diff --git a/setup.py b/setup.py
index 7132eea088..997745f1a8 100644
--- a/setup.py
+++ b/setup.py
@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import importlib.util
+import os
 from pathlib import Path
 
 import sys
 
 from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension
+from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension, ROCM_HOME
 
 
 ROOT = Path(__file__).parent
+KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include")
 
 
 def _check_toolchain() -> None:
@@ -20,6 +22,43 @@ def _check_toolchain() -> None:
     module.check_nvcc_matches_torch()
 
 
+def _is_rocm() -> bool:
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
+def _rocm_paths() -> tuple[list[str], list[str], str]:
+    candidates: list[Path] = []
+    if os.getenv("ROCM_HOME"):
+        candidates.append(Path(os.environ["ROCM_HOME"]))
+    if ROCM_HOME:
+        candidates.append(Path(ROCM_HOME))
+
+    # ROCm 7.14 PyTorch images ship the SDK as a Python package instead of
+    # installing it at /opt/rocm.
+    spec = importlib.util.find_spec("_rocm_sdk_core")
+    if spec and spec.submodule_search_locations:
+        candidates.append(Path(next(iter(spec.submodule_search_locations))))
+    candidates.append(Path("/opt/rocm"))
+
+    for rocm_home in dict.fromkeys(candidates):
+        include_dir = rocm_home / "include"
+        library_dir = rocm_home / "lib"
+        if not (include_dir / "hip" / "hip_runtime.h").exists():
+            continue
+        if (library_dir / "libamdhip64.so").exists():
+            return [str(include_dir)], [str(library_dir)], "amdhip64"
+        versioned = sorted(library_dir.glob("libamdhip64.so.*"))
+        if versioned:
+            return [str(include_dir)], [str(library_dir)], f":{versioned[-1].name}"
+
+    searched = ", ".join(str(path) for path in dict.fromkeys(candidates))
+    raise RuntimeError(
+        "A ROCm SDK with HIP headers and libamdhip64 is required to build on ROCm; "
+        f"searched: {searched}. Set ROCM_HOME to override."
+    )
+
+
 def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
     if CUDA_HOME is None:
         raise RuntimeError(
@@ -33,7 +72,21 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
     return [str(cuda_home / "include")], library_dirs
 
 
-cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths()
+IS_ROCM = _is_rocm()
+
+if IS_ROCM:
+    runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths()
+    runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"]
+    # These extensions contain host code only. BuildExtension supplies the ROCm
+    # platform defines to the C++ compiler; offload architecture flags belong on
+    # HIP device sources and would be rejected by the host compiler here.
+    extra_compile = ["-O3", "-std=c++17"]
+else:
+    runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths()
+    runtime_lib = "cudart"
+    runtime_link_args = []
+    extra_compile = ["-O3", "-std=c++17"]
+
 _check_toolchain()
 
 
@@ -44,12 +97,13 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
             sources=[
                 "python/freetoken/kernel/csrc/pinned_tensor.cpp",
             ],
-            include_dirs=cuda_include_dirs,
-            library_dirs=cuda_library_dirs,
-            libraries=["cudart"],
-            extra_compile_args=["-O3", "-std=c++17"],
+            include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs],
+            library_dirs=runtime_library_dirs,
+            libraries=[runtime_lib],
+            extra_compile_args=extra_compile,
+            extra_link_args=runtime_link_args,
         ),
-        # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the
+        # CPU-compute MoE executor for --moe-backend cpu. Links cudart/amdhip64 for the
         # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels
         # use per-function target attributes (avx512bf16/avx512f) + a runtime
         # __builtin_cpu_supports dispatch, so the single binary stays portable
@@ -59,10 +113,11 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
             sources=[
                 "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp",
             ],
-            include_dirs=cuda_include_dirs,
-            library_dirs=cuda_library_dirs,
-            libraries=["cudart"],
-            extra_compile_args=["-O3", "-std=c++17", "-pthread"],
+            include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs],
+            library_dirs=runtime_library_dirs,
+            libraries=[runtime_lib],
+            extra_compile_args=extra_compile + ["-pthread"],
+            extra_link_args=runtime_link_args,
         ),
         # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies
         *([
diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py
index 8bcb8e5810..fa8a1f1d32 100644
--- a/tests/kernels/test_e4m3_compat.py
+++ b/tests/kernels/test_e4m3_compat.py
@@ -41,6 +41,29 @@ def _native_cc() -> bool:
     return torch.cuda.get_device_capability() >= (8, 9)
 
 
+@pytest.mark.skipif(torch.version.hip is None, reason="needs ROCm")
+def test_rocm_host_and_compile_time_native_gates_match():
+    """The host wrapper and Triton constexpr must choose the same buffer ABI."""
+    import triton
+    import triton.language as tl
+
+    from freetoken.kernel.triton.e4m3_compat import e4m3_native, e4m3_native_cx
+
+    @triton.jit
+    def gate_kernel(output):
+        if e4m3_native_cx():
+            value = 1
+        else:
+            value = 0
+        tl.store(output, value)
+
+    output = torch.empty(1, dtype=torch.int32, device="cuda")
+    gate_kernel[(1,)](output)
+
+    assert e4m3_native() is False
+    assert output.item() == int(e4m3_native())
+
+
 # ======================================================================================
 # 1. Primitives vs the native fp8 unit (needs sm_89+ hardware for the reference).
 # ======================================================================================
diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py
index e61108fd53..2fee4f25c8 100644
--- a/tests/kernels/test_pinned_tensor.py
+++ b/tests/kernels/test_pinned_tensor.py
@@ -122,7 +122,12 @@ def test_host_device_ptr_is_identity_under_uva():
         pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead")
     # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer
     # (no registration validation); rejection of pageable memory only exists on
-    # non-identity platforms (Windows/WDDM), where the translation is real.
+    # non-identity CUDA platforms (Windows/WDDM), where the translation is real.
+    # HIP validates registration even though registered/pinned memory uses the
+    # identity address on Linux. Calling it with pageable memory also leaves a
+    # sticky HIP error, so the pinned identity case above is the relevant check.
+    if torch.version.hip is not None:
+        return
     pageable = torch.empty(64, dtype=torch.uint8)
     ext = _load_pinned_extension()
     assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr()
diff --git a/tests/kernels/test_rocm_launch_kwargs.py b/tests/kernels/test_rocm_launch_kwargs.py
new file mode 100644
index 0000000000..52fac54519
--- /dev/null
+++ b/tests/kernels/test_rocm_launch_kwargs.py
@@ -0,0 +1,51 @@
+import pytest
+import torch
+
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is None,
+    reason="a ROCm GPU is required",
+)
+
+
+class _KernelRecorder:
+    def __init__(self):
+        self.kwargs = None
+
+    def __getitem__(self, _grid):
+        def launch(*_args, **kwargs):
+            self.kwargs = kwargs
+
+        return launch
+
+
+def test_rocm_activation_launch_omits_cuda_pdl_attribute(monkeypatch):
+    import freetoken.kernel.triton.activation as activation
+
+    recorder = _KernelRecorder()
+    monkeypatch.setattr(activation, "_act_and_mul_kernel", recorder)
+    monkeypatch.setattr(activation, "_pdl_supported", lambda: True)
+    monkeypatch.setattr(activation, "is_rocm", lambda: True)
+
+    activation.silu_and_mul(torch.ones((2, 128), device="cuda"))
+
+    assert recorder.kwargs is not None
+    assert recorder.kwargs["ENABLE_PDL"] is True
+    assert "launch_pdl" not in recorder.kwargs
+
+
+def test_rocm_norm_launch_omits_cuda_pdl_attribute(monkeypatch):
+    import freetoken.kernel.triton.norm as norm
+
+    recorder = _KernelRecorder()
+    monkeypatch.setattr(norm, "_rmsnorm_kernel", recorder)
+    monkeypatch.setattr(norm, "is_sm90_supported", lambda: True)
+    monkeypatch.setattr(norm, "is_rocm", lambda: True)
+
+    x = torch.ones((2, 64), device="cuda")
+    weight = torch.ones(64, device="cuda")
+    norm.rmsnorm(x, weight)
+
+    assert recorder.kwargs is not None
+    assert recorder.kwargs["ENABLE_PDL"] is True
+    assert "launch_pdl" not in recorder.kwargs
diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py
new file mode 100644
index 0000000000..fcfab75812
--- /dev/null
+++ b/tests/utils/test_rocm_arch.py
@@ -0,0 +1,127 @@
+import importlib
+import os
+import pathlib
+import sys
+import types
+from types import SimpleNamespace
+
+import torch
+
+from freetoken.utils import arch
+
+
+def _clear_arch_caches() -> None:
+    arch.get_rocm_gfx_arch.cache_clear()
+
+
+def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch):
+    monkeypatch.setattr(arch, "is_rocm", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _device: SimpleNamespace(gcnArchName="gfx1201:sramecc-:xnack-"),
+    )
+    monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1200")
+    _clear_arch_caches()
+
+    assert arch.get_rocm_gfx_arch() == "gfx1201"
+
+    _clear_arch_caches()
+
+
+def test_rocm_arch_falls_back_to_cross_compile_env(monkeypatch):
+    monkeypatch.setattr(arch, "is_rocm", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201")
+    _clear_arch_caches()
+
+    assert arch.get_rocm_gfx_arch() == "gfx1200"
+
+    _clear_arch_caches()
+
+
+def test_hip_cflags_target_only_resolved_arch(monkeypatch):
+    from freetoken.kernel import utils
+
+    monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201")
+    monkeypatch.setattr(arch, "get_rocm_gfx_arch", lambda: "gfx1201")
+
+    arch_list = utils._rocm_arch_list()
+    flags = utils._hip_cflags(["-Wno-unused-command-line-argument"], arch_list)
+
+    assert arch_list == ["gfx1201"]
+    assert "--offload-arch=gfx1201" in flags
+    assert "--offload-arch=gfx1200" not in flags
+
+
+def test_rocm_loaders_pin_tvm_ffi_to_resolved_arch(monkeypatch):
+    from freetoken.kernel import utils
+
+    builds = []
+
+    def record_build(*_args, **kwargs):
+        builds.append(
+            (
+                os.environ[utils.ROCM_ARCH_LIST_ENV],
+                kwargs["extra_cuda_cflags"],
+            )
+        )
+        return object()
+
+    tvm_ffi = types.ModuleType("tvm_ffi")
+    tvm_ffi.__path__ = []
+    tvm_ffi_cpp = types.ModuleType("tvm_ffi.cpp")
+    tvm_ffi_cpp.load = record_build
+    tvm_ffi_cpp.load_inline = record_build
+    monkeypatch.setitem(sys.modules, "tvm_ffi", tvm_ffi)
+    monkeypatch.setitem(sys.modules, "tvm_ffi.cpp", tvm_ffi_cpp)
+    monkeypatch.setenv(utils.DISABLE_KERNEL_CACHE_ENV, "1")
+    monkeypatch.setenv(utils.ROCM_ARCH_LIST_ENV, "gfx1100 gfx1200")
+    monkeypatch.setattr(utils, "_is_rocm", lambda: True)
+    monkeypatch.setattr(utils, "_rocm_arch_list", lambda: ["gfx1201"])
+    monkeypatch.setattr(utils, "_rocm_link_flags", lambda: [])
+
+    utils.load_aot("test_rocm_aot_arch", cuda_files=["unused.cu"])
+    utils.load_jit("test_rocm_jit_arch", cuda_files=["unused.cu"])
+
+    assert builds == [
+        ("gfx1201", [*utils.DEFAULT_HIP_CFLAGS, "--offload-arch=gfx1201"]),
+        ("gfx1201", [*utils.DEFAULT_HIP_CFLAGS, "--offload-arch=gfx1201"]),
+    ]
+    assert os.environ[utils.ROCM_ARCH_LIST_ENV] == "gfx1100 gfx1200"
+
+
+def test_rocm_link_flags_support_versioned_modular_sdk(monkeypatch, tmp_path):
+    import torch.utils.cpp_extension as cpp_extension
+
+    from freetoken.kernel import utils
+
+    sdk = tmp_path / "sdk"
+    library_dir = sdk / "lib"
+    library_dir.mkdir(parents=True)
+    versioned_runtime = library_dir / "libamdhip64.so.7"
+    versioned_runtime.write_bytes(b"")
+    real_find_spec = importlib.util.find_spec
+
+    def find_spec(name: str):
+        if name == "_rocm_sdk_core":
+            return SimpleNamespace(submodule_search_locations=[str(sdk)])
+        return real_find_spec(name)
+
+    monkeypatch.delenv("ROCM_HOME", raising=False)
+    monkeypatch.setattr(cpp_extension, "ROCM_HOME", None)
+    monkeypatch.setattr(importlib.util, "find_spec", find_spec)
+    monkeypatch.setattr(pathlib.Path, "home", lambda: tmp_path)
+    utils._rocm_link_flags.cache_clear()
+
+    flags = utils._rocm_link_flags()
+
+    compat_dir = tmp_path / ".cache" / "freetoken" / "rocm-lib"
+    compat_link = compat_dir / "libamdhip64.so"
+    assert f"-L{compat_dir}" in flags
+    assert f"-Wl,-rpath,{library_dir}" in flags
+    assert compat_link.resolve() == versioned_runtime.resolve()
+
+    utils._rocm_link_flags.cache_clear()

From c2091c6fc42617ee6a31e8d636257a7e031a31ff Mon Sep 17 00:00:00 2001
From: Sanjay Sundaram <ssundaram@berkeley.edu>
Date: Mon, 14 Sep 2026 01:55:04 -0700
Subject: [PATCH 054/249] fix(engine): reject fa attention backend below sm_90
 at config time

Declare requires_sm90 on the fa BackendInfo entry so an explicit
--attention-backend fa fails during config validation on Ada and
consumer Blackwell instead of after model load, during CUDA-graph
capture.

Fixes #442
---
 python/freetoken/attention/__init__.py      |  2 ++
 python/freetoken/engine/engine.py           | 11 +++++++
 tests/engine/test_attention_backend_arch.py | 36 +++++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py
index 8a410ffe4a..a52a2e084e 100644
--- a/python/freetoken/attention/__init__.py
+++ b/python/freetoken/attention/__init__.py
@@ -26,6 +26,7 @@ class BackendInfo:
     supported_types: frozenset[AttnType]
     requires_flashinfer: bool = False
     requires_sgl_kernel: bool = False
+    requires_sm90: bool = False
     requires_sm100: bool = False
     # Allowed page sizes (None -> any). Config-time resolution coerces to the last
     # entry when the resolved page_size is not in the list.
@@ -71,6 +72,7 @@ def create_fi_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL}),
         requires_sgl_kernel=True,
+        requires_sm90=True,
     ),
 )
 def create_fa_backend(config: ModelConfig):
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index cc14cf5339..f3fa1eb6e9 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -115,6 +115,10 @@ def _backend_requirements_met(name: str) -> bool:
         return False
     if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available():
         return False
+    # FA3 is Hopper (sm_90a) cubins; FA4 is datacenter Blackwell (sm_100a). Ada / Ampere /
+    # consumer Blackwell have neither, even when sgl_kernel is installed.
+    if any(i.requires_sm90 for i in infos) and not (is_sm90_family() or is_sm100_family()):
+        return False
     if any(i.requires_sm100 for i in infos) and not is_sm100_family():
         return False
     return True
@@ -215,6 +219,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
                 "not installed. Install it with `pip install 'freetoken[sgl]'` (or "
                 "'freetoken[accel]'), or use --attention-backend triton."
             )
+        if info.requires_sm90 and not (is_sm90_family() or is_sm100_family()):
+            raise RuntimeError(
+                f"Attention backend {config.attention_backend!r} requires a compute capability "
+                "9.x GPU (FA3, sm_90 cubins) or 10.x GPU (FA4). Ada (sm_89) and consumer "
+                "Blackwell (sm_12x) are not in the published sgl_kernel flash_ops set. "
+                "Use --attention-backend fi (or triton) instead."
+            )
         if info.requires_sm100 and not is_sm100_family():
             raise RuntimeError(
                 f"Attention backend {config.attention_backend!r} requires a compute capability "
diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py
index 9271bd4254..44f9f87777 100644
--- a/tests/engine/test_attention_backend_arch.py
+++ b/tests/engine/test_attention_backend_arch.py
@@ -110,6 +110,42 @@ def test_explicit_trtllm_allowed_on_sm100_family(monkeypatch):
     assert config.page_size in (16, 32, 64)
 
 
+def test_explicit_fa_rejected_outside_sm90_family(monkeypatch):
+    from freetoken.engine.engine import _adjust_config
+
+    _patch_env(monkeypatch, major=8)
+    config = _engine_config(attention_backend="fa")
+    with pytest.raises(RuntimeError, match="9.x"):
+        _adjust_config(config)
+
+
+def test_explicit_fa_allowed_on_sm90_family(monkeypatch):
+    from freetoken.engine.engine import _adjust_config
+
+    _patch_env(monkeypatch, major=9)
+    config = _engine_config(attention_backend="fa")
+    _adjust_config(config)
+    assert config.attention_backend == "fa"
+
+
+def test_explicit_fa_allowed_on_sm100_family(monkeypatch):
+    from freetoken.engine.engine import _adjust_config
+
+    _patch_env(monkeypatch, major=10)
+    config = _engine_config(attention_backend="fa")
+    _adjust_config(config)
+    assert config.attention_backend == "fa"
+
+
+def test_explicit_fa_rejected_on_consumer_blackwell(monkeypatch):
+    from freetoken.engine.engine import _adjust_config
+
+    _patch_env(monkeypatch, major=12)
+    config = _engine_config(attention_backend="fa")
+    with pytest.raises(RuntimeError, match="9.x"):
+        _adjust_config(config)
+
+
 def test_fa_backend_rejects_cc12(monkeypatch):
     from freetoken.attention.fa import FlashAttentionBackend
 

From f29be24172b2e1109dba3adb3f946a12d391ae99 Mon Sep 17 00:00:00 2001
From: Sanjay Sundaram <ssundaram@berkeley.edu>
Date: Mon, 14 Sep 2026 02:02:51 -0700
Subject: [PATCH 055/249] fix(engine): drop the tutorial comment on the fa
 sm_90 gate

---
 python/freetoken/engine/engine.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index f3fa1eb6e9..32de9bba99 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -115,8 +115,7 @@ def _backend_requirements_met(name: str) -> bool:
         return False
     if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available():
         return False
-    # FA3 is Hopper (sm_90a) cubins; FA4 is datacenter Blackwell (sm_100a). Ada / Ampere /
-    # consumer Blackwell have neither, even when sgl_kernel is installed.
+    # requires_sm90 also admits sm_100 (FA4). Hopper-only would reject B200.
     if any(i.requires_sm90 for i in infos) and not (is_sm90_family() or is_sm100_family()):
         return False
     if any(i.requires_sm100 for i in infos) and not is_sm100_family():

From 6e63b8f0bf736474bf453118fec614fd33d5dd00 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Tue, 29 Sep 2026 03:44:08 +0300
Subject: [PATCH 056/249] fix(server): keep a silent stream alive and notice a
 client that left

---
 python/freetoken/server/api_server.py |  25 +++++-
 tests/server/test_stream_keepalive.py | 115 ++++++++++++++++++++++++++
 2 files changed, 138 insertions(+), 2 deletions(-)
 create mode 100644 tests/server/test_stream_keepalive.py

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 5592a2b0d5..494f659f01 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -57,6 +57,10 @@
 # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor.
 _SHUTTING_DOWN = threading.Event()
 BACKEND_DEATH_EXIT_GRACE_S = 10.0
+# The longest a stream stays silent before it sends an SSE comment. A long prompt prefills for
+# minutes before its first token, and clients time an idle body out (Node's fetch after 300 s).
+STREAM_KEEPALIVE_SECONDS = 15.0
+_STREAM_KEEPALIVE = b": keep-alive\n\n"
 
 
 def get_global_state() -> FrontendManager:
@@ -359,16 +363,33 @@ async def stream_generate(self, uid: int):
         logger.debug("Finished streaming response for user %s", uid)
 
     async def stream_with_cancellation(self, generator, request: Request, uid: int):
+        """Relay ``generator``'s SSE chunks; while it is silent (a long prefill), send a comment every
+        ``STREAM_KEEPALIVE_SECONDS`` and abort the request if the client has gone."""
+        chunks = generator.__aiter__()
+        pending: asyncio.Future | None = None
         try:
-            async for chunk in generator:
-                # detect if the client has disconnected
+            while True:
+                if pending is None:
+                    pending = asyncio.ensure_future(chunks.__anext__())
+                done, _ = await asyncio.wait({pending}, timeout=STREAM_KEEPALIVE_SECONDS)
                 if await request.is_disconnected():
                     logger.info("Client disconnected for user %s", uid)
                     raise asyncio.CancelledError
+                if not done:
+                    yield _STREAM_KEEPALIVE
+                    continue
+                try:
+                    chunk = pending.result()
+                except StopAsyncIteration:
+                    return
+                pending = None
                 yield chunk
         except asyncio.CancelledError:
             asyncio.create_task(self.abort_user(uid))
             raise
+        finally:
+            if pending is not None and not pending.done():
+                pending.cancel()
 
     async def abort_user(self, uid: int):
         await asyncio.sleep(0.1)
diff --git a/tests/server/test_stream_keepalive.py b/tests/server/test_stream_keepalive.py
new file mode 100644
index 0000000000..7925d2a390
--- /dev/null
+++ b/tests/server/test_stream_keepalive.py
@@ -0,0 +1,115 @@
+"""A stream that stays silent -- the prefill of a long prompt -- keeps its client, and notices when
+the client leaves instead of prefilling for nobody."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from freetoken.server import api_server
+from freetoken.server.api_server import FrontendManager
+
+# Converts a stream that never ends into a failure; the streams below finish in milliseconds.
+HANG_GUARD_SECONDS = 10.0
+
+
+class _Request:
+    def __init__(self) -> None:
+        self.gone = False
+
+    async def is_disconnected(self) -> bool:
+        return self.gone
+
+
+def _manager(aborted: list[int]) -> SimpleNamespace:
+    async def abort_user(uid: int) -> None:
+        aborted.append(uid)
+
+    return SimpleNamespace(abort_user=abort_user)
+
+
+def _is_sse_comment(chunk: bytes) -> bool:
+    return chunk.startswith(b":") and chunk.endswith(b"\n\n")
+
+
+def test_a_silent_stream_sends_comments_until_its_next_chunk(monkeypatch):
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01, raising=False)
+    released = asyncio.Event()
+
+    async def chunks():
+        yield b"data: first\n\n"
+        await released.wait()
+        yield b"data: last\n\n"
+
+    async def consume() -> list[bytes]:
+        out = []
+        request = _Request()
+        async for chunk in FrontendManager.stream_with_cancellation(_manager([]), chunks(), request, 7):
+            out.append(chunk)
+            if _is_sse_comment(chunk):
+                released.set()
+        return out
+
+    out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS))
+    assert out[0] == b"data: first\n\n" and out[-1] == b"data: last\n\n"
+    assert len(out) > 2 and all(_is_sse_comment(chunk) for chunk in out[1:-1])
+
+
+def test_a_client_that_leaves_during_the_silence_aborts_its_request(monkeypatch):
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01, raising=False)
+    prefill_done = asyncio.Event()
+
+    async def chunks():
+        yield b"data: first\n\n"
+        await prefill_done.wait()
+        yield b"data: never read\n\n"
+
+    async def scenario() -> tuple[list[bytes], list[int]]:
+        aborted = asyncio.Event()
+        seen: list[int] = []
+
+        async def abort_user(uid: int) -> None:
+            seen.append(uid)
+            aborted.set()
+
+        request = _Request()
+        out: list[bytes] = []
+
+        async def consume() -> None:
+            manager = SimpleNamespace(abort_user=abort_user)
+            async for chunk in FrontendManager.stream_with_cancellation(manager, chunks(), request, 9):
+                out.append(chunk)
+                request.gone = True
+
+        consumer = asyncio.ensure_future(consume())
+        # Only the wrapper noticing the departure can end this wait; the guard below fails it.
+        await aborted.wait()
+        with pytest.raises(asyncio.CancelledError):
+            await consumer
+        return out, seen
+
+    out, seen = asyncio.run(asyncio.wait_for(scenario(), HANG_GUARD_SECONDS))
+    assert out == [b"data: first\n\n"]
+    assert seen == [9]
+    assert not prefill_done.is_set()
+
+
+def test_a_stream_that_never_goes_quiet_carries_no_comments(monkeypatch):
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS, raising=False)
+
+    async def chunks():
+        for index in range(5):
+            yield f"data: {index}\n\n".encode()
+
+    async def consume() -> list[bytes]:
+        return [
+            chunk
+            async for chunk in FrontendManager.stream_with_cancellation(
+                _manager([]), chunks(), _Request(), 3
+            )
+        ]
+
+    out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS))
+    assert out == [f"data: {index}\n\n".encode() for index in range(5)]

From 5031e3d578a3b49572776c3536b3ff4b69769c9b Mon Sep 17 00:00:00 2001
From: YevheniiKotyrlo <44449990+YevheniiKotyrlo@users.noreply.github.com>
Date: Wed, 30 Sep 2026 02:29:17 +0300
Subject: [PATCH 057/249] perf(qsa): reduce block top-k candidates in levels
 past 16 splits (#571)

---
 python/freetoken/kernel/triton/qsa/topk.py | 226 ++++++++++++++++++---
 tests/models/qwen4_exp/test_qsa_kernels.py |  39 +++-
 2 files changed, 225 insertions(+), 40 deletions(-)

diff --git a/python/freetoken/kernel/triton/qsa/topk.py b/python/freetoken/kernel/triton/qsa/topk.py
index ae57566cab..0f37a0f9d9 100644
--- a/python/freetoken/kernel/triton/qsa/topk.py
+++ b/python/freetoken/kernel/triton/qsa/topk.py
@@ -17,10 +17,18 @@
 top-k is a subset of that union -- a global winner beats at most k-1 columns, so it beats all
 but at most k-1 of its own slice -- so the answer stays exact. ``_split_plan`` derives the
 geometry from the buffer width alone, which is fixed when a graph is captured.
+
+A row wide enough to leave more candidates than one program holds resident (262,144 columns,
+a 1M-token window, leaves 32 x k) gets reduction levels between the two phases: each program
+of a level takes ``_MAX_RESIDENT`` consecutive candidates of the level before and writes their
+own top-k, until the survivors fit the merge. The same subset argument holds per level, and
+candidates stay in ascending column order throughout, so ties still go to the lowest column.
 """
 
 from __future__ import annotations
 
+from dataclasses import dataclass
+
 import torch
 import triton
 import triton.language as tl
@@ -125,6 +133,44 @@ def _compact(
     return above, equal
 
 
+@triton.jit
+def _live_candidates(
+    visible,
+    CHUNK: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    TOP_K: tl.constexpr,
+    LEVEL: tl.constexpr,
+    RESIDENT: tl.constexpr,
+):
+    """Candidate slots level ``LEVEL`` filled for a row: its splits past the row's visible tail
+    are one skipped suffix, so every later stage keeps costing what the live part costs."""
+    live = tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K
+    for _ in tl.static_range(LEVEL):
+        live = tl.cdiv(live, RESIDENT) * TOP_K
+    return live
+
+
+@triton.jit
+def _emit_candidates(
+    key,
+    column,
+    k_eff,
+    key_out,
+    col_out,
+    TOP_K: tl.constexpr,
+    BINS: tl.constexpr,
+    RADIX: tl.constexpr,
+    PASSES: tl.constexpr,
+):
+    """Write a resident tile's own top-k as ranked (key, column) candidates; returns how many."""
+    prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES)
+    rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0)
+    write = take & (rank < TOP_K)
+    tl.store(key_out + rank, key.to(tl.int32, bitcast=True), mask=write)
+    tl.store(col_out + rank, column, mask=write)
+    return above + tl.minimum(equal, ties)
+
+
 @triton.jit
 def _qsa_block_topk_kernel(
     logits_ptr,
@@ -225,17 +271,95 @@ def _qsa_topk_split_kernel(
     if k_eff > 0:
         offsets = tl.arange(0, CHUNK)
         key = _load_keys(logits_ptr + row.to(tl.int64) * stride_logits_row + base, offsets, limit)
-        prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES)
-        rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0)
-        write = take & (rank < TOP_K)
-        tl.store(key_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write)
-        tl.store(col_ptr + slot + rank, (base + offsets).to(tl.int32), mask=write)
-        emitted = above + tl.minimum(equal, ties)
+        emitted = _emit_candidates(
+            key, (base + offsets).to(tl.int32), k_eff, key_ptr + slot, col_ptr + slot,
+            TOP_K, BINS, RADIX, PASSES,
+        )
     # A key of 0 is the dead-column sentinel, so the merge needs no separate count per slot.
     pad = tl.arange(0, PAD_K)
     tl.store(key_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K))
 
 
+@triton.jit
+def _reduce_tile(
+    in_key,
+    in_col,
+    out_key,
+    out_col,
+    count,
+    k_eff,
+    TOP_K: tl.constexpr,
+    BLOCK: tl.constexpr,
+    BINS: tl.constexpr,
+    RADIX: tl.constexpr,
+    PASSES: tl.constexpr,
+):
+    """Top-k of ``count`` candidates held as one resident tile; returns how many it wrote."""
+    tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half")
+    offsets = tl.arange(0, BLOCK)
+    inside = offsets < count
+    key = tl.load(in_key + offsets, mask=inside, other=0).to(tl.uint32, bitcast=True)
+    column = tl.load(in_col + offsets, mask=inside, other=-1)
+    return _emit_candidates(key, column, k_eff, out_key, out_col, TOP_K, BINS, RADIX, PASSES)
+
+
+@triton.jit
+def _qsa_topk_reduce_kernel(
+    visible_ptr,
+    in_key_ptr,
+    in_col_ptr,
+    out_key_ptr,
+    out_col_ptr,
+    stride_scratch_row,
+    num_columns,
+    TOP_K: tl.constexpr,
+    PAD_K: tl.constexpr,
+    CHUNK: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    LEVEL: tl.constexpr,
+    RESIDENT: tl.constexpr,
+    BLOCK_SMALL: tl.constexpr,
+    BLOCK_MID: tl.constexpr,
+    BINS: tl.constexpr,
+    RADIX: tl.constexpr,
+    PASSES: tl.constexpr,
+) -> None:
+    """A reduction level: one program per (row, ``RESIDENT`` candidates of level ``LEVEL - 1``),
+    writing their own top-k as this level's candidates."""
+    row = tl.program_id(0)
+    split = tl.program_id(1)
+    visible = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0)
+    k_eff = tl.minimum(visible, TOP_K)
+    live = _live_candidates(visible, CHUNK, N_SPLITS, TOP_K, LEVEL - 1, RESIDENT)
+    base = split * RESIDENT
+    count = tl.minimum(tl.maximum(live - base, 0), RESIDENT)
+    source = row.to(tl.int64) * stride_scratch_row + base
+    slot = row.to(tl.int64) * stride_scratch_row + split * TOP_K
+    in_key, in_col = in_key_ptr + source, in_col_ptr + source
+    out_key, out_col = out_key_ptr + slot, out_col_ptr + slot
+
+    emitted = 0
+    # As in the merge: residency costs the whole tile, so a short range takes a narrower one.
+    if count > 0:
+        if count <= BLOCK_SMALL:
+            emitted = _reduce_tile(
+                in_key, in_col, out_key, out_col, count, k_eff,
+                TOP_K, BLOCK_SMALL, BINS, RADIX, PASSES,
+            )
+        elif count <= BLOCK_MID:
+            emitted = _reduce_tile(
+                in_key, in_col, out_key, out_col, count, k_eff,
+                TOP_K, BLOCK_MID, BINS, RADIX, PASSES,
+            )
+        else:
+            emitted = _reduce_tile(
+                in_key, in_col, out_key, out_col, count, k_eff,
+                TOP_K, RESIDENT, BINS, RADIX, PASSES,
+            )
+    pad = tl.arange(0, PAD_K)
+    tl.store(out_key + pad, 0, mask=(pad >= emitted) & (pad < TOP_K))
+
+
 @triton.jit
 def _merge_tile(
     key_row,
@@ -274,6 +398,8 @@ def _qsa_topk_merge_kernel(
     PAD_K: tl.constexpr,
     CHUNK: tl.constexpr,
     N_SPLITS: tl.constexpr,
+    LEVEL: tl.constexpr,
+    RESIDENT: tl.constexpr,
     BLOCK_SMALL: tl.constexpr,
     BLOCK_MID: tl.constexpr,
     BLOCK_FULL: tl.constexpr,
@@ -281,13 +407,11 @@ def _qsa_topk_merge_kernel(
     RADIX: tl.constexpr,
     PASSES: tl.constexpr,
 ) -> None:
-    """Phase 2: one program per row over the candidates phase 1 left behind."""
+    """Phase 2: one program per row over the candidates the last level ``LEVEL`` left behind."""
     row = tl.program_id(0)
     limit = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0)
     k_eff = tl.minimum(limit, TOP_K)
-    # Candidates sit chunk-major, so the splits past the visible tail are one skipped suffix
-    # and the merge keeps costing what the live part of the row costs.
-    candidates = tl.minimum(tl.cdiv(limit, CHUNK), N_SPLITS) * TOP_K
+    candidates = _live_candidates(limit, CHUNK, N_SPLITS, TOP_K, LEVEL, RESIDENT)
     key_row = key_ptr + row.to(tl.int64) * stride_scratch_row
     col_row = col_ptr + row.to(tl.int64) * stride_scratch_row
     out_row = out_ptr + row.to(tl.int64) * stride_out_row
@@ -325,27 +449,37 @@ def _qsa_topk_merge_kernel(
     tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K))
 
 
-def _split_plan(columns: int, top_k: int) -> tuple[int, int] | None:
-    """``(chunk, n_splits)`` for the split+merge path, or None to keep the one-program path."""
+@dataclass(frozen=True)
+class _SplitPlan:
+    """``chunk`` columns per first-level program, and the program count of every level: the
+    first splits the row, each later one reduces the candidates of the level before it."""
+
+    chunk: int
+    levels: tuple[int, ...]
+
+
+def _split_plan(columns: int, top_k: int) -> _SplitPlan | None:
+    """The split geometry of a ``columns``-wide buffer, or None to keep the one-program path."""
     if top_k <= 0 or columns <= _MIN_CHUNK:
         return None
     max_splits = _MAX_RESIDENT // triton.next_power_of_2(top_k)
     if max_splits < 2:
         return None
-    chunk = max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits)))
-    if chunk > _MAX_RESIDENT:
-        return None
+    chunk = min(_MAX_RESIDENT, max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits))))
     n_splits = -(-columns // chunk)
     # Merging n_splits*top_k candidates has to be cheaper than scanning the row once.
     if n_splits < 2 or n_splits * top_k >= columns:
         return None
-    return chunk, n_splits
+    levels = [n_splits]
+    while levels[-1] * top_k > _MAX_RESIDENT:
+        levels.append(-(-levels[-1] * top_k // _MAX_RESIDENT))
+    return _SplitPlan(chunk, tuple(levels))
 
 
 def qsa_block_topk_scratch_width(columns: int, top_k: int) -> int:
     """int32 columns of scratch ``qsa_block_topk`` wants per row; 0 when it needs none."""
     plan = _split_plan(columns, top_k)
-    return 0 if plan is None else 2 * plan[1] * top_k
+    return 0 if plan is None else 2 * top_k * sum(plan.levels)
 
 
 def qsa_block_topk(
@@ -396,24 +530,31 @@ def qsa_block_topk(
         )
         return out
 
-    chunk, n_splits = plan
-    half = n_splits * top_k
+    width = 2 * top_k * sum(plan.levels)
     if scratch is None:
-        scratch = torch.empty((rows, 2 * half), dtype=torch.int32, device=logits.device)
+        scratch = torch.empty((rows, width), dtype=torch.int32, device=logits.device)
     elif (
         scratch.ndim != 2
         or scratch.shape[0] < rows
-        or scratch.shape[1] < 2 * half
+        or scratch.shape[1] < width
         or scratch.dtype != torch.int32
         or scratch.stride(1) != 1
     ):
         raise ValueError(
             f"QSA block top-k needs a row-contiguous int32 scratch of at least "
-            f"[{rows}, {2 * half}], got {tuple(scratch.shape)} {scratch.dtype}"
+            f"[{rows}, {width}], got {tuple(scratch.shape)} {scratch.dtype}"
         )
-    keys = scratch[:rows, :half]
-    cols = scratch[:rows, half : 2 * half]
-    _qsa_topk_split_kernel[(rows, n_splits)](
+    # Each level's candidates: its keys, then their columns, level after level along the row.
+    candidates: list[tuple[torch.Tensor, torch.Tensor]] = []
+    offset = 0
+    for programs in plan.levels:
+        half = programs * top_k
+        candidates.append(
+            (scratch[:rows, offset : offset + half], scratch[:rows, offset + half : offset + 2 * half])
+        )
+        offset += 2 * half
+    keys, cols = candidates[0]
+    _qsa_topk_split_kernel[(rows, plan.levels[0])](
         logits,
         visible,
         keys,
@@ -423,14 +564,39 @@ def qsa_block_topk(
         columns,
         TOP_K=top_k,
         PAD_K=pad_k,
-        CHUNK=chunk,
+        CHUNK=plan.chunk,
         BINS=_BINS,
         RADIX=_RADIX,
         PASSES=_PASSES,
         num_warps=8,
         num_stages=1,
     )
-    merge_block = triton.next_power_of_2(half)
+    for level in range(1, len(plan.levels)):
+        (in_keys, in_cols), (out_keys, out_cols) = candidates[level - 1], candidates[level]
+        _qsa_topk_reduce_kernel[(rows, plan.levels[level])](
+            visible,
+            in_keys,
+            in_cols,
+            out_keys,
+            out_cols,
+            keys.stride(0),
+            columns,
+            TOP_K=top_k,
+            PAD_K=pad_k,
+            CHUNK=plan.chunk,
+            N_SPLITS=plan.levels[0],
+            LEVEL=level,
+            RESIDENT=_MAX_RESIDENT,
+            BLOCK_SMALL=min(1024, _MAX_RESIDENT),
+            BLOCK_MID=min(4096, _MAX_RESIDENT),
+            BINS=_BINS,
+            RADIX=_RADIX,
+            PASSES=_PASSES,
+            num_warps=8,
+            num_stages=1,
+        )
+    keys, cols = candidates[-1]
+    merge_block = triton.next_power_of_2(keys.shape[1])
     _qsa_topk_merge_kernel[(rows,)](
         visible,
         keys,
@@ -441,8 +607,10 @@ def qsa_block_topk(
         columns,
         TOP_K=top_k,
         PAD_K=pad_k,
-        CHUNK=chunk,
-        N_SPLITS=n_splits,
+        CHUNK=plan.chunk,
+        N_SPLITS=plan.levels[0],
+        LEVEL=len(plan.levels) - 1,
+        RESIDENT=_MAX_RESIDENT,
         BLOCK_SMALL=min(1024, merge_block),
         BLOCK_MID=min(4096, merge_block),
         BLOCK_FULL=merge_block,
diff --git a/tests/models/qwen4_exp/test_qsa_kernels.py b/tests/models/qwen4_exp/test_qsa_kernels.py
index 4b678b3bdc..6715519f05 100644
--- a/tests/models/qwen4_exp/test_qsa_kernels.py
+++ b/tests/models/qwen4_exp/test_qsa_kernels.py
@@ -331,7 +331,17 @@ def test_torch_topk_env_picks_the_fallback(monkeypatch):
 # Block top-k, split + merge path (wide buffers)
 # --------------------------------------------------------------------------------------
 
-SPLIT_CHUNK = 4096  # _split_plan's chunk for every buffer these tests use
+# (columns, _split_plan's chunk). A 1M-token window scores 262,144 compressed rows, and its 32
+# splits leave more candidates than one program holds, so a reduction level runs before the merge.
+SPLIT_SHAPES = [(65536, 4096), (262144, 8192)]
+
+
+def _split_chunk(columns: int, width: int) -> int:
+    from freetoken.kernel.triton.qsa.topk import _split_plan
+
+    plan = _split_plan(columns, width)
+    assert plan is not None, f"{columns} columns must take the split path"
+    return plan.chunk
 
 
 def _policy_topk_blocks(logits, visible, width):
@@ -348,29 +358,34 @@ def _policy_topk_blocks(logits, visible, width):
     return out
 
 
-def _split_topk_case(n_blocks: int, width: int, bs: int, mode: str, seed: int):
+def _split_topk_case(n_blocks: int, chunk: int, width: int, bs: int, mode: str, seed: int):
     if mode != "boundary":
         return _topk_case(n_blocks, bs, mode, seed)
     # width - 212 columns beat the tie, so the 212 remaining winners start 100 columns below
-    # a chunk boundary and run past it into a chunk that is all tie.
+    # a chunk boundary and run past it into a chunk that is all tie. The last row cuts at the
+    # middle split boundary, which is where a reduction level divides the candidates.
     logits = torch.zeros(bs, n_blocks, device="cuda")
+    splits = n_blocks // chunk
     for row in range(bs):
-        cut = SPLIT_CHUNK * (1 + row % (n_blocks // SPLIT_CHUNK - 1))
+        cut = chunk * (splits // 2 if row == bs - 1 else 1 + row % (splits - 1))
         logits[row, : width - 212] = 2.0
         logits[row, cut - 100 :] = 1.0
     return logits, torch.full((bs,), n_blocks, dtype=torch.int32, device="cuda")
 
 
 @requires_cuda
-@pytest.mark.parametrize("n_blocks", [65536])
+@pytest.mark.parametrize("n_blocks,chunk", SPLIT_SHAPES)
 @pytest.mark.parametrize("bs", [4])
 @pytest.mark.parametrize("mode", ["random", "boundary", "ragged", "dead"])
-def test_block_topk_split_path_matches_torch_topk(n_blocks: int, bs: int, mode: str):
+def test_block_topk_split_path_matches_torch_topk(n_blocks: int, chunk: int, bs: int, mode: str):
     from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width
 
     width = 512
     assert qsa_block_topk_scratch_width(n_blocks, width) > 0, "case must take the split path"
-    logits, visible = _split_topk_case(n_blocks, width, bs, mode, seed=n_blocks + bs + len(mode))
+    assert _split_chunk(n_blocks, width) == chunk
+    logits, visible = _split_topk_case(
+        n_blocks, chunk, width, bs, mode, seed=n_blocks + bs + len(mode)
+    )
     blocks = torch.empty(bs, width, dtype=torch.int32, device=logits.device)
     qsa_block_topk(logits, visible, blocks)
 
@@ -431,8 +446,10 @@ def test_block_topk_split_path_cost_tracks_live_blocks():
     """A wide buffer with a short row must not pay for the splits past its visible tail."""
     from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width
 
-    # wide enough that the split work dwarfs the 20-launch floor even at boosted clocks
-    rows, columns, width = 1, 262144, 512
+    # enough rows that the live splits fill the card, so their work outweighs the launch floor
+    rows, columns, width = 8, 262144, 512
+    assert qsa_block_topk_scratch_width(columns, width) > 0, "case must take the split path"
+    chunk = _split_chunk(columns, width)
     device = torch.device("cuda")
     logits = torch.randn(rows, columns, device=device)
     visible = torch.full((rows,), columns, dtype=torch.int32, device=device)
@@ -463,8 +480,8 @@ def replay_us(live: int) -> float:
             best = min(best, start.elapsed_time(stop) * 1000.0 / 20)
         return best
 
-    full, short = replay_us(columns), replay_us(SPLIT_CHUNK)
-    assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {SPLIT_CHUNK} live {short:.1f}us"
+    full, short = replay_us(columns), replay_us(chunk)
+    assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {chunk} live {short:.1f}us"
 
 
 @requires_cuda

From eba7120eabde70c7414e0e14c551e1caa3b0f174 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 29 Sep 2026 19:32:36 -0700
Subject: [PATCH 058/249] feat(install): record the install root as
 FREETOKEN_HOME in environment.d (#576)

---
 install.sh | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/install.sh b/install.sh
index 2a9a3dec30..f623ffece7 100755
--- a/install.sh
+++ b/install.sh
@@ -20,7 +20,8 @@
 #                                 Defaults to $DEFAULT_KERNEL_CACHE_WHEEL_URL; if
 #                                 unset and FREETOKEN_WHEEL is local, the script
 #                                 auto-detects a sibling freetoken_kernel_cache-*.whl.
-#   FREETOKEN_HOME                install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv
+#   FREETOKEN_HOME                install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv;
+#                                 a successful install writes it to environment.d
 #   FREETOKEN_PY_VERSION          python for the venv (default: 3.12 — must match the wheel tag)
 #   FREETOKEN_BIN_DIR             where to symlink `ft` (default: ~/.local/bin)
 #   FREETOKEN_ENV_DIR             environment.d dir (default: ~/.config/environment.d)
@@ -248,8 +249,8 @@ ln -sf "$FT_BIN" "$BIN_DIR/ft"
 say "symlinked $BIN_DIR/ft -> $FT_BIN"
 
 mkdir -p "$ENV_DIR"
-printf 'FREETOKEN_FT_BIN=%s\n' "$FT_BIN" > "$ENV_DIR/50-freetoken.conf"
-say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_FT_BIN) — GUI picks it up after next login"
+printf 'FREETOKEN_HOME=%s\n' "$(cd "$FT_HOME" && pwd)" >"$ENV_DIR/50-freetoken.conf"
+say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_HOME) - FreeToken Desktop reads the file, other programs see the variable after next login"
 
 # --- 5. Self-check ---------------------------------------------------------
 if "$FT_BIN" --help >/dev/null 2>&1; then
@@ -264,10 +265,9 @@ ${C_GREEN}FreeToken engine installed.${C_RESET}
 
   ft binary        $FT_BIN
   on PATH as       $BIN_DIR/ft   (ensure $BIN_DIR is on PATH)
-  Desktop env      FREETOKEN_FT_BIN via environment.d (re-login to apply)
+  Desktop env      FREETOKEN_HOME in $ENV_DIR/50-freetoken.conf
 
-Run in this shell without re-login:
-  export FREETOKEN_FT_BIN="$FT_BIN"
+Run:
   ft serve --model <path> --port 1919
 
 EOF

From 09ab44ea0274cdd4ec0ba4d40af1cb99d369ce8c Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Thu, 1 Oct 2026 06:24:45 +0300
Subject: [PATCH 059/249] test(server): patch the keep-alive interval strictly

STREAM_KEEPALIVE_SECONDS exists, so the tests set it without raising=False: a
renamed constant now fails them instead of leaving them on the 15 s default. The
manager helper drops the abort list no test read.
---
 tests/server/test_stream_keepalive.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/server/test_stream_keepalive.py b/tests/server/test_stream_keepalive.py
index 7925d2a390..e86d4365fc 100644
--- a/tests/server/test_stream_keepalive.py
+++ b/tests/server/test_stream_keepalive.py
@@ -23,9 +23,9 @@ async def is_disconnected(self) -> bool:
         return self.gone
 
 
-def _manager(aborted: list[int]) -> SimpleNamespace:
+def _manager() -> SimpleNamespace:
     async def abort_user(uid: int) -> None:
-        aborted.append(uid)
+        pass
 
     return SimpleNamespace(abort_user=abort_user)
 
@@ -35,7 +35,7 @@ def _is_sse_comment(chunk: bytes) -> bool:
 
 
 def test_a_silent_stream_sends_comments_until_its_next_chunk(monkeypatch):
-    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01, raising=False)
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01)
     released = asyncio.Event()
 
     async def chunks():
@@ -46,7 +46,7 @@ async def chunks():
     async def consume() -> list[bytes]:
         out = []
         request = _Request()
-        async for chunk in FrontendManager.stream_with_cancellation(_manager([]), chunks(), request, 7):
+        async for chunk in FrontendManager.stream_with_cancellation(_manager(), chunks(), request, 7):
             out.append(chunk)
             if _is_sse_comment(chunk):
                 released.set()
@@ -58,7 +58,7 @@ async def consume() -> list[bytes]:
 
 
 def test_a_client_that_leaves_during_the_silence_aborts_its_request(monkeypatch):
-    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01, raising=False)
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01)
     prefill_done = asyncio.Event()
 
     async def chunks():
@@ -97,7 +97,7 @@ async def consume() -> None:
 
 
 def test_a_stream_that_never_goes_quiet_carries_no_comments(monkeypatch):
-    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS, raising=False)
+    monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS)
 
     async def chunks():
         for index in range(5):
@@ -107,7 +107,7 @@ async def consume() -> list[bytes]:
         return [
             chunk
             async for chunk in FrontendManager.stream_with_cancellation(
-                _manager([]), chunks(), _Request(), 3
+                _manager(), chunks(), _Request(), 3
             )
         ]
 

From 235e2016ddb64c8eb1f0d877237062f087fb2ab2 Mon Sep 17 00:00:00 2001
From: YevheniiKotyrlo <44449990+YevheniiKotyrlo@users.noreply.github.com>
Date: Thu, 1 Oct 2026 10:40:27 +0300
Subject: [PATCH 060/249] feat(rotary): size the rope table for YaRN and add
 --hf-overrides (#573)

* feat(rotary): size the rope table for YaRN and add --hf-overrides

* fix(rotary): guard mrope's table width and copy --hf-overrides values

- MRotaryEmbedding asserts its cos/sin table is rotary_dim wide. A proportional
  rope's table spans the head, and the mrope kernel indexes the rotary_dim // 2
  section table with the table's width.
- _with_overrides deep-copies every value it stores: transformers fills rope
  defaults into the rope_parameters it is handed, which was the caller's dict.
- The rope-table tests the engine test covers fold into it, the two shared-table
  tests are one, and the --hf-overrides test moves to tests/utils/test_hf.py,
  where it loads the override against the same values written into config.json,
  for a model type transformers knows and one it does not.
- docs/cli.md lists --hf-overrides.

* fix(server): give --hf-overrides a fresh dict per parse

argparse hands its default object to every parse that omits the flag, so
default={} put one shared dict in every ServerArgs. The flag defaults to None and
becomes a new dict where the arguments are assembled, as --mm-processor-kwargs
does.

* docs(rotary): an indexer's heads differ in width from the attention's, not wider
---
 docs/cli.md                                   |   1 +
 python/freetoken/attention/qsa_sparse.py      |  17 +-
 python/freetoken/engine/config.py             |   8 +-
 python/freetoken/engine/engine.py             |   9 +-
 python/freetoken/layers/rotary.py             | 193 +++++++++---------
 python/freetoken/models/config.py             |   7 +
 python/freetoken/server/args.py               |  12 ++
 python/freetoken/utils/hf.py                  |  31 ++-
 tests/engine/test_attention_backend_matrix.py |   4 +-
 tests/engine/test_cache_budget.py             |  21 +-
 tests/kernels/test_mrope.py                   |  13 ++
 tests/kernels/test_rotary.py                  |  61 +++++-
 tests/models/qwen4_exp/test_config.py         |  51 +++++
 tests/models/test_qwen3_vl.py                 |   2 +-
 tests/utils/test_hf.py                        |  34 +++
 15 files changed, 328 insertions(+), 136 deletions(-)
 create mode 100644 tests/utils/test_hf.py

diff --git a/docs/cli.md b/docs/cli.md
index 0c33943987..fb5e613f7b 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -33,6 +33,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--model-path`, `--model` | required | Local dir, HF repo id, or an FTW dir (auto-detected) |
 | `--served-model-name` | basename of `--model` | Model id reported by `/v1/models` |
+| `--hf-overrides` | — | JSON object applied to the checkpoint's config as vLLM's `--hf-overrides`: a nested config section updates key by key, any other value is replaced whole. A YaRN `rope_parameters` override serves `original_max_position_embeddings * factor` positions |
 
 ### Server & runtime
 
diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index b07c952323..8ce63c4242 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -167,23 +167,22 @@ def _cmp_pages(self, slot: int) -> torch.Tensor:
 
     def _index_rope_cache(self) -> torch.Tensor:
         """cos/sin table of the indexer rope: same rotary_dim and frequencies as the main
-        attention, ``head_size`` 128 instead of 256, so it is a separate get_rope instance.
+        attention with ``head_size`` 128 instead of 256, so the attention's own table.
 
         The table itself (not RotaryEmbedding.forward) because the indexer's norm+rope is one
         fused kernel and the compressed keys rope at their group's position, not the query's."""
         if self._index_cos_sin is None:
-            from freetoken.layers.rotary import get_rope
+            from freetoken.layers.rotary import rope_cos_sin_table
 
             rotary = self.rotary_config
             with torch.device(self.device):
-                rope = get_rope(
-                    head_dim=self.index_head_dim,
-                    rotary_dim=rotary.rotary_dim,
-                    max_position=rotary.max_position,
-                    base=rotary.base,
-                    rope_scaling=tuple(rotary.scaling.items()) if rotary.scaling else None,
+                self._index_cos_sin = rope_cos_sin_table(
+                    self.index_head_dim,
+                    rotary.rotary_dim,
+                    rotary.max_position,
+                    rotary.base,
+                    tuple(rotary.scaling.items()) if rotary.scaling else None,
                 )
-            self._index_cos_sin = rope._cos_sin_cache.to(self.device)
         return self._index_cos_sin
 
     def _index_rope_rows(self, positions: torch.Tensor) -> torch.Tensor:
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 44c5a38e73..3e38cb2165 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -3,7 +3,7 @@
 import copy
 from dataclasses import dataclass, field, replace
 from functools import cached_property
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, Any, List, Mapping
 
 import torch
 from freetoken.distributed import DistributedInfo
@@ -23,6 +23,8 @@ class EngineConfig:
     model_path: str
     tp_info: DistributedInfo
     dtype: torch.dtype
+    # --hf-overrides: applied to the checkpoint config the model is built from (cached_load_hf_config)
+    hf_overrides: Mapping[str, Any] = field(default_factory=dict)
     max_running_req: int = 4
     attention_backend: str = "auto"
     moe_strategy: str = "auto"
@@ -103,7 +105,7 @@ def __post_init__(self):
 
     @cached_property
     def hf_config(self):
-        return cached_load_hf_config(self.model_path)
+        return cached_load_hf_config(self.model_path, self.hf_overrides)
 
     @cached_property
     def model_spec(self) -> ModelSpec:
@@ -142,7 +144,7 @@ def model_config(self) -> ModelConfig:
     def max_seq_len(self) -> int:
         if self.max_seq_len_override is not None:
             return self.max_seq_len_override
-        return self.model_config.rotary_config.max_position
+        return self.model_config.rotary_config.table_positions
 
     @property
     def max_forward_len(self) -> int:
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index cc14cf5339..e855bcd32f 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1770,18 +1770,17 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             )
         override("num_page_override", config.num_token_override // config.page_size)
 
-    # The rope cos/sin table is baked to rotary_config.max_position, and neither rope kernel
+    # The rope cos/sin table covers rotary_config.table_positions, and neither rope kernel
     # bounds-checks the position it gathers with -- a longer ceiling reads past the table.
     # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config).
     rotary = getattr(model_config, "rotary_config", None)
     seq_override = getattr(config, "max_seq_len_override", None)
     if seq_override is not None and rotary is not None and not is_dsv4:
-        if seq_override > rotary.max_position:
+        if seq_override > rotary.table_positions:
             raise ValueError(
                 f"--max-seq-len-override {seq_override} exceeds the model's "
-                f"rope table ({rotary.max_position} positions). Serving past it would read "
-                "out of bounds; extend the checkpoint's rope_scaling / "
-                "max_position_embeddings in config.json instead."
+                f"rope table ({rotary.table_positions} positions). Serving past it would read "
+                "out of bounds; extend the rope with --hf-overrides (YaRN rope_parameters) instead."
             )
 
     # The startup ServerArgs dump is the *requested* config, printed in the frontend process
diff --git a/python/freetoken/layers/rotary.py b/python/freetoken/layers/rotary.py
index 0d04645d99..d0178f3bc1 100644
--- a/python/freetoken/layers/rotary.py
+++ b/python/freetoken/layers/rotary.py
@@ -2,7 +2,7 @@
 
 import functools
 import math
-from typing import Any, Callable, Dict, Tuple
+from typing import Any, Mapping, Tuple
 
 import torch
 
@@ -14,50 +14,21 @@ def __init__(
         self,
         head_size: int,
         rotary_dim: int,
-        max_position_embeddings: int,
-        base: float,
-        post_process: None | Callable[[torch.Tensor], torch.Tensor] = None,
-        proportional: bool = False,
-        attention_factor: float = 1.0,
+        cos_sin_cache: torch.Tensor,
         is_neox: bool = True,
     ) -> None:
         super().__init__()
+        assert 0 < rotary_dim <= head_size
+        assert rotary_dim % 2 == 0
+        assert head_size in [64, 128, 256, 512]
         self.head_size = head_size
         self.rotary_dim = rotary_dim
         # NeoX (half-rotation, HF default) vs GPT-J interleaved (adjacent pairs,
         # ``rope_interleave`` models: GLM MLA lineage). Both underlying kernels
         # accept the flag; the cos/sin cache layout is identical.
         self.is_neox = is_neox
-        if proportional:
-            assert 0 < rotary_dim <= head_size
-            assert rotary_dim % 2 == 0
-            inv_freq = 1.0 / (
-                base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size)
-            )
-            if rotary_dim < head_size:
-                inv_freq[rotary_dim // 2 :] = 0.0
-        else:
-            # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size):
-            # rope is applied to the first ``rotary_dim`` dims of each head, the rest pass
-            # through. Frequencies are spaced over ``rotary_dim`` (matches HF default
-            # partial rope, e.g. Qwen3.5 partial_rotary_factor, MiniMax-M2's
-            # ``apply_rotary_pos_emb``). Full rope is rotary_dim == head_size and is
-            # unaffected. ``head_size`` is passed to flashinfer separately so it rotates
-            # only the first ``rotary_dim`` dims.
-            assert 0 < rotary_dim <= head_size
-            assert rotary_dim % 2 == 0
-            inv_freq = 1.0 / (
-                base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)
-            )
-        if post_process is not None:
-            inv_freq = post_process(inv_freq)
-        t = torch.arange(max_position_embeddings, dtype=torch.float)
-        freqs = torch.einsum("i,j -> ij", t, inv_freq)
-        cos = freqs.cos() * attention_factor
-        sin = freqs.sin() * attention_factor
-        # buffer, so don't load/save
-        self._cos_sin_cache = torch.cat((cos, sin), dim=-1)
-        assert self.head_size in [64, 128, 256, 512]
+        # shared with every rope of the same frequencies (rope_cos_sin_table), so never written
+        self._cos_sin_cache = cos_sin_cache
 
         from freetoken.kernel.backend import is_flashinfer_installed
 
@@ -160,6 +131,7 @@ def __init__(
     ) -> None:
         super().__init__(*args, **kwargs)
         assert self.is_neox, "mrope is defined on the NeoX half-rotation layout"
+        assert self._cos_sin_cache.shape[1] == self.rotary_dim, "mrope reads a rotary_dim-wide table, not a proportional rope's"
         half = self.rotary_dim // 2
         assert sum(mrope_section) == half, (mrope_section, half)
         self._section_table = build_section_table(tuple(mrope_section), layout)
@@ -193,56 +165,50 @@ def forward(
         return query, key
 
 
-def _get_rope(
-    head_dim: int,
-    rotary_dim: int,
-    max_position: int,
-    base: float,
-    rope_scaling: Dict[str, Any] | None = None,
-    is_neox: bool = True,
-) -> RotaryEmbedding:
-    if rope_scaling is None:
-        return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox)
-    # need to test some cases:
-    match rope_scaling["rope_type"]:
+def rope_table_positions(max_position: int, rope_scaling: Mapping[str, Any] | None) -> int:
+    """Rows the cos/sin table needs: YaRN serves ``original_max_position_embeddings * factor``,
+    above ``max_position_embeddings`` when only the rope parameters are overridden."""
+    if rope_scaling is None or rope_scaling.get("rope_type") != "yarn":
+        return max_position
+    extended = int(rope_scaling["original_max_position_embeddings"] * rope_scaling["factor"])
+    return max(max_position, extended)
+
+
+def _rope_frequencies(
+    head_size: int, rotary_dim: int, base: float, rope_scaling: Mapping[str, Any] | None
+) -> Tuple[torch.Tensor, float]:
+    """``inv_freq`` over the table's half width, and the factor cos and sin are both scaled by."""
+    rope_type = "default" if rope_scaling is None else rope_scaling["rope_type"]
+    if rope_type == "proportional":
+        # Frequencies spaced over the whole head; the pairs past rotary_dim do not rotate.
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size))
+        inv_freq[rotary_dim // 2 :] = 0.0
+        return inv_freq, 1.0
+    # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size): rope is applied
+    # to the first ``rotary_dim`` dims of each head, the rest pass through, and frequencies
+    # are spaced over ``rotary_dim`` (HF default partial rope, e.g. Qwen3.5
+    # partial_rotary_factor, MiniMax-M2's ``apply_rotary_pos_emb``).
+    inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
+    match rope_type:
         case "default":
-            return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox)
-
-        case "proportional":
-            return RotaryEmbedding(
-                head_dim,
-                rotary_dim,
-                max_position,
-                base,
-                proportional=True,
-                is_neox=is_neox,
-            )
+            return inv_freq, 1.0
 
         case "llama3":
             scaling_factor: float = rope_scaling["factor"]
             low_freq_factor: float = rope_scaling["low_freq_factor"]
             high_freq_factor: float = rope_scaling["high_freq_factor"]
             original_max_position: int = rope_scaling["original_max_position_embeddings"]
-
-            def post_process(inv_freq: torch.Tensor) -> torch.Tensor:
-                # no smooth if low_freq_factor == high_freq_factor
-                wave_len = 2 * math.pi / inv_freq
-                if low_freq_factor == high_freq_factor:
-                    return torch.where(
-                        wave_len < original_max_position / high_freq_factor,
-                        inv_freq,
-                        inv_freq / scaling_factor,
-                    )
-
-                delta = high_freq_factor - low_freq_factor
-                smooth = (original_max_position / wave_len - low_freq_factor) / delta
-                smooth = torch.clamp(smooth, 0, 1)
-                factor = (1 - smooth) / scaling_factor + smooth
-                return factor * inv_freq
-
-            return RotaryEmbedding(
-                head_dim, rotary_dim, max_position, base, post_process, is_neox=is_neox
-            )
+            wave_len = 2 * math.pi / inv_freq
+            # no smooth if low_freq_factor == high_freq_factor
+            if low_freq_factor == high_freq_factor:
+                return torch.where(
+                    wave_len < original_max_position / high_freq_factor,
+                    inv_freq,
+                    inv_freq / scaling_factor,
+                ), 1.0
+            delta = high_freq_factor - low_freq_factor
+            smooth = torch.clamp((original_max_position / wave_len - low_freq_factor) / delta, 0, 1)
+            return ((1 - smooth) / scaling_factor + smooth) * inv_freq, 1.0
 
         case "yarn":
             factor: float = rope_scaling["factor"]
@@ -288,27 +254,54 @@ def _find_correction_dim(num_rotations: float) -> float:
             high = min(high, rotary_dim - 1)
             if low == high:  # HF nudges instead of flooring the gap at 1 ("truncate": false)
                 high += 0.001
-
-            def post_process(inv_freq: torch.Tensor) -> torch.Tensor:
-                ramp = torch.clamp(
-                    (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low),
-                    0, 1,
-                )
-                return (inv_freq / factor) * ramp + inv_freq * (1 - ramp)
-
-            return RotaryEmbedding(
-                head_dim,
-                rotary_dim,
-                max_position,
-                base,
-                post_process,
-                attention_factor=float(attention_factor),
-                is_neox=is_neox,
+            ramp = torch.clamp(
+                (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low), 0, 1
             )
+            return (inv_freq / factor) * ramp + inv_freq * (1 - ramp), float(attention_factor)
 
     raise ValueError(f"Unsupported {rope_scaling = }")
 
 
+@functools.cache
+def _cos_sin_table(
+    head_size: int,
+    rotary_dim: int,
+    positions: int,
+    base: float,
+    rope_scaling: Tuple[Tuple[str, Any], ...] | None,
+    device: torch.device,
+) -> torch.Tensor:
+    inv_freq, attention_factor = _rope_frequencies(
+        head_size, rotary_dim, base, dict(rope_scaling) if rope_scaling is not None else None
+    )
+    freqs = torch.arange(positions, dtype=torch.float)[:, None] * inv_freq[None, :]
+    return torch.cat((freqs.cos() * attention_factor, freqs.sin() * attention_factor), dim=-1)
+
+
+def rope_cos_sin_table(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: float,
+    rope_scaling: Tuple[Tuple[str, Any], ...] | None = None,
+) -> torch.Tensor:
+    """The ``[positions, width]`` cos|sin table, one per distinct rope: only proportional rope
+    depends on the head size, so an indexer whose heads differ in width reads the attention's table."""
+    rope_map = dict(rope_scaling) if rope_scaling is not None else None
+    proportional = rope_map is not None and rope_map["rope_type"] == "proportional"
+    device = torch.get_default_device()
+    if device.type == "cuda" and device.index is None:
+        device = torch.device("cuda", torch.cuda.current_device())
+    return _cos_sin_table(
+        head_size if proportional else rotary_dim,
+        rotary_dim,
+        rope_table_positions(max_position, rope_map),
+        base,
+        rope_scaling,
+        device,
+    )
+
+
 _ROPE_DEVICE: torch.device | None = None
 
 
@@ -328,17 +321,15 @@ def get_rope(
     mrope_section: Tuple[int, ...] | None = None,
     mrope_layout: str = "interleaved",
 ) -> RotaryEmbedding:
-    rope_map = dict(rope_scaling) if rope_scaling is not None else None
-
     def build() -> RotaryEmbedding:
+        table = rope_cos_sin_table(head_dim, rotary_dim, max_position, base, rope_scaling)
         if mrope_section is not None:
-            assert rope_map is None or rope_map.get("rope_type", "default") == "default"
             return MRotaryEmbedding(
-                head_dim, rotary_dim, max_position, base,
+                head_dim, rotary_dim, table,
                 is_neox=is_neox, mrope_section=tuple(mrope_section),
                 layout=mrope_layout,
             )
-        return _get_rope(head_dim, rotary_dim, max_position, base, rope_map, is_neox)
+        return RotaryEmbedding(head_dim, rotary_dim, table, is_neox=is_neox)
 
     t = torch.tensor([])
     if t.device == torch.device("meta"):
@@ -358,5 +349,7 @@ def build() -> RotaryEmbedding:
     "RotaryEmbedding",
     "build_section_table",
     "get_rope",
+    "rope_cos_sin_table",
+    "rope_table_positions",
     "set_rope_device",
 ]
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index d7b504a769..c5943dce1b 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -92,6 +92,13 @@ class RotaryConfig:
     mrope_section: list | None = None
     mrope_layout: str = "contiguous"  # see freetoken.layers.rotary.build_section_table
 
+    @property
+    def table_positions(self) -> int:
+        """Positions the cos/sin table covers: the longest sequence this rope can serve."""
+        from freetoken.layers.rotary import rope_table_positions
+
+        return rope_table_positions(self.max_position, self.scaling)
+
 
 def mrope_layout_from_rope_params(rope_params: Any) -> str:
     """rope_parameters flags -> layout name: mrope_interleaved_glm, else mrope_interleaved, else contiguous."""
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..23a62f8419 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -272,6 +272,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.",
     )
 
+    parser.add_argument(
+        "--hf-overrides",
+        type=_json_object,
+        default=None,
+        metavar="JSON",
+        help="JSON object applied to the checkpoint config the model is built from, as vLLM's "
+        "--hf-overrides: a nested config section is updated key by key, any other value is "
+        "replaced whole. A YaRN rope_parameters override extends the servable context to "
+        "original_max_position_embeddings * factor.",
+    )
+
     parser.add_argument(
         "--tensor-parallel-size",
         "--tp-size",
@@ -892,6 +903,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         image_max_tokens=image_max_tokens,
         processor_kwargs=kwargs.pop("mm_processor_kwargs") or {},
     )
+    kwargs["hf_overrides"] = kwargs["hf_overrides"] or {}
     result = ServerArgs(**kwargs)
     logger.info(f"Parsed arguments:\n{result}")
     return result, run_shell
diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py
index b93d3b1203..489585e8c2 100644
--- a/python/freetoken/utils/hf.py
+++ b/python/freetoken/utils/hf.py
@@ -1,9 +1,8 @@
+import copy
 import functools
 import json
 import os
-from typing import Any
-
-from typing import FrozenSet
+from typing import Any, FrozenSet, Mapping
 
 from huggingface_hub import hf_hub_download, snapshot_download
 from huggingface_hub.utils import EntryNotFoundError
@@ -235,20 +234,40 @@ def _load_hf_config(model_path: str) -> Any:
     return config
 
 
-def cached_load_hf_config(model_path: str) -> PretrainedConfig:
+def _with_overrides(config: Any, data: dict, overrides: Mapping[str, Any]) -> dict:
+    merged = dict(data)
+    for key, value in overrides.items():
+        section = getattr(config, key, None)
+        if isinstance(value, Mapping) and isinstance(section, (PretrainedConfig, RawConfigShim)):
+            merged[key] = _with_overrides(section, merged.get(key) or {}, value)
+        else:
+            merged[key] = copy.deepcopy(value)
+    return merged
+
+
+def cached_load_hf_config(
+    model_path: str, overrides: Mapping[str, Any] | None = None
+) -> PretrainedConfig:
+    """A fresh copy of the checkpoint's config with ``overrides`` applied as vLLM's --hf-overrides
+    does: nested config sections update key by key, every other value is replaced whole."""
     # A .gguf file (or an FTW dir converted from one) carries its own metadata (no HF
     # config.json); return a shim the model registry dispatches on instead of a
     # PretrainedConfig.
     from freetoken.models.gguf.reader import gguf_config_source
 
     if (gguf_src := gguf_config_source(model_path)) is not None:
+        if overrides:
+            raise ValueError("--hf-overrides applies to a HF config.json; a GGUF carries its own metadata")
         from freetoken.models.gguf.config import build_gguf_shim
 
         return build_gguf_shim(gguf_src)
     config = _load_hf_config(model_path)
+    data = config.to_dict()
+    if overrides:
+        data = _with_overrides(config, data, overrides)
     if isinstance(config, RawConfigShim):
-        return RawConfigShim(config.to_dict())
-    return type(config)(**config.to_dict())
+        return RawConfigShim(data)
+    return type(config)(**data)
 
 
 def _weight_allow_patterns(repo_id: str) -> list[str]:
diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py
index ef7dc7ff85..ad07a9dd09 100644
--- a/tests/engine/test_attention_backend_matrix.py
+++ b/tests/engine/test_attention_backend_matrix.py
@@ -40,6 +40,8 @@ def _spec(name, attn_type, *, mla=False, sliding_window=None, index_head_dim=0,
 
 
 def _model_config(kind):
+    from freetoken.models.config import RotaryConfig
+
     mc = SimpleNamespace(
         model_type=kind,
         single_stream_only=False,
@@ -48,7 +50,7 @@ def _model_config(kind):
         has_swa_attention=False,
         has_linear_attention=False,
         num_layers=4,
-        rotary_config=SimpleNamespace(max_position=1024),
+        rotary_config=RotaryConfig(head_dim=128, rotary_dim=128, max_position=1024, base=1e4, scaling=None),
     )
     if kind == "full":
         specs = (_spec("full", AttnType.FULL),)
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index a4e47aa02f..073112bb5a 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -430,13 +430,17 @@ def test_page_table_width_covers_whole_trailing_pages():
             assert w > last_col and w % 32 == 0
 
 
-def _generic_rotary_cfg(max_position, override):
+def _generic_rotary_cfg(max_position, override, scaling=None):
     from types import SimpleNamespace
 
+    from freetoken.models.config import RotaryConfig
+
     model_config = SimpleNamespace(
         single_stream_only=False, is_moe=False, expert_quant="none",
         has_swa_attention=False, has_linear_attention=False,
-        rotary_config=SimpleNamespace(max_position=max_position),
+        rotary_config=RotaryConfig(
+            head_dim=128, rotary_dim=128, max_position=max_position, base=1e4, scaling=scaling
+        ),
     )
 
     class Cfg:
@@ -476,6 +480,19 @@ def test_adjust_config_allows_override_at_rope_table_boundary():
     _adjust_config(_generic_rotary_cfg(max_position=1024, override=1024))  # must not raise
 
 
+YARN_4X = {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 1024}
+
+
+def test_adjust_config_allows_override_to_the_yarn_extended_table():
+    # A YaRN override that leaves max_position_embeddings at the trained length (Qwen3.8's 1M
+    # recipe) serves original * factor positions; the gate must follow the table, not the config.
+    from freetoken.engine.engine import _adjust_config
+
+    _adjust_config(_generic_rotary_cfg(max_position=1024, override=4096, scaling=YARN_4X))
+    with pytest.raises(ValueError, match=r"rope table \(4096 positions\)"):
+        _adjust_config(_generic_rotary_cfg(max_position=1024, override=4097, scaling=YARN_4X))
+
+
 def test_adjust_config_rope_gate_exempts_dsv4():
     # DSV4 sizes its own rope table from the resolved max_seq_len (_adjust_dsv4_config),
     # so the generic gate must not fire even when the override dwarfs max_position.
diff --git a/tests/kernels/test_mrope.py b/tests/kernels/test_mrope.py
index f8402f8754..08dff5cb57 100644
--- a/tests/kernels/test_mrope.py
+++ b/tests/kernels/test_mrope.py
@@ -32,6 +32,19 @@ def test_section_tables():
         build_section_table(SECTION, "diagonal")
 
 
+def test_refuses_a_table_wider_than_the_rotary_dims():
+    # A proportional rope's table spans the whole head; mrope takes its width from the table
+    # and indexes the rotary_dim // 2 section table with it, reading past the end.
+    from freetoken.layers.rotary import get_rope
+
+    get_rope.cache_clear()
+    with pytest.raises(AssertionError, match="rotary_dim"):
+        get_rope(
+            head_dim=HEAD, rotary_dim=ROT, max_position=8, base=1e7,
+            rope_scaling=(("rope_type", "proportional"),), mrope_section=SECTION,
+        )
+
+
 def _make(mrope: bool, layout: str = "interleaved"):
     from freetoken.layers.rotary import get_rope
 
diff --git a/tests/kernels/test_rotary.py b/tests/kernels/test_rotary.py
index f3e8654fc8..82ee01537d 100644
--- a/tests/kernels/test_rotary.py
+++ b/tests/kernels/test_rotary.py
@@ -101,7 +101,7 @@ def test_yarn_correction_range_matches_hf_where_the_clamp_and_gap_bind():
     longest-wavelength dims; flooring the gap at 1 flattens the ramp for a sub-1 gap."""
     from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 
-    from freetoken.layers.rotary import get_rope
+    from freetoken.layers.rotary import _rope_frequencies
 
     class _Shim:  # duck-typed config for HF's real _compute_yarn_parameters
         def __init__(self, head_dim, max_position, rope_parameters):
@@ -122,16 +122,18 @@ def standardize_rope_params(self):
         (128, 1e6, {"factor": 16.0, "beta_fast": 1.15, "beta_slow": 1.0,
                     "original_max_position_embeddings": 32768, "truncate": False}),
     ]
+    def table(inv_freq: torch.Tensor, attention_factor: float) -> torch.Tensor:
+        freqs = torch.outer(torch.arange(max_position, dtype=torch.float), inv_freq.float())
+        return torch.cat((freqs.cos() * attention_factor, freqs.sin() * attention_factor), dim=-1)
+
+    # The frequencies, not get_rope's table: YaRN sizes that to original * factor (up to 2M
+    # rows here), and the correction range is a property of the frequencies alone.
     for rotary_dim, base, scaling in cases:
-        get_rope.cache_clear()
-        rope = get_rope(head_dim=rotary_dim, rotary_dim=rotary_dim, max_position=max_position,
-                        base=base, rope_scaling=(("rope_type", "yarn"), *scaling.items()))
+        ours = table(*_rope_frequencies(rotary_dim, rotary_dim, base, {"rope_type": "yarn", **scaling}))
         params = {"rope_type": "yarn", "rope_theta": base, "partial_rotary_factor": 1.0, **scaling}
-        inv_freq, attn_factor = ROPE_INIT_FUNCTIONS["yarn"](
-            _Shim(rotary_dim, max_position, params), device=torch.device("cpu"))
-        freqs = torch.outer(torch.arange(max_position, dtype=torch.float), inv_freq.float())
-        expected = torch.cat((freqs.cos() * attn_factor, freqs.sin() * attn_factor), dim=-1)
-        torch.testing.assert_close(rope._cos_sin_cache, expected, rtol=0, atol=1e-6)
+        expected = table(*ROPE_INIT_FUNCTIONS["yarn"](
+            _Shim(rotary_dim, max_position, params), device=torch.device("cpu")))
+        torch.testing.assert_close(ours, expected, rtol=0, atol=1e-6)
 
 
 def test_yarn_mscale_all_dim_zero_falls_back_like_hf():
@@ -160,3 +162,44 @@ def test_yarn_mscale_all_dim_zero_falls_back_like_hf():
     expected = 0.1 * math.log(factor) + 1.0  # 1.277, not 1.196 (= 1 + 0.1*0.707*ln 16)
     cos = rope._cos_sin_cache[:, :32]
     torch.testing.assert_close(cos[0], torch.full_like(cos[0], expected))
+
+
+def _yarn(factor: float, original: int) -> tuple:
+    return (("rope_type", "yarn"), ("factor", factor), ("original_max_position_embeddings", original))
+
+
+def test_yarn_table_covers_the_extended_context():
+    # A checkpoint extended by overriding only its rope parameters (Qwen3.8's 1M recipe) keeps
+    # max_position_embeddings at the trained length; YaRN serves original * factor positions.
+    from freetoken.layers.rotary import get_rope
+
+    get_rope.cache_clear()
+    rope = get_rope(head_dim=64, rotary_dim=64, max_position=64, base=10000.0, rope_scaling=_yarn(4.0, 64))
+    assert rope._cos_sin_cache.shape == (256, 64)
+
+    longer = get_rope(head_dim=64, rotary_dim=64, max_position=512, base=10000.0, rope_scaling=_yarn(4.0, 64))
+    assert longer._cos_sin_cache.shape == (512, 64)
+    torch.testing.assert_close(longer._cos_sin_cache[:256], rope._cos_sin_cache, rtol=0, atol=0)
+
+
+def test_one_table_serves_every_rope_of_the_same_frequencies():
+    # The QSA indexer ropes 128-wide heads with the attention's 64 rotary dims, and mrope reads
+    # the same rows per axis: one table, not a second copy of a 1M-row fp32 tensor.
+    from freetoken.layers.rotary import MRotaryEmbedding, get_rope
+
+    get_rope.cache_clear()
+    attention = get_rope(head_dim=256, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64))
+    indexer = get_rope(head_dim=128, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64))
+    mrope = get_rope(
+        head_dim=256, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64),
+        mrope_section=(11, 11, 10), mrope_layout="interleaved",
+    )
+    assert indexer.head_size == 128 and isinstance(mrope, MRotaryEmbedding)
+    assert attention._cos_sin_cache.shape == (256, 64)
+    table = attention._cos_sin_cache.data_ptr()
+    assert indexer._cos_sin_cache.data_ptr() == table and mrope._cos_sin_cache.data_ptr() == table
+
+    proportional = (("rope_type", "proportional"),)
+    wide = get_rope(head_dim=128, rotary_dim=64, max_position=4, base=1e4, rope_scaling=proportional)
+    narrow = get_rope(head_dim=64, rotary_dim=64, max_position=4, base=1e4, rope_scaling=proportional)
+    assert wide._cos_sin_cache.shape == (4, 128) and narrow._cos_sin_cache.shape == (4, 64)
diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py
index be532eeca9..f850927b99 100644
--- a/tests/models/qwen4_exp/test_config.py
+++ b/tests/models/qwen4_exp/test_config.py
@@ -215,3 +215,54 @@ def test_released_checkpoints_keep_the_dense_projections_bf16(quantization_confi
     quant = _quant(_hf_config(quantization_config), tmp_path)
     for prefix in DENSE_PREFIXES + BF16_PREFIXES:
         assert quant.scheme_for(prefix) is None, prefix
+
+
+# Qwen3.8-Flash-Next model card, "Processing Ultra-Long Texts": the 1M-token override.
+MODEL_CARD_YARN = {
+    "mrope_interleaved": True,
+    "mrope_section": [11, 11, 10],
+    "rope_type": "yarn",
+    "rope_theta": 10000000,
+    "partial_rotary_factor": 0.25,
+    "factor": 4.0,
+    "original_max_position_embeddings": 262144,
+}
+
+
+def _write_checkpoint_config(path, **fields) -> None:
+    import json
+
+    text = {k: v for k, v in vars(_text_config()).items()}
+    config = {
+        "model_type": "qwen4_exp",
+        "architectures": ["Qwen4ExpForConditionalGeneration"],
+        "image_token_id": 248056,
+        "text_config": text,
+        "quantization_config": RADIXARK_NVFP4,
+        **fields,
+    }
+    (path / "config.json").write_text(json.dumps(config))
+
+
+def test_the_engine_serves_the_model_card_1m_recipe(tmp_path):
+    import torch
+
+    from freetoken.distributed import DistributedInfo
+    from freetoken.engine.config import EngineConfig
+    from freetoken.mm.config import ENCODER_KINDS, MultimodalConfig
+
+    _write_checkpoint_config(tmp_path)
+    config = EngineConfig(
+        model_path=str(tmp_path),
+        tp_info=DistributedInfo(rank=0, size=1),
+        dtype=torch.bfloat16,
+        hf_overrides={"text_config": {"rope_parameters": MODEL_CARD_YARN}},
+        mm=MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS)),
+    )
+    rotary = config.model_config.rotary_config
+    assert rotary.scaling["rope_type"] == "yarn" and rotary.scaling["factor"] == 4.0 and rotary.mrope_section is None
+    assert rotary.max_position == 262144 and rotary.rotary_dim == 64 and rotary.base == 10000000
+    assert config.max_seq_len == 4 * 262144
+    plain = EngineConfig(model_path=str(tmp_path), tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16,
+                         mm=MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS)))
+    assert plain.model_config.rotary_config.scaling is None and plain.max_seq_len == 262144
diff --git a/tests/models/test_qwen3_vl.py b/tests/models/test_qwen3_vl.py
index c45a9f150e..6c773d3deb 100644
--- a/tests/models/test_qwen3_vl.py
+++ b/tests/models/test_qwen3_vl.py
@@ -39,7 +39,7 @@ def _engine_config(monkeypatch, hf, processor, **overrides):
     import freetoken.engine.config as engine_config
     import freetoken.mm.processor as mm_processor
 
-    monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf)
+    monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, overrides: hf)
     monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None)
     monkeypatch.setattr(mm_processor, "get_mm_processor", lambda path, mm=None: processor)
     return EngineConfig(
diff --git a/tests/utils/test_hf.py b/tests/utils/test_hf.py
new file mode 100644
index 0000000000..28630b1431
--- /dev/null
+++ b/tests/utils/test_hf.py
@@ -0,0 +1,34 @@
+"""cached_load_hf_config with --hf-overrides, against the same config written into config.json."""
+
+import copy
+import json
+
+import pytest
+
+from freetoken.utils.hf import cached_load_hf_config
+
+
+def _checkpoint(path, model_type: str, dtype: str, rope_parameters: dict) -> str:
+    path.mkdir()
+    text_config = {"hidden_size": 64, "max_position_embeddings": 4096, "rope_parameters": rope_parameters}
+    (path / "config.json").write_text(json.dumps({"model_type": model_type, "dtype": dtype, "text_config": text_config}))
+    return str(path)
+
+
+@pytest.mark.parametrize("model_type", ["qwen3_5", "a_model_type_transformers_does_not_know"])
+def test_hf_overrides_load_as_if_the_checkpoint_config_said_so(tmp_path, model_type):
+    # vLLM's merge: a nested config section updates key by key, any other value is replaced whole
+    yarn = {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 4096}
+    checkpoint = _checkpoint(tmp_path / "checkpoint", model_type, "bfloat16", {"rope_type": "default", "rope_theta": 1e6})
+    edited = _checkpoint(tmp_path / "edited", model_type, "float16", yarn)
+    overrides = {"dtype": "float16", "text_config": {"rope_parameters": yarn}}
+    requested = copy.deepcopy(overrides)
+
+    hf = cached_load_hf_config(checkpoint, overrides)
+    expected = cached_load_hf_config(edited)
+    assert hf.dtype == expected.dtype
+    assert hf.text_config.to_dict() == expected.text_config.to_dict()
+    # transformers fills rope defaults into the rope_parameters it is handed
+    assert overrides == requested
+    # a fresh copy each time: the override never reaches the cached checkpoint config
+    assert cached_load_hf_config(checkpoint).text_config.rope_parameters["rope_type"] == "default"

From d3512b43affe981465e03ee28cbd88f49c39b9aa Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 1 Oct 2026 16:59:45 -0700
Subject: [PATCH 061/249] fix(ple): fence disk staging against in-flight
 lookups (#588)

Co-authored-by: andy-yang-1 <andy_yang@berkeley.edu>
---
 python/freetoken/models/qwen4_exp/ple_disk.py | 31 ++++++++++++++++---
 tests/models/qwen4_exp/test_ple_disk.py       | 18 +++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index ec7da80d56..69ee9fa93a 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -150,9 +150,18 @@ def __init__(
         self._graph_dev = torch.empty(
             max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device
         )
+        # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the
+        # dispatch context after the step's fill: a record inside the capture would not be observable)
+        self._graph_consumed = torch.cuda.Event()
         eager_bytes = max_extend_tokens * self._token_bytes
-        self._eager_pinned = alloc_pinned_tensor(eager_bytes, dtype=torch.uint8)
-        self._eager_pinned.zero_()  # the warmup prefill stages nothing and reads whatever sits here
+        # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the
+        # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy
+        # that last read it has run (``_eager_read[i]``, recorded right after that copy)
+        self._eager_pinned = [alloc_pinned_tensor(eager_bytes, dtype=torch.uint8) for _ in range(2)]
+        for buf in self._eager_pinned:
+            buf.zero_()  # the warmup prefill stages nothing and reads whatever sits here
+        self._eager_read = [torch.cuda.Event(), torch.cuda.Event()]
+        self._eager_slot = 0
         self._eager_dev = torch.empty(eager_bytes, dtype=torch.uint8, device=self._device)
         # probe picks flag-sync (graph WAITs at the consume, host fills then signals) or launch-gating
         from freetoken.kernel.row_store import probe_wait_sync
@@ -170,7 +179,13 @@ def __init__(
 
     def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
         """Stage per-request token runs (two context ids, then the new tokens) in batch order."""
-        pinned = self._graph_pinned if graph else self._eager_pinned
+        if graph:
+            self._graph_consumed.synchronize()
+            pinned = self._graph_pinned
+        else:
+            self._eager_slot ^= 1
+            self._eager_read[self._eager_slot].synchronize()
+            pinned = self._eager_pinned[self._eager_slot]
         offset = 0
         for run in runs:
             self._store.stage(run.data_ptr(), run.numel() - 2, pinned.data_ptr() + offset * self._token_bytes)
@@ -236,21 +251,27 @@ def forward_host_ctx(self, batch: Batch, use_graph: bool):
         # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run
         if deferred is not None:
             deferred()
+        if use_graph:
+            # only after this step's fill: the next fill waits on this launch, not the fill on it
+            self._graph_consumed.record(torch.cuda.current_stream(self._device))
 
     # ---------------- device side (PLETableBackend protocol) ----------------
 
     def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
         rows = row_ids.shape[0]
+        stream = torch.cuda.current_stream(self._device)
         capturing = torch.cuda.is_current_stream_capturing()
         if capturing and self._wait_sync:
             from freetoken.kernel.row_store import wait_reset
 
-            wait_reset(torch.cuda.current_stream(self._device), self._flag)
+            wait_reset(stream, self._flag)
         pinned, dev = (
-            (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned, self._eager_dev)
+            (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev)
         )
         nbytes = rows * self._token_bytes
         dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True)
+        if not capturing:
+            self._eager_read[self._eager_slot].record(stream)
         values = dev[:nbytes].view(torch.float8_e4m3fn).to(self.dtype)
         if self.scale != 1.0:
             values = values * self.scale
diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py
index 1112210b84..ba9cd98c63 100644
--- a/tests/models/qwen4_exp/test_ple_disk.py
+++ b/tests/models/qwen4_exp/test_ple_disk.py
@@ -287,6 +287,24 @@ def test_disk_table_matches_oracle(tmp_path):
     assert _bitwise_equal(disk.lookup(ids), oracle.lookup(ids)), "hook prefill"
 
 
+@requires_cuda
+def test_eager_fill_does_not_overwrite_an_in_flight_lookup(tmp_path):
+    """The overlap scheduler fills batch k+1 while batch k's lookup copy is still queued behind the
+    layers before it (a sleep kernel here): batch k must read its own rows."""
+    disk, oracle, args = _make_table(tmp_path)
+    emb = _embedding()
+    steps = [(3, 4, 7), (4, 7, 9), (7, 9, 13), (9, 13, 5)]
+    ids = [emb.row_ids(_meta([[token]], [[older, newer]], decode=True)).cuda() for older, newer, token in steps]
+    got = []
+    for (older, newer, token), step_ids in zip(steps, ids):
+        disk.fill([torch.tensor([older, newer, token])], graph=False)
+        torch.cuda._sleep(50_000_000)
+        got.append(disk.lookup(step_ids))
+    torch.cuda.synchronize()
+    for i, (out, step_ids) in enumerate(zip(got, ids)):
+        assert _bitwise_equal(out, oracle.lookup(step_ids)), f"step {i} read another step's rows"
+
+
 @requires_cuda
 def test_graph_sync_protocol(tmp_path, monkeypatch):
     disk, oracle, args = _make_table(tmp_path)

From 1d5ce5459b6632adc60feb50b07cd4bbf73af1c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= <Avicennasis@gmail.com>
Date: Sat, 3 Oct 2026 09:45:33 -0400
Subject: [PATCH 062/249] fix(checkpoint): refuse converting into a directory
 that already holds an FTW

---
 python/freetoken/checkpoint/convert.py   | 20 ++++++++
 tests/checkpoint/test_convert_refusal.py | 60 ++++++++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 tests/checkpoint/test_convert_refusal.py

diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py
index a875d77f34..2e6ebbad93 100644
--- a/python/freetoken/checkpoint/convert.py
+++ b/python/freetoken/checkpoint/convert.py
@@ -103,6 +103,25 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]:
     return copied
 
 
+def _refuse_occupied_output(out_dir: str) -> None:
+    """Refuse an output dir that already holds FTW output: the index of a finished conversion,
+    or shards left by an interrupted one. There is no resume; the writer would open shard
+    00000 with "wb" and a stale index could then describe the wrong bytes."""
+    from .ftw import INDEX_NAME, is_ftw_checkpoint
+
+    if is_ftw_checkpoint(out_dir):
+        found = INDEX_NAME
+    else:
+        shards = glob.glob(os.path.join(glob.escape(out_dir), "freetoken-*.ftw"))
+        if not shards:
+            return
+        found = os.path.basename(min(shards))
+    raise SystemExit(
+        f"{out_dir} already holds an FTW checkpoint ({found}); "
+        f"convert into an empty directory or remove the old one first"
+    )
+
+
 class _ConvertSink:
     """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each
     completed layer's banks as their own FTW entries immediately (name
@@ -184,6 +203,7 @@ def convert_checkpoint(
 
     if is_ftw_checkpoint(model_path):
         raise SystemExit(f"{model_path} is already an FTW checkpoint")
+    _refuse_occupied_output(out_dir)
     tp = try_get_tp_info()
     if tp is None:
         set_tp_info(rank=0, size=1)
diff --git a/tests/checkpoint/test_convert_refusal.py b/tests/checkpoint/test_convert_refusal.py
new file mode 100644
index 0000000000..6c09f79ad8
--- /dev/null
+++ b/tests/checkpoint/test_convert_refusal.py
@@ -0,0 +1,60 @@
+"""convert_checkpoint refuses an output directory that already holds FTW output, before it writes anything.
+
+No GPU: the refusal sits ahead of the TP/CUDA setup, so a CPU box reaches it with an empty source dir."""
+
+import os
+
+import pytest
+import torch
+
+from freetoken.checkpoint.convert import _refuse_occupied_output, convert_checkpoint
+from freetoken.checkpoint.ftw import INDEX_NAME, FTWWriter
+
+
+def _listing(d):
+    return sorted((name, os.stat(os.path.join(d, name)).st_size, os.stat(os.path.join(d, name)).st_mtime_ns)
+                  for name in os.listdir(d))
+
+
+def _refused(src, out):
+    before = _listing(out)
+    with pytest.raises(SystemExit, match="already holds an FTW checkpoint") as ei:
+        convert_checkpoint(str(src), str(out))
+    assert _listing(out) == before  # nothing written, nothing touched
+    return str(ei.value)
+
+
+def test_finished_ftw_in_the_output_dir_is_refused_by_name(tmp_path):
+    out = tmp_path / "out"
+    writer = FTWWriter(str(out))
+    writer.add_tensor("model.a.weight", torch.zeros(4, dtype=torch.bfloat16))
+    writer.finalize({})
+    src = tmp_path / "src"
+    src.mkdir()
+    msg = _refused(src, out)
+    assert str(out) in msg and INDEX_NAME in msg
+
+
+def test_shards_from_an_interrupted_conversion_are_refused_by_name(tmp_path):
+    out = tmp_path / "out"
+    out.mkdir()
+    (out / "freetoken-00001.ftw").write_bytes(b"\0" * 64)
+    (out / "freetoken-00000.ftw").write_bytes(b"\0" * 64)
+    src = tmp_path / "src"
+    src.mkdir()
+    msg = _refused(src, out)
+    assert str(out) in msg and "freetoken-00000.ftw" in msg and INDEX_NAME not in msg
+
+
+def test_empty_absent_or_unrelated_output_dirs_pass(tmp_path):
+    empty = tmp_path / "empty"
+    empty.mkdir()
+    assert _refuse_occupied_output(str(empty)) is None
+    absent = tmp_path / "absent"
+    assert _refuse_occupied_output(str(absent)) is None
+    assert not absent.exists()  # the check creates nothing
+    other = tmp_path / "other"
+    other.mkdir()
+    (other / "config.json").write_text("{}")  # copied metadata alone is not FTW output
+    (other / "notes.ftw.bak").write_bytes(b"")
+    assert _refuse_occupied_output(str(other)) is None

From 189c8ad0c3350e30252209d40780274976cb7527 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 4 Oct 2026 13:51:55 -0700
Subject: [PATCH 063/249] feat(moe): serve NVFP4 and MXFP4 experts fused
 through the expert banks (#601)

* feat(moe): serve NVFP4 and MXFP4 experts fused through the expert banks

Resident experts skip the state dict and load through the offload expert-bank loader (HF pieces or FTW banks) into per-layer GPU banks, so fused serves bf16, fp8_block, NVFP4 and MXFP4 alike. Unified-memory GPUs auto-select fused for these formats and ignore offload-only flags with a warning.

BREAKING CHANGE: ft checkpoint always writes MoE experts as banks; an FTW converted with --moe-backend fused or triton (dense experts) no longer loads and must be reconverted.

* docs(moe): name only GB10 / DGX Spark for unified memory
---
 docs/cli.md                                   |   9 +-
 docs/models.md                                |  17 ++-
 python/freetoken/checkpoint/__main__.py       |   2 +-
 python/freetoken/checkpoint/convert.py        |   4 +-
 python/freetoken/checkpoint/ftw.py            | 140 ++++++++++++++----
 python/freetoken/engine/engine.py             |  59 ++++++--
 python/freetoken/layers/__init__.py           |   3 +-
 python/freetoken/layers/moe.py                |  17 ++-
 .../freetoken/layers/quantization/moe/base.py |  15 +-
 .../layers/quantization/moe/fp8_block.py      |  16 +-
 .../layers/quantization/moe/mxfp4.py          |  36 +----
 .../layers/quantization/moe/nvfp4.py          |  14 +-
 .../layers/quantization/moe/unquantized.py    |  13 +-
 python/freetoken/models/qwen3_5_moe/weight.py |  39 +----
 python/freetoken/moe/expert_banks.py          |  71 +++++++--
 python/freetoken/moe/offload_cache.py         |  15 +-
 python/freetoken/server/args.py               |   4 +-
 tests/engine/test_cache_budget.py             |   7 +-
 tests/models/qwen4_exp/test_skeleton.py       |   9 +-
 tests/models/test_quant_config.py             |  20 +--
 20 files changed, 298 insertions(+), 212 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index fb5e613f7b..deeda18419 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -80,7 +80,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 
 | Flag | Default | Meaning |
 |---|---|---|
-| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile. `--moe-backend` is the deprecated old spelling |
+| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling |
 | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name |
 | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=<marlin\|b12x\|triton>` (`flashinfer` means b12x); cannot be combined with `--quant-backend` |
 | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) |
@@ -179,9 +179,10 @@ ft checkpoint --model <hf_dir> --out <ftw_dir> [--dtype bfloat16] [--moe-backend
 ```
 
 Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained
-fast-load format; point `ft serve --model` at the output dir. `--moe-backend
-offload` (default) packs experts into offload banks; `--moe-backend triton`
-keeps them dense for resident serving. See the FTW caveats in
+fast-load format; point `ft serve --model` at the output dir. MoE experts are
+always packed into expert banks, which serve the fused and the offload
+strategies alike (GGUF and DeepSeek-V4 experts serve offload only);
+`--moe-backend` is kept for compatibility. See the FTW caveats in
 [models.md](models.md#notes); FTW files from older builds can be repaired with
 [scripts/ftw_hotfix.py](ftw-hotfix.md) instead of reconverting.
 
diff --git a/docs/models.md b/docs/models.md
index 548a363c67..efd5268b37 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -38,7 +38,8 @@ These families accept image input by default; pass `--text-model-only` to skip t
 
 `ft serve --moe-strategy {auto,fused,offload,cpu,hybrid}` (`--moe-backend` is the deprecated old spelling):
 
-- **fused** — experts resident on GPU (needs the VRAM); never auto-selected.
+- **fused** — experts resident on GPU (needs the VRAM). Serves every expert
+  format except GGUF and DeepSeek-V4's.
 - **offload** — experts live in host RAM, an LRU cache of expert slots on GPU;
   misses stream over PCIe.
 - **cpu** — misses are computed on the CPU instead of fetched.
@@ -46,12 +47,24 @@ These families accept image input by default; pass `--text-model-only` to skip t
   CPU, overlapped. Run `ft bench bw` once per machine to calibrate the split.
 - **auto** — dense models always resolve to `fused`; MoE models resolve to
   `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile
-  recommends it.
+  recommends it. Unified-memory GPUs differ, see below.
+
+### Unified-memory GPUs (GB10 / DGX Spark)
+
+The GPU and the CPU share one DRAM, so offload only copies experts between two
+names for the same memory. On these GPUs `auto` resolves every MoE model to
+`fused`: the default is the resident path only. GGUF and DeepSeek-V4 experts
+have no resident path and stay on `offload`. Offload-only flags
+(`--moe-cpu-layers`, `--moe-cache-size` / `-rate` / `-auto`,
+`--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a
+warning; pass `--moe-strategy offload` to use them.
 
 ## Notes
 
 - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into
   FreeToken's fast-load format, and `ft serve --model` auto-detects the result.
+- An FTW converted by an older build with `--moe-backend fused` or `triton` keeps
+  its experts as dense weights and no longer loads; reconvert it with `ft checkpoint`.
 - FTW files converted by builds before the quantization refactor may fail to load;
   see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool.
 - An FTW converted before its family served images holds no vision encoder: `ft serve`
diff --git a/python/freetoken/checkpoint/__main__.py b/python/freetoken/checkpoint/__main__.py
index 2b7e35f1d1..0d7c2d7ed7 100644
--- a/python/freetoken/checkpoint/__main__.py
+++ b/python/freetoken/checkpoint/__main__.py
@@ -37,7 +37,7 @@ def main(argv: list[str] | None = None, prog: str = "freetoken.checkpoint") -> i
     p.add_argument("--out", required=True, help="output FTW checkpoint dir")
     p.add_argument("--dtype", choices=sorted(_DTYPES), default="bfloat16")
     p.add_argument("--moe-backend", default="offload",
-                   help="offload (experts -> banks) or e.g. triton (experts stay dense)")
+                   help="kept for compatibility: MoE experts always go to banks, which every strategy reads")
     p.add_argument("--quant-backend", type=_parse_quant_backend, default=None,
                    help="kernel per quantized layer type, as for ft serve; the expert banks are packed for the "
                         "chosen MoE kernel and the server picks that kernel back up from the FTW")
diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py
index a875d77f34..152f3871c0 100644
--- a/python/freetoken/checkpoint/convert.py
+++ b/python/freetoken/checkpoint/convert.py
@@ -200,7 +200,8 @@ def convert_checkpoint(
     cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size),
                        dtype=dtype, moe_strategy=moe_backend, quant_backend=quant_backend)
     mc = cfg.model_config
-    offload = moe_backend == "offload" and getattr(mc, "is_moe", False)
+    # every MoE strategy reads the routed experts from the banks, fused included
+    offload = getattr(mc, "is_moe", False)
     include_moe_experts = not offload
     method = None
     if offload:
@@ -208,6 +209,7 @@ def convert_checkpoint(
         from freetoken.layers import set_rope_device
 
         set_rope_device(dev)
+        object.__setattr__(cfg, "moe_strategy", "offload")
         method = offload_expert_method(cfg)
 
     from freetoken.utils.progress import byte_bar, count_bar
diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index e2bd0de24e..e31e9d8278 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -432,6 +432,43 @@ def _producer():
         raise err[0]
 
 
+def _split_bank_entries(reader: FTWReader, path: str, num_layers: int):
+    """``(all, alphas, flat regions, {base: {layer: entry}})`` of the ``experts_bank`` entries; None when there are none."""
+    bank_entries = reader.entries("experts_bank")
+    if not bank_entries:
+        return None
+
+    alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES]
+    row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES]
+
+    meta_layers = reader.meta("expert_bank_num_layers")
+    if meta_layers is not None and meta_layers != num_layers:
+        raise RuntimeError(
+            f"{path!r} was converted with {meta_layers} expert-bank layers but the "
+            f"model config says num_moe_layers={num_layers}; the checkpoint does not "
+            "match its config"
+        )
+
+    # Split row entries into the two layouts by name.
+    flat_entries: list[dict] = []
+    per_layer_groups: dict[str, dict[int, dict]] = {}
+    for e in row_entries:
+        m = _LAYER_ENTRY_RE.match(e["name"])
+        if m is None:
+            flat_entries.append(e)
+            continue
+        per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e
+
+    mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys()
+    assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+    for base, by_layer in per_layer_groups.items():
+        assert sorted(by_layer) == list(range(num_layers)), (
+            f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
+            f"expected exactly range({num_layers})"
+        )
+    return bank_entries, alpha_entries, flat_entries, per_layer_groups
+
+
 def load_ftw_banks(
     path: str, *, num_layers: int, workers: int = 8, chunk: int = _DEFAULT_CHUNK,
     layer_residency: list[str] | None = None,
@@ -482,40 +519,20 @@ def _backing(layer_id: int) -> str:
         return "mmap"
 
     reader = FTWReader(path)
-    bank_entries = reader.entries("experts_bank")
-    if not bank_entries:
+    try:
+        split = _split_bank_entries(reader, path, num_layers)
+    except BaseException:
         reader.close()
-        return None
-
-    alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES]
-    row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES]
-
-    meta_layers = reader.meta("expert_bank_num_layers")
-    if meta_layers is not None and meta_layers != num_layers:
+        raise
+    if split is None:
         reader.close()
-        raise RuntimeError(
-            f"{path!r} was converted with {meta_layers} expert-bank layers but the "
-            f"model config says num_moe_layers={num_layers}; the checkpoint does not "
-            "match its config"
-        )
+        return None
+    bank_entries, alpha_entries, flat_entries, per_layer_groups = split
 
     # Alphas: unchanged, one flat HostBank per entry.
     alpha_specs = {e["name"]: (tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries}
     alpha_hb = alloc_banks(alpha_specs)
 
-    # Split row entries into the two layouts by name.
-    flat_entries: list[dict] = []
-    per_layer_groups: dict[str, dict[int, dict]] = {}
-    for e in row_entries:
-        m = _LAYER_ENTRY_RE.match(e["name"])
-        if m is None:
-            flat_entries.append(e)
-            continue
-        per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e
-
-    mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys()
-    assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
-
     # Row banks: one padded-window HostBank per (name, layer_id) for the flat layout, plus
     # how to carve the real [num_experts, *row_shape] tensor out of its head; ``None`` marks
     # a per-layer entry (direct view, no carving needed).
@@ -546,10 +563,6 @@ def _backing(layer_id: int) -> str:
             row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
 
     for base, by_layer in per_layer_groups.items():
-        assert sorted(by_layer) == list(range(num_layers)), (
-            f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
-            f"expected exactly range({num_layers})"
-        )
         row_hb[base] = []
         row_view_args[base] = []
         for layer_id in range(num_layers):
@@ -659,8 +672,71 @@ def _read_layer(job):
     )
 
 
+def load_ftw_banks_to_device(path: str, *, num_layers: int, device: torch.device, workers: int = 8, chunk: int = _DEFAULT_CHUNK):
+    """:func:`load_ftw_banks` for resident experts: every bank lands on ``device``, one (bank, layer) at a time through a transient host buffer."""
+    from freetoken.moe.expert_banks import ExpertBanks
+    from freetoken.moe.legacy_format import canonical_role, kind_kernel_for
+    from freetoken.utils.progress import byte_bar
+
+    reader = FTWReader(path)
+    try:
+        split = _split_bank_entries(reader, path, num_layers)
+        if split is None:
+            return None
+        bank_entries, alpha_entries, flat_entries, per_layer_groups = split
+
+        # (name, layer_id or None for an alpha, read offset, read length, head pad, shape, dtype)
+        jobs = [(e["name"], None, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries]
+        for e in flat_entries:
+            total, *row_shape = e["shape"]
+            assert total % num_layers == 0, (e["name"], total, num_layers)
+            layer_bytes = total // num_layers * math.prod(row_shape) * _elsize(_dtype_of(e["dtype"]))
+            assert layer_bytes * num_layers == e["nbytes"], (e["name"], layer_bytes, num_layers, e["nbytes"])
+            for layer_id in range(num_layers):
+                off = e["global_off"] + layer_id * layer_bytes
+                win_off = (off // ALIGN) * ALIGN
+                jobs.append((e["name"], layer_id, win_off, _align_up(off + layer_bytes) - win_off, off - win_off, (total // num_layers, *row_shape), _dtype_of(e["dtype"])))
+        for base, by_layer in per_layer_groups.items():
+            for layer_id in range(num_layers):
+                e = by_layer[layer_id]
+                jobs.append((base, layer_id, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"])))
+
+        def _read(job) -> torch.Tensor:
+            _name, _layer, off, nbytes, head_pad, shape, dtype = job
+            buf = _transient_buffer(nbytes)
+            reader.read_into(memoryview(buf), {"global_off": off, "nbytes": nbytes}, workers=workers, chunk=chunk)
+            count = math.prod(shape) * _elsize(dtype)
+            return torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + count)[head_pad:].view(dtype).view(shape)
+
+        sources: dict[str, list] = {}
+        alphas: dict[str, torch.Tensor] = {}
+        bar = byte_bar(sum(e["nbytes"] for e in bank_entries), "Loading expert banks (FTW)")
+        try:
+            # read the next (bank, layer) while the current one copies to the device
+            with ThreadPoolExecutor(1) as ex:
+                ahead = ex.submit(_read, jobs[0])
+                for k, (name, layer_id, *_rest) in enumerate(jobs):
+                    host = ahead.result()
+                    if k + 1 < len(jobs):
+                        ahead = ex.submit(_read, jobs[k + 1])
+                    if layer_id is None:
+                        alphas[name] = host.to(device)
+                    else:
+                        sources.setdefault(canonical_role(name), [None] * num_layers)[layer_id] = host.to(device)
+                    bar.update(host.numel() * host.element_size())
+                    del host
+        finally:
+            bar.close()
+        quant_format = reader.meta("quant_format")
+    finally:
+        reader.close()
+
+    kind, kernel = kind_kernel_for(quant_format) if quant_format is not None else (None, None)
+    return ExpertBanks(quant_format, sources, **alphas, kind=kind, kernel=kernel)
+
+
 __all__ = [
     "INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT",
     "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader",
-    "iter_ftw_weights", "load_ftw_banks", "layer_bank_entry_name",
+    "iter_ftw_weights", "load_ftw_banks", "load_ftw_banks_to_device", "layer_bank_entry_name",
 ]
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index e855bcd32f..e5ccf71050 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -13,14 +13,15 @@
 from freetoken.core import Batch, Context, Req, set_global_ctx
 from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
 from freetoken.gpu_select import gpu_identity
-from freetoken.layers import set_rope_device
+from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint
+from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device
 from freetoken.layers.quantization import LayerKind, QuantBackend, finalize_quant, set_quant_backend
 from freetoken.moe.offload_cache import iter_offload_moe_layers
 from freetoken.mm.config import ENCODER_SECTIONS
 from freetoken.models import create_model, load_weight
 from freetoken.models.weight import ftw_lacks_vision
 from freetoken.moe import is_offload_moe_strategy
-from freetoken.moe.expert_banks import load_expert_banks
+from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks
 from freetoken.moe.host_banks import PinFailed
 from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache
 from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype
@@ -531,9 +532,14 @@ def _load_weights(self, config: EngineConfig) -> None:
                 "family served images. Reconvert it with `ft checkpoint`, add the encoder in place with "
                 "scripts/ftw_hotfix.py (docs/ftw-hotfix.md), or start with --text-model-only"
             )
+        resident = [layer for layer in iter_moe_layers(self.model) if not isinstance(layer, OffloadMoELayer)]
+        if resident and not config.use_dummy_weight and is_ftw_checkpoint(config.model_path) and not ftw_tensor_names(config.model_path, "experts_bank"):
+            raise ValueError(f"{config.model_path} keeps its routed experts as dense weights (an older `ft checkpoint --moe-backend fused`); reconvert it with `ft checkpoint`")
         with _weight_load_context():
             self.model.load_state_dict(self._load_weight_state_dict(config))
         finalize_quant(self.model)
+        if resident:
+            self._load_resident_experts(config, resident)
 
     def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]:
         model_state = self.model.state_dict()
@@ -541,18 +547,32 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso
             return _make_dummy_weight_state_dict(model_state, device=self.device)
         # _materialize casts each loaded tensor to its model-param dtype (model_state), so
         # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved;
-        # offload models exclude experts (served from the offload cache, not dense weights).
+        # MoE models exclude routed experts (they load as expert banks, not dense weights).
         return _materialize_loaded_weight_state_dict(
             model_state,
             load_weight(
                 config.model_path,
                 self.device,
-                include_moe_experts=not is_offload_moe_strategy(config.moe_strategy),
+                include_moe_experts=not getattr(config.model_config, "is_moe", False),
                 include_vision=bool(config.active_encoders),
             ),
             device=self.device,
         )
 
+    def _load_resident_experts(self, config: EngineConfig, layers) -> None:
+        with _weight_load_context():
+            banks = load_expert_banks(
+                config.model_path,
+                config.model_config,
+                method=layers[0].quant_method,
+                device=self.device,
+                dtype=self.dtype,
+                dummy=config.use_dummy_weight,
+                parallel={"serial": False, "parallel": True}.get(config.expert_load, None),
+                resident=True,
+            )
+        attach_resident_banks(layers, banks)
+
     @torch.inference_mode()
     def _warmup_encoders(self) -> None:
         for item in self.mm_processor.dummy_items(self.dtype, self.device):
@@ -1135,15 +1155,14 @@ def _is_unified_memory_gpu(index: "int | None" = None) -> bool:
         return False
 
 
+_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4")
+
+
 def _fused_resident_ok(model_config) -> bool:
-    """Whether the resident ('fused') MoE path can hold this model's experts.
-   
-       FIXME: auto resolves to fused only for bf16 and fp8_block experts; drop this gate once the other quant formats support fused.
-       """
-    expert_quant = getattr(model_config, "expert_quant", "none")
-    if expert_quant not in ("none", "fp8_block"):
+    """Whether the resident ('fused') MoE path can hold this model's experts; GGUF q4_0 and DeepSeek-V4 experts stay offload."""
+    if getattr(model_config, "expert_quant", "none") not in _RESIDENT_EXPERT_QUANTS:
         return False
-    return getattr(model_config, "moe_weight_format", None) in (None, "bf16")
+    return getattr(model_config, "moe_weight_format", None) in (None, "bf16", "mxfp4")
 
 
 def _ensure_expandable_segments() -> None:
@@ -1594,8 +1613,8 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     if is_moe and config.moe_strategy == "auto":
         # A MoE model always defaults to the offload family: experts stream from pinned host
         # banks into an auto-sized GPU slot cache, which is the only default that serves a model
-        # bigger than the GPU. The resident 'fused' path (bf16 / block-fp8 experts, the two
-        # formats MoELayer can allocate) is still reachable, but only when asked for explicitly
+        # bigger than the GPU. The resident 'fused' path (every expert format but GGUF q4_0 and
+        # DeepSeek-V4's) is still reachable, but only when asked for explicitly
         # -- auto never picks it, because nothing here knows whether the experts would fit in
         # HBM and a wrong guess is a weight-load OOM rather than a slower-but-working run.
         default_backend = "offload"
@@ -1614,6 +1633,18 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
                 "Unified-memory GPU detected; auto-selecting 'fused' MoE strategy "
                 "(resident experts) instead of offload"
             )
+            # the cache sizing flags get their own warning in the 'fused' block below
+            inert = [flag for flag, on in (
+                (f"--moe-cpu-layers={config.moe_cpu_layers}", config.moe_cpu_layers),
+                ("--moe-prefill-hit-d2d", config.moe_prefill_hit_d2d),
+                ("--disable-moe-prefill-overlap", not config.moe_prefill_overlap),
+            ) if on]
+            if inert:
+                logger.warning_rank0(
+                    f"Unified-memory GPU serves MoE experts resident; ignoring {', '.join(inert)} "
+                    "(use --moe-strategy offload to keep the offload flags)"
+                )
+                override("moe_cpu_layers", None)
         # Hardware-adaptive config: a cached `ft bench bw` profile can upgrade
         # the offload default to hybrid when this machine's CPU MoE bandwidth clears its PCIe
         # gather bandwidth by the bench threshold (default 2x). hybrid is VRAM-equivalent to
@@ -1733,7 +1764,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
 
     if (
         is_moe
-        and expert_quant not in ("none", "fp8_block")
+        and expert_quant not in _RESIDENT_EXPERT_QUANTS
         and not is_offload_moe_strategy(config.moe_strategy)
     ):
         raise ValueError(
diff --git a/python/freetoken/layers/__init__.py b/python/freetoken/layers/__init__.py
index d868751a6b..a07484620e 100644
--- a/python/freetoken/layers/__init__.py
+++ b/python/freetoken/layers/__init__.py
@@ -16,7 +16,7 @@
     LinearReplicated,
     LinearRowParallel,
 )
-from .moe import MoELayer, OffloadMoELayer, make_moe_layer
+from .moe import MoELayer, OffloadMoELayer, iter_moe_layers, make_moe_layer
 from .norm import (
     GatedRMSNorm,
     GemmaPlusOneRMSNorm,
@@ -57,5 +57,6 @@
     "LinearReplicated",
     "MoELayer",
     "OffloadMoELayer",
+    "iter_moe_layers",
     "make_moe_layer",
 ]
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index c209454996..996beaa796 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Iterator, Tuple
 
 import torch
 from freetoken.core import get_global_ctx
@@ -435,6 +435,19 @@ def _expert_gemm(
         raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}")
 
 
+def iter_moe_layers(model) -> Iterator[MoELayer]:
+    if isinstance(model, MoELayer):
+        yield model
+    if not isinstance(model, BaseOP):
+        return
+    for value in model.__dict__.values():
+        if isinstance(value, BaseOP):
+            yield from iter_moe_layers(value)
+        elif isinstance(value, (list, tuple)):
+            for item in value:
+                yield from iter_moe_layers(item)
+
+
 def make_moe_layer(
     config: "ModelConfig",
     *,
@@ -483,10 +496,10 @@ def make_moe_layer(
         has_bias=has_bias,
         quant_config=quant_config,
         prefix=prefix,
+        layer_id=layer_id,
     )
     if offload:
         assert layer_id is not None, "offload MoE backends need the layer_id"
-        kwargs["layer_id"] = layer_id
         kwargs["strategy"] = config.moe_strategy
         kwargs["decode_target"] = config.decode_target
     return layer_cls(**kwargs)
diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py
index 89f13ed09d..e4a428d3e4 100644
--- a/python/freetoken/layers/quantization/moe/base.py
+++ b/python/freetoken/layers/quantization/moe/base.py
@@ -32,6 +32,7 @@ class MoEConfig:
     apply_router_weight_on_input: bool = False
     strategy: str = "resident"
     decode_target: str = "gpu"
+    dtype: torch.dtype = torch.bfloat16
 
     @classmethod
     def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig":
@@ -52,6 +53,7 @@ def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig":
             apply_router_weight_on_input=bool(layer.apply_router_weight_on_input),
             strategy=layer.strategy,
             decode_target=layer.decode_target,
+            dtype=torch.get_default_dtype(),
         )
 
     @property
@@ -147,9 +149,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None:
         """Most GPU cache slots the kernel can address for ``cfg``; None for no limit."""
         return self.max_slots
 
-    def _common_reject(self, cfg: MoEConfig, *, resident_ok: bool, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None:
-        if not resident_ok and cfg.strategy == "resident":
-            return "not served resident; use --moe-strategy offload or cpu"
+    def _common_reject(self, cfg: MoEConfig, *, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None:
         if not tp_ok and cfg.tp_size > 1:
             return "TP > 1 is not supported for this expert format"
         if not cpu_ok and cfg.decode_target != "gpu":
@@ -190,12 +190,15 @@ def apply(self, x: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Ten
     def cpu_format(self) -> str | None:
         return self.kernel.cpu_format
 
-    @abstractmethod
-    def create_weights(self, layer: Any) -> None: ...
+    def create_weights(self, layer: Any) -> None:
+        # resident experts stay out of the state dict; attach_resident_banks fills these after the dense weights load
+        layer._expert_banks = None
+        layer._expert_alphas = None
 
     def finalize(self, layer: Any) -> None:
         pass
 
-    @abstractmethod
     def resident_view(self, layer: Any) -> ExpertView:
         """Layout roles -> the resident layer's tensors, for apply()."""
+        assert layer._expert_banks is not None, "resident expert banks were never loaded"
+        return ExpertView(layer._expert_banks, n=self.cfg.num_experts, alphas=layer._expert_alphas)
diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py
index 6c0167b71b..20c3dde9dc 100644
--- a/python/freetoken/layers/quantization/moe/fp8_block.py
+++ b/python/freetoken/layers/quantization/moe/fp8_block.py
@@ -21,7 +21,7 @@ class TritonFp8BlockMoEKernel(MoEKernel):
     name = "triton"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, resident_ok=True, tp_ok=False, cpu_ok=False, plain_silu_only=False)
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False)
         if reason:
             return reason
         reason = gated_epilogue_reason(cfg)
@@ -64,17 +64,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 @register_method(QuantKind.FP8_BLOCK, LayerKind.MOE)
 class Fp8BlockMoEMethod(MoEMethod):
     candidates = (TritonFp8BlockMoEKernel,)
-
-    def create_weights(self, layer) -> None:
-        g = self.cfg
-        e, i, h, b = g.num_experts, g.intermediate, g.hidden, BLOCK
-        layer.gate_up_proj = torch.empty(e, 2 * i, h, dtype=FP8)
-        layer.gate_up_scale_inv = torch.empty(e, 2 * i // b, h // b, dtype=torch.bfloat16)
-        layer.down_proj = torch.empty(e, h, i, dtype=FP8)
-        layer.down_scale_inv = torch.empty(e, h // b, i // b, dtype=torch.bfloat16)
-
-    def resident_view(self, layer) -> ExpertView:
-        return ExpertView({
-            "gate_up": layer.gate_up_proj, "gate_up_scale": layer.gate_up_scale_inv,
-            "down": layer.down_proj, "down_scale": layer.down_scale_inv,
-        })
diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py
index aa020ee33a..13cd97466c 100644
--- a/python/freetoken/layers/quantization/moe/mxfp4.py
+++ b/python/freetoken/layers/quantization/moe/mxfp4.py
@@ -22,7 +22,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
             return "standard MXFP4 kernel reads the concatenated gate|up row order"
         if (cfg.alpha, cfg.beta) != (1.0, 0.0):
             return "standard MXFP4 kernel has no alpha / beta in its swiglu"
-        return self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+        return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         i, h = cfg.intermediate, cfg.hidden
@@ -106,37 +106,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 @register_method(QuantKind.MXFP4, LayerKind.MOE)
 class Mxfp4MoEMethod(MoEMethod):
     candidates = (TritonMxfp4MoEKernel, TritonGptossMxfp4MoEKernel)
-
-    def create_weights(self, layer) -> None:
-        if not self.cfg.has_bias:
-            raise NotImplementedError("standard MXFP4 experts are served from the offload cache, not resident")
-        g = self.cfg
-        e, i, h = g.num_experts, g.local_intermediate, g.hidden
-        if h % GROUP:
-            raise ValueError(f"MXFP4 hidden size must be divisible by {GROUP}")
-        layer.gate_up_proj_blocks = torch.empty(e, 2 * i, h // GROUP, 16, dtype=torch.uint8)
-        layer.gate_up_proj_scales = torch.empty(e, 2 * i, h // GROUP, dtype=torch.uint8)
-        layer.gate_up_proj_bias = torch.empty(e, 2 * i, dtype=torch.bfloat16)
-        layer.down_proj_blocks = torch.empty(e, h, i // GROUP, 16, dtype=torch.uint8)
-        layer.down_proj_scales = torch.empty(e, h, i // GROUP, dtype=torch.uint8)
-        layer.down_proj_bias = torch.empty(e, h, dtype=torch.bfloat16)
-
-    def finalize(self, layer) -> None:
-        if getattr(layer, "gate_up_proj_blocks", None) is None:
-            return
-        from freetoken.moe.fused_mxfp4 import _transpose_mxfp4_for_decode
-
-        # one transposed copy serves prefill and decode; the HF blocks are freed so 120B fits
-        layer._gu_blocks_t, layer._gu_scales_t = _transpose_mxfp4_for_decode(layer.gate_up_proj_blocks, layer.gate_up_proj_scales)
-        layer._dn_blocks_t, layer._dn_scales_t = _transpose_mxfp4_for_decode(layer.down_proj_blocks, layer.down_proj_scales)
-        layer.gate_up_proj_blocks = None
-        layer.gate_up_proj_scales = None
-        layer.down_proj_blocks = None
-        layer.down_proj_scales = None
-        torch.cuda.empty_cache()
-
-    def resident_view(self, layer) -> ExpertView:
-        return ExpertView({
-            "gate_up": layer._gu_blocks_t, "gate_up_scale": layer._gu_scales_t, "gate_up_bias": layer.gate_up_proj_bias,
-            "down": layer._dn_blocks_t, "down_scale": layer._dn_scales_t, "down_bias": layer.down_proj_bias,
-        })
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index bb341cdfb7..cf66e2071b 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -1,4 +1,4 @@
-"""NVFP4 experts: served from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x.
+"""NVFP4 experts: served resident or from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x.
 
 FreeToken stores MoE experts as ModelOpt NVFP4 (packed e2m1 codes + fp8-e4m3 per-16
 block scales + per-tensor global scale) in pinned host banks and gathers the routed
@@ -41,7 +41,7 @@ class TritonNvfp4MoEKernel(MoEKernel):
     cpu_format = "nvfp4"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
         if reason:
             return reason
         reason = gated_epilogue_reason(cfg)
@@ -230,7 +230,7 @@ class MarlinNvfp4MoEKernel(MoEKernel):
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
         if not backend.is_vllm_installed():
             return "vLLM is not installed"
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True)
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True)
         if reason:
             return reason
         if not _marlin_symbols_ok():
@@ -496,7 +496,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
             return f"b12x requires sm_120+, got sm_{cc[0]}{cc[1]}"
         if not backend.is_flashinfer_installed():
             return "flashinfer is not installed"
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True)
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True)
         if reason:
             return reason
         return _b12x_unusable_reason(cc)
@@ -567,9 +567,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 @register_method(QuantKind.NVFP4, LayerKind.MOE)
 class Nvfp4MoEMethod(MoEMethod):
     candidates = (TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel)
-
-    def create_weights(self, layer) -> None:
-        raise NotImplementedError("NVFP4 experts are served from the offload cache, not resident")
-
-    def resident_view(self, layer) -> ExpertView:
-        raise NotImplementedError("NVFP4 experts are not resident")
diff --git a/python/freetoken/layers/quantization/moe/unquantized.py b/python/freetoken/layers/quantization/moe/unquantized.py
index ef435686c5..ded7d8135c 100644
--- a/python/freetoken/layers/quantization/moe/unquantized.py
+++ b/python/freetoken/layers/quantization/moe/unquantized.py
@@ -19,7 +19,9 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         i = cfg.local_intermediate
-        return {"gate_up": BankSpec((2 * i, cfg.hidden), torch.bfloat16), "down": BankSpec((cfg.hidden, i), torch.bfloat16)}
+        # resident experts compute in the model dtype; offload banks stay bf16, the FTW and CPU executor format
+        dt = cfg.dtype if cfg.strategy == "resident" else torch.bfloat16
+        return {"gate_up": BankSpec((2 * i, cfg.hidden), dt), "down": BankSpec((cfg.hidden, i), dt)}
 
     def pack(self, pieces, cfg: MoEConfig, out):
         out["gate_up"].copy_(fused_piece(pieces, "gate_up"))
@@ -37,12 +39,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 @register_method(QuantKind.NONE, LayerKind.MOE)
 class UnquantizedMoEMethod(MoEMethod):
     candidates = (FusedMoEKernel,)
-
-    def create_weights(self, layer) -> None:
-        g = self.cfg
-        i = g.local_intermediate
-        layer.gate_up_proj = torch.empty(g.num_experts, 2 * i, g.hidden)
-        layer.down_proj = torch.empty(g.num_experts, g.hidden, i)
-
-    def resident_view(self, layer) -> ExpertView:
-        return ExpertView({"gate_up": layer.gate_up_proj, "down": layer.down_proj})
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index 1bac8227b0..de5b9e8983 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -1,6 +1,6 @@
 """Qwen3.5 / 3.6 / 3.8 checkpoint reader.
 
-The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the offload cache (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only bf16 stacked experts and resident block-fp8 experts come from here.
+The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the expert-bank loader (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only the bf16 stacked experts it reads come from here.
 """
 
 from __future__ import annotations
@@ -241,9 +241,9 @@ def iter_weights(
     include_non_moe: bool,
     include_vision: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
-    """Yield the dense weights fused to the model's buffers, and the routed experts only where a resident path takes them from here: bf16 stacked experts as stored, block-fp8 experts restacked per layer.
+    """Yield the dense weights fused to the model's buffers, and with ``include_moe_experts`` the bf16 stacked experts as stored.
 
-    Per-expert NVFP4 experts always come from the offload cache's expert reader.
+    Block-fp8 and NVFP4 experts always come from the expert-bank reader.
     """
     if get_tp_info().size > 1:
         raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only")
@@ -253,8 +253,6 @@ def iter_weights(
     if include_non_moe or stacked:
         reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0])) if include_non_moe else None
         yield from _iter_shards(model_path, device, reader, stacked=stacked, include_vision=include_vision)
-    if include_moe_experts and config.is_moe and config.expert_quant == "fp8_block":
-        yield from _resident_fp8_experts(model_path, config)
 
 
 def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | None, *, stacked: bool, include_vision: bool):
@@ -327,7 +325,7 @@ def _is_expert(raw_name: str) -> bool:
 
 
 # ======================================================================================
-# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): offload expert pieces and resident stacks.
+# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): expert pieces.
 # ======================================================================================
 
 # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the
@@ -338,35 +336,6 @@ def _is_expert(raw_name: str) -> bool:
 )
 
 
-def _resident_fp8_experts(model_path, config):
-    from freetoken.kernel.triton.fp8_block_linear import FP8
-
-    B = 128
-    L, E, H, I, dense = _moe_dims(config)
-    shapes = {
-        "gate_up_proj": ((E, 2 * I, H), FP8),
-        "gate_up_scale_inv": ((E, 2 * I // B, H // B), torch.bfloat16),
-        "down_proj": ((E, H, I), FP8),
-        "down_scale_inv": ((E, H // B, I // B), torch.bfloat16),
-    }
-    layers: dict[int, dict[str, torch.Tensor]] = {}
-    placed = [0] * L
-    for li, e0, e1, piece in iter_expert_pieces(model_path, config, QuantKind.FP8_BLOCK, parallel=None):
-        stack = layers.setdefault(li, {n: torch.empty(shape, dtype=dt) for n, (shape, dt) in shapes.items()})
-        stack["gate_up_proj"][e0:e1, :I] = piece["gate"]
-        stack["gate_up_proj"][e0:e1, I:] = piece["up"]
-        stack["gate_up_scale_inv"][e0:e1, : I // B] = piece["gate_scale"]
-        stack["gate_up_scale_inv"][e0:e1, I // B :] = piece["up_scale"]
-        stack["down_proj"][e0:e1] = piece["down"]
-        stack["down_scale_inv"][e0:e1] = piece["down_scale"]
-        placed[li] += e1 - e0
-        if placed[li] == E:
-            pre = f"model.layers.{dense + li}.mlp.experts"
-            for name, tensor in layers.pop(li).items():
-                yield f"{pre}.{name}", tensor
-    assert not layers, f"incomplete resident fp8 experts for layers {sorted(layers)}"
-
-
 def _moe_dims(model_config):
     L = model_config.num_moe_layers
     return (
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index e3a768b6d0..7500174ae4 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -75,6 +75,7 @@ def build_expert_banks(
     device: torch.device,
     layer_sink=None,
     dummy: bool = False,
+    resident: bool = False,
 ) -> ExpertBanks:
     """Fill host banks in the kernel's layout from a stream of expert pieces.
 
@@ -83,6 +84,7 @@ def build_expert_banks(
     complete once its ``num_experts`` rows have arrived: with ``layer_sink=None`` its banks
     are pinned in the background, otherwise the sink receives them (converter). ``dummy``
     skips the pieces and fills the banks with finite random contents.
+    ``resident`` packs into banks on ``device`` instead of host banks (the resident 'fused' experts).
     """
     from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, pin_banks
     from freetoken.moe.legacy_format import legacy_format_for
@@ -91,8 +93,13 @@ def build_expert_banks(
     layout = method.layout()
     E = method.cfg.num_experts
     specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident}
-    hb = alloc_layer_banks(specs, num_layers)
-    banks = {role: [b.tensor for b in hb[role]] for role in specs}
+    if resident:
+        assert layer_sink is None
+        hb = None
+        banks = {role: [torch.empty(shape, dtype=dtype, device=device) for _ in range(num_layers)] for role, (shape, dtype) in specs.items()}
+    else:
+        hb = alloc_layer_banks(specs, num_layers)
+        banks = {role: [b.tensor for b in hb[role]] for role in specs}
     alphas = {
         role: torch.empty(num_layers * E, dtype=spec.dtype, device=device)
         for role, spec in layout.items() if spec.resident
@@ -104,7 +111,7 @@ def build_expert_banks(
                 _dummy_fill(role, tensor)
         for alpha in alphas.values():
             alpha.fill_(1.0)
-        if torch.cuda.is_available():
+        if hb is not None and torch.cuda.is_available():
             pin_banks(hb)
         return ExpertBanks(
             legacy_format_for(method.kind, kernel.name), banks,
@@ -136,7 +143,7 @@ def _fill(sink) -> None:
 
     if layer_sink is not None:
         _fill(layer_sink)
-    elif torch.cuda.is_available():
+    elif torch.cuda.is_available() and not resident:
         with PinPipeline() as pins:
             _fill(pins)
     else:
@@ -192,16 +199,16 @@ def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, paralle
     )
 
 
-def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks:
+def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None, resident=False) -> ExpertBanks:
     from freetoken.moe.expert_pieces import iter_expert_pieces
 
     num_layers = model_config.num_moe_layers
     if dummy:
-        return build_expert_banks(method, num_layers, None, device=device, dummy=True)
+        return build_expert_banks(method, num_layers, None, device=device, dummy=True, resident=resident)
     pieces = iter_expert_pieces(
         model_path, model_config, method.kind, parallel=parallel, workers=workers, chunk=chunk
     )
-    return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink)
+    return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink, resident=resident)
 
 
 def _host_ram_fits_parallel(model_path: str) -> bool:
@@ -286,6 +293,7 @@ def load_expert_banks(
     decode_target: str = "gpu",
     layer_sink=None,
     layer_residency: list[str] | None = None,
+    resident: bool = False,
 ) -> ExpertBanks:
     """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning
     the same normalized ``ExpertBanks`` and both pinning after fill:
@@ -310,14 +318,21 @@ def load_expert_banks(
 
     ``layer_residency``: per-layer ``HostResidency`` labels applied at settle time -- explicitly on the FTW fast path, ambiently (``requested_residency``) in the slow-path providers.
     Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota).
+
+    ``resident`` (with ``method``): the banks land on ``device`` for the resident 'fused' experts; no host bank outlives the load.
     """
-    from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks
+    from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks, load_ftw_banks_to_device
 
+    if resident:
+        assert method is not None and layer_sink is None and layer_residency is None
     if model_path and is_ftw_checkpoint(model_path) and not dummy:
-        banks = load_ftw_banks(
-            model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk,
-            layer_residency=layer_residency,
-        )
+        if resident:
+            banks = load_ftw_banks_to_device(model_path, num_layers=model_config.num_moe_layers, device=device, workers=workers, chunk=chunk)
+        else:
+            banks = load_ftw_banks(
+                model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk,
+                layer_residency=layer_residency,
+            )
         if banks is not None:
             logger.info_rank0(f"expert banks: FTW fast path (FTW checkpoint {model_path})")
             return banks
@@ -354,7 +369,7 @@ def load_expert_banks(
 
     def _build(par: bool) -> ExpertBanks:
         if method is not None:
-            return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink)
+            return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident)
         return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink)
 
     with requested_residency(layer_residency) as residency_plan:
@@ -393,3 +408,33 @@ def _echo_residency(banks: ExpertBanks, requested, plan) -> ExpertBanks:
             "but saves no pinned quota"
         )
     return banks
+
+
+def attach_resident_banks(layers, banks: ExpertBanks) -> None:
+    """Give each resident MoE layer its rows of the device ``banks``; bank layer ``i`` goes to the layer whose ``layer_id`` is ``i``."""
+    ids = sorted(layer.layer_id for layer in layers if layer.layer_id is not None)
+    num_layers = len(next(iter(banks.sources.values())))
+    if ids != list(range(num_layers)) or len(ids) != len(layers):
+        raise ValueError(f"resident MoE layer ids {ids} do not cover the {num_layers} expert bank layers")
+    for layer in layers:
+        i, e = layer.layer_id, layer.num_experts
+        layout = layer.quant_method.layout()
+        rows = {role: (e, *spec.shape) for role, spec in layout.items() if not spec.resident}
+        if set(banks.sources) != set(rows):
+            raise ValueError(f"expert banks hold {sorted(banks.sources)}, the {layer.quant_method.kernel.name} kernel reads {sorted(rows)}")
+        for role, shape in rows.items():
+            bank = banks.sources[role][i]
+            if tuple(bank.shape) != shape or bank.dtype != layout[role].dtype:
+                raise ValueError(
+                    f"bank {role!r} of layer {i} is {tuple(bank.shape)} {bank.dtype} but the {layer.quant_method.kernel.name} kernel "
+                    f"wants {shape} {layout[role].dtype}; the banks were packed for another kernel or model"
+                )
+        alphas = []
+        # apply() reads view.alphas as (gate_up, down)
+        for role in (role for role in ("gate_up_alpha", "down_alpha") if role in layout):
+            a = getattr(banks, role)
+            if a is None or a.numel() != num_layers * e or a.dtype != layout[role].dtype:
+                raise ValueError(f"the {layer.quant_method.kernel.name} kernel needs {num_layers * e} {layout[role].dtype} {role} values, the expert banks carry {None if a is None else (a.numel(), a.dtype)}")
+            alphas.append(a)
+        layer._expert_banks = {role: per_layer[i] for role, per_layer in banks.sources.items()}
+        layer._expert_alphas = tuple(a[i * e : (i + 1) * e] for a in alphas) or None
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 32e0f3cc1e..c6f975df7e 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -1054,20 +1054,9 @@ def copy_missing(self) -> None:
 
 
 def iter_offload_moe_layers(model) -> Iterator:
-    from freetoken.layers import BaseOP, OffloadMoELayer
+    from freetoken.layers import OffloadMoELayer, iter_moe_layers
 
-    if isinstance(model, OffloadMoELayer):
-        yield model
-
-    if not isinstance(model, BaseOP):
-        return
-
-    for value in model.__dict__.values():
-        if isinstance(value, BaseOP):
-            yield from iter_offload_moe_layers(value)
-        elif isinstance(value, (list, tuple)):
-            for item in value:
-                yield from iter_offload_moe_layers(item)
+    return (layer for layer in iter_moe_layers(model) if isinstance(layer, OffloadMoELayer))
 
 
 def attach_offload_moe_cache(model, cache: OffloadMoeCache) -> list:
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 23a62f8419..4dd57b1ae5 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -600,8 +600,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         choices=["auto", *MOE_STRATEGIES],
         help=(
             "How the routed experts are served. 'auto' resolves a MoE model to the offload family "
-            "(offload, or hybrid when a `ft bench bw` profile recommends it); resident "
-            "'fused' experts must be requested explicitly."
+            "(offload, or hybrid when a `ft bench bw` profile recommends it), and to resident "
+            "'fused' experts on unified-memory GPUs (GB10 / DGX Spark)."
         ),
     )
 
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 073112bb5a..a77409f013 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -367,7 +367,7 @@ def _offload_engine_config(**overrides):
             num_layers=10,
             num_moe_layers=10,
             num_experts=8,
-            expert_quant="nvfp4",  # quantized experts -> must resolve to an offload backend
+            expert_quant="nvfp4",
             moe_strategy="auto",
         ),
     )
@@ -389,7 +389,7 @@ def test_guard_raises_actionable_error_when_too_small():
     assert "128" in msg and "moe-cache" in msg
 
 
-def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend():
+def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(monkeypatch):
     """Bare `ft serve <FTW MoE checkpoint>`: no --moe-backend, no --moe-cache-* flags at all.
 
     args.py's parse-time default only fires when the backend is *already*
@@ -400,9 +400,12 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend
     moe_cache_auto=True, so _init_offload_moe_cache's _require_offload_cache_size guard is never
     reached with moe_cache_size still 0.
     """
+    import freetoken.engine.engine as engine_module
     from freetoken.engine.engine import _adjust_config
     from freetoken.moe import is_offload_moe_strategy
 
+    # a unified-memory host (GB10) resolves auto to fused instead
+    monkeypatch.setattr(engine_module, "_is_unified_memory_gpu", lambda index=None: False)
     config = _offload_engine_config()
     _adjust_config(config)
 
diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py
index ac598e65df..8fb9112df2 100644
--- a/tests/models/qwen4_exp/test_skeleton.py
+++ b/tests/models/qwen4_exp/test_skeleton.py
@@ -16,7 +16,7 @@
 import torch
 import torch.nn.functional as F
 
-from freetoken.layers import BaseOP, LinearReplicated
+from freetoken.layers import BaseOP, LinearReplicated, iter_moe_layers
 from freetoken.models.config import ModelConfig
 from freetoken.models.qwen4_exp.config import parse_config
 from freetoken.models.qwen4_exp.hc import GatedResidual
@@ -30,12 +30,17 @@ def _config(num_layers: int = 4) -> ModelConfig:
 
 
 def _fill(op, gen: torch.Generator, scale: float = 0.05) -> None:
-    """Random floats / zeroed ints for every state-dict tensor of an op tree."""
+    """Random floats / zeroed ints for every state-dict tensor of an op tree, and random expert banks for its resident MoE layers."""
     for tensor in op.state_dict().values():
         if tensor.is_floating_point():
             tensor.normal_(0.0, scale, generator=gen)
         else:
             tensor.zero_()
+    for layer in iter_moe_layers(op):
+        layer._expert_banks = {
+            role: torch.randn((layer.num_experts, *spec.shape), generator=gen, device=gen.device, dtype=spec.dtype) * scale
+            for role, spec in layer.quant_method.layout().items()
+        }
 
 
 def _group_norm(x, weight, eps, groups):
diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py
index 6ac1b84b17..084a2e4aa7 100644
--- a/tests/models/test_quant_config.py
+++ b/tests/models/test_quant_config.py
@@ -121,14 +121,6 @@ class Case:
     check: Callable[[Any, dict], None] | None = None  # family extras beyond the method table
 
 
-def _installed(module: str) -> bool:
-    try:
-        __import__(module)
-        return True
-    except ImportError:
-        return False
-
-
 def _qwen35_fp8(model, idx):
     assert model.model.layers.op_list[idx["lin"]].linear_attn._split_in_proj
 
@@ -162,13 +154,18 @@ def _dsv4(model, idx):
     assert "model.head.weight" in keys and not any(".ffn.experts." in k for k in keys)
 
 
+def _banked_experts(model, idx):
+    experts = op_at(model, f"model.layers.{idx['lin']}.mlp.experts")
+    assert not experts.state_dict()
+
+
 def _gpt_oss(cls):
     def check(model, idx):
         experts = op_at(model, "model.layers.0.mlp.experts")
         assert type(experts) is cls
         assert experts.interleaved and experts.alpha == pytest.approx(1.702) and experts.limit == 7.0
         if cls is GptOssMoELayer:
-            assert experts.gate_up_proj_blocks.dtype is torch.uint8 and experts.gate_up_proj_blocks.dim() == 4
+            assert not experts.state_dict()
 
     return check
 
@@ -198,9 +195,8 @@ def _gemma_tied_head(model, idx):
         "model.layers.{lin}.mlp.shared_expert.down_proj": NVFP4, "model.layers.{lin}.mlp.experts": Nvfp4MoEMethod,
         "lm_head": NVFP4,
     }, check=_qwen35_nvfp4),
-    # resident NVFP4 experts have no kernel without vLLM's Marlin
-    Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {}, strategy="fused", raises=KernelSelectionError,
-         skip_if=lambda: not _installed("flashinfer") or _installed("vllm")),
+    # resident NVFP4 experts load through the expert-bank loader, not the state dict
+    Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {"model.layers.{lin}.mlp.experts": (Nvfp4MoEMethod, "triton")}, strategy="fused", check=_banked_experts),
     Case("Qwen3.8-27B-NVFP4", ModelOptConfig, {
         "model.layers.{lin}.linear_attn.in_proj_qkvz": FP8T, "model.layers.{full}.self_attn.o_proj": FP8T,
         "model.layers.{lin}.mlp.gate_up_proj": NVFP4, "model.layers.{lin}.mlp.down_proj": NVFP4, "lm_head": NVFP4,

From 581fdca310d429cca74ac9d94ed2b893b86d9320 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=91=E5=BE=AE?= <1067852565@qq.com>
Date: Mon, 5 Oct 2026 09:24:33 +0200
Subject: [PATCH 064/249] fix(server): report the backend kv page size in stats
 (#599)

* fix(server): report the backend kv page size in stats

Use resolved backend cache metadata, retaining the frontend fallback.

Refs FlashML-org/FreeToken#598
Assisted-by: Codex

* test(server): cover resolved kv page size and fallback

Refs FlashML-org/FreeToken#598
Assisted-by: Codex
---
 python/freetoken/server/stats.py           |  3 ++-
 tests/server/test_generation_accounting.py | 13 +++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py
index 1857d3d7c0..d01828e21b 100644
--- a/python/freetoken/server/stats.py
+++ b/python/freetoken/server/stats.py
@@ -138,9 +138,10 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict:
     config = state.config
     ready_at = getattr(state, "ready_at", None)
     uptime_s = max(0, int(time.monotonic() - ready_at)) if ready_at is not None else 0
+    pools = getattr(state, "cache_pools", None) or {}
     kv = (
         {"used_pages": tr.kv_used_pages, "total_pages": tr.kv_total_pages,
-         "page_size": getattr(config, "page_size", 1)}
+         "page_size": pools.get("page_size") or getattr(config, "page_size", 1)}
         if tr.kv_total_pages > 0 else None
     )
     mamba = (
diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py
index 8276163b0c..a6d2bdb916 100644
--- a/tests/server/test_generation_accounting.py
+++ b/tests/server/test_generation_accounting.py
@@ -91,6 +91,19 @@ def _last_row() -> dict:
     return rows[-1]
 
 
+@pytest.mark.parametrize("pools, expected", [(None, 1), ({"page_size": 64}, 64)])
+def test_stats_reports_backend_page_size(pools, expected):
+    from freetoken.server.stats import StatsTracker, build_stats
+
+    tracker = StatsTracker()
+    tracker.kv_total_pages = 130
+    config = SimpleNamespace(served_model_name="unit-model", max_seq_len=262144,
+                             page_size=1, served_modalities=(), model_config=SimpleNamespace())
+    state = SimpleNamespace(stats=tracker, config=config, cache_pools=pools)
+    doc = build_stats(state, 0, 0)
+    assert doc["kv"] == {"used_pages": 0, "total_pages": 130, "page_size": expected}
+
+
 def test_non_stream_records_the_request_with_real_token_totals():
     request_ring.reset()
     st = FakeState([_ack(prompt=5, completion=1, out="a"), _ack(completion=2, out="bc", finished=True)])

From 93b0a59164bb6d68749298f868ca3738fd464143 Mon Sep 17 00:00:00 2001
From: William Xinze Zheng <wxzheng@berkeley.edu>
Date: Mon, 5 Oct 2026 14:59:59 -0700
Subject: [PATCH 065/249] fix(server): keep Anthropic inline system messages in
 place for prefix reuse (#608)

* Implement fix for Claude Code's inline system message

* Fix for Muse
---
 docs/cli.md                                 |   2 +-
 python/freetoken/message/tokenizer.py       |   1 +
 python/freetoken/server/anthropic_api.py    |  19 +-
 python/freetoken/server/args.py             |   9 +
 python/freetoken/server/generation.py       |   9 +
 python/freetoken/tokenizer/inline_system.py | 173 +++++++
 python/freetoken/tokenizer/tokenize.py      |  31 +-
 tests/README.md                             |   5 +-
 tests/server/test_anthropic_api.py          | 103 ++++-
 tests/server/test_message_wire.py           |  12 +
 tests/tokenizer/test_tokenize.py            | 476 ++++++++++++++++++++
 11 files changed, 822 insertions(+), 18 deletions(-)
 create mode 100644 python/freetoken/tokenizer/inline_system.py

diff --git a/docs/cli.md b/docs/cli.md
index deeda18419..b3cec99826 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -99,6 +99,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--tool-call-parser` | auto | Tool-call format; auto-inferred from the model family |
 | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables |
 | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block |
+| `--anthropic-inline-system` | auto | Placement of late Anthropic system instructions: `auto`, `preserve`, or `fold` |
 
 ### Image input
 
@@ -203,4 +204,3 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then
 - What to measure: `--dtype`, `--model`, `--formats`, `--isa`.
 - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe
   by that factor.
-
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index 9442ee059e..d52d658e63 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -73,6 +73,7 @@ class TokenizeMsg(BaseTokenizerMsg):
     chat_template_kwargs: Dict[str, Any] | None = None
     tools: List[Dict[str, Any]] | None = None
     images: List[bytes] | None = None
+    inline_system_policy: str | None = None
 
 
 @dataclass
diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 0d60940c97..d9a6972f09 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -120,6 +120,7 @@ async def handle_anthropic_messages(
             default_max_tokens=(
                 getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS
             ),
+            inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"),
         )
         uid = await submit_generation(spec, state)
     except ValueError as exc:
@@ -164,7 +165,10 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state:
             400, "invalid_request_error", "messages: no tokenizable content"
         )
     try:
-        n_tokens = await count_prompt_tokens(messages, template_tools, ctk, state)
+        n_tokens = await count_prompt_tokens(
+            messages, template_tools, ctk, state,
+            inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"),
+        )
     except GenerationError as exc:
         # The chat template could not render this conversation (bad role ordering, an unmatched
         # tool_result, ...) — a client error, exactly as /v1/messages classifies the same failure.
@@ -186,9 +190,8 @@ def convert_anthropic_prompt(
     """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt
     side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so
     a counted prompt is exactly the prompt a generation would tokenize."""
-    # Collect all system content (top-level `system` + any system-role messages
-    # Claude Code interleaves in the array) and emit ONE system message at the
-    # front: strict chat templates (e.g. Qwen3.5) require system at the beginning.
+    # Hoisting turn-specific reminders rewrites the cached prompt prefix.
+    # Preserve their position; the tokenizer adapts them to the model's roles.
     system_texts: list[str] = []
     if req.system:
         if isinstance(req.system, str):
@@ -201,7 +204,11 @@ def convert_anthropic_prompt(
     other: list[dict[str, Any]] = []
     for msg in req.messages:
         if msg.role == "system":
-            system_texts.append(_content_text(msg.content))
+            text = _content_text(msg.content)
+            if not other:
+                system_texts.append(text)
+            elif text:
+                other.append({"role": "system", "content": text})
             continue
 
         if isinstance(msg.content, str):
@@ -308,6 +315,7 @@ def convert_anthropic_to_genspec(
     model_sampling: dict[str, Any],
     reasoning_parser: str | None = None,
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    inline_system_policy: str = "auto",
 ) -> GenSpec:
     messages, template_tools, parser_tools, ctk = convert_anthropic_prompt(
         req, reasoning_parser=reasoning_parser
@@ -327,6 +335,7 @@ def convert_anthropic_to_genspec(
         chat_template_kwargs=ctk,
         template_tools=template_tools,
         parser_tools=parser_tools,
+        inline_system_policy=inline_system_policy,
     )
 
 
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 4dd57b1ae5..2194b82e92 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -59,6 +59,7 @@ class ServerArgs(SchedulerConfig):
     # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses
     # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report.
     enable_cache_report: bool = False
+    anthropic_inline_system: str = "auto"
     # Comma-separated hostname allowlist for client-supplied image URLs; empty admits any domain.
     allowed_media_domains: str = ""
     # Directory file:// image refs may be read from; empty rejects local files.
@@ -535,6 +536,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--anthropic-inline-system",
+        choices=("auto", "preserve", "fold"),
+        default=ServerArgs.anthropic_inline_system,
+        help="Preserve inline system instructions when supported by the renderer, "
+        "or fold them into nearby user/tool content without hoisting the prompt prefix.",
+    )
+
     parser.add_argument(
         "--sampling-defaults",
         type=str,
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 91056e8d1f..ac988eddc7 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -25,6 +25,7 @@
 from freetoken.message import TokenizeMsg
 from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason
 from freetoken.tokenizer.tokenize import resolve_thinking_mode
+from freetoken.tokenizer.inline_system import InlineSystemError
 
 try:
     # Chat templates render through jinja2 (a transformers dependency): a TemplateError means
@@ -142,6 +143,7 @@ class GenSpec:
     chat_template_kwargs: dict[str, Any] = field(default_factory=dict)
     template_tools: list[dict[str, Any]] | None = None   # tools the model sees (TokenizeMsg.tools)
     parser_tools: list[dict[str, Any]] | None = None     # tools for FunctionCallParser; None disables parsing
+    inline_system_policy: str | None = None
 
     @property
     def parse_tools(self) -> bool:
@@ -302,6 +304,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int:
             chat_template_kwargs=spec.chat_template_kwargs,
             tools=spec.template_tools,
             images=images,
+            inline_system_policy=spec.inline_system_policy,
         )
     )
     return uid
@@ -312,6 +315,8 @@ async def count_prompt_tokens(
     tools: list[dict[str, Any]] | None,
     chat_template_kwargs: dict[str, Any],
     state: Any,
+    *,
+    inline_system_policy: str | None = None,
 ) -> int:
     """Token count of an already-converted (messages, tools, chat_template_kwargs) prompt,
     using the frontend's own tokenizer (``state.frontend_tokenizer()``) so the count equals the
@@ -334,10 +339,13 @@ class the generation path maps to a 400. A tokenizer *initialization* failure (m
         chat_template_kwargs=chat_template_kwargs,
         tools=tools,
         images=images,
+        inline_system_policy=inline_system_policy,
     )
     manager = await asyncio.to_thread(state.frontend_tokenizer)  # init failure -> server fault
     try:
         (user_msg,) = await asyncio.to_thread(manager.tokenize, [msg])
+    except InlineSystemError as exc:
+        raise GenerationError(str(exc)) from exc
     except _TemplateError as exc:
         raise GenerationError(str(exc)) from exc
     return int(user_msg.input_ids.numel())
@@ -361,6 +369,7 @@ async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None:
         sampling_params=SamplingParams(),
         chat_template_kwargs=spec.chat_template_kwargs,
         tools=spec.template_tools,
+        inline_system_policy=spec.inline_system_policy,
     )
     try:
         manager = await asyncio.to_thread(build)
diff --git a/python/freetoken/tokenizer/inline_system.py b/python/freetoken/tokenizer/inline_system.py
new file mode 100644
index 0000000000..7e6f2c06f9
--- /dev/null
+++ b/python/freetoken/tokenizer/inline_system.py
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# Renderer-probe approach adapted from vLLM PR #58772.
+"""Preserve or fold late system instructions without hoisting the prompt head."""
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, Callable
+
+
+# Allow small role-framing differences, not a repeated environment header.
+_INLINE_ROLE_SLACK_TOKENS = 8
+
+
+class InlineSystemError(ValueError):
+    """An input instruction cannot be placed without corrupting tool ordering."""
+
+
+def has_inline_system(messages: list[dict]) -> bool:
+    started = False
+    for message in messages:
+        if message.get("role") == "system":
+            if started:
+                return True
+        else:
+            started = True
+    return False
+
+
+def _add_text(message: dict, text: str, *, prepend: bool = False) -> dict:
+    result = dict(message)
+    content = result.get("content")
+    def join(a, b):
+        return "\n\n".join(part for part in (a, b) if part)
+    if content is None or isinstance(content, str):
+        result["content"] = join(text, content) if prepend else join(content, text)
+    elif isinstance(content, list):
+        parts = deepcopy(content)
+        index = 0 if prepend else -1
+        if parts and parts[index].get("type") == "text":
+            previous = parts[index].get("text") or ""
+            parts[index]["text"] = join(text, previous) if prepend else join(previous, text)
+        else:
+            parts.insert(0 if prepend else len(parts), {"type": "text", "text": text})
+        result["content"] = parts
+    else:
+        raise InlineSystemError("inline system folding requires text or content blocks")
+    return result
+
+
+def _after_tool_results(messages: list[dict]) -> list[dict]:
+    out, held = [], []
+    pending_ids: set[str] = set()
+    for message in messages:
+        role = message.get("role")
+        if role == "system" and pending_ids:
+            held.append(message)
+            continue
+        if role == "tool":
+            out.append(message)
+            pending_ids.discard(message.get("tool_call_id", ""))
+            if not pending_ids:
+                out.extend(held)
+                held.clear()
+            continue
+        if pending_ids and held:
+            raise InlineSystemError("inline system instruction interrupts incomplete tool results")
+        out.append(message)
+        pending_ids = {call.get("id", "") for call in message.get("tool_calls", [])} if role == "assistant" else set()
+    if held:
+        raise InlineSystemError("inline system instruction is waiting for missing tool results")
+    return out
+
+
+def normalize_inline_system(messages: list[dict], mode: str) -> list[dict]:
+    if mode not in ("preserve", "fold"):
+        raise InlineSystemError(f"invalid inline system mode: {mode}")
+    if not has_inline_system(messages):
+        return messages
+    ordered = _after_tool_results(deepcopy(messages))
+    # Without a stable leading system, some renderers attach tools to the first
+    # late system message. Folding avoids moving tools when that message changes.
+    if mode == "preserve" and ordered[0].get("role") == "system":
+        return ordered
+    out, pending = [], []
+    started = False
+    for message in ordered:
+        role = message.get("role")
+        if role == "system" and started:
+            text = message.get("content") or ""
+            if not isinstance(text, str):
+                raise InlineSystemError("inline system instructions must be text")
+            if text:
+                if out and out[-1].get("role") in ("user", "tool") and not pending:
+                    out[-1] = _add_text(out[-1], text)
+                else:
+                    pending.append(text)
+            continue
+        if pending:
+            text = "\n\n".join(pending)
+            pending.clear()
+            if role == "user":
+                message = _add_text(message, text, prepend=True)
+            else:
+                out.append({"role": "user", "content": text})
+        out.append(message)
+        started = started or role != "system"
+    if pending:
+        out.append({"role": "user", "content": "\n\n".join(pending)})
+    return out
+
+
+def _insertion(base: list[int], new: list[int]) -> tuple[int, list[int]] | None:
+    start = 0
+    while start < min(len(base), len(new)) and base[start] == new[start]:
+        start += 1
+    suffix = 0
+    while suffix < min(len(base), len(new)) - start and base[-suffix-1] == new[-suffix-1]:
+        suffix += 1
+    if start + suffix != len(base):
+        return None
+    return start, new[start:len(new)-suffix]
+
+
+def probe_inline_system(render: Callable[[list[dict]], str], tokenizer: Any) -> tuple[str, str]:
+    """Conservatively require a distinct, prefix-stable system segment."""
+    note, reasoning = "Echo inline instruction", "Foxtrot previous reasoning"
+    system = {"role": "system", "content": "Alpha stable instructions"}
+    user = {"role": "user", "content": "Bravo first question"}
+    cases = [
+        ("middle", [system, user, {"role": "assistant", "content": "Charlie first answer"},
+                    {"role": "user", "content": "Delta second question"}], 3),
+        ("trailing", [system, user], 2),
+        ("tool result", [system, user,
+            {"role": "assistant", "content": "", "reasoning_content": reasoning, "thinking": reasoning,
+             "reasoning": reasoning, "tool_calls": [{"id": "probe_call", "type": "function",
+                "function": {"name": "read", "arguments": {}}}]},
+            {"role": "tool", "tool_call_id": "probe_call", "content": "Golf tool output"}], 4),
+    ]
+    max_system_tokens = None
+    try:
+        special_ids = set(tokenizer.all_special_ids) | set(tokenizer.get_added_vocab().values())
+        if not special_ids:
+            return "fold", "renderer has no identifiable special tokens"
+        for name, messages, index in cases:
+            base_text = render(deepcopy(messages))
+            base = tokenizer.encode(base_text, add_special_tokens=False)
+            with_note = messages[:index] + [{"role": "system", "content": note}] + messages[index:]
+            new_text = render(deepcopy(with_note))
+            new = tokenizer.encode(new_text, add_special_tokens=False)
+            found = _insertion(base, new)
+            if found is None:
+                return "fold", f"{name}: other prompt tokens change"
+            position, segment = found
+            if index == len(messages) and position >= len(base):
+                return "fold", f"{name}: instruction follows the generation opening"
+            if tokenizer.decode(segment).count(note) != 1 or not special_ids.intersection(segment):
+                return "fold", f"{name}: no distinct role-marked instruction"
+            if (reasoning in base_text) != (reasoning in new_text):
+                return "fold", f"{name}: reasoning retention changes"
+            if name == "middle":
+                as_user = messages[:index] + [{"role": "user", "content": note}] + messages[index:]
+                user_tokens = tokenizer.encode(render(deepcopy(as_user)), add_special_tokens=False)
+                user_segment = _insertion(base, user_tokens)
+                if user_segment is None:
+                    return "fold", "middle: cannot isolate a user instruction"
+                if user_segment[1] == segment:
+                    return "fold", "system instructions render as ordinary user turns"
+                max_system_tokens = len(user_segment[1]) + _INLINE_ROLE_SLACK_TOKENS
+            if max_system_tokens is not None and len(segment) > max_system_tokens:
+                return "fold", f"{name}: system insertion overhead ({len(segment)} > {max_system_tokens} tokens)"
+    except Exception as exc:
+        return "fold", f"probe could not render: {type(exc).__name__}"
+    return "preserve", "renderer preserves distinct system turns and existing prompt tokens"
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index 758ffd4dcf..f473665267 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -4,6 +4,7 @@
 import json
 import os
 import threading
+from copy import deepcopy
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, List
 
@@ -22,6 +23,7 @@
     probe_thinking_profile,
     quantize_effort,
 )
+from .inline_system import has_inline_system, normalize_inline_system, probe_inline_system
 
 logger = init_logger(__name__)
 
@@ -58,6 +60,8 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor
         self._thinking_profile: ThinkingProfile | None = None
         self._effort_lock = threading.Lock()
         self._logged_effort_maps: set[tuple[Any, str | None]] = set()
+        self._inline_system_lock = threading.Lock()
+        self._inline_system_modes: dict[str, str] = {}
 
     def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
         results: List[UserMsg] = []
@@ -103,9 +107,30 @@ def render_prompt(self, msg: TokenizeMsg) -> str:
         validation, count_tokens) must quantize identically."""
         if not isinstance(msg.text, list):
             return msg.text
-        return self._render(
-            msg.text, msg.tools, self._sanitize_effort(msg.chat_template_kwargs or {})
-        )
+        kwargs = self._sanitize_effort(msg.chat_template_kwargs or {})
+        messages = msg.text
+        if msg.inline_system_policy is not None:
+            if msg.inline_system_policy not in ("auto", "preserve", "fold"):
+                raise ValueError(f"invalid inline system policy: {msg.inline_system_policy}")
+            if has_inline_system(messages):
+                mode = msg.inline_system_policy
+                if mode == "auto":
+                    mode = self.inline_system_mode(msg.tools, kwargs)
+                messages = normalize_inline_system(messages, mode)
+        return self._render(messages, msg.tools, kwargs)
+
+    def inline_system_mode(self, tools: list[dict] | None, kwargs: dict[str, Any]) -> str:
+        key = json.dumps([tools, kwargs], sort_keys=True)
+        with self._inline_system_lock:
+            if key not in self._inline_system_modes:
+                mode, reason = probe_inline_system(
+                    lambda messages: self._render(messages, deepcopy(tools), deepcopy(kwargs)), self.tokenizer
+                )
+                if len(self._inline_system_modes) >= 128:
+                    self._inline_system_modes.pop(next(iter(self._inline_system_modes)))
+                self._inline_system_modes[key] = mode
+                logger.info("Anthropic inline system policy: %s (%s)", mode, reason)
+            return self._inline_system_modes[key]
 
     def _render(
         self,
diff --git a/tests/README.md b/tests/README.md
index 26dba63391..be22c1fb20 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -53,7 +53,10 @@ checkpoint is set:
 | `FREETOKEN_GEMMA4_MODEL`   | `models/test_gemma4_vision.py`, `tokenizer/test_mm_tokenize.py` — local Gemma-4 `gemma4` checkpoint (tower parity, boi/eoi expansion; the tokenizer test also takes the 12B `gemma4_unified` checkpoint) |
 | `FREETOKEN_GEMMA4_UNIFIED_MODEL` | `models/test_gemma4_vision.py` — local gemma-4-12B `gemma4_unified` checkpoint (embedder parity) |
 | `FREETOKEN_GLM53_MODEL`    | `tokenizer/test_mm_tokenize.py` — local GLM-5.3-Flash checkpoint (image-token expansion) |
-| `FREETOKEN_MUSE_MODEL`     | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py` — local Muse-Glimmer-30B checkpoint (tower parity, image_start/image_end expansion) |
+| `FREETOKEN_DSV4_MODEL` | `tokenizer/test_tokenize.py` - local DSV4 checkpoint with its native `encoding/encoding_dsv4.py`; prefix stability and reminder retention, tokenizer only (no GPU or model weights loaded) |
+| `FREETOKEN_INLINE_GLM_MODEL` | `tokenizer/test_tokenize.py` - local GLM-5.3 tokenizer/template; inline-system preserve policy, prefix stability and reasoning retention (no weights loaded) |
+| `FREETOKEN_INLINE_QWEN_MODEL` | `tokenizer/test_tokenize.py` - local Qwen3.6 tokenizer/template; inline-system fold policy, prefix stability and reasoning retention (no weights loaded) |
+| `FREETOKEN_MUSE_MODEL`     | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py`, `tokenizer/test_tokenize.py` - local Muse-Glimmer-30B checkpoint (vision tests; tokenizer-only inline-system regressions) |
 | `FREETOKEN_MINIMAX_M3_MODEL` | `tokenizer/test_mm_tokenize.py` — local MiniMax-M3 checkpoint (start/end expansion, token budget) |
 | `FREETOKEN_TEST_MOE_CACHE_SIZE` | `e2e/test_aime.py` — >0 switches to the offload MoE backend with this cache size |
 | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) |
diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py
index 327bfcd87d..f382093fc0 100644
--- a/tests/server/test_anthropic_api.py
+++ b/tests/server/test_anthropic_api.py
@@ -231,10 +231,7 @@ def test_convert_thinking_replay_in_tool_loop():
     assert spec.messages[2]["role"] == "tool"
 
 
-def test_convert_hoists_and_merges_system_messages():
-    # Claude Code interleaves system messages mid-array; strict chat templates
-    # (e.g. Qwen3.5: "System message must be at the beginning") require ONE system
-    # message at the front. Merge top-level system + in-array system, hoist to front.
+def test_convert_preserves_mid_conversation_system_reminders():
     req = AnthropicMessagesRequest.model_validate(
         {
             "model": "claude-x",
@@ -247,11 +244,69 @@ def test_convert_hoists_and_merges_system_messages():
             ],
         }
     )
+    original = req.model_dump()
     spec = A.convert_anthropic_to_genspec(req, {})
-    assert [m["role"] for m in spec.messages] == ["system", "user", "assistant"]
-    assert sum(1 for m in spec.messages if m["role"] == "system") == 1
-    assert "top-level sys" in spec.messages[0]["content"]
-    assert "mid-stream sys" in spec.messages[0]["content"]
+    assert spec.messages == [
+        {"role": "system", "content": "top-level sys"},
+        {"role": "user", "content": "hello"},
+        {"role": "system", "content": "mid-stream sys"},
+        {"role": "assistant", "content": "hi"},
+    ]
+    assert req.model_dump() == original
+
+
+def test_convert_merges_only_leading_system_messages():
+    req = AnthropicMessagesRequest.model_validate({
+        "model": "claude-x", "max_tokens": 64,
+        "system": [{"type": "text", "text": "top-level sys"}],
+        "messages": [
+            {"role": "system", "content": "leading sys"},
+            {"role": "system", "content": [{"type": "text", "text": "second sys"}]},
+            {"role": "user", "content": "hello"},
+            {"role": "system", "content": [{"type": "text", "text": "turn reminder"}]},
+            {"role": "assistant", "content": "hi"},
+            {"role": "system", "content": ""},
+            {"role": "system", "content": "next reminder"},
+        ],
+    })
+    original = req.model_dump()
+    spec = A.convert_anthropic_to_genspec(req, {})
+    assert spec.messages == [
+        {"role": "system", "content": "top-level sys\n\nleading sys\n\nsecond sys"},
+        {"role": "user", "content": "hello"},
+        {"role": "system", "content": "turn reminder"},
+        {"role": "assistant", "content": "hi"},
+        {"role": "system", "content": "next reminder"},
+    ]
+    assert req.model_dump() == original
+
+
+@pytest.mark.parametrize("system", [None, "stable instructions"])
+@pytest.mark.parametrize("count_tokens", [False, True])
+def test_appending_system_reminder_preserves_converted_prefix(system, count_tokens):
+    from freetoken.server.anthropic_models import AnthropicCountTokensRequest
+
+    data = {
+        "model": "claude-x",
+        "messages": [
+            {"role": "user", "content": "inspect the source"},
+            {"role": "assistant", "content": "I found the function"},
+        ],
+        "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}],
+    }
+    if system is not None:
+        data["system"] = system
+    if not count_tokens:
+        data["max_tokens"] = 64
+    request_type = AnthropicCountTokensRequest if count_tokens else AnthropicMessagesRequest
+    first, first_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data))
+    data["messages"].append({"role": "system", "content": "<total_tokens>1000 tokens left</total_tokens>"})
+    second, second_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data))
+    assert second[:-1] == first
+    assert second[-1] == {
+        "role": "system", "content": "<total_tokens>1000 tokens left</total_tokens>"
+    }
+    assert second_tools == first_tools
 
 
 # --------------------------------------------------------------------------- #
@@ -664,6 +719,7 @@ def test_count_tokens_route():
     assert msg.text == spec.messages
     assert msg.tools == spec.template_tools
     assert msg.chat_template_kwargs == spec.chat_template_kwargs
+    assert msg.inline_system_policy == spec.inline_system_policy == "auto"
 
 
 def test_count_tokens_empty_messages_400():
@@ -673,6 +729,28 @@ def test_count_tokens_empty_messages_400():
     assert r.json()["error"]["type"] == "invalid_request_error"
 
 
+@pytest.mark.parametrize("policy", ["auto", "preserve", "fold"])
+def test_inline_system_override_reaches_counting_and_generation(policy):
+    manager = _FakeTokenizeManager()
+    fake = FakeState([("ok", True, 3, 1)])
+    fake._count_manager = manager
+    fake.config.anthropic_inline_system = policy
+    submitted = []
+
+    async def record(msg):
+        submitted.append(msg)
+
+    fake.send_one = record
+    client = _client(fake)
+    body = {**_COUNT_BODY, "messages": [
+        {"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}]}
+    assert client.post("/v1/messages/count_tokens", json=body).status_code == 200
+    assert client.post("/v1/messages", json={**body, "max_tokens": 32}).status_code == 200
+    assert manager.msgs[0].inline_system_policy == submitted[0].inline_system_policy == policy
+    assert manager.msgs[0].text == submitted[0].text
+    assert submitted[0].text[0]["content"] == body["system"]
+
+
 def test_count_tokens_works_while_not_serving():
     # Counting never touches the engine, so it must not 503 during load/rebuild.
     client, _ = _count_client(_FakeTokenizeManager(), maintenance_state="rebuilding")
@@ -719,6 +797,15 @@ def test_count_tokens_template_render_error_400():
     assert r.json()["error"]["type"] == "invalid_request_error"
 
 
+def test_count_tokens_inline_system_placement_error_is_400():
+    from freetoken.tokenizer.inline_system import InlineSystemError
+
+    client, _ = _count_client(_RaisingManager(InlineSystemError("missing tool results")))
+    r = client.post("/v1/messages/count_tokens", json=_COUNT_BODY)
+    assert r.status_code == 400
+    assert r.json()["error"]["type"] == "invalid_request_error"
+
+
 def test_count_tokens_excluded_from_request_ring():
     # count_tokens never enters generation accounting; its latency (incl. first-touch tokenizer
     # load) must not land in the /v1/requests ring or the /v1/stats p95.
diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py
index e228667154..fdc095921d 100644
--- a/tests/server/test_message_wire.py
+++ b/tests/server/test_message_wire.py
@@ -28,6 +28,18 @@
 from freetoken.core import SamplingParams
 
 
+def test_inline_system_policy_tokenizer_roundtrip():
+    for policy in (None, "auto", "preserve", "fold"):
+        msg = TokenizeMsg(uid=1, text=[{"role": "system", "content": "instruction"}],
+                          sampling_params=SamplingParams(), inline_system_policy=policy)
+        decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg))
+        assert decoded.inline_system_policy == policy
+        assert decoded.text == msg.text
+    legacy = BaseTokenizerMsg.encoder(msg)
+    legacy.pop("inline_system_policy")
+    assert BaseTokenizerMsg.decoder(legacy).inline_system_policy is None
+
+
 def test_cache_rebuild_msg_roundtrip():
     msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle")
     out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg))
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 369576a18b..16ef6d7419 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import json
+import os
+from copy import deepcopy
 
 import pytest
 import torch
@@ -8,14 +10,303 @@
 from freetoken.core import SamplingParams
 from freetoken.message import TokenizeMsg
 from freetoken.tokenizer.tokenize import TokenizeManager, _dsv4_arguments_str
+from freetoken.tokenizer.inline_system import InlineSystemError, normalize_inline_system, probe_inline_system
+
+
+def test_fold_keeps_leading_system_and_appends_text_without_mutation():
+    messages = [
+        {"role": "system", "content": "stable"},
+        {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "image"}},
+                                     {"type": "text", "text": "question"}]},
+        {"role": "system", "content": "instruction 1"},
+        {"role": "system", "content": "instruction 2"},
+    ]
+    original = deepcopy(messages)
+    result = normalize_inline_system(messages, "fold")
+    assert result == [messages[0], {"role": "user", "content": [
+        messages[1]["content"][0], {"type": "text", "text": "question\n\ninstruction 1\n\ninstruction 2"}]}]
+    assert messages == original
+
+
+@pytest.mark.parametrize("mode", ["preserve", "fold"])
+def test_inline_system_waits_for_all_parallel_tool_results(mode):
+    messages = [
+        {"role": "system", "content": "stable"}, {"role": "user", "content": "question"},
+        {"role": "assistant", "reasoning_content": "reason", "tool_calls": [{"id": "a"}, {"id": "b"}]},
+        {"role": "system", "content": "before results"},
+        {"role": "tool", "tool_call_id": "a", "content": "result a"},
+        {"role": "system", "content": "between results"},
+        {"role": "tool", "tool_call_id": "b", "content": "result b"},
+        {"role": "system", "content": "after results"},
+    ]
+    original = deepcopy(messages)
+    result = normalize_inline_system(messages, mode)
+    assert [m["tool_call_id"] for m in result if m["role"] == "tool"] == ["a", "b"]
+    assert result[2] == messages[2]
+    assert [m["role"] for m in result[3:5]] == ["tool", "tool"]
+    if mode == "fold":
+        assert len(result) == 5
+        assert result[-1]["content"] == "result b\n\nbefore results\n\nbetween results\n\nafter results"
+    else:
+        assert [m["content"] for m in result[5:]] == ["before results", "between results", "after results"]
+    assert messages == original
+
+
+@pytest.mark.parametrize("mode", ["preserve", "fold"])
+def test_inline_system_does_not_hide_missing_tool_results(mode):
+    messages = [{"role": "assistant", "tool_calls": [{"id": "missing"}]},
+                {"role": "system", "content": "instruction"}]
+    with pytest.raises(InlineSystemError, match="missing tool results"):
+        normalize_inline_system(messages, mode)
+
+
+def test_fold_after_assistant_prepends_next_user_or_supplies_a_user_turn():
+    messages = [{"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"},
+                {"role": "system", "content": "instruction"}]
+    assert normalize_inline_system(messages, "fold")[-1] == {"role": "user", "content": "instruction"}
+    result = normalize_inline_system(messages + [{"role": "user", "content": "next"}], "fold")
+    assert result == messages[:2] + [{"role": "user", "content": "instruction\n\nnext"}]
+
+
+def test_preserve_keeps_instruction_role_and_only_folds_without_a_stable_system():
+    messages = [{"role": "system", "content": "stable"}, {"role": "user", "content": "question"},
+                {"role": "system", "content": "late"}]
+    assert normalize_inline_system(messages, "preserve") == messages
+    assert normalize_inline_system(messages[1:], "preserve") == [{"role": "user", "content": "question\n\nlate"}]
+
+
+class ProbeTokenizer:
+    vocabulary = {text: 10000+i for i, text in enumerate(
+        ["<system>", "</system>", "<user>", "</user>", "<assistant>", "</assistant>",
+         "<tool>", "</tool>", "<think>", "</think>"])}
+    all_special_ids = list(vocabulary.values())
+
+    def get_added_vocab(self):
+        return self.vocabulary
+
+    def encode(self, text, **kwargs):
+        import re
+        parts = re.split("(" + "|".join(re.escape(x) for x in self.vocabulary) + ")", text)
+        return [token for part in parts for token in ([self.vocabulary[part]] if part in self.vocabulary else map(ord, part))]
+
+    def decode(self, ids):
+        inverse = {v: k for k, v in self.vocabulary.items()}
+        return "".join(inverse[x] if x in inverse else chr(x) for x in ids)
+
+
+@pytest.mark.parametrize("behavior,expected", [("preserve", "preserve"), ("reject", "fold"),
+    ("drop", "fold"), ("bare", "fold"), ("user", "fold"), ("rewrite", "fold"),
+    ("after_generation", "fold"), ("drop_reasoning", "fold")])
+def test_renderer_probe_checks_more_than_template_acceptance(behavior, expected):
+    def render(messages):
+        inline = any(m["role"] == "system" for m in messages[1:])
+        output, late = "", ""
+        for i, message in enumerate(messages):
+            role, content = message["role"], message.get("content", "")
+            if role == "system" and i:
+                if behavior == "reject":
+                    raise ValueError("system first only")
+                if behavior == "drop":
+                    continue
+                if behavior == "bare":
+                    output += content
+                    continue
+                if behavior == "user":
+                    role = "user"
+                if behavior == "after_generation":
+                    late += content
+                    continue
+            if behavior == "rewrite" and inline and i == 0:
+                content = "changed head"
+            reason = message.get("reasoning_content")
+            if reason and not (behavior == "drop_reasoning" and inline):
+                content = "<think>" + reason + "</think>" + content
+            output += f"<{role}>{content}</{role}>"
+        return output + "<assistant>" + late
+    assert probe_inline_system(render, ProbeTokenizer())[0] == expected
+
+
+@pytest.mark.parametrize("position", ["middle", "trailing", "tool result"])
+@pytest.mark.parametrize("metadata", ["Reasoning strength: high. Valid recipients: self, user.", "tool schema " * 1000],
+                         ids=["metadata", "tools"])
+def test_inline_probe_rejects_repeated_system_metadata(position, metadata):
+    def render(messages):
+        output = ""
+        for i, message in enumerate(messages):
+            role = message["role"]
+            content = message.get("content", "")
+            if role == "system" and i:
+                if i < len(messages)-1:
+                    location = "middle"
+                else:
+                    location = "tool result" if messages[i-1]["role"] == "tool" else "trailing"
+                if location == position:
+                    content += metadata
+            content = message.get("reasoning_content", "") + content
+            output += f"<{role}>{content}</{role}>"
+        return output + "<assistant>"
+
+    mode, reason = probe_inline_system(render, ProbeTokenizer())
+    assert mode == "fold"
+    assert "overhead" in reason
+
+
+def test_inline_probe_allows_small_role_framing_overhead():
+    def render(messages):
+        return "".join(f'<{m["role"]}>' + ("role: " if m["role"] == "system" else "")
+                       + m.get("reasoning_content", "") + m.get("content", "")
+                       + f'</{m["role"]}>' for m in messages) + "<assistant>"
+
+    assert probe_inline_system(render, ProbeTokenizer())[0] == "preserve"
+
+
+def test_inline_probe_folds_when_user_insertion_cannot_be_measured():
+    def render(messages):
+        rewrite = any(m["role"] == "user" and m.get("content") == "Echo inline instruction" for m in messages)
+        return ("changed head" if rewrite else "stable head") + "".join(
+            f'<{m["role"]}>' + m.get("reasoning_content", "") + m.get("content", "")
+            + f'</{m["role"]}>' for m in messages) + "<assistant>"
+
+    mode, reason = probe_inline_system(render, ProbeTokenizer())
+    assert mode == "fold"
+    assert "cannot isolate" in reason
+
+
+def test_inline_probe_cache_is_scoped_to_tools_and_render_kwargs(monkeypatch):
+    import freetoken.tokenizer.tokenize as module
+    calls = []
+    monkeypatch.setattr(module, "probe_inline_system", lambda *_: (calls.append(1) or ("fold", "test")))
+    manager = TokenizeManager(FakeTokenizer())
+    for _ in range(3):
+        assert manager.inline_system_mode(None, {}) == "fold"
+    manager.inline_system_mode([], {})
+    manager.inline_system_mode(None, {"enable_thinking": True})
+    assert len(calls) == 3
+
+
+def test_inline_policy_is_opt_in_and_not_forwarded_to_templates():
+    tokenizer = FakeTokenizer()
+    manager = TokenizeManager(tokenizer)
+    messages = [{"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}]
+    msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams())
+    manager.render_prompt(msg)
+    assert tokenizer.chat_template_messages == messages
+    msg.inline_system_policy = "fold"
+    manager.render_prompt(msg)
+    assert tokenizer.chat_template_messages == [{"role": "user", "content": "question\n\ninstruction"}]
+    assert "inline_system_policy" not in tokenizer.chat_template_kwargs
+
+
+@pytest.mark.needs_weights
+@pytest.mark.parametrize("variable,expected_mode", [
+    ("FREETOKEN_DSV4_MODEL", "fold"),
+    ("FREETOKEN_INLINE_GLM_MODEL", "preserve"),
+    ("FREETOKEN_INLINE_QWEN_MODEL", "fold"),
+])
+def test_native_model_inline_system_tool_history(variable, expected_mode):
+    from transformers import AutoTokenizer
+    from freetoken.server.anthropic_api import convert_anthropic_to_genspec
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+
+    path = os.environ.get(variable)
+    if not path:
+        pytest.skip(f"set {variable} to a local checkpoint")
+    manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True))
+    data = {
+        "model": "local", "max_tokens": 64, "system": "Stable project instructions.",
+        "thinking": {"type": "enabled"},
+        "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}],
+        "messages": [
+            {"role": "user", "content": "Inspect source:\n" + "value = 1\n" * 512},
+            {"role": "system", "content": "Keep all existing tests."},
+        ],
+    }
+    ids, prompts = [], []
+    for turn in (1, 2):
+        data["messages"].extend([
+            {"role": "assistant", "content": [
+                {"type": "thinking", "thinking": f"Retained thought {turn}.", "signature": ""},
+                {"type": "tool_use", "id": f"read_{turn}", "name": "read", "input": {}},
+            ]},
+            {"role": "system", "content": f"During tool call {turn}."},
+            {"role": "user", "content": [
+                {"type": "tool_result", "tool_use_id": f"read_{turn}", "content": f"Source result {turn}."},
+            ]},
+            {"role": "system", "content": f"Remaining budget: {1000-turn}."},
+        ])
+        req = AnthropicMessagesRequest.model_validate(data)
+        spec = convert_anthropic_to_genspec(req, {})
+        assert manager.inline_system_mode(spec.template_tools, spec.chat_template_kwargs) == expected_mode
+        original = deepcopy(spec.messages)
+        msg = TokenizeMsg(uid=turn, text=spec.messages, tools=spec.template_tools,
+                          sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs,
+                          inline_system_policy=spec.inline_system_policy)
+        prompts.append(manager.render_prompt(msg))
+        ids.append(manager.tokenize([msg])[0].input_ids.tolist())
+        assert spec.messages == original
+        assert prompts[-1].count("Keep all existing tests.") == 1
+        assert "Retained thought 1." in prompts[-1]
+        for previous in range(1, turn+1):
+            assert prompts[-1].count(f"During tool call {previous}.") == 1
+            assert prompts[-1].count(f"Remaining budget: {1000-previous}.") == 1
+    common = next((i for i, pair in enumerate(zip(*ids)) if pair[0] != pair[1]), min(map(len, ids)))
+    assert common >= len(ids[0]) - 64
+    assert common > 2048
+
+
+@pytest.mark.needs_weights
+@pytest.mark.parametrize("tool_count", [0, 20])
+def test_native_muse_inline_system_does_not_repeat_environment(tool_count):
+    from transformers import AutoTokenizer
+
+    path = os.environ.get("FREETOKEN_MUSE_MODEL")
+    if not path:
+        pytest.skip("set FREETOKEN_MUSE_MODEL to a local checkpoint")
+    manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True))
+    tools = [{"type": "function", "function": {"name": f"read_{i}",
+              "description": "Inspect the repository source. " * 100,
+              "parameters": {"type": "object", "properties": {}}}} for i in range(tool_count)] or None
+    messages = [{"role": "system", "content": "Stable instructions."},
+                {"role": "user", "content": "Inspect this file.\n" + "value = 1\n" * 256}]
+    assert manager.inline_system_mode(tools, {}) == "fold"
+    for turn in range(3):
+        messages.extend([{"role": "assistant", "content": f"Answer {turn}."},
+                         {"role": "user", "content": f"Question {turn}."},
+                         {"role": "system", "content": f"Budget reminder {turn}."}])
+        original = deepcopy(messages)
+        msg = TokenizeMsg(uid=turn, text=messages, tools=tools, sampling_params=SamplingParams(),
+                          inline_system_policy="auto")
+        prompt = manager.render_prompt(msg)
+        assert prompt.count("Reasoning strength:") == 1
+        assert prompt.count("# Valid recipients:") == 1
+        assert prompt.count("// Function schemas") == bool(tools)
+        for previous in range(turn+1):
+            assert prompt.count(f"Budget reminder {previous}.") == 1
+        ids = manager.tokenize([msg])[0].input_ids.tolist()
+        msg.inline_system_policy = "fold"
+        assert ids == manager.tokenize([msg])[0].input_ids.tolist()
+        assert messages == original
+
+    latest = [m for i, m in enumerate(messages) if m["role"] != "system" or i in (0, len(messages)-1)]
+    updated = deepcopy(latest)
+    updated[-1]["content"] = "Budget reminder 900."
+    tokens = []
+    for conversation in (latest, updated):
+        msg = TokenizeMsg(uid=0, text=conversation, tools=tools, sampling_params=SamplingParams(),
+                          inline_system_policy="auto")
+        tokens.append(manager.tokenize([msg])[0].input_ids.tolist())
+    common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens)))
+    assert common >= min(map(len, tokens)) - 16
 
 
 class FakeTokenizer:
     def __init__(self) -> None:
         self.chat_template_kwargs = None
+        self.chat_template_messages = None
 
     def apply_chat_template(self, messages, **kwargs):
         self.chat_template_kwargs = kwargs
+        self.chat_template_messages = deepcopy(messages)
         return "rendered prompt"
 
     def encode(self, prompt, return_tensors=None, add_special_tokens=True):
@@ -74,6 +365,31 @@ def test_tokenize_manager_passes_tools_to_chat_template():
     assert input_ids.tolist() == [1, 2, 3]
 
 
+@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]])
+def test_generic_template_folds_system_messages_without_mutating_input(initial_system):
+    tokenizer = FakeTokenizer()
+    manager = TokenizeManager(tokenizer)
+    messages = initial_system + [
+        {"role": "user", "content": "inspect files"},
+        {"role": "system", "content": "first reminder"},
+        {"role": "assistant", "content": "found the bug"},
+        {"role": "system", "content": "second reminder"},
+    ]
+    original = deepcopy(messages)
+    tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}]
+    original_tools = deepcopy(tools)
+    manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(),
+                                 inline_system_policy="auto")])
+    assert tokenizer.chat_template_messages == initial_system + [
+        {"role": "user", "content": "inspect files\n\nfirst reminder"},
+        {"role": "assistant", "content": "found the bug"},
+        {"role": "user", "content": "second reminder"},
+    ]
+    assert tokenizer.chat_template_kwargs["tools"] == original_tools
+    assert messages == original
+    assert tools == original_tools
+
+
 class FakeDsv4Tokenizer:
     chat_template = None
 
@@ -126,6 +442,166 @@ def encode_messages(messages, thinking_mode, reasoning_effort=None):
     assert input_ids.tolist() == [4, 5, 6]
 
 
+@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]])
+def test_dsv4_fold_keeps_instructions_and_tools_on_stable_prefix(tmp_path, initial_system):
+    encoding_dir = tmp_path / "encoding"
+    encoding_dir.mkdir()
+    (encoding_dir / "encoding_dsv4.py").write_text(
+        "import json\n"
+        "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n"
+        "    return json.dumps(messages)\n"
+    )
+    tokenizer = FakeDsv4Tokenizer(tmp_path)
+    manager = TokenizeManager(tokenizer)
+    tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}]
+    messages = initial_system + [
+        {"role": "user", "content": "inspect files"},
+        {"role": "system", "content": "first reminder"},
+        {"role": "assistant", "content": "found the bug"},
+    ]
+    original, original_tools = deepcopy(messages), deepcopy(tools)
+    manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(),
+                                 inline_system_policy="fold")])
+    first = json.loads(tokenizer.prompt)
+    extended = messages + [{"role": "system", "content": "next reminder"}]
+    original_extended = deepcopy(extended)
+    manager.tokenize([TokenizeMsg(uid=2, text=extended, tools=tools, sampling_params=SamplingParams(),
+                                 inline_system_policy="fold")])
+    second = json.loads(tokenizer.prompt)
+    assert [m["role"] for m in first] == ["system", "user", "assistant"]
+    assert first[0]["content"] == (initial_system[0]["content"] if initial_system else "")
+    assert first[0]["tools"] == original_tools
+    assert first[1]["content"] == "inspect files\n\nfirst reminder"
+    assert second[:-1] == first
+    assert second[-1] == {"role": "user", "content": "next reminder"}
+    assert all("tools" not in m for m in second[1:])
+    assert messages == original and extended == original_extended and tools == original_tools
+
+
+@pytest.fixture(scope="module")
+def native_dsv4_manager():
+    model_path = os.environ.get("FREETOKEN_DSV4_MODEL")
+    if not model_path:
+        pytest.skip("set FREETOKEN_DSV4_MODEL to a local DSV4 checkpoint")
+    from transformers import AutoTokenizer
+
+    manager = TokenizeManager(AutoTokenizer.from_pretrained(model_path, local_files_only=True))
+    assert manager._dsv4_encoder is not None, "checkpoint must supply the native DSV4 encoder"
+    return manager
+
+
+def test_dsv4_no_tools_retains_reminders_when_encoder_drops_developer_turns(tmp_path):
+    encoding_dir = tmp_path / "encoding"
+    encoding_dir.mkdir()
+    (encoding_dir / "encoding_dsv4.py").write_text(
+        "import json\n"
+        "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n"
+        "    if thinking_mode == 'thinking':\n"
+        "        messages = [m for m in messages if m['role'] != 'developer']\n"
+        "    return json.dumps(messages)\n"
+    )
+    manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path))
+    messages = [
+        {"role": "user", "content": "First task."},
+        {"role": "system", "content": "Keep this earlier instruction."},
+        {"role": "assistant", "content": "First answer."},
+        {"role": "user", "content": "Second task."},
+    ]
+    original = deepcopy(messages)
+    msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams(),
+                      chat_template_kwargs={"enable_thinking": True}, inline_system_policy="fold")
+    rendered = json.loads(manager.render_prompt(msg))
+    assert all(m["role"] in ("system", "user", "assistant") for m in rendered)
+    assert sum(m["content"].count("Keep this earlier instruction.") for m in rendered) == 1
+    assert messages == original
+
+
+@pytest.mark.needs_weights
+@pytest.mark.parametrize("system", [None, "You are a coding assistant."])
+@pytest.mark.parametrize("change", ["update", "append"])
+def test_native_dsv4_reminder_update_preserves_token_prefix(native_dsv4_manager, system, change):
+    from freetoken.server.anthropic_api import convert_anthropic_to_genspec
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+
+    data = {
+        "model": "dsv4", "max_tokens": 64,
+        "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}],
+        "messages": [
+            {"role": "user", "content": "Inspect this file:\n" + "value = 1\n" * 512},
+            {"role": "system", "content": "Keep the changes minimal."},
+            {"role": "assistant", "content": [
+                {"type": "thinking", "thinking": "I should read the source.", "signature": ""},
+                {"type": "tool_use", "id": "call_read", "name": "read", "input": {}},
+            ]},
+            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_read", "content": "value = 1"}]},
+            {"role": "system", "content": "Remaining token budget: 1000"},
+        ],
+    }
+    if system is not None:
+        data["system"] = system
+    tokens = []
+    for budget in (1000, 900):
+        if budget == 900 and change == "append":
+            data["messages"].extend([
+                {"role": "assistant", "content": [
+                    {"type": "thinking", "thinking": "The file is simple.", "signature": ""},
+                    {"type": "text", "text": "I inspected the file."},
+                ]},
+                {"role": "system", "content": "Remaining token budget: 900"},
+            ])
+        else:
+            data["messages"][-1]["content"] = f"Remaining token budget: {budget}"
+        req = AnthropicMessagesRequest.model_validate(data)
+        original = req.model_dump()
+        spec = convert_anthropic_to_genspec(req, {})
+        msg = TokenizeMsg(uid=budget, text=spec.messages, tools=spec.template_tools,
+                          sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs,
+                          inline_system_policy=spec.inline_system_policy)
+        prompt = native_dsv4_manager.render_prompt(msg)
+        assert prompt.count("Keep the changes minimal.") == 1
+        assert prompt.index("Inspect this file:") < prompt.index(f"Remaining token budget: {budget}")
+        tokens.append(native_dsv4_manager.tokenize([msg])[0].input_ids.tolist())
+        assert req.model_dump() == original
+    common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens)))
+    assert common > 2048
+    assert common >= min(map(len, tokens)) - 64
+
+
+@pytest.mark.needs_weights
+@pytest.mark.parametrize("disable_declared_tools", [False, True])
+@pytest.mark.parametrize("thinking_enabled", [False, True])
+def test_native_dsv4_no_tools_keeps_earlier_reminder(
+    native_dsv4_manager, disable_declared_tools, thinking_enabled
+):
+    from freetoken.server.anthropic_api import convert_anthropic_to_genspec
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+
+    data = {
+        "model": "dsv4", "max_tokens": 64, "system": "Stable instructions.",
+        "thinking": {"type": "enabled" if thinking_enabled else "disabled"},
+        "messages": [
+            {"role": "user", "content": "First task."},
+            {"role": "system", "content": "Keep this earlier instruction."},
+            {"role": "assistant", "content": "First answer."},
+            {"role": "user", "content": "Second task."},
+        ],
+    }
+    if disable_declared_tools:
+        data["tools"] = [{"name": "read", "input_schema": {"type": "object", "properties": {}}}]
+        data["tool_choice"] = {"type": "none"}
+    req = AnthropicMessagesRequest.model_validate(data)
+    original = req.model_dump()
+    spec = convert_anthropic_to_genspec(req, {})
+    assert not spec.template_tools
+    msg = TokenizeMsg(uid=1, text=spec.messages, tools=spec.template_tools,
+                      sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs,
+                      inline_system_policy=spec.inline_system_policy)
+    prompt = native_dsv4_manager.render_prompt(msg)
+    assert prompt.count("Keep this earlier instruction.") == 1
+    assert "Second task." in prompt
+    assert req.model_dump() == original
+
+
 def test_dsv4_encoder_gets_tool_call_arguments_as_json_string(tmp_path):
     """Regression: render_messages hands the template dict arguments; the dsv4
     encoder contract is a JSON-object STRING -- a dict trips its fallback that

From 555efd89447232a555e05d187256b76a51c1aaa0 Mon Sep 17 00:00:00 2001
From: cherry77-cloud <1615405@qq.com>
Date: Tue, 6 Oct 2026 06:29:44 +0800
Subject: [PATCH 066/249] fix(kvcache): clone inserted radix keys to release
 request buffers (#559)

---
 python/freetoken/kvcache/hybrid_radix_cache.py | 5 ++++-
 python/freetoken/kvcache/radix_cache.py        | 5 ++++-
 python/freetoken/kvcache/swa_radix_cache.py    | 3 ++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py
index 1bd9e0b9cd..0b66253b17 100644
--- a/python/freetoken/kvcache/hybrid_radix_cache.py
+++ b/python/freetoken/kvcache/hybrid_radix_cache.py
@@ -98,7 +98,10 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor,
         node, prefix_len = self._walk(input_ids)
         if prefix_len != insert_len:
             new_node = RadixTreeNode(self.key_fn)
-            new_node.set_key_value(input_ids[prefix_len:], kv_indices[prefix_len:].clone())
+            # A view would keep the request's full output-budget token buffer alive.
+            new_node.set_key_value(
+                input_ids[prefix_len:].clone(), kv_indices[prefix_len:].clone()
+            )
             new_node.set_parent(node)
             self.full_evictable += new_node.length
             node = new_node
diff --git a/python/freetoken/kvcache/radix_cache.py b/python/freetoken/kvcache/radix_cache.py
index 5437286d12..ec0bf75d35 100644
--- a/python/freetoken/kvcache/radix_cache.py
+++ b/python/freetoken/kvcache/radix_cache.py
@@ -167,7 +167,10 @@ def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> Inser
         node, prefix_len = self._tree_walk(input_ids)
         if prefix_len != insert_len:  # NOTE: prefix_len < insert_len
             new_node = RadixTreeNode(self.key_fn)
-            new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone())
+            # A view would keep the request's full output-budget token buffer alive.
+            new_node.set_key_value(
+                input_ids[prefix_len:].clone(), indices[prefix_len:].clone()
+            )
             new_node.set_parent(node)
             self.evictable_size += new_node.length
             node = new_node
diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py
index 260439ad66..d2cc89c26f 100644
--- a/python/freetoken/kvcache/swa_radix_cache.py
+++ b/python/freetoken/kvcache/swa_radix_cache.py
@@ -215,7 +215,8 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor,
         # then a live (non-tombstone) leaf for the in-window remainder (never a tombstone leaf --
         # the clamp + the free_swa -page_size margin guarantee a live tail).
         if total < insert_len:
-            suffix_ids = input_ids[total:]
+            # Both suffix nodes can share this copy without retaining the request buffer.
+            suffix_ids = input_ids[total:].clone()
             suffix_kv = kv_indices[total:].clone()
             boundary = max(0, min(swa_evicted_seqlen, insert_len) - total)
             boundary = min(boundary, max(0, len(suffix_ids) - self.page_size))

From a8498c7e96e25fad7ed5e61968188fdbb66b35c6 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Tue, 6 Oct 2026 13:27:13 +0300
Subject: [PATCH 067/249] fix(server): leak no semaphores when a signal stops
 the server

---
 python/freetoken/server/launch.py     |  8 ++++++
 python/freetoken/server/supervisor.py |  7 +++++
 tests/server/test_launch.py           | 37 +++++++++++++++++++++++++++
 tests/server/test_supervisor.py       | 31 ++++++++++++++++++++++
 4 files changed, 83 insertions(+)
 create mode 100644 tests/server/test_launch.py

diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py
index bb0c256f46..5179e8e238 100644
--- a/python/freetoken/server/launch.py
+++ b/python/freetoken/server/launch.py
@@ -59,6 +59,14 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None:
     if args.shell_mode:
         _detach_process_group()
 
+    import threading
+
+    from tqdm import tqdm
+
+    # Only rank 0 draws a bar, so no bar lock spans processes; tqdm's default one is a multiprocessing
+    # semaphore, which this worker leaves to resource_tracker as leaked when a signal ends it.
+    tqdm.set_lock(threading.RLock())
+
     # published (not bound) here: the engine binds it after the allocator setup
     from freetoken.gpu_select import set_assigned_gpu
 
diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py
index 322419a65f..20deb7d132 100644
--- a/python/freetoken/server/supervisor.py
+++ b/python/freetoken/server/supervisor.py
@@ -55,6 +55,12 @@ class BackendHandle:
     processes: List[Any] = field(default_factory=list)
     expected_acks: int = 0
 
+    def release_acks(self) -> None:
+        """Drop the ack queue once readiness has read it for the last time. A multiprocessing queue's
+        semaphores are unlinked only when it is collected, and a main process that a signal ends never
+        collects it, so resource_tracker reports them as leaked."""
+        self.ack_queue = None
+
 
 class WorkerDied(Exception):
     def __init__(self, message: str) -> None:
@@ -172,6 +178,7 @@ def _shutting_down() -> bool:
             on_failure(exc.message)
         return
 
+    handle.release_acks()
     on_ready()
 
     while True:
diff --git a/tests/server/test_launch.py b/tests/server/test_launch.py
new file mode 100644
index 0000000000..defba4f307
--- /dev/null
+++ b/tests/server/test_launch.py
@@ -0,0 +1,37 @@
+"""The scheduler worker's own process setup, before it builds the engine."""
+
+from __future__ import annotations
+
+import queue
+import threading
+
+import pytest
+import torch
+from tqdm import tqdm
+
+import freetoken.gpu_select
+import freetoken.scheduler
+from freetoken.distributed import DistributedInfo
+from freetoken.server import launch
+from freetoken.server.args import ServerArgs
+from freetoken.utils import progress
+
+
+def test_the_scheduler_worker_draws_its_bars_under_a_lock_of_its_own(monkeypatch):
+    """Only rank 0 draws a bar, so no bar lock spans processes. tqdm's default one is a multiprocessing
+    semaphore, which a worker ended by a signal leaves to resource_tracker as leaked."""
+
+    class EngineBuilt(Exception):
+        pass
+
+    def scheduler(args):
+        raise EngineBuilt(tqdm.get_lock())
+
+    monkeypatch.setattr(freetoken.scheduler, "Scheduler", scheduler)
+    monkeypatch.setattr(freetoken.gpu_select, "set_assigned_gpu", lambda target: None)
+    monkeypatch.setattr(progress, "_PROGRESS_SINK", None)
+    monkeypatch.setattr(tqdm, "_lock", tqdm.get_lock())
+    args = ServerArgs(model_path="unused", tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16)
+    with pytest.raises(EngineBuilt) as built:
+        launch._run_scheduler(args, queue.Queue())
+    assert type(built.value.args[0]) is type(threading.RLock())
diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py
index 3568d90505..1ee1ca17bf 100644
--- a/tests/server/test_supervisor.py
+++ b/tests/server/test_supervisor.py
@@ -211,6 +211,37 @@ def on_ready() -> None:
     assert "failure" not in seen  # graceful stop: the death was not reported
 
 
+def test_supervisor_releases_the_ack_queue_once_the_backend_is_ready():
+    """Readiness is the ack queue's last read. A multiprocessing queue's semaphores are unlinked only when
+    it is collected, and a main process ended by a signal never collects it."""
+    import multiprocessing
+    import weakref
+
+    from freetoken.server.supervisor import run_backend_supervisor
+
+    class Proc:
+        name = "freetoken-TP0-scheduler"
+        alive = True
+
+        def is_alive(self) -> bool:
+            return self.alive
+
+    proc = Proc()
+    ack_queue = multiprocessing.get_context("spawn").Queue()
+    ack_queue.put("scheduler ready")
+    queue_at_ready = weakref.ref(ack_queue)
+    handle = BackendHandle(ack_queue=ack_queue, processes=[proc], expected_acks=1)
+    del ack_queue
+    seen: dict = {}
+
+    def on_ready() -> None:
+        seen["queue at ready"] = queue_at_ready()
+        proc.alive = False
+
+    run_backend_supervisor(handle, LoadProgress(), on_ready=on_ready, poll=0.01, is_shutting_down=lambda: True)
+    assert seen == {"queue at ready": None}
+
+
 def test_supervisor_silent_on_startup_death_during_shutdown():
     """A worker dying mid-load while an orderly stop is already in progress must not be
     reported as a load failure either."""

From c1a504a50b6ced2c28c5628d29df9885f0f3ae74 Mon Sep 17 00:00:00 2001
From: SuperHotDogCat <115141367+SuperHotDogCat@users.noreply.github.com>
Date: Wed, 7 Oct 2026 11:54:38 +0900
Subject: [PATCH 068/249] perf(moe): Improved `fast_index_copy_multi_jit` for
 NVLink-C2C hosts (#610)

* perf(moe): default fused gather to 32 blocks/bank on NVLink-C2C hosts

* build(aot): prebuild fast_index_copy_multi (1024, 32) variant

* perf(moe): make fused H2D gather grid configurable via env var

* docs: document FREETOKEN_H2D_BLOCKS_PER_BANK for NVLink-C2C hosts

* fix(moe): auto-detect H2D gather grid per device in the copy plan

* chore(docs): revised docs for NVLinkC2C
---
 docs/cli.md                                |  3 +++
 docs/models.md                             | 14 ++++++++++++++
 python/freetoken/kernel/aot.py             |  1 +
 python/freetoken/kernel/fast_index_copy.py | 17 +++++++++++++++++
 python/freetoken/moe/offload_cache.py      |  4 ++++
 5 files changed, 39 insertions(+)

diff --git a/docs/cli.md b/docs/cli.md
index b3cec99826..b9be878ce7 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -204,3 +204,6 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then
 - What to measure: `--dtype`, `--model`, `--formats`, `--isa`.
 - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe
   by that factor.
+
+On NVLink-C2C hosts, see [models.md](models.md#nvlink-c2c-hosts-gh200--gb200)
+for `FREETOKEN_H2D_BLOCKS_PER_BANK` if `PCIe-gather` sits well below the linear H2D ceiling.
diff --git a/docs/models.md b/docs/models.md
index efd5268b37..2cd3baa317 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -59,6 +59,20 @@ have no resident path and stay on `offload`. Offload-only flags
 `--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a
 warning; pass `--moe-strategy offload` to use them.
 
+### NVLink-C2C hosts (GH200 / GB200)
+
+The host link is ~450 GB/s per direction instead of PCIe's ~32-64 GB/s. The
+offload expert gather reads host memory zero-copy and is latency-bound, so it
+needs a wider grid to keep enough loads in flight. Set
+`FREETOKEN_H2D_BLOCKS_PER_BANK`:
+```
+    FREETOKEN_H2D_BLOCKS_PER_BANK=32 ft serve --model <model> --moe-strategy offload
+```
+On GH200 this raises the gather from ~220 to ~410 GB/s (`PCIe-gather` in
+`ft bench bw`), matching the DMA ceiling. Run `ft bench bw` with the variable
+set so the hybrid split is calibrated against the faster gather. Leave it
+unset on PCIe GPUs: wider grids add no bandwidth there.
+
 ## Notes
 
 - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into
diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py
index 5e87f8c923..bbc2d2c669 100644
--- a/python/freetoken/kernel/aot.py
+++ b/python/freetoken/kernel/aot.py
@@ -153,6 +153,7 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]:
         for feature_size in DEFAULT_FAST_INDEX_COPY_FEATURE_SIZES
     )
     specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=8))
+    specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32))
     # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding.
     specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64))
     specs.append(_batch_memcpy_spec())
diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py
index 1aaa1303d2..b2030620ee 100644
--- a/python/freetoken/kernel/fast_index_copy.py
+++ b/python/freetoken/kernel/fast_index_copy.py
@@ -149,6 +149,23 @@ def _jit_fast_index_copy_multi_module(*, num_threads: int, blocks_per_bank: int)
         cuda_wrappers=[("launch", f"&MultiIndexCopyKernel<{args}>::run")],
     )
 
+def default_blocks_per_bank(device: torch.device) -> int:
+    """H2D gather grid width per bank. The link's bandwidth-delay product sets how many
+    16B loads must be in flight: PCIe (~31 GB/s) saturates at ~4K threads/bank, NVLink-C2C
+    (GH200, ~450 GB/s) needs ~32K (measured: 8x1024 -> 222, 32x1024 -> 412 GB/s)."""
+    env = os.environ.get("FREETOKEN_H2D_BLOCKS_PER_BANK")
+    if env is not None:
+        blocks_per_bank = int(
+            env
+        )
+        assert blocks_per_bank > 0, (
+            "FREETOKEN_H2D_BLOCKS_PER_BANK must be greater than 0, "
+            f"got {blocks_per_bank}"
+        )
+        return blocks_per_bank  # manual override (validate > 0)
+    if any(k in torch.cuda.get_device_name(device) for k in ("GH200", "GB200", "GB300")):
+        return 32
+    return 8
 
 def fast_index_copy_multi_jit(
     dst_ptrs: torch.Tensor,
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index c6f975df7e..0969f73829 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -25,6 +25,7 @@
 # entry the batch sees is >= this size.
 _SMALL_BANK_FEAT_BYTES = 256 * 1024
 
+from freetoken.kernel.fast_index_copy import default_blocks_per_bank
 from freetoken.utils import init_logger
 
 logger = init_logger(__name__)
@@ -257,6 +258,7 @@ def __post_init__(self) -> None:
         # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device
         # tensor); dst/feat are layer-invariant.
         self._copy_fused_ok = False
+        self._copy_blocks_per_bank = 8
         self._copy_dst_ptrs: torch.Tensor | None = None
         self._copy_src_ptrs: list[torch.Tensor] | None = None
         self._copy_feat_bytes: torch.Tensor | None = None
@@ -438,6 +440,7 @@ def _build_fused_copy_plan(self) -> None:
         elif self._gather_bank_ids:
             self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous()
             self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous()
+        self._copy_blocks_per_bank = default_blocks_per_bank(self.device)
         self._copy_fused_ok = True
 
     def validate_rebuild(self, cache_size: int) -> None:
@@ -1038,6 +1041,7 @@ def copy_missing(self) -> None:
                 self.evict_slots,
                 self.src_indices,
                 self.num_indices,
+                blocks_per_bank=self._copy_blocks_per_bank,
             )
             return
 

From 1c183e74ae39b6f2708394056c00dda556f89bea Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 6 Oct 2026 21:38:10 -0700
Subject: [PATCH 069/249] perf(moe): spin-wait the cpu moe pool between layers
 (#619)

Co-authored-by: shubham3-ucb <229030647+shubham3-ucb@users.noreply.github.com>
---
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       | 71 +++++++++++++++----
 python/freetoken/moe/cpu_executor.py          | 18 ++++-
 2 files changed, 75 insertions(+), 14 deletions(-)

diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index e2cb0d9c33..fb6a33f497 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -49,6 +49,14 @@
 
 namespace {
 
+inline void cpu_relax() {
+#if CPU_MOE_X86
+  _mm_pause();
+#elif defined(__aarch64__)
+  asm volatile("yield" ::: "memory");
+#endif
+}
+
 using bf16_t = uint16_t;
 
 inline float bf16_to_f32(bf16_t v) {
@@ -1290,13 +1298,19 @@ struct CpuMoeExecutor {
   std::mutex sync_mtx;
   std::condition_variable sync_cv;
 
-  bool stop = false;
+  std::atomic<bool> stop{false};  // spinning workers read it without the lock
   uint64_t cur_gen = 0;
   MoeTask* cur_task = nullptr;
-  std::atomic<uint64_t> submitted{0};
-  std::atomic<uint64_t> completed{0};
+  // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next.
+  alignas(64) std::atomic<uint64_t> submitted{0};
+  alignas(64) std::atomic<uint64_t> completed{0};
+
+  // Spin-then-park waits. Off unless the Python wrapper finds spare CPUs for the spinners.
+  static constexpr auto kSpinWindow = std::chrono::milliseconds(50);
+  std::atomic<bool> spin_wait{false};
+  bool coord_pinned = false;
 
-  std::atomic<int64_t> p1_next{0};
+  alignas(64) std::atomic<int64_t> p1_next{0};
   std::atomic<int64_t> p2_next{0};
   std::atomic<int64_t> prt_next{0};  // ds_fp4 intermediate fp8 round-trip phase
   int64_t p1_total = 0, p2_total = 0, prt_total = 0;
@@ -1866,19 +1880,34 @@ struct CpuMoeExecutor {
     }
   }
 
+  // Tasks arrive once per MoE layer, sub-ms apart. Spinning skips a futex wakeup per worker per layer; parking after kSpinWindow keeps an idle server off the CPU.
   void worker_loop(int tid) {
     pin_self(tid);
     uint64_t my_gen = 0;
+    using wclk = std::chrono::steady_clock;
+    auto last_active = wclk::now();
     for (;;) {
-      MoeTask* t;
-      {
+      uint64_t gen = submitted.load(std::memory_order_acquire);
+      if (gen == my_gen && spin_wait.load(std::memory_order_relaxed)) {
+        unsigned polls = 0;
+        while (gen == my_gen) {
+          cpu_relax();
+          if (stop.load(std::memory_order_relaxed)) return;
+          if ((++polls & 1023u) == 0 && wclk::now() - last_active > kSpinWindow) break;
+          gen = submitted.load(std::memory_order_acquire);
+        }
+      }
+      if (gen == my_gen) {
         std::unique_lock<std::mutex> lk(task_mtx);
-        task_cv.wait(lk, [&] { return stop || cur_gen != my_gen; });
-        if (stop) return;
-        my_gen = cur_gen;
-        t = cur_task;
+        task_cv.wait(lk, [&] { return stop.load() || cur_gen != my_gen; });
+        if (stop.load()) return;
+        gen = cur_gen;
       }
+      // submit() writes cur_task before the release store to `submitted`.
+      my_gen = gen;
+      MoeTask* t = cur_task;
       run_task_body(t);
+      last_active = wclk::now();
       if (done_count.fetch_add(1) + 1 == num_threads) {
         completed.store(my_gen, std::memory_order_release);
         {
@@ -1961,8 +1990,19 @@ struct CpuMoeExecutor {
     task_cv.notify_all();
   }
 
-  void sync() {
+  // Only a coordinator on its own core may spin: host-func callbacks and eager calls share cores with the workers.
+  void sync(bool spin = false) {
     const uint64_t target = submitted.load(std::memory_order_acquire);
+    if (spin) {
+      using sclk = std::chrono::steady_clock;
+      const auto t0 = sclk::now();
+      unsigned polls = 0;
+      while (completed.load(std::memory_order_acquire) < target) {
+        cpu_relax();
+        if ((++polls & 1023u) == 0 && sclk::now() - t0 > kSpinWindow) break;
+      }
+      if (completed.load(std::memory_order_acquire) >= target) return;
+    }
     std::unique_lock<std::mutex> lk(sync_mtx);
     sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; });
   }
@@ -2003,6 +2043,7 @@ struct CpuMoeExecutor {
     }
     flag_served.assign(num_slots, 0);
     coord_stop.store(false);
+    coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0;
     coord_thread = std::thread([this, pin_core] {
 #if CPU_MOE_HAS_AFFINITY
       if (pin_core >= 0) {
@@ -2050,7 +2091,7 @@ struct CpuMoeExecutor {
           }
           if (t != nullptr) {
             submit(t);
-            sync();
+            sync(coord_pinned && spin_wait.load(std::memory_order_relaxed));
           }
           // Release: the workers' y stores are visible before the GPU sees done.
           flag_store_release(&done_flags[L], 1);
@@ -2144,7 +2185,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
       .def("set_input_prequant",
            [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; },
            py::arg("value"))
-      .def("isa_name", &CpuMoeExecutor::isa_name);
+      .def("isa_name", &CpuMoeExecutor::isa_name)
+      .def("set_spin_wait",
+           [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); },
+           py::arg("enabled"))
+      .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); });
   m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr"));
   m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"),
         py::arg("ready_addr"), py::arg("slot"));
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index e2e48d1461..a46806740d 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -53,6 +53,8 @@
 # sizes plus any eager padded sizes); more than that is unheard of, and the overflow
 # just keeps the host-func path for the extra combos.
 _FLAG_SLOTS_PER_LAYER = 16
+# Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out.
+_SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0"
 
 # Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the
 # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue,
@@ -116,6 +118,17 @@ def physical_core_cpus() -> list[int]:
     return reps or allowed or [0]
 
 
+def _spin_wait_ok(core_ids: list[int], coord_core: int) -> bool:
+    """Spin only if each pool thread owns a distinct CPU and two CPUs stay free for the main and CUDA callback threads."""
+    if not _SPIN_WAIT or not hasattr(os, "sched_getaffinity"):
+        return False
+    pool = list(core_ids) + ([coord_core] if coord_core >= 0 else [])
+    if len(set(pool)) != len(pool):
+        return False
+    spare = set(os.sched_getaffinity(0)) - set(pool)
+    return len(spare) >= 2
+
+
 def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]:
     """Return (num_threads, core_ids) for the worker pool.
 
@@ -246,6 +259,8 @@ def __init__(
         self.num_threads = nthreads
         self.core_ids = core_ids
         self.isa = self._ext.isa_name()
+        self.spin_wait = _spin_wait_ok(core_ids, coord_core)
+        self._ext.set_spin_wait(self.spin_wait)
 
         spare = len(physical_core_cpus()) - nthreads - (1 if coord_core >= 0 else 0) - 1
         clamp = max(1, min(torch.get_num_threads(), spare))
@@ -319,7 +334,8 @@ def __init__(
             f"CPU MoE executor ready: threads={nthreads} (pinned to cores "
             f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} "
             f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} "
-            f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens}"
+            f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} "
+            f"spin_wait={self.spin_wait}"
         )
 
     def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor:

From 127c85665b49e39d78a5d09b95a92169b5794eb0 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 7 Oct 2026 00:15:11 -0700
Subject: [PATCH 070/249] feat(deepseek-v41): support text and vision serving
 (#603)

* test(deepseek-v41): add independent text and vision reference

* feat(quant): support block-32 and unrounded inference outputs

* feat(attention): add csa2 packed cache and sparse backend

* feat(deepseek-v41): serve text and images through shared runtime

* perf(mhc): tune batch-stable reduction tiles

* perf(moe): autotune fp4 decode launch geometry

* perf(moe): admit recurring misses against actual lru victims

* revert: perf(moe): admit recurring misses against actual lru victims

This reverts commit 1652d200f0ebb825379cf44c28baf52971e3a461.

* perf(mhc): derive the reduction tile from the width, not a timing

* perf(moe): replace the ds_fp4 decode autotune with a fixed two-warp launch

* fix(deepseek-v41): fence engram staging against in-flight lookups

* fix(moe): return the cpu routed sum in fp32 and round it in python

* fix(deepseek-v41): keep the vision tower and aligner unquantized

* fix(scheduler): admit bounded-replay hits by the history their resume reads

* refactor(deepseek-v41): name the attention family after the model, not csa2

* refactor(kvcache): import the dsv41 pool only for dsv41 models

* refactor(deepseek-v41): name the replay switch --swa-decoder-replay

* fix(server): keep reasoning_effort a string on the wire

* refactor(models): keep the prefill replay hooks off BaseLLMModel

* refactor(deepseek-v41): drop the fp32 routed sum and batch-invariant prefill

* fix(moe): weight ds_fp4 routes in the down epilogue again

* docs(models): anchor the v4.1 replay paragraph on the tech report

* refactor(deepseek-v41): run the decoder source's narrowed prefill on Block

* refactor(kvcache): let the dsv41 pool own the history a resume keeps live

* refactor(deepseek-v41): resume bounded-replay hits by capping the prefix match

* refactor(deepseek-v41): describe the attention group as a dsv4 variant and derive its geometry from the args

* refactor(deepseek-v41): keep the engram table locator in the model package

* refactor(deepseek-v41): drop the fp32 lm head output

- return the head's bf16 logits, as DeepSeek-V4 does
- revert ParallelLMHead(out_dtype=) and LinearKernel.apply_out_dtype, which only this call used
- remove test_head_preserves_fp32_logit_margin
- cast test_vision's prefill logits to fp32 before comparing, as the harness does

* refactor(quant): carry the w4a8 activation block on MoEConfig

- drop ActDesc and QuantScheme.act; the scheme only describes stored tensors
- Fp8BlockConfig sets MoEConfig.act_block from weight_block_size for its fp4 experts via a layer_config hook
- the ds_fp4 GPU kernel and CPU executor read the block from the layer config and require it

* refactor(deepseek-v41): size the kv pool from the args, not a geometry copy

- delete kvcache/dsv41_geometry.py; the pool and cost model read dsv41_args by attribute, as the DSV4 pool reads dsv4_args
- source / ratio aliasing reads args.roles; the duplicated source and ratio validation is gone
- the cost model owns the tier row formats, row bytes and the private window layers
- the pool keeps the format and window/ratio checks, and exposes the formats and is_private_window to the attention backend

* refactor(kvcache): move the dsv4 pool into kvcache/dsv4

* refactor(kvcache): move the dsv41 pool into kvcache/dsv4

* chore(benchmarks): drop the dsv41 selection and fp8 linear benchmarks

* test(deepseek-v41): drop tests that repeat other coverage

* test(deepseek-v41): load the tiny model's resident experts as banks

* fix(server): parse deepseek-v4.1's space-led dsml tool calls

---------

Co-authored-by: andy-yang-1 <andy_yang@berkeley.edu>
---
 docs/models.md                                |   32 +
 python/freetoken/attention/__init__.py        |   10 +
 python/freetoken/attention/base.py            |    3 +
 python/freetoken/attention/dsv41_sparse.py    |  438 +++++
 python/freetoken/core.py                      |    3 +
 python/freetoken/engine/cache_budget.py       |    2 +-
 python/freetoken/engine/config.py             |    3 +
 python/freetoken/engine/engine.py             |   44 +-
 python/freetoken/kernel/aot_models.py         |   11 +
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       |   20 +-
 .../kernel/triton/dsv4/fp8_linear.py          |  321 ++--
 .../freetoken/kernel/triton/dsv4/fused_moe.py |   35 +-
 python/freetoken/kernel/triton/dsv4/router.py |   83 +
 .../freetoken/kernel/triton/dsv41/__init__.py |    9 +
 .../freetoken/kernel/triton/dsv41/indexer.py  |  176 ++
 python/freetoken/kernel/triton/dsv41/pack.py  |  183 +++
 .../kernel/triton/dsv41/row_format.py         |   92 ++
 .../kernel/triton/dsv41/sparse_attn.py        |  323 ++++
 python/freetoken/kernel/triton/dsv41/topk.py  |  265 +++
 python/freetoken/kernel/triton/e4m3_compat.py |   14 +
 python/freetoken/kernel/triton/mhc.py         |  185 ++-
 python/freetoken/kvcache/__init__.py          |   27 +-
 python/freetoken/kvcache/base.py              |   17 +
 python/freetoken/kvcache/cache_status.py      |   44 +-
 python/freetoken/kvcache/dsa_pool.py          |    4 +-
 python/freetoken/kvcache/dsv4/__init__.py     |    1 +
 .../freetoken/kvcache/dsv4/v41_cost_model.py  |  216 +++
 python/freetoken/kvcache/dsv4/v41_pool.py     |  292 ++++
 .../freetoken/kvcache/dsv4/v41_row_format.py  |   69 +
 .../v4_cost_model.py}                         |   14 +-
 .../{dsv4_paged_pool.py => dsv4/v4_pool.py}   |  313 +---
 python/freetoken/kvcache/dsv4/window_tier.py  |  331 ++++
 python/freetoken/kvcache/hybrid_swa_pool.py   |   12 +-
 python/freetoken/layers/mhc.py                |  119 +-
 python/freetoken/layers/moe.py                |   23 +-
 .../layers/quantization/configs/base.py       |    6 +-
 .../layers/quantization/configs/fp8.py        |   34 +-
 .../layers/quantization/linear/fp8_block.py   |   24 +-
 .../freetoken/layers/quantization/moe/base.py |    6 +
 .../layers/quantization/moe/fp8_block.py      |    4 +-
 .../layers/quantization/moe/mxfp4.py          |   11 +-
 .../freetoken/layers/quantization/scheme.py   |   16 +-
 .../freetoken/mm/processors/deepseek_v41.py   |   88 +
 python/freetoken/models/config.py             |   12 +-
 .../freetoken/models/deepseek_v4/__init__.py  |    2 +-
 python/freetoken/models/deepseek_v4/model.py  |    2 +-
 .../freetoken/models/deepseek_v41/__init__.py |   14 +
 python/freetoken/models/deepseek_v41/args.py  |  225 +++
 .../models/deepseek_v41/attention.py          |  305 ++++
 .../models/deepseek_v41/compressor.py         |  131 ++
 .../freetoken/models/deepseek_v41/config.py   |   95 ++
 .../freetoken/models/deepseek_v41/engram.py   |  262 +++
 .../models/deepseek_v41/engram_table.py       |  296 ++++
 .../freetoken/models/deepseek_v41/indexer.py  |  127 ++
 python/freetoken/models/deepseek_v41/model.py |  289 ++++
 python/freetoken/models/deepseek_v41/moe.py   |   90 +
 python/freetoken/models/deepseek_v41/rope.py  |   58 +
 .../freetoken/models/deepseek_v41/vision.py   |  139 ++
 .../freetoken/models/deepseek_v41/weight.py   |  242 +++
 python/freetoken/models/register.py           |   14 +
 python/freetoken/moe/benchbw.py               |    8 +-
 python/freetoken/moe/cpu_executor.py          |   12 +-
 python/freetoken/moe/fused_ds_fp4.py          |   50 +-
 python/freetoken/scheduler/cache.py           |    9 +-
 python/freetoken/scheduler/prefill.py         |    6 +-
 python/freetoken/scheduler/scheduler.py       |   17 +-
 python/freetoken/server/api_server.py         |   27 +-
 python/freetoken/server/args.py               |   14 +
 .../freetoken/server/function_call_parser.py  |   49 +-
 python/freetoken/server/stats.py              |   10 +-
 python/freetoken/tokenizer/effort.py          |   35 +-
 python/freetoken/tokenizer/tokenize.py        |   34 +-
 scripts/ftw_hotfix.py                         |    4 +-
 tests/README.md                               |    1 +
 tests/attention/test_dsv41_backend.py         |  208 +++
 tests/checkpoint/test_ftw_hotfix_vision.py    |    6 +-
 tests/dsv4/test_dsv4_attn_backend.py          |    4 +-
 tests/kernels/test_dsv41_indexer.py           |  139 ++
 tests/kernels/test_dsv41_pack.py              |  112 ++
 tests/kernels/test_dsv41_sparse_attn.py       |  107 ++
 tests/kernels/test_dsv41_topk.py              |  159 ++
 tests/kernels/test_dsv4_fp8_block_linear.py   |  101 ++
 tests/kernels/test_dsv4_router_select.py      |   52 +
 tests/kvcache/test_cache_unit_bytes.py        |    4 +-
 tests/kvcache/test_dsv41_pool.py              |  205 +++
 tests/kvcache/test_dsv4_free_list.py          |    2 +-
 tests/kvcache/test_dsv4_pool.py               |    4 +-
 tests/kvcache/test_kv_cache_rebuild.py        |   12 +-
 tests/kvcache/test_pool_sizing_surface.py     |   18 +-
 tests/layers/test_mhc.py                      |   97 ++
 tests/models/deepseek_v41/__init__.py         |    0
 tests/models/deepseek_v41/common.py           |  176 ++
 tests/models/deepseek_v41/harness.py          |  182 ++
 tests/models/deepseek_v41/reference/LICENSE   |   21 +
 .../models/deepseek_v41/reference/__init__.py |    2 +
 tests/models/deepseek_v41/reference/engram.py |  186 +++
 tests/models/deepseek_v41/reference/kernel.py |  113 ++
 tests/models/deepseek_v41/reference/model.py  | 1458 +++++++++++++++++
 tests/models/deepseek_v41/reference/vision.py |  121 ++
 tests/models/deepseek_v41/test_config.py      |  169 ++
 .../models/deepseek_v41/test_engine_config.py |   81 +
 tests/models/deepseek_v41/test_engram.py      |  281 ++++
 .../models/deepseek_v41/test_forward_smoke.py |  184 +++
 .../deepseek_v41/test_reference_parity.py     |  219 +++
 tests/models/deepseek_v41/test_vision.py      |  214 +++
 tests/models/test_quant_config.py             |   26 +-
 tests/moe/test_cpu_moe.py                     |    3 +-
 tests/moe/test_dsfp4_prequant.py              |    2 +-
 tests/moe/test_dsfp4_reference_semantics.py   |  170 ++
 tests/scheduler/test_dsv4_generic_manager.py  |   31 +-
 tests/scheduler/test_scheduler_kv_usage.py    |   11 +-
 tests/server/test_effort_dialect.py           |   16 +
 tests/server/test_function_call_parser.py     |   28 +
 tests/server/test_parser_auto_selection.py    |    7 +
 tests/server/test_streaming_model_matrix.py   |    7 +
 tests/tokenizer/test_effort.py                |   36 +
 tests/tokenizer/test_tokenize.py              |   55 +
 117 files changed, 10881 insertions(+), 650 deletions(-)
 create mode 100644 python/freetoken/attention/dsv41_sparse.py
 create mode 100644 python/freetoken/kernel/triton/dsv4/router.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/__init__.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/indexer.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/pack.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/row_format.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/sparse_attn.py
 create mode 100644 python/freetoken/kernel/triton/dsv41/topk.py
 create mode 100644 python/freetoken/kvcache/dsv4/__init__.py
 create mode 100644 python/freetoken/kvcache/dsv4/v41_cost_model.py
 create mode 100644 python/freetoken/kvcache/dsv4/v41_pool.py
 create mode 100644 python/freetoken/kvcache/dsv4/v41_row_format.py
 rename python/freetoken/kvcache/{dsv4_cost_model.py => dsv4/v4_cost_model.py} (96%)
 rename python/freetoken/kvcache/{dsv4_paged_pool.py => dsv4/v4_pool.py} (52%)
 create mode 100644 python/freetoken/kvcache/dsv4/window_tier.py
 create mode 100644 python/freetoken/mm/processors/deepseek_v41.py
 create mode 100644 python/freetoken/models/deepseek_v41/__init__.py
 create mode 100644 python/freetoken/models/deepseek_v41/args.py
 create mode 100644 python/freetoken/models/deepseek_v41/attention.py
 create mode 100644 python/freetoken/models/deepseek_v41/compressor.py
 create mode 100644 python/freetoken/models/deepseek_v41/config.py
 create mode 100644 python/freetoken/models/deepseek_v41/engram.py
 create mode 100644 python/freetoken/models/deepseek_v41/engram_table.py
 create mode 100644 python/freetoken/models/deepseek_v41/indexer.py
 create mode 100644 python/freetoken/models/deepseek_v41/model.py
 create mode 100644 python/freetoken/models/deepseek_v41/moe.py
 create mode 100644 python/freetoken/models/deepseek_v41/rope.py
 create mode 100644 python/freetoken/models/deepseek_v41/vision.py
 create mode 100644 python/freetoken/models/deepseek_v41/weight.py
 create mode 100644 tests/attention/test_dsv41_backend.py
 create mode 100644 tests/kernels/test_dsv41_indexer.py
 create mode 100644 tests/kernels/test_dsv41_pack.py
 create mode 100644 tests/kernels/test_dsv41_sparse_attn.py
 create mode 100644 tests/kernels/test_dsv41_topk.py
 create mode 100644 tests/kernels/test_dsv4_fp8_block_linear.py
 create mode 100644 tests/kernels/test_dsv4_router_select.py
 create mode 100644 tests/kvcache/test_dsv41_pool.py
 create mode 100644 tests/models/deepseek_v41/__init__.py
 create mode 100644 tests/models/deepseek_v41/common.py
 create mode 100644 tests/models/deepseek_v41/harness.py
 create mode 100644 tests/models/deepseek_v41/reference/LICENSE
 create mode 100644 tests/models/deepseek_v41/reference/__init__.py
 create mode 100644 tests/models/deepseek_v41/reference/engram.py
 create mode 100644 tests/models/deepseek_v41/reference/kernel.py
 create mode 100644 tests/models/deepseek_v41/reference/model.py
 create mode 100644 tests/models/deepseek_v41/reference/vision.py
 create mode 100644 tests/models/deepseek_v41/test_config.py
 create mode 100644 tests/models/deepseek_v41/test_engine_config.py
 create mode 100644 tests/models/deepseek_v41/test_engram.py
 create mode 100644 tests/models/deepseek_v41/test_forward_smoke.py
 create mode 100644 tests/models/deepseek_v41/test_reference_parity.py
 create mode 100644 tests/models/deepseek_v41/test_vision.py
 create mode 100644 tests/moe/test_dsfp4_reference_semantics.py

diff --git a/docs/models.md b/docs/models.md
index 2cd3baa317..b999d29e41 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -5,6 +5,7 @@ for them; other checkpoints of the same architectures work too.
 
 | Model | HF checkpoints |
 |---|---|
+| DeepSeek-V4.1-Flash | [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) |
 | DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) |
 | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) |
 | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) |
@@ -27,6 +28,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 
 | Family | Image tokens | `--image-min-tokens` / `--image-max-tokens` | `--mm-processor-kwargs` example |
 | --- | --- | --- | --- |
+| DeepSeek-V4.1-Flash (ViT tower, streamed under `--mm-encoder-weights host`) | one feature per 42x42 pixels, plus a learned row break per row and two delimiters; checkpoint cap 1024 total span tokens | maximum caps the entire span; minimum maps to a pixel-area floor before the cap; checkpoint minimum area is 544x544 | `{"max_image_tokens": 512, "min_pixels": 295936}` |
 | Qwen3.6 (both variants, every listed weight format), Qwen3.8-Flash-Next, Qwen3-VL | one token per 32x32 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 64 to 16384 tokens | `{"size": {"longest_edge": 1048576}}` |
 | Gemma-4 26B-A4B, 31B (`gemma4`: ViT tower, streamed under `--mm-encoder-weights host`) | one of the soft-token budgets 70 / 140 / 280 / 560 / 1120, every image scaled to its budget as far as the aspect ratio allows | the maximum picks the largest budget within it, below 70 is refused at start-up; the minimum has no effect | `{"max_soft_tokens": 1120}` |
 | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same |
@@ -86,4 +88,34 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md).
 - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the
   authoritative model args are read from there.
+- DeepSeek-V4.1-Flash serves text and images from the HF checkpoint as shipped (fp8 block-32
+  dense weights, fp4 experts on the offload cache, 890 B/token global KV in packed
+  fp4/fp8 pools). Its two 98 GiB Engram tables stream from the checkpoint shards on
+  demand (keep them on a fast NVMe; the 6 GiB of table scales stay in host RAM), using
+  bounded pinned staging buffers; an FTW conversion copies the shards that hold them next
+  to the checkpoint.
+  SWA bounded replay (DeepSeek_V41_Tech_Report.pdf, shipped in the checkpoint, §3.2.2): a replayed
+  segment recomputes only the SWA KV of its last `n_win` tokens and truncates each query's window to
+  the segment. *Encoder* replay rebuilds the encoder SWA KV behind a prefix hit from the global KV
+  alone (approximate; the report's fallback when the SWA KV of a hit has been evicted). *Decoder*
+  replay runs the decoder layers on each prompt's last `n_win` tokens only; their SWA KV is never
+  prefix-cached and post-training simulated it. FreeToken keeps encoder SWA KV in the radix cache (no
+  encoder replay); `--swa-decoder-replay bounded` (default) is the report's decoder replay, with a
+  prefix hit stopping at least `n_win` tokens before the prompt end so those tokens are prefilled;
+  `exact` runs the decoder on every prompt token (reference numerics).
+  Bounded output differs from exact by construction and does not depend on the prefill chunk size.
+  Under expert offload the prefill time is bounded by streaming each layer's experts, so decoder
+  replay saves decoder-layer compute, not prefill time.
+  Run `ft bench bw --model dsv4.1-flash` to measure local PCIe and CPU bandwidth before
+  choosing a MoE strategy; `--moe-strategy hybrid` splits expert misses using those
+  measurements. With `--moe-cache-auto`, `--kv-reserve-tokens` sets the token capacity reserved
+  before the expert cache is allocated, subject to the attention pool's structural minimum; size it
+  for the intended context and concurrency. `--max-extend-length` controls the prefill chunk size
+  and its workspace. Image input uses the shared multimodal processor and encoder cache, including
+  chunked prefill and prefix replay; `--text-model-only` skips the vision weights. DSpark
+  speculative decoding is not served.
+  `reasoning_effort` takes `low`, `high` or `max`; the 1-100 integer budget goes through the
+  template kwargs with thinking on, `"chat_template_kwargs": {"enable_thinking": true,
+  "reasoning_effort": 37}`, and the checkpoint's encoder validates it. Runtime window-cache
+  controls use 128-token pages.
 - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM.
diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py
index 8a410ffe4a..eda3b61f55 100644
--- a/python/freetoken/attention/__init__.py
+++ b/python/freetoken/attention/__init__.py
@@ -102,6 +102,16 @@ def create_dsv4_sparse_backend(config: ModelConfig):
     return DSV4SparseAttnBackend(config)
 
 
+@SUPPORTED_ATTENTION_BACKENDS.register(
+    "dsv41_sparse",
+    BackendInfo(supported_types=frozenset({AttnType.DSV41})),
+)
+def create_dsv41_sparse_backend(config: ModelConfig):
+    from .dsv41_sparse import DSV41SparseAttnBackend
+
+    return DSV41SparseAttnBackend(config)
+
+
 @SUPPORTED_ATTENTION_BACKENDS.register(
     "dsa",
     BackendInfo(supported_types=frozenset({AttnType.MLA, AttnType.DSA})),
diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py
index 9115b80929..846e99447a 100644
--- a/python/freetoken/attention/base.py
+++ b/python/freetoken/attention/base.py
@@ -20,6 +20,9 @@ class AttnType(str, Enum):
     MLA = "mla"  # plain latent-KV MLA -> MLAKVCache
     DSA = "dsa"  # latent-KV MLA + DSA sparse indexer -> DSAKVCache
     DSV4 = "dsv4"  # DSV4 window+compressed sparse -> DSV4PagedKVCache
+    # DSV41 (DeepSeek-V4.1): window + cross-layer-shared compressed KV in packed fp8/fp4
+    # rows, Full/Reindex/Reuse layer modes -> DSV41PagedKVCache
+    DSV41 = "dsv41"
     LINEAR = "linear"  # GDN/mamba state layers -> LinearStatePool
     # GQA block-sparse (MiniMax-M3): paged GQA K/V + a per-sparse-layer index-key
     # slab; the indexer picks top-k 128-token blocks per query -> BSAKVCache
diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py
new file mode 100644
index 0000000000..baaa2cf019
--- /dev/null
+++ b/python/freetoken/attention/dsv41_sparse.py
@@ -0,0 +1,438 @@
+"""DeepSeek-V4.1 sparse-attention backend.
+
+Same contract and division of labour as ``dsv4_sparse``: the backend owns the per-forward KV
+ADDRESSING over the shared page table (window ring slots, compressed rows, compress-state carry,
+the decode snapshot and the CUDA-graph staging buffers) plus the selection helpers that are pure
+index arithmetic (causal top-k, the hierarchical candidate pool); the model computes projections
+and hands the picks back to be resolved into pool rows.
+
+What DSV41 adds on top of DSV4's vocabulary:
+
+* **cross-layer sharing** -- the compressed tiers belong to the kv-source layer; a consumer layer
+  addresses ``source_of(layer)``'s pools. The Top-K rows and the candidate pool the Full / Reindex
+  layers produce ride on the per-forward ``SharedSelection`` for the Reuse layers to read.
+* **packed rows** -- writes quantize into the pool (``pool.store_*``), reads dequantize in-kernel.
+* **prefill passes with a window floor** -- a ``PrefillSegment`` carries ``window_floor``: the
+  absolute position a query's sliding window may not reach below (Decoder SWA Bounded Replay
+  truncates the decoder's window at ``L - W``, the start of the prompt's last window). The metadata
+  carries the encoder pass (every new token) and the decoder pass (the new tokens inside the
+  prompt's last window, or every new token in exact mode).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List
+
+import torch
+from freetoken.core import Batch, get_global_ctx
+
+from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
+
+if TYPE_CHECKING:
+    from freetoken.models import ModelConfig
+
+
+@dataclass(frozen=True)
+class PrefillSegment:
+    """One request's slice of a flat prefill token stream.
+
+    ``offset``/``n`` tile the stream; ``start_pos`` is the absolute position of the first token;
+    ``window_floor`` bounds every query's sliding window from below (0 = the exact window; bounded
+    replay sets it to the start of the prompt's last window).
+    """
+
+    offset: int
+    n: int
+    table_idx: int
+    start_pos: int
+    window_floor: int = 0
+
+    @property
+    def end(self) -> int:
+        return self.start_pos + self.n
+
+
+@dataclass
+class SharedSelection:
+    """What the index-producing layers hand down the stack within ONE forward.
+
+    ``topk_rows`` are GLOBAL main-pool rows (``-1`` masked) for the current kv source, laid out
+    exactly as the attention kernel consumes them; ``cmp_counts`` is the decode-only per-row live
+    count that bounds the kernel's loop; ``candidates`` is the Hierarchical Sparse Indexer's
+    compressed-position pool (``-1`` = empty slot). Layers run in order and every producer writes
+    before its consumers read, so one slot each suffices.
+    """
+
+    source: int | None = None
+    topk_rows: torch.Tensor | None = None
+    cmp_counts: torch.Tensor | None = None
+    candidates: torch.Tensor | None = None
+
+
+@dataclass
+class DSV41AttnMetadata(BaseAttnMetadata):
+    last_indices: torch.Tensor
+    # Prefill: the encoder pass tiles the batch's new tokens; the decoder pass tiles the tokens the
+    # decoder layers process (``decoder_rows`` gathers them out of the encoder stream; None = all of
+    # it). ``decoder_pad``: some request has no decoder row this chunk, and its ``last_indices`` entry
+    # points at one placeholder row the model appends after the decoder stream (its logits are unused).
+    segments: List[PrefillSegment] | None = None
+    decoder_segments: List[PrefillSegment] | None = None
+    decoder_rows: torch.Tensor | None = None
+    decoder_pad: bool = False
+    # Decode: the whole-history full-loc snapshot (captured buffer under a replay, lazy eager copy).
+    full_snap: torch.Tensor | None = None
+    table_rows: torch.Tensor | None = None
+    window_ar: torch.Tensor | None = None
+    selection: SharedSelection = field(default_factory=SharedSelection)
+
+    def get_last_indices(self, bs: int) -> torch.Tensor:
+        return self.last_indices[:bs]
+
+    @property
+    def stage_width(self) -> int:
+        assert self.full_snap is not None, "stage_width is capture/replay-only"
+        return self.full_snap.shape[1]
+
+    def full_snapshot(self) -> torch.Tensor:
+        if self.full_snap is None:
+            assert self.table_rows is not None, "snapshot is decode-only"
+            pool = get_global_ctx().kv_cache
+            self.full_snap = pool.full_loc_map.index_select(0, self.table_rows).to(torch.int64)
+        return self.full_snap
+
+    def window_ctx(self, pos: torch.Tensor, rows: torch.Tensor):
+        """Layer-invariant decode ring context ``(window_slots, prev_window_slots,
+        window_slots_topk [B, 1, win])``; computed fresh on every call (never cached -- a capture
+        must record these gathers), off the snapshot so a concurrent allocation cannot redirect
+        an in-flight replay."""
+        snap = self.full_snapshot()
+        translate = get_global_ctx().kv_cache.translate_full_to_window
+        bs = pos.shape[0]
+        j = self.window_ar
+        assert j is not None, "window_ctx is decode-only"
+        win = j.shape[0]
+        window_slots = translate(snap[rows, pos])
+        prev_window_slots = translate(snap[rows, (pos - 1).clamp_min(0)])
+        p = pos[:, None] - ((pos[:, None] - j[None, :]) % win)
+        ws = translate(snap[rows[:, None], p.clamp(min=0)])
+        window_slots_topk = torch.where((p >= 0) & (j[None, :] <= pos[:, None]), ws, -1).view(bs, 1, win)
+        return window_slots, prev_window_slots, window_slots_topk
+
+    def private_window_ctx(self, pos: torch.Tensor, rows: torch.Tensor):
+        """The private-ring counterpart of ``window_ctx`` for request-private window layers:
+        ``(window_slots [B], window_slots_topk [B, 1, win])`` in the layer's ring, where a request's
+        row holds positions by ``pos % win`` (candidate ``j`` is the position congruent to ``j``)."""
+        assert self.table_rows is not None and self.window_ar is not None, "private_window_ctx is decode-only"
+        pool = get_global_ctx().kv_cache
+        bs = pos.shape[0]
+        j = self.window_ar
+        win = j.shape[0]
+        table = self.table_rows[rows]
+        window_slots = pool.ring_slots(table, pos)
+        p = pos[:, None] - ((pos[:, None] - j[None, :]) % win)  # the position at ring index j (< 0: none yet)
+        window_slots_topk = pool.ring_slots(table[:, None], p).view(bs, 1, win)
+        return window_slots, window_slots_topk
+
+
+@dataclass
+class DSV41CaptureData:
+    full_snap: torch.Tensor
+    last_indices: torch.Tensor
+    table_rows: torch.Tensor  # each batch row's page-table row (the private window rings key on it)
+
+    @classmethod
+    def create(cls, max_bs: int, width: int, device: torch.device) -> "DSV41CaptureData":
+        return cls(
+            full_snap=torch.full((max_bs, width), -1, dtype=torch.int64, device=device),
+            last_indices=torch.arange(max_bs, dtype=torch.int32, device=device),
+            table_rows=torch.zeros(max_bs, dtype=torch.int64, device=device),
+        )
+
+
+class DSV41SparseAttnBackend(BaseAttnBackend):
+    def __init__(self, config: ModelConfig):
+        self.config = config
+        self.device = get_global_ctx().kv_cache.device
+        args = config.dsv41_args
+        self.window_size = args.window_size
+        self.swa_decoder_replay: str = args.swa_decoder_replay
+        self.capture: DSV41CaptureData | None = None
+        self.capture_bs: List[int] = []
+        self.max_graph_bs = 0
+        self._window_ar = torch.arange(self.window_size, device=self.device)
+
+    @property
+    def pool(self):
+        return get_global_ctx().kv_cache
+
+    # ----- generic contract -------------------------------------------------------------
+    def forward(self, q, k, v, layer_id, batch, attn_spec: AttentionSpec | None = None):
+        raise NotImplementedError("DSV41 attention is driven per-tier from the model module; use DSV41SparseAttnBackend.attend().")
+
+    def prepare_metadata(self, batch: Batch) -> None:
+        if not batch.is_decode:
+            batch.attn_metadata = self._prefill_metadata(batch)
+            return
+        last = torch.tensor([r.extend_len for r in batch.padded_reqs], dtype=torch.int32, device=self.device).cumsum_(0) - 1
+        batch.attn_metadata = DSV41AttnMetadata(last_indices=last, table_rows=self._table_rows(batch), window_ar=self._window_ar)
+
+    def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata:
+        """The encoder pass runs every request's new tokens ``[cached_len, device_len)``. Exact mode runs
+        the decoder on the same stream. Decoder SWA Bounded Replay runs it on the new tokens inside the
+        prompt's last window ``[L - W, L)`` (``L = prompt_len``) with every query's window floored at
+        ``L - W``: the chunks of a chunked prefill each take their part of that window (earlier rows
+        come from the request's private ring), and a chunk that ends before it runs no decoder row."""
+        segments: List[PrefillSegment] = []
+        off = 0
+        for r in batch.reqs:
+            segments.append(PrefillSegment(off, r.extend_len, r.table_idx, r.cached_len))
+            off += r.extend_len
+        if self.swa_decoder_replay == "exact":
+            last = torch.tensor([s.n for s in segments], dtype=torch.int32, device=self.device).cumsum_(0) - 1
+            return DSV41AttnMetadata(last_indices=last, segments=segments, decoder_segments=segments)
+        win = self.window_size
+        dec: List[PrefillSegment] = []
+        rows: list[torch.Tensor] = []
+        last: list[int] = []
+        doff = 0
+        for r, s in zip(batch.reqs, segments):
+            floor = max(0, r.prompt_len - win)
+            lo = max(s.start_pos, floor)
+            n = max(0, s.end - lo)
+            if n:
+                # the decoder's window KV goes to the request's private rings, never to radix-shared pages
+                dec.append(PrefillSegment(doff, n, s.table_idx, lo, window_floor=floor))
+                rows.append(torch.arange(s.offset + lo - s.start_pos, s.offset + s.n))
+                doff += n
+            last.append(doff - 1 if n else -1)
+        pad = -1 in last
+        last_indices = torch.tensor([doff if i < 0 else i for i in last], dtype=torch.int32, device=self.device)
+        decoder_rows = torch.cat(rows).to(self.device, non_blocking=True) if rows else torch.empty(0, dtype=torch.int64, device=self.device)
+        return DSV41AttnMetadata(last_indices=last_indices, segments=segments, decoder_segments=dec, decoder_rows=decoder_rows, decoder_pad=pad)
+
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+        assert self.capture is None, "Capture already initialized."
+        self.max_graph_bs = max(bs_list)
+        self.capture = DSV41CaptureData.create(self.max_graph_bs, max_seq_len, self.device)
+        self.capture_bs = sorted(bs_list)
+
+    def prepare_for_capture(self, batch: Batch) -> None:
+        bs = batch.size
+        rows = torch.full((bs,), batch.padded_reqs[0].table_idx, dtype=torch.int64, device=self.device)
+        self._point_to_capture(batch, bs, rows)
+
+    def prepare_for_replay(self, batch: Batch) -> None:
+        self._point_to_capture(batch, batch.padded_size, self._table_rows(batch))
+
+    def _table_rows(self, batch: Batch) -> torch.Tensor:
+        assert batch.active_table_idx is not None, "decode batch is missing its page-table rows"
+        return batch.active_table_idx.to(torch.int64)
+
+    def _point_to_capture(self, batch: Batch, bs: int, rows_ti: torch.Tensor) -> None:
+        assert self.capture is not None and bs <= self.max_graph_bs
+        cap = self.capture
+        src = self.pool.full_loc_map.index_select(0, rows_ti)
+        w = min(src.shape[1], cap.full_snap.shape[1])
+        cap.full_snap[:bs, :w].copy_(src[:bs, :w])
+        cap.table_rows[:bs].copy_(rows_ti[:bs])
+        batch.attn_metadata = DSV41AttnMetadata(
+            last_indices=cap.last_indices[:bs], full_snap=cap.full_snap[:bs], table_rows=cap.table_rows[:bs], window_ar=self._window_ar,
+        )
+
+    # ----- metadata / selection access -------------------------------------------------
+    @property
+    def metadata(self) -> DSV41AttnMetadata:
+        md = get_global_ctx().batch.attn_metadata
+        assert isinstance(md, DSV41AttnMetadata)
+        return md
+
+    @property
+    def selection(self) -> SharedSelection:
+        return self.metadata.selection
+
+    def snapshot(self) -> torch.Tensor:
+        return self.metadata.full_snapshot()
+
+    # ----- window tier -------------------------------------------------------------------
+    def window_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor:
+        """SHARED window slots of positions ``[lo, hi)`` off the request's LIVE full locs (slid-out -> -1)."""
+        return self.pool.translate_full_to_window(self.pool.full_loc_map[ti, lo:hi])
+
+    def ring_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor:
+        """PRIVATE-ring slots of positions ``[lo, hi)`` of page-table row ``ti`` (``hi - lo <= win``)."""
+        assert hi - lo <= self.window_size, f"a private ring holds one window, not {hi - lo} positions"
+        return self.pool.ring_slots(torch.tensor(ti, device=self.device), torch.arange(lo, hi, device=self.device))
+
+    def layer_window_slots_of(self, layer_id: int, ti: int, lo: int, hi: int) -> torch.Tensor:
+        """Where ``layer_id`` keeps the window KV of positions ``[lo, hi)``: its private ring or the shared pool."""
+        if self.pool.is_private_window(layer_id):
+            return self.ring_slots_of(ti, lo, hi)
+        return self.window_slots_of(ti, lo, hi)
+
+    def store_window(self, kv: torch.Tensor, layer_id: int, window_slots: torch.Tensor) -> None:
+        self.pool.store_window(kv, layer_id, window_slots)
+
+    def window_topk_prefill(self, seg: PrefillSegment, layer_id: int | None = None) -> torch.Tensor:
+        """Per-query window candidates for a prefill segment as GLOBAL window slots ``[1, n, win]``
+        (``-1`` where empty) in the tier ``layer_id`` reads (shared pool by default). Query at
+        absolute ``p`` sees ``[max(floor, first_retained, p - win + 1), p]``."""
+        win, device = self.window_size, self.device
+        lo = max(0, seg.start_pos - win + 1, seg.window_floor)
+        private = layer_id is not None and self.pool.is_private_window(layer_id)
+        ws_pool = (self.ring_slots_of if private else self.window_slots_of)(seg.table_idx, lo, seg.end)  # [end - lo]
+        abs_p = seg.start_pos + torch.arange(seg.n, device=device).unsqueeze(1)
+        cand = (abs_p - win + 1).clamp(min=lo) + torch.arange(win, device=device)
+        cols = torch.where(cand > abs_p, -1, cand - lo)
+        g = ws_pool[cols.clamp_min(0)]
+        return torch.where(cols < 0, -1, g).unsqueeze(0)
+
+    # ----- compressed tiers (per kv source) ----------------------------------------------
+    def source_of(self, layer_id: int) -> int:
+        return self.pool.source_of(layer_id)
+
+    def compressed_rows_of(self, ti: int, group_starts: torch.Tensor, ratio: int) -> torch.Tensor:
+        """Main / index rows of the compressed groups whose ABSOLUTE first positions are
+        ``group_starts``, off the request's LIVE full locs (a page is ratio-divisible, so every
+        position of a group shares one row)."""
+        return self.pool.cmp_rows(self.pool.full_loc_map[ti, group_starts], ratio)
+
+    def locs_prefill(self, ti: int, end: int) -> torch.Tensor:
+        """``[1, end]`` int32: the request's live full locs for positions ``[0, end)`` -- what the indexer
+        derives compressed rows from (``loc // ratio``) for the positions its queries may see."""
+        return self.pool.full_loc_map[ti : ti + 1, :end]
+
+    def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None:
+        self.pool.store_main(latent, source, rows)
+
+    def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None:
+        self.pool.store_index(k, source, rows)
+
+    def decode_store_rows(self, rows: torch.Tensor, pos: torch.Tensor, ratio: int, source: int, completed: torch.Tensor) -> torch.Tensor:
+        """Per-row decode destination: the completed group's arithmetic row, or the row's OWN scratch
+        row when this step did not complete a group (graph-safe, collision-free masked store)."""
+        row_of_group = self.pool.cmp_rows(self.snapshot()[rows, pos], ratio)
+        scratch = rows + self.pool.scratch_base[source]
+        return torch.where(completed, row_of_group, scratch)
+
+    # ----- compress-state ring (ratio > 1 sources) ---------------------------------------
+    def ring_page_base(self, window_slots: torch.Tensor, ring_size: int) -> torch.Tensor:
+        return torch.div(window_slots, self.window_size, rounding_mode="floor") * ring_size
+
+    def carry_state_loc(self, window_slot: int, ring_size: int) -> torch.Tensor:
+        base = (window_slot // self.window_size) * ring_size
+        return torch.arange(base, base + ring_size, device=self.device, dtype=torch.int64)
+
+    def read_carry(self, source: int, window_slot: int) -> torch.Tensor:
+        """The ``[ring_size, 2 * head_dim]`` carry block at ``window_slot``'s page (kv | score)."""
+        ring = self.pool.state_ring[source]
+        return ring.get(self.carry_state_loc(window_slot, ring.ring_size))
+
+    def write_carry(self, source: int, window_slot: int, kv_score: torch.Tensor) -> None:
+        ring = self.pool.state_ring[source]
+        ring.set(self.carry_state_loc(window_slot, ring.ring_size), kv_score)
+
+    def read_carry_blocks(self, source: int, window_slots: torch.Tensor) -> torch.Tensor:
+        ring = self.pool.state_ring[source]
+        return ring.get_blocks(self.ring_page_base(window_slots, ring.ring_size))
+
+    def write_carry_blocks(self, source: int, window_slots: torch.Tensor, blocks: torch.Tensor) -> None:
+        ring = self.pool.state_ring[source]
+        ring.set_blocks(self.ring_page_base(window_slots, ring.ring_size), blocks)
+
+    def write_boundary_carries(self, source: int, *, lo: int, hi: int, window_slots: torch.Tensor) -> None:
+        """Persist the compressor carry at every window-page boundary ``B`` in ``(lo, hi]`` so a
+        page-aligned radix match can resume by value. A page holds whole groups (``P % ratio == 0``),
+        so the carry at a boundary is the empty group -- the reset block is written so a resume never
+        reads a stale one. ``window_slots`` is indexed by ``pos - lo``. One batched ring write, no host syncs."""
+        ring = self.pool.state_ring.get(source)
+        if ring is None:
+            return
+        P = self.window_size
+        first = (lo // P + 1) * P
+        if first > hi:
+            return
+        bounds = torch.arange(first, hi + 1, P, device=self.device)
+        page_slots = window_slots[bounds - 1 - lo].to(torch.int64)  # the last slot of each page
+        locs = (torch.div(page_slots, P, rounding_mode="floor") * ring.ring_size)[:, None] + torch.arange(ring.ring_size, device=self.device)
+        empty = torch.cat([
+            torch.zeros(ring.ring_size, ring.item_size, dtype=torch.float32, device=self.device),
+            torch.full((ring.ring_size, ring.item_size), float("-inf"), dtype=torch.float32, device=self.device),
+        ], dim=-1)
+        ring.set(locs.flatten(), empty.repeat(bounds.numel(), 1))
+
+    # ----- indexer ----------------------------------------------------------------------
+    def indexer_logits(
+        self, q: torch.Tensor, weights: torch.Tensor, source: int, locs: torch.Tensor, ratio: int, live: torch.Tensor,
+        T: int | None = None, candidates: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """``[B, S, T]`` full-range scores (live columns written) or ``[B, S, NC]`` candidate scores; the
+        row of compressed position ``t`` is ``locs[b, t * ratio] // ratio``."""
+        from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed
+
+        return indexer_logits_packed(q, weights, self.pool.idx_pool[source], self.pool.idx_fmt, locs, ratio, live, T=T, candidates=candidates)
+
+    @staticmethod
+    def select_topk(scores: torch.Tensor, live: torch.Tensor, topk: int) -> torch.Tensor:
+        """Exact causal top-k over ``[B, S, T]`` full-range scores (columns at or past ``live`` unread):
+        compressed positions ``[B, S, topk]`` int32 in ascending order, ``-1`` in the tail for picks
+        that do not exist (fewer live or finite positions than ``topk``). Ties: lowest position wins."""
+        from freetoken.kernel.triton.dsv41.topk import dsv41_topk
+
+        b, s, t = scores.shape
+        return dsv41_topk(scores.reshape(b * s, t), live.reshape(b * s), topk).view(b, s, topk)
+
+    @staticmethod
+    def select_topk_in_candidates(scores: torch.Tensor, candidates: torch.Tensor, topk: int) -> torch.Tensor:
+        """Top-k over ``[B, S, NC]`` candidate-aligned scores -> the picked compressed positions
+        ``[B, S, topk]`` int32. The candidate list is a sorted valid prefix (``-1`` tail) and empty /
+        unreachable slots score ``-inf``, so the ascending winning slots map to ascending positions."""
+        from freetoken.kernel.triton.dsv41.topk import dsv41_topk
+
+        b, s, nc = scores.shape
+        live = torch.full((b * s,), nc, dtype=torch.int32, device=scores.device)
+        slots = dsv41_topk(scores.reshape(b * s, nc), live, topk).view(b, s, topk)
+        pos = candidates.gather(-1, slots.clamp_min(0).to(torch.int64)).to(torch.int32)
+        return torch.where(slots < 0, -1, pos)
+
+    @staticmethod
+    def select_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor:
+        """Level one of the Hierarchical Sparse Indexer: the ``topk_blocks`` best live ``block_size``
+        blocks by their best position, the block holding the newest live position always kept,
+        expanded to the compact candidate list ``[B, S, topk_blocks * block_size]`` int32 of
+        compressed positions (ascending, ``-1`` tail). Mirrors the reference ``select_candidate_blocks``
+        (which returns the equivalent boolean mask)."""
+        from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks
+
+        b, s, t = scores.shape
+        return dsv41_candidate_blocks(scores.reshape(b * s, t), live.reshape(b * s), topk_blocks, block_size).view(b, s, -1)
+
+    @staticmethod
+    def positions_to_rows(positions: torch.Tensor, locs: torch.Tensor, ratio: int) -> torch.Tensor:
+        """Compressed positions ``[B, S, K]`` (``-1`` masked) -> GLOBAL main rows ``locs[b, p * ratio] // ratio``."""
+        b = positions.shape[0]
+        idx = (positions.to(torch.int64) * ratio).clamp_min(0).flatten(1)
+        full = locs.to(torch.int64).gather(1, idx.clamp(max=locs.shape[1] - 1)).view_as(positions)
+        rows = torch.div(full, ratio, rounding_mode="floor")
+        return torch.where((positions < 0) | (full < 0), -1, rows)
+
+    # ----- attention -----------------------------------------------------------------------
+    def attend(
+        self, q: torch.Tensor, layer_id: int, topk_idxs: torch.Tensor, n_window: int, attn_sink: torch.Tensor,
+        softmax_scale: float, cmp_counts: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Paged sparse attention over ``[window | compressed]`` global rows; the compressed half reads
+        the layer's kv source (window-only layers pass ``n_window == topk``)."""
+        from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed
+
+        pool, args = self.pool, self.config.dsv41_args
+        src = args.roles[layer_id].kv_source
+        cmp = pool.main_pool[src] if src is not None else pool.main_pool[args.backbone_kv_sources[0]]
+        return sparse_attn_packed(
+            q, pool.window_pool[layer_id], pool.win_fmt, cmp, pool.main_fmt, attn_sink,
+            topk_idxs, n_window, softmax_scale, cmp_counts=cmp_counts,
+        )
+
+
+__all__ = ["DSV41SparseAttnBackend", "DSV41AttnMetadata", "PrefillSegment", "SharedSelection"]
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index fd691d180c..45a763c0fe 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -43,6 +43,8 @@ class Req:
     mm_items: list | None = None
     mrope_positions_full: torch.Tensor | None = None  # [3, prompt_len] int32, CPU
     mrope_delta: int = 0
+    # The whole prompt's length; 0 = this request's input_ids (a ChunkedReq's end at its chunk).
+    prompt_len: int = 0
 
     # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or
     # until allocated from LinearStatePool. Set by the scheduler (P2). ---
@@ -69,6 +71,7 @@ def __post_init__(self) -> None:
         assert self.input_ids.is_cpu
         self.device_len = len(self.input_ids)
         self.max_device_len = len(self.input_ids) + self.output_len
+        self.prompt_len = self.prompt_len or self.device_len
         assert 0 <= self.cached_len < self.device_len <= self.max_device_len
         self._alloc_ids_buf()
 
diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py
index d522be4991..b5f07ee8ab 100644
--- a/python/freetoken/engine/cache_budget.py
+++ b/python/freetoken/engine/cache_budget.py
@@ -24,7 +24,7 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int:
     # marlin/b12x gate_up/down alpha scales are fixed [L*E] residency (do not scale
     # with cache_size), so they are intentionally excluded from the per-slot growth term.
     # tensor[0].numel() is the per-row element count (one expert slot); see the matching
-    # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4_paged_pool.py.
+    # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py.
     return sum(t[0][0].numel() * t[0].element_size() for t in sources.values())
 
 
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 3e38cb2165..b6fa8b8949 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -34,6 +34,9 @@ class EngineConfig:
     quant_backend: str | None = None
     # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM.
     ple_backend: str = "disk"
+    # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each
+    # prompt's last window; "exact" runs them on every token.
+    swa_decoder_replay: str = "bounded"
     # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered
     # experts in parallel but falls back to serial when free RAM can't cover the banks + the
     # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index e5ccf71050..bf3960ab8d 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -129,6 +129,8 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str:
     candidates: list[tuple[str, bool]] = []
     if AttnType.DSV4 in required:
         candidates.append(("dsv4_sparse", True))
+    if AttnType.DSV41 in required:
+        candidates.append(("dsv41_sparse", True))
     if required & {AttnType.MLA, AttnType.DSA}:
         candidates.append(("dsa", True))
     if AttnType.BSA in required:
@@ -812,6 +814,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         # Decode batches never exceed max_running_req, but CUDA-graph padding can
         # round a batch up to the largest captured size; cover both.
         max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1)
+        method = sample.quant_method
         executor = CpuMoeExecutor(
             cache,
             top_k=sample.top_k,
@@ -823,7 +826,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
             swiglu_alpha=float(sample.alpha),
             swiglu_limit=sample.limit,
             # FIXME: the None branch serves GGUF q4_0 banks, which have no quant method yet; drop it once GGUF joins the quant path
-            fmt=sample.quant_method.cpu_format if sample.quant_method is not None else None,
+            fmt=method.cpu_format if method is not None else None,
+            act_block=method.cfg.act_block if method is not None else None,
         )
         cache.set_cpu_executor(executor)
         self.cpu_moe_executor = executor
@@ -1200,6 +1204,39 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str:
     return requested
 
 
+def _adjust_dsv41_config(config: EngineConfig, override) -> None:
+    """DeepSeek-V4.1 config reconciliation, the DSV4 policy with the replay knob: sync the
+    runtime into ``dsv41_args``, page_size = the window page P, radix -> swa_radix, decode graph
+    batches <= max_running_req. Unlike DSV4 the prefill chunk keeps ``max_extend_tokens`` (whole
+    window pages): at a 1M ceiling the window pool alone would admit ~40K-token chunks whose
+    activations (64 x 512 latent queries per token) do not fit next to the expert cache."""
+    model_config = config.model_config
+    args = model_config.dsv41_args
+    args.max_seq_len = config.max_seq_len
+    args.max_batch_size = config.max_running_req + 1  # +1 dummy
+    args.swa_decoder_replay = config.swa_decoder_replay
+    P = args.window_size
+    override("page_size", P)
+    logger.info_rank0(
+        f"DSV41 KV pages are {P}-token window pages; page_size set to {P}; SWA decoder replay: {args.swa_decoder_replay}"
+        + ("" if args.swa_decoder_replay == "exact" else f" (prefix hits stop {P} tokens before the prompt end)")
+    )
+    if getattr(config, "cache_type", "radix") != "naive":
+        override("cache_type", "swa_radix")
+    current = getattr(config, "max_extend_tokens", None)  # a SchedulerConfig field
+    if current is not None and (extend := max(P, current // P * P)) != current:
+        override("max_extend_tokens", extend)
+    mr = config.max_running_req
+    if config.cuda_graph_max_bs is not None and config.cuda_graph_max_bs > mr:
+        logger.warning_rank0(f"cuda_graph_max_bs {config.cuda_graph_max_bs} exceeds max_running_req {mr}; clamping (larger decode batches never occur).")
+        override("cuda_graph_max_bs", mr)
+    if config.cuda_graph_bs is not None:
+        kept = [bs for bs in config.cuda_graph_bs if bs <= mr]
+        if kept != list(config.cuda_graph_bs):
+            logger.warning_rank0(f"dropping cuda_graph_bs entries above max_running_req {mr}: {[bs for bs in config.cuda_graph_bs if bs > mr]}")
+            override("cuda_graph_bs", kept)
+
+
 def _adjust_dsv4_config(config: EngineConfig, override) -> None:
     """DSV4 engine-config reconciliation at config-resolution time (before the pool exists).
     Syncs the resolved runtime config into the opaque ``dsv4_args`` payload, sets
@@ -1477,6 +1514,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     model_config = config.model_config
     single_stream_only = getattr(model_config, "single_stream_only", False)
     is_dsv4 = getattr(model_config, "dsv4_args", None) is not None
+    is_dsv41 = getattr(model_config, "dsv41_args", None) is not None
     has_swa_attention = getattr(model_config, "has_swa_attention", False)
     has_linear_attention = getattr(model_config, "has_linear_attention", False)
     is_moe = getattr(model_config, "is_moe", False)
@@ -1520,6 +1558,8 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
 
     if is_dsv4:
         _adjust_dsv4_config(config, override)
+    if is_dsv41:
+        _adjust_dsv41_config(config, override)
 
     if has_swa_attention:
         # Both SWA cache paths use the global-paged swa pool (page_size==1 only for now).
@@ -1806,7 +1846,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config).
     rotary = getattr(model_config, "rotary_config", None)
     seq_override = getattr(config, "max_seq_len_override", None)
-    if seq_override is not None and rotary is not None and not is_dsv4:
+    if seq_override is not None and rotary is not None and not (is_dsv4 or is_dsv41):
         if seq_override > rotary.table_positions:
             raise ValueError(
                 f"--max-seq-len-override {seq_override} exceeds the model's "
diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index a8ef85d432..ecd1f8502b 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -332,6 +332,17 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         expert_formats=("ds_fp4",),
         embed_indexing=False,  # plain nn.Embedding
     ),
+    AotModel(
+        # DSV41 pools are packed byte rows written by the dsv41 pack kernel, not store_cache;
+        # the routed experts share DSV4's ds_fp4 bank layout (hidden 5120, intermediate 2304).
+        name="deepseek-ai/DeepSeek-V4.1-Flash",
+        architecture="DeepseekV41ForCausalLM",
+        hidden_size=5120,
+        kv_groups=(),
+        top_k=6,
+        moe_intermediate_size=2304,
+        expert_formats=("ds_fp4",),
+    ),
     # ---- dense checkpoints (store/index only, no expert banks) ----
     AotModel(
         name="Qwen/Qwen3.6-27B",
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index fb6a33f497..86765013e3 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -1063,11 +1063,12 @@ inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, in
   }
 }
 
-// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block,
-// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s.
-void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) {
-  for (int b0 = 0; b0 < K; b0 += 128) {
-    const int b1 = std::min(K, b0 + 128);
+// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's
+// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448));
+// y = round_e4m3(clamp(x/s,+-448)) * s.
+void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) {
+  for (int b0 = 0; b0 < K; b0 += block) {
+    const int b1 = std::min(K, b0 + block);
     float amax = 1e-4f;
     for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i])));
     const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f)));
@@ -1275,6 +1276,8 @@ struct CpuMoeExecutor {
   // it to a captured GPU elementwise kernel removes it while keeping the official
   // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py).
   bool input_prequant = false;
+  // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block.
+  int act_block = 128;
   // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I).
   int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0;
   float e2m1_lut[16];
@@ -1817,7 +1820,7 @@ struct CpuMoeExecutor {
                  gas_scratch.data() + (size_t)r * (I / 32));
       return;
     }
-    if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I);
+    if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block);
     float* ge = ge_scratch.data() + (size_t)r * (I / 2);
     float* go = go_scratch.data() + (size_t)r * (I / 2);
     deinterleave_bf16_f32(g, ge, go, I);
@@ -1957,7 +1960,7 @@ struct CpuMoeExecutor {
         const bf16_t* src = t->x + (size_t)tok * H;
         if (ds) {  // DSV4 FP8-round-trips the input before the gate_up GEMV
           bf16_t* xq = xq_scratch.data() + (size_t)tok * H;
-          fp8_roundtrip_bf16(src, xq, H);
+          fp8_roundtrip_bf16(src, xq, H, act_block);
           src = xq;
         }
         float* xe = xe_scratch.data() + (size_t)tok * (H / 2);
@@ -2185,6 +2188,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
       .def("set_input_prequant",
            [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; },
            py::arg("value"))
+      .def("set_act_block",
+           [](CpuMoeExecutor& e, int v) { e.act_block = v; },
+           py::arg("value"))
       .def("isa_name", &CpuMoeExecutor::isa_name)
       .def("set_spin_wait",
            [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); },
diff --git a/python/freetoken/kernel/triton/dsv4/fp8_linear.py b/python/freetoken/kernel/triton/dsv4/fp8_linear.py
index 47e3c9f58a..fb17a35175 100644
--- a/python/freetoken/kernel/triton/dsv4/fp8_linear.py
+++ b/python/freetoken/kernel/triton/dsv4/fp8_linear.py
@@ -1,23 +1,25 @@
-"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4, matching the reference numerics.
+"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4 / V4.1, matching the reference numerics.
 
 The reference (``inference/model.py`` ``linear`` + ``inference/kernel.py``
-``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-128 block
+``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-``block``
 power-of-two (ue8m0) scale, then runs an FP8xFP8 block-scaled GEMM against the FP8
-weight (which carries its own 128x128 ue8m0 block scale). Both operands' scales are
-applied per 128-K block to a separate FP32 accumulator. This module reproduces that:
+weight (which carries its own ``block x block`` ue8m0 scale). Both operands' scales are
+applied per ``block``-K slab to a separate FP32 accumulator. This module reproduces that:
 
-  ``y = fp8_gemm(act_quant(x, 128, ue8m0), weight_fp8, weight_scale_e8m0)``
+  ``y = fp8_gemm(act_quant(x, block, ue8m0), weight_fp8, weight_scale_e8m0)``
+
+``block`` is the checkpoint's ``weight_block_size``: 128 for DeepSeek-V4, 32 for V4.1.
 
 ``act_quant`` (reference): per block ``s = 2**ceil(log2(max(|x|,1e-4)/448))`` (exact
 via IEEE bit ops -> matches ``fast_round_scale``), ``x_fp8 = round_e4m3(clamp(x/s,
 +-448))``, scale stored e8m0. The GEMM accumulates ``sum_k (A_fp8 @ B_fp8) * s_a * s_b``
-per 128-K block in FP32.
+per ``block``-K slab in FP32.
 
 Also provides ``act_quant_fp8_inplace`` -- the fused FP8 quant+dequant round-trip the
 reference applies in-place to the window / compressor KV (``act_quant(..., 64, ...,
 inplace=True)``), returning BF16.
 
-Assumes ``K % 128 == 0`` and ``N % 128 == 0`` (true for every DeepSeek-V4 projection).
+Assumes ``K % block == 0`` and ``N % block == 0`` (true for every DeepSeek-V4 / V4.1 projection).
 """
 
 from __future__ import annotations
@@ -218,8 +220,8 @@ def fp4_act_quant_inplace(x: torch.Tensor, block: int = 32) -> torch.Tensor:
 def _fp8_act_gemm_kernel(
     a_ptr,            # [M, K] float8_e4m3fn (quantized activation)
     w_ptr,            # [N, K] float8_e4m3fn
-    sa_ptr,           # [M, K//128] uint8 (e8m0 act codes)
-    sb_ptr,           # [N//128, K//128] uint8 (e8m0 weight codes)
+    sa_ptr,           # [M, K//BLOCK_K] uint8 (e8m0 act codes)
+    sb_ptr,           # [N//BLOCK_K, K//BLOCK_K] uint8 (e8m0 weight codes; square BLOCK_K blocks)
     c_ptr,            # [M, N] compute dtype
     M, N, K,
     stride_am, stride_ak, stride_wn, stride_wk,
@@ -228,6 +230,11 @@ def _fp8_act_gemm_kernel(
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
     compute_type: tl.constexpr,
 ):
+    """``BLOCK_K`` is the weight's square block edge, so each K step maps to exactly one scale
+    column. When the N tile IS one scale row (``BLOCK_N == BLOCK_K``, the DeepSeek-V4 block-128
+    case) the weight scale is a scalar per step and the N masks vanish -- the original kernel;
+    otherwise (block 32) the tile spans ``BLOCK_N // BLOCK_K`` scale rows and N is masked."""
+    SCALAR_SCALE: tl.constexpr = BLOCK_N == BLOCK_K
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
@@ -238,112 +245,217 @@ def _fp8_act_gemm_kernel(
     w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     num_k = tl.cdiv(K, BLOCK_K)
-    for k in range(num_k):
-        a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
-        w = tl.load(w_ptrs)
-        if e4m3_native_cx():
-            p = tl.dot(a, tl.trans(w), out_dtype=tl.float32)
-        else:
-            # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc
-            p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32)
-        sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0)
-        sca = tl.exp2(sa_code.to(tl.float32) - 127.0)            # [BLOCK_M]
-        sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk)
-        scb = tl.exp2(sb_code.to(tl.float32) - 127.0)            # scalar (one 128-N block)
-        acc += p * sca[:, None] * scb
-        a_ptrs += BLOCK_K * stride_ak
-        w_ptrs += BLOCK_K * stride_wk
-    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
-    tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None])
+    if SCALAR_SCALE:
+        for k in range(num_k):
+            a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
+            w = tl.load(w_ptrs)
+            if e4m3_native_cx():
+                p = tl.dot(a, tl.trans(w), out_dtype=tl.float32)
+            else:
+                # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc
+                p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32)
+            sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0)
+            sca = tl.exp2(sa_code.to(tl.float32) - 127.0)            # [BLOCK_M]
+            sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk)
+            scb = tl.exp2(sb_code.to(tl.float32) - 127.0)            # scalar (one N block)
+            acc += p * sca[:, None] * scb
+            a_ptrs += BLOCK_K * stride_ak
+            w_ptrs += BLOCK_K * stride_wk
+        c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+        tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None])
+    else:
+        n_mask = offs_n < N  # N is a multiple of BLOCK_K, not necessarily of BLOCK_N
+        sn = offs_n // BLOCK_K
+        for k in range(num_k):
+            a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
+            w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
+            if e4m3_native_cx():
+                p = tl.dot(a, tl.trans(w), out_dtype=tl.float32)
+            else:
+                p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32)
+            sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0)
+            sca = tl.exp2(sa_code.to(tl.float32) - 127.0)            # [BLOCK_M]
+            sb_code = tl.load(sb_ptr + sn * stride_sbn + k * stride_sbk, mask=n_mask, other=0)
+            scb = tl.exp2(sb_code.to(tl.float32) - 127.0)            # [BLOCK_N]
+            acc += p * sca[:, None] * scb[None, :]
+            a_ptrs += BLOCK_K * stride_ak
+            w_ptrs += BLOCK_K * stride_wk
+        c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+        tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None] & n_mask[None, :])
 
 
 @triton.jit
-def _fp8_act_gemv_splitk_kernel(
-    a_ptr,            # [K] float8_e4m3fn
-    sa_ptr,           # [K//128] uint8 (e8m0 act codes)
-    w_ptr,            # [N, K] float8_e4m3fn
-    sb_ptr,           # [N//128, K//128] uint8 (e8m0 weight codes)
-    part_ptr,         # [SPLIT_K, N] fp32
-    N, K,
-    stride_ak, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_pk, stride_pn,
-    BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
+def _fp8_gemv_kernel(
+    x_ptr,            # [M, G, K] activation in the compute dtype (bf16); QUANT_ACT: quantized in-kernel per SLAB
+    w_ptr,            # [G * N, K] float8_e4m3fn (or its uint8 view on pre-Ada parts)
+    sb_ptr,           # [G * N // SLAB, K // SLAB] uint8 (e8m0 weight codes)
+    out_ptr,          # [M, G * N] OUT dtype (SPLIT_K == 1)
+    part_ptr,         # [SPLIT_K, M, G * N] fp32 partials (SPLIT_K > 1)
+    M, N, K,
+    stride_xm, stride_xg, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_om, stride_pk, stride_pm,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SLAB: tl.constexpr, SPLIT_K: tl.constexpr,
+    QUANT_ACT: tl.constexpr, OUT: tl.constexpr,
 ):
+    """Small-M (decode) block-scaled fp8 GEMV/GEMM: one pass over the weight at near HBM bandwidth.
+
+    Each program owns ``BLOCK_N`` rows and ``K // SPLIT_K`` contiguous K, walked ``BLOCK_K`` (several
+    ``SLAB``-wide scale slabs) at a time: a wide tile keeps enough bytes in flight per SM to saturate
+    the memory system, which a slab-at-a-time walk (512 B per load at the V4.1 slab of 32) cannot.
+    Every weight tile is applied to all ``M`` (<= ``BLOCK_M``) activation rows, so a decode batch
+    costs one weight read like a single token does (the tensor-core GEMM path is 3-6x slower at
+    M = 2..8: too few tiles, tiny K steps). ``QUANT_ACT`` quantizes the activation here rather than
+    by a separate ``act_quant`` launch: every program recomputes the ue8m0 scale and the e4m3 rounding
+    of the slabs it touches (the reference formula, bit-identical to ``act_quant_fp8``) from the
+    L2-resident bf16 rows, so the path is one kernel (plus a small reduce when ``SPLIT_K > 1``).
+    Without it the activation stays bf16 (W8A16: the reference's bf16 einsum over a dequantized
+    weight, e.g. ``wo_a``). The third grid axis is an independent group (block-diagonal weights:
+    group ``g`` reads ``x[:, g]`` and rows ``[g * N, (g + 1) * N)``)."""
+    tl.static_assert(BLOCK_K % SLAB == 0, "the K tile holds whole scale slabs")
+    NSLAB: tl.constexpr = BLOCK_K // SLAB
     pid_n = tl.program_id(0)
     pid_k = tl.program_id(1)
+    g = tl.program_id(2).to(tl.int64)
+    x_ptr += g * stride_xg
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     n_mask = offs_n < N
-    sn = offs_n // 128
+    sn = (g * N + offs_n) // SLAB
+    offs_n = g * N + offs_n
+    rows = tl.arange(0, BLOCK_M)
     k_per = K // SPLIT_K
     k_start = pid_k * k_per
-    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
-    for k0 in range(0, k_per, 128):
-        offs_k = k_start + k0 + tl.arange(0, 128)
-        a = tl.load(a_ptr + offs_k * stride_ak).to(tl.float32)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, k_per, BLOCK_K):
+        offs_k = k_start + k0 + tl.arange(0, BLOCK_K)
+        # --- the weight tile, dequantized per slab (read once for every activation row) ---
         w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
         if e4m3_native_cx():
             w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32)
         else:
             w = e4m3_u8_to_f32(tl.load(w_ptrs, mask=n_mask[:, None], other=0))
-        kb = (k_start + k0) // 128
-        sb_code = tl.load(sb_ptr + sn * stride_sbn + kb * stride_sbk, mask=n_mask, other=0)
-        scb = tl.exp2(sb_code.to(tl.float32) - 127.0)
-        sa_code = tl.load(sa_ptr + kb)
-        sca = tl.exp2(sa_code.to(tl.float32) - 127.0)
-        acc += tl.sum(w * a[None, :], axis=1) * scb * sca
-    tl.store(part_ptr + pid_k * stride_pk + offs_n * stride_pn, acc, mask=n_mask)
+        kb = (k_start + k0) // SLAB + tl.arange(0, NSLAB)
+        sb_code = tl.load(sb_ptr + sn[:, None] * stride_sbn + kb[None, :] * stride_sbk, mask=n_mask[:, None], other=0)
+        scb = tl.exp2(sb_code.to(tl.float32) - 127.0)  # [BLOCK_N, NSLAB]
+        for m in tl.static_range(BLOCK_M):
+            if m < M:
+                xs = tl.reshape(tl.load(x_ptr + m * stride_xm + offs_k).to(tl.float32), (NSLAB, SLAB))
+                if QUANT_ACT:
+                    # --- reference act_quant on this row's slabs: s = 2**ceil(log2(max(|x|, 1e-4) / 448)) ---
+                    amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4)
+                    sa = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32))  # [NSLAB]
+                    q = tl.clamp(xs / sa[:, None], -448.0, 448.0)
+                    if e4m3_native_cx():
+                        q = q.to(tl.float8e4nv).to(tl.float32)
+                    else:
+                        q = round_e4m3(q).to(tl.float32)
+                else:
+                    sa = tl.full((NSLAB,), 1.0, tl.float32)
+                    q = xs
+                a = tl.reshape(q, (BLOCK_K,))
+                slab = tl.sum(tl.reshape(w * a[None, :], (BLOCK_N, NSLAB, SLAB)), axis=2)  # [BLOCK_N, NSLAB]
+                contrib = tl.sum(slab * scb * sa[None, :], axis=1)  # [BLOCK_N]
+                acc = tl.where(rows[:, None] == m, acc + contrib[None, :], acc)
+    row_mask = rows < M
+    if SPLIT_K == 1:
+        tl.store(out_ptr + rows[:, None] * stride_om + offs_n[None, :], acc.to(OUT), mask=row_mask[:, None] & n_mask[None, :])
+    else:
+        tl.store(part_ptr + pid_k * stride_pk + rows[:, None] * stride_pm + offs_n[None, :], acc, mask=row_mask[:, None] & n_mask[None, :])
 
 
 @triton.jit
-def _splitk_reduce_kernel(part_ptr, out_ptr, N, SPLIT_K: tl.constexpr,
-                          stride_pk, stride_pn, BLOCK: tl.constexpr, OUT: tl.constexpr):
+def _splitk_reduce_kernel(part_ptr, out_ptr, M, N, SPLIT_K: tl.constexpr,
+                          stride_pk, stride_pm, stride_om, BLOCK: tl.constexpr, OUT: tl.constexpr):
+    m = tl.program_id(1)
     offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
     mask = offs < N
     acc = tl.zeros((BLOCK,), dtype=tl.float32)
     for k in tl.static_range(SPLIT_K):
-        acc += tl.load(part_ptr + k * stride_pk + offs * stride_pn, mask=mask, other=0.0)
-    tl.store(out_ptr + offs, acc.to(OUT), mask=mask)
-
-
-# Swept-best decode GEMV config per (N, K) -> (BLOCK_N, SPLIT_K, num_warps).
-_DECODE_FP8_CFG = {
-    (1024, 4096): (16, 16, 1),   # wq_a
-    (32768, 1024): (16, 1, 2),   # wq_b
-    (512, 4096): (16, 16, 1),    # attn wkv
-    (4096, 8192): (16, 8, 1),    # wo_b
-    (2048, 4096): (16, 8, 1),    # shared w1 / w3
-    (4096, 2048): (16, 8, 1),    # shared w2
-    (8192, 1024): (16, 2, 4),    # indexer wq_b
-}
-
-
-def _decode_cfg(N: int, K: int) -> tuple[int, int, int]:
-    cfg = _DECODE_FP8_CFG.get((N, K))
-    if cfg is not None:
-        bn, sk, nw = cfg
-        return bn, max(1, min(sk, K // 128)), nw
-    bn = 16
-    n_tiles = triton.cdiv(N, bn)
-    sk = max(1, 1536 // n_tiles)
-    sk = 1 << (sk.bit_length() - 1)
-    return bn, max(1, min(sk, K // 128)), 1
-
-
-def _fp8_act_gemv(a_fp8: torch.Tensor, sa: torch.Tensor, weight: torch.Tensor,
-                  sb: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
-    N, K = weight.shape
-    BLOCK_N, split_k, num_warps = _decode_cfg(N, K)
-    n_tiles = triton.cdiv(N, BLOCK_N)
-    part = torch.empty((split_k, N), dtype=torch.float32, device=a_fp8.device)
-    _fp8_act_gemv_splitk_kernel[(n_tiles, split_k)](
-        a_fp8, sa, weight, sb, part, N, K,
-        a_fp8.stride(0), weight.stride(0), weight.stride(1),
-        sb.stride(0), sb.stride(1), part.stride(0), part.stride(1),
-        BLOCK_N=BLOCK_N, SPLIT_K=split_k, num_warps=num_warps,
+        acc += tl.load(part_ptr + k * stride_pk + m * stride_pm + offs, mask=mask, other=0.0)
+    tl.store(out_ptr + m * stride_om + offs, acc.to(OUT), mask=mask)
+
+
+_GEMV_BLOCK_N = 32
+_GEMV_NARROW_BLOCK_N = 8
+def _gemv_cfg(N: int, K: int, slab: int, sm_count: int) -> tuple[int, int, int, int]:
+    """Choose whole scale tiles and enough independent programs to cover the device."""
+    block_k = next((bk for bk in (512, 256, 128) if bk >= slab and K % bk == 0), slab)
+    if triton.cdiv(N, _GEMV_BLOCK_N) >= sm_count:
+        return _GEMV_BLOCK_N, block_k, 1, 4
+    n_tiles = triton.cdiv(N, _GEMV_NARROW_BLOCK_N)
+    split = 1
+    while n_tiles * split < sm_count and (K // block_k) % (split * 2) == 0:
+        split *= 2
+    return _GEMV_NARROW_BLOCK_N, block_k, split, 4
+
+
+GEMV_MAX_M = 4  # bound the unrolled row loop's register footprint
+
+
+def _fp8_gemv(
+    x: torch.Tensor, weight: torch.Tensor, sb: torch.Tensor, out_dtype: torch.dtype, block: int, *, groups: int = 1, quant_act: bool = True,
+) -> torch.Tensor:
+    """``x [M, groups, K]`` bf16 x block-diagonal ``weight [groups * N, K]`` fp8 -> ``[M, groups * N]``
+    in ``out_dtype`` (``M <= GEMV_MAX_M``). ``quant_act`` applies the reference fp8 activation quant
+    (W8A8); off, the activation stays bf16 (W8A16). ``groups == 1`` is the plain GEMV."""
+    GN, K = weight.shape
+    N = GN // groups
+    x = x.reshape(-1, groups, K)
+    M = x.shape[0]
+    assert 1 <= M <= GEMV_MAX_M, M
+    block_m = triton.next_power_of_2(M)
+    block_n, block_k, split_k, num_warps = _gemv_cfg(GN, K, block, torch.cuda.get_device_properties(x.device).multi_processor_count)
+    out = torch.empty((M, GN), dtype=out_dtype, device=x.device)
+    part = out if split_k == 1 else torch.empty((split_k, M, GN), dtype=torch.float32, device=x.device)
+    _fp8_gemv_kernel[(triton.cdiv(N, block_n), split_k, groups)](
+        x, weight, sb, out, part, M, N, K,
+        x.stride(0), x.stride(1), weight.stride(0), weight.stride(1), sb.stride(0), sb.stride(1), out.stride(0),
+        part.stride(0) if split_k > 1 else 0, part.stride(1) if split_k > 1 else 0,
+        BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, SLAB=block, SPLIT_K=split_k, QUANT_ACT=quant_act, OUT=_TL_DTYPE[out_dtype],
+        num_warps=num_warps,
     )
-    out = torch.empty(N, dtype=out_dtype, device=a_fp8.device)
-    _splitk_reduce_kernel[(triton.cdiv(N, 256),)](
-        part, out, N, split_k, part.stride(0), part.stride(1),
-        BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2,
+    if split_k > 1:
+        _splitk_reduce_kernel[(triton.cdiv(GN, 256), M)](
+            part, out, M, GN, split_k, part.stride(0), part.stride(1), out.stride(0),
+            BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2,
+        )
+    return out
+
+
+def grouped_w8a16_gemv(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor:
+    """Decode form of the reference's grouped bf16 einsum ``tgd,grd->tgr`` over a block-diagonal fp8
+    weight: ``x [t, G, d]`` bf16 (``t <= GEMV_MAX_M``), ``weight [G * r, d]`` fp8 with ``[G * r //
+    block, d // block]`` e8m0 ``scale`` -> ``[t, G * r]`` bf16. The weight is dequantized in-kernel
+    (exact in bf16: e4m3 x 2^k), the activation is not quantized, so this is the einsum's math with a
+    different fp32 summation order."""
+    assert x.ndim == 3 and weight.dtype == FP8
+    G = x.shape[1]
+    sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous()
+    return _fp8_gemv(x.contiguous(), e4m3_kernel_view(weight), sb, x.dtype, block, groups=G, quant_act=False)
+
+
+@triton.jit
+def _dequant_block_fp8_kernel(w_ptr, sb_ptr, out_ptr, N, K, stride_wn, stride_sbn, stride_on, SLAB: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
+    mask = (offs_n < N)[:, None] & (offs_k < K)[None, :]
+    if e4m3_native_cx():
+        w = tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0.0).to(tl.float32)
+    else:
+        w = e4m3_u8_to_f32(tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0))
+    code = tl.load(sb_ptr + (offs_n // SLAB)[:, None] * stride_sbn + (offs_k // SLAB)[None, :], mask=mask, other=0)
+    tl.store(out_ptr + offs_n[:, None] * stride_on + offs_k[None, :], (w * tl.exp2(code.to(tl.float32) - 127.0)).to(tl.bfloat16), mask=mask)
+
+
+def dequant_block_fp8(weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor:
+    """``[N, K]`` block-scaled e4m3 -> bf16 (exact), one pass; the prefill-side companion of
+    :func:`grouped_w8a16_gemv` for weights the reference consumes through a bf16 einsum."""
+    N, K = weight.shape
+    sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous()
+    out = torch.empty((N, K), dtype=torch.bfloat16, device=weight.device)
+    _dequant_block_fp8_kernel[(triton.cdiv(N, 64), triton.cdiv(K, 256))](
+        e4m3_kernel_view(weight), sb, out, N, K, weight.stride(0), sb.stride(0), out.stride(0),
+        SLAB=block, BLOCK_N=64, BLOCK_K=256, num_warps=4,
     )
     return out
 
@@ -353,36 +465,41 @@ def block_fp8_linear(
     weight: torch.Tensor,
     scale: torch.Tensor,
     bias: torch.Tensor | None = None,
+    *,
+    block: int = 128,
 ) -> torch.Tensor:
     """``y = act_quant(x) @ weight^T`` (reference FP8 path).
 
     ``x``: ``[..., K]`` bf16; ``weight``: ``[N, K]`` float8_e4m3fn; ``scale``:
-    ``[N//128, K//128]`` float8_e8m0fnu (weight block scale). Activation is quantized
-    to FP8 with a per-128 ue8m0 scale; the GEMM applies both scales per 128-K block.
+    ``[N//block, K//block]`` float8_e8m0fnu (weight block scale). Activation is quantized
+    to FP8 with a per-``block`` ue8m0 scale; the GEMM applies both scales per ``block``-K slab.
     """
     assert weight.dtype == FP8
     *lead, K = x.shape
     N = weight.shape[0]
     assert weight.shape[1] == K
-    assert K % 128 == 0 and N % 128 == 0, (N, K)
+    assert K % block == 0 and N % block == 0, (N, K, block)
+    assert tuple(scale.shape) == (N // block, K // block), (tuple(scale.shape), N, K, block)
     compute_dtype = x.dtype if x.dtype in _TL_DTYPE else torch.bfloat16
     sb = scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale
     sb = sb.contiguous()
     w = e4m3_kernel_view(weight)
 
-    a_fp8, sa = act_quant_fp8(x, 128)  # [M,K] fp8, [M,K//128] e8m0 codes
-    M = a_fp8.shape[0]
-
-    if M == 1:
-        out = _fp8_act_gemv(a_fp8[0], sa[0], w, sb, compute_dtype).reshape(*lead, N)
+    x2d = x.reshape(-1, K)
+    if x2d.shape[0] == 1:
+        # decode: one kernel quantizes the activation rows and streams the weight once
+        out = _fp8_gemv(x2d.contiguous(), w, sb, compute_dtype, block).reshape(*lead, N)
         if bias is not None:
             out = out + bias.to(out.dtype)
         return out
 
+    a_fp8, sa = act_quant_fp8(x, block)  # [M,K] fp8, [M,K//block] e8m0 codes
+    M = a_fp8.shape[0]
+
     out = torch.empty((M, N), dtype=compute_dtype, device=x.device)
     BLOCK_M = 32
     BLOCK_N = 128
-    BLOCK_K = 128
+    BLOCK_K = block
     grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
     _fp8_act_gemm_kernel[grid](
         a_fp8, w, sa, sb, out,
@@ -399,4 +516,4 @@ def block_fp8_linear(
     return out
 
 
-__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace"]
+__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace", "grouped_w8a16_gemv", "dequant_block_fp8", "GEMV_MAX_M"]
diff --git a/python/freetoken/kernel/triton/dsv4/fused_moe.py b/python/freetoken/kernel/triton/dsv4/fused_moe.py
index 6eea606c36..0adfbdd34b 100644
--- a/python/freetoken/kernel/triton/dsv4/fused_moe.py
+++ b/python/freetoken/kernel/triton/dsv4/fused_moe.py
@@ -24,6 +24,9 @@
 import triton
 import triton.language as tl
 
+from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil
+from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, round_e4m3
+
 _E2M1_VALUES = [
     0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
     -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
@@ -71,6 +74,12 @@ def _decode_dsfp4_moe_kernel(
     n_mask = offs_n < N
 
     slot = tl.load(topk_ids_ptr + token_id * stride_tid_m + route_k * stride_tid_k).to(tl.int64)
+    c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn
+    if slot < 0:
+        # an inactive route (hybrid decode computes it on the CPU): zero contribution, and no read of
+        # any slot -- a slot's storage may be uninitialized, and 0 * NaN is not 0
+        tl.store(c_ptrs, tl.zeros((BLOCK_SIZE_N,), dtype=compute_type), mask=(route_id < total_routes) & n_mask)
+        return
     a_row = route_id if A_ROW_IS_ROUTE else token_id
     a_base = a_ptr + a_row * stride_am
 
@@ -113,7 +122,6 @@ def _decode_dsfp4_moe_kernel(
         weight = tl.load(topk_weights_ptr + token_id * stride_tw_m + route_k * stride_tw_k)
         accumulator = accumulator * weight
 
-    c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn
     tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask)
 
 
@@ -124,11 +132,15 @@ def _swiglu_kernel(
     R, I, limit,
     stride_gr, stride_gi, stride_or, stride_oi,
     BLOCK: tl.constexpr, HAS_LIMIT: tl.constexpr, compute_type: tl.constexpr,
+    ACT_BLOCK: tl.constexpr,  # > 0: fp8 (e4m3, ue8m0 per ACT_BLOCK) round-trip of the stored value, fused
 ):
     """Fused SwiGLU: ``out = silu(min(gate, limit)) * clamp(up, -limit, limit)``.
 
     Collapses the gate/up split + two clamps + silu + mul + the fp32 round-trip
-    (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass."""
+    (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass. With
+    ``ACT_BLOCK`` the down projection's activation quant (``act_quant_fp8_inplace`` on the bf16
+    value: ``s = 2**ceil(log2(max(|x|, 1e-4) / 448))``, e4m3 rounding, dequant) is applied to the
+    same tile before the store -- one launch and one memory pass fewer, bit-identical."""
     row = tl.program_id(0)
     cb = tl.program_id(1)
     offs = cb * BLOCK + tl.arange(0, BLOCK)
@@ -139,6 +151,17 @@ def _swiglu_kernel(
         g = tl.minimum(g, limit)
         u = tl.minimum(tl.maximum(u, -limit), limit)
     act = (g * tl.sigmoid(g)) * u
+    if ACT_BLOCK > 0:
+        tl.static_assert(BLOCK % ACT_BLOCK == 0, "the tile holds whole quant blocks")
+        xs = tl.reshape(act.to(compute_type).to(tl.float32), (BLOCK // ACT_BLOCK, ACT_BLOCK))  # the value the store would hold
+        amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4)
+        sc = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32))
+        q = tl.clamp(xs / sc[:, None], -448.0, 448.0)
+        if e4m3_native_cx():
+            q = q.to(tl.float8e4nv).to(tl.float32)
+        else:
+            q = round_e4m3(q)
+        act = tl.reshape(q * sc[:, None], (BLOCK,))
     tl.store(out_ptr + row * stride_or + offs * stride_oi, act.to(compute_type), mask=mask)
 
 
@@ -274,20 +297,22 @@ def _compute_type(dtype: torch.dtype):
     return {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32}[dtype]
 
 
-def fused_swiglu(gate_up: torch.Tensor, limit: float) -> torch.Tensor:
-    """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``)."""
+def fused_swiglu(gate_up: torch.Tensor, limit: float, *, act_block: int = 0) -> torch.Tensor:
+    """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``). ``act_block > 0``
+    also applies the down projection's fp8 activation round-trip (block ``act_block``) in the same pass."""
     *lead, two_I = gate_up.shape
     I = two_I // 2
     gu = gate_up.reshape(-1, two_I)
     R = gu.shape[0]
     out = torch.empty((R, I), dtype=gate_up.dtype, device=gate_up.device)
     BLOCK = 1024
+    assert act_block == 0 or (I % act_block == 0 and BLOCK % act_block == 0), (I, act_block)
     grid = (R, triton.cdiv(I, BLOCK))
     _swiglu_kernel[grid](
         gu, out, R, I, float(limit),
         gu.stride(0), gu.stride(1), out.stride(0), out.stride(1),
         BLOCK=BLOCK, HAS_LIMIT=limit > 0,
-        compute_type=_compute_type(gate_up.dtype), num_warps=4,
+        compute_type=_compute_type(gate_up.dtype), ACT_BLOCK=act_block, num_warps=4,
     )
     return out.reshape(*lead, I)
 
diff --git a/python/freetoken/kernel/triton/dsv4/router.py b/python/freetoken/kernel/triton/dsv4/router.py
new file mode 100644
index 0000000000..17e3506a19
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv4/router.py
@@ -0,0 +1,83 @@
+"""Fused MoE gate selection (DeepSeek-V4 / V4.1 ``Gate``): score activation, selection bias, top-k,
+weight normalization and route scale in one program per token.
+
+    s       = act(scores)                       act: softmax | sigmoid | sqrt(softplus)
+    ids     = topk(s + bias)                    selection is biased, the weights are not
+    w       = s[ids] [/ (sum + 1e-20)] * route_scale
+
+One launch replaces the eight or so elementwise / top-k / gather launches of the torch expression --
+at decode each of those moves a few hundred floats and costs its launch floor. Ties in ``s + bias``
+go to the lowest expert id (torch.topk leaves them unspecified).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+_SCORE = {"softmax": 0, "sigmoid": 1, "sqrtsoftplus": 2}
+
+
+@triton.jit
+def _router_select_kernel(
+    scores_ptr, bias_ptr, alternate_bias_ptr, alternate_mask_ptr, w_ptr, id_ptr, E, route_scale,
+    stride_s, stride_w, stride_id,
+    SCORE: tl.constexpr, TOPK: tl.constexpr, BLOCK_K: tl.constexpr, NORMALIZE: tl.constexpr, BLOCK_E: tl.constexpr, ALTERNATE: tl.constexpr,
+):
+    t = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK_E)
+    mask = offs < E
+    s = tl.load(scores_ptr + t * stride_s + offs, mask=mask, other=float("-inf")).to(tl.float32)
+    if SCORE == 0:
+        s = tl.exp(s - tl.max(s))
+        s = s / tl.sum(tl.where(mask, s, 0.0))
+    elif SCORE == 1:
+        s = tl.sigmoid(s)
+    else:
+        s = tl.sqrt(tl.where(s > 20.0, s, libdevice.log1p(tl.exp(s))))  # torch softplus (threshold 20)
+    if ALTERNATE:
+        bias_ptr = tl.where(tl.load(alternate_mask_ptr + t), alternate_bias_ptr, bias_ptr)
+    bias = tl.load(bias_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    key = tl.where(mask, s + bias, float("-inf"))
+    slots = tl.arange(0, BLOCK_K)
+    picked_w = tl.zeros((BLOCK_K,), dtype=tl.float32)
+    picked_id = tl.zeros((BLOCK_K,), dtype=tl.int32)
+    for k in tl.static_range(TOPK):
+        best = tl.max(key)
+        idx = tl.min(tl.where(key == best, offs, BLOCK_E))  # lowest id among ties
+        w = tl.sum(tl.where(offs == idx, s, 0.0))
+        picked_w = tl.where(slots == k, w, picked_w)
+        picked_id = tl.where(slots == k, idx.to(tl.int32), picked_id)
+        key = tl.where(offs == idx, float("-inf"), key)
+    if NORMALIZE:
+        picked_w = picked_w / (tl.sum(picked_w) + 1e-20)
+    picked_w = picked_w * route_scale
+    tl.store(w_ptr + t * stride_w + slots, picked_w, mask=slots < TOPK)
+    tl.store(id_ptr + t * stride_id + slots, picked_id, mask=slots < TOPK)
+
+
+def router_select(
+    scores: torch.Tensor, bias: torch.Tensor, topk: int, *, score_func: str, normalize: bool, route_scale: float,
+    alternate_bias: torch.Tensor | None = None, alternate_mask: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``scores [T, E]`` fp32 (pre-activation gate logits), ``bias [E]`` -> ``(weights [T, topk] fp32,
+    ids [T, topk] int32)`` in selection order (best first)."""
+    T, E = scores.shape
+    assert scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.dtype, scores.stride())
+    weights = torch.empty((T, topk), dtype=torch.float32, device=scores.device)
+    ids = torch.empty((T, topk), dtype=torch.int32, device=scores.device)
+    if T == 0:
+        return weights, ids
+    _router_select_kernel[(T,)](
+        scores, bias, alternate_bias if alternate_bias is not None else bias,
+        alternate_mask if alternate_mask is not None else scores, weights, ids, E, float(route_scale),
+        scores.stride(0), weights.stride(0), ids.stride(0),
+        SCORE=_SCORE[score_func], TOPK=topk, BLOCK_K=triton.next_power_of_2(topk), NORMALIZE=normalize and topk > 1,
+        BLOCK_E=triton.next_power_of_2(E), ALTERNATE=alternate_mask is not None, num_warps=4,
+    )
+    return weights, ids
+
+
+__all__ = ["router_select"]
diff --git a/python/freetoken/kernel/triton/dsv41/__init__.py b/python/freetoken/kernel/triton/dsv41/__init__.py
new file mode 100644
index 0000000000..00ec72daba
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/__init__.py
@@ -0,0 +1,9 @@
+"""Triton kernels for DeepSeek-V4.1 Compressed Sparse Attention 2 over packed KV pools.
+
+``row_format``  device helpers: load a tile of packed rows (bf16 / fp8+e8m0 / fp4+e4m3 / fp4+e8m0)
+                and dequantize it to fp32, encode / decode e2m1 nibbles.
+``pack``        host wrappers: quantize rows into a packed pool (``pack_rows``) and read them back
+                (``unpack_rows``), matching the reference ``act_quant`` / ``fp4_act_quant`` numerics.
+``sparse_attn`` paged sparse attention gathering window (fp8) and compressed (fp4) rows.
+``indexer``     Lightning-indexer logits over fp4 index keys, full-range or candidate-restricted.
+"""
diff --git a/python/freetoken/kernel/triton/dsv41/indexer.py b/python/freetoken/kernel/triton/dsv41/indexer.py
new file mode 100644
index 0000000000..3721315837
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/indexer.py
@@ -0,0 +1,176 @@
+"""Lightning-Indexer logits over PACKED index keys (DeepSeek-V4.1, Full and Reindex modes).
+
+    logits[b, s, j] = sum_h relu(q[b, s, h, :] . k(b, pos_j)) * weights[b, s, h]
+
+``q`` carries the ``index_n_heads`` fp4-baked index queries, ``k`` is the single shared index key of
+a compressed position, gathered from the kv-source's packed pool and dequantized in-kernel. The
+row of compressed position ``t`` is derived in-kernel from the request's full-token locs (the page
+table row at prefill, the decode snapshot): ``locs[b, t * ratio] // ratio`` (negative = no row) --
+no ``[B, T]`` row table is materialized per layer. ``weights`` folds ``softmax_scale *
+n_heads**-0.5`` in.
+
+Two position sets, one kernel:
+
+* **full range** (Full mode): position ``j`` is the compressed position ``j``; the output is
+  ``[B, S, T]`` over the staged width ``T``, but only the live tiles (columns below ``live[b, s]``
+  rounded up to ``BLOCK_T``, the rest of the last tile ``-inf``) are WRITTEN. A fixed grid of worker
+  programs strides over the live tiles, so both the work and the launch geometry follow the live
+  history rather than the staged width; consumers (``dsv41/topk.py``) read below ``live`` only, so
+  the dead area needs no fill.
+* **candidate pool** (Reindex mode, Hierarchical Sparse Indexer): position ``j`` is
+  ``candidates[b, s, j]`` (a compressed position, ``-1`` = empty slot); the output is ``[B, S, NC]``
+  aligned with the candidate list and fully written (``-inf`` for empty / unreachable slots), so
+  a top-k over it indexes straight back into the list.
+
+Causality lives in the kernel: ``live[b, s]`` is the number of compressed positions query ``s`` may
+see, and every position at or past it (or with no row) scores ``-inf``. The dot product, weighted
+scores and head reduction round to bf16 at the same boundaries as the reference expression.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.kvcache.dsv4.v41_row_format import RowFormat
+
+from .row_format import load_rows
+
+BLOCK_T = 64
+
+
+@triton.jit
+def _score_tile(
+    q, w, pool_ptr, locs_row, live, pos, store_mask, out_ptrs,
+    n_locs, RATIO: tl.constexpr, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr,
+):
+    """Score one tile of compressed positions ``pos`` against the query (``q [BLOCK_H, D]`` bf16,
+    ``w [BLOCK_H]`` fp32) and store ``[BLOCK_T]`` logits (``-inf`` where unreachable)."""
+    pos_ok = store_mask & (pos >= 0) & (pos < live)
+    full_idx = pos * RATIO
+    full = tl.load(locs_row + tl.maximum(full_idx, 0), mask=pos_ok & (full_idx < n_locs), other=-1)
+    rows = tl.where(full >= 0, full // RATIO, -1).to(tl.int32)
+    valid = pos_ok & (rows >= 0)
+    offs_d = tl.arange(0, D)
+    k = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT).to(tl.bfloat16)  # [BLOCK_T, D]
+    score = tl.dot(q, tl.trans(k)).to(tl.bfloat16).to(tl.float32)
+    score = (tl.maximum(score, 0.0) * w[:, None]).to(tl.bfloat16).to(tl.float32)
+    logits = tl.sum(score, axis=0).to(tl.bfloat16).to(tl.float32)
+    tl.store(out_ptrs, tl.where(valid, logits, float("-inf")), mask=store_mask)
+
+
+@triton.jit
+def _indexer_logits_packed_kernel(
+    q_ptr, w_ptr, pool_ptr, locs_ptr, live_ptr, cand_ptr, out_ptr,
+    T, NC, n_locs,
+    stride_qb, stride_qs, stride_qh, stride_qd,
+    stride_wb, stride_ws, stride_wh,
+    stride_lb, stride_lt,
+    stride_vb, stride_vs,
+    stride_cb, stride_cs, stride_cn,
+    stride_ob, stride_os, stride_on,
+    RATIO: tl.constexpr,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    ROW_BYTES: tl.constexpr,
+    FMT: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    HAS_CAND: tl.constexpr,
+    N_WORKERS: tl.constexpr,
+):
+    pid_s = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_t = tl.program_id(2)
+
+    live = tl.load(live_ptr + pid_b * stride_vb + pid_s * stride_vs)
+    locs_row = locs_ptr + pid_b * stride_lb
+    out_row = out_ptr + pid_b * stride_ob + pid_s * stride_os
+    offs_d = tl.arange(0, D)
+    offs_h = tl.arange(0, BLOCK_H)
+    h_mask = offs_h < H
+    q = tl.load(
+        q_ptr + pid_b * stride_qb + pid_s * stride_qs + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd,
+        mask=h_mask[:, None], other=0.0,
+    ).to(tl.bfloat16)
+    w = tl.load(w_ptr + pid_b * stride_wb + pid_s * stride_ws + offs_h * stride_wh, mask=h_mask, other=0.0).to(tl.bfloat16).to(tl.float32)
+    offs = tl.arange(0, BLOCK_T)
+
+    if HAS_CAND:
+        # one program per candidate tile; the list is fixed-width and fully written
+        offs_t = pid_t * BLOCK_T + offs
+        store_mask = offs_t < NC
+        pos = tl.load(cand_ptr + pid_b * stride_cb + pid_s * stride_cs + offs_t * stride_cn, mask=store_mask, other=-1)
+        _score_tile(q, w, pool_ptr, locs_row, live, pos, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT)
+    else:
+        # worker pid_t scores live tiles pid_t, pid_t + N_WORKERS, ...; nothing past live is written
+        for tile in range(pid_t, tl.cdiv(tl.minimum(live, T), BLOCK_T), N_WORKERS):
+            offs_t = tile * BLOCK_T + offs
+            store_mask = offs_t < T
+            _score_tile(q, w, pool_ptr, locs_row, live, offs_t, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT)
+
+
+def indexer_logits_packed(
+    q: torch.Tensor,           # [B, S, H, D] bf16 (fp4 round-tripped index queries)
+    weights: torch.Tensor,     # [B, S, H] (softmax_scale * n_heads**-0.5 folded in)
+    k_pool: torch.Tensor,      # [R, fmt.row_bytes(D)] uint8, the kv-source's packed index keys
+    fmt: RowFormat,
+    locs: torch.Tensor,        # [B, W] int32/int64 full-token locs of each request (page-table row / decode snapshot)
+    ratio: int,                # compressed position t -> full loc locs[b, t * ratio]; its row is loc // ratio
+    live: torch.Tensor,        # [B, S] int32: compressed positions query (b, s) may see
+    T: int | None = None,      # staged width of the full-range output (default: W // ratio)
+    candidates: torch.Tensor | None = None,  # [B, S, NC] int32 compressed positions (-1 = empty)
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Head-reduced indexer logits: ``[B, S, T]`` over the full staged range (columns below ``live``
+    written, the rest untouched), or ``[B, S, NC]`` aligned with ``candidates`` (fully written).
+    Unreachable / empty positions are ``-inf``."""
+    B, S, H, D = q.shape
+    assert locs.ndim == 2 and locs.shape[0] == B and live.shape == (B, S), (locs.shape, live.shape, (B, S))
+    assert weights.shape == (B, S, H), (weights.shape, (B, S, H))
+    assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}"
+    assert k_pool.dtype == torch.uint8 and k_pool.shape[1] == fmt.row_bytes(D) and k_pool.is_contiguous()
+    assert locs.dtype in (torch.int32, torch.int64) and locs.stride(1) == 1, (locs.dtype, locs.stride())
+    n_locs = locs.shape[1]
+    if T is None:
+        T = n_locs // ratio
+    live = live.to(torch.int32)
+    has_cand = candidates is not None
+    if has_cand:
+        assert candidates.shape[:2] == (B, S), (candidates.shape, (B, S))
+        cand = candidates.to(torch.int32)
+        NC = cand.shape[2]
+        cs = cand.stride()
+    else:
+        cand, NC, cs = live, 0, (0, 0, 0)
+    n_out = NC if has_cand else T
+    if out is None:
+        out = torch.empty((B, S, n_out), dtype=torch.float32, device=q.device)
+    assert out.shape == (B, S, n_out) and out.dtype == torch.float32
+    if n_out == 0 or S == 0:
+        return out
+
+    n_tiles = triton.cdiv(n_out, BLOCK_T)
+    # A device-sized worker grid strides over live history without launching a
+    # program for every reserved tile. Independent queries share the SM budget.
+    sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count
+    workers = n_tiles if has_cand else min(n_tiles, max(1, triton.cdiv(sm_count, B * S)))
+    grid = (S, B, workers)
+    _indexer_logits_packed_kernel[grid](
+        q, weights, k_pool, locs, live, cand, out,
+        T, NC, n_locs,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        weights.stride(0), weights.stride(1), weights.stride(2),
+        locs.stride(0), locs.stride(1),
+        live.stride(0), live.stride(1),
+        cs[0], cs[1], cs[2],
+        out.stride(0), out.stride(1), out.stride(2),
+        RATIO=ratio, H=H, D=D, ROW_BYTES=fmt.row_bytes(D), FMT=fmt.code,
+        BLOCK_H=triton.next_power_of_2(H), BLOCK_T=BLOCK_T, HAS_CAND=has_cand, N_WORKERS=workers,
+        num_warps=4, num_stages=2,
+    )
+    return out
+
+
+__all__ = ["indexer_logits_packed", "BLOCK_T"]
diff --git a/python/freetoken/kernel/triton/dsv41/pack.py b/python/freetoken/kernel/triton/dsv41/pack.py
new file mode 100644
index 0000000000..47a9164e6e
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/pack.py
@@ -0,0 +1,183 @@
+"""Quantize rows into a packed pool and read them back (DeepSeek-V4.1 KV storage formats).
+
+``pack_rows`` reproduces the reference quantizers (``inference/kernel.py``) exactly:
+
+* fp8_e8m0_b32  ``act_quant(x, 32, "ue8m0", e8m0)``: per 32-group ``s = 2**ceil(log2(max(amax, 1e-4) / 448))``,
+                ``q = round_e4m3(clamp(x / s, +-448))``
+* fp4_e4m3_b16  ``fp4_act_quant(x, 16, scale_dtype=e4m3)``: ``s = e4m3(max(amax, 6 * 2^-9) / 6)`` (round to
+                nearest e4m3), ``q = round_e2m1(clamp(x / s, +-6))``
+* fp4_e8m0_b32  ``fp4_act_quant(x, 32)``: ``s = 2**ceil(log2(max(amax, 6 * 2^-126) / 6))``, same e2m1 rounding
+* bf16          a plain byte copy
+
+Rows land at ``pool[row_ids[m]]`` (or ``pool[m]`` without ids), so the same launch serves the KV write
+path (scatter into slots) and plain buffer conversion. ``unpack_rows`` is the inverse gather + dequant
+(bf16-rounded like the reference cache); the attention / indexer kernels use the same device helper
+(``row_format.load_rows``) in-line instead of materializing rows.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil, _round_fp4
+from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8, e4m3_native_cx, round_e4m3
+from freetoken.kvcache.dsv4.v41_row_format import BF16, RowFormat
+
+from .row_format import fp4_f32_to_nibble, load_rows
+
+
+@triton.jit
+def _e4m3_codes(q):
+    """fp32 (|q| <= 448) -> e4m3 bit pattern, RNE."""
+    if e4m3_native_cx():
+        return q.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
+    else:
+        return e4m3_f32_to_u8(round_e4m3(q))
+
+
+@triton.jit
+def _e4m3_round(q):
+    """fp32 -> nearest e4m3 value, as fp32."""
+    if e4m3_native_cx():
+        return q.to(tl.float8e4nv).to(tl.float32)
+    else:
+        return round_e4m3(q)
+
+
+@triton.jit
+def _pack_rows_kernel(
+    x_ptr, out_ptr, idx_ptr, stride_xm,
+    D: tl.constexpr, GROUP: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, HAS_IDX: tl.constexpr,
+):
+    """One program per input row ``m``; writes packed bytes to pool row ``idx[m]`` (or ``m``)."""
+    m = tl.program_id(0).to(tl.int64)
+    if HAS_IDX:
+        row = tl.load(idx_ptr + m).to(tl.int64)
+    else:
+        row = m
+    dst = out_ptr + row * ROW_BYTES
+    src = x_ptr + m * stride_xm
+
+    if FMT == 0:
+        offs_d = tl.arange(0, D)
+        bits = tl.load(src + offs_d).to(tl.bfloat16).to(tl.uint16, bitcast=True)
+        tl.store(dst + 2 * offs_d, (bits & 0xFF).to(tl.uint8))
+        tl.store(dst + 2 * offs_d + 1, (bits >> 8).to(tl.uint8))
+    elif FMT == 1:
+        NG: tl.constexpr = D // GROUP
+        offs_d = tl.arange(0, D)
+        offs_g = tl.arange(0, NG)
+        x = tl.reshape(tl.load(src + offs_d).to(tl.float32), [NG, GROUP])
+        amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4)
+        e = _log2_ceil(amax * (1.0 / 448.0))
+        s = tl.exp2(e.to(tl.float32))
+        q = tl.clamp(tl.math.div_rn(x, s[:, None]), -448.0, 448.0)
+        tl.store(dst + offs_d, tl.reshape(_e4m3_codes(q), [D]))
+        tl.store(dst + D + offs_g, (e + 127).to(tl.uint8))
+    else:
+        # e2m1 pairs: even channels from one strided load, odd from another, so a byte is (lo | hi << 4)
+        NG: tl.constexpr = D // GROUP
+        HALF: tl.constexpr = GROUP // 2
+        offs_h = tl.arange(0, D // 2)
+        offs_g = tl.arange(0, NG)
+        xe = tl.reshape(tl.load(src + 2 * offs_h).to(tl.float32), [NG, HALF])
+        xo = tl.reshape(tl.load(src + 2 * offs_h + 1).to(tl.float32), [NG, HALF])
+        amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1))
+        if FMT == 2:
+            amax = tl.maximum(amax, 6.0 * (2.0 ** -9))
+            s = _e4m3_round(amax * (1.0 / 6.0))
+            tl.store(dst + D // 2 + offs_g, _e4m3_codes(s))
+        else:
+            amax = tl.maximum(amax, 6.0 * (2.0 ** -126))
+            e = _log2_ceil(amax * (1.0 / 6.0))
+            s = tl.exp2(e.to(tl.float32))
+            tl.store(dst + D // 2 + offs_g, (e + 127).to(tl.uint8))
+        # IEEE division: the reference's CUDA ``x / s`` is correctly rounded, and with an e4m3
+        # (non power-of-two) scale triton's default approximate div would land a hair off the
+        # e2m1 midpoints and flip their round-to-even
+        qe = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xe, s[:, None]), -6.0, 6.0)))
+        qo = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xo, s[:, None]), -6.0, 6.0)))
+        tl.store(dst + offs_h, tl.reshape(qe | (qo << 4), [D // 2]))
+
+
+@triton.jit
+def _unpack_rows_kernel(
+    pool_ptr, idx_ptr, out_ptr, M,
+    D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, BLOCK_T: tl.constexpr, HAS_IDX: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_t = pid * BLOCK_T + tl.arange(0, BLOCK_T)
+    valid = offs_t < M
+    if HAS_IDX:
+        rows = tl.load(idx_ptr + offs_t, mask=valid, other=0)
+    else:
+        rows = offs_t
+    offs_d = tl.arange(0, D)
+    vals = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT)
+    tl.store(
+        out_ptr + offs_t[:, None].to(tl.int64) * D + offs_d[None, :],
+        vals.to(out_ptr.dtype.element_ty), mask=valid[:, None],
+    )
+
+
+def pack_rows(
+    x: torch.Tensor, fmt: RowFormat, pool: torch.Tensor | None = None, row_ids: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize ``x [M, D]`` (bf16 / fp16 / fp32) into packed rows.
+
+    With ``pool [rows, row_bytes] uint8`` and ``row_ids [M]``, row ``m`` is written to ``pool[row_ids[m]]``
+    (the KV write path). Without a pool a fresh ``[M, row_bytes]`` uint8 tensor is returned.
+    """
+    assert x.dim() == 2, x.shape
+    m, d = x.shape
+    fmt.validate_dim(d)
+    row_bytes = fmt.row_bytes(d)
+    assert d & (d - 1) == 0 and d >= 16, f"pack_rows wants a power-of-two D >= 16, got {d}"
+    if pool is None:
+        assert row_ids is None
+        pool = torch.empty((m, row_bytes), dtype=torch.uint8, device=x.device)
+    else:
+        assert pool.dtype == torch.uint8 and pool.shape[1] == row_bytes and pool.is_contiguous(), (pool.shape, row_bytes)
+        assert row_ids is not None and row_ids.shape == (m,), (None if row_ids is None else row_ids.shape, m)
+    if m == 0:
+        return pool
+    if x.stride(1) != 1:
+        x = x.contiguous()
+    _pack_rows_kernel[(m,)](
+        x, pool, row_ids if row_ids is not None else pool, x.stride(0),
+        D=d, GROUP=fmt.group or d, ROW_BYTES=row_bytes, FMT=fmt.code, HAS_IDX=row_ids is not None,
+        num_warps=4,
+    )
+    return pool
+
+
+def unpack_rows(
+    pool: torch.Tensor, fmt: RowFormat, dim: int, row_ids: torch.Tensor | None = None,
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Gather ``pool[row_ids]`` (or every row) and dequantize to ``[M, dim]`` in ``out_dtype``."""
+    fmt.validate_dim(dim)
+    assert pool.dtype == torch.uint8 and pool.shape[1] == fmt.row_bytes(dim) and pool.is_contiguous(), pool.shape
+    m = pool.shape[0] if row_ids is None else row_ids.numel()
+    out = torch.empty((m, dim), dtype=out_dtype, device=pool.device)
+    if m == 0:
+        return out
+    block_t = 16
+    _unpack_rows_kernel[(triton.cdiv(m, block_t),)](
+        pool, row_ids if row_ids is not None else pool, out, m,
+        D=dim, ROW_BYTES=fmt.row_bytes(dim), FMT=fmt.code, BLOCK_T=block_t, HAS_IDX=row_ids is not None,
+        num_warps=4,
+    )
+    return out
+
+
+def quant_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor:
+    """``unpack(pack(x))`` -- the value the cache will hand back; the reference's ``inplace=True`` semantics."""
+    if fmt is BF16:
+        return x.to(torch.bfloat16)
+    return unpack_rows(pack_rows(x, fmt), fmt, x.shape[1], out_dtype=torch.bfloat16)
+
+
+__all__ = ["pack_rows", "unpack_rows", "quant_roundtrip"]
diff --git a/python/freetoken/kernel/triton/dsv41/row_format.py b/python/freetoken/kernel/triton/dsv41/row_format.py
new file mode 100644
index 0000000000..032d77e970
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/row_format.py
@@ -0,0 +1,92 @@
+"""Device-side row-format helpers shared by the DSV41 kernels.
+
+The Python-side descriptors live in ``freetoken.kvcache.dsv4.v41_row_format``; the ``FMT`` constexpr a
+kernel receives is ``RowFormat.code``:
+
+  0  bf16          2 B / channel, no scale
+  1  fp8_e8m0_b32  e4m3 byte / channel, ue8m0 scale byte per 32 channels
+  2  fp4_e4m3_b16  e2m1 nibble / channel (even channel low nibble), e4m3 scale byte per 16
+  3  fp4_e8m0_b32  e2m1 nibble / channel, ue8m0 scale byte per 32
+
+A row is ``[values | scales]``; ``ROW_BYTES`` is the pool's row stride. Every dequantized value
+is rounded through bf16 before use -- the reference bakes its quant round-trip into a bf16 cache,
+so this is what its attention / indexer kernels actually read.
+"""
+
+from __future__ import annotations
+
+import triton
+import triton.language as tl
+
+from freetoken.kernel.triton.e4m3_compat import e4m3_u8_to_f32
+
+FMT_BF16 = 0
+FMT_FP8_E8M0_B32 = 1
+FMT_FP4_E4M3_B16 = 2
+FMT_FP4_E8M0_B32 = 3
+
+
+@triton.jit
+def fp4_nibble_to_f32(nib):
+    """e2m1 code (uint8, 0..15) -> fp32: bit 3 sign, bits 2-1 exponent, bit 0 mantissa;
+    codes 0/1 are the subnormals 0 and 0.5, the rest ``(1 + m/2) * 2^(e-1)``."""
+    c = (nib & 7).to(tl.int32)
+    neg = (nib & 8) != 0
+    sub = c.to(tl.float32) * 0.5
+    normal = (1.0 + 0.5 * (c & 1).to(tl.float32)) * tl.exp2(((c >> 1) - 1).to(tl.float32))
+    mag = tl.where(c < 2, sub, normal)
+    return tl.where(neg, -mag, mag)
+
+
+@triton.jit
+def fp4_f32_to_nibble(q):
+    """fp32 value already on the e2m1 grid ({0, .5, 1, 1.5, 2, 3, 4, 6}, signed) -> e2m1 code.
+    Thresholds sit between grid points, so a value a hair off the grid still maps to its point."""
+    a = tl.abs(q)
+    c = tl.where(
+        a < 1.75,
+        tl.where(a < 0.75, tl.where(a < 0.25, 0, 1), tl.where(a < 1.25, 2, 3)),
+        tl.where(a < 3.5, tl.where(a < 2.5, 4, 5), tl.where(a < 5.0, 6, 7)),
+    )
+    return (c + tl.where(q < 0, 8, 0)).to(tl.uint8)
+
+
+@triton.jit
+def load_rows(pool_ptr, rows, valid, offs_d, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr):
+    """Gather + dequantize ``[T, D]`` fp32 from a packed pool.
+
+    ``pool_ptr`` uint8 pool base; ``rows`` [T] row ids (any int width; masked-out rows may be -1);
+    ``valid`` [T] bool; ``offs_d`` [D] = ``tl.arange(0, D)``. Invalid rows read as 0.
+    """
+    row_i64 = tl.maximum(rows, 0).to(tl.int64)
+    base = pool_ptr + row_i64[:, None] * ROW_BYTES
+    m = valid[:, None] & (offs_d >= 0)[None, :]
+    if FMT == 0:
+        lo = tl.load(base + 2 * offs_d[None, :], mask=m, other=0).to(tl.uint16)
+        hi = tl.load(base + 2 * offs_d[None, :] + 1, mask=m, other=0).to(tl.uint16)
+        val = (lo | (hi << 8)).to(tl.bfloat16, bitcast=True).to(tl.float32)
+    elif FMT == 1:
+        v = tl.load(base + offs_d[None, :], mask=m, other=0)
+        s = tl.load(base + D + (offs_d // 32)[None, :], mask=m, other=127)
+        val = e4m3_u8_to_f32(v) * tl.exp2(s.to(tl.float32) - 127.0)
+    else:
+        b = tl.load(base + (offs_d // 2)[None, :], mask=m, other=0)
+        nib = tl.where((offs_d % 2 == 0)[None, :], b & 0xF, b >> 4)
+        q = fp4_nibble_to_f32(nib)
+        if FMT == 2:
+            s = e4m3_u8_to_f32(tl.load(base + D // 2 + (offs_d // 16)[None, :], mask=m, other=0))
+        else:
+            s = tl.exp2(tl.load(base + D // 2 + (offs_d // 32)[None, :], mask=m, other=127).to(tl.float32) - 127.0)
+        val = q * s
+    return val.to(tl.bfloat16).to(tl.float32)
+
+
+__all__ = [
+    "FMT_BF16",
+    "FMT_FP8_E8M0_B32",
+    "FMT_FP4_E4M3_B16",
+    "FMT_FP4_E8M0_B32",
+    "fp4_nibble_to_f32",
+    "fp4_f32_to_nibble",
+    "load_rows",
+]
diff --git a/python/freetoken/kernel/triton/dsv41/sparse_attn.py b/python/freetoken/kernel/triton/dsv41/sparse_attn.py
new file mode 100644
index 0000000000..9d6b0816db
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/sparse_attn.py
@@ -0,0 +1,323 @@
+"""Sparse gathered-KV flash attention over PACKED pools (DeepSeek-V4.1).
+
+Same contract as ``kernel/triton/dsv4/sparse_attn.py`` -- each query attends to a per-query list of
+GLOBAL slots laid out ``[window part | compressed part]`` plus a per-head attention sink, K == V is one
+shared latent head -- but the two pools are byte pools in (possibly different) ``RowFormat``s: the
+window ring holds fp8 rows, the compressed KV holds fp4 rows, and every gathered row is dequantized
+in-kernel (``row_format.load_rows``) to the bf16 value the reference cache would have held.
+
+Because the two pools dequantize differently, a tile never mixes them: the window part is
+``n_window`` columns with ``n_window % BLOCK_T == 0`` (the caller pads it with ``-1``), so the kernel
+walks the window tiles and then the compressed tiles in two loops, each with a single load site.
+Both dots take bf16 operands with fp32 accumulation, like the reference ``sparse_attn`` kernel.
+
+Two implementations, picked from the launch shape (no knob):
+  * prefill (``m`` > 1): one program per (query, request, head block)
+  * decode (``m`` == 1): flash-decoding split over the candidate axis, merged through log-sum-exp
+    (the sink joins in the merge). Not bit-identical to the single-program path.
+
+``cmp_counts`` (device ``[b, m]`` int32) bounds the compressed columns a CUDA-graph replay visits;
+without it the full ``topk_idxs`` width is walked.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.kvcache.dsv4.v41_row_format import RowFormat
+
+from .row_format import load_rows
+
+BLOCK_H = 16
+BLOCK_T = 32
+MAX_SPLITS = 32
+# Decode splits at tile boundaries, bounded by the device's SM count below.
+MIN_TILES_PER_SPLIT = 1
+
+
+@triton.jit
+def _attend_tile(q, kv, valid, scale, m_i, l_i, acc):
+    """One online-softmax step over a ``[BLOCK_T, D]`` fp32 KV tile (bf16-exact values)."""
+    kv16 = kv.to(tl.bfloat16)
+    scores = tl.dot(q, tl.trans(kv16)) * scale  # [BLOCK_H, BLOCK_T] fp32
+    scores = tl.where(valid[None, :], scores, -float("inf"))
+    m_new = tl.maximum(m_i, tl.max(scores, axis=1))
+    # an all-masked tile keeps m_new at -inf; short-circuit so -inf - -inf never forms
+    alpha = tl.where(m_new == -float("inf"), 1.0, tl.exp(m_i - m_new))
+    p = tl.where(valid[None, :], tl.exp(scores - m_new[:, None]), 0.0)
+    l_i = l_i * alpha + tl.sum(p, axis=1)
+    acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv16)
+    return m_new, l_i, acc
+
+
+@triton.jit
+def _sparse_attn_packed_kernel(
+    q_ptr, win_ptr, cmp_ptr, o_ptr, sink_ptr, idx_ptr, cnt_ptr,
+    scale,
+    H, TOPK, N_WINDOW,
+    stride_qb, stride_qm, stride_qh, stride_qd,
+    stride_ob, stride_om, stride_oh, stride_od,
+    stride_ib, stride_im, stride_it,
+    stride_nb, stride_nm,
+    D: tl.constexpr,
+    WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr,
+    CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    HAS_COUNTS: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_h = tl.program_id(2)
+
+    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = offs_h < H
+    offs_d = tl.arange(0, D)
+
+    q_ptrs = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd
+    q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16)  # [BLOCK_H, D]
+
+    m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, D), dtype=tl.float32)
+
+    n_active = TOPK
+    if HAS_COUNTS:
+        n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm)
+
+    idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im
+    for start in range(0, N_WINDOW, BLOCK_T):
+        offs_t = start + tl.arange(0, BLOCK_T)
+        idxs = tl.load(idx_base + offs_t * stride_it)
+        valid = idxs >= 0
+        kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT)
+        m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc)
+    for start in range(N_WINDOW, n_active, BLOCK_T):
+        offs_t = start + tl.arange(0, BLOCK_T)
+        t_mask = offs_t < n_active
+        idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1)
+        valid = idxs >= 0
+        kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT)
+        m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc)
+
+    sink = tl.load(sink_ptr + offs_h, mask=h_mask, other=0.0).to(tl.float32)
+    l_i = l_i + tl.exp(sink - m_i)
+    o = acc / l_i[:, None]
+
+    o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_d[None, :] * stride_od
+    tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None])
+
+
+@triton.jit
+def _sparse_attn_packed_splitk_kernel(
+    q_ptr, win_ptr, cmp_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr,
+    scale,
+    H, TOPK, N_WINDOW,
+    stride_qb, stride_qm, stride_qh, stride_qd,
+    stride_mb, stride_mm, stride_mh, stride_ms, stride_md,
+    stride_lb, stride_lm, stride_lh, stride_ls,
+    stride_ib, stride_im, stride_it,
+    stride_nb, stride_nm,
+    D: tl.constexpr,
+    WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr,
+    CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    HAS_COUNTS: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+):
+    """Stage 1: each program reduces one BLOCK_T-aligned slice of the candidate list and writes
+    its normalized partial output + log-sum-exp. A slice may straddle the window / compressed
+    boundary; it is walked as (window tiles, compressed tiles) since N_WINDOW is tile-aligned."""
+    pid_ms = tl.program_id(0)
+    pid_m = pid_ms // NUM_SPLITS
+    split_id = pid_ms % NUM_SPLITS
+    pid_b = tl.program_id(1)
+    pid_h = tl.program_id(2)
+
+    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = offs_h < H
+    offs_d = tl.arange(0, D)
+
+    n_active = TOPK
+    if HAS_COUNTS:
+        n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm)
+
+    per_split = tl.cdiv(tl.cdiv(n_active, NUM_SPLITS), BLOCK_T) * BLOCK_T
+    split_start = per_split * split_id
+    split_end = tl.minimum(split_start + per_split, n_active)
+
+    m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, D), dtype=tl.float32)
+
+    if split_end > split_start:
+        q_ptrs = (
+            q_ptr + pid_b * stride_qb + pid_m * stride_qm
+            + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd
+        )
+        q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16)
+        idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im
+
+        win_end = tl.minimum(split_end, N_WINDOW)
+        for start in range(split_start, win_end, BLOCK_T):
+            offs_t = start + tl.arange(0, BLOCK_T)
+            idxs = tl.load(idx_base + offs_t * stride_it)
+            valid = idxs >= 0
+            kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT)
+            m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc)
+        cmp_start = tl.maximum(split_start, N_WINDOW)
+        for start in range(cmp_start, split_end, BLOCK_T):
+            offs_t = start + tl.arange(0, BLOCK_T)
+            t_mask = offs_t < split_end
+            idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1)
+            valid = idxs >= 0
+            kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT)
+            m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc)
+
+    out = tl.where(l_i[:, None] == 0.0, 0.0, acc / l_i[:, None])
+    lse = tl.where(l_i == 0.0, -float("inf"), m_i + tl.log(l_i))
+
+    mid_base = (
+        mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm
+        + offs_h[:, None] * stride_mh + split_id * stride_ms + offs_d[None, :] * stride_md
+    )
+    tl.store(mid_base, out, mask=h_mask[:, None])
+    lse_base = (
+        mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm
+        + offs_h * stride_lh + split_id * stride_ls
+    )
+    tl.store(lse_base, lse, mask=h_mask)
+
+
+@triton.jit
+def _sparse_attn_splitk_merge_kernel(
+    mid_o_ptr, mid_lse_ptr, o_ptr, sink_ptr,
+    stride_mb, stride_mm, stride_mh, stride_ms, stride_md,
+    stride_lb, stride_lm, stride_lh, stride_ls,
+    stride_ob, stride_om, stride_oh, stride_od,
+    D: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+):
+    """Stage 2: log-sum-exp merge over the splits; the attention sink joins here once, as a null
+    key with logit ``attn_sink[h]`` and zero value. One program per (query, request, head)."""
+    pid_m = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_h = tl.program_id(2)
+    offs_d = tl.arange(0, D)
+
+    m_i = tl.load(sink_ptr + pid_h).to(tl.float32)
+    l_i = 1.0
+    acc = tl.zeros((D,), dtype=tl.float32)
+
+    mid_base = mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm + pid_h * stride_mh + offs_d * stride_md
+    lse_base = mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm + pid_h * stride_lh
+
+    for split_id in tl.range(0, NUM_SPLITS, num_stages=2):
+        partial = tl.load(mid_base + split_id * stride_ms)
+        lse = tl.load(lse_base + split_id * stride_ls)
+        m_new = tl.maximum(m_i, lse)
+        alpha = tl.exp(m_i - m_new)
+        beta = tl.where(lse == -float("inf"), 0.0, tl.exp(lse - m_new))
+        acc = acc * alpha + partial * beta
+        l_i = l_i * alpha + beta
+        m_i = m_new
+
+    o = acc / l_i
+    o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_d * stride_od
+    tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty))
+
+
+def split_count(b: int, m: int, h: int, topk: int, device) -> int:
+    """How many ways to split the candidate axis; 0 means run the single-program kernel. Decode
+    splits until every SM has a program or every split is down to one tile."""
+    if m != 1:
+        return 0
+    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    n_splits = min(
+        MAX_SPLITS,
+        triton.cdiv(topk, MIN_TILES_PER_SPLIT * BLOCK_T),
+        max(1, sm_count // (b * triton.cdiv(h, BLOCK_H))),
+    )
+    return n_splits if n_splits > 1 else 0
+
+
+def sparse_attn_packed(
+    q: torch.Tensor,             # [b, m, h, d] bf16
+    window_pool: torch.Tensor,   # [n_win_slots, win_fmt.row_bytes(d)] uint8
+    win_fmt: RowFormat,
+    cmp_pool: torch.Tensor,      # [n_cmp, cmp_fmt.row_bytes(d)] uint8
+    cmp_fmt: RowFormat,
+    attn_sink: torch.Tensor,     # [h]
+    topk_idxs: torch.Tensor,     # [b, m, topk] int32, GLOBAL rows, layout [window | compressed]
+    n_window: int,               # window columns (a multiple of BLOCK_T; pad with -1)
+    softmax_scale: float,
+    cmp_counts: torch.Tensor | None = None,  # [b, m] int32 live compressed columns per query
+) -> torch.Tensor:
+    """Paged sparse MLA attention over packed KV pools; see the module docstring for the contract."""
+    b, m, h, d = q.shape
+    topk = topk_idxs.shape[-1]
+    assert n_window % BLOCK_T == 0 and 0 <= n_window <= topk, (n_window, topk)
+    assert window_pool.dtype == torch.uint8 and window_pool.shape[1] == win_fmt.row_bytes(d), (window_pool.shape, win_fmt)
+    assert cmp_pool.dtype == torch.uint8 and cmp_pool.shape[1] == cmp_fmt.row_bytes(d), (cmp_pool.shape, cmp_fmt)
+    assert window_pool.is_contiguous() and cmp_pool.is_contiguous()
+    q = q.contiguous()
+    idx = topk_idxs.contiguous().to(torch.int32)
+    sink = attn_sink.contiguous().to(torch.float32)
+    o = torch.empty_like(q)
+
+    has_counts = cmp_counts is not None
+    if has_counts:
+        cnt = cmp_counts.contiguous().to(torch.int32).view(b, m)
+        stride_nb, stride_nm = cnt.stride()
+    else:
+        cnt, stride_nb, stride_nm = idx, 0, 0
+
+    fmt_args = dict(
+        WIN_FMT=win_fmt.code, WIN_ROW_BYTES=win_fmt.row_bytes(d),
+        CMP_FMT=cmp_fmt.code, CMP_ROW_BYTES=cmp_fmt.row_bytes(d),
+    )
+    n_splits = split_count(b, m, h, topk, q.device)
+    if n_splits:
+        head_blocks = triton.cdiv(h, BLOCK_H)
+        mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device)
+        mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device)
+        _sparse_attn_packed_splitk_kernel[(m * n_splits, b, head_blocks)](
+            q, window_pool, cmp_pool, mid_o, mid_lse, idx, cnt,
+            float(softmax_scale),
+            h, topk, int(n_window),
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+            mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4),
+            mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3),
+            idx.stride(0), idx.stride(1), idx.stride(2),
+            stride_nb, stride_nm,
+            D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, NUM_SPLITS=n_splits,
+            **fmt_args, num_warps=8, num_stages=2,
+        )
+        _sparse_attn_splitk_merge_kernel[(m, b, h)](
+            mid_o, mid_lse, o, sink,
+            mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4),
+            mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3),
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+            D=d, NUM_SPLITS=n_splits, num_warps=4,
+        )
+        return o
+
+    grid = (m, b, triton.cdiv(h, BLOCK_H))
+    _sparse_attn_packed_kernel[grid](
+        q, window_pool, cmp_pool, o, sink, idx, cnt,
+        float(softmax_scale),
+        h, topk, int(n_window),
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+        idx.stride(0), idx.stride(1), idx.stride(2),
+        stride_nb, stride_nm,
+        D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts,
+        **fmt_args, num_warps=8, num_stages=2,
+    )
+    return o
+
+
+__all__ = ["sparse_attn_packed", "split_count", "BLOCK_T"]
diff --git a/python/freetoken/kernel/triton/dsv41/topk.py b/python/freetoken/kernel/triton/dsv41/topk.py
new file mode 100644
index 0000000000..4d638e4d49
--- /dev/null
+++ b/python/freetoken/kernel/triton/dsv41/topk.py
@@ -0,0 +1,265 @@
+"""Exact top-k and block-candidate selection whose work follows the live history (DeepSeek-V4.1).
+
+The indexer scores a request's compressed positions into a buffer as wide as the staged history
+(``max_seq_len / ratio``, fixed while a CUDA graph is captured), but only the first ``live[row]``
+columns are meaningful -- and at 1M staged positions a 4K-token request must not pay for the
+other 99.6%. Every kernel here takes the live count per row on the device, keeps its output
+addresses and shapes fixed, and restricts what it READS and reduces to the live prefix; nothing
+past ``live`` is ever read, so the score buffer's dead area needs no initialization.
+
+Selection is a fixed pipeline of small kernels with explicit workspace contracts:
+
+* ``dsv41_topk``: exact top-k of ``[R, T]`` fp32 scores below ``live``. Level 0 gives every
+  ``CHUNK``-wide slice of a row its own program, which radix-selects the slice's top-k into a
+  ``(key, column)`` candidate workspace (a global winner beats at most k-1 columns, so it is
+  inside its own slice's top-k: the union is exact). Further levels apply the same program to the
+  candidates until one register-resident tile remains; the final program ranks it and emits the
+  winners in ascending column order, ``-1`` padded. Slices past ``live`` exit at once and a
+  slice with no more entries than ``k`` is copied through, so a short history costs a handful of
+  near-empty launches. The radix primitives are the QSA block top-k's (``kernel/triton/qsa/topk``).
+* ``dsv41_block_scores``: per-``block_size`` block maxima over the live prefix (the Hierarchical
+  Sparse Indexer's level one), the block holding the newest live position pinned to ``+inf`` so it
+  is always kept; a fixed worker grid strides over the live blocks.
+* ``dsv41_candidate_blocks``: the two above plus a fixed-width expansion into the compact candidate
+  list ``[R, topk_blocks * block_size]`` of compressed positions, ascending with ``-1`` in the tail.
+
+Tie rule (documented, deterministic): scores order descending, equal scores by ascending column,
+so the lowest positions win among exact ties. ``torch.topk`` leaves ties unspecified.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.kernel.triton.qsa.topk import _BINS, _PASSES, _RADIX, _monotone_key, _resident_prefix, _tile_ranks
+
+CHUNK = 4096  # register-resident slice per level program (fp32 keys + int32 columns)
+FINAL_MAX = 4096  # the final program's resident candidate tile
+BLOCK_SCORE_TILE = 512  # blocks per block-max program tile
+
+
+@triton.jit
+def _load_entries(src_ptr, col_ptr, base, offsets, limit, SRC_LOGITS: tl.constexpr):
+    """One slice of a level's input as ``(monotone key, column)``; key 0 = dead (past ``limit``,
+    ``-inf``, or a padding slot of the previous level)."""
+    live = offsets < limit
+    if SRC_LOGITS:
+        value = tl.load(src_ptr + base + offsets, mask=live, other=-float("inf"))
+        key = tl.where(live & (value > -float("inf")), _monotone_key(value), 0)
+        cols = (base + offsets).to(tl.int32)
+    else:
+        key = tl.load(src_ptr + base + offsets, mask=live, other=0).to(tl.uint32, bitcast=True)
+        cols = tl.load(col_ptr + base + offsets, mask=live, other=-1)
+    return key, cols
+
+
+@triton.jit
+def _dsv41_topk_level_kernel(
+    src_ptr, src_col_ptr, live_in_ptr, key_out_ptr, col_out_ptr, live_out_ptr,
+    stride_src_row, stride_out_row, num_columns,
+    TOP_K: tl.constexpr, PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr,
+    SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr,
+):
+    """One program per (row, slice): the slice's top-k into candidate slot ``split * TOP_K``. Slot
+    ``split`` is read downstream only when ``split < cdiv(live, CHUNK)``; program 0 publishes that
+    live candidate count for the next level."""
+    tl.static_assert(CHUNK <= 0xFFFF, "packed cumsum keeps 16 bits per half")
+    tl.static_assert(TOP_K <= CHUNK, "a slice must hold its own top-k")
+    row = tl.program_id(0)
+    split = tl.program_id(1)
+    visible = tl.maximum(tl.minimum(tl.load(live_in_ptr + row), num_columns), 0)
+    if split == 0:
+        tl.store(live_out_ptr + row, tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K)
+    base = split * CHUNK
+    limit = tl.minimum(tl.maximum(visible - base, 0), CHUNK)
+    if limit <= 0:
+        return  # a dead slice: nothing downstream reads its slot
+    slot = row.to(tl.int64) * stride_out_row + split * TOP_K
+    offsets = tl.arange(0, CHUNK)
+    key, cols = _load_entries(src_ptr + row.to(tl.int64) * stride_src_row, src_col_ptr + row.to(tl.int64) * stride_src_row, base, offsets, limit, SRC_LOGITS)
+    if limit <= TOP_K:
+        # every live entry is a winner: copy in order (dead keys stay 0 for the next level)
+        inside = offsets < TOP_K
+        tl.store(key_out_ptr + slot + offsets, key.to(tl.int32, bitcast=True), mask=inside)
+        tl.store(col_out_ptr + slot + offsets, cols, mask=inside)
+        return
+    prefix, ties = _resident_prefix(key, TOP_K, BINS, RADIX, PASSES)
+    rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0)
+    write = take & (rank < TOP_K)
+    tl.store(key_out_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write)
+    tl.store(col_out_ptr + slot + rank, cols, mask=write)
+    emitted = above + tl.minimum(equal, ties)
+    pad = tl.arange(0, PAD_K)
+    tl.store(key_out_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K))
+
+
+@triton.jit
+def _final_tile(src_row, col_row, out_row, limit, TOP_K: tl.constexpr, BLOCK: tl.constexpr, SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr):
+    tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half")
+    offsets = tl.arange(0, BLOCK)
+    key, cols = _load_entries(src_row, col_row, 0, offsets, limit, SRC_LOGITS)
+    prefix, ties = _resident_prefix(key, tl.minimum(limit, TOP_K), BINS, RADIX, PASSES)
+    rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0)
+    tl.store(out_row + rank, cols, mask=take & (rank < TOP_K))
+    return above + tl.minimum(equal, ties)
+
+
+@triton.jit
+def _dsv41_topk_final_kernel(
+    src_ptr, src_col_ptr, live_ptr, out_ptr,
+    stride_src_row, stride_out_row, num_columns,
+    TOP_K: tl.constexpr, PAD_K: tl.constexpr,
+    BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr,
+    SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr,
+):
+    """One program per row over the last level's live candidates (or the scores themselves when the
+    row fits one tile): winners at their ascending-column ranks, ``-1`` from ``emitted`` on."""
+    row = tl.program_id(0)
+    limit = tl.maximum(tl.minimum(tl.load(live_ptr + row), num_columns), 0)
+    src_row = src_ptr + row.to(tl.int64) * stride_src_row
+    col_row = src_col_ptr + row.to(tl.int64) * stride_src_row
+    out_row = out_ptr + row.to(tl.int64) * stride_out_row
+    emitted = 0
+    if limit > 0:
+        # register residency costs the whole tile, so a short row takes a narrower one
+        if limit <= BLOCK_SMALL:
+            emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_SMALL, SRC_LOGITS, BINS, RADIX, PASSES)
+        elif limit <= BLOCK_MID:
+            emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_MID, SRC_LOGITS, BINS, RADIX, PASSES)
+        else:
+            emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_FULL, SRC_LOGITS, BINS, RADIX, PASSES)
+    pad = tl.arange(0, PAD_K)
+    tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K))
+
+
+MAX_K = CHUNK // 2  # a slice must shed at least half its entries per level, or the plan cannot converge
+
+
+def topk_plan(columns: int, k: int) -> list[tuple[int, int]]:
+    """``[(input width, n_splits), ...]`` of the reduction levels before the final tile; empty when
+    the row already fits the final tile. Derived from the buffer width alone (graph-static). Each
+    level maps ``width`` to ``cdiv(width, CHUNK) * k`` candidates, which shrinks the row by at least
+    ``CHUNK / k >= 2`` per level for ``k <= MAX_K``; larger ``k`` would stall (``k == CHUNK`` never
+    shrinks) and is rejected here, before any launch."""
+    if not 0 < k <= MAX_K:
+        raise ValueError(f"dsv41_topk supports 1 <= k <= {MAX_K}, got {k}")
+    levels: list[tuple[int, int]] = []
+    width = columns
+    while width > FINAL_MAX:
+        n_splits = -(-width // CHUNK)
+        levels.append((width, n_splits))
+        nxt = n_splits * k
+        assert nxt < width, (width, k)  # guaranteed by k <= CHUNK // 2 for width > FINAL_MAX == CHUNK
+        width = nxt
+    return levels
+
+
+def dsv41_topk(scores: torch.Tensor, live: torch.Tensor, k: int, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Exact top-``k`` columns of every ``scores`` row among its first ``live[row]`` columns.
+
+    ``scores [R, T]`` fp32 row-contiguous (columns at or past ``live`` are never read; ``-inf``
+    columns are dead), ``live [R]`` int32. Returns ``out [R, k]`` int32: the winners in ascending
+    column order packed at the front, ``-1`` from the ``emitted`` count on (fewer live or finite
+    columns than ``k`` leave a ``-1`` tail). Ties: lower column wins."""
+    assert scores.ndim == 2 and scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.shape, scores.dtype)
+    R, T = scores.shape
+    assert live.shape == (R,), (live.shape, R)
+    live = live.to(torch.int32).contiguous()
+    if out is None:
+        out = torch.empty((R, k), dtype=torch.int32, device=scores.device)
+    assert out.shape == (R, k) and out.dtype == torch.int32 and out.stride(1) == 1
+    if R == 0 or k == 0:
+        return out
+    if T == 0:
+        out.fill_(-1)
+        return out
+    pad_k = triton.next_power_of_2(k)
+    src, src_col, src_live, width, src_logits = scores, scores, live, T, True
+    for level_width, n_splits in topk_plan(T, k):
+        keys = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device)
+        cols = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device)
+        live_out = torch.empty((R,), dtype=torch.int32, device=scores.device)
+        _dsv41_topk_level_kernel[(R, n_splits)](
+            src, src_col, src_live, keys, cols, live_out,
+            src.stride(0), keys.stride(0), level_width,
+            TOP_K=k, PAD_K=pad_k, CHUNK=CHUNK, N_SPLITS=n_splits, SRC_LOGITS=src_logits,
+            BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1,
+        )
+        src, src_col, src_live, width, src_logits = keys, cols, live_out, n_splits * k, False
+    block = triton.next_power_of_2(max(width, 1))
+    _dsv41_topk_final_kernel[(R,)](
+        src, src_col, src_live, out,
+        src.stride(0), out.stride(0), width,
+        TOP_K=k, PAD_K=pad_k,
+        BLOCK_SMALL=min(1024, block), BLOCK_MID=min(4096, block), BLOCK_FULL=block,
+        SRC_LOGITS=src_logits, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1,
+    )
+    return out
+
+
+@triton.jit
+def _dsv41_block_scores_kernel(
+    scores_ptr, live_ptr, blk_ptr,
+    stride_s_row, stride_b_row, num_blocks,
+    BLOCK_SIZE: tl.constexpr, TILE: tl.constexpr, N_WORKERS: tl.constexpr,
+):
+    """Worker ``w`` of a row scores block tiles ``w, w + N_WORKERS, ...`` below the live block count:
+    the block's best live position, ``+inf`` for the block holding the newest live position."""
+    row = tl.program_id(0)
+    worker = tl.program_id(1)
+    live = tl.maximum(tl.load(live_ptr + row), 0)
+    nb_live = tl.minimum(tl.cdiv(live, BLOCK_SIZE), num_blocks)
+    s_row = scores_ptr + row.to(tl.int64) * stride_s_row
+    b_row = blk_ptr + row.to(tl.int64) * stride_b_row
+    offs_b = tl.arange(0, TILE)
+    offs_p = tl.arange(0, BLOCK_SIZE)
+    for tile in range(worker, tl.cdiv(nb_live, TILE), N_WORKERS):
+        blocks = tile * TILE + offs_b
+        b_mask = blocks < nb_live
+        pos = blocks[:, None] * BLOCK_SIZE + offs_p[None, :]
+        v = tl.load(s_row + pos, mask=b_mask[:, None] & (pos < live), other=-float("inf"))
+        best = tl.max(v, axis=1)
+        best = tl.where(blocks == nb_live - 1, float("inf"), best)
+        tl.store(b_row + blocks, best, mask=b_mask)
+
+
+def dsv41_block_scores(scores: torch.Tensor, live: torch.Tensor, block_size: int, out: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(block scores [R, cdiv(T, block_size)] fp32, live blocks [R] int32)``: each live block's best
+    live position (``+inf`` for the newest block); blocks at or past the live count are not written."""
+    R, T = scores.shape
+    nb = -(-T // block_size)
+    live = live.to(torch.int32).contiguous()
+    if out is None:
+        out = torch.empty((R, nb), dtype=torch.float32, device=scores.device)
+    live_blocks = torch.div(live + block_size - 1, block_size, rounding_mode="floor").clamp_(min=0).to(torch.int32)
+    if R == 0 or nb == 0:
+        return out, live_blocks
+    assert block_size == triton.next_power_of_2(block_size), block_size
+    sm_count = torch.cuda.get_device_properties(scores.device).multi_processor_count
+    workers = min(triton.cdiv(nb, BLOCK_SCORE_TILE), max(1, triton.cdiv(sm_count, R)))
+    _dsv41_block_scores_kernel[(R, workers)](
+        scores, live, out, scores.stride(0), out.stride(0), nb,
+        BLOCK_SIZE=block_size, TILE=BLOCK_SCORE_TILE, N_WORKERS=workers, num_warps=4,
+    )
+    return out, live_blocks
+
+
+def dsv41_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor:
+    """Level one of the Hierarchical Sparse Indexer over ``[R, T]`` scores: the ``topk_blocks`` best
+    live blocks (the newest always among them) expanded to the compact candidate list ``[R,
+    topk_blocks * block_size]`` of compressed positions -- ascending, ``-1`` for positions at or past
+    ``live`` and for unfilled blocks, all in the tail (a sorted valid prefix)."""
+    R, T = scores.shape
+    blk, live_blocks = dsv41_block_scores(scores, live, block_size)
+    keep = dsv41_topk(blk, live_blocks, min(topk_blocks, blk.shape[1]))  # [R, KB] ascending, -1 tail
+    pos = keep.to(torch.int64).unsqueeze(-1) * block_size + torch.arange(block_size, device=scores.device)
+    dead = (keep < 0).unsqueeze(-1) | (pos >= live.to(torch.int64).view(R, 1, 1))
+    out = torch.where(dead, -1, pos).flatten(1).to(torch.int32)
+    if out.shape[1] < topk_blocks * block_size:
+        out = torch.nn.functional.pad(out, (0, topk_blocks * block_size - out.shape[1]), value=-1)
+    return out
+
+
+__all__ = ["dsv41_topk", "dsv41_block_scores", "dsv41_candidate_blocks", "topk_plan", "CHUNK", "FINAL_MAX", "MAX_K"]
diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py
index 6ea0c7537d..14a7acd835 100644
--- a/python/freetoken/kernel/triton/e4m3_compat.py
+++ b/python/freetoken/kernel/triton/e4m3_compat.py
@@ -109,6 +109,20 @@ def e4m3_u8_to_f16_x128(v):
     return h.to(tl.float16, bitcast=True) * 32768.0
 
 
+@jit
+def e4m3_f32_to_u8(x):
+    """Encode an fp32 value that already sits on the e4m3 grid (see :func:`round_e4m3`, |x| <= 448)
+    into its e4m3 bit pattern (uint8). Normals: fp32 exponent rebased 127 -> 7 with the top three
+    mantissa bits; subnormals (|x| < 2^-6, grid 2^-9): mantissa = |x| / 2^-9, exponent field 0."""
+    b = x.to(tl.uint32, bitcast=True)
+    sign = ((b >> 31) & 1).to(tl.uint8) << 7
+    exp = ((b >> 23) & 0xFF).to(tl.int32) - 127 + 7
+    man = ((b >> 20) & 0x7).to(tl.uint8)
+    normal = (exp.to(tl.uint8) << 3) | man
+    sub = (tl.abs(x) * 512.0).to(tl.uint8)  # 0..7 on the 2^-9 grid
+    return sign | tl.where(tl.abs(x) >= 0.015625, normal, sub)
+
+
 @jit
 def round_e4m3(x):
     """Round fp32 onto the e4m3 value grid (RNE), fp32 -> fp32, in a SINGLE
diff --git a/python/freetoken/kernel/triton/mhc.py b/python/freetoken/kernel/triton/mhc.py
index 0782efe37d..876643c164 100644
--- a/python/freetoken/kernel/triton/mhc.py
+++ b/python/freetoken/kernel/triton/mhc.py
@@ -27,7 +27,7 @@ def _mhc_stage1_kernel(
     H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr,
     SPLIT: tl.constexpr,     # hidden elems per split (multiple of BLOCK_H)
     BLOCK_H: tl.constexpr,
-    NS: tl.constexpr,
+    NS: tl.constexpr,        # splits per token (the grid's second axis; need not be a power of two)
     HAS_POST: tl.constexpr,
 ):
     """Split-K stage of the fused mHC: each program owns one hidden slice of one
@@ -79,24 +79,24 @@ def _mhc_stage1_kernel(
 
 
 @triton.jit
-def _mhc_stage2_kernel(
-    sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr,
+def _mhc_gates(
+    t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr,
     post_out_ptr, comb_out_ptr, pre_out_ptr,
     rms_eps, hc_eps, post_mult,
     SINKHORN: tl.constexpr,
     H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr,
-    NS: tl.constexpr,
+    NS: tl.constexpr, BLK_NS: tl.constexpr,
 ):
-    """Reduce the split partials and run the tiny gate math (sigmoid gates,
-    row-softmax + in-register 4x4 Sinkhorn); emits pre gates for stage 3."""
-    t = tl.program_id(0).to(tl.int64)
+    """Reduce token ``t``'s split partials and run the tiny gate math (sigmoid gates,
+    row-softmax + in-register 4x4 Sinkhorn); writes post / comb / pre for the next sublayer."""
     offs_mix = tl.arange(0, BLK_MIX)
     mix_mask = offs_mix < MIX
-    offs_s = tl.arange(0, NS)
+    offs_s = tl.arange(0, BLK_NS)
+    s_mask = offs_s < NS
 
-    sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s))
+    sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s, mask=s_mask, other=0.0))
     acc = tl.sum(
-        tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :]),
+        tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :], mask=s_mask[:, None], other=0.0),
         axis=0,
     )
     inv_rms = tl.math.rsqrt(sqsum / (N * H) + rms_eps)
@@ -142,6 +142,23 @@ def _mhc_stage2_kernel(
     tl.store(comb_out_ptr + t * N * N + offs_n2[:, None] * N + offs_n2[None, :], comb)
 
 
+
+@triton.jit
+def _mhc_stage2_kernel(
+    sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr,
+    post_out_ptr, comb_out_ptr, pre_out_ptr,
+    rms_eps, hc_eps, post_mult,
+    SINKHORN: tl.constexpr,
+    H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr,
+    NS: tl.constexpr, BLK_NS: tl.constexpr,
+):
+    """Gate math alone (the two-pass variant, whose stage 3 consumes these pre gates)."""
+    _mhc_gates(
+        tl.program_id(0).to(tl.int64), sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr,
+        rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS,
+    )
+
+
 @triton.jit
 def _mhc_stage3_kernel(
     res_out_ptr, pre_ptr, li_out_ptr,
@@ -161,7 +178,46 @@ def _mhc_stage3_kernel(
     tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask)
 
 
-def mhc_fused_post_pre_triton(
+@triton.jit
+def _mhc_stage23_kernel(
+    res_out_ptr, pre_ptr, li_out_ptr,
+    sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr,
+    post_out_ptr, comb_out_ptr, pre_out_ptr,
+    rms_eps, hc_eps, post_mult,
+    SINKHORN: tl.constexpr,
+    H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr,
+    NS: tl.constexpr, BLK_NS: tl.constexpr, BLOCK_H: tl.constexpr,
+):
+    """Single-pass mHC: the layer input mixes with the PREVIOUS sublayer's pre gates, so the
+    gate math for the next sublayer (stage 2) does not gate stage 3 -- one launch does both:
+    every program mixes one hidden chunk, program 0 of each token also reduces the split
+    partials and writes the next sublayer's gates."""
+    t = tl.program_id(0).to(tl.int64)
+    hb = tl.program_id(1)
+    offs_h = hb * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = offs_h < H
+    offs_n = tl.arange(0, N)
+    pre = tl.load(pre_ptr + t * N + offs_n)
+    li = tl.zeros([BLOCK_H], dtype=tl.float32)
+    for n in tl.static_range(N):
+        r = tl.load(res_out_ptr + (t * N + n) * H + offs_h, mask=h_mask, other=0.0).to(tl.float32)
+        li += tl.sum(tl.where(offs_n == n, pre, 0.0)) * r
+    tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask)
+    if hb == 0:
+        _mhc_gates(
+            t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr,
+            rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS,
+        )
+
+
+def _stage1_hidden_tile(h: int) -> int:
+    """The stage-1 hidden split. It fixes the fp32 reduction order over the hidden width, so it
+    depends on the width alone and results reproduce across processes, machines and chunk sizes.
+    128 gives 32 / 40 programs per token at widths 4096 / 5120."""
+    return min(128, triton.next_power_of_2(h))
+
+
+def _mhc_fused_launch(
     x: torch.Tensor,
     residual: torch.Tensor,
     post_mix: torch.Tensor | None,
@@ -173,11 +229,12 @@ def mhc_fused_post_pre_triton(
     hc_eps: float,
     post_mult: float,
     sinkhorn_repeat: int,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage
-    split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns
-    (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32,
-    layer_input [T,H] bf16)."""
+    pre_mix_in: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The three stages behind both variants. ``pre_mix_in`` selects single-pass: the
+    layer input is mixed with it instead of with the pre gates stage 2 just computed
+    (those are still returned, for the next sublayer). Returns (residual_new, post,
+    comb, pre_new [T,N] fp32, layer_input)."""
     t, n, h = residual.shape
     mix = 2 * n + n * n
     assert fn.shape == (mix, n * h) and fn.dtype == torch.float32
@@ -190,12 +247,7 @@ def mhc_fused_post_pre_triton(
     comb_out = torch.empty(t, n, n, dtype=torch.float32, device=dev)
     li_out = torch.empty(t, h, dtype=residual.dtype, device=dev)
 
-    block_h = min(512, triton.next_power_of_2(h))
-    # NS feeds a tl.arange in stage 2 -> keep it a power of two.
-    ns = 1
-    while ns * 2 <= min(16, h // block_h):
-        ns *= 2
-    split = triton.cdiv(triton.cdiv(h, ns), block_h) * block_h
+    block_h = split = _stage1_hidden_tile(h)
     ns = triton.cdiv(h, split)
     blk_mix = triton.next_power_of_2(mix)
     sq_part = torch.empty(t, ns, dtype=torch.float32, device=dev)
@@ -213,20 +265,83 @@ def mhc_fused_post_pre_triton(
         HAS_POST=has_post,
         num_warps=4, num_stages=2,
     )
-    _mhc_stage2_kernel[(t,)](
-        sq_part, mix_part, hc_scale, hc_base,
-        post_out, comb_out, pre_out,
-        rms_eps, hc_eps, post_mult,
-        SINKHORN=sinkhorn_repeat,
-        H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns,
-        num_warps=1,
+    block_h3 = min(1024, triton.next_power_of_2(h))
+    if pre_mix_in is None:
+        # two-pass: stage 3 mixes with the gates stage 2 just produced
+        _mhc_stage2_kernel[(t,)](
+            sq_part, mix_part, hc_scale, hc_base,
+            post_out, comb_out, pre_out,
+            rms_eps, hc_eps, post_mult,
+            SINKHORN=sinkhorn_repeat,
+            H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns),
+            num_warps=1,
+        )
+        _mhc_stage3_kernel[(t, triton.cdiv(h, block_h3))](
+            res_out, pre_out, li_out,
+            H=h, N=n, BLOCK_H=block_h3,
+            num_warps=4,
+        )
+    else:
+        # single-pass: the mix uses the previous sublayer's gates, so stages 2 and 3 share a launch
+        assert pre_mix_in.shape == (t, n), (pre_mix_in.shape, (t, n))
+        _mhc_stage23_kernel[(t, triton.cdiv(h, block_h3))](
+            res_out, pre_mix_in.to(torch.float32).contiguous(), li_out,
+            sq_part, mix_part, hc_scale, hc_base,
+            post_out, comb_out, pre_out,
+            rms_eps, hc_eps, post_mult,
+            SINKHORN=sinkhorn_repeat,
+            H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns), BLOCK_H=block_h3,
+            num_warps=4,
+        )
+    return res_out, post_out.view(t, n, 1), comb_out, pre_out, li_out
+
+
+def mhc_fused_post_pre_triton(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_mix: torch.Tensor | None,
+    comb_mix: torch.Tensor | None,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage
+    split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns
+    (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32,
+    layer_input [T,H] bf16)."""
+    res_out, post_out, comb_out, _pre_out, li_out = _mhc_fused_launch(
+        x, residual, post_mix, comb_mix, fn, hc_scale, hc_base,
+        rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=None,
     )
-    _mhc_stage3_kernel[(t, triton.cdiv(h, 1024))](
-        res_out, pre_out, li_out,
-        H=h, N=n, BLOCK_H=min(1024, triton.next_power_of_2(h)),
-        num_warps=4,
+    return res_out, post_out, comb_out, li_out
+
+
+def mhc_fused_post_pre_single_pass_triton(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_mix: torch.Tensor | None,
+    comb_mix: torch.Tensor | None,
+    pre_mix_in: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Single-pass mHC (DeepSeek-V4.1): same fused stages, but the layer input is mixed
+    with the PREVIOUS sublayer's ``pre_mix_in`` and this sublayer's pre gates come back
+    as ``pre_new`` for the next one. Returns (residual_new, post [T,N,1], comb [T,N,N],
+    pre_new [T,N] fp32, layer_input)."""
+    return _mhc_fused_launch(
+        x, residual, post_mix, comb_mix, fn, hc_scale, hc_base,
+        rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=pre_mix_in,
     )
-    return res_out, post_out.view(t, n, 1), comb_out, li_out
 
 
-__all__ = ["mhc_fused_post_pre_triton"]
+__all__ = ["mhc_fused_post_pre_triton", "mhc_fused_post_pre_single_pass_triton"]
diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py
index 1792da2b3c..e0cf606db1 100644
--- a/python/freetoken/kvcache/__init__.py
+++ b/python/freetoken/kvcache/__init__.py
@@ -34,7 +34,7 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]:
     specs_fn = getattr(model_config, "kv_cache_group_specs", None)
     if specs_fn is None:
         if getattr(model_config, "dsv4_args", None) is not None:
-            from .dsv4_paged_pool import DSV4PagedKVCache
+            from .dsv4.v4_pool import DSV4PagedKVCache
 
             return DSV4PagedKVCache
         from .mha_pool import MHAKVCache
@@ -43,9 +43,13 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]:
     specs = list(specs_fn())
     types = {spec.attn_type for spec in specs}
     if AttnType.DSV4 in types:
-        from .dsv4_paged_pool import DSV4PagedKVCache
+        from .dsv4.v4_pool import DSV4PagedKVCache
 
         return DSV4PagedKVCache
+    if AttnType.DSV41 in types:
+        from .dsv4.v41_pool import DSV41PagedKVCache
+
+        return DSV41PagedKVCache
     if AttnType.SWA in types:
         from .hybrid_swa_pool import HybridSWAKVCache
 
@@ -80,9 +84,9 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt
     """Build the engine's KV pool for ``num_pages`` USABLE pages (the dummy page and every
     secondary tier -- window pool, index slab, state rings -- are derived here or inside
     the pool). Single factory entry for all pool families, DSV4 included."""
-    from .dsv4_cost_model import _dsv4_pool_sizes
+    from .dsv4.v4_cost_model import _dsv4_pool_sizes
     from .hybrid_swa_pool import _naive_swa_num_tokens, _swa_paged_num_tokens
-    from .dsv4_paged_pool import DSV4PagedKVCache
+    from .dsv4.v4_pool import DSV4PagedKVCache
 
     model_config = config.model_config
     if resolve_pool_class(model_config) is DSV4PagedKVCache:
@@ -99,6 +103,21 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt
         )
         pool._init_paged_state(config.max_running_req, config.cache_type != "naive")
         return pool
+    if getattr(model_config, "dsv41_args", None) is not None:
+        from .dsv4.v41_cost_model import _dsv41_pool_sizes
+        from .dsv4.v41_pool import DSV41PagedKVCache
+
+        # Same route as DSV4: the generic CacheManager over the shared page table; the pool is the
+        # swa_pool plug-in (window tier + per-source packed main / index pools + state rings).
+        pool = DSV41PagedKVCache(
+            sizes=_dsv41_pool_sizes(config, num_pages + 1),  # +1 for dummy page
+            args=model_config.dsv41_args,
+            device=device,
+            dtype=dtype,
+            n_scratch=config.max_running_req + 1,
+        )
+        pool._init_paged_state(config.max_running_req, config.cache_type != "naive")
+        return pool
 
     num_swa_tokens = None
     # Both the naive and radix SWA paths share the global-paged swa pool; radix sizes it by
diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py
index 95669e8c81..569da1cb6c 100644
--- a/python/freetoken/kvcache/base.py
+++ b/python/freetoken/kvcache/base.py
@@ -35,6 +35,14 @@ def spec_kv_bytes_per_token(spec, config) -> int:
     return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio
 
 
+@dataclass(frozen=True)
+class WindowPoolSpec:
+    """Window resize units and minimum usable capacity, excluding reserved pages."""
+
+    page_size: int
+    min_pages: int
+
+
 class BaseKVCachePool(ABC):
     """
     Base class for key-value caches.
@@ -45,6 +53,15 @@ class BaseKVCachePool(ABC):
     # model re-bound after a rebuild; the engine asks before it resizes.
     needs_rebind_on_rebuild: ClassVar[bool] = False
 
+    @classmethod
+    def window_spec(cls, config) -> WindowPoolSpec | None:
+        return None
+
+    @property
+    def window_pages(self) -> int:
+        """Current usable window pages in the units declared by window_spec."""
+        return 0
+
     # ---- sizing/cost classmethods: run BEFORE the pool exists (startup budget solve,
     # --moe-cache-auto). The engine measures memory and passes bytes in; each pool family
     # implements kv_cost for ITS OWN buffers only (the engine sums families, e.g. adds
diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py
index 6aa13e7f75..5b531ce265 100644
--- a/python/freetoken/kvcache/cache_status.py
+++ b/python/freetoken/kvcache/cache_status.py
@@ -3,13 +3,14 @@
 from typing import Any, Dict
 
 
+def window_pool_spec(config):
+    from . import resolve_pool_class
+
+    return resolve_pool_class(config.model_config).window_spec(config)
+
+
 def _supports_swa_ratio(config) -> bool:
-    """Whether ``swa_full_tokens_ratio`` sizes a separate window pool for this model -- DSV4
-    (always) or a radix-SWA model (Gemma). Gates the ratio in telemetry and rebuild."""
-    mc = config.model_config
-    if mc.dsv4_args is not None:
-        return True
-    return mc.has_swa_attention and config.cache_type == "swa_radix"
+    return window_pool_spec(config) is not None
 
 
 def compute_cache_unit_bytes(engine: "Engine") -> Dict[str, int]:
@@ -120,18 +121,8 @@ def _mamba() -> int:
         return int(_linear_pool_min_slots(config) - 1)
 
     def _swa() -> int:
-        # Report usable tokens, matching num_swa_pages in rebuild requests.
-        # DSV4's physical floor includes a dummy page; the SWA floor already excludes slot 0.
-        from .dsv4_cost_model import _dsv4_window_floor_pages
-        from .hybrid_swa_pool import _swa_pool_floor
-
-        mc = config.model_config
-        if mc.dsv4_args is not None:
-            P = mc.dsv4_args.window_size
-            return int((_dsv4_window_floor_pages(config, P) - 1) * P)
-        if not (mc.has_swa_attention and config.cache_type == "swa_radix"):
-            return 0
-        return int(_swa_pool_floor(config))
+        spec = window_pool_spec(config)
+        return spec.page_size * spec.min_pages if spec is not None else 0
 
     for key, fn in (("kv_tokens", _kv), ("moe_experts", _moe), ("mamba_slots", _mamba),
                     ("swa_tokens", _swa)):
@@ -157,19 +148,10 @@ def compute_cache_pools(engine: "Engine") -> Dict[str, int]:
         pools["num_pages"] = int(engine.num_pages or 0)
         if config is not None:
             pools["page_size"] = int(config.page_size or 0)
-            # Window pool: its own page unit (swa_page_size -- DSV4 windows are P-token pages,
-            # radix-SWA is token-granular page_size 1) and the concrete current size in that unit
-            # (num_swa_pages, usable count). Same source as the scheduler's _current_cache_geometry.
-            # Both 0 for models without a window pool. Lets a client denominate the swa control.
-            mc = config.model_config
-            if mc.dsv4_args is not None:
-                pools["swa_page_size"] = int(mc.dsv4_args.window_size or 0)
-                sizes = getattr(engine.kv_cache, "sizes", None)  # usable = physical minus dummy
-                if sizes is not None:
-                    pools["num_swa_pages"] = max(0, int(sizes.n_win_pages) - 1)
-            elif mc.has_swa_attention and config.cache_type == "swa_radix":
-                pools["swa_page_size"] = 1  # usable = pool tokens minus the slot-0 sentinel
-                pools["num_swa_pages"] = max(0, int(getattr(engine.kv_cache, "swa_num_tokens", 0) or 0) - 1)
+            spec = window_pool_spec(config)
+            if spec is not None:
+                pools["swa_page_size"] = spec.page_size
+                pools["num_swa_pages"] = max(0, int(engine.kv_cache.window_pages))
         moe = engine.moe_offload_cache
         if moe is not None:
             pools["moe_cache_size"] = int(moe.cache_size or 0)
diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py
index 13225cec06..d7c36fd734 100644
--- a/python/freetoken/kvcache/dsa_pool.py
+++ b/python/freetoken/kvcache/dsa_pool.py
@@ -2,7 +2,7 @@
 
 ``MLAKVCache`` is the MHA pool's sibling for latent-KV models: ONE slab holding the
 per-token latent ``ckv (kv_lora_rank) | kpe (qk_rope_head_dim)`` -- there is no
-separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4_paged_pool's
+separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4/v4_pool's
 single-latent tiers). ``DSAKVCache`` extends it with the DeepSeek-Sparse-Attention
 index-key slab: one ``index_head_dim``-wide bf16 row per token per full-indexer
 layer, addressed by the SAME physical rows as the latent slab (GLM-5.2, page 1;
@@ -75,7 +75,7 @@ def k_cache(self, layer_id: int) -> torch.Tensor:
         )
 
     def v_cache(self, layer_id: int) -> torch.Tensor:
-        # MLA: K == V (single latent); same buffer, dsv4_paged_pool precedent.
+        # MLA: K == V (single latent); same buffer, dsv4/v4_pool precedent.
         return self.k_cache(layer_id)
 
     def latent_rows(self, layer_id: int) -> torch.Tensor:
diff --git a/python/freetoken/kvcache/dsv4/__init__.py b/python/freetoken/kvcache/dsv4/__init__.py
new file mode 100644
index 0000000000..7531470f0f
--- /dev/null
+++ b/python/freetoken/kvcache/dsv4/__init__.py
@@ -0,0 +1 @@
+"""Paged KV pools for the DeepSeek-V4 family."""
diff --git a/python/freetoken/kvcache/dsv4/v41_cost_model.py b/python/freetoken/kvcache/dsv4/v41_cost_model.py
new file mode 100644
index 0000000000..4db1fb18c1
--- /dev/null
+++ b/python/freetoken/kvcache/dsv4/v41_cost_model.py
@@ -0,0 +1,216 @@
+"""Cost model + per-tier sizing for the DSV41 paged KV pool (DeepSeek-V4.1).
+
+Same shape as ``v4_cost_model``: one bytes-per-P-token number for the budget division, and
+independent per-tier sizing from the anchor ``full_token = num_pages * P``. The tiers:
+
+* window pool, every layer          -- ``swa_ratio`` of the full history, packed fp8 rows
+* main KV pool, per kv source       -- ``full_token // ratio`` packed fp4 rows
+* index-key pool, per kv source     -- ``full_token // ratio`` packed fp4 rows
+* compress-state ring, ratio>1 sources -- ``ring_size`` fp32 ``kv|score`` slots per WINDOW page
+
+Sizing reads ``ModelConfig.dsv41_args`` by attribute, like ``v4_cost_model`` reads ``dsv4_args``:
+the kvcache package never imports the model package. All byte widths come from the tier row formats
+below, so there is exactly one place that knows what a row costs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32
+from .window_tier import reserved_window_pages
+
+_INT64_BYTES = 8
+_FP32_BYTES = 4
+
+WIN_FMT = FP8_E8M0_B32
+MAIN_FMT = FP4_E4M3_B16
+IDX_FMT = FP4_E8M0_B32
+
+
+def _win_row_bytes(args) -> int:
+    return WIN_FMT.row_bytes(args.head_dim)
+
+
+def _main_row_bytes(args) -> int:
+    return MAIN_FMT.row_bytes(args.head_dim)
+
+
+def _idx_row_bytes(args) -> int:
+    return IDX_FMT.row_bytes(args.index_head_dim)
+
+
+def _state_bytes(args) -> int:
+    """One compress-state ring slot: fp32 ``kv | score`` of the latent width."""
+    return 2 * args.head_dim * _FP32_BYTES
+
+
+def private_window_layer_ids(args) -> tuple[int, ...]:
+    """Layers whose window KV is REQUEST-PRIVATE: kept in a per-page-table-row ring of ``window``
+    slots (slot = row * window + pos % window) instead of the shared, page-bound window pool, so
+    it never lands in radix-shared pages. The decoder layers under Decoder SWA Bounded Replay:
+    their KV is recomputed per request over the prompt's last window (positions before it are
+    never computed), so a shared page could hold rows another request never wrote. Exact mode
+    computes every position and keeps the decoder on the shared pool."""
+    if args.swa_decoder_replay == "exact":
+        return ()
+    return tuple(range(args.decoder_start_layer, args.n_layers))
+
+
+@dataclass
+class DSV41PoolSizes:
+    """Per-tier row counts derived from the budget anchor ``full_token``; the global tiers are
+    keyed by kv-source layer id."""
+
+    P: int
+    swa_ratio: float
+    full_token: int
+    n_win_slots: int
+    n_win_pages: int
+    main_rows: dict[int, int] = field(default_factory=dict)
+    idx_rows: dict[int, int] = field(default_factory=dict)
+    state_slots: dict[int, int] = field(default_factory=dict)  # ratio>1 sources only
+
+
+def dsv41_pool_sizes(
+    num_pages: int, args, swa_ratio: float, P: int, n_win_pages: int | None = None
+) -> DSV41PoolSizes:
+    full_token = num_pages * P
+    if n_win_pages is None:
+        n_win_pages = (round(swa_ratio * full_token) + P - 1) // P
+    n_win_pages = min(n_win_pages, num_pages)
+    sizes = DSV41PoolSizes(P=P, swa_ratio=swa_ratio, full_token=full_token, n_win_slots=n_win_pages * P, n_win_pages=n_win_pages)
+    for src in args.backbone_kv_sources:
+        ratio = args.compress_ratios[src]
+        sizes.main_rows[src] = full_token // ratio
+        sizes.idx_rows[src] = full_token // ratio
+        if ratio > 1:
+            sizes.state_slots[src] = n_win_pages * ratio  # ring size == ratio
+    return sizes
+
+
+def dsv41_pool_bytes(sizes: DSV41PoolSizes, args, n_scratch: int = 1) -> int:
+    """Exact bytes a ``DSV41PagedKVCache`` built from ``sizes`` allocates (mirror of ``total_bytes``)."""
+    n_private = len(private_window_layer_ids(args))
+    total = (args.n_layers - n_private) * sizes.n_win_slots * _win_row_bytes(args)
+    total += n_private * n_scratch * sizes.P * _win_row_bytes(args)  # per-request rings
+    total += (sizes.full_token + 1) * _INT64_BYTES  # full_to_window (+ sentinel row)
+    for src in args.backbone_kv_sources:
+        total += (sizes.main_rows[src] + n_scratch) * _main_row_bytes(args)
+        total += (sizes.idx_rows[src] + n_scratch) * _idx_row_bytes(args)
+        if src in sizes.state_slots:
+            total += (sizes.state_slots[src] + 1) * _state_bytes(args)
+    return int(total)
+
+
+def dsv41_cache_per_page(args, swa_ratio: float, P: int) -> int:
+    """Marginal bytes per P-token page across all tiers (window scaled by ``swa_ratio``)."""
+    n_shared = args.n_layers - len(private_window_layer_ids(args))
+    total = n_shared * round(swa_ratio * P) * _win_row_bytes(args)
+    for src in args.backbone_kv_sources:
+        ratio = args.compress_ratios[src]
+        total += (P // ratio) * (_main_row_bytes(args) + _idx_row_bytes(args))
+        if ratio > 1:
+            total += round(swa_ratio * ratio) * _state_bytes(args)
+    return int(total)
+
+
+def dsv41_kv_unit_bytes(args, P: int) -> int:
+    """FULL-tier bytes per full-history token (main + index pools + the mapping); the slider's
+    ``kv_bytes_per_token``. 890 B/token for DeepSeek-V4.1-Flash before the mapping."""
+    per_page = P * _INT64_BYTES
+    for src in args.backbone_kv_sources:
+        per_page += (P // args.compress_ratios[src]) * (_main_row_bytes(args) + _idx_row_bytes(args))
+    return -(-per_page // P)
+
+
+def dsv41_window_unit_bytes(args, P: int) -> int:
+    """WINDOW-tier bytes per window token (sliding KV on every shared-window layer + the state rings)."""
+    per_page = (args.n_layers - len(private_window_layer_ids(args))) * P * _win_row_bytes(args)
+    for src in args.backbone_kv_sources:
+        ratio = args.compress_ratios[src]
+        if ratio > 1:
+            per_page += ratio * _state_bytes(args)
+    return -(-per_page // P)
+
+
+def dsv41_solve_num_pages(
+    available_bytes: int, args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1
+) -> DSV41PoolSizes:
+    """Largest budget-respecting pool; the window floor is honored in pages and the total is
+    byte-checked. Raises ``ValueError`` when even the minimal pool does not fit."""
+
+    def _sizes(num: int) -> DSV41PoolSizes:
+        win = max(floor_win_pages, (round(swa_ratio * num * P) + P - 1) // P)
+        return dsv41_pool_sizes(num, args, swa_ratio, P, n_win_pages=win)
+
+    lo = max(floor_win_pages, 2)
+    if dsv41_pool_bytes(_sizes(lo), args, n_scratch) > available_bytes:
+        raise ValueError(
+            f"DSV41 KV budget {available_bytes} bytes cannot fit the minimal pool ({lo} pages incl. the "
+            f"window working-set floor {floor_win_pages}); raise memory_ratio or lower max_running_req/max_seq_len"
+        )
+    hi = max(lo, available_bytes // max(1, dsv41_cache_per_page(args, 0.0, P)))
+    while dsv41_pool_bytes(_sizes(hi), args, n_scratch) <= available_bytes:
+        hi *= 2
+    while lo < hi - 1:
+        mid = (lo + hi) // 2
+        if dsv41_pool_bytes(_sizes(mid), args, n_scratch) <= available_bytes:
+            lo = mid
+        else:
+            hi = mid
+    return _sizes(lo)
+
+
+def dsv41_auto_cost_model(args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1):
+    """Affine ``(cache_per_page, fixed_cache_size, min_reserve_tokens)`` for the MoE-first auto planner."""
+    per_page = dsv41_cache_per_page(args, swa_ratio, P) + P * _INT64_BYTES
+    n0 = max(floor_win_pages, 2)
+    win0 = max(floor_win_pages, (round(swa_ratio * n0 * P) + P - 1) // P)
+    base = dsv41_pool_bytes(dsv41_pool_sizes(n0, args, swa_ratio, P, n_win_pages=win0), args, n_scratch)
+    # The engine combines this structural floor with the configured KV reserve.
+    return per_page, max(0, base - n0 * per_page), n0 * P
+
+
+# ---- config-facing sizing (EngineConfig in, sizes out) ----
+
+
+def _dsv41_swa_ratio(config) -> float:
+    return float(config.swa_full_tokens_ratio)
+
+
+def _dsv41_window_floor_pages(config, args) -> int:
+    """Minimum window pages: one prefill chunk's reach (capped at 8 pages) + the running requests'
+    working set (see ``reserved_window_pages``)."""
+    P = args.window_size
+    prefill_reach_pages = (config.max_seq_len + P - 1) // P
+    radix = config.cache_type != "naive"
+    return min(prefill_reach_pages, 8) + reserved_window_pages(config.max_running_req, radix)
+
+
+def _dsv41_pool_sizes(config, num_pages: int, num_swa_pages: int | None = None) -> DSV41PoolSizes:
+    """Sizes for ``num_pages`` PHYSICAL pages (dummy included). Window precedence: an explicit
+    ``num_swa_pages`` > ``config.swa_num_pages_override`` > ``swa_ratio`` x full, floored ONCE in pages."""
+    args = config.model_config.dsv41_args
+    P = args.window_size
+    swa_ratio = _dsv41_swa_ratio(config)
+    floor_pages = _dsv41_window_floor_pages(config, args)
+    target = num_swa_pages if num_swa_pages is not None else config.swa_num_pages_override
+    if target is not None:
+        win = min(num_pages, max(floor_pages, int(target) + 1))
+    else:
+        win = max(floor_pages, (round(swa_ratio * num_pages * P) + P - 1) // P)
+    return dsv41_pool_sizes(num_pages, args, swa_ratio, P, n_win_pages=win)
+
+
+__all__ = [
+    "DSV41PoolSizes",
+    "dsv41_auto_cost_model",
+    "dsv41_cache_per_page",
+    "dsv41_kv_unit_bytes",
+    "dsv41_pool_bytes",
+    "dsv41_pool_sizes",
+    "dsv41_solve_num_pages",
+    "dsv41_window_unit_bytes",
+    "private_window_layer_ids",
+]
diff --git a/python/freetoken/kvcache/dsv4/v41_pool.py b/python/freetoken/kvcache/dsv4/v41_pool.py
new file mode 100644
index 0000000000..36d3986ae5
--- /dev/null
+++ b/python/freetoken/kvcache/dsv4/v41_pool.py
@@ -0,0 +1,292 @@
+"""DSV41 paged KV pool (DeepSeek-V4.1): packed byte tiers over the shared page table.
+
+Tiers, sized from a budget (``v41_cost_model``) not from ``num_requests``:
+
+* ``window_pool[L]``      -- every layer; the P-sliding KV ring in ``win_fmt`` (fp8) rows, page-granular.
+* ``main_pool[src]``      -- per kv-source layer; compressed KV latents in ``main_fmt`` (fp4) rows.
+* ``idx_pool[src]``       -- per kv-source layer; indexer keys in ``idx_fmt`` (fp4) rows.
+* ``state_ring[src]``     -- ratio>1 sources; the per-window-page compress-state ring (fp32 ``kv|score``).
+
+Consumer layers (Reindex / Reuse modes) alias their source's pools: ``main_pool_of(layer)``.
+
+Addressing is the shared ``WindowTierPagedPool``'s (``kvcache/dsv4/window_tier.py``, also under the DSV4
+pool): the page table maps ``(table_idx, pos)`` to a full-token loc; ``full_to_window`` maps that to a
+window slot (page-bound by the free list); the main / index rows are pure arithmetic ``full_loc //
+ratio``; the ring slot derives from the window slot. Scratch rows past every main / index pool's
+arithmetic range give a graph-safe destination for the decode rows whose compressed group did not
+complete this step.
+
+Rows are written through ``pack_rows`` (quantize + scatter) and read inside the attention /
+indexer kernels; ``read_*`` helpers dequantize for tests and torch reference paths.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from freetoken.utils import init_logger
+
+from .v41_cost_model import (
+    IDX_FMT,
+    MAIN_FMT,
+    WIN_FMT,
+    DSV41PoolSizes,
+    _idx_row_bytes,
+    _main_row_bytes,
+    _win_row_bytes,
+    dsv41_kv_unit_bytes,
+    dsv41_window_unit_bytes,
+    private_window_layer_ids,
+)
+from .window_tier import CompressStateRing, WindowTierPagedPool
+
+logger = init_logger(__name__)
+
+
+class DSV41PagedKVCache(WindowTierPagedPool):
+    win_fmt = WIN_FMT
+    main_fmt = MAIN_FMT
+    idx_fmt = IDX_FMT
+
+    def __init__(
+        self,
+        sizes: DSV41PoolSizes,
+        args,
+        device: torch.device,
+        dtype: torch.dtype = torch.bfloat16,
+        n_scratch: int = 1,
+    ) -> None:
+        for fmt, dim in ((self.win_fmt, args.head_dim), (self.main_fmt, args.head_dim), (self.idx_fmt, args.index_head_dim)):
+            fmt.validate_dim(dim)
+        for src in args.backbone_kv_sources:
+            if args.window_size % args.compress_ratios[src]:
+                raise ValueError(f"window {args.window_size} must be a multiple of compress ratio {args.compress_ratios[src]}")
+        self.args = args
+        self.private_window_layer_ids = private_window_layer_ids(args)
+        self.sizes = sizes
+        self._device = device
+        self._dtype = dtype  # the model's compute dtype; the tiers themselves are packed bytes
+        self.P = args.window_size
+        self.head_dim = args.head_dim
+        self.index_head_dim = args.index_head_dim
+        self.n_scratch = int(n_scratch)
+        self._paged_params: tuple[int, bool] | None = None
+        self.full_loc_map: torch.Tensor | None = None
+        self._alloc_buffers()
+
+    # ----- buffers -----
+    def _alloc_buffers(self) -> None:
+        sizes, args, dev = self.sizes, self.args, self._device
+        self.full_to_window = torch.full((sizes.full_token + 1,), -1, dtype=torch.int64, device=dev)
+        # shared layers: the page-bound window pool; private layers: one ring of P slots per
+        # page-table row (n_scratch rows, the dummy included), addressed row * P + pos % P
+        self.window_pool: list[torch.Tensor] = [
+            torch.zeros(self.n_scratch * self.P if self.is_private_window(l) else sizes.n_win_slots, _win_row_bytes(args), dtype=torch.uint8, device=dev)
+            for l in range(args.n_layers)
+        ]
+        self.main_pool: dict[int, torch.Tensor] = {}
+        self.idx_pool: dict[int, torch.Tensor] = {}
+        self.state_ring: dict[int, CompressStateRing] = {}
+        self.scratch_base: dict[int, int] = {}
+        for src in args.backbone_kv_sources:
+            rows = sizes.main_rows[src]
+            self.scratch_base[src] = rows
+            self.main_pool[src] = torch.zeros(rows + self.n_scratch, _main_row_bytes(args), dtype=torch.uint8, device=dev)
+            self.idx_pool[src] = torch.zeros(sizes.idx_rows[src] + self.n_scratch, _idx_row_bytes(args), dtype=torch.uint8, device=dev)
+            if src in sizes.state_slots:
+                self.state_ring[src] = CompressStateRing(
+                    n_slots=sizes.state_slots[src], ring_size=args.compress_ratios[src], overlap=False,
+                    head_dim=args.head_dim, device=dev,
+                )
+
+    def total_bytes(self) -> int:
+        n = self.full_to_window.numel() * self.full_to_window.element_size()
+        n += sum(t.numel() for t in self.window_pool)
+        n += sum(t.numel() for t in self.main_pool.values())
+        n += sum(t.numel() for t in self.idx_pool.values())
+        n += sum(r.buffer.numel() * r.buffer.element_size() for r in self.state_ring.values())
+        return int(n)
+
+    @property
+    def prefix_replay_tokens(self) -> int:
+        """Prompt-tail tokens a prefix hit must leave to the prefill: bounded replay runs the decoder on
+        the prompt's last window from that window's encoder outputs, which no cache keeps."""
+        return self.P if self.private_window_layer_ids else 0
+
+    # ----- private window rings -----
+    def is_private_window(self, layer_id: int) -> bool:
+        return layer_id in self.private_window_layer_ids
+
+    def ring_slots(self, table_rows: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Ring slots of ``positions`` (int64, broadcastable with ``table_rows``) for private window layers;
+        negative positions map to ``-1``."""
+        slots = table_rows.to(torch.int64) * self.P + positions % self.P
+        return torch.where(positions < 0, -1, slots)
+
+    # ----- source aliasing -----
+    def source_of(self, layer_id: int) -> int:
+        src = self.args.roles[layer_id].kv_source
+        assert src is not None, f"layer {layer_id} is window-only"
+        return src
+
+    def main_pool_of(self, layer_id: int) -> torch.Tensor:
+        return self.main_pool[self.source_of(layer_id)]
+
+    def idx_pool_of(self, layer_id: int) -> torch.Tensor:
+        return self.idx_pool[self.source_of(layer_id)]
+
+    def ratio_of(self, layer_id: int) -> int:
+        return self.args.roles[layer_id].ratio
+
+    # ----- engine-facing sizing / rebuild surface -----
+    @classmethod
+    def kv_cost(cls, config) -> tuple[int, int, int, int]:
+        from .v41_cost_model import _dsv41_swa_ratio, _dsv41_window_floor_pages, dsv41_auto_cost_model
+
+        args = config.model_config.dsv41_args
+        P = args.window_size
+        per_page, fixed, min_reserve = dsv41_auto_cost_model(
+            args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P, n_scratch=config.max_running_req + 1
+        )
+        return per_page, fixed, config.page_size, min_reserve
+
+    @classmethod
+    def solve_num_pages(cls, config, available_memory: int) -> int:
+        from freetoken.utils import mem_GB
+
+        from .v41_cost_model import (
+            _dsv41_pool_sizes,
+            _dsv41_swa_ratio,
+            _dsv41_window_floor_pages,
+            dsv41_pool_bytes,
+            dsv41_solve_num_pages,
+        )
+
+        args = config.model_config.dsv41_args
+        P = args.window_size
+        num_pages = config.num_page_override
+        if num_pages is None:
+            sizes = dsv41_solve_num_pages(
+                available_memory, args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P,
+                n_scratch=config.max_running_req + 1,
+            )
+            num_pages = sizes.full_token // P - 1  # one physical page is the dummy
+        else:
+            floor = _dsv41_window_floor_pages(config, args)
+            if num_pages < floor:
+                raise ValueError(
+                    f"--num-pages {num_pages} ({num_pages * P} tokens) is below the DSV41 window working-set "
+                    f"floor {floor} pages ({floor * P} tokens); raise --num-pages or lower max_running_req/max_seq_len"
+                )
+            sizes = _dsv41_pool_sizes(config, num_pages + 1)
+        assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages"
+        real = dsv41_pool_bytes(sizes, args, config.max_running_req + 1)
+        logger.info(
+            f"Allocating {num_pages * P} tokens for DSV41 KV cache ({sizes.n_win_pages} window pages), total = {mem_GB(real)}"
+        )
+        return num_pages
+
+    @classmethod
+    def window_spec(cls, config):
+        from ..base import WindowPoolSpec
+        from .v41_cost_model import _dsv41_window_floor_pages
+
+        args = config.model_config.dsv41_args
+        return WindowPoolSpec(args.window_size, _dsv41_window_floor_pages(config, args) - 1)
+
+    @classmethod
+    def min_kv_tokens(cls, config) -> int:
+        from .v41_cost_model import _dsv41_window_floor_pages
+
+        args = config.model_config.dsv41_args
+        return _dsv41_window_floor_pages(config, args) * args.window_size
+
+    def validate_rebuild(
+        self, config, *, num_pages: int | None, target_moe: int, per_expert_bytes: int,
+        baseline_free: int, weights_bytes: int, current_num_pages: int,
+        extra_fixed_bytes: int = 0, extra_note: str = "",
+        num_swa_pages: int | None = None, **targets,
+    ) -> None:
+        from freetoken.engine.cache_budget import net_cache_budget_bytes
+        from freetoken.utils import mem_GB
+
+        from ..base import CacheRebuildRejected
+        from .v41_cost_model import _dsv41_pool_sizes, _dsv41_window_floor_pages, dsv41_pool_bytes
+
+        if num_pages is not None:
+            floor = _dsv41_window_floor_pages(config, self.args)
+            if num_pages < floor:
+                raise CacheRebuildRejected(
+                    f"num_pages {num_pages} is below the DSV41 window working-set floor {floor} "
+                    f"(max_running_req={config.max_running_req}); admission would deadlock"
+                )
+        if num_pages is not None or num_swa_pages is not None:
+            target_pages = num_pages if num_pages is not None else current_num_pages
+            kv_sizes = _dsv41_pool_sizes(config, target_pages + 1, num_swa_pages=num_swa_pages)
+        else:
+            kv_sizes = self.sizes
+        budget = net_cache_budget_bytes(config.memory_ratio, baseline_free, weights_bytes, 0)
+        need = target_moe * per_expert_bytes + dsv41_pool_bytes(kv_sizes, self.args, config.max_running_req + 1)
+        if need > budget:
+            kv_part = f"kv={num_pages} P-pages" if num_pages is not None else "kv=current pool"
+            raise CacheRebuildRejected(
+                f"requested cache (moe={target_moe} slots, {kv_part}) needs {mem_GB(need)} > budget {mem_GB(budget)}; "
+                "old cache kept, still serving"
+            )
+
+    def rebuild_from_config(self, config, num_pages: int, *, num_swa_pages: int | None = None) -> None:
+        from .v41_cost_model import _dsv41_pool_sizes
+
+        self.rebuild(_dsv41_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages))
+
+    def _drop_buffers(self) -> None:
+        self.window_pool = self.main_pool = self.idx_pool = self.state_ring = None  # type: ignore[assignment]
+
+    def unit_bytes(self) -> tuple[int, int]:
+        return dsv41_kv_unit_bytes(self.args, self.P), dsv41_window_unit_bytes(self.args, self.P)
+
+    # ----- writes (quantize + scatter) -----
+    def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None:
+        from freetoken.kernel.triton.dsv41.pack import pack_rows
+
+        pack_rows(kv, self.win_fmt, pool=self.window_pool[layer_id], row_ids=window_slot)
+
+    def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None:
+        from freetoken.kernel.triton.dsv41.pack import pack_rows
+
+        pack_rows(latent, self.main_fmt, pool=self.main_pool[source], row_ids=rows)
+
+    def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None:
+        from freetoken.kernel.triton.dsv41.pack import pack_rows
+
+        pack_rows(k, self.idx_fmt, pool=self.idx_pool[source], row_ids=rows)
+
+    # ----- reads (gather + dequantize; tests and torch reference paths) -----
+    def read_window(self, layer_id: int, window_slot: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel.triton.dsv41.pack import unpack_rows
+
+        return unpack_rows(self.window_pool[layer_id], self.win_fmt, self.head_dim, window_slot)
+
+    def read_main(self, source: int, rows: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel.triton.dsv41.pack import unpack_rows
+
+        return unpack_rows(self.main_pool[source], self.main_fmt, self.head_dim, rows)
+
+    def read_index(self, source: int, rows: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel.triton.dsv41.pack import unpack_rows
+
+        return unpack_rows(self.idx_pool[source], self.idx_fmt, self.index_head_dim, rows)
+
+    # ----- compress-state ring accessors -----
+    def get_state(self, source: int, state_loc: torch.Tensor) -> torch.Tensor:
+        return self.state_ring[source].get(state_loc)
+
+    def set_state(self, source: int, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None:
+        self.state_ring[source].set(state_loc, kv_score)
+
+    @property
+    def num_layers(self) -> int:
+        return self.args.n_layers
+
+
+__all__ = ["DSV41PagedKVCache"]
diff --git a/python/freetoken/kvcache/dsv4/v41_row_format.py b/python/freetoken/kvcache/dsv4/v41_row_format.py
new file mode 100644
index 0000000000..f273be4f56
--- /dev/null
+++ b/python/freetoken/kvcache/dsv4/v41_row_format.py
@@ -0,0 +1,69 @@
+"""RowFormat: how one DeepSeek-V4.1 KV row (a head_dim-wide vector) is stored in a packed byte pool.
+
+DeepSeek-V4.1 keeps three cache tiers in three storage formats (tech report sec. 2.4.4 and
+``inference/model.py``):
+
+* SWA KV      -- ``fp8_e8m0_b32``: e4m3 values, one ue8m0 (power-of-two) scale per 32 channels
+* main KV     -- ``fp4_e4m3_b16``: e2m1 values, one e4m3 scale per 16 channels (NVFP4 without the global scale)
+* indexer K   -- ``fp4_e8m0_b32``: e2m1 values, one ue8m0 scale per 32 channels (OCP MXFP4)
+
+A row is laid out as ``[values | scales]``: ``value_bytes(dim)`` bytes of (packed) values followed by
+``scale_bytes(dim)`` scale bytes, so one gather fetches both. e2m1 pairs pack the even channel into the low
+nibble and the odd channel into the high nibble (``torch.float4_e2m1fn_x2`` / ``convert.py`` order).
+
+The Triton side (``kernel/triton/dsv41/row_format.py``) selects the dequant path from ``RowFormat.code``
+(a constexpr); the pack / unpack wrappers in ``kernel/triton/dsv41/pack.py`` reproduce the reference
+``act_quant`` / ``fp4_act_quant`` numerics exactly, including the bf16 rounding of the dequantized value
+the reference bakes into its cache.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class RowFormat:
+    name: str
+    elem_bits: int  # 16 (bf16), 8 (e4m3) or 4 (e2m1)
+    group: int  # channels per scale; 0 = unscaled
+    scale: str | None  # None | "e8m0" | "e4m3"
+    code: int  # Triton constexpr selector
+
+    def value_bytes(self, dim: int) -> int:
+        assert (dim * self.elem_bits) % 8 == 0, (dim, self.elem_bits)
+        return dim * self.elem_bits // 8
+
+    def scale_bytes(self, dim: int) -> int:
+        if not self.group:
+            return 0
+        assert dim % self.group == 0, (dim, self.group)
+        return dim // self.group
+
+    def row_bytes(self, dim: int) -> int:
+        return self.value_bytes(dim) + self.scale_bytes(dim)
+
+    def num_groups(self, dim: int) -> int:
+        return dim // self.group if self.group else 0
+
+    def validate_dim(self, dim: int) -> None:
+        if dim <= 0 or (dim * self.elem_bits) % 8 or (self.group and dim % self.group):
+            raise ValueError(f"row format {self.name} cannot hold a {dim}-wide row")
+
+
+BF16 = RowFormat("bf16", elem_bits=16, group=0, scale=None, code=0)
+FP8_E8M0_B32 = RowFormat("fp8_e8m0_b32", elem_bits=8, group=32, scale="e8m0", code=1)
+FP4_E4M3_B16 = RowFormat("fp4_e4m3_b16", elem_bits=4, group=16, scale="e4m3", code=2)
+FP4_E8M0_B32 = RowFormat("fp4_e8m0_b32", elem_bits=4, group=32, scale="e8m0", code=3)
+
+ROW_FORMATS: dict[str, RowFormat] = {f.name: f for f in (BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32)}
+
+
+def row_format(name: str) -> RowFormat:
+    try:
+        return ROW_FORMATS[name]
+    except KeyError:
+        raise ValueError(f"unknown row format {name!r}; one of {sorted(ROW_FORMATS)}") from None
+
+
+__all__ = ["RowFormat", "BF16", "FP8_E8M0_B32", "FP4_E4M3_B16", "FP4_E8M0_B32", "ROW_FORMATS", "row_format"]
diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4/v4_cost_model.py
similarity index 96%
rename from python/freetoken/kvcache/dsv4_cost_model.py
rename to python/freetoken/kvcache/dsv4/v4_cost_model.py
index 2c084aff84..c6f5b7ce41 100644
--- a/python/freetoken/kvcache/dsv4_cost_model.py
+++ b/python/freetoken/kvcache/dsv4/v4_cost_model.py
@@ -28,14 +28,12 @@
 
 
 def dsv4_reserved_window_pages(max_running_req: int, radix: bool) -> int:
-    """Window pages the sliding pool must always keep for the concurrent working set: each
-    running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent
-    request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the
-    window is 2 pages here because the retention gap page-aligns to a whole extra page at
-    P==window==128, so a distinct follow-up per running request can re-lock 2 pages each.
-    Shared by the engine's window-floor and the manager's prefill_chunk_budget so both
-    reserve the same set."""
-    return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1
+    """Window pages the sliding pool must always keep for the concurrent working set; the
+    formula is shared with every sliding-window paged pool (``window_tier.reserved_window_pages``)
+    and by the engine's window-floor and the manager's prefill_chunk_budget."""
+    from .window_tier import reserved_window_pages
+
+    return reserved_window_pages(max_running_req, radix)
 
 
 def ring_size_for_ratio(ratio: int) -> int:
diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4/v4_pool.py
similarity index 52%
rename from python/freetoken/kvcache/dsv4_paged_pool.py
rename to python/freetoken/kvcache/dsv4/v4_pool.py
index 3ce23eea7d..d5eae6d080 100644
--- a/python/freetoken/kvcache/dsv4_paged_pool.py
+++ b/python/freetoken/kvcache/dsv4/v4_pool.py
@@ -23,135 +23,22 @@
 
 from freetoken.utils import init_logger
 
-from .base import BaseKVCachePool
-from .dsv4_cost_model import (
+from .v4_cost_model import (
     DSV4PoolSizes,
     dsv4_kv_unit_bytes,
     dsv4_window_unit_bytes,
     ring_size_for_ratio,
 )
+# The window-tier building blocks and the pool base live in window_tier.py (shared with the DSV41
+# pool); the first two are re-exported here for the DSV4 callers and tests that import them from
+# this module.
+from .window_tier import CompressStateRing, FreeListAllocator, WindowTierPagedPool  # noqa: F401
 
 
 logger = init_logger(__name__)
 
 
-# LIFO free-list allocator for the window tier: one allocated unit spans ``page_unit = P``
-# slots, so a unit base is always a multiple of ``P`` (the spec's ``G*P`` page-base
-# invariant). The cmp/idx tiers have no allocator -- their rows are ``full_loc // ratio``
-# pure arithmetic.
-class FreeListAllocator:
-    def __init__(
-        self,
-        capacity: int,
-        device: torch.device,
-        page_unit: int = 1,
-    ) -> None:
-        assert capacity % page_unit == 0, (
-            f"capacity {capacity} must be a multiple of page_unit {page_unit}"
-        )
-        self._capacity = int(capacity)
-        self._page_unit = int(page_unit)
-        self._device = device
-        # Unit base slots: [0, page_unit, 2*page_unit, ...]. LIFO -> pop the tail.
-        self._free = self._fresh_free()
-
-    def _fresh_free(self) -> torch.Tensor:
-        n_units = self._capacity // self._page_unit
-        return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit
-
-    def alloc(self, n_units: int) -> torch.Tensor:
-        """Return ``n_units`` unit base slots (each a multiple of ``page_unit``)."""
-        if n_units < 0:
-            raise ValueError(f"n_units must be non-negative, got {n_units}")
-        if n_units > self._free.numel():
-            raise RuntimeError(
-                f"FreeListAllocator out of slots: requested {n_units} units, "
-                f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})"
-            )
-        if n_units == 0:
-            return self._free[:0].clone()
-        taken = self._free[-n_units:].clone()
-        self._free = self._free[:-n_units]
-        return taken
-
-    def free(self, units: torch.Tensor) -> None:
-        """Return previously-allocated unit base slots to the pool (LIFO recycle)."""
-        if units.numel() == 0:
-            return
-        units = units.to(device=self._device, dtype=torch.int64).reshape(-1)
-        self._free = torch.cat([self._free, units])
-
-    def available(self) -> int:
-        """Free capacity in slots (units * page_unit)."""
-        return int(self._free.numel()) * self._page_unit
-
-    @property
-    def capacity(self) -> int:
-        return self._capacity
-
-    def reset(self) -> None:
-        self._free = self._fresh_free()
-
-
-
-class CompressStateRing:
-    """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``.
-
-    Last row (index ``-1``) is a permanent scratch slot, re-cleared on every
-    write. Last dim is split ``kv | score``; ``set_state`` writes both halves.
-    """
-
-    def __init__(
-        self,
-        n_slots: int,
-        ring_size: int,
-        overlap: bool,
-        head_dim: int,
-        device: torch.device,
-        dtype: torch.dtype = torch.float32,
-    ) -> None:
-        self.ring_size = ring_size
-        self.head_dim = head_dim
-        self._item_size = (1 + int(overlap)) * head_dim  # width of kv (== width of score)
-        last_dim = 2 * self._item_size
-        # +1 trailing scratch row at index -1.
-        self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device)
-        self.n_slots = n_slots
-        self._clear_scratch()
-
-    def _clear_scratch(self) -> None:
-        self.buffer[-1, : self._item_size].zero_()
-        self.buffer[-1, self._item_size :].fill_(float("-inf"))
-
-    @property
-    def item_size(self) -> int:
-        return self._item_size
-
-    def get(self, state_loc: torch.Tensor) -> torch.Tensor:
-        """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row)."""
-        return self.buffer[state_loc]
-
-    def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None:
-        """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row."""
-        self.buffer[state_loc] = kv_score
-        self._clear_scratch()
-
-    def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor:
-        """Batched per-row carry-block read. ``page_base`` is ``[B]`` (the ring block base
-        row per row, == ``(window_slot // P) * ring_size``). Returns ``[B, ring_size, 2*item]``
-        — each row's whole rolling carry block. Distinct pages -> disjoint blocks (isolation)."""
-        rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device)
-        return self.buffer[rows]
-
-    def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None:
-        """Batched per-row carry-block write: ``buffer[base+arange(ring_size)] = blocks[row]``.
-        ``page_base`` ``[B]``, ``blocks`` ``[B, ring_size, 2*item]``. Re-clears scratch."""
-        rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device)
-        self.buffer[rows] = blocks
-        self._clear_scratch()
-
-
-class DSV4PagedKVCache(BaseKVCachePool):
+class DSV4PagedKVCache(WindowTierPagedPool):
     def __init__(
         self,
         sizes: DSV4PoolSizes,
@@ -273,50 +160,11 @@ def _alloc_buffers(self) -> None:
                 )
             )
 
-    # ----- full-loc translation (gather-only -1 safety; see field comment) -----
-    def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor:
-        # int64 gather indices: the shared page_table stores full locs as int32.
-        return self.full_to_window[full_locs.to(dtype=torch.int64)]
-
-    @staticmethod
-    def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor:
-        return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor")
-
-    def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None:
-        assert full_page_base % self.P == 0 and window_page_base % self.P == 0
-        self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange(
-            window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device
-        )
-
-    def unbind_window_pages(self, full_locs: torch.Tensor) -> None:
-        self.full_to_window[full_locs[full_locs >= 0]] = -1
-
-    # ----- generic swa_pool duck-type (the CacheManager plug-in surface) -----
-    # ShadowRadix layering: the shared page_table is the virtual full-token coordinate; this pool
-    # projects it into the physical tiers. The window tier is the managed "second currency" --
-    # token-face signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window
-    # pages are 1:1 page-bound to full pages; the per-page state ring requires it). alloc_swa
-    # receives whole ascending pages (allocate_paged's _page_to_token expansion) and free paths
-    # are page-complete by construction (padded finish tails, align_down frontiers, page-aligned
-    # tree nodes) -- asserted here, not assumed.
-    swa_paged = True
-
-    @property
-    def sliding_window_size(self) -> int:
-        return self.P
-
-    @property
-    def prefill_chunk_budget(self) -> int:
-        return self._chunk_budget
-
     # ----- engine-facing rebuild surface -----
-    # The tier buffers are bound into per-forward model scratch, invalid after a realloc.
-    needs_rebind_on_rebuild = True
-
     @classmethod
     def kv_cost(cls, config) -> tuple[int, int, int, int]:
-        from .dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages
-        from .dsv4_cost_model import dsv4_auto_cost_model
+        from .v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages
+        from .v4_cost_model import dsv4_auto_cost_model
 
         dsv4_args = config.model_config.dsv4_args
         P = dsv4_args.window_size
@@ -335,8 +183,8 @@ def solve_num_pages(cls, config, available_memory: int) -> int:
         # graceful config error, not a late OOM.
         from freetoken.utils import mem_GB
 
-        from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages
-        from .dsv4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages
+        from .v4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages
+        from .v4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages
 
         dsv4_args = config.model_config.dsv4_args
         P = dsv4_args.window_size
@@ -369,11 +217,19 @@ def solve_num_pages(cls, config, available_memory: int) -> int:
         )
         return num_pages
 
+    @classmethod
+    def window_spec(cls, config):
+        from ..base import WindowPoolSpec
+        from .v4_cost_model import _dsv4_window_floor_pages
+
+        P = config.model_config.dsv4_args.window_size
+        return WindowPoolSpec(P, _dsv4_window_floor_pages(config, P) - 1)
+
     @classmethod
     def min_kv_tokens(cls, config) -> int:
         # The full anchor must cover the window working-set floor (full >= window always), so
         # that floor -- the value validate_rebuild enforces -- is the pool's floor in tokens.
-        from .dsv4_cost_model import _dsv4_window_floor_pages
+        from .v4_cost_model import _dsv4_window_floor_pages
 
         P = config.model_config.dsv4_args.window_size
         return _dsv4_window_floor_pages(config, P) * P
@@ -387,9 +243,9 @@ def validate_rebuild(
         from freetoken.engine.cache_budget import net_cache_budget_bytes
         from freetoken.utils import mem_GB
 
-        from .base import CacheRebuildRejected
-        from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages
-        from .dsv4_cost_model import dsv4_pool_bytes
+        from ..base import CacheRebuildRejected
+        from .v4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages
+        from .v4_cost_model import dsv4_pool_bytes
 
         dsv4_args = config.model_config.dsv4_args
         if num_pages is not None:
@@ -426,120 +282,19 @@ def validate_rebuild(
     def rebuild_from_config(
         self, config, num_pages: int, *, num_swa_pages: int | None = None
     ) -> None:
-        from .dsv4_cost_model import _dsv4_pool_sizes
+        from .v4_cost_model import _dsv4_pool_sizes
 
         # +1 for the dummy page
         self.rebuild(_dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages))
 
-    def attach_page_table(self, page_table: torch.Tensor) -> None:
-        # The model reads full locs through full_loc_map; under the shared route that IS the
-        # page table. The attention backend re-allocates its decode snapshot on re-capture.
-        self.full_loc_map = page_table
-
     def unit_bytes(self) -> tuple[int, int]:
         # No measurable flat buffer (owned paged pool): the full (cmp/idx + mapping) and window
         # (sliding KV + state rings) per-token costs come from the per-tier cost model.
         return dsv4_kv_unit_bytes(self.args, self.P), dsv4_window_unit_bytes(self.args, self.P)
 
-    def rebuild(self, sizes) -> None:
-        """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's
-        tree/page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine
-        re-attaches the page table via attach_page_table afterwards."""
-        import gc
-
-        assert self._paged_params is not None, "rebuild before _init_paged_state"
-        self.sizes = sizes
+    def _drop_buffers(self) -> None:
         self.window_pool = self.cmp_pool = self.idx_pool = None
         self.state_ring = self.indexer_state_ring = None
-        self.full_to_window = None
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        self._alloc_buffers()
-        self._init_paged_state(*self._paged_params)
-
-    def _init_paged_state(self, max_running_req: int, radix: bool) -> None:
-        """Build the pool-owned window free-list + the tail dummy binding. The LAST full page and
-        LAST window page are the reserved dummy region: page_table's dummy row points at
-        ``full_token - P`` (== the generic ``fill_(num_tokens)`` convention with num_tokens = the
-        allocatable token count), permanently bound so graph-padded rows scatter to a real slot."""
-        from .dsv4_cost_model import dsv4_reserved_window_pages
-        P = self.P
-        self._paged_params = (int(max_running_req), bool(radix))
-        self.full_to_window.fill_(-1)
-        self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P)
-        self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P)
-        # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees
-        # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and
-        # halve the rest -- the same formula the bespoke manager used.
-        n_win_pages = (self.sizes.n_win_slots // P) - 1
-        reserved = dsv4_reserved_window_pages(max_running_req, radix)
-        self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P)
-
-    @property
-    def swa_num_tokens(self) -> int:
-        # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a
-        # sentinel (cap == swa_num_tokens - 1); DSV4's reserved unit is the tail dummy page,
-        # already excluded from the free-list, so +1 re-encodes the same cap.
-        return (self.sizes.n_win_slots - self.P) + 1
-
-    def swa_available_size(self) -> int:
-        return int(self._win_alloc.available())
-
-    def alloc_swa(self, full_indices: torch.Tensor) -> None:
-        """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending
-        pages (the ``_page_to_token`` expansion); the in-page offsets are preserved
-        (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires."""
-        n = int(full_indices.numel())
-        if n == 0:
-            return
-        P = self.P
-        assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots"
-        fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P)
-        fbases = fi[:, 0]
-        assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), (
-            "alloc_swa pages must be contiguous ascending"
-        )
-        wbases = self._win_alloc.alloc(fbases.numel())  # raises when exhausted (caller gated)
-        offsets = torch.arange(P, dtype=torch.int64, device=self._device)
-        self.full_to_window[(fbases[:, None] + offsets).flatten()] = (
-            wbases[:, None] + offsets
-        ).flatten()
-
-    def free_swa(self, full_indices: torch.Tensor) -> None:
-        """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic:
-        the incoming locs must cover each touched page completely (guaranteed by the padded
-        finish tails / aligned frontiers / page-aligned tree values). Idempotent over already
-        unbound (slid/tombstoned) pages."""
-        if full_indices.numel() == 0:
-            return
-        P = self.P
-        fi = full_indices.to(device=self._device, dtype=torch.int64)
-        fi = fi[fi >= 0]
-        if fi.numel() == 0:
-            return
-        fbases, counts = torch.unique(
-            torch.div(fi, P, rounding_mode="floor") * P, return_counts=True
-        )
-        assert bool((counts == P).all()), (
-            f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})"
-        )
-        ws = self.full_to_window[fbases]
-        live = ws[ws >= 0]
-        offsets = torch.arange(P, dtype=torch.int64, device=self._device)
-        self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1
-        if live.numel():
-            self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P)
-
-    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
-        return self.full_to_window[kv_indices.to(dtype=torch.int64)]
-
-    # ----- state_loc derivation (vectorized, LongTensor in/out) -----
-    @staticmethod
-    def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor:
-        pages = torch.div(window_slot, P, rounding_mode="floor")
-        loc = pages * ring_size + (window_slot % ring_size)
-        return torch.where(window_slot < 0, torch.full_like(loc, -1), loc)
 
     def ring_size(self, layer_id: int) -> int:
         return ring_size_for_ratio(self.compress_ratios[layer_id])
@@ -586,26 +341,6 @@ def total_bytes(self) -> int:
         )
         return int(n)
 
-    # ----- BaseKVCachePool interface -----
-    def k_cache(self, index: int) -> torch.Tensor:
-        return self.window_pool[index]
-
-    def v_cache(self, index: int) -> torch.Tensor:  # MLA: K == V (single latent)
-        return self.window_pool[index]
-
-    def store_kv(self, k, v, out_loc, layer_id) -> None:
-        # Thin window-write shim for ABC compat; DSV4 writes via the specialized
-        # setters above. K == V (single latent), so the window slot is out_loc.
-        self.store_window(k, layer_id, out_loc)
-
-    @property
-    def device(self) -> torch.device:
-        return self._device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self._dtype
-
     @property
     def num_layers(self) -> int:
         return self._n_layers
diff --git a/python/freetoken/kvcache/dsv4/window_tier.py b/python/freetoken/kvcache/dsv4/window_tier.py
new file mode 100644
index 0000000000..0a25c05872
--- /dev/null
+++ b/python/freetoken/kvcache/dsv4/window_tier.py
@@ -0,0 +1,331 @@
+"""Shared pieces of the sliding-window paged pools (DSV4, DSV41).
+
+* ``FreeListAllocator``   -- LIFO page allocator for the window tier (unit = one P-slot page).
+* ``CompressStateRing``   -- per-window-page fp32 ``kv | score`` carry ring with a scratch row.
+* ``reserved_window_pages`` -- the window working set every such pool keeps for its running requests.
+* ``window_state_loc``    -- the derived (never stored) ring slot of a window slot.
+* ``WindowTierPagedPool`` -- the pool base: the full-loc -> window-slot mapping over the shared page
+  table, the page-atomic window allocator behind the CacheManager's ``swa_pool`` duck-type, the
+  in-place rebuild, and the ABC glue every such pool shares. Subclasses allocate their own tiers
+  (``_alloc_buffers`` / ``_drop_buffers``), size themselves (the ``kv_cost`` family) and write rows.
+
+``DSV4PagedKVCache`` (the origin of this code) and ``DSV41PagedKVCache`` both build on these, so a
+fix to allocation, freeing or the mapping lands in one place.
+"""
+
+from __future__ import annotations
+
+import gc
+from abc import abstractmethod
+
+import torch
+
+from ..base import BaseKVCachePool
+
+
+class FreeListAllocator:
+    """LIFO free list over ``capacity // page_unit`` units; a unit's base slot is a multiple of
+    ``page_unit`` (the window tier's ``G*P`` page-base invariant)."""
+
+    def __init__(self, capacity: int, device: torch.device, page_unit: int = 1) -> None:
+        assert capacity % page_unit == 0, f"capacity {capacity} must be a multiple of page_unit {page_unit}"
+        self._capacity = int(capacity)
+        self._page_unit = int(page_unit)
+        self._device = device
+        self._free = self._fresh_free()
+
+    def _fresh_free(self) -> torch.Tensor:
+        n_units = self._capacity // self._page_unit
+        return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit
+
+    def alloc(self, n_units: int) -> torch.Tensor:
+        """Return ``n_units`` unit base slots (each a multiple of ``page_unit``)."""
+        if n_units < 0:
+            raise ValueError(f"n_units must be non-negative, got {n_units}")
+        if n_units > self._free.numel():
+            raise RuntimeError(
+                f"FreeListAllocator out of slots: requested {n_units} units, "
+                f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})"
+            )
+        if n_units == 0:
+            return self._free[:0].clone()
+        taken = self._free[-n_units:].clone()
+        self._free = self._free[:-n_units]
+        return taken
+
+    def free(self, units: torch.Tensor) -> None:
+        """Return previously-allocated unit base slots to the pool (LIFO recycle)."""
+        if units.numel() == 0:
+            return
+        units = units.to(device=self._device, dtype=torch.int64).reshape(-1)
+        self._free = torch.cat([self._free, units])
+
+    def available(self) -> int:
+        """Free capacity in slots (units * page_unit)."""
+        return int(self._free.numel()) * self._page_unit
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    def reset(self) -> None:
+        self._free = self._fresh_free()
+
+
+class CompressStateRing:
+    """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``.
+
+    Last row (index ``-1``) is a permanent scratch slot, re-cleared on every write. Last dim is
+    split ``kv | score``; ``set`` writes both halves.
+    """
+
+    def __init__(
+        self,
+        n_slots: int,
+        ring_size: int,
+        overlap: bool,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        self.ring_size = ring_size
+        self.head_dim = head_dim
+        self._item_size = (1 + int(overlap)) * head_dim  # width of kv (== width of score)
+        last_dim = 2 * self._item_size
+        self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device)
+        self.n_slots = n_slots
+        self._clear_scratch()
+
+    def _clear_scratch(self) -> None:
+        self.buffer[-1, : self._item_size].zero_()
+        self.buffer[-1, self._item_size :].fill_(float("-inf"))
+
+    @property
+    def item_size(self) -> int:
+        return self._item_size
+
+    def get(self, state_loc: torch.Tensor) -> torch.Tensor:
+        """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row)."""
+        return self.buffer[state_loc]
+
+    def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None:
+        """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row."""
+        self.buffer[state_loc] = kv_score
+        self._clear_scratch()
+
+    def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor:
+        """Batched per-row carry-block read: ``page_base`` ``[B]`` (``(window_slot // P) * ring_size``)
+        -> ``[B, ring_size, 2*item]``. Distinct pages -> disjoint blocks."""
+        rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device)
+        return self.buffer[rows]
+
+    def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None:
+        """Batched per-row carry-block write of ``[B, ring_size, 2*item]``; re-clears scratch."""
+        rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device)
+        self.buffer[rows] = blocks
+        self._clear_scratch()
+
+
+def reserved_window_pages(max_running_req: int, radix: bool) -> int:
+    """Window pages the sliding pool must always keep for the concurrent working set: each
+    running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent
+    request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the
+    window is 2 pages here because the retention gap page-aligns to a whole extra page at
+    P == window."""
+    return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1
+
+
+def window_state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor:
+    """Ring slot of a window slot: ``(ws // P) * ring_size + ws % ring_size``; ``-1`` stays ``-1``
+    (the scratch row). ``ring_size | P`` keeps distinct pages on disjoint blocks."""
+    pages = torch.div(window_slot, P, rounding_mode="floor")
+    loc = pages * ring_size + (window_slot % ring_size)
+    return torch.where(window_slot < 0, torch.full_like(loc, -1), loc)
+
+
+class WindowTierPagedPool(BaseKVCachePool):
+    """A paged KV pool whose window tier is page-bound to the shared full-token page table.
+
+    Subclasses set ``P`` (the window page), ``_device``, ``_dtype`` and ``sizes`` (with
+    ``full_token`` and ``n_win_slots``) before calling ``_alloc_buffers``, which must create
+    ``self.full_to_window`` (``[full_token + 1]`` int64, ``-1`` = unmapped, the trailing row a
+    permanent ``-1`` sentinel so a gather at ``-1`` is safe) and ``self.window_pool`` (one tensor
+    per layer, indexed by window slot). Scatter paths must never see a negative slot.
+
+    ShadowRadix layering: the shared page table is the virtual full-token coordinate; the pool
+    projects it into physical tiers. The window tier is the managed "second currency" -- token-face
+    signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window pages are 1:1
+    page-bound to full pages; the per-page state ring requires it). ``alloc_swa`` receives whole
+    ascending pages and free paths are page-complete by construction -- asserted here, not assumed.
+    """
+
+    swa_paged = True
+    # the tier buffers are bound into per-forward model scratch, invalid after a realloc
+    needs_rebind_on_rebuild = True
+
+    P: int
+    _device: torch.device
+    _dtype: torch.dtype
+    sizes: object
+    full_to_window: torch.Tensor
+    window_pool: list
+
+    # ----- tiers (subclass) -----
+    @abstractmethod
+    def _alloc_buffers(self) -> None: ...
+
+    @abstractmethod
+    def _drop_buffers(self) -> None:
+        """Release every tier buffer before a rebuild reallocates them."""
+
+    # ----- full-loc translation (gather-only -1 safety; see the class docstring) -----
+    def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor:
+        # int64 gather indices: the shared page_table stores full locs as int32
+        return self.full_to_window[full_locs.to(dtype=torch.int64)]
+
+    @staticmethod
+    def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor:
+        """The compressed row of a full loc: ``full_loc // ratio`` (negative stays negative)."""
+        return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor")
+
+    def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None:
+        assert full_page_base % self.P == 0 and window_page_base % self.P == 0
+        self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange(
+            window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device
+        )
+
+    def unbind_window_pages(self, full_locs: torch.Tensor) -> None:
+        self.full_to_window[full_locs[full_locs >= 0]] = -1
+
+    @staticmethod
+    def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor:
+        return window_state_loc(window_slot, ring_size, P)
+
+    # ----- generic swa_pool duck-type (the CacheManager plug-in surface) -----
+    @property
+    def sliding_window_size(self) -> int:
+        return self.P
+
+    @property
+    def prefill_chunk_budget(self) -> int:
+        return self._chunk_budget
+
+    def _init_paged_state(self, max_running_req: int, radix: bool) -> None:
+        """Build the pool-owned window free list + the tail dummy binding. The LAST full page and
+        LAST window page are the reserved dummy region: page_table's dummy row points at
+        ``full_token - P`` (the generic ``fill_(num_tokens)`` convention), permanently bound so
+        graph-padded rows scatter to a real slot."""
+        P = self.P
+        self._paged_params = (int(max_running_req), bool(radix))
+        self.full_to_window.fill_(-1)
+        self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P)
+        self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P)
+        # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees
+        # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and
+        # halve the rest.
+        n_win_pages = (self.sizes.n_win_slots // P) - 1
+        reserved = reserved_window_pages(max_running_req, radix)
+        self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P)
+
+    @property
+    def window_pages(self) -> int:
+        return self.sizes.n_win_pages - 1
+
+    @property
+    def swa_num_tokens(self) -> int:
+        # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a
+        # sentinel (cap == swa_num_tokens - 1); here the reserved unit is the tail dummy page,
+        # already excluded from the free list, so +1 re-encodes the same cap.
+        return (self.sizes.n_win_slots - self.P) + 1
+
+    def swa_available_size(self) -> int:
+        return int(self._win_alloc.available())
+
+    def alloc_swa(self, full_indices: torch.Tensor) -> None:
+        """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending
+        pages (the ``_page_to_token`` expansion); the in-page offsets are preserved
+        (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires."""
+        n = int(full_indices.numel())
+        if n == 0:
+            return
+        P = self.P
+        assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots"
+        fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P)
+        fbases = fi[:, 0]
+        assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), "alloc_swa pages must be contiguous ascending"
+        wbases = self._win_alloc.alloc(fbases.numel())  # raises when exhausted (caller gated)
+        offsets = torch.arange(P, dtype=torch.int64, device=self._device)
+        self.full_to_window[(fbases[:, None] + offsets).flatten()] = (wbases[:, None] + offsets).flatten()
+
+    def free_swa(self, full_indices: torch.Tensor) -> None:
+        """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic: the
+        incoming locs must cover each touched page completely (guaranteed by the padded finish
+        tails / aligned frontiers / page-aligned tree values). Idempotent over already unbound
+        (slid / tombstoned) pages."""
+        if full_indices.numel() == 0:
+            return
+        P = self.P
+        fi = full_indices.to(device=self._device, dtype=torch.int64)
+        fi = fi[fi >= 0]
+        if fi.numel() == 0:
+            return
+        fbases, counts = torch.unique(torch.div(fi, P, rounding_mode="floor") * P, return_counts=True)
+        assert bool((counts == P).all()), f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})"
+        ws = self.full_to_window[fbases]
+        live = ws[ws >= 0]
+        offsets = torch.arange(P, dtype=torch.int64, device=self._device)
+        self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1
+        if live.numel():
+            self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P)
+
+    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
+        return self.full_to_window[kv_indices.to(dtype=torch.int64)]
+
+    # ----- rebuild / page table -----
+    def rebuild(self, sizes) -> None:
+        """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's
+        tree / page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine
+        re-attaches the page table via attach_page_table afterwards."""
+        assert self._paged_params is not None, "rebuild before _init_paged_state"
+        self.sizes = sizes
+        self._drop_buffers()
+        self.full_to_window = None  # type: ignore[assignment]
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        self._alloc_buffers()
+        self._init_paged_state(*self._paged_params)
+
+    def attach_page_table(self, page_table: torch.Tensor) -> None:
+        # the model reads full locs through full_loc_map; under the shared route that IS the page table
+        self.full_loc_map = page_table
+
+    # ----- BaseKVCachePool glue: the window tier is the K == V latent cache -----
+    def k_cache(self, index: int) -> torch.Tensor:
+        return self.window_pool[index]
+
+    def v_cache(self, index: int) -> torch.Tensor:
+        return self.window_pool[index]
+
+    def store_kv(self, k, v, out_loc, layer_id) -> None:
+        self.store_window(k, layer_id, out_loc)
+
+    @abstractmethod
+    def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: ...
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self._dtype
+
+
+__all__ = [
+    "FreeListAllocator",
+    "CompressStateRing",
+    "WindowTierPagedPool",
+    "reserved_window_pages",
+    "window_state_loc",
+]
diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py
index 41c3880e34..022eb49918 100644
--- a/python/freetoken/kvcache/hybrid_swa_pool.py
+++ b/python/freetoken/kvcache/hybrid_swa_pool.py
@@ -8,7 +8,7 @@
 from freetoken.models.config import KVCacheGroupSpec
 from freetoken.utils import align_ceil, div_even
 
-from .base import BaseKVCachePool
+from .base import BaseKVCachePool, WindowPoolSpec
 
 
 @dataclass(frozen=True)
@@ -295,6 +295,16 @@ def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> Non
             # by CacheManager.rebuild is then consistent by construction (idle-only).
             self._init_swa_paged_state()
 
+    @classmethod
+    def window_spec(cls, config) -> WindowPoolSpec | None:
+        if config.cache_type != "swa_radix":
+            return None
+        return WindowPoolSpec(1, _swa_pool_floor(config))
+
+    @property
+    def window_pages(self) -> int:
+        return self.swa_num_tokens - 1
+
     @classmethod
     def kv_cost(cls, config, *, num_swa_pages: int | None = None) -> tuple[int, int, int, int]:
         """Full groups ride cache_per_page; the window pool is either a pinned absolute
diff --git a/python/freetoken/layers/mhc.py b/python/freetoken/layers/mhc.py
index 1cd6459d7b..dcf4915140 100644
--- a/python/freetoken/layers/mhc.py
+++ b/python/freetoken/layers/mhc.py
@@ -19,6 +19,16 @@
 correctness baseline; the fused triton kernel (``kernel/triton/mhc.py``,
 dispatched by ``mhc_fused_post_pre`` below) replaces it on CUDA and is
 validated against this file.
+
+Two variants share the gate math:
+
+* **multi-pass** (GLM-5.3, DeepSeek-V4): ``pre`` is computed from the residual the
+  sublayer is about to read and applied immediately -- ``mhc_pre`` / ``mhc_fused_post_pre``.
+* **single-pass** (DeepSeek-V4.1 "Single-Pass mHC", tech report sec. 2.4.1): every
+  sublayer consumes the ``pre`` gates the *previous* sublayer produced, so the input
+  mix no longer waits on this sublayer's own reduction -- ``mhc_pre_single_pass`` /
+  ``mhc_fused_post_pre_single_pass``. The returned ``pre_new`` is what the next
+  sublayer applies; the model seeds the chain with a one-hot pre on stream 0.
 """
 
 from __future__ import annotations
@@ -36,8 +46,8 @@ def hc_contract(x: torch.Tensor) -> torch.Tensor:
     return x.mean(dim=1)
 
 
-def mhc_pre(
-    residual: torch.Tensor,  # [T, n, hidden] bf16
+def mhc_gates(
+    residual: torch.Tensor,  # [T, n, hidden]
     fn: torch.Tensor,  # [2n + n^2, n*hidden] fp32
     hc_scale: torch.Tensor,  # [3] fp32
     hc_base: torch.Tensor,  # [2n + n^2] fp32
@@ -46,8 +56,8 @@ def mhc_pre(
     post_mult: float,
     sinkhorn_repeat: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Returns (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32,
-    layer_input [T, hidden] bf16)."""
+    """The coefficient predictor H(X): returns (pre_mix [T, n] fp32, post_mix [T, n, 1]
+    fp32, comb_mix [T, n, n] fp32) from one fp32 GEMM over the flattened streams."""
     n, hidden = residual.shape[-2], residual.shape[-1]
     t = residual.shape[0]
 
@@ -68,11 +78,46 @@ def mhc_pre(
     for _ in range(sinkhorn_repeat - 1):
         comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_eps)
         comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps)
+    return pre_mix, post_mix.view(t, n, 1), comb
 
-    layer_input = (
-        (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype)
-    )
-    return post_mix.view(t, n, 1), comb, layer_input
+
+def mhc_mix_input(residual: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
+    """The sublayer input ``sum_i pre_i * res_i``: [T, n, hidden] x [T, n] -> [T, hidden]."""
+    return (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype)
+
+
+def mhc_pre(
+    residual: torch.Tensor,  # [T, n, hidden] bf16
+    fn: torch.Tensor,  # [2n + n^2, n*hidden] fp32
+    hc_scale: torch.Tensor,  # [3] fp32
+    hc_base: torch.Tensor,  # [2n + n^2] fp32
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Multi-pass: gates from ``residual``, pre applied to it right away. Returns
+    (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden] bf16)."""
+    pre_mix, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat)
+    return post_mix, comb, mhc_mix_input(residual, pre_mix)
+
+
+def mhc_pre_single_pass(
+    residual: torch.Tensor,  # [T, n, hidden]
+    pre_mix_in: torch.Tensor,  # [T, n] fp32 -- the previous sublayer's pre gates
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Single-pass: gates from ``residual`` for THIS sublayer's post/comb and the NEXT
+    sublayer's pre; the input is mixed with ``pre_mix_in``. Returns (pre_new [T, n] fp32,
+    post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden])."""
+    pre_new, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat)
+    return pre_new, post_mix, comb, mhc_mix_input(residual, pre_mix_in.to(torch.float32))
 
 
 def mhc_post(
@@ -140,11 +185,69 @@ def mhc_fused_post_pre(
     )
 
 
+def mhc_fused_post_pre_single_pass_torch(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_mix: torch.Tensor | None,
+    comb_mix: torch.Tensor | None,
+    pre_mix_in: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Decomposed single-pass reference: the previous sublayer's hc_post (skipped when
+    ``post_mix is None`` -- the model's first sublayer), then this sublayer's gates with
+    the input mixed by ``pre_mix_in``. Returns (residual_new, post_new [T, n, 1],
+    comb_new [T, n, n], pre_new [T, n], layer_input)."""
+    residual_new = residual if post_mix is None else mhc_post(x, residual, post_mix, comb_mix)
+    pre_new, post_new, comb_new, layer_input = mhc_pre_single_pass(
+        residual_new, pre_mix_in, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat
+    )
+    return residual_new, post_new, comb_new, pre_new, layer_input
+
+
+def mhc_fused_post_pre_single_pass(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_mix: torch.Tensor | None,
+    comb_mix: torch.Tensor | None,
+    pre_mix_in: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+    post_mult: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Single-pass sublayer boundary (DeepSeek-V4.1): fused triton on CUDA, torch elsewhere."""
+    if residual.is_cuda:
+        from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton
+
+        return mhc_fused_post_pre_single_pass_triton(
+            x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base,
+            rms_eps, hc_eps, post_mult, sinkhorn_repeat,
+        )
+    return mhc_fused_post_pre_single_pass_torch(
+        x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base,
+        rms_eps, hc_eps, post_mult, sinkhorn_repeat,
+    )
+
+
 __all__ = [
     "hc_expand",
     "hc_contract",
+    "mhc_gates",
+    "mhc_mix_input",
     "mhc_pre",
+    "mhc_pre_single_pass",
     "mhc_post",
     "mhc_fused_post_pre",
     "mhc_fused_post_pre_torch",
+    "mhc_fused_post_pre_single_pass",
+    "mhc_fused_post_pre_single_pass_torch",
 ]
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 996beaa796..4918a032fe 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -308,18 +308,21 @@ def _decode_hybrid(
         The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so
         the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the
         routing split is device-side elementwise and the CPU submit/sync are host nodes.
-        Each route is computed exactly once -- the GPU weights are zeroed for CPU-assigned
-        routes and the CPU ids are -1 for GPU-assigned routes (the C++ kernel skips id<0).
+        Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with
+        weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads
+        slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for
+        GPU-assigned routes (the C++ kernel skips id<0).
         """
         executor = cache.cpu_executor
         assert executor is not None, "CPU MoE executor was not initialized"
-        raw = topk_ids.clone()  # raw expert ids for the CPU partial
-        cache.ensure_experts_hybrid(self.layer_id, topk_ids)  # -> slot (hit/fetched) or -1
+        raw_ids = topk_ids.clone()
+        cache.ensure_experts_hybrid(self.layer_id, topk_ids)
+        on_gpu = topk_ids >= 0
+        cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous()
+        gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous()
+        gpu_slots = topk_ids
         if cache.collect_stats:
             cache.record_decode_stats_hybrid(self.layer_id)
-        on_gpu = topk_ids >= 0
-
-        cpu_ids = torch.where(on_gpu, raw.new_full((), -1), raw).contiguous()
         pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids)
 
         # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the
@@ -329,8 +332,10 @@ def _decode_hybrid(
         )
 
         cache.copy_missing()
-        gpu_slots = topk_ids.clamp_min(0)  # -1 -> slot 0 (zero-weighted below)
-        gpu_w = torch.where(on_gpu, topk_weights, topk_weights.new_zeros(())).contiguous()
+        if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots):
+            # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
+            # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
+            gpu_slots = gpu_slots.clamp_min(0)
         gpu_routed = self._expert_gemm(
             cache,
             hidden_states,
diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py
index e84dcbcbba..a64136103b 100644
--- a/python/freetoken/layers/quantization/configs/base.py
+++ b/python/freetoken/layers/quantization/configs/base.py
@@ -89,9 +89,13 @@ def get_quant_method(self, layer: Any, prefix: str):
         layer_kind = layer.quant_layer_kind
         kind = scheme.kind if scheme else QuantKind.NONE
         cls = method_class(kind, layer_kind)
-        cfg = LAYER_CONFIGS[layer_kind].from_layer(layer, scheme)
+        cfg = self.layer_config(layer, layer_kind, scheme)
         return cls(cfg, get_quant_backend().select(layer_kind, kind))
 
+    def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig:
+        """The layer's Linear / MoE config; a dialect overrides it to add what only its checkpoint knows."""
+        return LAYER_CONFIGS[layer_kind].from_layer(layer, scheme)
+
     @staticmethod
     def from_hf(
         hf_config: Any,
diff --git a/python/freetoken/layers/quantization/configs/fp8.py b/python/freetoken/layers/quantization/configs/fp8.py
index c23cbfdee8..3e74f0c4d0 100644
--- a/python/freetoken/layers/quantization/configs/fp8.py
+++ b/python/freetoken/layers/quantization/configs/fp8.py
@@ -1,20 +1,25 @@
 from __future__ import annotations
 
+import dataclasses
 from typing import Any, ClassVar
 
+from ..linear import LinearConfig
+from ..moe import MoEConfig
 from ..names import is_routed_expert, name_set, substr_set
-from ..registry import register_dialect
+from ..registry import LayerKind, register_dialect
 from ..scheme import QuantKind, QuantScheme
-from ..scheme import FP8_BLOCK, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme
+from ..scheme import FP8_BLOCK_SIZES, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme
 from .base import QuantConfig, Stored, cfg_get
 
 
 @register_dialect
 class Fp8BlockConfig(QuantConfig):
-    """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert variant."""
+    """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert
+    variant; DeepSeek-V4.1 keeps that dialect with 32x32 blocks (``weight_block_size: [32, 32]``)."""
 
     dialect = "fp8"
 
+    # the 128-block schemes every dialect table names; ``block_scheme`` carries the checkpoint's actual block edge
     SCHEMES: ClassVar[dict[str, QuantScheme]] = {
         "BLOCK": fp8_block_scheme("float"),
         "BLOCK_E8M0": fp8_block_scheme("e8m0"),
@@ -32,15 +37,23 @@ class Fp8BlockConfig(QuantConfig):
     def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()):
         super().__init__(name_map, unquantized)
         block = tuple(int(x) for x in (q.get("weight_block_size") or ()))
-        if q.get("weight_per_tensor") or block != (FP8_BLOCK, FP8_BLOCK):
-            raise NotImplementedError(f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; only 128x128 blocks are")
+        square = len(block) == 2 and block[0] == block[1] and block[0] in FP8_BLOCK_SIZES
+        if q.get("weight_per_tensor") or not square:
+            raise NotImplementedError(
+                f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; "
+                f"only square blocks of {FP8_BLOCK_SIZES} are"
+            )
+        self.block = block[0]
         # transformers skips lm_head when the checkpoint gives no list
         not_convert = tuple(q.get("modules_to_not_convert") or ("lm_head",))
         self.not_convert = name_set(not_convert)
         self.not_convert_substr = substr_set(not_convert)
         self.convert_tables = name_set(tuple(q.get("modules_to_convert") or ()))
         self.e8m0 = str(q.get("scale_fmt") or "").lower() == "ue8m0"
-        self.expert_fp4 = str(cfg_get(hf_config, "expert_dtype") or "").lower() == "fp4"
+        # ``expert_dtype`` sits at the config top level (V4) or inside quantization_config (V4.1)
+        expert_dtype = q.get("expert_dtype") or cfg_get(hf_config, "expert_dtype")
+        self.expert_fp4 = str(expert_dtype or "").lower() == "fp4"
+        self.block_scheme = fp8_block_scheme("e8m0" if self.e8m0 else "float", self.block)
 
     def storage(self, scheme: QuantScheme) -> dict[str, Stored]:
         names = super().storage(scheme)
@@ -55,4 +68,11 @@ def scheme_for_name(self, name: str) -> QuantScheme | None:
             return None
         if self.expert_fp4 and is_routed_expert(name):
             return self.SCHEMES["EXPERT_MXFP4"]
-        return self.SCHEMES["BLOCK_E8M0" if self.e8m0 else "BLOCK"]
+        return self.block_scheme
+
+    def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig:
+        cfg = super().layer_config(layer, layer_kind, scheme)
+        if layer_kind is LayerKind.MOE and scheme is not None and scheme.kind is QuantKind.MXFP4:
+            # the fp4 experts quantize their activations at the same block as the dense linears
+            cfg = dataclasses.replace(cfg, act_block=self.block)
+        return cfg
diff --git a/python/freetoken/layers/quantization/linear/fp8_block.py b/python/freetoken/layers/quantization/linear/fp8_block.py
index 3b78e9776d..3eafdfdf5b 100644
--- a/python/freetoken/layers/quantization/linear/fp8_block.py
+++ b/python/freetoken/layers/quantization/linear/fp8_block.py
@@ -1,4 +1,4 @@
-"""fp8 e4m3 weight with 128x128 block scales (DeepSeek-V3 style)."""
+"""fp8 e4m3 weight with square block scales: 128x128 (DeepSeek-V3 style) or 32x32 (DeepSeek-V4.1)."""
 
 from __future__ import annotations
 
@@ -7,7 +7,7 @@
 import torch
 
 from ..registry import LayerKind, register_method
-from ..scheme import FP8_BLOCK as BLOCK, QuantKind
+from ..scheme import FP8_BLOCK, QuantKind, fp8_block_size
 from .base import LinearConfig, LinearKernel, LinearMethod
 
 FP8 = torch.float8_e4m3fn
@@ -18,6 +18,11 @@ def _e8m0(cfg: LinearConfig) -> bool:
     return cfg.scheme is not None and cfg.scheme.weight.scale == "e8m0"
 
 
+def layer_block_size(layer: Any) -> int:
+    """The block edge a finalized layer was declared with: ``K / (scale columns)``, exact by construction."""
+    return layer.weight.shape[1] // layer.weight_scale_inv.shape[1]
+
+
 class Dsv4Fp8BlockLinearKernel(LinearKernel):
     """DeepSeek-V4's reference path: activations quantized to fp8 with power-of-two block scales, e8m0 weight scales read as codes."""
 
@@ -29,7 +34,7 @@ def unusable_reason(self, cfg: LinearConfig) -> str | None:
     def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
         from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear
 
-        return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias)
+        return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias, block=layer_block_size(layer))
 
 
 class TritonFp8BlockLinearKernel(LinearKernel):
@@ -38,7 +43,11 @@ class TritonFp8BlockLinearKernel(LinearKernel):
     name = "triton"
 
     def unusable_reason(self, cfg: LinearConfig) -> str | None:
-        return "reads float block scales; e8m0 codes go to the dsv4 kernel" if _e8m0(cfg) else None
+        if _e8m0(cfg):
+            return "reads float block scales; e8m0 codes go to the dsv4 kernel"
+        if fp8_block_size(cfg.scheme) != FP8_BLOCK:
+            return f"float-scale kernel serves {FP8_BLOCK}x{FP8_BLOCK} blocks only"
+        return None
 
     def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
         from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear
@@ -52,9 +61,10 @@ class Fp8BlockLinearMethod(LinearMethod):
 
     def create_weights(self, layer: Any) -> None:
         g = self.cfg
-        if g.in_features % BLOCK or any(o % BLOCK for o in g.output_sizes):
-            raise ValueError(f"block-fp8 needs in/out sizes divisible by {BLOCK}, got K={g.in_features} N={g.output_sizes}")
+        block = fp8_block_size(g.scheme)
+        if g.in_features % block or any(o % block for o in g.output_sizes):
+            raise ValueError(f"block-fp8 needs in/out sizes divisible by {block}, got K={g.in_features} N={g.output_sizes}")
         layer.weight = torch.empty(g.out_features, g.in_features, dtype=FP8)
         # e8m0 codes stay codes for the dsv4 kernel; float scales are bf16 as the readers push them today
         scale_dtype = E8M0 if _e8m0(g) else torch.bfloat16
-        layer.weight_scale_inv = torch.empty(g.out_features // BLOCK, g.in_features // BLOCK, dtype=scale_dtype)
+        layer.weight_scale_inv = torch.empty(g.out_features // block, g.in_features // block, dtype=scale_dtype)
diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py
index e4a428d3e4..de1ce004fe 100644
--- a/python/freetoken/layers/quantization/moe/base.py
+++ b/python/freetoken/layers/quantization/moe/base.py
@@ -28,6 +28,8 @@ class MoEConfig:
     beta: float = 0.0
     limit: float | None = None
     interleaved: bool = False
+    # dynamic fp8 activation quant block of a W4A8 expert kernel, set by the quant dialect; None = no activation quant
+    act_block: int | None = None
     has_bias: bool = False
     apply_router_weight_on_input: bool = False
     strategy: str = "resident"
@@ -138,6 +140,10 @@ class MoEKernel(ABC):
     name: ClassVar[str]
     cpu_format: ClassVar[str | None] = None
     max_slots: ClassVar[int | None] = None
+    # Whether ``apply`` treats a slot of -1 as an inactive route (zero output, no read of any slot).
+    # Hybrid decode hands the GPU a route split with such holes; kernels without this get the holes
+    # pointed at slot 0 with weight 0 (slot storage is zero-filled at allocation, so that is exact).
+    supports_inactive_slots: ClassVar[bool] = False
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
         return None
diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py
index 20c3dde9dc..7552978d3f 100644
--- a/python/freetoken/layers/quantization/moe/fp8_block.py
+++ b/python/freetoken/layers/quantization/moe/fp8_block.py
@@ -5,7 +5,7 @@
 import torch
 
 from ..registry import LayerKind, register_method
-from ..scheme import FP8_BLOCK as BLOCK, QuantKind
+from ..scheme import FP8_BLOCK as BLOCK, QuantKind, fp8_block_size
 from .base import BankSpec, ExpertView, fused_piece, gated_epilogue_reason, limit_or_inf, MoEConfig, MoEKernel, MoEMethod
 
 FP8 = torch.float8_e4m3fn
@@ -24,6 +24,8 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
         reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False)
         if reason:
             return reason
+        if cfg.scheme is not None and fp8_block_size(cfg.scheme) != BLOCK:
+            return f"fp8 block MoE kernel serves {BLOCK}x{BLOCK} expert blocks only"
         reason = gated_epilogue_reason(cfg)
         if reason:
             return f"fp8 block MoE kernel: {reason}"
diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py
index 13cd97466c..1af6d9f43e 100644
--- a/python/freetoken/layers/quantization/moe/mxfp4.py
+++ b/python/freetoken/layers/quantization/moe/mxfp4.py
@@ -12,16 +12,20 @@
 
 
 class TritonMxfp4MoEKernel(MoEKernel):
-    """Standard OCP MXFP4 experts (DeepSeek-V4 ds_fp4): e2m1 pairs + e8m0 scales, no bias."""
+    """Standard OCP MXFP4 experts (DeepSeek-V4 / V4.1 ds_fp4): e2m1 pairs + e8m0 scales, no bias;
+    W4A8 with the fp8 activation round-trips at the checkpoint's block."""
 
     name = "triton"
     cpu_format = "ds_fp4"
+    supports_inactive_slots = True
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
         if cfg.interleaved:
             return "standard MXFP4 kernel reads the concatenated gate|up row order"
         if (cfg.alpha, cfg.beta) != (1.0, 0.0):
             return "standard MXFP4 kernel has no alpha / beta in its swiglu"
+        if cfg.act_block is None:
+            return "W4A8 ds_fp4 experts need the fp8 activation quant block"
         return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
@@ -44,13 +48,14 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         t = view.tensors
         banks = (t["gate_up"], t["gate_up_scale"], t["down"], t["down_scale"])
         limit = limit_or_inf(layer)
+        act_block = layer.quant_method.cfg.act_block
         if is_prefill and view.n is not None:
             from freetoken.moe.fused_ds_fp4 import routed_experts_fp4_prefill
 
-            return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n)
+            return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n, act_block=act_block)
         from freetoken.moe.fused_ds_fp4 import routed_experts_fp4
 
-        return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit)
+        return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit, act_block=act_block)
 
 
 class TritonGptossMxfp4MoEKernel(MoEKernel):
diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py
index 668eb6fb37..89f968afc8 100644
--- a/python/freetoken/layers/quantization/scheme.py
+++ b/python/freetoken/layers/quantization/scheme.py
@@ -25,7 +25,8 @@ def __str__(self) -> str:
         return self.value
 
 
-FP8_BLOCK = 128
+FP8_BLOCK = 128  # DeepSeek-V3 style square block; V4.1 exports 32x32
+FP8_BLOCK_SIZES = (32, 128)
 NVFP4_GROUP = 16
 MX_GROUP = 32
 
@@ -80,8 +81,17 @@ def fp8_tensor_scheme(scale: str, *, per_row: bool = False, input_scale: bool =
     return QuantScheme(QuantKind.FP8_TENSOR, WeightDesc("e4m3", (1, -1) if per_row else (-1, -1), scale), roles)
 
 
-def fp8_block_scheme(scale: str) -> QuantScheme:
-    return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (FP8_BLOCK, FP8_BLOCK), scale), {"weight", "weight_scale_inv"})
+def fp8_block_scheme(scale: str, block: int = FP8_BLOCK) -> QuantScheme:
+    """Square ``block x block`` e4m3 weight blocks, one scale each; ``block`` is the checkpoint's ``weight_block_size``."""
+    if block not in FP8_BLOCK_SIZES:
+        raise NotImplementedError(f"fp8 block size {block} is not supported; one of {FP8_BLOCK_SIZES}")
+    return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (block, block), scale), {"weight", "weight_scale_inv"})
+
+
+def fp8_block_size(scheme: QuantScheme) -> int:
+    """The square block edge of an FP8_BLOCK scheme (128 for DeepSeek-V3 style exports, 32 for DeepSeek-V4.1)."""
+    assert scheme.kind is QuantKind.FP8_BLOCK, scheme
+    return scheme.weight.group[0]
 
 
 def mxfp8_scheme() -> QuantScheme:
diff --git a/python/freetoken/mm/processors/deepseek_v41.py b/python/freetoken/mm/processors/deepseek_v41.py
new file mode 100644
index 0000000000..e40c961569
--- /dev/null
+++ b/python/freetoken/mm/processors/deepseek_v41.py
@@ -0,0 +1,88 @@
+"""DeepSeek-V4.1 image resize/patch layout and learned image-span delimiters."""
+
+from __future__ import annotations
+
+import math
+import struct
+
+import numpy as np
+import torch
+from PIL import ImageOps
+
+from freetoken.message import MMItem
+from freetoken.mm import mm_pad_value
+from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash
+
+
+def image_grid(width, height, patch, ratio, min_pixels, max_tokens, max_wh_ratio=None):
+    if max_wh_ratio is not None:
+        width = min(width, height * max_wh_ratio)
+    if width * height < min_pixels:
+        scale = (min_pixels / (width * height)) ** 0.5
+        width, height = int(width * scale), int(height * scale)
+    w, h = math.ceil(width / patch) * patch, math.ceil(height / patch) * patch
+    if math.ceil(h / (patch * ratio)) * (math.ceil(w / (patch * ratio)) + 1) + 2 > max_tokens:
+        aspect, cell = height / width, patch * ratio
+        wf = math.sqrt((max_tokens - 2) / aspect + 0.25) - 0.5
+        hf = wf * aspect
+        if wf < 1:
+            h, w = (max_tokens - 2) // 2 * cell, cell
+        elif hf < 1:
+            h, w = cell, (max_tokens - 3) * cell
+        else:
+            scale = min(math.floor(wf) * cell / width, math.floor(hf) * cell / height)
+            h, w = math.floor(height * scale / patch) * patch, math.floor(width * scale / patch) * patch
+    return h, w
+
+
+class DeepseekV41MMProcessor(MMProcessor):
+    def __init__(self, hf_config, model_path, mm):
+        super().__init__(model_path, mm)
+        self.vc = hf_config.vision_config
+        self.placeholder = [hf_config.image_token_id]
+        self.min_pixels = self.vc.min_pixels
+        self.max_tokens = mm.image_max_tokens if mm.image_max_tokens is not None else self.vc.max_image_tokens
+        if mm.image_min_tokens is not None:
+            # Delimiters and row breaks count against the total span budget.
+            self.min_pixels = max(1, mm.image_min_tokens - 2) * (self.vc.patch_size * self.vc.downsample_ratio)**2
+        allowed = {"min_pixels", "max_image_tokens", "max_wh_ratio"}
+        unknown = set(mm.processor_kwargs) - allowed
+        if unknown:
+            raise ValueError(f"unsupported DeepSeek image processor options: {sorted(unknown)}")
+        self.min_pixels = mm.processor_kwargs.get("min_pixels", self.min_pixels)
+        self.max_tokens = mm.processor_kwargs.get("max_image_tokens", self.max_tokens)
+        self.max_wh_ratio = mm.processor_kwargs.get("max_wh_ratio", self.vc.max_wh_ratio)
+        if self.max_tokens < 4 or self.min_pixels < 0:
+            raise ValueError("image token budget must be >= 4 and min_pixels must be nonnegative")
+        if self.max_wh_ratio is not None and self.max_wh_ratio <= 0:
+            raise ValueError("max_wh_ratio must be positive")
+
+    def process(self, images):
+        p, r = self.vc.patch_size, self.vc.downsample_ratio
+        items = []
+        for image in images:
+            h, w = image_grid(image.width, image.height, p, r, self.min_pixels, self.max_tokens, self.max_wh_ratio)
+            if self.max_wh_ratio is not None and image.width >= self.max_wh_ratio * image.height:
+                image = image.resize((w, h))
+            else:
+                image = ImageOps.pad(image, (w, h), color=(127, 127, 127))
+            x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
+            x = ((x - 0.5) / 0.5).to(torch.bfloat16)
+            nh, nw = h // p, w // p
+            patches = x.reshape(3, nh, p, nw, p).permute(1, 3, 0, 2, 4).reshape(nh * nw, 3, p, p).contiguous()
+            item_hash = content_hash(patches, struct.pack("<3i", 1, nh, nw))
+            items.append(MMItem(modality="image", hash=item_hash, pad_value=mm_pad_value(item_hash), offsets=[], feature=patches,
+                                model_specific_data={"grid_thw": [1, nh, nw]}))
+        return items
+
+    def prompt_replacement(self, item):
+        _, h, w = item.grid_thw
+        r = self.vc.downsample_ratio
+        n = ((h + r - 1) // r) * ((w + r - 1) // r + 1) + 2
+        # Every span row, including delimiters/newlines, has a learned visual embedding.
+        return PromptReplacement([self.placeholder[0]] * n)
+
+    def dummy_items(self, dtype, device):
+        p = self.vc.patch_size
+        return [MMItem(modality="image", hash=0, pad_value=0, offsets=[[0, 4]],
+                       feature=torch.zeros(1, 3, p, p, dtype=dtype, device=device), model_specific_data={"grid_thw": [1, 1, 1]})]
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index c5943dce1b..55b90bb58c 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -205,7 +205,8 @@ class DSV4AttentionGroupConfig(BaseAttentionGroupConfig):
     """DSV4 sparse attention (window + compressed tiers + Lightning Indexer). Standalone
     class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attention
     and reroute DSV4 through the SWA gates. Geometry beyond the latent width lives in
-    dsv4_args; the pool prices itself from there, not from this spec."""
+    dsv4_args (``dsv41_args`` for the ``v41`` variant); the pool prices itself from there,
+    not from this spec."""
 
     kind: ClassVar[Literal["dsv4"]] = "dsv4"
     cache_kind: ClassVar[Literal["dsv4_paged"]] = "dsv4_paged"
@@ -213,6 +214,8 @@ class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attenti
     num_kv_heads: int
     head_dim: int
     sliding_window: int  # the P-token window page
+    # "v41": DeepSeek-V4.1's cross-layer-shared packed tiers (AttnType.DSV41, its own pool and backend)
+    variant: Literal["v4", "v41"] = "v4"
 
 
 AttentionGroupConfig: TypeAlias = (
@@ -327,6 +330,9 @@ class ModelConfig:
     # CSA/HCA compressors, Lightning Indexer, manifold-constrained Hyper-Connections,
     # hash routing). Opaque to model-agnostic engine code; None for non-DSV4 models.
     dsv4_args: Any | None = None
+    # DeepSeek-V4.1 payload (DeepseekV41Args): DSV41 layer roles, Engram, replay policy. Opaque
+    # to model-agnostic engine code; None for every other model.
+    dsv41_args: Any | None = None
     # GLM-5.2 (glm_moe_dsa) MLA/DSA payload (GlmMoeDsaArgs): the MLA low-rank dims and the
     # DSA indexer geometry the model module needs. Opaque to model-agnostic engine code;
     # None for every other model.
@@ -448,7 +454,7 @@ def attn_type_for_layer(self, layer_id: int) -> AttnType:
         if isinstance(group, SWAAttentionGroupConfig):
             return AttnType.SWA
         if isinstance(group, DSV4AttentionGroupConfig):
-            return AttnType.DSV4
+            return AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4
         return _full_group_attn_type(group)
 
     def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]:
@@ -502,7 +508,7 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]:
                         num_kv_heads=group.num_kv_heads,
                         head_dim=group.head_dim,
                         sliding_window=group.sliding_window,
-                        attn_type=AttnType.DSV4,
+                        attn_type=AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4,
                     )
                 )
         return tuple(specs)
diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py
index 37f5076ba7..e19438d1ba 100644
--- a/python/freetoken/models/deepseek_v4/__init__.py
+++ b/python/freetoken/models/deepseek_v4/__init__.py
@@ -10,7 +10,7 @@
 DeepSeek-V4-Flash is a first-class registered model on the shared paged-KV engine:
 its window / compressed-attention / compressed-index KV live in DSV4-owned pools
 addressed by page tables, and sparse attention is a physical-slot gather (see
-:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4_paged_pool`).
+:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4.v4_pool`).
 """
 
 from .args import DeepseekV4Args, load_args
diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py
index a2d4387ade..9faefb83fe 100644
--- a/python/freetoken/models/deepseek_v4/model.py
+++ b/python/freetoken/models/deepseek_v4/model.py
@@ -4,7 +4,7 @@
 + stateful KV compressors / Lightning Indexer, manifold-constrained Hyper-Connections,
 sqrtsoftplus / hash MoE), wired onto FreeToken's shared paged engine:
 
-  - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4_paged_pool.DSV4PagedKVCache`):
+  - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4.v4_pool.DSV4PagedKVCache`):
     each layer's window-ring + compressed KV is a region of a shared global pool, addressed by
     per-layer slot maps. Sparse attention is a PAGED physical-slot gather:
     ``sparse_attn_paged`` reads KV directly from the two global pools (window / compressed) at
diff --git a/python/freetoken/models/deepseek_v41/__init__.py b/python/freetoken/models/deepseek_v41/__init__.py
new file mode 100644
index 0000000000..bcdef078c7
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/__init__.py
@@ -0,0 +1,14 @@
+"""DeepSeek-V4.1-Flash: Causal Encoder-Decoder MoE with Compressed Sparse Attention 2.
+
+40 layers (20 encoder + 20 decoder), every layer a sliding window (128) plus -- from layer 2 -- global
+sparse attention over cross-layer-shared compressed KV in packed fp4 (Full / Reindex / Reuse
+modes, a hierarchical candidate pool in the decoder), single-pass mHC residual streams, Engram n-gram
+memory at layers 1 and 14 streamed from disk, one shared + 384 routed MXFP4 experts per layer served
+from the offload cache. Images use the DeepSeek ViT; DSpark speculative decoding is not served.
+"""
+
+from .config import parse_config
+from .model import DeepseekV41ForCausalLM
+from .weight import ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights
+
+__all__ = ["DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files"]
diff --git a/python/freetoken/models/deepseek_v41/args.py b/python/freetoken/models/deepseek_v41/args.py
new file mode 100644
index 0000000000..867dcb21a9
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/args.py
@@ -0,0 +1,225 @@
+"""DeepSeek-V4.1-Flash hyperparameters and the static per-layer attention roles.
+
+``DeepseekV41Args`` mirrors the reference ``ModelArgs`` (``inference/model.py``) field for field,
+read from the HF ``config.json``'s ``text_config`` (the ``inference/config.json`` dialect is also
+accepted). Everything the engine reconciles at runtime (``max_seq_len``, the replay policy) is
+attached by ``parse_config`` / the engine, not read from the checkpoint.
+
+``LayerRole`` is the attention mode table (tech report sec. 2.3.1): every compressing layer is Full
+(owns the compressed KV + indexer K), Reindex (own indexer over a shared K) or Reuse (shared
+Top-K); the candidate pool of the Hierarchical Sparse Indexer (sec. 2.3.2) is produced by the
+first decoder Full layer and consumed by the decoder Reindex layers after it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+# text_config key -> ModelArgs key, where the two dialects differ
+_HF_TO_ARGS = {
+    "hidden_size": "dim",
+    "moe_intermediate_size": "moe_inter_dim",
+    "num_hidden_layers": "n_layers",
+    "num_nextn_predict_layers": "n_mtp_layers",
+    "num_attention_heads": "n_heads",
+    "num_experts_per_tok": "n_activated_experts",
+    "scoring_func": "score_func",
+    "routed_scaling_factor": "route_scale",
+    "qk_rope_head_dim": "rope_head_dim",
+    "rms_norm_eps": "norm_eps",
+    "sliding_window": "window_size",
+    "kv_source_layer_ids": "kv_source_layers",
+    "index_source_layer_ids": "index_source_layers",
+    "candidate_source_layer_id": "candidate_source_layer",
+    "engram_pad_token_id": "engram_pad_id",
+    "dspark_num_experts_per_tok": "dspark_n_activated_experts",
+}
+
+
+class Mode(str, Enum):
+    WINDOW = "window"  # sliding window only (compress_ratio == 0)
+    FULL = "full"  # computes main KV + indexer K, runs its indexer
+    REINDEX = "reindex"  # reuses main KV + indexer K, runs its own indexer
+    REUSE = "reuse"  # reuses main KV and the latest Top-K
+
+
+@dataclass(frozen=True)
+class LayerRole:
+    layer_id: int
+    ratio: int
+    mode: Mode
+    kv_source: int | None  # the layer whose main KV / indexer K this layer reads (itself when Full)
+    is_candidate_source: bool  # builds the hierarchical candidate pool
+    uses_candidates: bool  # indexes only inside the candidate pool
+
+    @property
+    def compresses(self) -> bool:
+        return self.ratio > 0
+
+    @property
+    def has_indexer(self) -> bool:
+        return self.mode in (Mode.FULL, Mode.REINDEX)
+
+
+@dataclass
+class DeepseekV41Args:
+    vocab_size: int = 129280
+    dim: int = 5120
+    moe_inter_dim: int = 2304
+    n_layers: int = 40
+    n_mtp_layers: int = 3
+    n_heads: int = 64
+    # moe
+    n_routed_experts: int = 384
+    n_shared_experts: int = 1
+    n_activated_experts: int = 6
+    score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus"
+    norm_topk_prob: bool = True
+    route_scale: float = 1.5
+    swiglu_limit: float = 10.0
+    # attention: latent q, one shared latent kv head, grouped low-rank output projection
+    q_lora_rank: int = 1280
+    head_dim: int = 512
+    rope_head_dim: int = 64
+    norm_eps: float = 1e-20
+    o_groups: int = 8
+    o_lora_rank: int = 1024
+    # sparse attention
+    window_size: int = 128
+    compress_ratios: tuple[int, ...] = ()  # one per layer, MTP layers included
+    kv_source_layers: tuple[int, ...] = ()
+    index_source_layers: tuple[int, ...] = ()
+    compress_rope_theta: float = 160000.0
+    original_seq_len: int = 65536
+    rope_theta: float = 10000.0
+    rope_factor: float = 16.0
+    beta_fast: int = 32
+    beta_slow: int = 1
+    index_n_heads: int = 32
+    index_head_dim: int = 128
+    index_topk: int = 512
+    candidate_source_layer: int = -1
+    candidate_topk_blocks: int = 0
+    candidate_block_size: int = 0
+    # hyper-connections (single-pass mHC)
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1e-6
+    # engram
+    engram_layer_ids: tuple[int, ...] = ()
+    engram_num_embeddings: tuple[int, ...] = ()
+    engram_max_ngram_size: int = 4
+    engram_vocab_size: int = 0
+    engram_n_heads: int = 0
+    engram_head_dim: int = 0
+    engram_pad_id: int = 2
+    engram_compressed_vocab_size: int = 0
+    # dspark (weights are not served in v1; kept so the config round-trips)
+    dspark_block_size: int = 0
+    dspark_target_layer_ids: tuple[int, ...] = ()
+    # ---- runtime (engine-reconciled, never read from the checkpoint) ----
+    max_seq_len: int = 4096
+    max_batch_size: int = 1
+    # "bounded": the decoder runs on the prompt's last window_size tokens with the sliding window
+    # truncated there (Decoder SWA Bounded Replay, tech report sec. 3.2.2); "exact": on every token.
+    swa_decoder_replay: Literal["bounded", "exact"] = "bounded"
+    roles: tuple[LayerRole, ...] = field(default_factory=tuple, repr=False)
+
+    def __post_init__(self) -> None:
+        for name in ("compress_ratios", "kv_source_layers", "index_source_layers", "engram_layer_ids",
+                     "engram_num_embeddings", "dspark_target_layer_ids"):
+            setattr(self, name, tuple(int(v) for v in getattr(self, name)))
+        if len(self.compress_ratios) < self.n_layers:
+            raise ValueError(f"compress_ratios has {len(self.compress_ratios)} entries for {self.n_layers} layers")
+        if self.n_shared_experts != 1:
+            raise ValueError("DeepSeek-V4.1 has exactly one shared expert")
+        if self.hc_mult < 1:
+            raise ValueError("hc_mult must be positive")
+        self.roles = tuple(self._role(layer) for layer in range(self.n_layers))
+        if len(self.engram_layer_ids) != len(self.engram_num_embeddings):
+            raise ValueError("engram_layer_ids and engram_num_embeddings differ in length")
+
+    # ---- derived ----
+    def _role(self, layer: int) -> LayerRole:
+        ratio = self.compress_ratios[layer]
+        if ratio == 0:
+            return LayerRole(layer, 0, Mode.WINDOW, None, False, False)
+        sources = [s for s in self.kv_source_layers if s <= layer]
+        if not sources:
+            raise ValueError(f"layer {layer} compresses but no kv source precedes it")
+        source = max(sources)
+        if self.compress_ratios[source] != ratio:
+            raise ValueError(f"layer {layer} (ratio {ratio}) reads kv source {source} (ratio {self.compress_ratios[source]})")
+        if layer == source:
+            mode = Mode.FULL
+        elif layer in self.index_source_layers:
+            mode = Mode.REINDEX
+        else:
+            mode = Mode.REUSE
+        cand = self.candidate_source_layer
+        is_cand_src = cand >= 0 and layer == cand
+        uses = cand >= 0 and cand < layer and mode is Mode.REINDEX
+        if uses and self.kv_source_layers and max(s for s in self.kv_source_layers if s <= cand) != source:
+            raise ValueError(f"layer {layer} would index the candidate pool of layer {cand} over a different kv source")
+        return LayerRole(layer, ratio, mode, source, is_cand_src, uses)
+
+    @property
+    def nope_head_dim(self) -> int:
+        return self.head_dim - self.rope_head_dim
+
+    @property
+    def backbone_kv_sources(self) -> tuple[int, ...]:
+        return tuple(s for s in self.kv_source_layers if s < self.n_layers)
+
+    @property
+    def decoder_start_layer(self) -> int:
+        """The first decoder layer: the last kv source (CED projects the decoder's global KV from
+        its input, the final encoder hidden state)."""
+        return max(self.backbone_kv_sources) if self.backbone_kv_sources else self.n_layers
+
+    @property
+    def engram_hash_cols(self) -> int:
+        return (self.engram_max_ngram_size - 1) * self.engram_n_heads
+
+    @property
+    def index_softmax_scale(self) -> float:
+        return self.index_head_dim**-0.5
+
+    def freqs_params(self, layer: int) -> tuple[int, int, float, float, int, int]:
+        """``(rotary_dim, original_seq_len, theta, factor, beta_fast, beta_slow)``: compressing layers
+        rotate at the compressed theta under YaRN, window-only layers at the base theta without it."""
+        if self.compress_ratios[layer]:
+            return self.rope_head_dim, self.original_seq_len, self.compress_rope_theta, self.rope_factor, self.beta_fast, self.beta_slow
+        return self.rope_head_dim, 0, self.rope_theta, self.rope_factor, self.beta_fast, self.beta_slow
+
+    @classmethod
+    def from_hf(cls, hf_config: Any) -> "DeepseekV41Args":
+        """From the HF config (``text_config`` of the multimodal wrapper, or the text config itself,
+        or the ``inference/config.json`` dialect)."""
+        text = getattr(hf_config, "text_config", None) or hf_config
+        if isinstance(text, dict):
+            raw = text
+        elif hasattr(text, "to_dict"):  # PretrainedConfig / RawConfigShim
+            raw = text.to_dict()
+        else:
+            raw = vars(text)
+        names = {f.name for f in cls.__dataclass_fields__.values()} - {"roles", "max_seq_len", "max_batch_size", "swa_decoder_replay"}
+        kwargs: dict[str, Any] = {}
+        for key, value in raw.items():
+            key = _HF_TO_ARGS.get(key, key)
+            if key in names:
+                kwargs[key] = value
+        # HF nests the YaRN parameters; the inference dialect keeps them flat
+        scaling = raw.get("rope_scaling") or raw.get("rope_parameters")
+        if isinstance(scaling, dict):
+            if "factor" in scaling:
+                kwargs["rope_factor"] = scaling["factor"]
+            for src, dst in (("beta_fast", "beta_fast"), ("beta_slow", "beta_slow"), ("original_max_position_embeddings", "original_seq_len")):
+                if src in scaling:
+                    kwargs[dst] = scaling[src]
+        return cls(**kwargs)
+
+
+__all__ = ["DeepseekV41Args", "LayerRole", "Mode"]
diff --git a/python/freetoken/models/deepseek_v41/attention.py b/python/freetoken/models/deepseek_v41/attention.py
new file mode 100644
index 0000000000..1643ddb6d1
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/attention.py
@@ -0,0 +1,305 @@
+"""DSV41 attention layer (reference ``Attention``): latent MLA over a sliding window plus, on
+compressing layers, ``index_topk`` compressed positions -- one paged sparse-attention call.
+
+Per layer mode (``LayerRole``):
+
+* Full     -- ``publish_*`` compresses its input into the kv source's main KV (fp4) and index keys
+              (fp4), then its indexer selects fresh Top-K rows and publishes them on the forward's
+              ``SharedSelection``; the decoder's first Full layer also builds the candidate pool.
+* Reindex  -- selects fresh Top-K rows with its own indexer over the shared keys (inside the
+              candidate pool in the decoder) and publishes them.
+* Reuse    -- attends over the latest published Top-K rows.
+* Window   -- the sliding window only.
+
+All modes compute their own query and sliding-window KV (fp8 into the layer's window ring). The
+attention output is rotated back (K == V share one rotated latent) and projected through the
+grouped low-rank ``wo_a`` (block-diagonal over ``o_groups``; fp8 payload kept, W8A16 at decode) and ``wo_b``.
+
+Prefill runs ragged over a flat token stream tiled by ``PrefillSegment``s; the compressor / indexer
+state is per request, so those run per segment. Decode is batched (one token per row) and
+CUDA-graph safe: every gather reads the backend's snapshot, never the live page table.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List
+
+import torch
+from freetoken.attention.dsv41_sparse import PrefillSegment
+from freetoken.core import get_global_ctx
+from freetoken.kernel.triton.dsv4.fp8_linear import GEMV_MAX_M, dequant_block_fp8, grouped_w8a16_gemv
+from freetoken.layers import BaseOP, LinearReplicated, RMSNorm
+from freetoken.layers.quantization import QuantKind
+from freetoken.layers.quantization.scheme import fp8_block_size
+
+from .args import DeepseekV41Args, LayerRole, Mode
+from .compressor import Compressor
+from .indexer import Indexer
+from .rope import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis
+
+
+class DSV41Attention(BaseOP):
+    def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""):
+        self.args = args
+        self.role = role
+        self.layer_id = role.layer_id
+        self.n_heads = args.n_heads
+        self.head_dim = args.head_dim
+        self.rope_dim = args.rope_head_dim
+        self.n_groups = args.o_groups
+        self.o_lora_rank = args.o_lora_rank
+        self.window = args.window_size
+        self.softmax_scale = args.head_dim**-0.5
+
+        self.wq_a = LinearReplicated(args.dim, args.q_lora_rank, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_a")
+        self.q_norm = RMSNorm(args.q_lora_rank, args.norm_eps)
+        self.wq_b = LinearReplicated(args.q_lora_rank, args.n_heads * args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b")
+        self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv")
+        self.kv_norm = RMSNorm(args.head_dim, args.norm_eps)
+        # Block-diagonal over groups. The reference dequantizes it to bf16 and runs a grouped einsum
+        # (the activation stays bf16); a quantized checkpoint keeps the e4m3 payload + e8m0 scale
+        # here and dequantizes in-kernel -- the same math, half the bytes per decode step, and no
+        # bf16 copy in VRAM. An unquantized checkpoint ships it bf16.
+        wo_a_scheme = quant_config.scheme_for(f"{prefix}.wo_a") if quant_config is not None else None
+        self._wo_a_block = fp8_block_size(wo_a_scheme) if wo_a_scheme is not None and wo_a_scheme.kind is QuantKind.FP8_BLOCK else None
+        rows, cols = args.o_groups * args.o_lora_rank, args.n_heads * args.head_dim // args.o_groups
+        if self._wo_a_block is not None:
+            self.wo_a = torch.empty(rows, cols, dtype=torch.float8_e4m3fn)
+            self.wo_a_scale = torch.empty(rows // self._wo_a_block, cols // self._wo_a_block, dtype=torch.float8_e8m0fnu)
+        else:
+            self.wo_a = torch.empty(rows, cols, dtype=torch.bfloat16)
+        self.wo_b = LinearReplicated(args.o_groups * args.o_lora_rank, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wo_b")
+        self.attn_sink = torch.empty(args.n_heads, dtype=torch.float32)
+
+        self.compressor = Compressor(args, self.layer_id, role.ratio, quant_config=quant_config, prefix=f"{prefix}.compressor") if role.mode is Mode.FULL else None
+        self.indexer = Indexer(args, role, quant_config=quant_config, prefix=f"{prefix}.indexer") if role.has_indexer else None
+        self._freqs_params = args.freqs_params(self.layer_id)
+        self._freqs: torch.Tensor | None = None
+
+    # addressing lives on the live backend, pool buffers on the live pool (both read per access)
+    @property
+    def attn(self):
+        return get_global_ctx().attn_backend
+
+    @property
+    def freqs_key(self) -> tuple:
+        """The RoPE table this layer rotates with (window layers: base theta; compressing layers: the
+        compressed theta under YaRN). Layers with the same key share one table."""
+        return self._freqs_params
+
+    def bind(self, device: torch.device, tables: dict | None = None) -> None:
+        """Attach the RoPE table; ``tables`` (key -> table) lets the model share one table per key
+        instead of every layer holding its own ``[max_seq_len, rd // 2]`` copy."""
+        if tables is not None and self.freqs_key in tables:
+            self._freqs = tables[self.freqs_key]
+            return
+        rd, original_seq_len, theta, factor, beta_fast, beta_slow = self._freqs_params
+        self._freqs = get_freqs_cis(rd, self.args.max_seq_len, original_seq_len, theta, factor, beta_fast, beta_slow, device)
+        if tables is not None:
+            tables[self.freqs_key] = self._freqs
+
+    @property
+    def freqs(self) -> torch.Tensor:
+        assert self._freqs is not None, "attention layer is not bound (call bind first)"
+        return self._freqs
+
+    # ----- shared pieces -------------------------------------------------------------------
+    def _q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        qr = self.q_norm.forward(self.wq_a.forward(x))
+        q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim))
+        return qr, q
+
+    def _kv(self, x: torch.Tensor) -> torch.Tensor:
+        return self.kv_norm.forward(self.wkv.forward(x))
+
+    def _wo(self, o: torch.Tensor) -> torch.Tensor:
+        """``wo_b(wo_a(o))``: up to ``GEMV_MAX_M`` rows the grouped W8A16 GEMV over the fp8 payload, else the
+        reference's grouped einsum over the bf16 (or dequantized) ``wo_a``."""
+        t = o.shape[0]
+        o = o.reshape(t, self.n_groups, -1)
+        if self._wo_a_block is not None and t <= GEMV_MAX_M:
+            return self.wo_b.forward(grouped_w8a16_gemv(o, self.wo_a, self.wo_a_scale, block=self._wo_a_block))
+        wo_a = self.wo_a if self._wo_a_block is None else dequant_block_fp8(self.wo_a, self.wo_a_scale, block=self._wo_a_block)
+        return self.wo_b.forward(torch.einsum("tgd,grd->tgr", o, wo_a.view(self.n_groups, self.o_lora_rank, -1)).flatten(1))
+
+    def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False) -> None:
+        key = f"{prefix}.wo_a" if prefix else "wo_a"
+        scale_key = f"{key}_scale"
+        if self._wo_a_block is None and scale_key in state_dict:
+            # a bf16 layer (no fp8 scheme) fed a quantized checkpoint dequantizes wo_a on the way in
+            from .weight import dequant_fp8_block
+
+            state_dict[key] = dequant_fp8_block(state_dict[key], state_dict.pop(scale_key))
+        elif self._wo_a_block is not None and scale_key not in state_dict:
+            # an FTW conversion made before wo_a kept its fp8 payload carries a dequantized bf16
+            # tensor and no scale: say so instead of failing on the missing key
+            raise RuntimeError(
+                f"{key}: the checkpoint has no {scale_key} for the fp8 wo_a layout; an FTW checkpoint "
+                "converted by an earlier build stores wo_a dequantized -- reconvert it with `ft checkpoint`"
+            )
+        super().load_state_dict(state_dict, prefix=prefix, _internal=_internal)
+
+    # ----- prefill -----------------------------------------------------------------------------
+    def publish_prefill(self, x: torch.Tensor, segments: List[PrefillSegment]) -> None:
+        """Full mode: compress this layer's input for every segment into the source's main KV and
+        index keys. Runs on EVERY new token (under bounded replay the decoder source publishes for
+        the whole prompt even though its attention runs on the prompt's last window only)."""
+        assert self.compressor is not None and self.indexer is not None
+        attn, ratio, src = self.attn, self.role.ratio, self.layer_id
+        for seg in segments:
+            slots = attn.window_slots_of(seg.table_idx, seg.start_pos, seg.end)
+            tail = None
+            if self.compressor.needs_tail_carry(seg.start_pos):
+                tail = int(attn.window_slots_of(seg.table_idx, seg.start_pos - 1, seg.start_pos).item())
+            latent, starts = self.compressor.forward_prefill(x[seg.offset : seg.offset + seg.n], seg.start_pos, slots, tail)
+            if not latent.shape[0]:
+                continue
+            rows = attn.compressed_rows_of(seg.table_idx, starts, ratio)
+            freqs = self.freqs.index_select(0, starts)
+            attn.store_index(self.indexer.index_keys(latent, freqs), src, rows)
+            apply_rotary_emb(latent[..., -self.rope_dim :], freqs)
+            attn.store_main(latent, src, rows)
+
+    def forward_prefill(self, x: torch.Tensor, segments: List[PrefillSegment], positions: torch.Tensor) -> torch.Tensor:
+        """Ragged prefill over ``x [T, dim]`` tiled by ``segments`` (``positions [T]`` absolute)."""
+        attn, role = self.attn, self.role
+        freqs = self.freqs.index_select(0, positions)
+        qr, q = self._q(x)
+        apply_rotary_emb(q[..., -self.rope_dim :], freqs)
+        kv = self._kv(x)
+        apply_rotary_emb(kv[..., -self.rope_dim :], freqs)
+
+        sel = attn.selection
+        win_parts, cmp_parts, pools = [], [], []
+        for seg in segments:
+            lo, hi = seg.offset, seg.offset + seg.n
+            attn.store_window(kv[lo:hi], self.layer_id, attn.layer_window_slots_of(self.layer_id, seg.table_idx, seg.start_pos, seg.end))
+            win_parts.append(attn.window_topk_prefill(seg, self.layer_id))
+            if not role.compresses:
+                continue
+            if role.has_indexer:
+                locs = attn.locs_prefill(seg.table_idx, seg.end // role.ratio * role.ratio)
+                cand = sel.candidates[:, lo:hi] if role.uses_candidates else None
+                rows, pool = self.indexer.select_prefill(
+                    x[lo:hi], qr[lo:hi], freqs[lo:hi], start_pos=seg.start_pos, ratio=role.ratio, locs=locs, candidates=cand,
+                )
+                cmp_parts.append(rows)
+                if pool is not None:
+                    pools.append(pool)
+            else:
+                assert sel.source == role.kv_source, "Reuse layer without a published selection for its source"
+                cmp_parts.append(sel.topk_rows[:, lo:hi])
+        if role.has_indexer:
+            sel.source, sel.topk_rows = role.kv_source, torch.cat(cmp_parts, dim=1)
+            if pools:
+                sel.candidates = torch.cat(pools, dim=1)
+        win = torch.cat(win_parts, dim=1)
+        topk = torch.cat([win, torch.cat(cmp_parts, dim=1)], dim=-1) if role.compresses else win
+
+        o = attn.attend(q.unsqueeze(0), self.layer_id, topk.to(torch.int32), self.window, self.attn_sink, self.softmax_scale)[0]
+        apply_rotary_emb(o[..., -self.rope_dim :], freqs, inverse=True)
+        return self._wo(o)
+
+    # ----- decode ------------------------------------------------------------------------------
+    def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, rows: torch.Tensor, dctx: "DecodeStepContext", cmp_stage_cap: int) -> torch.Tensor:
+        """Batched single-token attention. ``pos [B]`` device positions, ``rows [B]`` local rows into
+        the decode snapshot, ``dctx`` the step's layer-invariant context (ring slots, window
+        candidates, RoPE rows), ``cmp_stage_cap`` the static compressed staging bound (max position
+        any row reaches; the capture width under a graph)."""
+        attn, role = self.attn, self.role
+        B = x.shape[0]
+        freqs_t = dctx.freqs[self.freqs_key]
+        qr, q = self._q(x)
+        apply_rotary_emb_decode(q[..., -self.rope_dim :].unsqueeze(1), freqs_t)
+        kv = self._kv(x)
+        apply_rotary_emb_decode(kv[..., -self.rope_dim :].unsqueeze(1), freqs_t)
+        # the shared ring context feeds the compressor's carry ring on every layer; a request-private
+        # window layer stores and reads its own KV through its ring context instead
+        window_slots, prev_window_slots = dctx.window_slots, dctx.prev_window_slots
+        private = dctx.private is not None and attn.pool.is_private_window(self.layer_id)
+        tier = dctx.private if private else dctx.shared
+        attn.store_window(kv, self.layer_id, tier.slots if private else window_slots)
+
+        cmp_counts = None
+        if role.compresses:
+            ratio, src = role.ratio, role.kv_source
+            n_stage = (cmp_stage_cap + 1) // ratio
+            sel = attn.selection
+            if role.mode is Mode.FULL:
+                latent, completed = self.compressor.forward_decode(x, pos, prev_window_slots, window_slots)
+                dst = attn.decode_store_rows(rows, pos, ratio, src, completed)
+                freqs_g = self.freqs.index_select(0, (pos + 1 - ratio).clamp_min(0))
+                k = self.indexer.k_norm.forward(self.indexer.wk.forward(latent))
+                apply_rotary_emb_decode(k[..., -self.rope_dim :].unsqueeze(1), freqs_g)
+                attn.store_index(k, src, dst)
+                apply_rotary_emb_decode(latent[..., -self.rope_dim :].unsqueeze(1), freqs_g)
+                attn.store_main(latent, src, dst)
+            if role.has_indexer:
+                cand = sel.candidates if role.uses_candidates else None
+                cmp_rows, counts, pool = self.indexer.select_decode(x, qr, freqs_t, pos=pos, ratio=ratio, locs=attn.snapshot(), T=n_stage, candidates=cand)
+                sel.source, sel.topk_rows, sel.cmp_counts = src, cmp_rows, counts
+                if pool is not None:
+                    sel.candidates = pool
+                # the step's [window | compressed] candidate buffer of this tier takes the new rows once;
+                # the Reuse layers that follow read them in place (no per-layer cat / cast)
+                tier.topk[..., self.window :].copy_(cmp_rows)
+            else:
+                assert sel.source == src, "Reuse layer without a published selection for its source"
+                counts = sel.cmp_counts
+            topk = tier.topk
+            cmp_counts = counts
+        else:
+            topk = tier.window_topk
+
+        o = attn.attend(q.view(B, 1, self.n_heads, self.head_dim), self.layer_id, topk, self.window, self.attn_sink, self.softmax_scale, cmp_counts=cmp_counts)
+        apply_rotary_emb_decode(o[..., -self.rope_dim :], freqs_t, inverse=True)
+        return self._wo(o.view(B, self.n_heads, self.head_dim))
+
+
+@dataclass
+class TierCandidates:
+    """One window tier's per-step decode candidates: the request's store slot and the int32
+    ``[B, 1, win]`` window candidates, plus the ``[B, 1, win + topk]`` ``[window | compressed]``
+    buffer whose compressed half the index layers of that tier fill as they publish."""
+
+    slots: torch.Tensor
+    window_topk: torch.Tensor
+    topk: torch.Tensor
+
+
+@dataclass
+class DecodeStepContext:
+    """Everything a decode step resolves once for all layers (a capture records it once): the shared
+    pool's ring slots (the compressor's carry ring keys on them), the shared / private tier
+    candidates, and each RoPE table gathered at the step's positions."""
+
+    window_slots: torch.Tensor
+    prev_window_slots: torch.Tensor
+    shared: TierCandidates
+    private: TierCandidates | None
+    freqs: dict
+
+    @classmethod
+    def build(cls, layers, md, pos: torch.Tensor, rows: torch.Tensor, window: int, topk: int) -> "DecodeStepContext":
+        B = pos.shape[0]
+
+        def tier(slots, window_topk):
+            win32 = window_topk.to(torch.int32)
+            buf = torch.empty((B, 1, window + topk), dtype=torch.int32, device=pos.device)
+            buf[..., :window].copy_(window_topk)
+            return TierCandidates(slots, win32, buf)
+
+        window_slots, prev_window_slots, window_topk = md.window_ctx(pos, rows)
+        private = None
+        if get_global_ctx().kv_cache.private_window_layer_ids:
+            private = tier(*md.private_window_ctx(pos, rows))
+        tables: dict = {}
+        for layer in layers:
+            key = layer.freqs_key
+            if key not in tables:
+                tables[key] = layer.freqs.index_select(0, pos)
+        return cls(window_slots, prev_window_slots, tier(window_slots, window_topk), private, tables)
+
+
+__all__ = ["DSV41Attention", "DecodeStepContext", "TierCandidates"]
diff --git a/python/freetoken/models/deepseek_v41/compressor.py b/python/freetoken/models/deepseek_v41/compressor.py
new file mode 100644
index 0000000000..5d5d07fe5b
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/compressor.py
@@ -0,0 +1,131 @@
+"""DSV41 KV compressor: ``compress_ratio`` consecutive tokens -> one latent (reference ``Compressor``).
+
+* ratio 1 (the CED decoder source): a plain bf16 projection + RMSNorm, one latent per token.
+* ratio > 1 (the encoder sources): ``wkv`` / ``wgate`` promoted to fp32, a softmax over the group's
+  gate scores pools the group's ``wkv`` rows; the trailing partial group is carried across chunks
+  and decode steps in the pool's compress-state ring (per window page, so a page-aligned radix hit
+  resumes it by value -- see ``DSV41SparseAttnBackend`` for the addressing).
+
+Returns the latent BEFORE RoPE: the indexer key is projected from the unrotated latent, the
+attention layer rotates and quantizes it afterwards.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from freetoken.core import get_global_ctx
+from freetoken.layers import BaseOP, LinearReplicated, RMSNorm
+
+from .args import DeepseekV41Args
+
+
+class Compressor(BaseOP):
+    def __init__(self, args: DeepseekV41Args, layer_id: int, ratio: int, *, quant_config=None, prefix: str = ""):
+        assert ratio >= 1
+        self.layer_id = layer_id
+        self.ratio = ratio
+        self.dim = args.dim
+        self.head_dim = args.head_dim
+        self.window = args.window_size
+        if ratio > 1:
+            # fp32 like the reference (the softmax pooling runs in fp32); the checkpoint ships bf16
+            self.wkv = torch.empty(args.head_dim, args.dim, dtype=torch.float32)
+            self.wgate = torch.empty(args.head_dim, args.dim, dtype=torch.float32)
+        else:
+            self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv")
+        self.norm = RMSNorm(args.head_dim, args.norm_eps)
+
+    @property
+    def attn(self):
+        return get_global_ctx().attn_backend
+
+    def _project(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """fp32 ``wkv`` / ``wgate`` projections."""
+        xf = x.float()
+        return F.linear(xf, self.wkv), F.linear(xf, self.wgate)
+
+    # ----- prefill -----------------------------------------------------------------------
+    def needs_tail_carry(self, start_pos: int) -> bool:
+        """Whether a prefill from ``start_pos`` resumes a partial group (never at a page-aligned start)."""
+        return self.ratio > 1 and start_pos % self.ratio != 0
+
+    def forward_prefill(
+        self, x: torch.Tensor, start_pos: int, window_slots: torch.Tensor, tail_window_slot: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compress the tokens ``[start_pos, start_pos + n)`` of one request.
+
+        ``window_slots`` are the tokens' window slots (the ring block of each window page carries the
+        partial group); ``tail_window_slot`` is the previous token's slot, needed only when
+        ``needs_tail_carry(start_pos)``. Returns ``(latent [G, head_dim] bf16 pre-RoPE, group_starts [G]
+        int64)`` for every group that completed in this call.
+        """
+        n = x.shape[0]
+        device = x.device
+        if self.ratio == 1:
+            latent = self.norm.forward(self.wkv.forward(x))
+            return latent, torch.arange(start_pos, start_pos + n, device=device)
+
+        ratio, d = self.ratio, self.head_dim
+        kv, score = self._project(x)
+        # the partial group in progress at start_pos: its rows live in the previous token's ring block
+        offset = start_pos % ratio
+        if offset:
+            assert tail_window_slot is not None, "resuming mid-group needs the previous token's window slot"
+            carry = self.attn.read_carry(self.layer_id, tail_window_slot)  # [ratio, 2d]
+            kv = torch.cat([carry[:offset, :d], kv], dim=0)
+            score = torch.cat([carry[:offset, d:], score], dim=0)
+        total = kv.shape[0]
+        groups = total // ratio
+        remainder = total % ratio
+        cut = groups * ratio
+        if groups:
+            kv_g = kv[:cut].view(groups, ratio, d)
+            score_g = score[:cut].view(groups, ratio, d)
+            pooled = (kv_g * score_g.softmax(dim=1)).sum(dim=1)
+            latent = self.norm.forward(pooled.to(x.dtype))
+        else:
+            latent = x.new_empty(0, d)
+        group_starts = (start_pos - offset) + ratio * torch.arange(groups, device=device)
+        # ring: the reset block at every page boundary crossed (a page holds whole groups), and the
+        # in-progress tail (the remainder rows) on the last page so decode / the next chunk resume
+        end = start_pos + n
+        self.attn.write_boundary_carries(self.layer_id, lo=start_pos, hi=end, window_slots=window_slots)
+        if end % self.window != 0:
+            block = self._carry_block(kv[cut:], score[cut:], remainder)
+            self.attn.write_carry(self.layer_id, int(window_slots[-1].item()), block)
+        return latent, group_starts
+
+    def _carry_block(self, kv_tail: torch.Tensor, score_tail: torch.Tensor, filled: int) -> torch.Tensor:
+        ratio, d = self.ratio, self.head_dim
+        ks = kv_tail.new_zeros(ratio, d)
+        ss = kv_tail.new_full((ratio, d), float("-inf"))
+        if filled:
+            ks[:filled] = kv_tail
+            ss[:filled] = score_tail
+        return torch.cat([ks, ss], dim=-1)
+
+    # ----- decode ------------------------------------------------------------------------
+    def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, prev_window_slots: torch.Tensor, window_slots: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """One token per row. Returns ``(latent [B, head_dim] bf16 pre-RoPE, completed [B] bool)``;
+        a row whose group did not complete carries a meaningless latent the caller discards."""
+        B = x.shape[0]
+        if self.ratio == 1:
+            latent = self.norm.forward(self.wkv.forward(x))
+            return latent, torch.ones(B, dtype=torch.bool, device=x.device)
+
+        from freetoken.kernel.triton.dsv4.compress import gated_pool
+
+        ratio, d = self.ratio, self.head_dim
+        kv, score = self._project(x)
+        idx = (pos % ratio).view(B, 1, 1).expand(B, 1, d)
+        block = self.attn.read_carry_blocks(self.layer_id, prev_window_slots)  # [B, ratio, 2d]
+        ks = block[..., :d].clone().scatter_(1, idx, kv.unsqueeze(1))
+        ss = block[..., d:].clone().scatter_(1, idx, score.unsqueeze(1))
+        self.attn.write_carry_blocks(self.layer_id, window_slots, torch.cat([ks, ss], dim=-1))
+        pooled = gated_pool(ks, ss, x.dtype).view(B, d)
+        completed = (pos + 1) % ratio == 0
+        return self.norm.forward(pooled), completed
+
+
+__all__ = ["Compressor"]
diff --git a/python/freetoken/models/deepseek_v41/config.py b/python/freetoken/models/deepseek_v41/config.py
new file mode 100644
index 0000000000..9e738939c8
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/config.py
@@ -0,0 +1,95 @@
+"""Engine-facing config for DeepSeek-V4.1-Flash.
+
+``parse_config`` maps the standard transformer fields into :class:`ModelConfig`, carries the full
+:class:`DeepseekV41Args` in ``ModelConfig.dsv41_args`` for the model module and the KV pool (which
+sizes its tiers from it), and declares the ``v41`` variant of the DSV4 attention group.
+The engine reconciles the runtime knobs (``max_seq_len``, ``swa_decoder_replay``) onto ``dsv41_args``
+at config resolution.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from typing import Any
+
+from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig
+
+from .args import DeepseekV41Args
+
+
+@dataclass(frozen=True)
+class VisionConfig:
+    num_hidden_layers: int = 32
+    hidden_size: int = 1024
+    num_attention_heads: int = 16
+    intermediate_size: int = 2816
+    patch_size: int = 14
+    rope_theta: float = 10000.0
+    downsample_ratio: int = 3
+    max_image_tokens: int = 1024
+    min_pixels: int = 544 * 544
+    max_wh_ratio: float | None = None
+
+
+def parse_config(hf_config: Any) -> ModelConfig:
+    args = DeepseekV41Args.from_hf(hf_config)
+    text = getattr(hf_config, "text_config", None) or hf_config
+    max_position = int(getattr(text, "max_position_embeddings", 0) or getattr(hf_config, "max_position_embeddings", 0) or 0)
+    if max_position <= 0:
+        max_position = int(args.original_seq_len * args.rope_factor)
+    vc = getattr(hf_config, "vision_config", None)
+    vision = VisionConfig(**{f.name: getattr(vc, f.name, f.default) for f in fields(VisionConfig)}) if vc is not None else None
+
+    return ModelConfig(
+        num_layers=args.n_layers,
+        num_qo_heads=args.n_heads,
+        num_kv_heads=1,  # one shared latent KV head (K == V)
+        head_dim=args.head_dim,
+        hidden_size=args.dim,
+        vocab_size=args.vocab_size,
+        intermediate_size=args.moe_inter_dim,
+        hidden_act="silu",
+        rms_norm_eps=args.norm_eps,
+        tie_word_embeddings=False,
+        rotary_config=RotaryConfig(
+            head_dim=args.head_dim,
+            rotary_dim=args.rope_head_dim,
+            max_position=max_position,
+            base=args.rope_theta,
+            scaling={
+                "rope_type": "yarn",
+                "factor": args.rope_factor,
+                "beta_fast": args.beta_fast,
+                "beta_slow": args.beta_slow,
+                "original_max_position_embeddings": args.original_seq_len,
+            },
+        ),
+        num_experts=args.n_routed_experts,
+        num_experts_per_tok=args.n_activated_experts,
+        moe_intermediate_size=args.moe_inter_dim,
+        norm_topk_prob=args.norm_topk_prob,
+        n_shared_experts=args.n_shared_experts,
+        routed_scaling_factor=args.route_scale,
+        swiglu_limit=args.swiglu_limit,
+        model_type="deepseek_v41",
+        architectures=["DeepseekV41ForCausalLM"],
+        moe_enabled=True,
+        expert_quant="ds_fp4",
+        attn_sm_scale=args.head_dim**-0.5,
+        dsv41_args=args,
+        vision_config=vision,
+        image_token_id=getattr(hf_config, "image_token_id", None),
+        attention_groups=(
+            DSV4AttentionGroupConfig(
+                name="dsv41",
+                layer_ids=tuple(range(args.n_layers)),
+                num_kv_heads=1,
+                head_dim=args.head_dim,
+                sliding_window=args.window_size,
+                variant="v41",
+            ),
+        ),
+    )
+
+
+__all__ = ["parse_config"]
diff --git a/python/freetoken/models/deepseek_v41/engram.py b/python/freetoken/models/deepseek_v41/engram.py
new file mode 100644
index 0000000000..1c997a7cfd
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/engram.py
@@ -0,0 +1,262 @@
+"""Engram conditional memory (reference ``inference/engram.py`` + ``Engram``): n-gram hash lookups
+added into the residual streams before layers 1 and 14.
+
+Hashing (host side, ``EngramHash``): token ids are first mapped through a compressed vocabulary
+(tokens that normalize alike collapse together), then each position is hashed with the
+``max_ngram_size - 1`` tokens before it -- a rolling XOR of ``id * multiplier`` over the lookback,
+one bucket range (a distinct prime modulus) per (n-gram order, head). Lookback stops at the sequence
+start (pad id) and never crosses an image span. Addresses depend
+only on the token ids, so the rows can be fetched before the forward starts.
+
+Lookup (``EngramTable`` protocol): the 384M-row fp8 tables live on disk; a backend gathers the
+``n_hash_cols`` rows of every token into a device tensor ``[T, n_hash_cols * head_dim]`` in bf16.
+
+Gate (``EngramLayer``): ``wkv`` (fp8 linear) turns the gathered rows into one key per residual
+stream plus a shared value; the gate is a normalized dot product of stream and key, signed-sqrt
+then sigmoid; ``stream += gate * value``.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+import numpy as np
+import torch
+from freetoken.layers import BaseOP, LinearReplicated
+
+from .args import DeepseekV41Args
+
+
+# ----- hashing ---------------------------------------------------------------------------------
+
+
+def _is_prime(n: int) -> bool:
+    """Deterministic Miller-Rabin for n < 3.3e24 (the bases below are a proven set), matching
+    ``sympy.isprime`` on every 64-bit input."""
+    if n < 2:
+        return False
+    small = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41)
+    for p in small:
+        if n % p == 0:
+            return n == p
+    d, s = n - 1, 0
+    while d % 2 == 0:
+        d //= 2
+        s += 1
+    for a in small:
+        x = pow(a, d, n)
+        if x in (1, n - 1):
+            continue
+        for _ in range(s - 1):
+            x = pow(x, 2, n)
+            if x == n - 1:
+                break
+        else:
+            return False
+    return True
+
+
+def find_next_prime(start: int, seen: set[int]) -> int:
+    """The smallest prime above ``start`` not handed out yet (reference ``find_next_prime``)."""
+    candidate = start + 1
+    while not _is_prime(candidate) or candidate in seen:
+        candidate += 1
+    return candidate
+
+
+def bucket_primes(args: DeepseekV41Args) -> list[list[list[int]]]:
+    """``[layer][ngram order - 2][head]`` bucket moduli: distinct primes drawn in order from
+    ``engram_vocab_size - 1`` upward, never reused across layers / orders / heads."""
+    primes: list[list[list[int]]] = []
+    seen: set[int] = set()
+    for _ in args.engram_layer_ids:
+        per_ngram = []
+        for _ in range(args.engram_max_ngram_size - 1):
+            sizes, current = [], args.engram_vocab_size - 1
+            for _ in range(args.engram_n_heads):
+                current = find_next_prime(current, seen)
+                seen.add(current)
+                sizes.append(current)
+            per_ngram.append(sizes)
+        primes.append(per_ngram)
+    return primes
+
+
+def hash_multipliers(layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int) -> torch.Tensor:
+    """``[n_layers, max_ngram_size]`` odd int64 multipliers from a per-layer RNG (reference
+    ``compute_hash_multipliers``), bounded so ``id * multiplier`` cannot overflow int64."""
+    max_long = np.iinfo(np.int64).max
+    bound = max(1, (max_long // compressed_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        gen = np.random.default_rng(10007 * layer_id)
+        values = gen.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64)
+        rows.append(torch.tensor(values * 2 + 1))
+    return torch.stack(rows)
+
+
+def build_compressed_token_map(tokenizer) -> tuple[list[int], int]:
+    """Token id -> compressed id: tokens that normalize alike (NFKC, accents stripped, lowercased,
+    whitespace collapsed) share one id. Reference ``build_compressed_token_map``; the size feeds the
+    multipliers, so it is asserted against ``engram_compressed_vocab_size``."""
+    from tokenizers import Regex, normalizers
+
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence([
+        normalizers.NFKC(),
+        normalizers.NFD(),
+        normalizers.StripAccents(),
+        normalizers.Lowercase(),
+        normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+        normalizers.Replace(Regex(r"^ $"), sentinel),
+        normalizers.Strip(),
+        normalizers.Replace(sentinel, " "),
+    ])
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+        new_id = key_to_new.get(key)
+        if new_id is None:
+            new_id = len(key_to_new)
+            key_to_new[key] = new_id
+        lookup[token_id] = new_id
+    return lookup, len(key_to_new)
+
+
+class EngramHash:
+    """Maps token runs to table rows: ``row_ids(tokens) -> [T, n_layers, n_hash_cols]`` int64.
+
+    ``token_map`` is the compressed vocabulary (``build_compressed_token_map``); ``pad_id`` is the
+    compressed id of ``engram_pad_id``. Runs on the host (numpy / CPU torch) since the addresses
+    are needed before the forward launches.
+    """
+
+    def __init__(self, args: DeepseekV41Args, token_map: list[int] | torch.Tensor, compressed_vocab_size: int):
+        if compressed_vocab_size != args.engram_compressed_vocab_size:
+            raise ValueError(
+                f"compressed vocab has {compressed_vocab_size} ids, config says {args.engram_compressed_vocab_size}; "
+                "the hash multipliers derive from it, so the tokenizer does not match the checkpoint"
+            )
+        self.args = args
+        self.max_ngram = args.engram_max_ngram_size
+        self.token_map = torch.as_tensor(token_map, dtype=torch.int64)
+        self.pad_id = int(self.token_map[args.engram_pad_id])
+        primes = bucket_primes(args)
+        flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes]
+        self.primes = torch.tensor(primes, dtype=torch.int64)  # [L, ngram-1, heads]
+        self.offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat]), dtype=torch.int64)  # [L, cols]
+        self.multipliers = hash_multipliers(args.engram_layer_ids, self.max_ngram, compressed_vocab_size)  # [L, ngram]
+        self.rows_per_layer = [int(o[-1] + s[-1]) for o, s in zip(self.offsets.tolist(), [[p for pn in layer for p in pn] for layer in primes])]
+
+    @property
+    def n_cols(self) -> int:
+        return (self.max_ngram - 1) * self.args.engram_n_heads
+
+    def row_ids(self, tokens: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor:
+        """Rows for ``tokens [T]`` (raw ids, one contiguous run). ``context`` holds the
+        ``max_ngram_size - 1`` raw ids immediately before the run (fewer -> the run starts a
+        sequence and the missing lookback is the pad id). Returns ``[T, n_layers, n_hash_cols]``."""
+        tokens = tokens.to(torch.int64).cpu()
+        ctx_len = self.max_ngram - 1
+        if context is None:
+            context = tokens.new_empty(0)
+        context = context.to(torch.int64).cpu()[-ctx_len:] if ctx_len else tokens.new_empty(0)
+        raw = torch.cat([context, tokens])
+        # Content-pad ids identify image spans; n-grams stop at them, including across chunks.
+        dead = (raw < 0) | (raw >= self.token_map.numel())
+        compressed = self.token_map[raw.clamp(0, self.token_map.numel() - 1)]
+        n_ctx = context.numel()
+        T = tokens.numel()
+        positions = torch.arange(n_ctx, n_ctx + T)
+        cols = []
+        blocked = torch.zeros(T, dtype=torch.bool)
+        for shift in range(self.max_ngram):
+            src = positions - shift
+            blocked |= (src < 0) | dead[src.clamp_min(0)]
+            cols.append(torch.where(blocked, torch.full_like(src, self.pad_id), compressed[src.clamp_min(0)]))
+        stack = torch.stack(cols, dim=-1)  # [T, max_ngram]; column s is the token s places back
+        products = stack.unsqueeze(1) * self.multipliers  # [T, L, max_ngram]
+        rolling, hashes = products[..., 0], []
+        for i in range(1, self.max_ngram):
+            rolling = torch.bitwise_xor(rolling, products[..., i])
+            hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1])  # [T, L, heads]
+        return torch.cat(hashes, dim=-1) + self.offsets  # [T, L, cols]
+
+
+# ----- lookup ------------------------------------------------------------------------------------
+
+
+class EngramTable(Protocol):
+    """Row store of one Engram layer: ``lookup`` returns the dequantized rows of the current
+    forward's tokens as ``[T, n_hash_cols * head_dim]`` bf16 (staged by the host before launch)."""
+
+    def lookup(self, num_tokens: int) -> torch.Tensor: ...
+
+
+class ZeroEngramTable:
+    """Stand-in when no table is attached (dummy weights / unit tests): every lookup reads zeros."""
+
+    def __init__(self, width: int, device: torch.device) -> None:
+        self.width = width
+        self.device = device
+
+    def lookup(self, num_tokens: int) -> torch.Tensor:
+        return torch.zeros(num_tokens, self.width, dtype=torch.bfloat16, device=self.device)
+
+
+# ----- the layer -----------------------------------------------------------------------------------
+
+
+class EngramLayer(BaseOP):
+    """``streams [T, hc, dim] += gate * value`` with key / value from the gathered n-gram rows."""
+
+    def __init__(self, args: DeepseekV41Args, layer_id: int, *, quant_config=None, prefix: str = ""):
+        self.layer_id = layer_id
+        self.dim = args.dim
+        self.hc_mult = args.hc_mult
+        self.eps = args.norm_eps
+        self.clamp_value = 1e-6
+        self.width = args.engram_hash_cols * args.engram_head_dim
+        self.wkv = LinearReplicated(self.width, args.dim * (args.hc_mult + 1), has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv")
+        self.q_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32)
+        self.k_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32)
+        self._table: EngramTable | None = None
+
+    def attach_table(self, table: EngramTable) -> None:
+        self._table = table
+
+    def forward(self, streams: torch.Tensor, rows: torch.Tensor | None = None, *, image_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """``streams [T, hc, dim]``; ``rows`` overrides the table lookup (tests / reference parity)."""
+        T = streams.shape[0]
+        if rows is None:
+            assert self._table is not None, "Engram table was never attached"
+            rows = self._table.lookup(T)
+        kv = self.wkv.forward(rows.to(streams.dtype))
+        key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
+        key = key.float().view(T, self.hc_mult, self.dim)
+        weight = self.q_weight * self.k_weight  # only ever used as a product
+        h = streams.float()
+        # normalized per (token, stream) over dim, not jointly over the streams
+        rstd = torch.rsqrt(h.square().mean(-1) + self.eps) * torch.rsqrt(key.square().mean(-1) + self.eps)
+        dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5
+        gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot))
+        out = (h + gate.unsqueeze(-1) * value.float().unsqueeze(1)).to(streams.dtype)
+        return out if image_mask is None else torch.where(image_mask[:, None, None], streams, out)
+
+
+__all__ = [
+    "EngramHash",
+    "EngramLayer",
+    "EngramTable",
+    "ZeroEngramTable",
+    "build_compressed_token_map",
+    "bucket_primes",
+    "hash_multipliers",
+]
diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py
new file mode 100644
index 0000000000..55450bbd80
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/engram_table.py
@@ -0,0 +1,296 @@
+"""Engram tables streamed directly from checkpoint shards.
+
+Each Engram layer owns a ``[384M, 256]`` fp8 table (98 GB) plus a ``[384M, 8]`` ue8m0 scale
+tensor (3 GB). The tables stay where the checkpoint put them: ``EngramDiskTable`` maps the shard
+tensor as one ``RowStore`` extent and reads the 24 rows a token needs per fill (io_uring, O_DIRECT,
+duplicates deduplicated), while the scales are host-resident and gathered next to the values, so
+the staged row is exactly the pool row format ``fp8_e8m0_b32`` and the device dequantizes it with
+the same kernel the KV pools use.
+
+Protocol (the Qwen3.8 PLE disk table's, ``models/qwen4_exp/ple_disk.py``): the engine wraps every
+dispatch in ``EngramHost.forward_host_ctx``. Eager: hash + stage + read before the launch. CUDA-graph
+decode: the sampled token lives on the device under overlap scheduling, so the host reads it back,
+fills the graph's pinned staging and releases a per-table flag the captured ``lookup`` WAITs on (or,
+without stream memops, fills before the replay). The fill runs on a worker thread that is started
+BEFORE the graph is launched: with a graph this large (40 layers) the driver's launch call itself
+blocks once the GPU sits on the WAIT with the rest of the graph queued behind it, so a fill scheduled
+after the launch returns would deadlock. The worker works from an immutable dispatch snapshot plus
+the readback event (which the GPU completes on its own), and the dispatch context waits for it after
+the launch so its failure surfaces before the step's output is consumed.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import struct
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Sequence
+
+import torch
+from freetoken.core import Batch
+from freetoken.kernel.pinned import alloc_pinned_tensor
+from freetoken.kvcache.dsv4.v41_row_format import FP8_E8M0_B32
+from freetoken.utils import init_logger
+
+from .engram import EngramHash
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class _TensorExtent:
+    """Where a safetensors tensor's bytes sit in its shard: a contiguous ``[offset, offset + nbytes)``."""
+
+    path: str
+    offset: int
+    nbytes: int
+    shape: tuple[int, ...]
+    dtype: str  # the safetensors dtype string (e.g. "F8_E4M3")
+
+
+def _tensor_extent(folder: str, name: str) -> _TensorExtent:
+    """Locate ``name`` in the checkpoint directory without reading it (the table is streamed in place)."""
+    from freetoken.models.loader import safetensors_weight_map
+
+    path = os.path.join(folder, safetensors_weight_map(folder)[name])
+    with open(path, "rb") as fh:
+        n = struct.unpack("<Q", fh.read(8))[0]
+        header = json.loads(fh.read(n))
+    meta = header[name]
+    start, end = meta["data_offsets"]
+    return _TensorExtent(path, 8 + n + start, end - start, tuple(meta["shape"]), meta["dtype"])
+
+
+@dataclass(frozen=True)
+class EngramRowSource:
+    """The fp8 table tensor in its shard (one extent) and its host-resident scales."""
+
+    path: str
+    base: int
+    num_rows: int
+    head_dim: int
+    scales: torch.Tensor  # [num_rows, head_dim // 32] uint8 (e8m0 codes), CPU
+
+
+def engram_row_source(folder: str, layer_id: int) -> EngramRowSource:
+    """Locate a layer's table in whatever safetensors files the checkpoint directory holds: the HF
+    shards (through the index) or the shards an FTW conversion copied next to it (through their headers)."""
+    from safetensors import safe_open
+
+    weight = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.weight")
+    if weight.dtype != "F8_E4M3":
+        raise ValueError(f"engram table of layer {layer_id} is {weight.dtype}, expected F8_E4M3")
+    rows, head_dim = weight.shape
+    scale = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.scale")
+    with safe_open(scale.path, framework="pt", device="cpu") as f:
+        scales = f.get_tensor(f"layers.{layer_id}.engram.embed.scale").view(torch.uint8).contiguous()
+    if tuple(scales.shape) != (rows, head_dim // 32):
+        raise ValueError(f"engram scales of layer {layer_id} are {tuple(scales.shape)}, expected {(rows, head_dim // 32)}")
+    return EngramRowSource(weight.path, weight.offset, rows, head_dim, scales)
+
+
+class EngramDiskTable:
+    """One layer's table: stages ``[T, n_cols]`` rows into pinned memory on the host, dequantizes on
+    the device. Implements the model's ``EngramTable`` protocol (``lookup``)."""
+
+    def __init__(
+        self, source: EngramRowSource, n_cols: int, device: torch.device, *, max_graph_rows: int, max_extend_tokens: int,
+        use_io_uring: bool = True,
+    ) -> None:
+        from freetoken.kernel.row_store import RowStore
+
+        self.head_dim = source.head_dim
+        self.n_cols = n_cols
+        self.num_rows = source.num_rows
+        self.row_bytes = FP8_E8M0_B32.row_bytes(source.head_dim)  # values | scales
+        self.scales = source.scales
+        self.device = device
+        self.store = RowStore(
+            paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows,
+            row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=use_io_uring,
+        )
+        # the graph staging is allocated up front (pinned alloc inside capture is illegal) and
+        # outlives any one capture so a re-capture keeps the same pointers; the eager staging
+        # grows to the largest prefill chunk seen (never inside a capture)
+        self.max_graph_rows = max_graph_rows
+        self._graph_pinned = alloc_pinned_tensor(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8)
+        self._graph_pinned.zero_()
+        self._graph_dev = torch.empty(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8, device=device)
+        # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the
+        # host coordinator after the launch: a record inside the capture would not be observable)
+        self.graph_consumed = torch.cuda.Event()
+        # the overlap scheduler stages batch k+1 while batch k's lookup copy may still be queued: the
+        # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy
+        # that last read it has run (``_eager_read[i]``, recorded right after that copy)
+        self._eager_pinned: list[torch.Tensor] = []
+        self._eager_read = [torch.cuda.Event(), torch.cuda.Event()]
+        self._eager_slot = 0
+        self._eager_dev = None
+        self._grow_eager(max_extend_tokens)
+        self.flag = alloc_pinned_tensor(1, dtype=torch.int64)
+        self.flag.zero_()
+        self.wait_sync = False  # set by the host coordinator after its probe
+
+    def _grow_eager(self, num_tokens: int) -> None:
+        nbytes = num_tokens * self.n_cols * self.row_bytes
+        if self._eager_pinned and self._eager_pinned[0].numel() >= nbytes:
+            return
+        for event in self._eager_read:
+            event.synchronize()
+        self._eager_pinned = [alloc_pinned_tensor(nbytes, dtype=torch.uint8) for _ in self._eager_read]
+        for buf in self._eager_pinned:
+            buf.zero_()  # a warmup prefill stages nothing and reads whatever sits here
+        self._eager_dev = torch.empty(nbytes, dtype=torch.uint8, device=self.device)
+
+    # ----- host -----
+    def stage(self, row_ids: torch.Tensor, *, graph: bool) -> None:
+        """Queue ``row_ids [T, n_cols]`` (int64, host): values through the store, scales from the resident tensor."""
+        n = row_ids.numel()
+        if graph:
+            self.graph_consumed.synchronize()
+            pinned = self._graph_pinned
+            assert n <= self.max_graph_rows * self.n_cols, f"graph staging holds {self.max_graph_rows} rows, batch has {n // self.n_cols}"
+        else:
+            self._grow_eager(n // self.n_cols)
+            self._eager_slot ^= 1
+            self._eager_read[self._eager_slot].synchronize()
+            pinned = self._eager_pinned[self._eager_slot]
+        ids = row_ids.reshape(-1).to(torch.int64).contiguous()
+        rows = pinned[: n * self.row_bytes].view(n, self.row_bytes)
+        self.store.stage_rows(ids.data_ptr(), n, rows.data_ptr(), self.row_bytes)
+        rows[:, self.head_dim :].copy_(self.scales.index_select(0, ids))
+
+    def flush(self, *, signal: bool) -> None:
+        self.store.flush(self.flag.data_ptr() if signal else 0)
+
+    # ----- device (EngramTable protocol) -----
+    def lookup(self, num_tokens: int) -> torch.Tensor:
+        from freetoken.kernel.row_store import wait_reset
+        from freetoken.kernel.triton.dsv41.pack import unpack_rows
+
+        stream = torch.cuda.current_stream(self.device)
+        capturing = torch.cuda.is_current_stream_capturing()
+        if capturing and self.wait_sync:
+            wait_reset(stream, self.flag)
+        pinned, dev = (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev)
+        n = num_tokens * self.n_cols
+        nbytes = n * self.row_bytes
+        dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True)
+        if not capturing:
+            self._eager_read[self._eager_slot].record(stream)
+        values = unpack_rows(dev[:nbytes].view(n, self.row_bytes), FP8_E8M0_B32, self.head_dim)
+        return values.view(num_tokens, self.n_cols * self.head_dim)
+
+
+def _context(ids: torch.Tensor, position: int, width: int) -> torch.Tensor:
+    """The up-to-``width`` raw ids before ``position`` (fewer at the sequence start)."""
+    return ids[max(0, position - width) : position]
+
+
+class EngramHost:
+    """The per-model coordinator: hashes each batch's token runs and fills every layer's table.
+
+    The token range a prefill hashes is the attention metadata's encoder segments -- the one
+    execution plan the scheduler, the attention layers and the model already agreed on. A graph
+    decode's fill runs on a worker thread from an immutable dispatch snapshot (position + the
+    ``max_ngram - 1`` preceding ids per row) plus the readback of the sampled token; the request
+    objects themselves advance on the main thread right after dispatch and are never read late.
+    """
+
+    def __init__(self, hash: EngramHash, tables: Sequence[EngramDiskTable], device: torch.device, *, sync_mode: str = "auto") -> None:
+        from freetoken.kernel.row_store import probe_wait_sync
+
+        self.hash = hash
+        self.tables = list(tables)
+        self.device = device
+        self.wait_sync = probe_wait_sync(sync_mode, device)
+        for t in self.tables:
+            t.wait_sync = self.wait_sync
+        self._token_readback = alloc_pinned_tensor(min(t.max_graph_rows for t in self.tables), dtype=torch.int32)
+        self._readback_event = torch.cuda.Event()
+        self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="engram-fill")
+        self._pending: Future | None = None
+        logger.info_rank0(f"Engram disk backend: {self.tables[0].store.io_backend()}, {'wait-sync' if self.wait_sync else 'launch-gating'}")
+
+    def close(self) -> None:
+        self._worker.shutdown(wait=True)
+
+    def __del__(self) -> None:
+        worker = getattr(self, "_worker", None)
+        if worker is not None:
+            worker.shutdown(wait=False)
+
+    def fill(self, runs: Sequence[tuple[torch.Tensor, torch.Tensor]], *, graph: bool) -> None:
+        """``runs``: per request ``(context ids, token ids)`` in batch order; stages every layer's rows."""
+        ids = torch.cat([self.hash.row_ids(tokens, ctx) for ctx, tokens in runs]) if runs else torch.empty(0, len(self.tables), self.hash.n_cols, dtype=torch.int64)
+        for l, table in enumerate(self.tables):
+            table.stage(ids[:, l], graph=graph)
+            table.flush(signal=graph and self.wait_sync)
+
+    @torch.inference_mode()
+    def _deferred_fill(self, bs: int, histories: list[tuple[int, ...]]) -> None:
+        """Worker side of a graph decode: wait for the sampled tokens, hash them against the dispatch
+        snapshot, stage, and release the tables' flags. Reads nothing that the main thread mutates.
+        Inference mode is thread-local; staging buffers were created under it by the server."""
+        try:
+            self._readback_event.synchronize()
+            tokens = self._token_readback[:bs].to(torch.int64)
+            runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)]
+            self.fill(runs, graph=True)
+        except BaseException:
+            from freetoken.kernel.row_store import signal
+
+            for t in self.tables:  # unblock the stream before surfacing; the step's output is discarded
+                signal(t.flag)
+            raise
+
+    def _await_pending(self) -> None:
+        """Surface the deferred fill's outcome (it depends only on GPU work already enqueued, so this
+        cannot wait on the host)."""
+        if self._pending is not None:
+            pending, self._pending = self._pending, None
+            pending.result()
+
+    def host_fill_batch(self, batch: Batch, use_graph: bool) -> None:
+        """Stage this batch's rows: eagerly, or -- a graph decode under flag-sync -- on the worker
+        thread once the sampled tokens have been read back."""
+        self._await_pending()
+        width = self.hash.max_ngram - 1
+        if batch.is_decode:
+            reqs = list(batch.padded_reqs)
+            # the dispatch snapshot: each row's position context, taken before anything advances
+            histories = [tuple(_context(r.input_ids, r.device_len - 1, width).tolist()) for r in reqs]
+            if use_graph and self.wait_sync:
+                bs = batch.padded_size
+                self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
+                self._readback_event.record(torch.cuda.current_stream(self.device))
+                self._pending = self._worker.submit(self._deferred_fill, bs, histories)
+                return
+            tokens = batch.input_ids.to("cpu").to(torch.int64)
+            runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)]
+            self.fill(runs, graph=use_graph)
+            return
+        # prefill: the encoder segments are the token range every consumer of this forward executes
+        segments = batch.attn_metadata.segments
+        assert segments is not None and len(segments) == len(batch.reqs)
+        runs = [(_context(r.input_ids, seg.start_pos, width), r.input_ids[seg.start_pos : seg.end]) for r, seg in zip(batch.reqs, segments)]
+        self.fill(runs, graph=False)
+
+    @contextmanager
+    def forward_host_ctx(self, batch: Batch, use_graph: bool):
+        self.host_fill_batch(batch, use_graph)
+        yield
+        # the launch has been enqueued: wait for the deferred fill here so a failure surfaces before
+        # this step's output is consumed (the fill only waits on the GPU readback, never on the host)
+        self._await_pending()
+        if use_graph:
+            # only after this step's fill: the next fill waits on this launch, not the fill on it
+            stream = torch.cuda.current_stream(self.device)
+            for t in self.tables:
+                t.graph_consumed.record(stream)
+
+
+__all__ = ["EngramDiskTable", "EngramHost", "EngramRowSource", "engram_row_source"]
diff --git a/python/freetoken/models/deepseek_v41/indexer.py b/python/freetoken/models/deepseek_v41/indexer.py
new file mode 100644
index 0000000000..4ab3da2bd4
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/indexer.py
@@ -0,0 +1,127 @@
+"""DSV41 Lightning Indexer (reference ``Indexer``): scores the compressed positions and keeps the
+``index_topk`` best per query.
+
+A Full-mode layer owns the index keys: ``k = k_norm(wk(latent))`` from its compressor's unrotated
+latent, RoPE'd at the group position, fp4-quantized (ue8m0 per 32) into the kv source's index pool.
+Every Full / Reindex layer scores its own fp4-baked queries ``wq_b(qr)`` against those keys,
+``sum_h relu(q_h . k) * w_h`` with ``w = weights_proj(x) * softmax_scale * n_heads**-0.5``.
+
+The Hierarchical Sparse Indexer (sec. 2.3.2): the candidate-source layer (the decoder's first Full
+layer) additionally keeps the ``candidate_topk_blocks`` best ``candidate_block_size``-position blocks
+as a candidate pool; the decoder's Reindex layers score only that pool. Logits are produced in query
+blocks bounded by ``LOGITS_BUDGET_BYTES`` so a long-context prefill never materializes ``[T, 1M]``.
+"""
+
+from __future__ import annotations
+
+import torch
+from freetoken.core import get_global_ctx
+from freetoken.kernel.triton.dsv4.fp8_linear import fp4_act_quant_inplace
+from freetoken.layers import BaseOP, LinearReplicated, RMSNorm
+
+from .args import DeepseekV41Args, LayerRole, Mode
+from .rope import apply_rotary_emb, apply_rotary_emb_decode
+
+LOGITS_BUDGET_BYTES = 512 << 20
+
+
+class Indexer(BaseOP):
+    def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""):
+        self.role = role
+        self.n_heads = args.index_n_heads
+        self.head_dim = args.index_head_dim
+        self.rope_dim = args.rope_head_dim
+        self.topk = args.index_topk
+        self.weights_scale = args.index_softmax_scale * args.index_n_heads**-0.5
+        self.candidate_topk_blocks = args.candidate_topk_blocks
+        self.candidate_block_size = args.candidate_block_size
+        self.wq_b = LinearReplicated(args.q_lora_rank, self.n_heads * self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b")
+        self.weights_proj = LinearReplicated(args.dim, self.n_heads, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.weights_proj")
+        if role.mode is Mode.FULL:
+            self.wk = LinearReplicated(args.head_dim, self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wk")
+            self.k_norm = RMSNorm(self.head_dim, args.norm_eps)
+
+    @property
+    def attn(self):
+        return get_global_ctx().attn_backend
+
+    # ----- keys (Full mode) ----------------------------------------------------------------
+    def index_keys(self, latent: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        """``[G, index_head_dim]`` keys from the unrotated latents, rotated at their group positions. The
+        fp4 quantization is the pool's packing."""
+        k = self.k_norm.forward(self.wk.forward(latent))
+        apply_rotary_emb(k[..., -self.rope_dim :], freqs)
+        return k
+
+    # ----- queries ---------------------------------------------------------------------------
+    def _queries(self, x: torch.Tensor, qr: torch.Tensor, rotate) -> tuple[torch.Tensor, torch.Tensor]:
+        q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim))
+        rotate(q[..., -self.rope_dim :])
+        fp4_act_quant_inplace(q, 32)
+        weights = self.weights_proj.forward(x) * self.weights_scale
+        return q, weights
+
+    # ----- selection -------------------------------------------------------------------------
+    def select_prefill(
+        self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, start_pos: int, ratio: int, locs: torch.Tensor,
+        candidates: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Per-query picks for one request's ``n`` new tokens at ``[start_pos, start_pos + n)``.
+
+        ``locs [1, W]`` are the request's live full locs for the positions its queries may see (the
+        source's row of compressed position ``t`` is ``locs[t * ratio] // ratio``); ``candidates [1, n,
+        NC]`` restricts a Reindex layer to the pool. Returns ``(rows [1, n, topk] int32 global main rows,
+        -1 tail; candidates [1, n, NC] | None -- the pool this layer built as candidate source)``.
+        """
+        n = x.shape[0]
+        T = locs.shape[1] // ratio
+        if T == 0:  # a prompt shorter than one compression block: no source row exists yet
+            rows = torch.full((1, n, self.topk), -1, dtype=torch.int32, device=x.device)
+            pool = None
+            if self.role.is_candidate_source:
+                pool = torch.full((1, n, self.candidate_topk_blocks * self.candidate_block_size), -1, dtype=torch.int32, device=x.device)
+            return rows, pool
+        q, w = self._queries(x, qr, lambda t: apply_rotary_emb(t, freqs))
+        live = ((start_pos + torch.arange(1, n + 1, device=x.device)) // ratio).to(torch.int32).view(1, n)
+        width = candidates.shape[-1] if candidates is not None else T
+        qb = max(1, LOGITS_BUDGET_BYTES // max(1, width * 4))
+        picks, pools = [], []
+        for s0 in range(0, n, qb):
+            s1 = min(n, s0 + qb)
+            cand = candidates[:, s0:s1] if candidates is not None else None
+            logits = self.attn.indexer_logits(q[None, s0:s1], w[None, s0:s1], self.role.kv_source, locs, ratio, live[:, s0:s1], T=T, candidates=cand)
+            if cand is not None:
+                pos = self.attn.select_topk_in_candidates(logits, cand, self.topk)
+            else:
+                pos = self.attn.select_topk(logits, live[:, s0:s1], self.topk)
+                if self.role.is_candidate_source:
+                    pools.append(self.attn.select_candidate_blocks(logits, live[:, s0:s1], self.candidate_topk_blocks, self.candidate_block_size))
+            picks.append(self.attn.positions_to_rows(pos, locs, ratio))
+        rows = torch.cat(picks, dim=1)
+        pool = torch.cat(pools, dim=1) if pools else None
+        return rows.to(torch.int32), pool
+
+    def select_decode(
+        self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, pos: torch.Tensor, ratio: int, locs: torch.Tensor,
+        T: int, candidates: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        """Batched one-token selection. ``locs [B, W]`` is the decode snapshot (graph-static width),
+        ``T`` the staged compressed width, ``pos [B]``. Returns ``(rows [B, 1, topk] int32, counts [B, 1]
+        int32 valid picks, candidates [B, 1, NC] | None)``."""
+        B = x.shape[0]
+        q, w = self._queries(x, qr, lambda t: apply_rotary_emb_decode(t.unsqueeze(1), freqs))
+        live = ((pos + 1) // ratio).to(torch.int32).view(B, 1)
+        logits = self.attn.indexer_logits(q.view(B, 1, self.n_heads, self.head_dim), w.view(B, 1, -1), self.role.kv_source, locs, ratio, live, T=T, candidates=candidates)
+        pool = None
+        if candidates is not None:
+            picks = self.attn.select_topk_in_candidates(logits, candidates, self.topk)
+        else:
+            picks = self.attn.select_topk(logits, live, self.topk)
+            if self.role.is_candidate_source:
+                pool = self.attn.select_candidate_blocks(logits, live, self.candidate_topk_blocks, self.candidate_block_size)
+        rows = self.attn.positions_to_rows(picks, locs, ratio)
+        counts = (picks >= 0).sum(dim=-1, dtype=torch.int32)  # valid picks form a prefix
+        return rows.to(torch.int32), counts, pool
+
+
+__all__ = ["Indexer", "LOGITS_BUDGET_BYTES"]
diff --git a/python/freetoken/models/deepseek_v41/model.py b/python/freetoken/models/deepseek_v41/model.py
new file mode 100644
index 0000000000..8112022f76
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/model.py
@@ -0,0 +1,289 @@
+"""DeepSeek-V4.1-Flash model (engine-native port of ``inference/model.py``).
+
+    embed -> hc_mult residual streams -> [Engram at 1, 14] -> 40 blocks (single-pass mHC around
+    DSV41 attention and MoE) -> collapse with the last block's pre gates -> norm -> head
+
+Single-pass mHC (``layers/mhc.py``): each sublayer boundary applies the PREVIOUS sublayer's post /
+comb to the streams, predicts this sublayer's (pre, post, comb) from the updated streams, and mixes
+the sublayer input with the pre gates the previous sublayer predicted. The chain state
+``(y, residual, post, comb, pre)`` threads through the layer loop; layer 0's attention starts it
+with a one-hot pre on stream 0 and no pending post.
+
+Causal Encoder-Decoder: the decoder's kv source (layer 20) compresses ITS INPUT -- the final encoder
+hidden state -- into the global KV every decoder layer reads. Under Decoder SWA Bounded Replay the
+decoder layers run on each request's last ``window_size`` prompt tokens only (their sliding window
+truncated at the start of that window, see the backend; a prefix hit stops before it, so those tokens
+are always prefilled); layer 20 still publishes global KV for every prompt token.
+``--swa-decoder-replay exact`` runs the decoder on every token (the reference numerics).
+
+KV addressing is the attention backend's; pool buffers are read off the live pool per access, so a
+runtime cache rebuild needs no unbind. Decode is batched and CUDA-graph safe (the DSV4 precedent).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, NamedTuple
+
+import torch
+from freetoken.attention.dsv41_sparse import DSV41AttnMetadata, PrefillSegment
+from freetoken.core import get_global_ctx
+from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding
+from freetoken.layers.mhc import mhc_fused_post_pre_single_pass, mhc_mix_input, mhc_post
+from freetoken.models.blocks import BaseLLMModel, embed_input_ids
+
+from .args import DeepseekV41Args
+from .attention import DSV41Attention, DecodeStepContext
+from .engram import EngramLayer
+from .moe import MoE
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+
+POST_MULT = 2.0  # post gates are 2 * sigmoid (reference hc_split_sinkhorn)
+
+
+class Streams(NamedTuple):
+    """The single-pass mHC chain state between sublayers: the last sublayer output ``y`` with its
+    not-yet-applied ``post`` / ``comb``, the residual streams ``[T, hc, dim]`` they apply to, and
+    the ``pre`` gates ``[T, hc]`` the next sublayer mixes its input with."""
+
+    y: torch.Tensor | None
+    residual: torch.Tensor
+    post: torch.Tensor | None
+    comb: torch.Tensor | None
+    pre: torch.Tensor
+
+    def materialize(self) -> torch.Tensor:
+        """The streams with the pending post applied."""
+        if self.post is None:
+            return self.residual
+        return mhc_post(self.y, self.residual, self.post, self.comb)
+
+
+class Block(BaseOP):
+    def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
+        args: DeepseekV41Args = config.dsv41_args
+        self.layer_id = layer_id
+        self.role = args.roles[layer_id]
+        self.norm_eps = args.norm_eps
+        self.hc_eps = args.hc_eps
+        self.sinkhorn = args.hc_sinkhorn_iters
+        hc, dim = args.hc_mult, args.dim
+        mix = (2 + hc) * hc
+        self.attn = DSV41Attention(args, self.role, quant_config=config.quant, prefix=f"{prefix}.attn")
+        self.ffn = MoE(config, layer_id, prefix=f"{prefix}.ffn")
+        self.attn_norm = RMSNorm(dim, args.norm_eps)
+        self.ffn_norm = RMSNorm(dim, args.norm_eps)
+        self.engram = EngramLayer(args, layer_id, quant_config=config.quant, prefix=f"{prefix}.engram") if layer_id in args.engram_layer_ids else None
+        self.hc_attn_fn = torch.empty(mix, hc * dim, dtype=torch.float32)
+        self.hc_ffn_fn = torch.empty(mix, hc * dim, dtype=torch.float32)
+        self.hc_attn_base = torch.empty(mix, dtype=torch.float32)
+        self.hc_ffn_base = torch.empty(mix, dtype=torch.float32)
+        self.hc_attn_scale = torch.empty(3, dtype=torch.float32)
+        self.hc_ffn_scale = torch.empty(3, dtype=torch.float32)
+
+    def _boundary(self, s: Streams, fn, scale, base):
+        residual, post, comb, pre_next, x = mhc_fused_post_pre_single_pass(
+            s.y if s.y is not None else s.residual[:, 0], s.residual, s.post, s.comb, s.pre, fn, scale, base,
+            self.norm_eps, self.hc_eps, POST_MULT, self.sinkhorn,
+        )
+        return residual, post, comb, pre_next, x
+
+    def attn_input(self, s: Streams, *, image_mask: torch.Tensor | None = None):
+        """Apply the pending post, predict the attention gates, mix and norm the attention input.
+        Returns ``(residual, post, comb, pre_next, x)``."""
+        if self.engram is not None:
+            s = Streams(None, self.engram.forward(s.materialize(), image_mask=image_mask), None, None, s.pre)
+        residual, post, comb, pre_next, x = self._boundary(s, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+        return residual, post, comb, pre_next, self.attn_norm.forward(x)
+
+    def ffn_step(self, y: torch.Tensor, residual, post, comb, pre, *, image_mask: torch.Tensor | None = None) -> Streams:
+        """The FFN sublayer after the attention output ``y``: boundary, norm, MoE; returns the chain state."""
+        residual, post, comb, pre_next, x = self._boundary(Streams(y, residual, post, comb, pre), self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+        return Streams(self.ffn.forward(self.ffn_norm.forward(x), image_mask=image_mask), residual, post, comb, pre_next)
+
+    def forward_prefill(self, s: Streams, segments: List[PrefillSegment], positions: torch.Tensor, image_mask: torch.Tensor | None = None) -> Streams:
+        residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask)
+        if self.attn.compressor is not None:
+            self.attn.publish_prefill(x, segments)
+        y = self.attn.forward_prefill(x, segments, positions)
+        return self.ffn_step(y, residual, post, comb, pre_next, image_mask=image_mask)
+
+    def forward_prefill_narrowed(
+        self, s: Streams, segments: List[PrefillSegment], rows: torch.Tensor, query_segments: List[PrefillSegment],
+        query_positions: torch.Tensor, image_mask: torch.Tensor | None = None, query_image_mask: torch.Tensor | None = None,
+    ) -> Streams | None:
+        """``forward_prefill`` with a row cut after the global KV is published: the attention input and
+        ``publish_prefill`` cover every token of ``segments``, the attention and the FFN only ``rows``
+        (tiled by ``query_segments``) -- the CED decoder's kv source under bounded replay. None when no
+        query segment is left (every chunk in the batch ends before its prompt's last window)."""
+        residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask)
+        if self.attn.compressor is not None:
+            self.attn.publish_prefill(x, segments)
+        if not query_segments:
+            return None
+        y = self.attn.forward_prefill(x[rows], query_segments, query_positions)
+        return self.ffn_step(y, residual[rows], post[rows], comb[rows], pre_next[rows], image_mask=query_image_mask)
+
+    def forward_decode(self, s: Streams, pos: torch.Tensor, rows: torch.Tensor, dctx, cmp_stage_cap: int) -> Streams:
+        residual, post, comb, pre_next, x = self.attn_input(s)
+        y = self.attn.forward_decode(x, pos, rows, dctx, cmp_stage_cap)
+        return self.ffn_step(y, residual, post, comb, pre_next)
+
+
+class Transformer(BaseOP):
+    def __init__(self, config: ModelConfig, *, prefix: str = "model"):
+        args: DeepseekV41Args = config.dsv41_args
+        self.args = args
+        self.embed = VocabParallelEmbedding(num_embeddings=args.vocab_size, embedding_dim=args.dim)
+        self.layers = OPList([Block(config, i, prefix=f"{prefix}.layers.{i}") for i in range(args.n_layers)])
+        self.norm = RMSNorm(args.dim, args.norm_eps)
+
+    def bind(self, device: torch.device) -> None:
+        tables: dict = {}  # one RoPE table per distinct parameter set, shared by the layers using it
+        for block in self.layers.op_list:
+            block.attn.bind(device, tables)
+
+    def _entry(self, input_ids: torch.Tensor) -> Streams:
+        h = embed_input_ids(self.embed, input_ids, get_global_ctx().batch)
+        residual = h.unsqueeze(1).expand(-1, self.args.hc_mult, -1).contiguous()
+        pre = torch.zeros(h.shape[0], self.args.hc_mult, dtype=torch.float32, device=h.device)
+        pre[:, 0] = 1.0
+        return Streams(None, residual, None, None, pre)
+
+    def _exit(self, s: Streams) -> torch.Tensor:
+        return self.norm.forward(mhc_mix_input(s.materialize(), s.pre))
+
+    def prefill(self, input_ids: torch.Tensor, positions: torch.Tensor, md: DSV41AttnMetadata) -> torch.Tensor:
+        """Ragged prefill over the flat new-token stream; returns the hidden states the head reads (the
+        decoder stream under bounded replay, where ``md.last_indices`` also point, plus the placeholder
+        row ``md.decoder_pad`` asks for)."""
+        s = self._entry(input_ids)
+        dtype = s.residual.dtype
+        image_mask = input_ids >= self.args.vocab_size if get_global_ctx().batch.mm_embeds is not None else None
+        segments, dec_segments, rows = md.segments, md.decoder_segments, md.decoder_rows
+        decoder_start = self.args.decoder_start_layer
+        if rows is not None:
+            dec_positions = positions[rows]
+            dec_image_mask = None if image_mask is None else image_mask[rows]
+        for block in self.layers.op_list:
+            if rows is None or block.layer_id < decoder_start:
+                s = block.forward_prefill(s, segments, positions, image_mask)
+            elif block.layer_id == decoder_start:
+                s = block.forward_prefill_narrowed(s, segments, rows, dec_segments, dec_positions, image_mask, dec_image_mask)
+                if s is None:
+                    break
+            else:
+                s = block.forward_prefill(s, dec_segments, dec_positions, dec_image_mask)
+        pad = [torch.zeros(1, self.args.dim, dtype=dtype, device=input_ids.device)] if md.decoder_pad else []
+        return torch.cat([self._exit(s)] + pad) if s is not None else pad[0]
+
+    def decode(self, input_ids: torch.Tensor, pos: torch.Tensor, md: DSV41AttnMetadata, cmp_stage_cap: int) -> torch.Tensor:
+        B = input_ids.shape[0]
+        rows = torch.arange(B, device=input_ids.device)
+        s = self._entry(input_ids)
+        # the layer-invariant step context, resolved once (a capture records it once): ring slots and
+        # window candidates of the shared pool -- and of the private rings when bounded replay keeps
+        # the decoder's window KV per request -- plus each RoPE table at the step's positions
+        dctx = DecodeStepContext.build([b.attn for b in self.layers.op_list], md, pos, rows, self.args.window_size, self.args.index_topk)
+        for block in self.layers.op_list:
+            s = block.forward_decode(s, pos, rows, dctx, cmp_stage_cap)
+        return self._exit(s)
+
+
+class DeepseekV41ForCausalLM(BaseLLMModel):
+    def __init__(self, config: ModelConfig):
+        self._config = config
+        self._args: DeepseekV41Args = config.dsv41_args
+        self.model = Transformer(config)
+        self.head = ParallelLMHead(num_embeddings=self._args.vocab_size, embedding_dim=self._args.dim, quant_config=config.quant, prefix="head")
+        self._bound = False
+        if config.is_multimodal:
+            from .vision import DeepseekV41Vision
+
+            self.visual = DeepseekV41Vision(config.vision_config, config.hidden_size,
+                                           num_layers=self._args.n_layers, num_experts=self._args.n_routed_experts)
+
+    def encode(self, item):
+        return self.visual.forward(item)
+
+    def place_encoder_weights(self, mode: str) -> None:
+        self.visual.place_weights(mode)
+
+    def _ensure_bound(self) -> None:
+        if not self._bound:
+            self.model.bind(get_global_ctx().kv_cache.device)
+            if self._config.is_multimodal:
+                for i, block in enumerate(self.model.layers.op_list):
+                    block.ffn.gate._bias_vl = self.visual.routing_bias[i]
+            self._bound = True
+
+    def mark_for_rebind(self) -> None:
+        """A runtime pool rebuild changes nothing the model caches (rope tables are pool-independent),
+        but the engine asks; re-bind is cheap and keeps the contract with the DSV4 sibling."""
+        self._bound = False
+
+    def engram_layers(self) -> List[EngramLayer]:
+        return [b.engram for b in self.model.layers.op_list if b.engram is not None]
+
+    def load_host_tables(self, engine_config) -> int:
+        """Attach the Engram tables (streamed from the checkpoint shards; zeros for dummy weights)
+        and take over ``forward_host_ctx`` so every dispatch stages its rows. Returns the pinned
+        host bytes the engine reserves from its pin budget."""
+        import os
+
+        layers = self.engram_layers()
+        if not layers:
+            return 0
+        device = torch.device("cuda", torch.cuda.current_device())
+        if getattr(engine_config, "use_dummy_weight", False):
+            from .engram import ZeroEngramTable
+
+            for layer in layers:
+                layer.attach_table(ZeroEngramTable(layer.width, device))
+            return 0
+        from freetoken.utils import download_hf_weight
+        from transformers import AutoTokenizer
+
+        from .engram import EngramHash, build_compressed_token_map
+        from .engram_table import EngramDiskTable, EngramHost, engram_row_source
+
+        folder = download_hf_weight(engine_config.model_path)
+        token_map, compressed_vocab = build_compressed_token_map(AutoTokenizer.from_pretrained(folder, trust_remote_code=True))
+        hash = EngramHash(self._args, token_map, compressed_vocab)
+        graph_rows = max(engine_config.max_running_req, engine_config.cuda_graph_max_bs or 0, 1)
+        tables = [
+            EngramDiskTable(
+                engram_row_source(folder, layer.layer_id), hash.n_cols, device,
+                max_graph_rows=graph_rows,
+                max_extend_tokens=engine_config.max_extend_tokens,
+                use_io_uring=os.getenv("FREETOKEN_ENGRAM_IO_URING", "1") != "0",
+            )
+            for layer in layers
+        ]
+        for layer, table in zip(layers, tables):
+            layer.attach_table(table)
+        self._engram_host = EngramHost(hash, tables, device, sync_mode=os.getenv("FREETOKEN_ENGRAM_SYNC", "auto"))
+        self.forward_host_ctx = self._engram_host.forward_host_ctx
+        return sum(t._graph_pinned.numel() + sum(b.numel() for b in t._eager_pinned) for t in tables)
+
+    def forward(self) -> torch.Tensor:
+        self._ensure_bound()
+        batch = get_global_ctx().batch
+        md = batch.attn_metadata
+        assert isinstance(md, DSV41AttnMetadata)
+        if batch.is_prefill:
+            hidden = self.model.prefill(batch.input_ids, batch.positions.long(), md)
+        else:
+            input_ids = batch.input_ids
+            B = batch.padded_size
+            pos = batch.positions.long().view(-1)[:B]
+            if torch.cuda.is_current_stream_capturing():
+                cmp_stage_cap = md.stage_width - 1
+            else:
+                cmp_stage_cap = int(pos.max().item())
+            hidden = self.model.decode(input_ids.view(B), pos, md, cmp_stage_cap)
+        return self.head.forward(hidden)
+
+
+__all__ = ["DeepseekV41ForCausalLM", "Transformer", "Block", "Streams"]
diff --git a/python/freetoken/models/deepseek_v41/moe.py b/python/freetoken/models/deepseek_v41/moe.py
new file mode 100644
index 0000000000..91c62ecb01
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/moe.py
@@ -0,0 +1,90 @@
+"""DeepSeek-V4.1 MoE block: sqrtsoftplus router with a selection-only correction bias, one shared
+clamped-SwiGLU expert, and the routed MXFP4 experts on the shared offload cache.
+
+Router (reference ``Gate``): ``scores = sqrt(softplus(x @ W^T))`` in fp32; the bias picks the experts,
+the unbiased scores weight them (renormalized, ``+1e-20`` as in training), times ``route_scale``.
+Image-span tokens select experts using the checkpoint's separate ``bias_vl``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32
+from freetoken.kernel.triton.dsv4.router import router_select
+from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu
+from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer
+
+from .args import DeepseekV41Args
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+
+TopK = Tuple[torch.Tensor, torch.Tensor]
+
+
+class Router(BaseOP):
+    def __init__(self, args: DeepseekV41Args):
+        self.topk = args.n_activated_experts
+        self.score_func = args.score_func
+        self.norm_topk_prob = args.norm_topk_prob
+        self.route_scale = args.route_scale
+        self.weight = torch.empty(args.n_routed_experts, args.dim, dtype=torch.bfloat16)
+        self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32)
+        self._bias_vl = None  # bound to the visual stack's per-layer routing bias after loading
+
+    def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> TopK:
+        if image_mask is not None:
+            assert self._bias_vl is not None, "image routing bias was not loaded"
+        scores = bf16_linear_fp32(x, self.weight)
+        return router_select(
+            scores.view(-1, scores.shape[-1]), self.bias, self.topk,
+            score_func=self.score_func, normalize=self.norm_topk_prob, route_scale=self.route_scale,
+            alternate_bias=self._bias_vl, alternate_mask=image_mask,
+        )
+
+
+class SharedExpert(BaseOP):
+    """The dense clamped-SwiGLU expert every token passes through (fp8 block-32 linears)."""
+
+    def __init__(self, args: DeepseekV41Args, *, quant_config=None, prefix: str = ""):
+        self.w1 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w1")
+        self.w2 = LinearReplicated(args.moe_inter_dim, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w2")
+        self.w3 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w3")
+        self.limit = args.swiglu_limit
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = fused_swiglu(self.w1.forward(x), self.w3.forward(x), self.limit, x.dtype)
+        return self.w2.forward(h)
+
+
+class MoE(BaseOP):
+    def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
+        args: DeepseekV41Args = config.dsv41_args
+        self.dim = args.dim
+        self.gate = Router(args)
+        self.shared_experts = SharedExpert(args, quant_config=config.quant, prefix=f"{prefix}.shared_experts")
+        # every layer is MoE: the offload cache indexes experts by layer id directly
+        self.experts = make_moe_layer(
+            config,
+            layer_id=layer_id,
+            activation="swiglu_clamp",
+            renormalize=args.norm_topk_prob,
+            limit=args.swiglu_limit,
+            quant_config=config.quant,
+            prefix=f"{prefix}.experts",
+        )
+
+    def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> torch.Tensor:
+        shape = x.shape
+        x = x.view(-1, self.dim)
+        weights, indices = self.gate.forward(x, image_mask=image_mask)
+        # the shared expert runs before the routed experts (same stream, so it does not overlap the
+        # hybrid decode's CPU work, whose submit follows it -- reordering that is follow-up work)
+        shared = self.shared_experts.forward(x)
+        routed = self.experts.routed_forward(x, weights, indices)
+        return (routed + shared).view(shape)
+
+
+__all__ = ["MoE", "Router", "SharedExpert"]
diff --git a/python/freetoken/models/deepseek_v41/rope.py b/python/freetoken/models/deepseek_v41/rope.py
new file mode 100644
index 0000000000..6d7df9f8d3
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/rope.py
@@ -0,0 +1,58 @@
+"""Rotary embeddings for DeepSeek-V4.1 (verbatim reference math, ``inference/model.py``).
+
+Interleaved complex rotation on the last ``rope_head_dim`` channels, with YaRN when
+``original_seq_len > 0``. Two regimes share two cached tables: window-only layers at the base
+theta without YaRN, compressing layers at ``compress_rope_theta`` under YaRN. A compressed latent
+stands for the first token of its group, so group ``j`` rotates at position ``j * ratio``; the
+attention output is rotated back (``inverse=True``) because K == V share one rotated latent.
+"""
+
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+
+import torch
+
+
+def precompute_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int) -> torch.Tensor:
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    if original_seq_len > 0:
+
+        def corrected_dim(rotations):
+            return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base))
+
+        low = max(math.floor(corrected_dim(beta_fast)), 0)
+        high = min(math.ceil(corrected_dim(beta_slow)), dim - 1)
+        ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1)
+        smooth = 1 - ramp
+        freqs = freqs / factor * (1 - smooth) + freqs * smooth
+    freqs = torch.outer(torch.arange(seqlen), freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+@lru_cache(2)  # the model's two rope regimes; a 1M table is ~256 MB, so never one per layer
+def get_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int, device: torch.device) -> torch.Tensor:
+    return precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow).to(device)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor:
+    """In-place rotation of ``x[..., :]`` (``[T, rd]`` or ``[T, H, rd]``) by per-token ``freqs_cis [T, rd//2]``."""
+    xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs_cis = freqs_cis.conj()
+    if xc.ndim == 3:
+        freqs_cis = freqs_cis.unsqueeze(1)
+    x.copy_(torch.view_as_real(xc * freqs_cis).flatten(-2))
+    return x
+
+
+def apply_rotary_emb_decode(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor:
+    """Batched decode rotation with PER-ROW freqs: ``x [B, 1, ..., rd]``, ``freqs_cis [B, rd//2]``
+    (the fused dsv4 triton kernel; bit-identical to ``apply_rotary_emb``)."""
+    from freetoken.kernel.triton.dsv4.rope import rope_decode_inplace
+
+    return rope_decode_inplace(x, freqs_cis, inverse)
+
+
+__all__ = ["precompute_freqs_cis", "get_freqs_cis", "apply_rotary_emb", "apply_rotary_emb_decode"]
diff --git a/python/freetoken/models/deepseek_v41/vision.py b/python/freetoken/models/deepseek_v41/vision.py
new file mode 100644
index 0000000000..f5fea969ea
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/vision.py
@@ -0,0 +1,139 @@
+"""DeepSeek ViT and 3x3 pixel-unshuffle aligner, with the shared encoder weight streamer."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from freetoken.layers import BaseOP, LinearReplicated, OPList
+from freetoken.models.weight_stream import BlockWeightStreamer
+
+
+class VisionNorm(BaseOP):
+    def __init__(self, dim: int):
+        self.weight = torch.empty(dim, dtype=torch.float32)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = x.float()
+        return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6) * self.weight).to(x.dtype)
+
+
+class VisionAttention(BaseOP):
+    def __init__(self, vc):
+        self.n_heads = vc.num_attention_heads
+        self.head_dim = vc.hidden_size // self.n_heads
+        self.wqkv = LinearReplicated(vc.hidden_size, 3 * vc.hidden_size, has_bias=True)
+        self.wo = LinearReplicated(vc.hidden_size, vc.hidden_size, has_bias=True)
+
+    def forward(self, x, cos, sin):
+        n = x.shape[0]
+        q, k, v = (t.reshape(n, self.n_heads, self.head_dim) for t in self.wqkv.forward(x).chunk(3, -1))
+        def rotary(t):
+            a, b = t.float().chunk(2, -1)
+            return torch.cat((a * cos - b * sin, b * cos + a * sin), -1).to(t.dtype)
+        q, k = rotary(q), rotary(k)
+        q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v))
+        y = F.scaled_dot_product_attention(q, k, v)
+        return self.wo.forward(y[0].transpose(0, 1).reshape(n, -1))
+
+
+class VisionMLP(BaseOP):
+    def __init__(self, vc):
+        self.w1 = LinearReplicated(vc.hidden_size, 2 * vc.intermediate_size, has_bias=False)
+        self.w2 = LinearReplicated(vc.intermediate_size, vc.hidden_size, has_bias=False)
+
+    def forward(self, x):
+        gate, up = self.w1.forward(x).chunk(2, -1)
+        return self.w2.forward(F.silu(gate) * up)
+
+
+class VisionBlock(BaseOP):
+    def __init__(self, vc):
+        self.norm1 = VisionNorm(vc.hidden_size)
+        self.attn = VisionAttention(vc)
+        self.norm2 = VisionNorm(vc.hidden_size)
+        self.mlp = VisionMLP(vc)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn.forward(self.norm1.forward(x), cos, sin)
+        return x + self.mlp.forward(self.norm2.forward(x))
+
+
+class PatchEmbed(BaseOP):
+    def __init__(self, vc):
+        self.proj = LinearReplicated(3 * vc.patch_size**2, vc.hidden_size, has_bias=True)
+
+    def forward(self, x):
+        return self.proj.forward(x.flatten(1))
+
+
+class VisionTransformer(BaseOP):
+    def __init__(self, vc):
+        self.patch_embed = PatchEmbed(vc)
+        self.blocks = OPList([VisionBlock(vc) for _ in range(vc.num_hidden_layers)])
+        self.norm = VisionNorm(vc.hidden_size)
+        self._vc = vc
+        self._streamer = None
+
+    def place_weights(self, mode: str) -> None:
+        if mode == "host" and self._streamer is None:
+            self._streamer = BlockWeightStreamer(self.blocks.op_list, self.patch_embed.proj.weight.device)
+        elif mode == "gpu" and self._streamer is not None:
+            self._streamer.unstream()
+            self._streamer = None
+        elif mode not in ("gpu", "host"):
+            raise ValueError(f"unknown vision weight placement {mode!r}")
+
+    def forward(self, patches, h: int, w: int):
+        x = self.patch_embed.forward(patches)
+        dim = self._vc.hidden_size // self._vc.num_attention_heads // 2
+        inv = 1.0 / (self._vc.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=x.device) / dim))
+        rows = torch.arange(h, device=x.device)[:, None].expand(h, w)
+        cols = torch.arange(w, device=x.device)[None, :].expand(h, w)
+        angles = (torch.stack((rows, cols), -1).reshape(-1, 2, 1).float() * inv).flatten(1)
+        cos, sin = angles.cos().unsqueeze(1), angles.sin().unsqueeze(1)
+        blocks = enumerate(self.blocks.op_list) if self._streamer is None else self._streamer.blocks(self.blocks.op_list)
+        for _, block in blocks:
+            x = block.forward(x, cos, sin)
+        return self.norm.forward(x)
+
+
+class VisionAligner(BaseOP):
+    def __init__(self, vc, dim: int):
+        self.ratio = vc.downsample_ratio
+        self.w1 = LinearReplicated(vc.hidden_size * self.ratio**2, dim, has_bias=True)
+        self.w2 = LinearReplicated(dim, dim, has_bias=True)
+
+    def forward(self, x, h: int, w: int):
+        r = self.ratio
+        x = x.view(h, w, -1).permute(2, 0, 1)
+        x = F.pad(x, (0, -w % r, 0, -h % r))
+        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).T
+        return self.w2.forward(F.gelu(self.w1.forward(x)))
+
+
+class DeepseekV41Vision(BaseOP):
+    def __init__(self, vc, dim: int, *, num_layers: int, num_experts: int):
+        self.vision = VisionTransformer(vc)
+        self.aligner = VisionAligner(vc, dim)
+        self.image_start = torch.empty(dim)
+        self.image_end = torch.empty(dim)
+        self.image_newline = torch.empty(dim)
+        # Keep all image-only parameters under visual, including FTW/text-only filtering.
+        self.routing_bias = torch.empty(num_layers, num_experts, dtype=torch.float32)
+
+    def place_weights(self, mode: str) -> None:
+        self.vision.place_weights(mode)
+
+    @torch.inference_mode()
+    def forward(self, item):
+        _, h, w = item.grid_thw
+        patches = item.feature.to(device=self.image_start.device, dtype=self.image_start.dtype)
+        x = self.aligner.forward(self.vision.forward(patches, h, w), h, w)
+        r = self.aligner.ratio
+        nh, nw = (h + r - 1) // r, (w + r - 1) // r
+        out = x.new_empty(nh * (nw + 1) + 2, x.shape[-1])
+        out[0], out[-1] = self.image_start, self.image_end
+        rows = out[1:-1].view(nh, nw + 1, -1)
+        rows[:, :nw] = x.view(nh, nw, -1)
+        rows[:, nw] = self.image_newline
+        return out
diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py
new file mode 100644
index 0000000000..1ccd23eaef
--- /dev/null
+++ b/python/freetoken/models/deepseek_v41/weight.py
@@ -0,0 +1,242 @@
+"""Weight loading for DeepSeek-V4.1-Flash.
+
+* :func:`iter_weights` streams the resident (non-expert) tensors keyed by the model's attribute
+  paths: checkpoint ``layers.N.*`` -> ``model.layers.N.*``, ``embed`` / ``norm`` -> ``model.*``,
+  ``head`` -> ``head``. fp8 linears keep their e4m3 payload and e8m0 ``scale`` (declared by the quant
+  method as ``weight_scale_inv``); ``wo_a`` keeps its e4m3 payload and ``wo_a_scale`` too (the layer
+  dequantizes in-kernel for the reference's bf16 grouped einsum). The engine casts everything else to
+  the declared parameter dtype (fp32 compressor weights, hc gates, engram gate weights).
+* :func:`iter_expert_pieces` streams the routed MXFP4 experts (e2m1 pairs + e8m0 per-32 scales) per
+  expert for the offload banks.
+
+Dropped: the DSpark ``mtp.*`` stage; vision weights and ``gate.bias_vl`` when serving text-only.
+The Engram tables (``layers.{1,14}.engram.embed.*``) are
+not state-dict tensors: the disk row store maps them in place (``engram_table.py``); an FTW conversion
+carries them as side files (:func:`ftw_side_files`).
+"""
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+import json
+import os
+import re
+
+import safetensors
+import torch
+from freetoken.distributed import get_tp_info
+from freetoken.layers.quantization import QuantKind
+from freetoken.models.loader import drop_page_cache
+from tqdm import tqdm
+
+from .args import DeepseekV41Args, Mode
+
+
+class _ShardReader:
+    def __init__(self, folder: str, weight_map: dict, device):
+        self._folder = folder
+        self._weight_map = weight_map
+        self._device = str(device)
+        self._handles: dict[str, object] = {}
+        self._stack = ExitStack()
+
+    def has(self, name: str) -> bool:
+        return name in self._weight_map
+
+    def get(self, name: str) -> torch.Tensor:
+        shard = self._weight_map[name]
+        handle = self._handles.get(shard)
+        if handle is None:
+            handle = self._stack.enter_context(safetensors.safe_open(os.path.join(self._folder, shard), framework="pt", device=self._device))
+            self._handles[shard] = handle
+        return handle.get_tensor(name)
+
+    def close(self) -> None:
+        try:
+            self._stack.close()
+        finally:
+            for shard in self._handles:
+                drop_page_cache(os.path.join(self._folder, shard))
+            self._handles.clear()
+
+
+def weight_map(model_path: str) -> dict:
+    with open(os.path.join(model_path, "model.safetensors.index.json")) as f:
+        return json.load(f)["weight_map"]
+
+
+def load_args(model_path: str) -> DeepseekV41Args:
+    from freetoken.utils.hf import cached_load_hf_config
+
+    return DeepseekV41Args.from_hf(cached_load_hf_config(model_path))
+
+
+def dequant_fp8_block(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """Square block-scaled e4m3 -> bf16; ``scale`` holds e8m0 codes (``2^(code-127)``), the block
+    edge follows from the shapes."""
+    n, k = weight.shape
+    block = n // scale.shape[0]
+    assert (n // block, k // block) == tuple(scale.shape), (weight.shape, scale.shape)
+    s = torch.exp2(scale.view(torch.uint8).to(torch.float32) - 127.0)
+    s = s.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1)
+    return (weight.to(torch.float32) * s).to(torch.bfloat16)
+
+
+def iter_weights(model_path: str, device, *, include_moe_experts: bool = True, include_non_moe: bool = True, include_vision: bool = True):
+    args = load_args(model_path)
+    reader = _ShardReader(model_path, weight_map(model_path), device)
+    get = reader.get
+    def linear(src: str, dst: str):
+        yield f"{dst}.weight", get(f"{src}.weight")
+        if reader.has(f"{src}.scale"):
+            yield f"{dst}.weight_scale_inv", get(f"{src}.scale")
+
+    try:
+        if include_moe_experts:
+            # Resident experts support unquantized checkpoints; FP4 uses the offload cache.
+            if reader.has("layers.0.ffn.experts.0.w1.scale"):
+                raise ValueError(
+                    "DeepSeek-V4.1's fp4 routed experts are served from the offload cache; run with "
+                    "--moe-strategy offload (include_moe_experts must be False)."
+                )
+            for L in range(args.n_layers):
+                p = f"layers.{L}.ffn.experts"
+                gate_up = torch.stack([torch.cat([get(f"{p}.{e}.w1.weight"), get(f"{p}.{e}.w3.weight")]) for e in range(args.n_routed_experts)])
+                down = torch.stack([get(f"{p}.{e}.w2.weight") for e in range(args.n_routed_experts)])
+                yield f"model.layers.{L}.ffn.experts.gate_up_proj", gate_up
+                yield f"model.layers.{L}.ffn.experts.down_proj", down
+        if not include_non_moe:
+            return
+        if include_vision and reader.has("vision.patch_embed.proj.weight"):
+            yield from iter_vision_weights(model_path, device)
+        yield "model.embed.weight", get("embed.weight")
+        yield "model.norm.weight", get("norm.weight")
+        yield "head.weight", get("head.weight")
+        for L in range(args.n_layers):
+            role = args.roles[L]
+            a, m = f"layers.{L}.attn", f"model.layers.{L}.attn"
+            yield from linear(f"{a}.wq_a", f"{m}.wq_a")
+            yield f"{m}.q_norm.weight", get(f"{a}.q_norm.weight")
+            yield from linear(f"{a}.wq_b", f"{m}.wq_b")
+            yield from linear(f"{a}.wkv", f"{m}.wkv")
+            yield f"{m}.kv_norm.weight", get(f"{a}.kv_norm.weight")
+            yield f"{m}.wo_a", get(f"{a}.wo_a.weight")
+            if reader.has(f"{a}.wo_a.scale"):  # the fp8 payload + scale; the layer dequantizes in-kernel
+                yield f"{m}.wo_a_scale", get(f"{a}.wo_a.scale")
+            yield from linear(f"{a}.wo_b", f"{m}.wo_b")
+            yield f"{m}.attn_sink", get(f"{a}.attn_sink")
+            if role.mode is Mode.FULL:
+                c = f"{a}.compressor"
+                yield f"{m}.compressor.norm.weight", get(f"{c}.norm.weight")
+                if role.ratio > 1:
+                    yield f"{m}.compressor.wkv", get(f"{c}.wkv.weight")
+                    yield f"{m}.compressor.wgate", get(f"{c}.wgate.weight")
+                else:
+                    yield from linear(f"{c}.wkv", f"{m}.compressor.wkv")
+                yield from linear(f"{a}.indexer.wk", f"{m}.indexer.wk")
+                yield f"{m}.indexer.k_norm.weight", get(f"{a}.indexer.k_norm.weight")
+            if role.has_indexer:
+                yield from linear(f"{a}.indexer.wq_b", f"{m}.indexer.wq_b")
+                yield from linear(f"{a}.indexer.weights_proj", f"{m}.indexer.weights_proj")
+
+            p = f"layers.{L}"
+            yield f"model.{p}.attn_norm.weight", get(f"{p}.attn_norm.weight")
+            yield f"model.{p}.ffn_norm.weight", get(f"{p}.ffn_norm.weight")
+            yield f"model.{p}.ffn.gate.weight", get(f"{p}.ffn.gate.weight")
+            yield f"model.{p}.ffn.gate.bias", get(f"{p}.ffn.gate.bias")
+            for proj in ("w1", "w2", "w3"):
+                yield from linear(f"{p}.ffn.shared_experts.{proj}", f"model.{p}.ffn.shared_experts.{proj}")
+            for nm in ("hc_attn_fn", "hc_ffn_fn", "hc_attn_base", "hc_ffn_base", "hc_attn_scale", "hc_ffn_scale"):
+                yield f"model.{p}.{nm}", get(f"{p}.{nm}")
+            if L in args.engram_layer_ids:
+                e = f"{p}.engram"
+                yield from linear(f"{e}.wkv", f"model.{e}.wkv")
+                yield f"model.{e}.q_weight", get(f"{e}.q_weight")
+                yield f"model.{e}.k_weight", get(f"{e}.k_weight")
+    finally:
+        reader.close()
+
+
+def iter_vision_weights(model_path: str, device):
+    names = weight_map(model_path)
+    reader = _ShardReader(model_path, names, device)
+    try:
+        for name in names:
+            if name.startswith(("vision.", "aligner.")) or name in ("image_start", "image_end", "image_newline"):
+                yield f"visual.{name}", reader.get(name)
+        args = load_args(model_path)
+        yield "visual.routing_bias", torch.stack([reader.get(f"layers.{i}.ffn.gate.bias_vl") for i in range(args.n_layers)])
+    finally:
+        reader.close()
+
+
+# ----- FTW side files: the Engram tables ----------------------------------------------------------
+def engram_table_names(layer_id: int) -> tuple[str, str]:
+    return f"layers.{layer_id}.engram.embed.weight", f"layers.{layer_id}.engram.embed.scale"
+
+
+def ftw_side_files(model_path: str, out_dir: str) -> list[str]:
+    """Copy the checkpoint shards holding the Engram tables (fp8 rows + ue8m0 scales) next to an FTW
+    checkpoint. The tables are never FTW entries: ``load_host_tables`` maps them in place from whatever
+    safetensors file holds them (``engram_table.engram_row_source`` reads the headers), and the FTW
+    loader never reads ``.safetensors`` as weights, so the shard's other tensors ride along unused."""
+    import shutil
+
+    from freetoken.models.loader import safetensors_weight_map
+    from freetoken.utils import download_hf_weight
+
+    folder = download_hf_weight(model_path)
+    args = load_args(folder)
+    shard_of = safetensors_weight_map(folder)
+    shards = sorted({shard_of[n] for layer_id in args.engram_layer_ids if layer_id < args.n_layers for n in engram_table_names(layer_id)})
+    for shard in shards:
+        shutil.copyfile(os.path.join(folder, shard), os.path.join(out_dir, shard))
+    return shards
+
+
+# ----- routed MXFP4 expert pieces ---------------------------------------------------------------
+_EXPERT_RE = re.compile(r"^layers\.(?P<layer>\d+)\.ffn\.experts\.(?P<expert>\d+)\.(?P<proj>w1|w2|w3)\.(?P<kind>weight|scale)$")
+_PROJ_ROLE = {"w1": "gate", "w3": "up", "w2": "down"}
+_KIND_SUFFIX = {"weight": "", "scale": "_scale"}
+
+
+def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
+    """One piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0 ``_scale`` companions.
+    The DSpark stage's experts (``mtp.*``) never match."""
+    if kind is not QuantKind.MXFP4:
+        return None
+    if get_tp_info().size > 1:
+        raise NotImplementedError("DeepSeek-V4.1 expert banks support TP=1 only")
+    from freetoken.models.weight import iter_expert_tensors_parallel
+    from freetoken.moe.expert_pieces import per_expert_pieces
+
+    args = load_args(model_path)
+    L, E = args.n_layers, args.n_routed_experts
+
+    def locate(raw_name: str):
+        m = _EXPERT_RE.match(raw_name)
+        if m is None or int(m["layer"]) >= L:
+            return None
+        return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]]
+
+    if parallel:
+        tensors = iter_expert_tensors_parallel(model_path, lambda n: locate(n) is not None, workers=workers, chunk=chunk)
+        return per_expert_pieces(tensors, locate, tensors_per_expert=6)
+
+    def _serial():
+        reader = _ShardReader(model_path, weight_map(model_path), torch.device("cpu"))
+        try:
+            for li in tqdm(range(L), desc="Loading DeepSeek-V4.1 experts (serial)", disable=not get_tp_info().is_primary()):
+                for e in range(E):
+                    base = f"layers.{li}.ffn.experts.{e}"
+                    for proj in ("w1", "w3", "w2"):
+                        for kind_ in ("weight", "scale"):
+                            name = f"{base}.{proj}.{kind_}"
+                            yield name, reader.get(name)
+        finally:
+            reader.close()
+
+    return per_expert_pieces(_serial(), locate, tensors_per_expert=6)
+
+
+__all__ = ["iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files", "load_args", "weight_map", "dequant_fp8_block"]
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 2a4ced10af..fe150ddc29 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -156,6 +156,20 @@ class ModelSpec:
         # the head, the KV compressors and the indexer's scorer ship bf16; the fp8 config has no modules_to_not_convert
         unquantized_modules=("head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj"),
     ),
+    # DeepSeek-V4.1-Flash (model_type deepseek_v41): multimodal wrapper config (text tower in
+    # text_config). The checkpoint has no ``model.`` root; the compressors, the indexer's key projection /
+    # scorer, the head and the vision tower + aligner ship bf16 (the fp8 config has no modules_to_not_convert).
+    "DeepseekV41ForCausalLM": ModelSpec(
+        "freetoken.models.deepseek_v41",
+        "DeepseekV41ForCausalLM",
+        mm_processor="freetoken.mm.processors.deepseek_v41:DeepseekV41MMProcessor",
+        encoders=(EncoderSpec("vision", "vision_config", ("image",)),),
+        checkpoint_roots=(("model.layers", "layers"), ("model.embed", "embed"), ("model.norm", "norm")),
+        packed_modules_mapping=_EXPERTS_W123_PACKED,
+        unquantized_modules=(
+            "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.wk", "*.indexer.weights_proj", "vision.*", "aligner.*",
+        ),
+    ),
     "Qwen3_5MoeForConditionalGeneration": ModelSpec(
         "freetoken.models.qwen3_5_moe",
         "Qwen3_5MoeForConditionalGeneration",
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index 83b3d19665..e9d6ffe54e 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -108,6 +108,7 @@ class Workload:
     activation: str = "silu"
     swiglu_alpha: float = 1.702
     swiglu_limit: float | None = None
+    act_block: int | None = None  # ds_fp4 W4A8 activation quant block
 
 
 # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py.
@@ -120,7 +121,8 @@ class Workload:
                              activation="gpt_oss_swiglu", swiglu_limit=7.0),
     "gpt-oss-20b": Workload("gpt-oss-20b", 2880, 2880, 32, 4, ("mxfp4_triton",),
                             activation="gpt_oss_swiglu", swiglu_limit=7.0),
-    "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0),
+    "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128),
+    "dsv4.1-flash": Workload("dsv4.1-flash", 5120, 2304, 384, 6, ("ds_fp4",), swiglu_limit=10.0, act_block=32),
     "glm4.7-nvfp4": Workload("glm4.7-nvfp4", 5120, 1536, 160, 8, ("nvfp4",)),
     "glm5.3-flash-nvfp4": Workload(
         "glm5.3-flash-nvfp4", 4096, 2048, 288, 8, ("nvfp4",),
@@ -142,7 +144,7 @@ class Workload:
     "fp8_block": Workload("dtype:fp8", 2048, 512, 128, 8, ("fp8_block",)),
     "mxfp4_triton": Workload("dtype:mxfp4", 2880, 2880, 128, 4, ("mxfp4_triton",),
                              activation="gpt_oss_swiglu", swiglu_limit=7.0),
-    "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0),
+    "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128),
 }
 
 
@@ -460,7 +462,7 @@ def _build_cpu_moe_executor(fmt: str, wl: Workload, banks: dict, num_threads: in
         cache, top_k=wl.top_k, activation=wl.activation,
         apply_router_weight_on_input=False, num_threads=num_threads, max_tokens=1,
         device=torch.device("cuda"), swiglu_alpha=wl.swiglu_alpha,
-        swiglu_limit=wl.swiglu_limit,
+        swiglu_limit=wl.swiglu_limit, act_block=wl.act_block,
     )
 
 
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index a46806740d..c4809075c8 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -171,6 +171,7 @@ def __init__(
         swiglu_alpha: float = 1.702,
         swiglu_limit: float | None = None,
         fmt: str | None = None,
+        act_block: int | None = None,
     ) -> None:
         from freetoken.kernel import _cpu_moe
         from freetoken.moe.legacy_format import canonical_role
@@ -323,6 +324,11 @@ def __init__(
         # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode:
         # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py).
         self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda"
+        # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1)
+        self._act_block = act_block
+        if fmt == "ds_fp4":
+            assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block"
+            self._ext.set_act_block(act_block)
         if self._gpu_prequant:
             self._ext.set_input_prequant(True)
             logger.info_rank0(
@@ -482,7 +488,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block
         scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row),
         so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the
-        activations (block 128) to match DSV4's W4A8 reference, hence the %128 dims."""
+        activations (the checkpoint's fp8 block) to match the DeepSeek W4A8 reference."""
         gup, gus = banks["gate_up"], banks["gate_up_scale"]
         dnp, dns = banks["down"], banks["down_scale"]
         assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype)
@@ -490,7 +496,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         I = int(gup[0].shape[1] // 2)
         H = int(gup[0].shape[2] * 2)
         assert gup[0].shape[1] == 2 * I
-        assert H % 128 == 0 and I % 128 == 0, (H, I)  # FP8 activation round-trip block=128
+        assert H % 32 == 0 and I % 32 == 0, (H, I)  # e8m0 scales per 32 along K
         assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I)
         assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H)
         assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I)
@@ -580,7 +586,7 @@ def decode_submit(
             # pre-quantized activations and skips its serial scalar pass.
             from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip
 
-            hidden_states = act_quant_fp8_roundtrip(hidden_states, block=128)
+            hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block)
 
         # D2H: ship this step's activations + routing to pinned host memory.
         io["x"].copy_(hidden_states, non_blocking=True)
diff --git a/python/freetoken/moe/fused_ds_fp4.py b/python/freetoken/moe/fused_ds_fp4.py
index 65fffb5cc9..b4dc5cd0d0 100644
--- a/python/freetoken/moe/fused_ds_fp4.py
+++ b/python/freetoken/moe/fused_ds_fp4.py
@@ -12,10 +12,7 @@
 import torch
 import triton
 
-from freetoken.kernel.triton.dsv4.fp8_linear import (
-    act_quant_fp8_inplace,
-    act_quant_fp8_roundtrip,
-)
+from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip
 from freetoken.kernel.triton.dsv4.fused_moe import (
     _decode_dsfp4_moe_kernel,
     _e2m1_lut,
@@ -50,7 +47,8 @@ def _grouped_decode(
     """Grouped per-route GEMM ``c[t,r,:] = a[row] @ dequant(W[slot])^T``.
 
     Returns ``[T, top_k, N]``. Used for both gate_up (a_row=token) and down
-    (a_row=route). FP4 weights are dequantized inline from the slot cache.
+    (a_row=route). FP4 weights are dequantized inline from the slot cache. A slot
+    of ``-1`` is an inactive route: its output is zero and its storage is never read.
     """
     T, top_k = slots.shape
     N = packed_cache.shape[1]
@@ -62,14 +60,13 @@ def _grouped_decode(
         topk_weights = out.new_empty((1, 1), dtype=torch.float32)
     scale_u8 = scale_cache.view(torch.uint8)
 
-    # Decode is a per-route GEMV, bound by the inline FP4 dequant (LUT gather + block
-    # scale), not weight HBM bandwidth -- so it tops out ~29% of HBM peak. BN=16/
-    # BKB=128/1 warp maximizes memory-level parallelism (many single-warp CTAs, no
-    # cross-warp reduction) -> ~950 GB/s on H100 (was ~830 at BN=8/BKB=256/2). K_BYTES
-    # must be a multiple of BLOCK_SIZE_KB (gate_up 2048, down 1024 -- both /128).
+    # One fixed launch geometry, never tuned per machine: the row / warp split changes the per-row
+    # K reduction layout (and so the output bits) on some Triton versions. 16 rows on 2 warps is at
+    # or near the fastest split for the V4 and V4.1 decode shapes. Each K tile contains whole
+    # packed FP4 scale groups.
     BLOCK_SIZE_N = 16
     BLOCK_SIZE_KB = 128
-    _NW = 1
+    _NW = 2
     assert (K // 2) % BLOCK_SIZE_KB == 0, (K, BLOCK_SIZE_KB)
     grid = (total_routes, triton.cdiv(N, BLOCK_SIZE_N))
     _decode_dsfp4_moe_kernel[grid](
@@ -103,28 +100,28 @@ def routed_experts_fp4(
     down_packed: torch.Tensor,     # [S, H, I//2] uint8
     down_scale: torch.Tensor,      # [S, H, I//32] e8m0
     swiglu_limit: float,
+    act_block: int = 128,
 ) -> torch.Tensor:
     """Full routed-expert output (summed over the top-k routes), excludes shared expert.
 
-    Precision matches the reference ``fp4_gemm(act_quant(x, 128), W_fp4)``: the gate_up
-    and down activations are FP8-round-tripped (block 128, ue8m0) before each GEMM. Since
-    an fp8 value x pow2 scale is exact in bf16, the round-tripped activation entering the
-    bf16 decode kernel is bit-identical to the reference's dequantized FP8 activation
-    (validated max diff = 0 vs the tilelang ``fp4_gemm`` reference)."""
+    Precision matches the reference ``fp4_gemm(act_quant(x, act_block), W_fp4)``: the gate_up
+    and down activations are FP8-round-tripped (block ``act_block`` -- the checkpoint's fp8 block,
+    128 on V4 and 32 on V4.1 -- ue8m0) before each GEMM. Since an fp8 value x pow2 scale is exact
+    in bf16, the round-tripped activation entering the bf16 decode kernel is bit-identical to the
+    reference's dequantized FP8 activation (validated max diff = 0 vs the tilelang ``fp4_gemm``
+    reference). The routing weight scales each route's down output in the GEMM epilogue."""
     T, top_k = slots.shape
     H = x.shape[1]
     two_I = gate_up_packed.shape[1]
     I = two_I // 2
 
-    x = act_quant_fp8_roundtrip(x, 128)  # gate_up activation -> FP8 round-trip (no clone)
+    x = act_quant_fp8_roundtrip(x, act_block)  # gate_up activation -> FP8 round-trip (no clone)
     gate_up = _grouped_decode(
         x, gate_up_packed, gate_up_scale, slots, None,
         a_row_is_route=False, mul_routed_weight=False,
     )  # [T, top_k, 2I]
-    act = fused_swiglu(gate_up, swiglu_limit)  # [T, top_k, I]
-
-    act = act.reshape(T * top_k, I)
-    act_quant_fp8_inplace(act, 128)  # down activation -> FP8 round-trip
+    # [T, top_k, I] with the down activation's FP8 round-trip in the same pass
+    act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(T * top_k, I)
     down = _grouped_decode(
         act, down_packed, down_scale, slots, topk_weights,
         a_row_is_route=True, mul_routed_weight=True,
@@ -191,6 +188,7 @@ def routed_experts_fp4_prefill(
     down_scale: torch.Tensor,      # [S, H, I//32] e8m0
     swiglu_limit: float,
     num_rows: int,
+    act_block: int = 128,
 ) -> torch.Tensor:
     """Grouped-GEMM counterpart of :func:`routed_experts_fp4` for dense prefill
     chunks: one moe_align sort shared by both GEMMs, each expert's weights
@@ -201,7 +199,7 @@ def routed_experts_fp4_prefill(
     if T * top_k < _GROUPED_MIN_ROUTES:
         return routed_experts_fp4(
             x, slots, topk_weights,
-            gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit,
+            gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit, act_block=act_block,
         )
     H = x.shape[1]
     two_I = gate_up_packed.shape[1]
@@ -217,16 +215,14 @@ def routed_experts_fp4_prefill(
     sorted_ids, expert_ids, ntpp = moe_align_block_size(slots, cfg["BLOCK_SIZE_M"], num_rows)
     tw = topk_weights.reshape(-1).contiguous()
 
-    x = act_quant_fp8_roundtrip(x, 128)  # gate_up activation -> FP8 round-trip (no clone)
+    x = act_quant_fp8_roundtrip(x, act_block)  # gate_up activation -> FP8 round-trip (no clone)
     gate_up = torch.empty((T, top_k, two_I), dtype=x.dtype, device=x.device)
     _grouped_prefill(
         x, gate_up_packed, gate_up_scale, gate_up, tw,
         sorted_ids, expert_ids, ntpp, routes, top_k, False, cfg,
     )
-    act = fused_swiglu(gate_up, swiglu_limit)  # [T, top_k, I]
-
-    act = act.reshape(routes, I)
-    act_quant_fp8_inplace(act, 128)  # down activation -> FP8 round-trip
+    # [T, top_k, I] with the down activation's FP8 round-trip in the same pass
+    act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(routes, I)
     down = torch.empty((T, top_k, H), dtype=x.dtype, device=x.device)
     _grouped_prefill(
         act, down_packed, down_scale, down, tw,
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 8f77be9330..772297e069 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -66,6 +66,11 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ
     supports_runtime_rebuild = True
     prefill_chunk_budget = None  # generic shared page pool: no per-model prefill chunk cap
 
+    @property
+    def prefix_replay_tokens(self) -> int:
+        """Prompt-tail tokens a prefix hit must leave to the prefill (a pool whose model recomputes them)."""
+        return getattr(self.swa_pool, "prefix_replay_tokens", 0) if self.swa_pool is not None else 0
+
     @property
     def prefill_chunk_align(self) -> int:
         """Granularity a non-final prefill chunk should end on. A hybrid snapshot is donated only
@@ -90,10 +95,10 @@ def _make_prefix_cache(self, device, page_size, type):
             return SWARadixCache(device, page_size, self.sliding_window_size)
         return create_prefix_cache(device=device, type=type, page_size=page_size)
 
-    def match_req(self, req: PendingReq) -> MatchResult:
+    def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult:
         input_len = req.input_len
         assert input_len > 0, "Input length must be greater than 0."
-        ids = req.input_ids[: input_len - 1]
+        ids = req.input_ids[: input_len - 1 if max_len is None else max(0, min(input_len - 1, max_len))]
         if self.is_swa:
             from freetoken.kvcache.swa_radix_cache import SWACacheHandle
             m = self.prefix_cache.match_prefix(ids)
diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index bd8129aae8..de4bddfb4a 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -71,8 +71,9 @@ def _try_allocate_one(self, req: PendingReq):
         if self.table_manager.available_size == 0:
             return None
 
-        # TODO: consider host cache match case
-        mr = self.cache_manager.match_req(req)
+        # TODO: consider host cache match case (it needs the same replay cap)
+        replay = self.cache_manager.prefix_replay_tokens
+        mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None)
         handle = mr.cuda_handle
         cached_len = handle.cached_len
         # TODO: better estimate policy
@@ -214,6 +215,7 @@ def _add_one_req(
             uid=pending_req.uid,
             cache_handle=cache_handle,
             sampling_params=pending_req.sampling_params,
+            prompt_len=pending_req.input_len,
         )
         req.mm_items = pending_req.mm_items
         req.mrope_positions_full = pending_req.mrope_positions_full
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index fa0d96a714..d057d3b908 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -349,7 +349,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                 next_token = int(next_token.item())
                 # EOS / stop-string -> "stop", output budget exhausted -> "length";
                 # EOS and stop strings win over length.
-                # Overlap can advance device_len ahead of the token delivered to the host.
+                # Overlap can advance device_len one token ahead of this host reply.
                 hit_length = req.input_ids.numel() >= req.max_device_len
                 hit_eos = (
                     not req.sampling_params.ignore_eos and next_token in self.eos_token_ids
@@ -737,17 +737,10 @@ def _current_cache_geometry(self) -> dict:
         the CONCRETE current window (usable pages) so a rollback restores it byte-for-byte,
         whether it was pinned or ratio-derived."""
         eng = self.engine
-        config = self.config
-        mc = config.model_config
-        num_swa_pages = None
-        if getattr(mc, "dsv4_args", None) is not None:
-            sizes = getattr(eng.kv_cache, "sizes", None)
-            if sizes is not None:  # usable window pages = physical n_win_pages minus the dummy page
-                num_swa_pages = max(0, sizes.n_win_pages - 1)
-        elif getattr(mc, "has_swa_attention", False) and (
-            getattr(config, "cache_type", None) == "swa_radix"
-        ):  # usable window tokens = pool tokens minus the slot-0 sentinel
-            num_swa_pages = max(0, int(getattr(eng.kv_cache, "swa_num_tokens", 0) or 0) - 1)
+        from freetoken.kvcache.cache_status import window_pool_spec
+
+        spec = window_pool_spec(self.config)
+        num_swa_pages = eng.kv_cache.window_pages if spec is not None else None
         return dict(
             num_pages=eng.num_pages,
             moe_cache_size=eng.moe_offload_cache.cache_size if eng.moe_offload_cache is not None else None,
diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 5592a2b0d5..34221cfc70 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -563,8 +563,12 @@ def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) ->
         or pools.get("num_pages", 0) or 0
     )
     page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1)
-    is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None
-    swa_page_size = page_size if is_dsv4 else 1
+    from freetoken.kvcache.cache_status import window_pool_spec
+
+    spec = window_pool_spec(config)
+    if spec is None:
+        raise ValueError("swa_full_tokens_ratio requested but this model has no resizable window pool")
+    swa_page_size = spec.page_size
     window_tokens = int(round(req.swa_full_tokens_ratio * num_pages * page_size))
     return max(1, -(-window_tokens // swa_page_size))  # ceil-div to the pool's page unit
 
@@ -604,12 +608,16 @@ async def cache_rebuild(req: CacheRebuildRequest):
             {"status": "failed", "error": "swa_full_tokens_ratio must be in (0, 1]"},
             status_code=422,
         )
+    try:
+        num_swa_pages = _resolve_num_swa_pages(state, req)
+    except ValueError as exc:
+        return JSONResponse({"status": "failed", "error": str(exc)}, status_code=422)
     result = await dispatch_rebuild(
         state,
         moe_cache_size=req.moe_cache_size,
         num_pages=req.num_pages,
         num_mamba_slots=req.num_mamba_slots,
-        num_swa_pages=_resolve_num_swa_pages(state, req),
+        num_swa_pages=num_swa_pages,
         mode=req.mode,
         timeout=req.timeout,
     )
@@ -739,11 +747,11 @@ def cache_geometry(state: Any) -> dict:
     # the internal currency), else the load-time meta value. 0.0 for models without a window pool.
     swa_full_tokens_ratio = float(getattr(state, "swa_full_tokens_ratio", 0.0) or 0.0)
     last_swa_pages = last.get("num_swa_pages")
-    if last_swa_pages:
-        is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None
-        full = num_pages if is_dsv4 else num_pages * page_size
+    swa_page_size = int(pools.get("swa_page_size", 0) or 0)
+    if last_swa_pages and swa_page_size:
+        full = num_pages * page_size
         if full > 0:
-            swa_full_tokens_ratio = min(1.0, last_swa_pages / full)
+            swa_full_tokens_ratio = min(1.0, last_swa_pages * swa_page_size / full)
     try:
         model_config = config.model_config
         num_experts = int(getattr(model_config, "num_experts", 0) or 0)
@@ -785,9 +793,8 @@ def cache_geometry(state: Any) -> dict:
         "moe_cache_policy": getattr(config, "moe_cache_policy", None),
         "unit_bytes": unit_bytes,
         "swa_full_tokens_ratio": swa_full_tokens_ratio,
-        # The window pool's page unit for num_swa_pages: DSV4 = P (== page_size), radix-SWA = 1
-        # token, 0 for models without a window pool. Lets a client denominate the swa control.
-        "swa_page_size": int(pools.get("swa_page_size", 0) or 0),
+        # The pool's window page unit; 0 for models without a resizable window pool.
+        "swa_page_size": swa_page_size,
         # Concrete current window size in that unit (usable pages): the last rebuild's value if it
         # pinned/derived one, else the load-time pool size. 0 for models without a window pool.
         "num_swa_pages": int(last.get("num_swa_pages") or pools.get("num_swa_pages", 0) or 0),
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 2194b82e92..b8d06018d8 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -204,6 +204,8 @@ def _infer_tool_call_parser(model_path: str) -> str:
             return "qwen3_coder"
         if "qwen" in marker:
             return "qwen25"
+        if "deepseek_v41" in marker or "deepseekv41" in marker:
+            return "deepseekv41"
         if "deepseek" in marker and ("v4" in marker or "deepseek_v4" in marker):
             return "deepseekv32"
         if "deepseek" in marker and ("v3.2" in marker or "v32" in marker):
@@ -576,6 +578,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
             "qwen3_coder",
             "mistral",
             "deepseekv32",
+            "deepseekv41",
             "gemma4",
             "glm47",
             "minimax",
@@ -645,6 +648,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--swa-decoder-replay",
+        default=ServerArgs.swa_decoder_replay,
+        choices=["bounded", "exact"],
+        help=(
+            "DeepSeek-V4.1 Decoder SWA Bounded Replay. 'bounded' (default) runs the 20 decoder layers "
+            "on each prompt's last 128 tokens with their sliding window truncated there, as in the "
+            "tech report; 'exact' runs them on every prompt token (the reference numerics)."
+        ),
+    )
+
     parser.add_argument(
         "--nvfp4-backend",
         action=_DeprecatedAlias,
diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py
index 0804dcc3f7..296cd560dc 100644
--- a/python/freetoken/server/function_call_parser.py
+++ b/python/freetoken/server/function_call_parser.py
@@ -78,6 +78,8 @@ def _should_forward_unknown_tool(name: Any) -> bool:
     "<｜DSML｜function_calls>",
     "<｜DSML｜tool_calls>",
     "<｜DSML｜invoke",
+    "<｜DSML｜ calls>",
+    "<｜DSML｜ invoke",
     "<atem:function_calls>",
 ]
 
@@ -1563,37 +1565,45 @@ class DeepSeekV32Detector(BaseFormatDetector):
     Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3.2
     """
 
+    # the DSML tag names: (block, alternative block), invoke, parameter
+    block_tags = ("function_calls", "tool_calls")
+    invoke_tag = "invoke"
+    param_tag = "parameter"
+
     def __init__(self):
         super().__init__()
         self.dsml_token = "｜DSML｜"
-        self.bot_token = f"<{self.dsml_token}function_calls>"
-        self.eot_token = f"</{self.dsml_token}function_calls>"
-        self.alt_bot_token = f"<{self.dsml_token}tool_calls>"
-        self.alt_eot_token = f"</{self.dsml_token}tool_calls>"
-        self.invoke_start_prefix = f"<{self.dsml_token}invoke"
-        self.invoke_end_token = f"</{self.dsml_token}invoke>"
-        self.param_end_token = f"</{self.dsml_token}parameter>"
+        block, alt_block = self.block_tags
+        self.bot_token = f"<{self.dsml_token}{block}>"
+        self.eot_token = f"</{self.dsml_token}{block}>"
+        self.alt_bot_token = f"<{self.dsml_token}{alt_block}>"
+        self.alt_eot_token = f"</{self.dsml_token}{alt_block}>"
+        self.invoke_start_prefix = f"<{self.dsml_token}{self.invoke_tag}"
+        self.invoke_end_token = f"</{self.dsml_token}{self.invoke_tag}>"
+        self.param_start_prefix = f"<{self.dsml_token}{self.param_tag}"
+        self.param_end_token = f"</{self.dsml_token}{self.param_tag}>"
 
         # Regex for complete invoke extraction
-        _de = re.escape(self.dsml_token)
+        _inv = re.escape(self.dsml_token + self.invoke_tag)
+        _par = re.escape(self.dsml_token + self.param_tag)
         self.invoke_regex = re.compile(
-            rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*?)</{_de}invoke>',
+            rf'<{_inv}\s+name="([^"]+)"\s*>(.*?)</{_inv}>',
             re.DOTALL,
         )
         # Regex for parameter extraction
         self.param_regex = re.compile(
-            rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)</{_de}parameter>',
+            rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)</{_par}>',
             re.DOTALL,
         )
         # Regex for partial invoke (name known, body still streaming)
         self.partial_invoke_regex = re.compile(
-            rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*)',
+            rf'<{_inv}\s+name="([^"]+)"\s*>(.*)',
             re.DOTALL,
         )
         # Streaming state machine tag regexes (anchored matches over the buffer).
-        self.invoke_open_regex = re.compile(rf'<{_de}invoke\s+name="([^"]+)"\s*>')
+        self.invoke_open_regex = re.compile(rf'<{_inv}\s+name="([^"]+)"\s*>')
         self.param_open_regex = re.compile(
-            rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>'
+            rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>'
         )
 
         self._last_arguments = ""
@@ -1687,7 +1697,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
 
         normal_parts: List[str] = []
         calls: List[ToolCallItem] = []
-        param_open_token = f"<{self.dsml_token}parameter"
+        param_open_token = self.param_start_prefix
 
         def _emit_args(fragment: str) -> None:
             if fragment:
@@ -3512,6 +3522,16 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
         return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls)
 
 
+class DeepSeekV41Detector(DeepSeekV32Detector):
+    """DeepSeek-V4.1's DSML: the V3.2 grammar with every tag name led by a space and the block
+    named ``calls`` (``<｜DSML｜ calls>`` / ``<｜DSML｜ invoke name=...>`` / ``<｜DSML｜ parameter
+    ...>``, the checkpoint's ``encoding/encoding.py``)."""
+
+    block_tags = (" calls", " calls")
+    invoke_tag = " invoke"
+    param_tag = " parameter"
+
+
 class FunctionCallParser:
     """
     Parser for function/tool calls in model outputs.
@@ -3523,6 +3543,7 @@ class FunctionCallParser:
 
     ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = {
         "deepseekv32": DeepSeekV32Detector,
+        "deepseekv41": DeepSeekV41Detector,
         "gemma4": Gemma4Detector,
         "gpt-oss": GptOssDetector,
         "gpt_oss": GptOssDetector,
diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py
index d01828e21b..3d48ee96a0 100644
--- a/python/freetoken/server/stats.py
+++ b/python/freetoken/server/stats.py
@@ -120,12 +120,10 @@ def derive_model_card(config: Any) -> dict:
 
 
 def _swa_page_size(config: Any) -> int:
-    """The window pool's own page unit: P (window_size) for DSV4, 1 token for radix-SWA.
-    Mirrors compute_cache_pools' swa_page_size."""
-    dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None)
-    if dsv4 is not None:
-        return int(getattr(dsv4, "window_size", 0) or 1)
-    return 1
+    from freetoken.kvcache.cache_status import window_pool_spec
+
+    spec = window_pool_spec(config)
+    return spec.page_size if spec is not None else 1
 
 
 def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict:
diff --git a/python/freetoken/tokenizer/effort.py b/python/freetoken/tokenizer/effort.py
index c81cd7bef1..2e9e3cd21f 100644
--- a/python/freetoken/tokenizer/effort.py
+++ b/python/freetoken/tokenizer/effort.py
@@ -84,7 +84,8 @@ class EffortProfile:
     template interpolates anything". ``strength_dialect`` True means the
     template reads the ``reasoning_strength`` spelling (muse-glimmer's family,
     whose documented ladder tops out at xhigh) rather than only
-    ``reasoning_effort``.
+    ``reasoning_effort``. ``integer_range`` records a validated 1-100 budget
+    accepted independently of the named effort vocabulary.
     """
 
     supported: frozenset[str]
@@ -92,6 +93,7 @@ class EffortProfile:
     consumes_effort: bool
     validates: bool = False
     strength_dialect: bool = False
+    integer_range: tuple[int, int] | None = None
 
 
 @dataclass(frozen=True)
@@ -118,7 +120,7 @@ class ThinkingProfile:
 _MAX_QUANTIZE_DISTANCE = 0.15
 
 
-def quantize_effort(value: Any, profile: EffortProfile) -> str | None:
+def quantize_effort(value: Any, profile: EffortProfile) -> str | int | None:
     """Map a client's effort onto ``profile``; ``None`` means "send nothing".
 
     In-vocabulary values pass through untouched. Other named levels land on the
@@ -128,6 +130,11 @@ def quantize_effort(value: Any, profile: EffortProfile) -> str | None:
     the template default. With "max" excluded the remaining scale values are
     unique, so quantization is deterministic across processes.
     """
+    if profile.integer_range is not None and isinstance(value, (int, float)):
+        low, high = profile.integer_range
+        if type(value) is not int or not low <= value <= high:
+            raise ValueError(f"reasoning_effort must be an integer in [{low}, {high}]")
+        return value
     supported = effective_efforts(profile)
     if not supported:
         return None
@@ -181,6 +188,7 @@ def probe_effort_profile(
     diverged: set[str] = set()
     matches_baseline: dict[str, bool] = {name: True for name in KNOWN_REASONING_EFFORTS}
     ran_rounds = 0
+    integer_range = None
 
     for base_kwargs, tools in _PROBE_ROUNDS:
         try:
@@ -198,6 +206,10 @@ def probe_effort_profile(
             if _renderings_differ(rendering, baseline):
                 diverged.add(name)
                 matches_baseline[name] = False
+        # Only a validated budget range is a numeric capability; arbitrary Jinja
+        # interpolation must not advertise support for numeric reasoning effort.
+        if _accepts_integer_budget(render, base_kwargs, tools):
+            integer_range = (1, 100)
 
     if ran_rounds == 0:
         # Nothing learnable: sending no effort is the only safe rendering.
@@ -214,7 +226,7 @@ def probe_effort_profile(
         strength_dialect = False
 
     supported = frozenset(name for name in KNOWN_REASONING_EFFORTS if name not in rejected)
-    consumes = bool(rejected or diverged)
+    consumes = bool(rejected or diverged or integer_range)
     default = None
     if consumes:
         defaults = [name for name in supported if matches_baseline[name]]
@@ -226,9 +238,26 @@ def probe_effort_profile(
         consumes_effort=consumes,
         validates=bool(rejected),
         strength_dialect=strength_dialect,
+        integer_range=integer_range,
     )
 
 
+def _accepts_integer_budget(render, base_kwargs, tools) -> bool:
+    for value in (0, 101):
+        try:
+            render({**base_kwargs, "reasoning_effort": value}, tools)
+        except Exception:  # noqa: BLE001 -- the encoder owns validation
+            continue
+        return False
+    outputs = []
+    for value in range(1, 101):
+        try:
+            outputs.append(render({**base_kwargs, "reasoning_effort": value}, tools))
+        except Exception:  # noqa: BLE001
+            return False
+    return _renderings_differ(outputs[0], outputs[-1])
+
+
 def probe_thinking_profile(
     render: Callable[[dict[str, Any], list[dict[str, Any]] | None], Any],
     efforts: EffortProfile,
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index f473665267..772472573a 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -59,7 +59,7 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor
         self._effort_profile: EffortProfile | None = None
         self._thinking_profile: ThinkingProfile | None = None
         self._effort_lock = threading.Lock()
-        self._logged_effort_maps: set[tuple[Any, str | None]] = set()
+        self._logged_effort_maps: set[tuple[Any, str | int | None]] = set()
         self._inline_system_lock = threading.Lock()
         self._inline_system_modes: dict[str, str] = {}
 
@@ -117,7 +117,7 @@ def render_prompt(self, msg: TokenizeMsg) -> str:
                 if mode == "auto":
                     mode = self.inline_system_mode(msg.tools, kwargs)
                 messages = normalize_inline_system(messages, mode)
-        return self._render(messages, msg.tools, kwargs)
+        return self._render(messages, msg.tools, kwargs, images=msg.images)
 
     def inline_system_mode(self, tools: list[dict] | None, kwargs: dict[str, Any]) -> str:
         key = json.dumps([tools, kwargs], sort_keys=True)
@@ -137,12 +137,13 @@ def _render(
         messages: list[dict[str, Any]],
         tools: list[dict[str, Any]] | None,
         chat_template_kwargs: dict[str, Any],
+        *, images: list[bytes] | None = None,
     ) -> str:
         """Raw render, no effort sanitation — the probe needs unsupported values
         to actually reach the template so rejection is observable."""
         if self._dsv4_encoder is not None:
             return _apply_dsv4_chat_encoder(
-                self._dsv4_encoder, messages, tools, chat_template_kwargs
+                self._dsv4_encoder, messages, tools, chat_template_kwargs, images=images
             )
         # Broadcast the effort in every spelling the ecosystem's templates read
         # (muse-glimmer grades ``reasoning_strength``; Jinja ignores undeclared
@@ -222,8 +223,12 @@ def _load_dsv4_encoder_if_needed(tokenizer: PreTrainedTokenizerBase) -> ModuleTy
     model_path = getattr(tokenizer, "name_or_path", None) or getattr(tokenizer, "_name_or_path", "")
     if not model_path:
         return None
-    encoder_path = os.path.join(str(model_path), "encoding", "encoding_dsv4.py")
-    if not os.path.isfile(encoder_path):
+    # DeepSeek-V4 ships encoding/encoding_dsv4.py, V4.1 encoding/encoding.py; same encode_messages contract
+    encoder_path = next(
+        (p for p in (os.path.join(str(model_path), "encoding", name) for name in ("encoding_dsv4.py", "encoding.py")) if os.path.isfile(p)),
+        None,
+    )
+    if encoder_path is None:
         return None
     spec = importlib.util.spec_from_file_location("encoding_dsv4", encoder_path)
     if spec is None or spec.loader is None:
@@ -240,9 +245,28 @@ def _apply_dsv4_chat_encoder(
     messages: list[dict],
     tools: list[dict] | None,
     chat_template_kwargs: dict,
+    *, images: list[bytes] | None = None,
 ) -> str:
     rendered_messages = [dict(message) for message in messages]
+    image_data = iter(images) if images is not None else None
     for message in rendered_messages:
+        if isinstance(message.get("content"), list):
+            content = []
+            for part in message["content"]:
+                if isinstance(part, dict) and part.get("type") == "image":
+                    # Prerender still has refs; the tokenizer worker receives fetched bytes.
+                    # The checkpoint encoder validates sources even when rendering only text.
+                    ref = part.get("freetoken_ref")
+                    if image_data is not None:
+                        try:
+                            part = {"type": "image", "data": next(image_data)}
+                        except StopIteration:
+                            raise ValueError("image blocks outnumber fetched images") from None
+                    elif ref is not None:
+                        key = "data" if ref["kind"] == "b64" else "url"
+                        part = {"type": "image", key: ref["data"]}
+                content.append(part)
+            message["content"] = content
         if message.get("tool_calls"):
             message["tool_calls"] = _dsv4_tool_calls(message["tool_calls"])
     if tools:
diff --git a/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py
index f9cb62d1a8..9640c7b7f7 100644
--- a/scripts/ftw_hotfix.py
+++ b/scripts/ftw_hotfix.py
@@ -263,7 +263,9 @@ def is_checkpoint_tower_name(name: str) -> bool:
     head = name.split(".")
     if head[0] == "model" and len(head) > 1:
         head = head[1:]
-    return "vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS
+    return ("vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS
+            or head[0] in ("aligner", "image_start", "image_end", "image_newline")
+            or name.endswith(".ffn.gate.bias_vl"))
 
 
 def read_tower(source: TensorSource, ftw_dir: str, checkpoint_names: list[str]) -> dict[str, torch.Tensor]:
diff --git a/tests/README.md b/tests/README.md
index be22c1fb20..76b2b1fb03 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -62,6 +62,7 @@ checkpoint is set:
 | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) |
 | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) |
 | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file |
+| `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) |
 
 `test_aime.py` takes its sampling protocol from the checkpoint's own
 `generation_config.json` (pass@N at the recommended temperature, or a single greedy
diff --git a/tests/attention/test_dsv41_backend.py b/tests/attention/test_dsv41_backend.py
new file mode 100644
index 0000000000..387742eb63
--- /dev/null
+++ b/tests/attention/test_dsv41_backend.py
@@ -0,0 +1,208 @@
+"""DSV41 attention backend: addressing and selection contracts, CPU-only (kernels live in tests/kernels).
+
+* prefill metadata: encoder segments, the decoder pass under exact vs bounded replay;
+* window candidates with a floor (bounded replay truncates the window at the prompt's last window);
+* decode snapshot staging and the layer-invariant ring context;
+* the selection helpers' reshape, candidate gather and position-to-row mapping.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx
+from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes
+from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache
+from freetoken.models.deepseek_v41.args import DeepseekV41Args
+
+P, MRR, DEVICE = 128, 4, torch.device("cpu")
+RATIOS = (0, 2, 2, 1, 1)
+SOURCES = (1, 3)
+
+
+def _ctx(pool):
+    try:
+        ctx = get_global_ctx()
+    except AssertionError:
+        ctx = Context(page_size=P)
+        set_global_ctx(ctx)
+    ctx.kv_cache = pool
+    return ctx
+
+
+def _stack(swa_decoder_replay="exact", num_pages=32, max_seq_len=8192):
+    from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend
+
+    args = DeepseekV41Args(n_layers=len(RATIOS), head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS,
+                           kv_source_layers=SOURCES, swa_decoder_replay=swa_decoder_replay)
+    pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, DEVICE, n_scratch=MRR + 1)
+    pool._init_paged_state(MRR, True)
+    pt = torch.zeros(MRR + 1, max_seq_len, dtype=torch.int32)
+    pt[MRR].fill_(num_pages * P)
+    pt[2, :300] = torch.arange(300, dtype=torch.int32)
+    for page in range(3):
+        pool.bind_window_pages(page * P, page * P)
+    pool.full_loc_map = pt
+    _ctx(pool)
+    backend = DSV41SparseAttnBackend(SimpleNamespace(dsv41_args=args))
+    return backend, pool, pt
+
+
+def _req(table_idx, cached_len, n_new, output_len=16, uid=0, prompt_len=0):
+    return Req(input_ids=torch.zeros(cached_len + n_new, dtype=torch.int32), table_idx=table_idx, cached_len=cached_len,
+               output_len=output_len, uid=uid, sampling_params=SamplingParams(), cache_handle=None, prompt_len=prompt_len)
+
+
+def _prefill_batch(reqs):
+    batch = Batch(reqs=reqs, phase="prefill")
+    batch.padded_reqs = reqs
+    return batch
+
+
+def test_prefill_metadata_exact_mode_runs_the_decoder_on_every_token():
+    backend, pool, _ = _stack("exact")
+    assert pool.prefix_replay_tokens == 0
+    batch = _prefill_batch([_req(0, 0, 300), _req(2, 256, 44, uid=1)])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    assert [(s.offset, s.n, s.table_idx, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0, 0), (300, 44, 2, 256, 0)]
+    assert md.decoder_segments is md.segments and md.decoder_rows is None
+    assert md.last_indices.tolist() == [299, 343]
+
+
+def test_prefill_metadata_bounded_mode_runs_the_decoder_on_the_last_window():
+    """The encoder runs every new token; the decoder runs each request's part of its prompt's last
+    window ``[L - P, L)``, floored at ``L - P``. A chunk ending before that window runs no decoder row
+    and its ``last_indices`` entry points at the placeholder row the model appends."""
+    backend, pool, _ = _stack("bounded")
+    assert pool.prefix_replay_tokens == P and pool.sliding_window_size == P
+    cold = _req(0, 0, 300)  # a whole prompt: L - P = 172
+    early = _req(1, 256, 128, uid=1, prompt_len=600)  # chunk [256, 384) of a 600-token prompt: before 472
+    straddle = _req(2, 384, 128, uid=2, prompt_len=600)  # chunk [384, 512): its rows [472, 512) are in the window
+    batch = _prefill_batch([cold, early, straddle])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    assert [(s.offset, s.n, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0), (300, 128, 256, 0), (428, 128, 384, 0)]
+    assert [(d.offset, d.n, d.table_idx, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 128, 0, 172, 172), (128, 40, 2, 472, 472)]
+    assert md.decoder_rows.tolist() == list(range(172, 300)) + list(range(428 + 88, 556))
+    assert md.decoder_pad and md.last_indices.tolist() == [127, 168, 167]  # 168: the placeholder after 168 decoder rows
+
+
+def test_prefill_metadata_bounded_mode_without_any_decoder_row():
+    backend, _, _ = _stack("bounded")
+    batch = _prefill_batch([_req(0, 0, 256, prompt_len=1024)])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    assert md.decoder_segments == [] and md.decoder_rows.numel() == 0
+    assert md.decoder_pad and md.last_indices.tolist() == [0]
+
+
+def test_prefill_metadata_short_prompt_runs_the_decoder_everywhere():
+    """A prompt shorter than the window floors the decoder's window at 0: it sees what exact mode sees."""
+    backend, _, _ = _stack("bounded")
+    batch = _prefill_batch([_req(0, 0, 50)])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    assert [(d.offset, d.n, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 50, 0, 0)]
+    assert not md.decoder_pad and md.last_indices.tolist() == [49]
+
+
+def test_private_window_layers_address_per_request_rings():
+    """Under bounded replay the decoder layers keep their window KV in per-request rings (slot =
+    table row * P + pos % P), never in the radix-shared window pages; the encoder layers stay on
+    the shared pool."""
+    from freetoken.attention.dsv41_sparse import PrefillSegment
+
+    backend, pool, _ = _stack("bounded")
+    assert pool.private_window_layer_ids == (3, 4) and not any(pool.is_private_window(l) for l in (0, 1, 2))
+    assert pool.window_pool[3].shape[0] == (MRR + 1) * P and pool.window_pool[0].shape[0] == pool.sizes.n_win_slots
+    # prefill: a decoder segment's slots are its ring row; an encoder segment's are page-bound
+    seg = PrefillSegment(0, 44, 2, 256, window_floor=256)
+    assert backend.layer_window_slots_of(3, 2, 256, 300).tolist() == [2 * P + p % P for p in range(256, 300)]
+    assert backend.layer_window_slots_of(0, 2, 256, 300).tolist() == list(range(256, 300))
+    g = backend.window_topk_prefill(seg, layer_id=3)[0]
+    assert g.shape == (44, P) and g[0, 0].item() == 2 * P + 0 and (g[0, 1:] == -1).all()  # position 256 sees itself only (floor)
+    assert set(g[43].tolist()) - {-1} == {2 * P + p % P for p in range(256, 300)}
+    # decode: the private ring context keys on the batch row's page-table row
+    batch = _decode_batch([2], [300])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    slots, topk = md.private_window_ctx(torch.tensor([300]), torch.arange(1))
+    assert slots.tolist() == [2 * P + 300 % P] and topk.shape == (1, 1, P)
+    assert set(topk[0, 0].tolist()) == {2 * P + j for j in range(P)}  # every ring index holds a position <= 300
+
+
+def test_window_candidates_honor_the_floor_and_the_retained_prefix():
+    from freetoken.attention.dsv41_sparse import PrefillSegment
+
+    backend, pool, _ = _stack()
+    # a cold segment at positions [0, 10): query p sees [0, p]
+    g = backend.window_topk_prefill(PrefillSegment(0, 10, 2, 0))
+    assert g.shape == (1, 10, P)
+    assert g[0, 3, :4].tolist() == [0, 1, 2, 3] and (g[0, 3, 4:] == -1).all()
+    # an extend at [256, 300) sees back into the retained prefix, but not below the floor
+    g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256, window_floor=200))
+    row = g[0, 0]  # query at 256 -> [200, 256]
+    assert row[0].item() == 200 and row[56].item() == 256 and (row[57:] == -1).all()
+    # without a floor the query at 256 sees [129, 256]
+    g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256))
+    assert g[0, 0, 0].item() == 129 and g[0, 0, 127].item() == 256
+
+
+def _decode_batch(rows, positions):
+    reqs = [_req(int(t), 0, 1, uid=i) for i, t in enumerate(rows)]
+    batch = Batch(reqs=reqs, phase="decode")
+    batch.padded_reqs = reqs
+    batch.active_table_idx = torch.tensor(rows, dtype=torch.int64)
+    batch.positions = torch.tensor(positions, dtype=torch.int64)
+    return batch
+
+
+def test_decode_snapshot_and_ring_context():
+    backend, pool, pt = _stack()
+    batch = _decode_batch([2, MRR], [259, 0])
+    backend.prepare_metadata(batch)
+    md = batch.attn_metadata
+    assert md.full_snap is None
+    snap = md.full_snapshot()
+    assert torch.equal(snap[0, :300], torch.arange(300))
+    pt[2, :10] = 7  # a later mutation must not reach the snapshot
+    assert torch.equal(snap[0, :300], torch.arange(300))
+    pos, rows = batch.positions, torch.arange(2)
+    ws, prev, topk = md.window_ctx(pos, rows)
+    dummy_ws = pool.sizes.n_win_slots - P  # the dummy row's full loc binds to the dummy window page
+    assert ws.tolist() == [259, dummy_ws] and prev.tolist() == [258, dummy_ws]
+    assert topk.shape == (2, 1, P) and (topk[0] >= 0).all()
+    assert set(topk[0, 0].tolist()) == set(range(259 - 127, 260))
+    assert topk[1, 0, 0].item() == dummy_ws and (topk[1, 0, 1:] == -1).all()
+    # the indexer derives compressed rows from the live locs: prefill off the page table row,
+    # decode off the snapshot (position t -> full_loc(t * ratio) // ratio)
+    assert torch.equal(backend.locs_prefill(2, 20), pt[2:3, :20])
+    assert snap[0, :16].tolist() == list(range(16))  # taken before the mutation above
+    # capture staging reuses one buffer
+    backend.init_capture_graph(512, [1, 2])
+    backend.prepare_for_replay(batch)
+    assert batch.attn_metadata.full_snap.data_ptr() == backend.capture.full_snap.data_ptr()
+    assert batch.attn_metadata.stage_width == 512
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="the selection kernels need CUDA")
+def test_topk_helpers():
+    from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend as B
+
+    scores = torch.tensor([[[0.1, 5.0, 3.0, 7.0, 9.0]]], device="cuda")  # columns past live are never read
+    live = torch.tensor([[3]], device="cuda", dtype=torch.int32)
+    assert B.select_topk(scores, live, 4).tolist() == [[[0, 1, 2, -1]]]
+    assert B.select_topk(scores, live, 2).tolist() == [[[1, 2]]]
+    # a candidate list is a sorted valid prefix; empty slots score -inf
+    cand = torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32)
+    cscores = torch.tensor([[[1.0, 4.0, -torch.inf, -torch.inf]]], device="cuda")
+    assert B.select_topk_in_candidates(cscores, cand, 3).tolist() == [[[2, 7, -1]]]
+    # rows come from the live locs: compressed position p -> locs[p * ratio] // ratio
+    locs = torch.arange(20, 40, device="cuda", dtype=torch.int32).view(1, 20)
+    locs[0, 14:16] = -1  # position 7 (ratio 2) has slid out
+    rows = B.positions_to_rows(torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32), locs, 2)
+    assert rows.tolist() == [[[12, -1, 19, -1]]]
diff --git a/tests/checkpoint/test_ftw_hotfix_vision.py b/tests/checkpoint/test_ftw_hotfix_vision.py
index d4a3eac68e..eb0c592f35 100644
--- a/tests/checkpoint/test_ftw_hotfix_vision.py
+++ b/tests/checkpoint/test_ftw_hotfix_vision.py
@@ -14,7 +14,7 @@
 
 # env var -> local checkpoint, as tests/README.md lists them
 _CHECKPOINTS = ("FREETOKEN_QWEN3VL_MODEL", "FREETOKEN_GEMMA4_MODEL", "FREETOKEN_GEMMA4_UNIFIED_MODEL",
-                "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL")
+                "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL", "FREETOKEN_DSV41_MODEL")
 
 
 @pytest.fixture(scope="module")
@@ -34,6 +34,10 @@ def hotfix():
     ("vision_tower.vision_model.embeddings.patch_embedding.weight", True),
     ("multi_modal_projector.linear_1.weight", True),
     ("patch_merge_mlp.linear_2.bias", True),
+    ("vision.blocks.0.norm1.weight", True),
+    ("aligner.w1.weight", True),
+    ("image_newline", True),
+    ("layers.3.ffn.gate.bias_vl", True),
     ("model.embed_audio.embedding_projection.weight", False),
     ("model.language_model.layers.0.self_attn.q_proj.weight", False),
     ("model.layers.3.mlp.experts.0.gate_proj.weight", False),
diff --git a/tests/dsv4/test_dsv4_attn_backend.py b/tests/dsv4/test_dsv4_attn_backend.py
index 8ffc6fd4d9..7bc72210f2 100644
--- a/tests/dsv4/test_dsv4_attn_backend.py
+++ b/tests/dsv4/test_dsv4_attn_backend.py
@@ -11,8 +11,8 @@
 import torch
 
 from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx
-from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
-from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes
+from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
+from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes
 from freetoken.models.deepseek_v4.args import DeepseekV4Args
 
 P, MRR, DEVICE = 128, 4, torch.device("cpu")
diff --git a/tests/kernels/test_dsv41_indexer.py b/tests/kernels/test_dsv41_indexer.py
new file mode 100644
index 0000000000..19b45ab8b9
--- /dev/null
+++ b/tests/kernels/test_dsv41_indexer.py
@@ -0,0 +1,139 @@
+"""DSV41 indexer logits over packed fp4 index keys: full-range (Full mode) and candidate-restricted
+(Reindex mode) against a torch reference on the dequantized keys, with causality from ``live``. The
+row of a compressed position comes from the request's full-token locs (``locs[t * ratio] // ratio``);
+in Full mode only the live columns are written."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kvcache.dsv4.v41_row_format import FP4_E8M0_B32
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+D, H = 128, 32
+ROWS = 4200
+
+
+RATIO = 2
+
+
+def _setup(b, s, t, seed):
+    from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows
+
+    torch.manual_seed(seed)
+    keys = torch.randn(ROWS, D, device="cuda", dtype=torch.bfloat16)
+    pool = pack_rows(keys, FP4_E8M0_B32)
+    deq = unpack_rows(pool, FP4_E8M0_B32, D)
+    q = torch.randn(b, s, H, D, device="cuda", dtype=torch.bfloat16)
+    w = (torch.randn(b, s, H, device="cuda") * 0.3).bfloat16()
+    # each request owns a random permutation of pool rows for its compressed positions; a few holes.
+    # The kernel derives rows from full-token locs: compressed position p -> loc = row * RATIO (+ any
+    # in-group offset) at locs[p * RATIO]; a hole is a negative loc.
+    k_rows = torch.stack([torch.randperm(ROWS, device="cuda")[:t] for _ in range(b)]).to(torch.int32)
+    k_rows[:, 5::11] = -1
+    locs = torch.full((b, t * RATIO), -1, device="cuda", dtype=torch.int32)
+    for i in range(b):
+        for p in range(t):
+            r = int(k_rows[i, p])
+            if r >= 0:
+                locs[i, p * RATIO : (p + 1) * RATIO] = torch.arange(r * RATIO, (r + 1) * RATIO, device="cuda", dtype=torch.int32)
+    return pool, deq, q, w, k_rows, locs
+
+
+def _reference(q, w, deq, k_rows, live, positions):
+    """positions: [B, S, N] compressed positions (-1 empty) -> logits [B, S, N]."""
+    b, s, n = positions.shape
+    out = torch.full((b, s, n), float("-inf"), device="cuda")
+    for i in range(b):
+        for j in range(s):
+            for c in range(n):
+                p = int(positions[i, j, c])
+                if p < 0 or p >= int(live[i, j]):
+                    continue
+                r = int(k_rows[i, p])
+                if r < 0:
+                    continue
+                score = (q[i, j] @ deq[r]).relu() * w[i, j]
+                out[i, j, c] = score.sum()
+    return out
+
+
+def test_full_range_matches_reference_with_causal_live():
+    from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed
+
+    b, s, t = 2, 5, 300
+    pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=1)
+    live = torch.tensor([[10, 64, 65, 200, 300], [0, 1, 128, 129, 250]], device="cuda", dtype=torch.int32)
+    out = torch.full((b, s, t), 12345.0, device="cuda")  # a sentinel: the dead area must stay untouched
+    got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, out=out)
+    positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t)
+    want = _reference(q, w, deq, k_rows, live, positions)
+    assert got.shape == (b, s, t)
+    written = positions < live.unsqueeze(-1)
+    touched = positions < (torch.div(live + 63, 64, rounding_mode="floor") * 64).unsqueeze(-1)  # whole live tiles
+    assert (got[~touched] == 12345.0).all() and torch.isneginf(got[touched & ~written]).all()
+    assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written]))
+    finite = written & ~torch.isinf(want)
+    torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2)
+
+
+def test_quantized_queries_preserve_reference_score_rounding():
+    from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed
+    from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows
+
+    torch.manual_seed(23)
+    t, s, topk = 4096, 32, 512
+    pool = pack_rows(torch.randn(t, D, device="cuda", dtype=torch.bfloat16), FP4_E8M0_B32)
+    keys = unpack_rows(pool, FP4_E8M0_B32, D)
+    q = torch.randn(s * H, D, device="cuda", dtype=torch.bfloat16)
+    q = unpack_rows(pack_rows(q, FP4_E8M0_B32), FP4_E8M0_B32, D).view(1, s, H, D)
+    w = torch.randn(1, s, H, device="cuda", dtype=torch.bfloat16) * (D * H) ** -0.5
+    locs = torch.arange(t, device="cuda", dtype=torch.int32).view(1, t)
+    live = torch.full((1, s), t, device="cuda", dtype=torch.int32)
+    got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, 1, live)
+    reference = torch.einsum("bshd,btd->bsht", q, keys.unsqueeze(0))
+    reference = (reference.relu() * w.unsqueeze(-1)).sum(dim=2).float()
+    torch.testing.assert_close(got, reference, atol=0, rtol=0)
+    # BF16 scores can tie at the cutoff; compare selected scores, not arbitrary tie order.
+    torch.testing.assert_close(got.topk(topk).values, reference.topk(topk).values, atol=0, rtol=0)
+
+
+def test_full_range_worker_grid_covers_wide_histories(monkeypatch):
+    """More live tiles than worker programs: each worker strides over its share."""
+    from freetoken.kernel.triton.dsv41 import indexer as mod
+
+    b, s, t = 1, 2, 4096  # 64 tiles
+    pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=4)
+    live = torch.tensor([[4096, 4000]], device="cuda", dtype=torch.int32)
+    from types import SimpleNamespace
+
+    properties = torch.cuda.get_device_properties(q.device)
+    with monkeypatch.context() as patch:
+        patch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace(
+            multi_processor_count=5, major=properties.major, minor=properties.minor,
+        ))
+        got = mod.indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live)
+    positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t)
+    want = _reference(q, w, deq, k_rows, live, positions)
+    written = positions < live.unsqueeze(-1)
+    assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written]))
+    finite = written & ~torch.isinf(want)
+    torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2)
+
+
+def test_candidate_pool_matches_reference():
+    from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed
+
+    b, s, t, nc = 2, 3, 400, 96
+    pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=2)
+    live = torch.tensor([[400, 350, 33], [1, 400, 200]], device="cuda", dtype=torch.int32)
+    cand = torch.randint(0, t, (b, s, nc), device="cuda", dtype=torch.int32)
+    cand[..., -10:] = -1  # empty tail of the pool
+    got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, candidates=cand)
+    want = _reference(q, w, deq, k_rows, live, cand)
+    assert got.shape == (b, s, nc)
+    assert torch.equal(torch.isinf(got), torch.isinf(want))
+    finite = ~torch.isinf(want)
+    torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2)
diff --git a/tests/kernels/test_dsv41_pack.py b/tests/kernels/test_dsv41_pack.py
new file mode 100644
index 0000000000..176889dc7f
--- /dev/null
+++ b/tests/kernels/test_dsv41_pack.py
@@ -0,0 +1,112 @@
+"""Packed KV row formats (``kvcache/dsv4/v41_row_format.py`` + ``kernel/triton/dsv41/pack.py``) against a torch
+transcription of the reference quantizers in DeepSeek-V4.1's ``inference/kernel.py``: the round-tripped
+value the pool hands back must equal what the reference bakes into its bf16 cache, bit for bit."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32, RowFormat
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
+
+
+def _round_fp4(x: torch.Tensor) -> torch.Tensor:
+    """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast)."""
+    a = x.abs()
+    grid = FP4_GRID.to(x.device)
+    # nearest grid point; ties go to the even-mantissa neighbour (0, 1, 2, 4 are even)
+    dist = (a[..., None] - grid).abs()
+    lo = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7)
+    below, above = grid[lo - 1], grid[lo]
+    mid = (below + above) / 2
+    even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device))
+    pick_above = (a > mid) | ((a == mid) & even_above)
+    r = torch.where(pick_above, above, below)
+    r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r)
+    del dist
+    return torch.copysign(r, x)
+
+
+def reference_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor:
+    """``act_quant`` / ``fp4_act_quant`` with ``inplace=True``: quant + dequant back to bf16."""
+    m, d = x.shape
+    if fmt is BF16:
+        return x.to(torch.bfloat16)
+    g = x.float().view(m, d // fmt.group, fmt.group)
+    amax = g.abs().amax(dim=-1)
+    if fmt is FP8_E8M0_B32:
+        s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(1e-4) / 448.0)))
+        q = (g / s[..., None]).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float()
+    elif fmt is FP4_E8M0_B32:
+        s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(6.0 * 2.0**-126) / 6.0)))
+        q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0))
+    elif fmt is FP4_E4M3_B16:
+        s = (amax.clamp_min(6.0 * 2.0**-9) / 6.0).to(torch.float8_e4m3fn).float()
+        q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0))
+    else:
+        raise AssertionError(fmt)
+    return (q * s[..., None]).view(m, d).to(torch.bfloat16)
+
+
+@pytest.mark.parametrize("fmt", [BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32], ids=lambda f: f.name)
+@pytest.mark.parametrize("dim", [128, 512])
+def test_roundtrip_matches_the_reference_quantizer(fmt: RowFormat, dim: int):
+    from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows
+
+    torch.manual_seed(0)
+    m = 300
+    # mixed magnitudes, exact zeros, a tiny group and a huge one
+    x = torch.randn(m, dim, device="cuda") * torch.exp2(torch.randint(-8, 6, (m, 1), device="cuda").float())
+    x[7] = 0
+    x[11, : fmt.group or 32] *= 1e-6
+    x[13, -32:] = 400.0
+    x = x.to(torch.bfloat16)
+
+    packed = pack_rows(x, fmt)
+    assert packed.shape == (m, fmt.row_bytes(dim)) and packed.dtype == torch.uint8
+    got = unpack_rows(packed, fmt, dim)
+    want = reference_roundtrip(x, fmt)
+    assert torch.equal(got, want), f"{fmt.name}: {(got.float() - want.float()).abs().max().item()}"
+
+
+def test_scatter_into_pool_rows_and_gather_back():
+    from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows
+
+    torch.manual_seed(1)
+    fmt, dim, rows = FP4_E4M3_B16, 512, 64
+    pool = torch.zeros((rows, fmt.row_bytes(dim)), dtype=torch.uint8, device="cuda")
+    x = torch.randn(5, dim, device="cuda", dtype=torch.bfloat16)
+    ids = torch.tensor([3, 60, 0, 17, 42], device="cuda", dtype=torch.int64)
+    assert pack_rows(x, fmt, pool=pool, row_ids=ids) is pool
+    got = unpack_rows(pool, fmt, dim, row_ids=ids.to(torch.int32))
+    assert torch.equal(got, reference_roundtrip(x, fmt))
+    untouched = torch.ones(rows, dtype=torch.bool, device="cuda")
+    untouched[ids] = False
+    assert int(pool[untouched].sum()) == 0
+
+
+def test_e4m3_encode_helper_matches_torch_bits():
+    """``e4m3_f32_to_u8`` (the pre-sm_89 scale encoder) on every finite e4m3 code."""
+    import triton
+    import triton.language as tl
+
+    from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8
+
+    @triton.jit
+    def _enc(x_ptr, o_ptr, N: tl.constexpr):
+        offs = tl.arange(0, N)
+        tl.store(o_ptr + offs, e4m3_f32_to_u8(tl.load(x_ptr + offs)))
+
+    codes = torch.arange(256, dtype=torch.uint8, device="cuda")
+    finite = (codes & 0x7F) != 0x7F  # skip the NaN codes
+    vals = codes.view(torch.float8_e4m3fn).float()
+    vals = torch.where(finite, vals, torch.zeros_like(vals))
+    out = torch.empty(256, dtype=torch.uint8, device="cuda")
+    _enc[(1,)](vals, out, N=256)
+    want = torch.where(finite, codes, torch.zeros_like(codes))
+    # -0.0 encodes to 0x80 in torch and 0x80 here; +0.0 -> 0x00
+    assert torch.equal(out, want)
diff --git a/tests/kernels/test_dsv41_sparse_attn.py b/tests/kernels/test_dsv41_sparse_attn.py
new file mode 100644
index 0000000000..8a4f5f2a4f
--- /dev/null
+++ b/tests/kernels/test_dsv41_sparse_attn.py
@@ -0,0 +1,107 @@
+"""DSV41 sparse attention over packed pools against a torch reference on the DEQUANTIZED pools.
+
+``m == 1`` takes the split-k decode path, ``m > 1`` the single-program prefill one. Tolerance follows
+the dsv4 sparse-attention tests (5e-2 on fp32-cast bf16): the online-softmax accumulation order differs
+from the reference and split-k merges through log-sum-exp on top of that."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP8_E8M0_B32
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+D, H = 128, 8
+N_WINDOW = 128  # a multiple of BLOCK_T
+N_WIN_SLOTS, N_CMP = 512, 1024
+TOL = dict(atol=5e-2, rtol=5e-2)
+
+
+def _reference(q, win, cmp, sink, idx, n_window, scale, counts):
+    """Per (request, query, head) softmax over the live columns; ``sink`` is a null key."""
+    b, m, h, d = q.shape
+    out = torch.zeros_like(q, dtype=torch.float32)
+    for i in range(b):
+        for j in range(m):
+            width = idx.shape[-1] if counts is None else n_window + int(counts[i, j])
+            kv = []
+            for c in range(min(width, idx.shape[-1])):
+                slot = int(idx[i, j, c])
+                if slot >= 0:
+                    kv.append((win if c < n_window else cmp)[slot].float())
+            if not kv:
+                continue
+            kv_t = torch.stack(kv)
+            logits = q[i, j].float() @ kv_t.T * scale
+            mx = torch.maximum(logits.max(dim=-1).values, sink.float())
+            probs = torch.exp(logits - mx[:, None])
+            denom = probs.sum(dim=-1) + torch.exp(sink.float() - mx)
+            out[i, j] = (probs @ kv_t) / denom[:, None]
+    return out
+
+
+def _pools(win_fmt, cmp_fmt, seed):
+    from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows
+
+    torch.manual_seed(seed)
+    win_vals = torch.randn(N_WIN_SLOTS, D, device="cuda", dtype=torch.bfloat16)
+    cmp_vals = torch.randn(N_CMP, D, device="cuda", dtype=torch.bfloat16)
+    win_pool = pack_rows(win_vals, win_fmt)
+    cmp_pool = pack_rows(cmp_vals, cmp_fmt)
+    return win_pool, cmp_pool, unpack_rows(win_pool, win_fmt, D), unpack_rows(cmp_pool, cmp_fmt, D)
+
+
+def _topk(b, m, n_cmp_cols, seed, ring_fill=None):
+    torch.manual_seed(seed)
+    win = torch.randint(0, N_WIN_SLOTS, (b, m, N_WINDOW), device="cuda", dtype=torch.int32)
+    if ring_fill is not None:  # a ring still filling: trailing window columns empty
+        win[..., ring_fill:] = -1
+    cmp = torch.randint(0, N_CMP, (b, m, n_cmp_cols), device="cuda", dtype=torch.int32)
+    cmp[..., ::7] = -1  # scattered unreachable picks
+    return torch.cat([win, cmp], dim=-1)
+
+
+@pytest.mark.parametrize("m", [1, 3])
+@pytest.mark.parametrize("fmts", [(FP8_E8M0_B32, FP4_E4M3_B16), (BF16, BF16)], ids=["fp8+fp4", "bf16"])
+def test_matches_reference_on_dequantized_pools(m, fmts):
+    from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed
+
+    win_fmt, cmp_fmt = fmts
+    b, n_cmp_cols = 2, 256
+    win_pool, cmp_pool, win_deq, cmp_deq = _pools(win_fmt, cmp_fmt, seed=1)
+    q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16)
+    sink = torch.randn(H, device="cuda") * 0.5
+    idx = _topk(b, m, n_cmp_cols, seed=2, ring_fill=100)
+    scale = D**-0.5
+
+    got = sparse_attn_packed(q, win_pool, win_fmt, cmp_pool, cmp_fmt, sink, idx, N_WINDOW, scale)
+    want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, None)
+    torch.testing.assert_close(got.float(), want, **TOL)
+
+
+def test_cmp_counts_bound_the_visited_columns():
+    from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed
+
+    b, m, n_cmp_cols = 3, 1, 512
+    win_pool, cmp_pool, win_deq, cmp_deq = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=3)
+    q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16)
+    sink = torch.zeros(H, device="cuda")
+    idx = _topk(b, m, n_cmp_cols, seed=4)
+    counts = torch.tensor([[0], [37], [512]], device="cuda", dtype=torch.int32)
+    scale = D**-0.5
+
+    got = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, sink, idx, N_WINDOW, scale, cmp_counts=counts)
+    want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, counts)
+    torch.testing.assert_close(got.float(), want, **TOL)
+
+
+def test_all_masked_query_yields_zeros():
+    from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed
+
+    win_pool, cmp_pool, _, _ = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=5)
+    q = torch.randn(1, 2, H, D, device="cuda", dtype=torch.bfloat16)
+    idx = torch.full((1, 2, N_WINDOW + 64), -1, device="cuda", dtype=torch.int32)
+    out = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, torch.zeros(H, device="cuda"), idx, N_WINDOW, 0.1)
+    assert torch.equal(out, torch.zeros_like(out))
diff --git a/tests/kernels/test_dsv41_topk.py b/tests/kernels/test_dsv41_topk.py
new file mode 100644
index 0000000000..e7fc66f8d8
--- /dev/null
+++ b/tests/kernels/test_dsv41_topk.py
@@ -0,0 +1,159 @@
+"""dsv41_topk / dsv41_candidate_blocks against an independent torch reference with the documented tie
+rule, over the contract's edge cases: empty and short histories, unequal live counts in one batch,
+-inf holes, widths crossing the slice and level boundaries, exact ties, newest-block inclusion,
+the sorted-valid-prefix / -1-tail output layout, and CUDA-graph replay across changing live lengths
+(stale workspace reads would surface as wrong picks)."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def ref_topk(scores: torch.Tensor, live: torch.Tensor, k: int) -> torch.Tensor:
+    """(score desc, column asc) over the live, finite columns; ascending columns, -1 padded."""
+    R, T = scores.shape
+    out = torch.full((R, k), -1, dtype=torch.int32)
+    for r in range(R):
+        n = int(live[r])
+        s = scores[r, :n].double().cpu()
+        cols = torch.arange(n)
+        alive = torch.isfinite(s) | (s == float("inf"))
+        cand = [(float(-s[c]), int(c)) for c in cols[alive]]
+        cand.sort()
+        picked = sorted(c for _, c in cand[:k])
+        out[r, : len(picked)] = torch.tensor(picked, dtype=torch.int32)
+    return out.to(scores.device)
+
+
+def ref_candidates(scores, live, topk_blocks, block_size):
+    R, T = scores.shape
+    nb = -(-T // block_size)
+    out = torch.full((R, topk_blocks * block_size), -1, dtype=torch.int32)
+    for r in range(R):
+        n = int(live[r])
+        if n == 0:
+            continue
+        nbl = -(-n // block_size)
+        s = torch.nn.functional.pad(scores[r, :n].float().cpu(), (0, nbl * block_size - n), value=float("-inf")).view(nbl, block_size).amax(-1)
+        s[nbl - 1] = float("inf")
+        cand = sorted((float(-s[b]), b) for b in range(nbl) if s[b] > float("-inf"))
+        keep = sorted(b for _, b in cand[:topk_blocks])
+        pos = [p for b in keep for p in range(b * block_size, (b + 1) * block_size) if p < n]
+        out[r, : len(pos)] = torch.tensor(pos, dtype=torch.int32)
+    return out.to(scores.device)
+
+
+def _scores(R, T, seed=0, ties=False):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    s = torch.randn(R, T, generator=g)
+    if ties:
+        s = torch.randint(0, 5, (R, T), generator=g).float()  # heavy exact ties
+    return s.cuda()
+
+
+@pytest.mark.parametrize("T,k", [(64, 8), (4096, 512), (4097, 512), (8192, 512), (100_000, 512), (262_144, 512), (1_048_576, 512), (131_072, 2048)])
+def test_topk_matches_reference_across_widths(T, k):
+    from freetoken.kernel.triton.dsv41.topk import dsv41_topk
+
+    R = 3
+    scores = _scores(R, T, seed=T)
+    live = torch.tensor([0, min(T, 5000), T], dtype=torch.int32, device="cuda")[:R]
+    # holes inside the live range
+    scores[2, ::7] = float("-inf")
+    got = dsv41_topk(scores, live, k)
+    want = ref_topk(scores, live, k)
+    assert torch.equal(got, want), (got[:, :8], want[:, :8])
+
+
+def test_topk_short_and_unequal_rows_and_exact_ties():
+    from freetoken.kernel.triton.dsv41.topk import dsv41_topk
+
+    T, k = 20_000, 512
+    scores = _scores(5, T, seed=3, ties=True)
+    live = torch.tensor([0, 1, 300, 4096, 20_000], dtype=torch.int32, device="cuda")
+    got = dsv41_topk(scores, live, k)
+    want = ref_topk(scores, live, k)
+    assert torch.equal(got, want)
+    # the layout contract: a sorted valid prefix, then -1
+    for r in range(5):
+        row = got[r].tolist()
+        n = sum(1 for v in row if v >= 0)
+        assert row[n:] == [-1] * (k - n) and row[:n] == sorted(row[:n]) and n == min(int(live[r]), k)
+
+
+def test_topk_all_dead_row_and_fewer_finite_than_k():
+    from freetoken.kernel.triton.dsv41.topk import dsv41_topk
+
+    scores = torch.full((2, 10_000), float("-inf"), device="cuda")
+    scores[1, [10, 4095, 4096, 9000]] = torch.tensor([1.0, 2.0, 3.0, 0.5], device="cuda")
+    live = torch.tensor([10_000, 9001], dtype=torch.int32, device="cuda")
+    got = dsv41_topk(scores, live, 8)
+    assert (got[0] == -1).all()
+    assert got[1].tolist() == [10, 4095, 4096, 9000, -1, -1, -1, -1]
+
+
+@pytest.mark.parametrize("T", [4000, 200_000])
+def test_candidate_blocks_match_reference_and_pin_the_newest_block(T):
+    from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks
+
+    R, KB, BS = 4, 2048, 8
+    scores = _scores(R, T, seed=T + 1)
+    scores[0, :] = -1000.0  # the newest block is kept even when every score is poor
+    live = torch.tensor([13, 8, min(T, 16_000), T], dtype=torch.int32, device="cuda")
+    got = dsv41_candidate_blocks(scores, live, KB, BS)
+    want = ref_candidates(scores, live, KB, BS)
+    assert got.shape == (R, KB * BS)
+    assert torch.equal(got, want)
+    # short histories: every live position is a candidate (the implicit fast path)
+    assert got[0].tolist()[:13] == list(range(13)) and (got[0, 13:] == -1).all()
+    assert got[1].tolist()[:8] == list(range(8)) and (got[1, 8:] == -1).all()
+
+
+def test_topk_under_graph_replay_across_live_lengths():
+    """Capture once at the staged width, replay with live lengths that move up and down: every replay
+    must equal the reference (a stale candidate slot or live count from a previous replay would not)."""
+    from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks, dsv41_topk
+
+    T, k, R = 300_000, 512, 2
+    scores = torch.empty((R, T), device="cuda")
+    live = torch.empty((R,), dtype=torch.int32, device="cuda")
+    out = torch.empty((R, k), dtype=torch.int32, device="cuda")
+    cand = torch.empty((R, 2048 * 8), dtype=torch.int32, device="cuda")
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        scores.copy_(_scores(R, T, seed=9))
+        live.fill_(T)
+        dsv41_topk(scores, live, k, out)  # warm-up compiles
+        cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8))
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            dsv41_topk(scores, live, k, out)
+            cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8))
+    torch.cuda.synchronize()
+    for step, (l0, l1) in enumerate([(T, 7), (4096, 4097), (0, 1), (250_000, 12_000), (3, T), (16_384, 16_385)]):
+        scores.copy_(_scores(R, T, seed=100 + step))
+        live.copy_(torch.tensor([l0, l1], dtype=torch.int32))
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(out, ref_topk(scores, live, k)), (step, l0, l1)
+        assert torch.equal(cand, ref_candidates(scores, live, 2048, 8)), (step, l0, l1)
+
+
+def test_topk_plan_terminates_or_rejects():
+    """Every accepted k converges (each level at least halves the row); k past CHUNK // 2 -- which
+    would stall the planner (k == CHUNK never shrinks, 3072 stalls at 6144) -- is rejected up front."""
+    from freetoken.kernel.triton.dsv41.topk import CHUNK, FINAL_MAX, MAX_K, topk_plan
+
+    for k in (1, 8, 512, 2048, MAX_K):
+        for columns in (1, 4096, 4097, 1 << 20, (1 << 20) + 3):
+            levels = topk_plan(columns, k)
+            widths = [w for w, _ in levels] + [(levels[-1][1] * k) if levels else columns]
+            assert all(a > b for a, b in zip(widths, widths[1:])) and widths[-1] <= FINAL_MAX
+            assert len(levels) <= 12
+    for k in (MAX_K + 1, 2049, 3072, CHUNK, 0, -1):
+        with pytest.raises(ValueError):
+            topk_plan(1 << 20, k)
diff --git a/tests/kernels/test_dsv4_fp8_block_linear.py b/tests/kernels/test_dsv4_fp8_block_linear.py
new file mode 100644
index 0000000000..2194a65e1b
--- /dev/null
+++ b/tests/kernels/test_dsv4_fp8_block_linear.py
@@ -0,0 +1,101 @@
+"""``kernel/triton/dsv4/fp8_linear.block_fp8_linear`` against a torch transcription of the reference
+``act_quant`` + ``fp8_gemm`` (ue8m0 activation scales, ``block x block`` e8m0 weight scales), for the
+DeepSeek-V4 block (128) and the DeepSeek-V4.1 block (32), on the GEMV (M=1) and GEMM (M>1) paths."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+FP8 = torch.float8_e4m3fn
+E8M0 = torch.float8_e8m0fnu
+
+
+def _pow2_ceil_scale(amax: torch.Tensor, qmax: float) -> torch.Tensor:
+    """``2 ** ceil(log2(amax / qmax))`` -- the reference ``fast_round_scale``."""
+    return torch.exp2(torch.ceil(torch.log2(amax / qmax)))
+
+
+def quantize_weight(w: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per ``block x block`` ue8m0-scaled e4m3 weight, scale as e8m0 codes."""
+    n, k = w.shape
+    blocks = w.float().view(n // block, block, k // block, block).permute(0, 2, 1, 3)
+    amax = blocks.abs().amax(dim=(-1, -2)).clamp_min(1e-4)
+    scale = _pow2_ceil_scale(amax, 448.0)  # [n//block, k//block]
+    q = (blocks / scale[..., None, None]).clamp(-448.0, 448.0).to(FP8)
+    w_fp8 = q.permute(0, 2, 1, 3).reshape(n, k)
+    return w_fp8, scale.to(E8M0)
+
+
+def act_quant_reference(x: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]:
+    m, k = x.shape
+    groups = x.float().view(m, k // block, block)
+    amax = groups.abs().amax(dim=-1).clamp_min(1e-4)
+    scale = _pow2_ceil_scale(amax, 448.0)
+    q = (groups / scale[..., None]).clamp(-448.0, 448.0).to(FP8)
+    return q.view(m, k), scale
+
+
+def linear_reference(x: torch.Tensor, w_fp8: torch.Tensor, w_scale: torch.Tensor, block: int) -> torch.Tensor:
+    m, k = x.shape
+    n = w_fp8.shape[0]
+    a_fp8, a_scale = act_quant_reference(x, block)
+    a = a_fp8.float().view(m, k // block, block)
+    w = w_fp8.float().view(n // block, block, k // block, block)
+    ws = w_scale.float()  # [n//block, k//block]
+    out = torch.zeros(m, n, dtype=torch.float32, device=x.device)
+    for kb in range(k // block):
+        wk = w[:, :, kb, :].reshape(n, block)  # [n, block]
+        partial = a[:, kb, :] @ wk.t()  # [m, n]
+        out += partial * a_scale[:, kb, None] * ws[:, kb].repeat_interleave(block)[None, :]
+    return out
+
+
+@pytest.mark.parametrize("block", [32, 128])
+@pytest.mark.parametrize("m", [1, 2, 5, 8, 9, 64])
+def test_block_fp8_linear_matches_reference(block: int, m: int):
+    from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear
+
+    torch.manual_seed(0)
+    # K // block = 5 slabs exercises the GEMV split that must divide the slab count; N spans two GEMM N-tiles
+    k, n = 5 * block, 256 if block == 128 else 224
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3
+    w_fp8, w_scale = quantize_weight(torch.randn(n, k, device="cuda", dtype=torch.bfloat16), block)
+
+    got = block_fp8_linear(x, w_fp8, w_scale, block=block).float()
+    want = linear_reference(x, w_fp8, w_scale, block)
+    # bf16 output; the fp32 reduction order inside the MMA may differ from the slab loop
+    torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2 * want.abs().max().item())
+
+
+def test_block_argument_must_match_the_scale_shape():
+    from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear
+
+    x = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16)
+    w_fp8, w_scale = quantize_weight(torch.randn(256, 256, device="cuda", dtype=torch.bfloat16), 32)
+    with pytest.raises(AssertionError):
+        block_fp8_linear(x, w_fp8, w_scale, block=128)
+
+
+@pytest.mark.parametrize("block", [32, 128])
+def test_grouped_w8a16_gemv_matches_the_dequantized_einsum(block):
+    """``wo_a`` at decode: the fp8 payload dequantized in-kernel against a bf16 activation must equal
+    the reference's bf16 einsum over the dequantized weight (fp32 accumulation, order aside)."""
+    from freetoken.kernel.triton.dsv4.fp8_linear import dequant_block_fp8, grouped_w8a16_gemv
+
+    G, r, d = 4, 256, 1024
+    torch.manual_seed(block)
+    w = torch.randn(G * r, d, device="cuda") * 0.05
+    w_fp8, scale = quantize_weight(w, block)
+    x = torch.randn(G, d, device="cuda", dtype=torch.bfloat16)
+    deq = dequant_block_fp8(w_fp8, scale, block=block)
+    ref_deq = (w_fp8.float() * torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(block, 0).repeat_interleave(block, 1)).to(torch.bfloat16)
+    assert torch.equal(deq, ref_deq)
+    for t in (1, 3):
+        xt = torch.randn(t, G, d, device="cuda", dtype=torch.bfloat16) if t > 1 else x.unsqueeze(0)
+        want = torch.einsum("tgd,grd->tgr", xt, ref_deq.view(G, r, d)).flatten(1)
+        got = grouped_w8a16_gemv(xt, w_fp8, scale, block=block)
+        assert got.shape == (t, G * r) and got.dtype == torch.bfloat16
+        torch.testing.assert_close(got.float(), want.float(), atol=2e-2, rtol=2e-2)
diff --git a/tests/kernels/test_dsv4_router_select.py b/tests/kernels/test_dsv4_router_select.py
new file mode 100644
index 0000000000..9829ed038a
--- /dev/null
+++ b/tests/kernels/test_dsv4_router_select.py
@@ -0,0 +1,52 @@
+"""Fused gate selection (``kernel/triton/dsv4/router``) against the torch expression it replaces, for
+every score function, with and without top-k weight normalization, at decode and prefill widths."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def reference(scores, bias, topk, score_func, normalize, route_scale):
+    if score_func == "softmax":
+        s = scores.softmax(dim=-1)
+    elif score_func == "sigmoid":
+        s = scores.sigmoid()
+    else:
+        s = F.softplus(scores).sqrt()
+    ids = (s + bias).topk(topk, dim=-1)[1]
+    w = s.gather(1, ids)
+    if normalize and topk > 1:
+        w = w / (w.sum(dim=-1, keepdim=True) + 1e-20)
+    return w * route_scale, ids
+
+
+@pytest.mark.parametrize("score_func", ["sqrtsoftplus", "sigmoid", "softmax"])
+@pytest.mark.parametrize("normalize", [True, False])
+@pytest.mark.parametrize("T", [1, 3, 300])
+def test_router_select_matches_torch(score_func, normalize, T):
+    from freetoken.kernel.triton.dsv4.router import router_select
+
+    g = torch.Generator(device="cpu").manual_seed(T)
+    E, K = 384, 6
+    scores = (torch.randn(T, E, generator=g) * 3).cuda()
+    bias = (torch.randn(E, generator=g) * 0.1).cuda()
+    w, ids = router_select(scores, bias, K, score_func=score_func, normalize=normalize, route_scale=2.5)
+    w_ref, ids_ref = reference(scores, bias, K, score_func, normalize, 2.5)
+    assert torch.equal(ids.sort(dim=-1).values, ids_ref.to(torch.int32).sort(dim=-1).values)
+    # same expert set; compare the weights per expert (selection order may differ only on exact ties)
+    order = ids.argsort(dim=-1)
+    order_ref = ids_ref.argsort(dim=-1)
+    torch.testing.assert_close(w.gather(1, order), w_ref.gather(1, order_ref), atol=1e-6, rtol=1e-5)
+
+
+def test_router_select_breaks_ties_toward_the_lowest_expert():
+    from freetoken.kernel.triton.dsv4.router import router_select
+
+    scores = torch.zeros(1, 384, device="cuda")  # every expert ties
+    bias = torch.zeros(384, device="cuda")
+    _, ids = router_select(scores, bias, 6, score_func="sigmoid", normalize=True, route_scale=1.0)
+    assert ids.tolist() == [[0, 1, 2, 3, 4, 5]]
diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py
index 496622f985..cce5df6246 100644
--- a/tests/kvcache/test_cache_unit_bytes.py
+++ b/tests/kvcache/test_cache_unit_bytes.py
@@ -16,7 +16,7 @@
     compute_cache_status_meta,
     compute_cache_unit_bytes,
 )
-from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages
+from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages
 
 
 def _init_tp() -> None:
@@ -204,7 +204,7 @@ def test_floors_dsv4_reports_real_window_floor():
             dsv4_args=SimpleNamespace(window_size=P), has_swa_attention=False
         ),
     )
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
 
     eng = SimpleNamespace(config=cfg, moe_offload_cache=None, linear_state_pool=None)
     eng.kv_cache = object.__new__(DSV4PagedKVCache)  # min_kv_tokens is a classmethod
diff --git a/tests/kvcache/test_dsv41_pool.py b/tests/kvcache/test_dsv41_pool.py
new file mode 100644
index 0000000000..f41df85174
--- /dev/null
+++ b/tests/kvcache/test_dsv41_pool.py
@@ -0,0 +1,205 @@
+"""DSV41 paged KV pool + cost model (CPU, no model): args validation, per-source tiers, sizing /
+byte accounting, the window free-list duck-type, full-loc translation; packed writes on CUDA."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kvcache.dsv4.v41_cost_model import (
+    dsv41_cache_per_page,
+    dsv41_kv_unit_bytes,
+    dsv41_pool_bytes,
+    dsv41_pool_sizes,
+    dsv41_solve_num_pages,
+    dsv41_window_unit_bytes,
+)
+from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache
+from freetoken.kvcache.dsv4.v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32
+from freetoken.models.deepseek_v41.args import DeepseekV41Args
+
+DEVICE = torch.device("cpu")
+P = 128
+# a 10-layer miniature of V4.1's layout: 2 window-only, ratio-2 encoder (source 2, source 5), ratio-1 decoder (source 7)
+RATIOS = (0, 0, 2, 2, 2, 2, 2, 1, 1, 1)
+SOURCES = (2, 5, 7)
+
+
+def _args(**over) -> DeepseekV41Args:
+    base = dict(n_layers=10, head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS,
+                kv_source_layers=SOURCES, swa_decoder_replay="exact")
+    base.update(over)
+    return DeepseekV41Args(**base)
+
+
+def _pool(num_pages=8, swa_ratio=0.5, n_scratch=1, **over):
+    args = _args(**over)
+    sizes = dsv41_pool_sizes(num_pages, args, swa_ratio, P)
+    return DSV41PagedKVCache(sizes, args, DEVICE, n_scratch=n_scratch), sizes, args
+
+
+def test_window_control_uses_pool_units_and_restores_concrete_capacity():
+    from types import SimpleNamespace
+
+    from freetoken.attention import AttnType
+    from freetoken.kvcache.cache_status import (
+        _supports_swa_ratio, compute_cache_floors, compute_cache_pools,
+    )
+    from freetoken.kvcache.dsv4.v41_cost_model import _dsv41_pool_sizes
+    from freetoken.scheduler.scheduler import Scheduler
+    from freetoken.server.api_server import CacheRebuildRequest, _resolve_num_swa_pages, cache_geometry
+    from freetoken.server.stats import _swa_page_size
+
+    pool, _, args = _pool(num_pages=64)
+    config = SimpleNamespace(
+        page_size=P, max_running_req=1, max_seq_len=1024, cache_type="swa_radix",
+        swa_full_tokens_ratio=0.5, swa_num_pages_override=None,
+        model_config=SimpleNamespace(
+            dsv4_args=None, dsv41_args=_args(), has_swa_attention=False,
+            kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.DSV41)],
+        ),
+    )
+    engine = SimpleNamespace(
+        config=config, kv_cache=pool, num_pages=63,
+        moe_offload_cache=None, linear_state_pool=None,
+    )
+    assert _supports_swa_ratio(config)
+    assert compute_cache_pools(engine)["swa_page_size"] == P
+    assert compute_cache_pools(engine)["num_swa_pages"] == 31
+    floor = _dsv41_pool_sizes(config, 64, num_swa_pages=1).n_win_pages - 1
+    assert compute_cache_floors(engine)["swa_tokens"] == floor * P
+    prior = Scheduler._current_cache_geometry(SimpleNamespace(engine=engine, config=config))
+    assert prior["num_swa_pages"] == 31
+    state = SimpleNamespace(config=config, cache_pools=compute_cache_pools(engine))
+    req = CacheRebuildRequest(num_pages=64, swa_full_tokens_ratio=0.5)
+    assert _resolve_num_swa_pages(state, req) == 32
+    assert _swa_page_size(config) == P
+    state.stats = SimpleNamespace(kv_total_pages=63, mamba_total_slots=0)
+    state.last_rebuild = {"num_pages": 64, "num_swa_pages": 32}
+    assert cache_geometry(state)["swa_full_tokens_ratio"] == 0.5
+    pool._init_paged_state(1, True)
+    pool.rebuild_from_config(config, 63, num_swa_pages=24)
+    assert pool.window_pages == 24
+    pool.rebuild_from_config(config, prior["num_pages"], num_swa_pages=prior["num_swa_pages"])
+    assert pool.window_pages == 31
+
+
+def test_pool_reads_sources_and_rings_from_args():
+    pool, _, _ = _pool()
+    assert [pool.source_of(l) for l in range(2, 10)] == [2, 2, 2, 5, 5, 7, 7, 7]
+    assert set(pool.state_ring) == {2, 5} and pool.state_ring[2].ring_size == 2
+    with pytest.raises(ValueError):  # the window is not a multiple of a compress ratio
+        _pool(window_size=P - 1)
+    with pytest.raises(ValueError):  # an indexer key the fp4 row format cannot hold
+        _pool(index_head_dim=48)
+
+
+@pytest.mark.parametrize("replay", ["bounded", "exact"])
+def test_v41_flash_row_bytes_and_private_window_layers(replay):
+    """The tech report's headline: 3 ratio-2 encoder sources + 1 ratio-1 decoder source = 890 B/token."""
+    ratios = (0, 0) + (2,) * 18 + (1,) * 20
+    args = DeepseekV41Args(n_layers=40, head_dim=512, index_head_dim=128, window_size=128, compress_ratios=ratios,
+                           kv_source_layers=(2, 8, 14, 20), swa_decoder_replay=replay)
+    pool = DSV41PagedKVCache(dsv41_pool_sizes(4, args, 1.0, 128), args, DEVICE)
+    win, main, idx = pool.window_pool[0].shape[1], pool.main_pool[2].shape[1], pool.idx_pool[2].shape[1]
+    assert (win, main, idx) == (528, 288, 68)
+    assert pool.state_ring[2].buffer.shape[1] * 4 == 4096
+    assert sum((main + idx) // pool.ratio_of(s) for s in args.backbone_kv_sources) == 890
+    assert dsv41_kv_unit_bytes(args, 128) == 890 + 8  # + the int64 full->window map slot
+    assert pool.private_window_layer_ids == (tuple(range(20, 40)) if replay == "bounded" else ())
+
+
+def test_pool_tiers_per_source_and_aliasing():
+    pool, sizes, args = _pool()
+    assert len(pool.window_pool) == 10 and pool.window_pool[0].shape == (sizes.n_win_slots, 528)
+    assert set(pool.main_pool) == set(SOURCES) == set(pool.idx_pool)
+    assert set(pool.state_ring) == {2, 5}  # the ratio-1 decoder source carries no partial group
+    assert pool.main_pool[2].shape == (sizes.full_token // 2 + 1, 288)
+    assert pool.main_pool[7].shape == (sizes.full_token + 1, 288)
+    assert pool.idx_pool[5].shape == (sizes.full_token // 2 + 1, 68)
+    assert pool.main_pool_of(4) is pool.main_pool[2] and pool.idx_pool_of(9) is pool.idx_pool[7]
+    assert pool.scratch_base[2] == sizes.full_token // 2
+    with pytest.raises(AssertionError):
+        pool.main_pool_of(0)
+    ring = pool.state_ring[2]
+    assert ring.ring_size == 2 and ring.item_size == 512 and ring.buffer.shape == (sizes.state_slots[2] + 1, 1024)
+    assert torch.all(ring.buffer[-1, :512] == 0) and torch.all(torch.isneginf(ring.buffer[-1, 512:]))
+
+
+def test_pool_bytes_match_allocation_and_solver_respects_budget():
+    pool, sizes, args = _pool(num_pages=16, swa_ratio=0.25)
+    assert pool.total_bytes() == dsv41_pool_bytes(sizes, args, n_scratch=1)
+    assert dsv41_cache_per_page(args, 0.25, P) > 0
+    assert dsv41_window_unit_bytes(args, P) == -(-(10 * P * 528 + 2 * 2 * 4096) // P)
+    budget = 64 << 20
+    solved = dsv41_solve_num_pages(budget, args, 0.25, floor_win_pages=4, P=P, n_scratch=3)
+    assert dsv41_pool_bytes(solved, args, 3) <= budget
+    # one more page, sized the way the solver sizes (window = max(floor, ceil(ratio * pages))), overflows
+    more = solved.full_token // P + 1
+    bigger = dsv41_pool_sizes(more, args, 0.25, P, n_win_pages=max(4, (round(0.25 * more * P) + P - 1) // P))
+    assert dsv41_pool_bytes(bigger, args, 3) > budget
+    assert solved.n_win_pages >= 4
+    with pytest.raises(ValueError):
+        dsv41_solve_num_pages(1 << 10, args, 0.25, floor_win_pages=4, P=P)
+
+
+def test_translation_state_loc_and_cmp_rows():
+    pool, sizes, args = _pool(num_pages=16)
+    pool.bind_window_pages(full_page_base=0, window_page_base=2 * P)
+    pool.bind_window_pages(full_page_base=3 * P, window_page_base=0)
+    assert pool.translate_full_to_window(torch.tensor([0, 1, 127])).tolist() == [2 * P, 2 * P + 1, 2 * P + 127]
+    assert pool.translate_full_to_window(torch.tensor([3 * P + 5])).item() == 5
+    assert pool.translate_full_to_window(torch.tensor([P + 7, -1])).tolist() == [-1, -1]
+    full = torch.tensor([0, 1, 2, 127, 128, 4 * P - 1])
+    assert pool.cmp_rows(full, 2).tolist() == [0, 0, 1, 63, 64, 2 * P - 1]
+    assert pool.cmp_rows(full, 1).tolist() == full.tolist()
+    assert pool.cmp_rows(torch.tensor([-1]), 2).item() < 0
+    top = pool.cmp_rows(torch.tensor([sizes.full_token - 1]), 2).item()
+    assert top < pool.scratch_base[2]
+    ws = pool.translate_full_to_window(torch.tensor([127, 3 * P]))
+    assert DSV41PagedKVCache.state_loc(ws, 2, P).tolist() == [2 * 2 + 1, 0]
+    assert DSV41PagedKVCache.state_loc(torch.tensor([-1]), 2, P).item() == -1
+
+
+def _expand(bases):
+    return (torch.tensor(bases, dtype=torch.int64)[:, None] + torch.arange(P)).flatten()
+
+
+def test_swa_duck_type_alloc_free_and_dummy():
+    pool, sizes, _ = _pool(num_pages=8)
+    pool._init_paged_state(max_running_req=2, radix=True)
+    cap = sizes.n_win_slots - P
+    assert pool.swa_available_size() == cap and pool.swa_num_tokens - 1 == cap
+    assert pool.sliding_window_size == P and pool.prefill_chunk_budget >= P
+    pool.alloc_swa(_expand([0, 2 * P]))
+    assert pool.swa_available_size() == cap - 2 * P
+    ws = pool.translate_loc_from_full_to_swa(torch.arange(2 * P, 3 * P))
+    assert (ws >= 0).all() and int(ws[0]) % P == 0 and torch.equal(ws - ws[0], torch.arange(P))
+    pool.free_swa(_expand([0]))
+    pool.free_swa(_expand([0]))  # idempotent
+    assert pool.swa_available_size() == cap - P
+    with pytest.raises(AssertionError):
+        pool.free_swa(torch.arange(2 * P, 2 * P + 5))
+    dummy = pool.translate_loc_from_full_to_swa(torch.arange(sizes.full_token - P, sizes.full_token))
+    assert int(dummy[0]) == sizes.n_win_slots - P
+    assert pool.k_cache(3) is pool.window_pool[3] and pool.num_layers == 10 and pool.unit_bytes()[0] > 0
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="packed writes run through triton")
+def test_packed_writes_round_trip_on_cuda():
+    from kernels.test_dsv41_pack import reference_roundtrip  # tests/ is on sys.path under pytest
+
+    args = _args()
+    sizes = dsv41_pool_sizes(4, args, 0.5, P)
+    pool = DSV41PagedKVCache(sizes, args, torch.device("cuda"), n_scratch=2)
+    kv = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16)
+    slots = torch.tensor([0, 130, 5], device="cuda")
+    pool.store_window(kv, 4, slots)
+    assert torch.equal(pool.read_window(4, slots), reference_roundtrip(kv, FP8_E8M0_B32))
+    pool.store_main(kv, 2, slots)
+    assert torch.equal(pool.read_main(2, slots), reference_roundtrip(kv, FP4_E4M3_B16))
+    k = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16)
+    pool.store_index(k, 7, slots)
+    assert torch.equal(pool.read_index(7, slots), reference_roundtrip(k, FP4_E8M0_B32))
+    pool.store_kv(kv, kv, slots, 0)
+    assert torch.equal(pool.read_window(0, slots), reference_roundtrip(kv, FP8_E8M0_B32))
diff --git a/tests/kvcache/test_dsv4_free_list.py b/tests/kvcache/test_dsv4_free_list.py
index a3ae8d0cd7..7b0dfb5566 100644
--- a/tests/kvcache/test_dsv4_free_list.py
+++ b/tests/kvcache/test_dsv4_free_list.py
@@ -5,7 +5,7 @@
 import pytest
 import torch
 
-from freetoken.kvcache.dsv4_paged_pool import FreeListAllocator
+from freetoken.kvcache.dsv4.v4_pool import FreeListAllocator
 
 DEVICE = torch.device("cpu")
 P = 128
diff --git a/tests/kvcache/test_dsv4_pool.py b/tests/kvcache/test_dsv4_pool.py
index c345b67c82..c8e21997ed 100644
--- a/tests/kvcache/test_dsv4_pool.py
+++ b/tests/kvcache/test_dsv4_pool.py
@@ -10,8 +10,8 @@
 import pytest
 import torch
 
-from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes
-from freetoken.kvcache.dsv4_paged_pool import CompressStateRing, DSV4PagedKVCache
+from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes
+from freetoken.kvcache.dsv4.v4_pool import CompressStateRing, DSV4PagedKVCache
 from freetoken.models.deepseek_v4.args import DeepseekV4Args
 
 DEVICE = torch.device("cpu")
diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py
index a5ec94a084..0e3b7050e8 100644
--- a/tests/kvcache/test_kv_cache_rebuild.py
+++ b/tests/kvcache/test_kv_cache_rebuild.py
@@ -178,13 +178,13 @@ def test_dsv4_rebuild_from_config_builds_the_pool_sizes_and_attaches_the_table()
     the engine used to do -- and re-points full_loc_map at the shared page table."""
     from types import SimpleNamespace
 
-    from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes
-    from freetoken.kvcache.dsv4_cost_model import (
+    from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes
+    from freetoken.kvcache.dsv4.v4_cost_model import (
         dsv4_kv_unit_bytes,
         dsv4_pool_sizes,
         dsv4_window_unit_bytes,
     )
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
     from freetoken.models.deepseek_v4.args import DeepseekV4Args
 
     P, mrr = 128, 1
@@ -222,7 +222,7 @@ def test_dsv4_refresh_seq_state_tracks_page_table_width():
 
     from freetoken.engine.engine import Engine
     from freetoken.utils import align_ceil
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
     from freetoken.scheduler.table import TableManager
 
     P, mrr = 128, 4
@@ -264,7 +264,7 @@ def test_every_kv_pool_answers_the_sizing_surface():
     the model resolved to. DSV4 must carry its OWN overrides (getattr would resolve to the
     base class and silently misprice the window floor if an override vanished)."""
     from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
     from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache
     from freetoken.kvcache.mha_pool import MHAKVCache
 
@@ -279,7 +279,7 @@ def test_every_kv_pool_answers_the_rebuild_surface():
     """The engine drives every pool through this one surface; a pool that skips an
     implementation must fail loudly at import/instantiation, not at rebuild time."""
     from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
     from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache
     from freetoken.kvcache.mha_pool import MHAKVCache
 
diff --git a/tests/kvcache/test_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py
index 516090b099..33f27983de 100644
--- a/tests/kvcache/test_pool_sizing_surface.py
+++ b/tests/kvcache/test_pool_sizing_surface.py
@@ -30,7 +30,7 @@ def _model_config(specs, **attrs):
 
 def test_resolve_pool_class_follows_attn_type():
     from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
     from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache
     from freetoken.kvcache.mha_pool import MHAKVCache
 
@@ -91,11 +91,11 @@ def _dsv4_config(num_page_override=None):
 
 
 def test_dsv4_kv_cost_and_floor_parity():
-    from freetoken.kvcache.dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
 
-    pytest.importorskip("freetoken.kvcache.dsv4_cost_model")
-    from freetoken.kvcache.dsv4_cost_model import dsv4_auto_cost_model
+    pytest.importorskip("freetoken.kvcache.dsv4.v4_cost_model")
+    from freetoken.kvcache.dsv4.v4_cost_model import dsv4_auto_cost_model
 
     config = _dsv4_config()
     P = 128
@@ -147,8 +147,8 @@ def check(pages, baseline):
 
 def test_dsv4_validate_rebuild_floor():
     from freetoken.kvcache.base import CacheRebuildRejected
-    from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages
-    from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+    from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages
+    from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
 
     config = _dsv4_config()
     pool = object.__new__(DSV4PagedKVCache)  # floor path reads no instance state
@@ -165,7 +165,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins():
     # The explicit parameter must take precedence over config.swa_num_pages_override --
     # the engine's override write would otherwise mask a dropped parameter forever.
     from freetoken.distributed import set_tp_info, try_get_tp_info
-    from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes
+    from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes
     from freetoken.kvcache.hybrid_swa_pool import _swa_paged_num_tokens, _swa_pool_floor
 
     if try_get_tp_info() is None:
@@ -183,7 +183,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins():
     assert _swa_paged_num_tokens(swa_cfg, 11, num_swa_pages=explicit) == explicit + 1
     assert _swa_paged_num_tokens(swa_cfg, 11) == 999_999 + 1  # override path unchanged
 
-    from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages
+    from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages
 
     d_cfg = _dsv4_config()
     d_floor = _dsv4_window_floor_pages(d_cfg, d_cfg.page_size)
diff --git a/tests/layers/test_mhc.py b/tests/layers/test_mhc.py
index 4cf144e720..0359bed9c9 100644
--- a/tests/layers/test_mhc.py
+++ b/tests/layers/test_mhc.py
@@ -152,6 +152,75 @@ def test_triton_fused_respects_input_dtype(dtype, tol):
         assert err < tol, f"{name} [{dtype}]: max abs err {err}"
 
 
+def test_single_pass_applies_the_previous_pre_and_hands_on_its_own():
+    """Single-pass mHC (DeepSeek-V4.1): the input mix uses ``pre_mix_in``, the gates
+    computed here are the next sublayer's pre; post/comb are unchanged from multi-pass."""
+    from freetoken.layers.mhc import (
+        mhc_fused_post_pre_single_pass,
+        mhc_gates,
+        mhc_mix_input,
+        mhc_pre_single_pass,
+    )
+
+    fn, scale, base = _weights(seed=21)
+    res = _residual(seed=22)
+    x = torch.randn(T, HIDDEN, dtype=torch.bfloat16)
+    pre_in = torch.rand(T, N, dtype=torch.float32)
+    post0, comb0, _ = mhc_pre(res, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN)
+
+    r1, p1, c1, pre1, li1 = mhc_fused_post_pre_single_pass(
+        x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN
+    )
+    r_ref = mhc_post(x, res, post0, comb0)
+    pre_ref, p_ref, c_ref = mhc_gates(r_ref, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN)
+    assert torch.equal(r1, r_ref) and torch.equal(p1, p_ref) and torch.equal(c1, c_ref)
+    assert torch.equal(pre1, pre_ref)
+    assert torch.equal(li1, mhc_mix_input(r_ref, pre_in))
+    # the multi-pass input would have used pre_ref, not pre_in
+    assert not torch.equal(li1, mhc_mix_input(r_ref, pre_ref))
+    # a one-hot pre picks one stream (the model's seed for the chain)
+    one_hot = torch.zeros(T, N)
+    one_hot[:, 0] = 1.0
+    _, _, _, li0 = mhc_pre_single_pass(res, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN)
+    assert torch.equal(li0, res[:, 0])
+    # no post on the first sublayer: the streams pass through untouched
+    r0, *_ = mhc_fused_post_pre_single_pass(
+        x, res, None, None, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN
+    )
+    assert torch.equal(r0, res)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("t,hidden,with_post", [(1, 64, True), (9, 64, True), (3, 5120, True), (3, 5120, False)])
+def test_triton_single_pass_matches_torch(t, hidden, with_post):
+    from freetoken.layers.mhc import mhc_fused_post_pre_single_pass_torch
+    from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton
+
+    torch.manual_seed(23)
+    mix = 2 * N + N * N
+    fn = torch.randn(mix, N * hidden, dtype=torch.float32, device="cuda") * 0.05
+    scale = torch.rand(3, dtype=torch.float32, device="cuda") + 0.5
+    base = torch.randn(mix, dtype=torch.float32, device="cuda") * 0.3
+    res = torch.randn(t, N, hidden, dtype=torch.bfloat16, device="cuda")
+    x = torch.randn(t, hidden, dtype=torch.bfloat16, device="cuda")
+    pre_in = torch.rand(t, N, dtype=torch.float32, device="cuda")
+    post0 = torch.rand(t, N, 1, dtype=torch.float32, device="cuda") * POST_MULT if with_post else None
+    comb0 = torch.softmax(torch.randn(t, N, N, device="cuda"), dim=-1) if with_post else None
+
+    ref = mhc_fused_post_pre_single_pass_torch(
+        x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN
+    )
+    got = mhc_fused_post_pre_single_pass_triton(
+        x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN
+    )
+    names = ("residual", "post", "comb", "pre_new", "layer_input")
+    tols = (2e-2, 2e-3, 2e-3, 2e-3, 2e-2)
+    for name, r, g, tol in zip(names, ref, got, tols):
+        assert g.shape == r.shape, (name, g.shape, r.shape)
+        err = (g.float() - r.float()).abs().max().item()
+        assert err < tol, f"{name}: max abs err {err}"
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 def test_triton_pre_only_matches_torch():
     """HAS_POST=False path (layer 0's standalone hc_pre through the fused kernel)."""
@@ -176,3 +245,31 @@ def test_triton_pre_only_matches_torch():
     assert (got_post.float() - ref_post.float()).abs().max().item() < 2e-3
     assert (got_comb.float() - ref_comb.float()).abs().max().item() < 2e-3
     assert (got_li.float() - ref_li.float()).abs().max().item() < 2e-2
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_mhc_split_preserves_chunked_results_and_graph_replay():
+    from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton as fused
+
+    torch.manual_seed(26)
+    hidden, tokens = 5120, 7
+    res = torch.randn(tokens, N, hidden, device="cuda", dtype=torch.bfloat16)
+    x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16)
+    post = torch.rand(tokens, N, 1, device="cuda")
+    comb = torch.softmax(torch.randn(tokens, N, N, device="cuda"), -1)
+    pre = torch.rand(tokens, N, device="cuda")
+    fn = torch.randn(2 * N + N * N, N * hidden, device="cuda") * 0.05
+    scale, base = torch.ones(3, device="cuda"), torch.zeros(2 * N + N * N, device="cuda")
+    def run(lo, hi):
+        return fused(x[lo:hi], res[lo:hi], post[lo:hi], comb[lo:hi], pre[lo:hi], fn, scale, base,
+                     RMS_EPS, EPS, POST_MULT, SINKHORN)
+    whole = run(0, tokens)
+    chunks = [run(0, 2), run(2, tokens)]
+    for i, expected in enumerate(whole):
+        assert torch.equal(expected, torch.cat([chunk[i] for chunk in chunks]))
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = run(0, tokens)
+    graph.replay()
+    for got, expected in zip(captured, whole):
+        assert torch.equal(got, expected)
diff --git a/tests/models/deepseek_v41/__init__.py b/tests/models/deepseek_v41/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/models/deepseek_v41/common.py b/tests/models/deepseek_v41/common.py
new file mode 100644
index 0000000000..d0f8efeef9
--- /dev/null
+++ b/tests/models/deepseek_v41/common.py
@@ -0,0 +1,176 @@
+"""Shared fixtures for the deepseek_v41 tests: a tiny V4.1-shaped config (2 window layers, a ratio-2
+encoder group with one Full + one Reuse layer, a ratio-1 decoder group with a Full candidate source and
+a Reindex layer) and synthetic bf16 and FP8/MXFP4 checkpoints in the real tensor naming. Holds no tests itself."""
+
+from __future__ import annotations
+
+import json
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+VOCAB = 512
+DIM = 256
+N_LAYERS = 6
+RATIOS = [0, 0, 2, 2, 1, 1, 0]  # + one trailing MTP layer the model ignores
+KV_SOURCES = [2, 4]
+INDEX_SOURCES = [2, 4, 5]
+ENGRAM_LAYERS = [1]
+ENGRAM_VOCAB = 500  # bucket prime start; 3 orders x 2 heads ~= 3000 rows
+ENGRAM_ROWS = 4096
+
+
+def tiny_text_config(**over) -> dict:
+    text = dict(
+        model_type="deepseek_v41_text",
+        vocab_size=VOCAB, hidden_size=DIM, moe_intermediate_size=64, num_hidden_layers=N_LAYERS, num_attention_heads=4,
+        num_key_value_heads=1, head_dim=64, qk_rope_head_dim=16, q_lora_rank=64, o_lora_rank=32, o_groups=2,
+        hidden_act="silu", swiglu_limit=10.0, rms_norm_eps=1e-20, max_position_embeddings=16384,
+        rope_theta=10000, rope_scaling=dict(rope_type="yarn", factor=16, beta_fast=32, beta_slow=1, original_max_position_embeddings=1024),
+        n_routed_experts=8, n_shared_experts=1, num_experts_per_tok=2, scoring_func="sqrtsoftplus", topk_method="noaux_tc",
+        norm_topk_prob=True, routed_scaling_factor=1.5, sliding_window=128,
+        compress_ratios=list(RATIOS), compress_rope_theta=160000, kv_source_layer_ids=list(KV_SOURCES),
+        index_source_layer_ids=list(INDEX_SOURCES), index_n_heads=2, index_head_dim=32, index_topk=8,
+        candidate_source_layer_id=4, candidate_topk_blocks=2, candidate_block_size=4,
+        hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-6,
+        engram_layer_ids=list(ENGRAM_LAYERS), engram_num_embeddings=[ENGRAM_ROWS], engram_max_ngram_size=4,
+        engram_vocab_size=ENGRAM_VOCAB, engram_n_heads=2, engram_head_dim=32, engram_pad_token_id=2,
+        engram_compressed_vocab_size=VOCAB,  # the synthetic tokenizer maps every id to itself
+        num_nextn_predict_layers=1, dspark_block_size=0, dspark_target_layer_ids=[],
+    )
+    text.update(over)
+    return text
+
+
+def tiny_hf_config(**over) -> SimpleNamespace:
+    return SimpleNamespace(architectures=["DeepseekV41ForCausalLM"], model_type="deepseek_v41", text_config=tiny_text_config(**over))
+
+
+def checkpoint_tensors(text: dict, seed: int = 0) -> dict[str, torch.Tensor]:
+    """Random bf16 weights under the checkpoint's names (the reference ``convert.py`` dialect without
+    quantization: no ``.scale`` tensors). hc / sink / router-bias / engram gates in fp32 as shipped."""
+    g = torch.Generator().manual_seed(seed)
+
+    def rnd(*shape, scale=0.02, dtype=torch.bfloat16):
+        return (torch.randn(*shape, generator=g) * scale).to(dtype)
+
+    dim, heads, hd, rd = text["hidden_size"], text["num_attention_heads"], text["head_dim"], text["qk_rope_head_dim"]
+    qr, ol, og = text["q_lora_rank"], text["o_lora_rank"], text["o_groups"]
+    inter, E, hc = text["moe_intermediate_size"], text["n_routed_experts"], text["hc_mult"]
+    ih, ihd = text["index_n_heads"], text["index_head_dim"]
+    mix = (2 + hc) * hc
+    t: dict[str, torch.Tensor] = {
+        "embed.weight": rnd(text["vocab_size"], dim, scale=1.0),
+        "norm.weight": torch.ones(dim, dtype=torch.bfloat16),
+        "head.weight": rnd(text["vocab_size"], dim),
+    }
+    for L in range(text["num_hidden_layers"]):
+        a = f"layers.{L}.attn"
+        t[f"{a}.wq_a.weight"] = rnd(qr, dim)
+        t[f"{a}.q_norm.weight"] = torch.ones(qr, dtype=torch.bfloat16)
+        t[f"{a}.wq_b.weight"] = rnd(heads * hd, qr)
+        t[f"{a}.wkv.weight"] = rnd(hd, dim)
+        t[f"{a}.kv_norm.weight"] = torch.ones(hd, dtype=torch.bfloat16)
+        t[f"{a}.wo_a.weight"] = rnd(og * ol, heads * hd // og)
+        t[f"{a}.wo_b.weight"] = rnd(dim, og * ol)
+        t[f"{a}.attn_sink"] = rnd(heads, scale=0.5, dtype=torch.float32)
+        ratio = text["compress_ratios"][L]
+        if L in text["kv_source_layer_ids"]:
+            t[f"{a}.compressor.wkv.weight"] = rnd(hd, dim)
+            if ratio > 1:
+                t[f"{a}.compressor.wgate.weight"] = rnd(hd, dim)
+            t[f"{a}.compressor.norm.weight"] = torch.ones(hd, dtype=torch.bfloat16)
+            t[f"{a}.indexer.wk.weight"] = rnd(ihd, hd)
+            t[f"{a}.indexer.k_norm.weight"] = torch.ones(ihd, dtype=torch.bfloat16)
+        if L in text["index_source_layer_ids"]:
+            t[f"{a}.indexer.wq_b.weight"] = rnd(ih * ihd, qr)
+            t[f"{a}.indexer.weights_proj.weight"] = rnd(ih, dim)
+        p = f"layers.{L}"
+        t[f"{p}.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16)
+        t[f"{p}.ffn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16)
+        t[f"{p}.ffn.gate.weight"] = rnd(E, dim)
+        t[f"{p}.ffn.gate.bias"] = rnd(E, scale=0.1, dtype=torch.float32)
+        t[f"{p}.ffn.gate.bias_vl"] = rnd(E, scale=0.1, dtype=torch.float32)
+        for w, shape in (("w1", (inter, dim)), ("w2", (dim, inter)), ("w3", (inter, dim))):
+            t[f"{p}.ffn.shared_experts.{w}.weight"] = rnd(*shape)
+            for e in range(E):
+                t[f"{p}.ffn.experts.{e}.{w}.weight"] = rnd(*shape)
+        for nm in ("hc_attn_fn", "hc_ffn_fn"):
+            t[f"{p}.{nm}"] = rnd(mix, hc * dim, scale=0.05, dtype=torch.float32)
+        for nm in ("hc_attn_base", "hc_ffn_base"):
+            t[f"{p}.{nm}"] = rnd(mix, scale=0.3, dtype=torch.float32)
+        for nm in ("hc_attn_scale", "hc_ffn_scale"):
+            t[f"{p}.{nm}"] = (torch.rand(3, generator=g) + 0.5).to(torch.float32)
+        if L in text["engram_layer_ids"]:
+            e = f"{p}.engram"
+            cols = (text["engram_max_ngram_size"] - 1) * text["engram_n_heads"]
+            t[f"{e}.wkv.weight"] = rnd(dim * (hc + 1), cols * text["engram_head_dim"])
+            t[f"{e}.q_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16)
+            t[f"{e}.k_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16)
+            rows = text["engram_num_embeddings"][text["engram_layer_ids"].index(L)]
+            t[f"{e}.embed.weight"] = rnd(rows, text["engram_head_dim"], scale=1.0).to(torch.float8_e4m3fn)
+            t[f"{e}.embed.scale"] = torch.full((rows, text["engram_head_dim"] // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu)
+    # the DSpark stage and the vision tower: present in a real checkpoint, dropped by the loader
+    t["mtp.0.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16)
+    t["vision.norm.weight"] = torch.ones(8, dtype=torch.bfloat16)
+    return t
+
+
+def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False, **over) -> dict[str, torch.Tensor]:
+    """A loadable checkpoint directory: HF-style config.json + one safetensors shard + index."""
+    from safetensors.torch import save_file
+
+    if quantized:
+        over.setdefault("moe_intermediate_size", 256)
+    text = tiny_text_config(**over)
+    os.makedirs(folder, exist_ok=True)
+    with open(os.path.join(folder, "config.json"), "w") as f:
+        config = {"architectures": ["DeepseekV41ForCausalLM"], "model_type": "deepseek_v41", "text_config": text}
+        if quantized:
+            config["quantization_config"] = {
+                "quant_method": "fp8", "activation_scheme": "dynamic",
+                "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4",
+            }
+        json.dump(config, f)
+    tensors = checkpoint_tensors(text, seed)
+    if quantized:
+        tensors = _quantize_checkpoint(tensors)
+    save_file(tensors, os.path.join(folder, "model-00001-of-00001.safetensors"))
+    with open(os.path.join(folder, "model.safetensors.index.json"), "w") as f:
+        json.dump({"metadata": {}, "weight_map": {k: "model-00001-of-00001.safetensors" for k in tensors}}, f)
+    return tensors
+
+
+def _quantize_checkpoint(tensors):
+    from .reference.kernel import _round_fp4
+
+    out = dict(tensors)
+    fp8_suffixes = (
+        ".attn.wq_a", ".attn.wq_b", ".attn.wkv", ".attn.wo_a", ".attn.wo_b",
+        ".indexer.wq_b", ".shared_experts.w1", ".shared_experts.w2", ".shared_experts.w3", ".engram.wkv",
+    )
+    grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6.])
+    for name, weight in tensors.items():
+        if not name.endswith(".weight"):
+            continue
+        stem = name.removesuffix(".weight")
+        if ".ffn.experts." in name:
+            groups = weight.float().unflatten(-1, (-1, 32))
+            scale = torch.exp2(torch.ceil(torch.log2(groups.abs().amax(-1).clamp_min(1e-12) / 6)))
+            q = _round_fp4(groups / scale.unsqueeze(-1)).flatten(-2)
+            codes = (q.abs().unsqueeze(-1) == grid).to(torch.int32).argmax(-1).to(torch.uint8)
+            codes |= torch.signbit(q).to(torch.uint8) * 8
+            out[name] = codes[:, 0::2] | (codes[:, 1::2] << 4)
+            out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu)
+        elif stem.endswith(fp8_suffixes):
+            n, k = weight.shape
+            blocks = weight.float().reshape(n // 32, 32, k // 32, 32).permute(0, 2, 1, 3)
+            scale = torch.exp2(torch.ceil(torch.log2(blocks.abs().amax((-1, -2)).clamp_min(1e-4) / 448)))
+            expanded = scale.repeat_interleave(32, 0).repeat_interleave(32, 1)
+            out[name] = (weight.float() / expanded).to(torch.float8_e4m3fn)
+            out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu)
+    return out
diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py
new file mode 100644
index 0000000000..98b19e2de4
--- /dev/null
+++ b/tests/models/deepseek_v41/harness.py
@@ -0,0 +1,182 @@
+"""A minimal in-process engine for the tiny DeepSeek-V4.1 model: builds the DSV41 pool, backend and
+context by hand (no scheduler), loads a synthetic bf16 or FP8/MXFP4 checkpoint with resident or offloaded experts, and drives
+ragged prefill / batched decode batches through ``model.forward()``. Shared by the forward smoke test
+and the reference-parity test."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import torch
+
+from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx
+from freetoken.distributed.info import set_tp_info, try_get_tp_info
+from freetoken.engine.engine import _materialize_loaded_weight_state_dict
+from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes
+from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache
+from freetoken.layers import set_rope_device
+from freetoken.layers.quantization import NoQuantConfig
+from freetoken.layers.quantization.method import finalize_quant
+from freetoken.models import create_model
+from freetoken.models.deepseek_v41.config import parse_config
+from freetoken.models.deepseek_v41.engram import ZeroEngramTable
+from freetoken.models.deepseek_v41.weight import iter_weights
+from freetoken.utils.hf import cached_load_hf_config
+from freetoken.utils.torch_utils import torch_dtype
+
+P = 128
+
+
+class TinyEngine:
+    def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: int = 2, swa_decoder_replay: str = "exact", engram_table=None, quantized: bool = False):
+        self.device = torch.device("cuda")
+        self.checkpoint = checkpoint
+        if try_get_tp_info() is None:
+            set_tp_info(0, 1)
+        set_rope_device(self.device)
+        hf_config = cached_load_hf_config(checkpoint)
+        mc = parse_config(hf_config)
+        args = mc.dsv41_args
+        args.max_seq_len = max_seq_len
+        args.max_batch_size = max_running_req + 1
+        args.swa_decoder_replay = swa_decoder_replay
+        self.args = args
+        quant = NoQuantConfig()
+        if quantized:
+            from freetoken.models.register import checkpoint_quant_config, get_model_spec
+
+            quant = checkpoint_quant_config(checkpoint, hf_config, get_model_spec("DeepseekV41ForCausalLM"))
+        self.config = replace(mc, moe_strategy="offload" if quantized else "fused", decode_target="gpu", quant=quant)
+        mc = self.config
+        with torch.device("meta"), torch_dtype(torch.bfloat16):
+            self.model = create_model(self.config)
+        state = _materialize_loaded_weight_state_dict(
+            self.model.state_dict(), iter_weights(checkpoint, self.device, include_moe_experts=False), device=self.device,
+        )
+        self.model.load_state_dict(state)
+        finalize_quant(self.model)
+        from freetoken.moe.expert_banks import load_expert_banks
+
+        method = self.model.model.layers.op_list[0].ffn.experts.quant_method
+        if quantized:
+            from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache
+
+            self.banks = load_expert_banks(
+                checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False,
+            )
+            self.expert_cache = OffloadMoeCache(
+                num_layers=mc.num_moe_layers, num_experts=mc.num_experts,
+                cache_size=mc.num_moe_layers * mc.num_experts, device=self.device,
+                quant_format=self.banks.quant_format, layout=method.layout(),
+            )
+            self.expert_cache.set_bank_sources(self.banks.sources, layer_residency=self.banks.layer_residency)
+            attach_offload_moe_cache(self.model, self.expert_cache)
+        else:  # resident experts load as banks, not through the state dict (as in Engine._load_resident_experts)
+            from freetoken.layers import iter_moe_layers
+            from freetoken.moe.expert_banks import attach_resident_banks
+
+            attach_resident_banks(list(iter_moe_layers(self.model)), load_expert_banks(
+                checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, resident=True,
+            ))
+        for layer in self.model.engram_layers():
+            layer.attach_table(engram_table if engram_table is not None else ZeroEngramTable(layer.width, self.device))
+
+        self.max_running_req = max_running_req
+        num_pages = max_seq_len // P
+        self.pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, self.device, n_scratch=max_running_req + 1)
+        self.pool._init_paged_state(max_running_req, radix=False)
+        self.page_table = torch.zeros(max_running_req + 1, max_seq_len, dtype=torch.int32, device=self.device)
+        self.page_table[max_running_req].fill_(num_pages * P)  # the dummy row -> the reserved tail page
+        self.pool.attach_page_table(self.page_table)
+        try:
+            ctx = get_global_ctx()
+        except AssertionError:
+            ctx = Context(page_size=P)
+            set_global_ctx(ctx)
+        from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend
+
+        self.ctx = ctx
+        ctx.page_table = self.page_table
+        ctx.kv_cache = self.pool  # the backend reads the pool's device from the context
+        self.backend = DSV41SparseAttnBackend(self.config)
+        self._next_page = 0
+        self._bind()
+
+    def _bind(self) -> None:
+        """Point the process-wide ``Context`` at this engine's page table, pool and backend. The
+        context is a singleton the model reads through ``get_global_ctx()``, so a test that holds two
+        engines must rebind before each forward; ``prefill`` / ``decode`` do this, and the model
+        re-resolves its backend when the binding changes."""
+        ctx = self.ctx
+        if (
+            getattr(ctx, "attn_backend", None) is self.backend
+            and getattr(ctx, "kv_cache", None) is self.pool
+            and getattr(ctx, "page_table", None) is self.page_table
+        ):
+            return
+        ctx.page_table = self.page_table
+        ctx.kv_cache = self.pool
+        ctx.attn_backend = self.backend
+        self.model.mark_for_rebind()
+
+    # ----- requests -----------------------------------------------------------------------
+    def new_request(self, table_idx: int, tokens: list[int], max_new: int = 64) -> Req:
+        """A request with pages for ``len(tokens) + max_new`` positions bound on row ``table_idx``."""
+        n_pages = -(-(len(tokens) + max_new) // P)
+        first = self._next_page
+        self._next_page += n_pages
+        locs = torch.arange(first * P, (first + n_pages) * P, device=self.device)
+        self.page_table[table_idx, : n_pages * P] = locs.to(torch.int32)
+        self.pool.alloc_swa(locs)
+        return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=0, output_len=max_new,
+                   uid=table_idx, sampling_params=SamplingParams(), cache_handle=None)
+
+    def new_request_on_prefix(self, table_idx: int, donor: Req, tokens: list[int], max_new: int = 64) -> Req:
+        """A request whose first ``hit`` positions alias ``donor``'s pages (what a radix prefix hit does
+        through the page table) and whose remaining positions get pages of their own. ``hit`` is what the
+        scheduler admits: the donor's prompt, capped by ``match_req`` (the last token, and the pool's
+        ``prefix_replay_tokens`` under bounded replay) and page-aligned. The caller prefills the rest."""
+        limit = len(tokens) - 1
+        if self.pool.prefix_replay_tokens:
+            limit = min(limit, max(0, len(tokens) - self.pool.prefix_replay_tokens))
+        hit = min(donor.device_len, limit) // P * P
+        assert tokens[:hit] == donor.input_ids[:hit].tolist(), "the prefix must match the donor's prompt"
+        self.page_table[table_idx, :hit] = self.page_table[donor.table_idx, :hit]
+        n_pages = -(-(len(tokens) + max_new - hit) // P)
+        first = self._next_page
+        self._next_page += n_pages
+        locs = torch.arange(first * P, (first + n_pages) * P, device=self.device)
+        self.page_table[table_idx, hit : hit + n_pages * P] = locs.to(torch.int32)
+        self.pool.alloc_swa(locs)
+        return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=hit, output_len=max_new,
+                   uid=table_idx, sampling_params=SamplingParams(), cache_handle=None)
+
+    def prefill(self, reqs: list[Req]) -> torch.Tensor:
+        """Run each request's ``[cached_len, device_len)`` tokens; returns ``[B, vocab]`` logits."""
+        batch = Batch(reqs=reqs, phase="prefill")
+        batch.padded_reqs = list(reqs)
+        batch.input_ids = torch.cat([r.input_ids[r.cached_len : r.device_len] for r in reqs]).to(self.device)
+        batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).to(self.device)
+        self._bind()
+        self.backend.prepare_metadata(batch)
+        with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False):
+            return self.model.forward().float()
+
+    def decode(self, reqs: list[Req], tokens: list[int]) -> torch.Tensor:
+        """Feed one new token per request (appended at ``device_len - 1``); returns ``[B, vocab]``."""
+        for r, t in zip(reqs, tokens):
+            r.append_host(torch.tensor([t], dtype=torch.int32))
+            r.complete_one()
+        batch = Batch(reqs=reqs, phase="decode")
+        batch.padded_reqs = list(reqs)
+        batch.input_ids = torch.tensor(tokens, dtype=torch.int32, device=self.device)
+        batch.positions = torch.tensor([r.device_len - 1 for r in reqs], device=self.device)
+        batch.active_table_idx = torch.tensor([r.table_idx for r in reqs], device=self.device)
+        self._bind()
+        self.backend.prepare_metadata(batch)
+        with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False):
+            return self.model.forward().float()
+
+    def finish_prefill(self, reqs: list[Req]) -> None:
+        for r in reqs:
+            r.cached_len = r.device_len
diff --git a/tests/models/deepseek_v41/reference/LICENSE b/tests/models/deepseek_v41/reference/LICENSE
new file mode 100644
index 0000000000..d42fae903e
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 DeepSeek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/tests/models/deepseek_v41/reference/__init__.py b/tests/models/deepseek_v41/reference/__init__.py
new file mode 100644
index 0000000000..79d0c33108
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/__init__.py
@@ -0,0 +1,2 @@
+"""The DeepSeek-V4.1 reference implementation (MIT, vendored from the HF repo's ``inference/``) with its
+tilelang kernels transcribed to torch -- the numerical oracle FreeToken's port is tested against."""
diff --git a/tests/models/deepseek_v41/reference/engram.py b/tests/models/deepseek_v41/reference/engram.py
new file mode 100644
index 0000000000..625553113f
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/engram.py
@@ -0,0 +1,186 @@
+# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/engram.py), MIT License
+# (see LICENSE next to this file); unmodified.
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from sympy import isprime
+from torch import nn
+
+
+def find_next_prime(start: int, seen_primes: set[int]) -> int:
+    """The smallest prime above `start` that has not been handed out yet."""
+    candidate = start + 1
+    while not isprime(candidate) or candidate in seen_primes:
+        candidate += 1
+    return candidate
+
+
+def build_compressed_token_map(tokenizer) -> tuple[list[int], int]:
+    """Map every token id onto a smaller id space where tokens that normalize alike collapse together.
+
+    N-grams are hashed over these compressed ids, so " The", "the" and "THE" all hash the same way.
+    Returns the lookup plus the size of the compressed vocab -- and that size matters beyond bounds
+    checking, because every hash multiplier is derived from it.
+    """
+    from tokenizers import Regex, normalizers
+
+    # a private-use char, so a token that is exactly one space survives Strip() instead of
+    # collapsing to the empty string and merging with unrelated tokens
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence(
+        [
+            normalizers.NFKC(),
+            normalizers.NFD(),
+            normalizers.StripAccents(),
+            normalizers.Lowercase(),
+            normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+            normalizers.Replace(Regex(r"^ $"), sentinel),
+            normalizers.Strip(),
+            normalizers.Replace(sentinel, " "),
+        ]
+    )
+
+    # the raw Rust tokenizer, matching what training decodes with (no clean_up_tokenization_spaces)
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            # a partial UTF-8 byte token: nothing to normalize, so key it by its raw form
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+
+        new_id = key_to_new.get(key)
+        if new_id is None:
+            new_id = len(key_to_new)
+            key_to_new[key] = new_id
+        lookup[token_id] = new_id
+
+    return lookup, len(key_to_new)
+
+
+def compute_hash_multipliers(
+    layer_ids: tuple[int, ...], max_ngram_size: int, tokenizer_vocab_size: int
+) -> torch.Tensor:
+    """One multiplier per (layer, lookback), from a per-layer RNG so layers hash differently.
+
+    Kept odd, and bounded so that `token_id * multiplier` cannot overflow int64.
+    """
+    max_long = np.iinfo(np.int64).max
+    multiplier_bound = max(1, (max_long // tokenizer_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        generator = np.random.default_rng(10007 * layer_id)
+        values = generator.integers(
+            low=0,
+            high=multiplier_bound,
+            size=(max_ngram_size,),
+            dtype=np.int64,
+        )
+        rows.append(torch.tensor(values * 2 + 1))
+    return torch.stack(rows)
+
+
+@dataclass(frozen=True)
+class EngramLayout:
+    """Bucket layout of the n-gram hash tables.
+
+    A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max_ngram_size-gram), each split
+    over `n_heads` heads. Every (n-gram size, head) pair owns its own prime-sized bucket range in the
+    layer's table; the primes are drawn in order and never reused, which keeps the ranges disjoint.
+    """
+
+    max_ngram_size: int
+    layer_ids: tuple[int, ...]
+    num_embeddings: tuple[int, ...]  # table rows, per engram layer
+    primes: tuple[tuple[tuple[int, ...], ...], ...]  # [layer][n-gram size][head] bucket modulus
+    n_heads: int
+    head_dim: int
+
+    @classmethod
+    def from_args(cls, args) -> "EngramLayout | None":
+        layer_ids = tuple(args.engram_layer_ids)
+        if not layer_ids:
+            return None
+        max_ngram_size, n_heads = args.engram_max_ngram_size, args.engram_n_heads
+        primes, seen = [], set()
+        for _ in layer_ids:
+            per_ngram = []
+            for _ in range(max_ngram_size - 1):
+                sizes, current = [], args.engram_vocab_size - 1
+                for _ in range(n_heads):
+                    current = find_next_prime(current, seen)
+                    seen.add(current)
+                    sizes.append(current)
+                per_ngram.append(tuple(sizes))
+            primes.append(tuple(per_ngram))
+        return cls(
+            max_ngram_size=max_ngram_size,
+            layer_ids=layer_ids,
+            num_embeddings=tuple(args.engram_num_embeddings),
+            primes=tuple(primes),
+            n_heads=n_heads,
+            head_dim=args.engram_head_dim,
+        )
+
+
+class NgramHashState(nn.Module):
+    """Maps each position to the hash ids of the n-grams ending there.
+
+    Ids go through the compressed table, then each position is hashed with the `max_ngram_size - 1`
+    tokens before it. Look-back stops at the start of the sequence and at any dead token (an image
+    span, cached as DEAD), so an n-gram never spans one. The cache carries all of this across the
+    prefill/decode split.
+    """
+
+    DEAD = -1
+
+    def __init__(self, args, layout: EngramLayout, tokenizer):
+        super().__init__()
+        self.layout = layout
+        # every hash multiplier derives from the compressed vocab size, so a mismatch there would
+        # silently rehash the whole table
+        token_map, vocab_size = build_compressed_token_map(tokenizer)
+        assert vocab_size == args.engram_compressed_vocab_size, (vocab_size, args.engram_compressed_vocab_size)
+        self.pad_id = token_map[args.engram_pad_id]
+        flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout.primes]
+        offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat]
+        multipliers = compute_hash_multipliers(layout.layer_ids, layout.max_ngram_size, vocab_size)
+        self.register_buffer("primes", torch.tensor(layout.primes), persistent=False)
+        self.register_buffer("offsets", torch.tensor(np.array(offsets)), persistent=False)
+        self.register_buffer("multipliers", multipliers, persistent=False)
+        self.register_buffer("token_map", torch.tensor(token_map), persistent=False)
+        self.register_buffer(
+            "cache", torch.empty(args.max_batch_size, args.max_seq_len, dtype=torch.int64), persistent=False
+        )
+
+    @torch.inference_mode()
+    def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """token_mask: [B, L], False for tokens that take no part in an n-gram (image spans).
+        Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols]."""
+        batch, seqlen = input_ids.shape
+        compressed = self.token_map[input_ids]
+        if token_mask is not None:
+            compressed = torch.where(token_mask, compressed, self.DEAD)
+        self.cache[:batch, start_pos : start_pos + seqlen] = compressed
+
+        positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen)
+        tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool)
+        for shift in range(self.layout.max_ngram_size):
+            source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0))
+            blocked = blocked | (positions < shift) | (source == self.DEAD)
+            tokens.append(torch.where(blocked, self.pad_id, source))
+        tokens = torch.stack(tokens, dim=-1)  # [B, L, max_ngram_size]
+
+        # XOR the multiplied ids together one lookback at a time, so the running value after step i
+        # is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range
+        products = tokens.unsqueeze(2) * self.multipliers  # [B, L, n_engram_layers, max_ngram_size]
+        rolling, hashes = products[..., 0], []
+        for i in range(1, self.layout.max_ngram_size):
+            rolling = torch.bitwise_xor(rolling, products[..., i])
+            hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1])
+        return torch.cat(hashes, dim=-1) + self.offsets
diff --git a/tests/models/deepseek_v41/reference/kernel.py b/tests/models/deepseek_v41/reference/kernel.py
new file mode 100644
index 0000000000..fab0ae08e2
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/kernel.py
@@ -0,0 +1,113 @@
+"""Pure-torch transcriptions of the reference tilelang kernels (inference/kernel.py), for the test oracle.
+
+Same numerics as the tilelang code: fp32 math, power-of-two (ue8m0) or e4m3 scales rounded the way
+the kernels round them, the e2m1 grid with the hardware's round-to-nearest-even, and the fp8/fp4
+in-place round trips that write the dequantized value back in the input dtype. Quantized GEMMs
+dequantize independently in torch and round each projection's output to bf16.
+"""
+
+from __future__ import annotations
+
+import torch
+
+_FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
+
+
+def _pow2_ceil(x: torch.Tensor) -> torch.Tensor:
+    """``fast_round_scale``: 2 ** ceil(log2(x)) for x > 0."""
+    return torch.exp2(torch.ceil(torch.log2(x)))
+
+
+def _round_fp4(x: torch.Tensor) -> torch.Tensor:
+    """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast)."""
+    a = x.abs()
+    grid = _FP4_GRID.to(x.device)
+    hi = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7)
+    below, above = grid[hi - 1], grid[hi]
+    mid = (below + above) / 2
+    even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device))
+    r = torch.where((a > mid) | ((a == mid) & even_above), above, below)
+    r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r)
+    return torch.copysign(r, x)
+
+
+def act_quant(x: torch.Tensor, block_size: int = 128, scale_fmt=None, scale_dtype=torch.float32, inplace: bool = False):
+    """Block-wise fp8 quantization; ``inplace=True`` writes the quant+dequant round trip back into ``x``."""
+    n = x.size(-1)
+    assert n % block_size == 0
+    g = x.float().unflatten(-1, (-1, block_size))
+    amax = g.abs().amax(dim=-1).clamp_min(1e-4)
+    s = _pow2_ceil(amax / 448.0) if scale_fmt is not None else amax / 448.0
+    q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
+    if inplace:
+        x.copy_((q.float() * s.unsqueeze(-1)).flatten(-2).to(x.dtype))
+        return x
+    return q.flatten(-2), s.to(scale_dtype)
+
+
+def fp4_act_quant(x: torch.Tensor, block_size: int = 32, inplace: bool = False, scale_dtype=torch.float8_e8m0fnu):
+    """Block-wise fp4 with ue8m0 (pow2) or e4m3 scales; ``inplace=True`` writes the round trip back."""
+    n = x.size(-1)
+    assert n % block_size == 0
+    g = x.float().unflatten(-1, (-1, block_size))
+    amax = g.abs().amax(dim=-1)
+    if scale_dtype == torch.float8_e4m3fn:
+        amax = amax.clamp_min(6 * 2**-9)
+        s = (amax / 6.0).to(torch.float8_e4m3fn).float()
+    else:
+        amax = amax.clamp_min(6 * 2**-126)
+        s = _pow2_ceil(amax / 6.0)
+    q = _round_fp4((g / s.unsqueeze(-1)).clamp(-6.0, 6.0))
+    if inplace:
+        x.copy_((q * s.unsqueeze(-1)).flatten(-2).to(x.dtype))
+        return x
+    return q.flatten(-2), s.to(scale_dtype)
+
+
+def fp8_gemm(a, a_s, b, b_s, scale_dtype=torch.float32, block_size=128):
+    """``C = (A * s_a) @ (B * s_b)^T`` over dequantized operands."""
+    k = a.size(-1)
+    a_deq = (a.float().unflatten(-1, (-1, block_size)) * a_s.float().unsqueeze(-1)).flatten(-2)
+    n = b.size(0)
+    bs = b_s.float().repeat_interleave(block_size, 0)[:n].repeat_interleave(block_size, 1)[:, :k]
+    return (a_deq @ (b.float() * bs).t()).bfloat16()
+
+
+def fp4_gemm(a, a_s, b, b_s, scale_dtype=torch.float8_e8m0fnu, act_block_size=32):
+    packed = b.view(torch.uint8)
+    codes = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2).long()
+    grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6., -0., -.5, -1., -1.5, -2., -3., -4., -6.], device=b.device)
+    weight = grid[codes] * b_s.float().repeat_interleave(32, -1)
+    x = (a.float().unflatten(-1, (-1, act_block_size)) * a_s.float().unsqueeze(-1)).flatten(-2)
+    return (x @ weight.t()).bfloat16()
+
+
+def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
+    """Per (batch, query): softmax over the gathered ``topk_idxs`` (-1 = absent) plus the sink null key.
+
+    ``q [b, m, h, d]``, ``kv [b, n, d]``, ``topk_idxs [b, m, topk]`` -> ``o [b, m, h, d]`` in q's dtype.
+    A row with no valid index yields zeros (the finite ``-1e30`` running max of the kernel)."""
+    b, m, h, d = q.shape
+    valid = topk_idxs >= 0
+    gathered = kv[torch.arange(b, device=q.device)[:, None, None], topk_idxs.clamp_min(0).long()]  # [b, m, topk, d]
+    scores = torch.einsum("bmhd,bmtd->bmht", q.float(), gathered.float()) * softmax_scale
+    scores = scores.masked_fill(~valid[:, :, None, :], float("-inf"))
+    mx = scores.amax(dim=-1).clamp_min(-1e30)
+    p = torch.exp(scores - mx.unsqueeze(-1))
+    denom = p.sum(-1) + torch.exp(attn_sink.float()[None, None, :] - mx)
+    o = torch.einsum("bmht,bmtd->bmhd", p, gathered.float()) / denom.unsqueeze(-1)
+    return o.to(q.dtype)
+
+
+def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6):
+    hc = hc_mult
+    mixes, sc, base = mixes.float(), hc_scale.float(), hc_base.float()
+    pre = torch.sigmoid(mixes[..., :hc] * sc[0] + base[:hc]) + eps
+    post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc])
+    comb = (mixes[..., 2 * hc :] * sc[2] + base[2 * hc :]).unflatten(-1, (hc, hc))
+    comb = comb.softmax(dim=-1) + eps
+    comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
+    for _ in range(sinkhorn_iters - 1):
+        comb = comb / (comb.sum(dim=-1, keepdim=True) + eps)
+        comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
+    return pre, post, comb
diff --git a/tests/models/deepseek_v41/reference/model.py b/tests/models/deepseek_v41/reference/model.py
new file mode 100644
index 0000000000..593c7518c4
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/model.py
@@ -0,0 +1,1458 @@
+# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/model.py), MIT License
+# (see LICENSE next to this file). FreeToken keeps this copy as the numerical ORACLE its port is tested
+# against. Edits, all marked "FreeToken:", are limited to (1) local imports with the tilelang kernels
+# replaced by the pure-torch transcriptions in kernel.py, and (2) forward_bounded(), a Decoder SWA
+# Bounded Replay oracle built from the unmodified layer math.
+import math
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Literal
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+# FreeToken: local imports.
+from .engram import EngramLayout, NgramHashState
+from .kernel import (
+    act_quant,
+    fp4_act_quant,
+    fp4_gemm,
+    fp8_gemm,
+    hc_split_sinkhorn,
+    sparse_attn,
+)
+
+from .vision import ViT, Aligner
+
+IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4)  # FreeToken: image_processor constants
+
+# Set once by Transformer.__init__; one model per process, so layers just read them.
+world_size = 1
+rank = 0
+default_dtype = torch.float8_e4m3fn  # storage dtype for Linear weights, from ModelArgs.dtype
+fp8_block_size = 32  # one fp8 scale per 32x32 weight block / 32 activations
+fp4_block_size = 32  # one fp4 scale per 32 elements along K
+scale_fmt = "ue8m0"
+scale_dtype = torch.float8_e8m0fnu
+
+
+@contextmanager
+def set_dtype(dtype):
+    """Temporarily override torch's default dtype, restoring it even if the body raises."""
+    prev = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(prev)
+
+
+@dataclass
+class ModelArgs:
+    """Field names are exactly the config JSON keys. The defaults are a small model that
+    `python model.py` can run, not the released shapes -- though the scale-independent
+    values (norm_eps, score_func, hc_*, engram_*) do match it."""
+
+    # runtime limits rather than model shape: they size the KV caches
+    max_batch_size: int = 4
+    max_seq_len: int = 4096
+    temperature: float = 1
+    dtype: Literal["bf16", "fp8"] = "fp8"
+    expert_dtype: Literal["fp4"] | None = "fp4"
+    vocab_size: int = 129280
+    dim: int = 1024
+    moe_inter_dim: int = 1024
+    n_layers: int = 5
+    n_mtp_layers: int = 1  # extra draft layers appended after the backbone, indices n_layers..
+    n_heads: int = 16
+    # moe
+    n_routed_experts: int = 8
+    n_shared_experts: int = 1
+    n_activated_experts: int = 2
+    score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus"
+    gate_temp: float = 1.0
+    norm_topk_prob: bool = True
+    route_scale: float = 1.0
+    swiglu_limit: float = 0.0
+    # attention: latent q/kv projections, plus a LoRA-factorised output projection over o_groups
+    q_lora_rank: int = 256
+    head_dim: int = 128
+    rope_head_dim: int = 32
+    norm_eps: float = 1e-20
+    o_groups: int = 8
+    o_lora_rank: int = 256
+    # sparse attention: every layer attends over a sliding window, and may add compressed KV on top
+    window_size: int = 128
+    # one entry per layer, MTP layers included: 0 = sliding window only, r = KV compressed r-to-1
+    compress_ratios: tuple[int, ...] = (0, 2, 2, 1, 1, 0)
+    # layers sharing a ratio also share one compressed KV and one indexer, produced by the first
+    kv_source_layers: tuple[int, ...] = (1, 3)
+    index_source_layers: tuple[int, ...] = (1, 3)
+    # rope, with YaRN extrapolation when original_seq_len > 0. Compressed KV rotates at its own
+    # theta because one latent stands for compress_ratio tokens, so its positions are further apart.
+    compress_rope_theta: float = 40000.0
+    original_seq_len: int = 0
+    rope_theta: float = 10000.0
+    rope_factor: float = 40
+    beta_fast: int = 32
+    beta_slow: int = 1
+    # the indexer: a small extra attention that scores compressed positions, so each query can keep
+    # just `index_topk` of them. Names match DeepSeek-V3.2-Exp, where this mechanism first appeared.
+    index_n_heads: int = 16
+    index_head_dim: int = 64
+    index_topk: int = 64
+    # candidate pre-filtering: candidate_source_layer < 0 turns it off and the other two are unused
+    candidate_source_layer: int = -1
+    candidate_topk_blocks: int = 0
+    candidate_block_size: int = 0
+    # hyper-connections: the residual stream is carried as hc_mult parallel copies
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1e-6
+    # engram: n-gram hash lookups added into the residual stream at a few layers
+    engram_layer_ids: tuple[int, ...] = ()
+    engram_num_embeddings: tuple[int, ...] = ()  # unpadded table rows; each rank allocates ceil(rows / world_size)
+    engram_max_ngram_size: int = 1
+    engram_vocab_size: int = 0  # bucket size each (n-gram size, head) starts searching primes from
+    engram_n_heads: int = 0
+    engram_head_dim: int = 0
+    engram_pad_id: int = 2  # token that fills n-gram slots with no history; matches training
+    # size of the compressed tokenizer vocab; every hash multiplier is derived from it
+    engram_compressed_vocab_size: int = 0
+    # vision (VL); vision_n_layers == 0 disables the vision path
+    vision_n_layers: int = 0
+    vision_dim: int = 1024
+    vision_n_heads: int = 16
+    vision_inter_dim: int = 2816
+    vision_patch_size: int = 14
+    vision_rope_theta: float = 10000.0
+    vision_downsample_ratio: int = 3
+    vision_max_n_token: int = 1024
+    vision_min_pixels: int = 544 * 544
+    vision_max_wh_ratio: int | None = None
+    # raw id of <｜deepseek_image｜>; every position of an image span carries this id in input_ids
+    image_token_id: int = 129264
+    # dspark draft head. Only the forward pass is implemented here -- nothing calls forward_spec,
+    # so these are read but the speculative-decoding loop itself is out of scope for this repo.
+    dspark_block_size: int = 0
+    dspark_noise_token_id: int = 0
+    dspark_target_layer_ids: tuple[int, ...] = ()
+    dspark_markov_rank: int = 256
+    dspark_n_routed_experts: int = 0
+    dspark_n_activated_experts: int = 0
+
+    @property
+    def vision_enabled(self) -> bool:
+        return self.vision_n_layers > 0
+
+    def get_moe_config(self, layer_id: int) -> tuple[int, int]:
+        """Return the routed/activated expert counts for a given layer."""
+        if layer_id < self.n_layers:
+            return self.n_routed_experts, self.n_activated_experts
+        return (
+            self.dspark_n_routed_experts or self.n_routed_experts,
+            self.dspark_n_activated_experts or self.n_activated_experts,
+        )
+
+
+class ParallelEmbedding(nn.Module):
+    """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows.
+    Out-of-range indices are zero-masked before all_reduce to combine partial embeddings."""
+
+    def __init__(self, vocab_size: int, dim: int):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        assert vocab_size % world_size == 0, (
+            f"Vocabulary size must be divisible by world size (world_size={world_size})"
+        )
+        self.part_vocab_size = vocab_size // world_size
+        self.vocab_start_idx = rank * self.part_vocab_size
+        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
+        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if world_size > 1:
+            # ids off this rank read row 0 then get zeroed, so the all_reduce sums one real row
+            mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
+            x = x - self.vocab_start_idx
+            x[mask] = 0
+        y = F.embedding(x, self.weight)
+        if world_size > 1:
+            y[mask] = 0
+            dist.all_reduce(y)
+        return y
+
+
+def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+    """Pick a GEMM from the weight dtype. Quantized weights need a quantized activation, and both
+    fp4 and fp8 weights take an fp8 one -- for fp4 the kernel handles the mixed precision."""
+    assert bias is None
+
+    if weight.dtype == torch.float4_e2m1fn_x2:
+        x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype)
+        return fp4_gemm(
+            x,
+            s,
+            weight,
+            weight.scale,
+            scale_dtype,
+            act_block_size=fp8_block_size,
+        )
+    elif weight.dtype == torch.float8_e4m3fn:
+        x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype)
+        return fp8_gemm(
+            x,
+            s,
+            weight,
+            weight.scale,
+            scale_dtype,
+            block_size=fp8_block_size,
+        )
+    else:
+        return F.linear(x, weight)
+
+
+class Linear(nn.Module):
+    """bf16, fp8 or fp4 weights. Quantized ones get a `scale`, also attached to `.weight` so that
+    `linear()` can reach it."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        dtype = dtype or default_dtype
+        if dtype == torch.float4_e2m1fn_x2:
+            # two values per byte: [out, in] logically, [out, in//2] stored
+            self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2))
+            self.weight.scale = self.scale = nn.Parameter(
+                torch.empty(out_features, in_features // fp4_block_size, dtype=torch.float8_e8m0fnu)
+            )
+        elif dtype == torch.float8_e4m3fn:
+            self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+            self.weight.scale = self.scale = nn.Parameter(
+                torch.empty(
+                    (out_features + fp8_block_size - 1) // fp8_block_size,
+                    (in_features + fp8_block_size - 1) // fp8_block_size,
+                    dtype=torch.float8_e8m0fnu,
+                )
+            )
+        else:
+            self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+            self.register_parameter("scale", None)
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return linear(x, self.weight, self.bias)
+
+
+class ColumnParallelLinear(Linear):
+    """Splits the output dim across ranks; each rank's slice of the output is already complete."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        assert out_features % world_size == 0, (
+            f"Output features must be divisible by world size (world_size={world_size})"
+        )
+        self.part_out_features = out_features // world_size
+        super().__init__(in_features, self.part_out_features, bias, dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return linear(x, self.weight, self.bias)
+
+
+class RowParallelLinear(Linear):
+    """Splits the reduction dim, so each rank holds a partial sum: hence the fp32 all_reduce, with
+    the bias added only after it."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        assert in_features % world_size == 0, (
+            f"Input features must be divisible by world size (world_size={world_size})"
+        )
+        self.part_in_features = in_features // world_size
+        super().__init__(self.part_in_features, out_features, bias, dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = linear(x, self.weight, None)
+        if world_size > 1:
+            y = y.float()
+            dist.all_reduce(y)
+        if self.bias is not None:
+            y += self.bias
+        return y.type_as(x)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor):
+        dtype = x.dtype
+        x = x.float()
+        var = x.square().mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class ParallelEngramEmbedding(nn.Module):
+    """The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup."""
+
+    def __init__(self, num_embeddings: int, dim: int):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.dim = dim
+        self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size
+        self.vocab_start_idx = rank * self.part_num_embeddings
+        self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings
+
+        self.block_size = fp8_block_size
+        # the table stays fp8 as stored: rows are dequantized with `scale` on lookup
+        self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn))
+        self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype))
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx)
+        local_indices = indices - self.vocab_start_idx
+        local_indices = local_indices.masked_fill(mask, 0)
+
+        values = F.embedding(local_indices, self.weight)
+        scales = F.embedding(local_indices, self.scale)
+        values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1)
+        values = values.flatten(-2).to(torch.bfloat16)
+        values = values.masked_fill(mask.unsqueeze(-1), 0)
+
+        if world_size > 1:
+            dist.all_reduce(values)
+        return values
+
+
+class Engram(nn.Module):
+    """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream.
+
+    The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared
+    value. The gate is a normalized dot product of stream against key.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout):
+        super().__init__()
+        self.layer_id = layer_id
+        self.layer_hash_index = layout.layer_ids.index(layer_id)
+        self.dim = args.dim
+        self.hc_mult = args.hc_mult
+        self.clamp_value = 1e-6
+
+        self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim)
+        n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
+        self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1))
+        self.eps = args.norm_eps
+        self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
+        self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
+
+    def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts
+        the gate so those positions pass through untouched."""
+        kv = self.wkv(self.embed(hash_ids).flatten(-2))
+        key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
+        key = key.float().unflatten(-1, (self.hc_mult, self.dim))
+        weight = self.q_weight.float() * self.k_weight.float()  # only ever used as a product
+        h, eps = x.float(), self.eps
+        # normalized per (token, hc copy) over `dim`, NOT jointly over the copies
+        rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps)
+        dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5
+        # signed sqrt before the sigmoid, matching the training kernel
+        gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot))
+        if token_mask is not None:
+            gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0)
+        return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)
+
+
+@lru_cache(2)
+def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor:
+    """Rotary frequencies as complex exponentials, one row per position.
+
+    With original_seq_len > 0 this applies YaRN: dimensions whose wavelength already fits inside the
+    training context keep their frequency, those far beyond it are divided by `factor`, and the
+    `beta_fast`..`beta_slow` band in between is faded across with a linear ramp.
+    """
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    if original_seq_len > 0:
+        # the dim whose wavelength completes `rotations` turns over the training context
+        def corrected_dim(rotations):
+            return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base))
+
+        low = max(math.floor(corrected_dim(beta_fast)), 0)
+        high = min(math.ceil(corrected_dim(beta_slow)), dim - 1)
+        ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1)
+        smooth = 1 - ramp
+        freqs = freqs / factor * (1 - smooth) + freqs * smooth
+
+    freqs = torch.outer(torch.arange(seqlen), freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor:
+    """Rotate `x` in place, taking adjacent element pairs as complex numbers. Accepts [b, s, d] and
+    [b, s, h, d]; `inverse` conjugates the rotation, which is how the attention output gets the
+    query's rotation removed again so the cache can stay in one shared rotated form."""
+    y = x
+    x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs_cis = freqs_cis.conj()
+    if x.ndim == 3:
+        freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1))
+    else:
+        freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+    x = torch.view_as_real(x * freqs_cis).flatten(-2)
+    y.copy_(x)
+    return y
+
+
+@lru_cache(1)
+def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int):
+    """Which sliding-window cache slots each query attends to; -1 marks a slot holding nothing.
+
+    The cache is a ring of `window_size` slots. Prefill needs one row per query, each seeing its own
+    causal window. A decode step has a single query that sees the whole ring, listed oldest first.
+    Order within a row does not matter to `sparse_attn`, which handles every slot independently.
+    """
+    if start_pos == 0:
+        end = torch.arange(seqlen).unsqueeze(1)
+        idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size))
+        idxs = torch.where(idxs > end, -1, idxs)  # before the sequence started
+    else:
+        oldest = start_pos % window_size + 1
+        idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)])
+        idxs = torch.where(idxs > start_pos, -1, idxs)  # ring still filling
+    # sparse_attn needs real [b, m, topk] int32 memory, hence the materializing expand
+    return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous()
+
+
+class Compressor(nn.Module):
+    """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate.
+
+    Returns the latent before RoPE, or None while a group is still filling up -- so during decode it
+    only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`.
+    Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int):
+        super().__init__()
+        compress_ratio = args.compress_ratios[layer_id]
+        head_dim = args.head_dim
+        self.compress_ratio = compress_ratio
+        self.head_dim = head_dim
+        self.norm = RMSNorm(head_dim, args.norm_eps)
+        # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling
+        # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match
+        self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16)
+        if compress_ratio > 1:
+            self.wgate = Linear(args.dim, head_dim, dtype=torch.float32)
+            # tail of an incomplete group, carried across decode steps
+            self.kv_state: torch.Tensor
+            self.score_state: torch.Tensor
+            state_shape = (args.max_batch_size, compress_ratio, head_dim)
+            self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False)
+            self.register_buffer(
+                "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False
+            )
+
+    def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None:
+        bsz, seqlen, _ = x.size()
+        ratio, dtype = self.compress_ratio, x.dtype
+        if ratio == 1:  # one token per group: nothing to pool, so no gate and no fp32
+            return self.norm(self.wkv(x))
+
+        x = x.float()
+        kv, score = self.wkv(x), self.wgate(x)
+        if start_pos == 0:
+            should_compress = seqlen >= ratio
+            remainder = seqlen % ratio
+            cutoff = seqlen - remainder
+            if remainder:  # trailing partial group waits in the state
+                kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1)
+                score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1)
+            kv = kv.unflatten(1, (-1, ratio))
+            score = score.unflatten(1, (-1, ratio))
+            kv = (kv * score.softmax(dim=2)).sum(dim=2)
+        else:  # one token per step: fill a slot, and pool only when the group just completed
+            should_compress = (start_pos + 1) % ratio == 0
+            slot = start_pos % ratio
+            self.kv_state[:bsz, slot] = kv.squeeze(1)
+            self.score_state[:bsz, slot] = score.squeeze(1)
+            if should_compress:
+                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
+        if not should_compress:
+            return None
+        return self.norm(kv.to(dtype))
+
+
+class Indexer(torch.nn.Module):
+    """Keeps the `index_topk` best compressed positions per query.
+
+    A small side attention: fp4 query heads against one shared key per compressed position, scores
+    rectified then combined by `weights_proj`. With a candidate source this is the second of two
+    levels; `select_candidate_blocks` is the first.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int):
+        super().__init__()
+        # the index keys are derived from the compressor's latent, so only a layer that compresses
+        # its own KV can produce them; every other indexer reads them from that layer's cache
+        self.owns_k = layer_id in args.kv_source_layers
+        self.compress_ratio = args.compress_ratios[layer_id]
+        self.is_candidate_source = layer_id == args.candidate_source_layer
+        self.uses_candidates = 0 <= args.candidate_source_layer < layer_id
+        self.candidate_topk_blocks = args.candidate_topk_blocks
+        self.candidate_block_size = args.candidate_block_size
+        self.dim = args.dim
+        self.n_heads = args.index_n_heads
+        self.n_local_heads = args.index_n_heads // world_size
+        self.index_head_dim = args.index_head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.index_topk = args.index_topk
+        self.q_lora_rank = args.q_lora_rank
+        self.softmax_scale = self.index_head_dim**-0.5
+        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.index_head_dim)
+        self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)
+        self.freqs_cis: torch.Tensor | None = None
+        if self.owns_k:
+            self.wk = Linear(args.head_dim, self.index_head_dim, dtype=torch.bfloat16)
+            self.k_norm = RMSNorm(self.index_head_dim, args.norm_eps)
+            self.k_cache: torch.Tensor
+            self.register_buffer(
+                "k_cache",
+                torch.zeros(args.max_batch_size, args.max_seq_len // self.compress_ratio, args.index_head_dim),
+                persistent=False,
+            )
+
+    def forward(self, x: torch.Tensor, qr: torch.Tensor, latent: torch.Tensor, start_pos: int, offset: int):
+        """`latent` is this layer's RoPE-free compressed latent, None when this layer does not
+        compress or when its current group is still incomplete. An index-key owner turns it into
+        index keys here, which has to happen before Attention overwrites that same storage with
+        the RoPE'd, quantized values."""
+        assert self.freqs_cis is not None
+        bsz, seqlen, _ = x.size()
+        ratio, rd, end_pos = self.compress_ratio, self.rope_head_dim, start_pos + seqlen
+
+        # latent is None while a group is still filling up, so there is nothing to publish yet
+        if self.owns_k and latent is not None:
+            # a latent stands for the first token of its group, so group j takes position j * ratio
+            freqs = (
+                self.freqs_cis[: seqlen - seqlen % ratio : ratio]
+                if start_pos == 0
+                else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0)
+            )
+            k = self.k_norm(self.wk(latent))
+            apply_rotary_emb(k[..., -rd:], freqs)
+            fp4_act_quant(k, fp4_block_size, True)
+            self.k_cache[:bsz, start_pos // ratio : start_pos // ratio + k.size(1)] = k
+            shared_attn.index_k = self.k_cache
+
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim))
+        apply_rotary_emb(q[..., -rd:], self.freqs_cis[start_pos:end_pos])
+        fp4_act_quant(q, fp4_block_size, True)
+
+        index_k = shared_attn.index_k[:bsz, : end_pos // ratio]
+        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5)
+        index_score = torch.einsum("bshd,btd->bsht", q, index_k)
+        index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
+        if world_size > 1:
+            dist.all_reduce(index_score)
+
+        # how many compressed positions each query can see: a block becomes visible once the query
+        # has passed its last token. One query per decode step, so there it is just a number.
+        if start_pos == 0:
+            compress_lens = (torch.arange(1, seqlen + 1, device=x.device) // ratio).unsqueeze(-1)
+            index_score.masked_fill_(torch.arange(seqlen // ratio, device=x.device) >= compress_lens, -torch.inf)
+        else:
+            compress_lens = end_pos // ratio
+
+        if self.is_candidate_source:
+            shared_attn.candidates = select_candidate_blocks(
+                index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size
+            )
+        elif self.uses_candidates:
+            # level two: score with our own weights, but only inside the source's candidate blocks
+            index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf)
+
+        # top-k by score, re-sorted into position order; unreachable -> -1, rest shifted by offset
+        topk = min(self.index_topk, end_pos // ratio)
+        idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values
+        return torch.where(idxs < compress_lens, idxs + offset, -1).int()
+
+    # FreeToken: Decoder SWA Bounded Replay oracle -- the indexer for a TAIL prefill: queries are the
+    # positions [s, s + n) of a prompt of L_all tokens; the keys (a source's) were published for all
+    # L_all tokens by forward()'s start_pos == 0 path. Same scoring / masking / top-k as forward().
+    def forward_tail(self, x, qr, s: int, n: int, offset: int):
+        assert self.freqs_cis is not None
+        bsz = x.size(0)
+        ratio, rd = self.compress_ratio, self.rope_head_dim
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim))
+        apply_rotary_emb(q[..., -rd:], self.freqs_cis[s : s + n])
+        fp4_act_quant(q, fp4_block_size, True)
+        end_pos = s + n
+        index_k = shared_attn.index_k[:bsz, : end_pos // ratio]
+        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5)
+        index_score = torch.einsum("bshd,btd->bsht", q, index_k)
+        index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
+        compress_lens = ((s + torch.arange(1, n + 1, device=x.device)) // ratio).unsqueeze(-1)
+        index_score.masked_fill_(torch.arange(end_pos // ratio, device=x.device) >= compress_lens, -torch.inf)
+        if self.is_candidate_source:
+            shared_attn.candidates = select_candidate_blocks(
+                index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size
+            )
+        elif self.uses_candidates:
+            index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf)
+        topk = min(self.index_topk, end_pos // ratio)
+        idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values
+        return torch.where(idxs < compress_lens, idxs + offset, -1).int()
+
+
+def select_candidate_blocks(
+    logits: torch.Tensor,
+    compress_lens: torch.Tensor | int,
+    topk_blocks: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Level one of the two-level top-k: keep the `topk_blocks` highest-scoring blocks per query.
+
+    `logits` is [..., n_positions] with positions the query cannot reach already at -inf, which is
+    what makes a block score of -inf mean "not reachable yet". `compress_lens` is a plain int during
+    decode, or broadcasts against logits' leading dims during prefill. Returns a bool mask shaped
+    like `logits`, so the layers consuming it just mask and never think about blocks again.
+    """
+    width = logits.size(-1)
+    # score each block by its best position; -inf pads the last one out to block_size
+    scores = F.pad(logits, (0, -width % block_size), value=-torch.inf)
+    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
+    num_blocks = scores.size(-1)
+
+    # the block with this query's newest position is only partly filled, so pin it in: it holds the
+    # most recent tokens but could otherwise be outscored by an older, full block
+    last = (compress_lens - 1) // block_size
+    scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf)
+
+    top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
+    # fewer reachable blocks than topk_blocks means leftover picks came back -inf: drop them
+    keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf)
+    return keep.repeat_interleave(block_size, dim=-1)[..., :width]
+
+
+class Attention(nn.Module):
+    """Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a
+    sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions
+    reaching further back. Q and the output projection are both low-rank, the latter grouped.
+
+    compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do,
+    the rest read that same cache.
+    """
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        self.layer_id = layer_id
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.n_local_heads = args.n_heads // world_size
+        self.q_lora_rank = args.q_lora_rank
+        self.o_lora_rank = args.o_lora_rank
+        self.head_dim = args.head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.nope_head_dim = args.head_dim - args.rope_head_dim
+        self.n_groups = args.o_groups
+        self.n_local_groups = self.n_groups // world_size
+        self.window_size = args.window_size
+        self.compress_ratio = args.compress_ratios[layer_id]
+        self.eps = args.norm_eps
+
+        self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
+        self.wq_a = Linear(self.dim, self.q_lora_rank)
+        self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
+        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
+        self.wkv = Linear(self.dim, self.head_dim)
+        self.kv_norm = RMSNorm(self.head_dim, self.eps)
+        self.wo_a = ColumnParallelLinear(
+            self.n_heads * self.head_dim // self.n_groups,
+            self.n_groups * args.o_lora_rank,
+            dtype=torch.bfloat16,
+        )
+        self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
+        self.softmax_scale = self.head_dim**-0.5
+
+        is_backbone = layer_id < args.n_layers
+        self.is_kv_source = is_backbone and layer_id in args.kv_source_layers
+        self.is_index_source = is_backbone and layer_id in args.index_source_layers
+        self.compressor: Compressor | None = None
+        self.indexer: Indexer | None = None
+        if self.is_kv_source:
+            self.compressor = Compressor(args, layer_id)
+        if self.is_index_source:
+            self.indexer = Indexer(args, layer_id)
+
+        self.window_kv_cache: torch.Tensor
+        self.register_buffer(
+            "window_kv_cache",
+            torch.zeros(args.max_batch_size, args.window_size, self.head_dim),
+            persistent=False,
+        )
+        if self.is_kv_source:
+            self.compress_kv_cache: torch.Tensor
+            self.register_buffer(
+                "compress_kv_cache",
+                torch.zeros(
+                    args.max_batch_size,
+                    args.max_seq_len // self.compress_ratio,
+                    self.head_dim,
+                ),
+                persistent=False,
+            )
+        if self.compress_ratio:
+            original_seq_len, rope_theta = (
+                args.original_seq_len,
+                args.compress_rope_theta,
+            )
+        else:
+            # disable YaRN and use base rope_theta in pure sliding-window attention
+            original_seq_len, rope_theta = 0, args.rope_theta
+        freqs_cis = precompute_freqs_cis(
+            self.rope_head_dim,
+            args.max_seq_len,
+            original_seq_len,
+            rope_theta,
+            args.rope_factor,
+            args.beta_fast,
+            args.beta_slow,
+        )
+        self.freqs_cis: torch.Tensor
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+    def _window_kv(self, x, freqs_cis, start_pos):
+        """This layer's sliding-window K and the window positions every query may attend to. The K
+        stays fp8, quantized over the whole post-RoPE vector, RoPE tail included."""
+        bsz, seqlen, _ = x.size()
+        win = self.window_size
+        kv = self.kv_norm(self.wkv(x))
+        apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis)
+        act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True)
+        if start_pos == 0:  # prefill: attend over this chunk, seeding the ring buffer for decode
+            if seqlen <= win:
+                self.window_kv_cache[:bsz, :seqlen] = kv
+            else:
+                cutoff = seqlen % win
+                self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = kv[:, -win:].split(
+                    [win - cutoff, cutoff], dim=1
+                )
+            window_kv = kv
+        else:  # decode: one token into the ring buffer, attend over the whole window
+            self.window_kv_cache[:bsz, start_pos % win] = kv.squeeze(1)
+            window_kv = self.window_kv_cache[:bsz]
+        return window_kv, get_window_topk_idxs(win, bsz, seqlen, start_pos)
+
+    def _compress_topk_idxs(self, x, qr, latent, start_pos, offset, compress_len):
+        """Which compressed positions each query attends to. Index sources run their own indexer;
+        the layers in between reuse the result their source published."""
+        if not self.is_index_source:
+            return shared_attn.topk_idxs
+
+        bsz, seqlen, _ = x.size()
+        if compress_len == 0:
+            idxs = torch.empty(bsz, seqlen, 0, dtype=torch.int32, device=x.device)
+        else:
+            assert self.indexer is not None
+            if self.indexer.freqs_cis is None:
+                self.indexer.freqs_cis = self.freqs_cis
+            idxs = self.indexer(x, qr, latent, start_pos, offset)
+        shared_attn.topk_idxs = idxs
+        return idxs
+
+    def _compress_kv(self, x, qr, start_pos, offset):
+        """The shared compressed KV and the compressed positions every query may attend to. This
+        layer compresses its own KV only when it is a source; otherwise it just reads the cache."""
+        bsz, seqlen, _ = x.size()
+        ratio = self.compress_ratio
+        compress_len = (start_pos + seqlen) // ratio
+        latent = None
+        if self.is_kv_source:
+            latent = self.compressor(x, start_pos)
+            shared_attn.compress_kv = self.compress_kv_cache
+        # the indexer needs the latent before RoPE, so it runs before the cache is written
+        idxs = self._compress_topk_idxs(x, qr, latent, start_pos, offset, compress_len)
+        if latent is not None:
+            # a latent stands for the first token of its group, so group j takes position j * ratio
+            freqs = (
+                self.freqs_cis[: seqlen - seqlen % ratio : ratio]
+                if start_pos == 0
+                else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0)
+            )
+            apply_rotary_emb(latent[..., -self.rope_head_dim :], freqs)
+            # Compressed KV uses groups of 16 with E4M3 scales; the indexer uses 32 with E8M0.
+            fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn)
+            self.compress_kv_cache[:bsz, start_pos // ratio : start_pos // ratio + latent.size(1)] = latent
+        # read after the write, so this does not depend on the slice aliasing the cache
+        return shared_attn.compress_kv[:bsz, :compress_len], idxs
+
+    def forward(self, x: torch.Tensor, start_pos: int):
+        bsz, seqlen, _ = x.size()
+        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+        rd = self.rope_head_dim
+
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
+        apply_rotary_emb(q[..., -rd:], freqs_cis)
+
+        kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos)
+        if self.compress_ratio:
+            compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1))
+            kv = torch.cat([kv, compress_kv], dim=1)
+            topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1)
+
+        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+        apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+
+        # wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not
+        # Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory.
+        o = o.view(bsz, seqlen, self.n_local_groups, -1)
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+        x = self.wo_b(o.flatten(2))
+        return x
+
+    # FreeToken: Decoder SWA Bounded Replay oracle. Prefill of the prompt TAIL [s, s + n) as its own
+    # window segment (a query sees SWA keys in [max(s, i - win + 1), i]) while the global branch
+    # covers the whole prompt: a kv source compresses ``x_all`` (every token, forward()'s start_pos
+    # == 0 path), the indexer scores the tail queries at their absolute positions.
+    def forward_tail(self, x: torch.Tensor, s: int, x_all: torch.Tensor | None = None):
+        bsz, n, _ = x.size()
+        win, rd, ratio = self.window_size, self.rope_head_dim, self.compress_ratio
+        freqs_cis = self.freqs_cis[s : s + n]
+
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
+        apply_rotary_emb(q[..., -rd:], freqs_cis)
+
+        kv = self.kv_norm(self.wkv(x))
+        apply_rotary_emb(kv[..., -rd:], freqs_cis)
+        act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True)
+        # seed the decode ring with the tail's keys at their absolute slots
+        for i in range(max(0, n - win), n):
+            self.window_kv_cache[:bsz, (s + i) % win] = kv[:, i]
+        topk_idxs = get_window_topk_idxs(win, bsz, n, 0).to(x.device)
+
+        if ratio:
+            end = s + n
+            compress_len = end // ratio
+            latent = None
+            if self.is_kv_source:
+                assert x_all is not None, "the kv source compresses every prompt token"
+                latent = self.compressor(x_all, 0)
+                shared_attn.compress_kv = self.compress_kv_cache
+                if self.indexer is not None:
+                    if self.indexer.freqs_cis is None:
+                        self.indexer.freqs_cis = self.freqs_cis
+                    # publish the index keys for every group (forward()'s start_pos == 0 path)
+                    L = x_all.size(1)
+                    k = self.indexer.k_norm(self.indexer.wk(latent))
+                    apply_rotary_emb(k[..., -rd:], self.freqs_cis[: L - L % ratio : ratio])
+                    fp4_act_quant(k, fp4_block_size, True)
+                    self.indexer.k_cache[:bsz, : k.size(1)] = k
+                    shared_attn.index_k = self.indexer.k_cache
+            if not self.is_index_source:
+                idxs = shared_attn.topk_idxs
+            elif compress_len == 0:
+                idxs = torch.empty(bsz, n, 0, dtype=torch.int32, device=x.device)
+            else:
+                if self.indexer.freqs_cis is None:
+                    self.indexer.freqs_cis = self.freqs_cis
+                idxs = self.indexer.forward_tail(x, qr, s, n, kv.size(1))
+                shared_attn.topk_idxs = idxs
+            if latent is not None:
+                L = x_all.size(1)
+                apply_rotary_emb(latent[..., -rd:], self.freqs_cis[: L - L % ratio : ratio])
+                fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn)
+                self.compress_kv_cache[:bsz, : latent.size(1)] = latent
+            kv = torch.cat([kv, shared_attn.compress_kv[:bsz, :compress_len]], dim=1)
+            topk_idxs = torch.cat([topk_idxs, idxs], dim=-1)
+
+        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+        apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+        o = o.view(bsz, n, self.n_local_groups, -1)
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+        return self.wo_b(o.flatten(2))
+
+
+class Gate(nn.Module):
+    """MoE gating. The correction bias steers expert selection only; the routing weights come from the
+    unbiased scores. Image-span tokens use a separate bias (training `noaux_tc_for_vl`)."""
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        n_routed_experts, n_activated_experts = args.get_moe_config(layer_id)
+        self.dim = args.dim
+        self.topk = n_activated_experts
+        self.score_func = args.score_func
+        self.gate_temp = args.gate_temp
+        self.norm_topk_prob = args.norm_topk_prob
+        self.route_scale = args.route_scale
+        self.weight = nn.Parameter(torch.empty(n_routed_experts, args.dim))
+        self.bias = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32))
+        self.bias_vl = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) if args.vision_enabled else None
+
+    def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """x: [n, dim]; image_mask: [n] bool, True for tokens inside an image span."""
+        scores = linear(x.float(), self.weight.float()) / self.gate_temp
+        if self.score_func == "softmax":
+            scores = scores.softmax(dim=-1)
+        elif self.score_func == "sigmoid":
+            scores = scores.sigmoid()
+        else:
+            scores = F.softplus(scores).sqrt()
+        bias = self.bias
+        if image_mask is not None and self.bias_vl is not None:
+            bias = torch.where(image_mask.unsqueeze(-1), self.bias_vl, bias)
+        # the bias picks experts but does not scale them: weights come from the raw scores
+        indices = (scores + bias).topk(self.topk, dim=-1)[1]
+        weights = scores.gather(1, indices)
+        if self.norm_topk_prob and self.topk > 1:
+            weights /= weights.sum(dim=-1, keepdim=True) + 1e-20  # not norm_eps, matches training
+        weights *= self.route_scale
+        return weights, indices
+
+
+class Expert(nn.Module):
+    """One SwiGLU FFN. The clamps come straight from training, where they keep fp8/fp4 activations in
+    range: the up branch is clamped on both sides, the gate branch only from above."""
+
+    def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0.0):
+        super().__init__()
+        self.w1 = Linear(dim, inter_dim, dtype=dtype)
+        self.w2 = Linear(inter_dim, dim, dtype=dtype)
+        self.w3 = Linear(dim, inter_dim, dtype=dtype)
+        self.swiglu_limit = swiglu_limit
+
+    def forward(self, x: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor:
+        dtype = x.dtype
+        gate = self.w1(x).float()
+        up = self.w3(x).float()
+        if self.swiglu_limit > 0:
+            up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit)
+            gate = torch.clamp(gate, max=self.swiglu_limit)
+        x = F.silu(gate) * up
+        if weights is not None:
+            x = weights * x
+        return self.w2(x.to(dtype))
+
+
+class MoE(nn.Module):
+    """Top-k routed experts plus one shared expert every token goes through. Experts are split
+    across ranks, so `self.experts` is None for those another rank owns."""
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        n_routed_experts, n_activated_experts = args.get_moe_config(layer_id)
+        self.layer_id = layer_id
+        self.dim = args.dim
+        assert n_routed_experts % world_size == 0, (
+            f"Number of experts must be divisible by world size (world_size={world_size})"
+        )
+        self.n_routed_experts = n_routed_experts
+        self.n_local_experts = n_routed_experts // world_size
+        self.n_activated_experts = n_activated_experts
+        self.experts_start_idx = rank * self.n_local_experts
+        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
+        self.gate = Gate(layer_id, args)
+        expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None
+        self.experts = nn.ModuleList(
+            [
+                Expert(
+                    args.dim,
+                    args.moe_inter_dim,
+                    dtype=expert_dtype,
+                    swiglu_limit=args.swiglu_limit,
+                )
+                if self.experts_start_idx <= i < self.experts_end_idx
+                else None
+                for i in range(self.n_routed_experts)
+            ]
+        )
+        assert args.n_shared_experts == 1
+        self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit)
+
+    def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> torch.Tensor:
+        shape = x.size()
+        x = x.view(-1, self.dim)
+        weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten())
+        y = torch.zeros_like(x, dtype=torch.float32)
+        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
+        for i in range(self.experts_start_idx, self.experts_end_idx):
+            if counts[i] == 0:
+                continue
+            expert = self.experts[i]
+            idx, top = torch.where(indices == i)
+            y[idx] += expert(x[idx], weights[idx, top, None])
+        if world_size > 1:
+            dist.all_reduce(y)
+        y += self.shared_experts(x)
+        return y.type_as(x).view(shape)
+
+
+class Block(nn.Module):
+    """A block whose residual stream is `hc_mult` parallel copies (Hyper-Connections).
+
+    Attention and FFN each sit between `hc_pre` (collapse the copies into one sublayer input) and
+    `hc_post` (expand back out, mixing the residual in through `comb`). `hc_mixes` derives all three
+    coefficient sets from the stream itself, `comb` made doubly stochastic by Sinkhorn.
+
+    The coefficients a sublayer computes are used by the *next* one -- see `forward`.
+    """
+
+    attention_cls = Attention
+
+    def __init__(
+        self,
+        layer_id: int,
+        args: ModelArgs,
+        engram_layout: EngramLayout | None = None,
+    ):
+        super().__init__()
+        self.layer_id = layer_id
+        self.norm_eps = args.norm_eps
+        self.attn = self.attention_cls(layer_id, args)
+        self.ffn = MoE(layer_id, args)
+        self.engram = None
+        if engram_layout is not None and layer_id in engram_layout.layer_ids:
+            self.engram = Engram(args, layer_id, engram_layout)
+        self.attn_norm = RMSNorm(args.dim, self.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, self.norm_eps)
+        self.hc_mult = hc_mult = args.hc_mult
+        self.hc_sinkhorn_iters = args.hc_sinkhorn_iters
+        self.hc_eps = args.hc_eps
+        mix_hc = (2 + hc_mult) * hc_mult
+        hc_dim = hc_mult * args.dim
+        with set_dtype(torch.float32):
+            self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+            self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+            self.hc_attn_base = nn.Parameter(torch.empty(mix_hc))
+            self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc))
+            self.hc_attn_scale = nn.Parameter(torch.empty(3))
+            self.hc_ffn_scale = nn.Parameter(torch.empty(3))
+
+    def hc_mixes(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor):
+        """x: [b,s,hc,d], hc_fn: [mix_hc, hc*d], hc_scale: [3], hc_base: [mix_hc]. Returns the
+        pre / post / comb coefficients, split out of one projection of the flattened stream."""
+        # normalized over the whole flattened hc*d stream, one statistic per token
+        x = x.flatten(2).float()
+        rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps)
+        mixes = F.linear(x, hc_fn) * rsqrt
+        return hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps)
+
+    def hc_pre(self, x: torch.Tensor, pre_mix: torch.Tensor):
+        """Collapse the hc copies into one, weighted by pre_mix. [b,s,hc,d] x [b,s,hc] -> [b,s,d]"""
+        y = torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2)
+        return y.to(x.dtype)
+
+    def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor):
+        """Expand the sublayer output back to hc copies and mix the residual in through `comb`.
+        x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc] -> [b,s,hc,d]"""
+        y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)
+        return y.type_as(x)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        pre_mix: torch.Tensor,
+        image_mask: torch.Tensor | None,
+        *attn_args,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`pre_mix` collapses the hc_mult copies down to one input for this block's attention. Each
+        sub-block's own `hc_mixes` produces the mix for the *next* one, so attention uses what the
+        previous layer's FFN produced and the FFN uses what this attention produced.
+
+        image_mask: [b, s] bool, True inside image spans (selects the VL routing bias)."""
+        residual = x
+        attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+        x = self.hc_pre(x, pre_mix)
+        x = self.attn_norm(x)
+        x = self.attn(x, start_pos, *attn_args)
+        x = self.hc_post(x, residual, attn_post, attn_comb)
+
+        residual = x
+        ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+        x = self.hc_pre(x, attn_pre)
+        x = self.ffn_norm(x)
+        x = self.ffn(x, image_mask)
+        x = self.hc_post(x, residual, ffn_post, ffn_comb)
+        return x, ffn_pre
+
+    # FreeToken: Decoder SWA Bounded Replay oracle. ``h`` / ``pre_mix`` are the tail [s, N) streams,
+    # or -- at the decoder's kv source -- ``h_all`` / ``pre_mix_all`` cover every token so the
+    # source compresses them all before the block narrows to the tail.
+    def forward_tail(self, h, s: int, pre_mix, h_all=None, pre_mix_all=None):
+        if h_all is not None:
+            attn_pre, attn_post, attn_comb = self.hc_mixes(h_all, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+            x_all = self.attn_norm(self.hc_pre(h_all, pre_mix_all))
+            residual = h_all[:, s:]
+            attn_pre, attn_post, attn_comb = attn_pre[:, s:], attn_post[:, s:], attn_comb[:, s:]
+            x = self.attn.forward_tail(x_all[:, s:], s, x_all)
+        else:
+            residual = h
+            attn_pre, attn_post, attn_comb = self.hc_mixes(h, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+            x = self.attn_norm(self.hc_pre(h, pre_mix))
+            x = self.attn.forward_tail(x, s)
+        x = self.hc_post(x, residual, attn_post, attn_comb)
+
+        residual = x
+        ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+        x = self.hc_pre(x, attn_pre)
+        x = self.ffn_norm(x)
+        x = self.ffn(x, None)
+        x = self.hc_post(x, residual, ffn_post, ffn_comb)
+        return x, ffn_pre
+
+
+class ParallelHead(nn.Module):
+    def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.norm_eps = norm_eps
+        self.hc_eps = hc_eps
+        self.part_vocab_size = vocab_size // world_size
+        # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly
+        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor, full_logits=False):
+        """x: [b, s, d]. Generation only needs the last position, so that is the default."""
+        if not full_logits:
+            x = x[:, -1]
+        logits = F.linear(x.float(), self.weight)
+        if world_size > 1:
+            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
+            dist.all_gather(all_logits, logits)
+            logits = torch.cat(all_logits, dim=-1)
+        return logits
+
+
+@lru_cache(1)
+def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int):
+    assert start_pos > 0
+    matrix = torch.cat(
+        [
+            torch.arange(min(window_size, start_pos + 1)),
+            window_size + torch.arange(block_size),
+        ]
+    )
+    return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1).contiguous()
+
+
+class DSparkAttention(Attention):
+    def forward(self, x: torch.Tensor, start_pos: int, main_x: torch.Tensor):
+        assert self.compress_ratio == 0
+        bsz, seqlen, _ = main_x.size()
+        win = self.window_size
+        rd = self.rope_head_dim
+
+        main_freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+        main_kv = self.kv_norm(self.wkv(main_x))
+        apply_rotary_emb(main_kv[..., -rd:], main_freqs_cis)
+        act_quant(main_kv, fp8_block_size, scale_fmt, scale_dtype, True)
+
+        if start_pos == 0:
+            if seqlen <= win:
+                self.window_kv_cache[:bsz, :seqlen] = main_kv
+            else:
+                cutoff = seqlen % win
+                self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = main_kv[:, -win:].split(
+                    [win - cutoff, cutoff], dim=1
+                )
+            return x
+
+        bsz, block_size, _ = x.size()
+        freqs_cis = self.freqs_cis[start_pos + seqlen : start_pos + seqlen + block_size]
+
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
+        apply_rotary_emb(q[..., -rd:], freqs_cis)
+        kv = self.kv_norm(self.wkv(x))
+        apply_rotary_emb(kv[..., -rd:], freqs_cis)
+        act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True)
+
+        topk_idxs = get_dspark_topk_idxs(win, bsz, block_size, start_pos)
+        self.window_kv_cache[:bsz, start_pos % win] = main_kv.squeeze(1)
+        kv = torch.cat([self.window_kv_cache[:bsz], kv], dim=1)
+        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+        apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+
+        o = o.view(bsz, block_size, self.n_local_groups, -1)
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+        x = self.wo_b(o.flatten(2))
+        return x
+
+
+class DSparkMarkovHead(nn.Module):
+    def __init__(self, vocab_size: int, dspark_markov_rank: int):
+        super().__init__()
+        self.embed = ParallelEmbedding(vocab_size, dspark_markov_rank)
+        self.head = ParallelHead(vocab_size, dspark_markov_rank)
+
+    def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        embed = self.embed(token_ids)
+        logits = self.head(embed, full_logits=True)
+        return logits, embed
+
+
+class DSparkConfidenceHead(nn.Module):
+    def __init__(self, input_dim: int):
+        super().__init__()
+        # proj in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for fp32 confidence score.
+        self.proj = Linear(input_dim, 1, dtype=torch.float32)
+
+    def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor):
+        hidden = torch.cat([hidden, markov_embed], dim=-1)
+        return self.proj(hidden.float()).squeeze(-1)
+
+
+class DSparkBlock(Block):
+    """DSpark stage stored under the mtp.* checkpoint namespace."""
+
+    attention_cls = DSparkAttention
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__(layer_id, args)
+        stage_id = layer_id - args.n_layers
+        self.block_size = args.dspark_block_size
+        self.noise_token_id = args.dspark_noise_token_id
+        self.temperature = args.temperature
+        if stage_id == 0:
+            assert len(args.dspark_target_layer_ids) > 0, "DSpark needs target layers"
+            self.main_proj = Linear(args.dim * len(args.dspark_target_layer_ids), args.dim)
+            self.main_norm = RMSNorm(args.dim, args.norm_eps)
+        if stage_id == args.n_mtp_layers - 1:
+            self.norm = RMSNorm(args.dim, args.norm_eps)
+            self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank)
+            self.confidence_head = DSparkConfidenceHead(args.dim + args.dspark_markov_rank)
+        self.embed: ParallelEmbedding | None = None
+        self.head: ParallelHead | None = None
+
+    def forward(self, x: torch.Tensor, start_pos: int, pre_mix: torch.Tensor, main_x: torch.Tensor):
+        if start_pos == 0:
+            self.attn(x, start_pos, main_x)  # prefill only seeds the window KV cache
+            return x, pre_mix
+        return super().forward(x, start_pos, pre_mix, None, main_x)  # drafts are text: no VL bias
+
+    def forward_embed(self, main_hidden: torch.Tensor, input_ids: torch.Tensor):
+        assert self.embed is not None
+        main_x = self.main_norm(self.main_proj(main_hidden))
+        draft_input_ids = input_ids.new_full([input_ids.size(0), self.block_size], self.noise_token_id)
+        draft_input_ids[:, 0] = input_ids
+        x = self.embed(draft_input_ids)
+        x = x.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+        return x, main_x
+
+    def forward_head(
+        self,
+        x: torch.Tensor,
+        pre_mix: torch.Tensor,
+        input_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert self.head is not None
+        x = self.hc_pre(x, pre_mix)
+        logits = self.head(self.norm(x), full_logits=True)
+        output_ids = input_ids.new_empty(input_ids.size(0), self.block_size + 1)
+        output_ids[:, 0] = input_ids
+        markov_embeds = []
+        for i in range(self.block_size):
+            logits_bias, markov_embed = self.markov_head(output_ids[:, i])
+            logits[:, i].add_(logits_bias)
+            markov_embeds.append(markov_embed)
+            output_ids[:, i + 1] = sample(logits[:, i], self.temperature)
+        markov_embed = torch.stack(markov_embeds, dim=1)
+        confidence = self.confidence_head(x, markov_embed)
+        return output_ids, logits, confidence
+
+
+def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor:
+    """initial one-hot mix"""
+    pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32)
+    pre_mix[:, :, 0] = 1.0
+    return pre_mix
+
+
+class SharedAttentionRuntime:
+    """What attention layers hand down the stack instead of recomputing. Layers run in order and
+    every source writes before its consumers read, so one slot each is enough and nothing needs
+    resetting between forwards. Sources: compress_kv and index_k from kv_source_layers, topk_idxs
+    from index_source_layers, candidates from candidate_source_layer."""
+
+    def __init__(self):
+        self.compress_kv: torch.Tensor | None = None
+        self.index_k: torch.Tensor | None = None
+        self.topk_idxs: torch.Tensor | None = None
+        self.candidates: torch.Tensor | None = None
+
+
+# Only ever one model per process, same as world_size / rank / default_dtype above.
+shared_attn = SharedAttentionRuntime()
+
+
+class Transformer(nn.Module):
+    """DeepSeek-V4.1: embed -> expand to hc_mult copies -> blocks -> collapse -> logits. Building
+    this sets the globals at the top of the file. The tokenizer only feeds the engram token map."""
+
+    def __init__(self, args: ModelArgs, tokenizer=None):
+        global world_size, rank, default_dtype
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
+        super().__init__()
+        self.max_seq_len = args.max_seq_len
+        self.temperature = args.temperature
+        self.norm_eps = args.norm_eps
+        self.hc_eps = args.hc_eps
+        self.engram_layout = EngramLayout.from_args(args)
+        self.engram_hash = (
+            NgramHashState(args, self.engram_layout, tokenizer) if self.engram_layout is not None else None
+        )
+        self.embed = ParallelEmbedding(args.vocab_size, args.dim)
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(args.n_layers):
+            self.layers.append(Block(layer_id, args, self.engram_layout))
+        self.norm = RMSNorm(args.dim, self.norm_eps)
+        self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps)
+        self.mtp = torch.nn.ModuleList()
+        self.target_layer_ids = args.dspark_target_layer_ids
+        if args.dspark_block_size:
+            for layer_id in range(args.n_mtp_layers):
+                self.mtp.append(DSparkBlock(args.n_layers + layer_id, args))
+                self.mtp[-1].embed = self.embed
+                self.mtp[-1].head = self.head
+        self.hc_mult = args.hc_mult
+        self.vision = None
+        if args.vision_enabled:
+            self.vision = ViT(args)
+            self.aligner = Aligner(args)
+            # learned embeddings for the image span delimiters
+            self.image_start = nn.Parameter(torch.empty(args.dim))
+            self.image_end = nn.Parameter(torch.empty(args.dim))
+            self.image_newline = nn.Parameter(torch.empty(args.dim))
+
+    @torch.inference_mode()
+    def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
+        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
+
+    def merge_image_embeddings(self, images, h: torch.Tensor):
+        """Overwrite each image's token span in h with its ViT/aligner features. The IMAGE slots take
+        the aligner rows in row-major order; the span delimiters take learned embeddings."""
+        for i, sample in enumerate(images):
+            for img in sample or ():
+                types = img.types.to(h.device)
+                span = h[i, img.start : img.start + types.numel()]
+                span[types == IMAGE_START] = self.image_start.to(h.dtype)
+                span[types == IMAGE_END] = self.image_end.to(h.dtype)
+                span[types == IMAGE_NEW_LINE] = self.image_newline.to(h.dtype)
+                embeds = self.encode_image(img.patches.to(h.device), img.n_vit_h, img.n_vit_w)
+                span[types == IMAGE] = embeds.to(h.dtype)
+
+    @torch.inference_mode()
+    def forward(
+        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
+    ):
+        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
+        has already filled, so the padding it uses internally never reaches here. token_types /
+        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
+        inside the first (start_pos 0) chunk."""
+        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
+        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
+        engram_mask = None if image_mask is None else ~image_mask
+        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
+        h = self.embed(input_ids)
+        if images is not None:
+            assert start_pos == 0, "image spans must be prefilled in a single chunk"
+            self.merge_image_embeddings(images, h)
+        # Expand to hc_mult copies for Hyper-Connections
+        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+        main_hiddens = []
+        pre_mix = make_identity_pre_mix(h, self.hc_mult)
+        for i, layer in enumerate(self.layers):
+            if layer.engram is not None:
+                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
+            # the MTP head reads the attention input of its target layers, not their output
+            if i in self.target_layer_ids:
+                main_hiddens.append(h.mean(dim=2))
+            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
+        h = layer.hc_pre(h, pre_mix)
+        logits = self.head(self.norm(h))
+        output_ids = sample(logits, self.temperature)
+        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
+        return output_ids, logits, main_hidden
+
+    # FreeToken: Decoder SWA Bounded Replay oracle (tech report sec. 3.2.2): the encoder runs on
+    # every prompt token, the decoder on the last ``n_win`` with its sliding window truncated
+    # there; the decoder's kv source (the last kv_source layer) compresses every token. Returns
+    # the last position's logits and leaves the caches ready for decode from position seqlen.
+    @torch.inference_mode()
+    def forward_bounded(self, input_ids: torch.Tensor, n_win: int):
+        engram_hashes = self.engram_hash(input_ids, 0, None) if self.engram_hash is not None else None
+        h = self.embed(input_ids).unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+        pre_mix = make_identity_pre_mix(h, self.hc_mult)
+        N = input_ids.size(1)
+        s = max(0, N - n_win)
+        dec_start = max(i for i, layer in enumerate(self.layers) if layer.attn.is_kv_source)
+        for i, layer in enumerate(self.layers):
+            if layer.engram is not None:
+                assert i < dec_start, "Engram sits in the encoder"
+                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], None)
+            if i < dec_start:
+                h, pre_mix = layer(h, 0, pre_mix, None)
+            elif i == dec_start:
+                h, pre_mix = layer.forward_tail(None, s, None, h_all=h, pre_mix_all=pre_mix)
+            else:
+                h, pre_mix = layer.forward_tail(h, s, pre_mix)
+        h = layer.hc_pre(h, pre_mix)
+        return self.head(self.norm(h))
+
+    @torch.inference_mode()
+    def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0):
+        h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids)
+        pre_mix = make_identity_pre_mix(h, self.hc_mult)
+        for layer in self.mtp:
+            h, pre_mix = layer(h, start_pos, pre_mix, main_x)
+        if start_pos == 0:
+            return None
+        return self.mtp[-1].forward_head(h, pre_mix, input_ids)
+
+
+def sample(logits, temperature: float = 1.0):
+    """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU,
+    since it avoids the GPU-to-CPU sync in torch.multinomial."""
+    if temperature == 0:
+        return logits.argmax(dim=-1)
+    logits = logits / max(temperature, 1e-5)
+    probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
+    return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
+
+
+if __name__ == "__main__":
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device("cuda")
+    torch.manual_seed(0)
+    args = ModelArgs(dspark_block_size=6, dspark_target_layer_ids=(3, 4))
+    x = torch.randint(0, args.vocab_size, (2, 150))
+    model = Transformer(args)
+
+    output_ids, logits, main_hidden = model(x[:, :128])
+    model.forward_spec(output_ids, main_hidden)
+    for i in range(128, 150):
+        output_ids, logits, main_hidden = model(x[:, i : i + 1], i)
+        result = model.forward_spec(output_ids, main_hidden, i)
+        assert result is not None
+        output_ids, logits, confidence = result
diff --git a/tests/models/deepseek_v41/reference/vision.py b/tests/models/deepseek_v41/reference/vision.py
new file mode 100644
index 0000000000..2fd86160f7
--- /dev/null
+++ b/tests/models/deepseek_v41/reference/vision.py
@@ -0,0 +1,121 @@
+# Vendored from deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a, inference/vision.py.
+# MIT License; see LICENSE in this directory.
+from functools import lru_cache
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@lru_cache(8)
+def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
+    wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
+    freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq
+    freqs = freqs.flatten(1)
+    return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
+
+
+def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    dtype = x.dtype
+    x1, x2 = x.float().chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.proj = nn.Linear(3 * args.vision_patch_size**2, args.vision_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x.flatten(1))
+
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.n_heads = args.vision_n_heads
+        self.head_dim = args.vision_dim // args.vision_n_heads
+        self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim)
+        self.wo = nn.Linear(args.vision_dim, args.vision_dim)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        n = x.size(0)
+        q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1))
+        q = apply_rotary(q, cos, sin)
+        k = apply_rotary(k, cos, sin)
+        o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1))
+        return self.wo(o.transpose(0, 1).reshape(n, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False)
+        self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.w1(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.norm1 = RMSNorm(args.vision_dim)
+        self.attn = Attention(args)
+        self.norm2 = RMSNorm(args.vision_dim)
+        self.mlp = MLP(args)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), cos, sin)
+        return x + self.mlp(self.norm2(x))
+
+
+class ViT(nn.Module):
+    """DeepSeek ViT: full bidirectional attention over one image with 2D RoPE."""
+
+    def __init__(self, args):
+        super().__init__()
+        self.rope_dim = args.vision_dim // args.vision_n_heads // 2
+        self.rope_theta = args.vision_rope_theta
+        self.patch_embed = PatchEmbed(args)
+        self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)])
+        self.norm = RMSNorm(args.vision_dim)
+
+    def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        x = self.patch_embed(patches)
+        cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        return self.norm(x)
+
+
+class Aligner(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.downsample_ratio = args.vision_downsample_ratio
+        in_dim = args.vision_dim * self.downsample_ratio**2
+        self.w1 = nn.Linear(in_dim, args.dim)
+        self.w2 = nn.Linear(args.dim, args.dim)
+
+    def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        r = self.downsample_ratio
+        x = x.view(n_h, n_w, -1).permute(2, 0, 1)
+        x = F.pad(x, (0, -n_w % r, 0, -n_h % r))
+        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
+        return self.w2(F.gelu(self.w1(x)))
diff --git a/tests/models/deepseek_v41/test_config.py b/tests/models/deepseek_v41/test_config.py
new file mode 100644
index 0000000000..fab385fb93
--- /dev/null
+++ b/tests/models/deepseek_v41/test_config.py
@@ -0,0 +1,169 @@
+"""Config parsing, the DSV41 layer-role table, registry dispatch and loader / model key parity."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import pytest
+import torch
+
+from freetoken.layers import set_rope_device
+from freetoken.layers.quantization import NoQuantConfig
+from freetoken.distributed.info import set_tp_info, try_get_tp_info
+from freetoken.models import create_model
+from freetoken.models.deepseek_v41.args import DeepseekV41Args, Mode
+from freetoken.models.deepseek_v41.config import parse_config
+from freetoken.models.register import get_model_spec
+from freetoken.utils.torch_utils import torch_dtype
+
+from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint
+
+
+def test_roles_follow_the_dsv41_mode_table():
+    args = parse_config(tiny_hf_config()).dsv41_args
+    modes = [r.mode for r in args.roles]
+    assert modes == [Mode.WINDOW, Mode.WINDOW, Mode.FULL, Mode.REUSE, Mode.FULL, Mode.REINDEX]
+    assert [r.kv_source for r in args.roles] == [None, None, 2, 2, 4, 4]
+    assert args.roles[4].is_candidate_source and args.roles[5].uses_candidates and not args.roles[2].uses_candidates
+    assert args.decoder_start_layer == 4 and args.backbone_kv_sources == (2, 4)
+    assert args.freqs_params(0)[1] == 0 and args.freqs_params(2)[1] == 1024  # yarn only on compressing layers
+    assert args.swa_decoder_replay == "bounded"
+
+
+def test_shipping_layout_roles():
+    """The real 40-layer layout: 3 encoder Full layers (ratio 2), the decoder Full at 20 building the
+    candidate pool, Reindex at 24/28/32/36 inside it, Reuse elsewhere."""
+    text = tiny_hf_config().text_config
+    text.update(
+        num_hidden_layers=40, compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0],
+        kv_source_layer_ids=[2, 8, 14, 20], index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36],
+        candidate_source_layer_id=20, num_nextn_predict_layers=3, engram_layer_ids=[1, 14], engram_num_embeddings=[10, 10],
+    )
+    args = DeepseekV41Args.from_hf(tiny_hf_config(**text))
+    full = [r.layer_id for r in args.roles if r.mode is Mode.FULL]
+    reindex = [r.layer_id for r in args.roles if r.mode is Mode.REINDEX]
+    assert full == [2, 8, 14, 20] and reindex == [24, 28, 32, 36]
+    assert sum(r.mode is Mode.REUSE for r in args.roles) == 40 - 2 - 4 - 4
+    assert [r.layer_id for r in args.roles if r.uses_candidates] == reindex
+    assert args.decoder_start_layer == 20
+
+
+def test_bad_layouts_are_rejected():
+    with pytest.raises(ValueError, match="no kv source"):
+        DeepseekV41Args.from_hf(tiny_hf_config(kv_source_layer_ids=[4]))
+    with pytest.raises(ValueError, match="reads kv source"):
+        DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 1, 1, 1, 0]))
+    with pytest.raises(ValueError, match="different kv source"):
+        DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 2, 2, 2, 0], kv_source_layer_ids=[2, 4], candidate_source_layer_id=2, index_source_layer_ids=[2, 4, 5]))
+
+
+def test_registry_and_kv_sources():
+    spec = get_model_spec("DeepseekV41ForCausalLM")
+    assert spec.module == "freetoken.models.deepseek_v41"
+    from freetoken.attention import AttnType
+    from freetoken.models.config import DSV4AttentionGroupConfig
+
+    mc = parse_config(tiny_hf_config())
+    args = mc.dsv41_args
+    assert args.backbone_kv_sources == (2, 4) and tuple(r.ratio for r in args.roles) == tuple(RATIOS[:N_LAYERS])
+    (group,) = mc.attention_groups
+    assert type(group) is DSV4AttentionGroupConfig and group.variant == "v41"
+    assert mc.kv_cache_group_specs()[0].attn_type is AttnType.DSV41 and mc.attn_type_for_layer(0) is AttnType.DSV41
+    assert mc.num_layers == N_LAYERS and mc.dsv41_args is not None
+
+
+def _meta_model(mc):
+    if try_get_tp_info() is None:
+        set_tp_info(0, 1)
+    set_rope_device(torch.device("cpu"))
+    mc = replace(mc, moe_strategy="offload", decode_target="gpu", quant=NoQuantConfig())
+    with torch.device("meta"), torch_dtype(torch.bfloat16):
+        return create_model(mc)
+
+
+def test_loader_names_every_model_parameter(tmp_path):
+    from freetoken.models.deepseek_v41.weight import iter_weights
+
+    write_tiny_checkpoint(str(tmp_path))
+    model = _meta_model(parse_config(tiny_hf_config()))
+    params = model.state_dict()
+    loaded = dict(iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False))
+    assert set(loaded) == set(params), (set(loaded) ^ set(params))
+    for name, tensor in loaded.items():
+        assert tuple(tensor.shape) == tuple(params[name].shape), name
+    # routed experts, the DSpark stage, the vision tower and the image routing bias are not resident weights
+    assert not any(".experts." in k or k.startswith(("mtp.", "vision.")) or k.endswith("bias_vl") for k in loaded)
+    # per-role structure
+    assert model.model.layers.op_list[3].attn.compressor is None and model.model.layers.op_list[3].attn.indexer is None
+    assert model.model.layers.op_list[4].attn.compressor.ratio == 1 and model.model.layers.op_list[2].attn.compressor.ratio == 2
+    assert [b.engram is not None for b in model.model.layers.op_list] == [L in ENGRAM_LAYERS for L in range(N_LAYERS)]
+    assert all(L in INDEX_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.indexer is not None)
+    assert all(L in KV_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.compressor is not None)
+
+
+@pytest.mark.parametrize("missing_weight", [False, True])
+def test_loader_closes_shards_when_expert_loading_stops(tmp_path, monkeypatch, missing_weight):
+    from freetoken.models.deepseek_v41 import weight
+
+    write_tiny_checkpoint(str(tmp_path))
+    closed = []
+    close, get = weight._ShardReader.close, weight._ShardReader.get
+
+    def close_reader(reader):
+        close(reader)
+        closed.append(True)
+
+    def read_weight(reader, name):
+        if missing_weight and name.endswith(".w3.weight"):
+            raise KeyError(name)
+        return get(reader, name)
+
+    monkeypatch.setattr(weight._ShardReader, "close", close_reader)
+    monkeypatch.setattr(weight._ShardReader, "get", read_weight)
+    tensors = weight.iter_weights(str(tmp_path), torch.device("cpu"))
+    if missing_weight:
+        with pytest.raises(KeyError):
+            next(tensors)
+    else:
+        next(tensors)
+        tensors.close()
+    assert closed == [True]
+
+
+def test_fp8_wo_a_layer_rejects_a_dequantized_ftw_layout_clearly():
+    """An fp8-declared wo_a fed a state dict without ``wo_a_scale`` (an FTW converted before wo_a kept
+    its payload) fails with a reconvert message, not a raw missing-key error; a bf16-declared layer fed
+    the fp8 payload + scale dequantizes it on the way in."""
+    from types import SimpleNamespace
+
+    from freetoken.layers.quantization.scheme import fp8_block_scheme
+    from freetoken.models.deepseek_v41.attention import DSV41Attention
+
+    if try_get_tp_info() is None:
+        set_tp_info(0, 1)
+    args = parse_config(tiny_hf_config()).dsv41_args
+    args.max_seq_len = 512
+    # wo_a is fp8; the LinearReplicated modules stay unquantized (the bf16 quant method)
+    plain = NoQuantConfig()
+    quant = SimpleNamespace(
+        scheme_for=lambda name: fp8_block_scheme("ue8m0", 32) if name.endswith(".wo_a") else None,
+        get_quant_method=lambda layer, prefix: plain.get_quant_method(layer, prefix),
+    )
+    with torch.device("meta"), torch_dtype(torch.bfloat16):
+        fp8_layer = DSV41Attention(args, args.roles[2], quant_config=quant, prefix="attn")
+        bf16_layer = DSV41Attention(args, args.roles[2], quant_config=None, prefix="attn")
+    assert fp8_layer.wo_a.dtype == torch.float8_e4m3fn and bf16_layer.wo_a.dtype == torch.bfloat16
+    rows, cols = fp8_layer.wo_a.shape
+    with torch.device("meta"):
+        state = {k: torch.empty_like(v) for k, v in fp8_layer.state_dict(prefix="attn").items()}
+    del state["attn.wo_a_scale"]
+    state["attn.wo_a"] = torch.empty(rows, cols, dtype=torch.bfloat16, device="meta")
+    with pytest.raises(RuntimeError, match="reconvert"):
+        fp8_layer.load_state_dict(dict(state), prefix="attn")
+    # the bf16 layer given the fp8 payload + scale dequantizes (real tensors: the dequant runs)
+    w = (torch.randn(rows, cols) * 0.05).to(torch.float8_e4m3fn)
+    scale = torch.full((rows // 32, cols // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu)
+    state2 = {k: torch.zeros(v.shape, dtype=v.dtype) for k, v in bf16_layer.state_dict(prefix="attn").items()}
+    state2["attn.wo_a"], state2["attn.wo_a_scale"] = w, scale
+    bf16_layer.load_state_dict(state2, prefix="attn")
+    assert torch.equal(bf16_layer.wo_a, w.to(torch.bfloat16))
diff --git a/tests/models/deepseek_v41/test_engine_config.py b/tests/models/deepseek_v41/test_engine_config.py
new file mode 100644
index 0000000000..116209ea51
--- /dev/null
+++ b/tests/models/deepseek_v41/test_engine_config.py
@@ -0,0 +1,81 @@
+"""Engine config resolution for DeepSeek-V4.1 on the tiny checkpoint: the DSV41 backend and pool are
+picked, the page size is the window page, the replay knob lands on the model args, and the fp8
+dialect (block 32, fp4 experts) reaches the quant layer with the matching activation block."""
+
+from __future__ import annotations
+
+import json
+import os
+
+import pytest
+import torch
+
+from freetoken.distributed import DistributedInfo
+from freetoken.scheduler.config import SchedulerConfig
+from freetoken.engine.engine import _adjust_config
+from freetoken.kvcache import resolve_pool_class
+from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache
+
+from .common import write_tiny_checkpoint
+
+
+def _engine_config(path, **over):
+    return SchedulerConfig(model_path=path, tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, **over)
+
+
+@pytest.mark.parametrize("replay", ["bounded", "exact"])
+def test_resolution_picks_dsv41_and_the_replay_knob(tmp_path, monkeypatch, replay):
+    from freetoken.engine import engine
+
+    monkeypatch.setattr(engine, "is_sm100_family", lambda: False)
+    monkeypatch.setattr(engine, "is_sm90_family", lambda: True)
+    write_tiny_checkpoint(str(tmp_path))
+    config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", swa_decoder_replay=replay, max_seq_len_override=2048)
+    _adjust_config(config)
+    assert config.attention_backend == "dsv41_sparse"
+    assert config.page_size == 128 and config.cache_type == "swa_radix"
+    assert resolve_pool_class(config.model_config) is DSV41PagedKVCache
+    args = config.model_config.dsv41_args
+    assert args.swa_decoder_replay == replay and args.max_seq_len == 2048 and args.max_batch_size == config.max_running_req + 1
+    assert config.max_extend_tokens == 8192  # the prefill chunk stays bounded (whole window pages)
+    # the cache contract follows the replay mode: a bounded-mode prefix hit must leave the prompt's last
+    # window to the prefill (the pool's prefix_replay_tokens); the history a resume reads stays one window
+    from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes
+
+    spec = config.model_config.kv_cache_group_specs()[0]
+    pool = DSV41PagedKVCache(dsv41_pool_sizes(16, args, 1.0, 128), args, torch.device("cpu"))
+    assert pool.prefix_replay_tokens == (128 if replay == "bounded" else 0)
+    assert pool.sliding_window_size == 128 and spec.sliding_window == 128 and not spec.is_swa
+    # bounded replay keeps the decoder's per-request window KV in private rings, off the shared pages
+    assert pool.private_window_layer_ids == (tuple(range(args.decoder_start_layer, args.n_layers)) if replay == "bounded" else ())
+    assert DSV41PagedKVCache.min_kv_tokens(config) // 128 == 8 + 3 * config.max_running_req + 2 * (config.max_running_req + 1) + 1
+
+
+def test_fp8_block32_dialect_reaches_the_quant_layer(tmp_path):
+    """A V4.1-style quantization_config on the tiny checkpoint: dense linears get the 32-block
+    e8m0 scheme, routed experts the MXFP4 scheme with a 32-wide activation block."""
+    from freetoken.distributed.info import set_tp_info, try_get_tp_info
+    from freetoken.layers import OffloadMoELayer
+    from freetoken.layers.quantization import Fp8BlockConfig, QuantKind
+    from freetoken.layers.quantization.scheme import fp8_block_size
+
+    write_tiny_checkpoint(str(tmp_path))
+    cfg_path = os.path.join(str(tmp_path), "config.json")
+    with open(cfg_path) as f:
+        cfg = json.load(f)
+    cfg["quantization_config"] = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"}
+    with open(cfg_path, "w") as f:
+        json.dump(cfg, f)
+    config = _engine_config(str(tmp_path), moe_strategy="offload")
+    quant = config.model_config.quant
+    assert type(quant) is Fp8BlockConfig and quant.block == 32
+    dense = quant.scheme_for("model.layers.3.attn.wq_a")
+    assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32
+    if try_get_tp_info() is None:
+        set_tp_info(0, 1)
+    experts = OffloadMoELayer(3, num_experts=4, top_k=2, hidden_size=64, intermediate_size=64, limit=10.0, quant_config=quant, prefix="model.layers.3.ffn.experts")
+    assert experts.quant_method.scheme.kind is QuantKind.MXFP4 and experts.quant_method.cfg.act_block == 32
+    # the family's bf16 modules stay unquantized although the fp8 config lists no exceptions
+    for name in ("model.layers.2.attn.compressor.wkv", "model.layers.2.attn.indexer.wk", "model.layers.5.attn.indexer.weights_proj", "head"):
+        assert quant.scheme_for(name) is None, name
+    assert quant.scheme_for("model.layers.1.engram.wkv").kind is QuantKind.FP8_BLOCK
diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py
new file mode 100644
index 0000000000..dec852ec03
--- /dev/null
+++ b/tests/models/deepseek_v41/test_engram.py
@@ -0,0 +1,281 @@
+"""Engram: the hash against a transcription of the reference ``NgramHashState``, the prime buckets,
+the disk table against a torch oracle over the synthetic checkpoint (eager, decode with context,
+CUDA-graph flag sync), and the layer's gate math against the reference ``Engram.forward``."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from freetoken.models.deepseek_v41.args import DeepseekV41Args
+from freetoken.models.deepseek_v41.engram import EngramHash, bucket_primes, hash_multipliers
+
+from .common import VOCAB, requires_cuda, tiny_hf_config, write_tiny_checkpoint
+
+
+def _args(**over) -> DeepseekV41Args:
+    return DeepseekV41Args.from_hf(tiny_hf_config(**over))
+
+
+def _reference_hashes(args: DeepseekV41Args, ids: torch.Tensor, start_pos: int, cache: torch.Tensor) -> torch.Tensor:
+    """``NgramHashState.forward`` for one sequence with an identity token map; ``cache`` is the
+    sequence's compressed-id history buffer (mutated like the reference's)."""
+    layout_primes = bucket_primes(args)
+    flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout_primes]
+    offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat]))
+    primes = torch.tensor(layout_primes)
+    multipliers = hash_multipliers(args.engram_layer_ids, args.engram_max_ngram_size, args.engram_compressed_vocab_size)
+    pad_id = args.engram_pad_id
+    seqlen = ids.numel()
+    cache[start_pos : start_pos + seqlen] = ids
+    positions = torch.arange(start_pos, start_pos + seqlen)
+    tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool)
+    for shift in range(args.engram_max_ngram_size):
+        source = cache[(positions - shift).clamp_min(0)]
+        blocked = blocked | (positions < shift)
+        tokens.append(torch.where(blocked, pad_id, source))
+    tokens = torch.stack(tokens, dim=-1)
+    products = tokens.unsqueeze(1) * multipliers
+    rolling, hashes = products[..., 0], []
+    for i in range(1, args.engram_max_ngram_size):
+        rolling = torch.bitwise_xor(rolling, products[..., i])
+        hashes.append(rolling.unsqueeze(-1) % primes[:, i - 1])
+    return torch.cat(hashes, dim=-1) + offsets
+
+
+def test_primes_are_distinct_and_ordered():
+    args = _args()
+    primes = bucket_primes(args)
+    flat = [p for layer in primes for per in layer for p in per]
+    assert len(flat) == len(set(flat)) and flat == sorted(flat) and min(flat) > args.engram_vocab_size - 1
+    from sympy import isprime
+
+    assert all(isprime(p) for p in flat)
+
+
+def test_hash_matches_the_reference_across_a_prefill_decode_split():
+    args = _args()
+    hash = EngramHash(args, list(range(VOCAB)), VOCAB)
+    ids = torch.randint(3, VOCAB, (37,), generator=torch.Generator().manual_seed(0))
+    cache = torch.empty(64, dtype=torch.int64)
+    want = _reference_hashes(args, ids[:30], 0, cache)
+    got = hash.row_ids(ids[:30])
+    assert torch.equal(got, want)
+    # decode continues with the three preceding tokens as context
+    for step in range(30, 37):
+        want = _reference_hashes(args, ids[step : step + 1], step, cache)
+        got = hash.row_ids(ids[step : step + 1], context=ids[:step])
+        assert torch.equal(got, want), step
+    # a run resumed mid-sequence (chunked prefill) hashes like the uninterrupted one
+    whole = hash.row_ids(ids)
+    assert torch.equal(torch.cat([hash.row_ids(ids[:20]), hash.row_ids(ids[20:], context=ids[:20])]), whole)
+    assert whole.shape == (37, 1, hash.n_cols) and int(whole.max()) < hash.rows_per_layer[0]
+    with pytest.raises(ValueError, match="compressed vocab"):
+        EngramHash(args, list(range(VOCAB)), VOCAB + 1)
+
+
+@requires_cuda
+def test_disk_table_matches_the_dequantized_shard(tmp_path):
+    from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source
+
+    tensors = write_tiny_checkpoint(str(tmp_path))
+    args = _args()
+    hash = EngramHash(args, list(range(VOCAB)), VOCAB)
+    src = engram_row_source(str(tmp_path), 1)
+    assert src.num_rows == tensors["layers.1.engram.embed.weight"].shape[0] and src.head_dim == 32
+    table = EngramDiskTable(src, hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16)
+    fp8 = tensors["layers.1.engram.embed.weight"]
+    scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0)
+
+    def oracle(rows: torch.Tensor) -> torch.Tensor:
+        vals = fp8[rows.flatten()].float().view(-1, 32 // 32, 32) * scale[rows.flatten()].unsqueeze(-1)
+        return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda()
+
+    ids = torch.randint(3, VOCAB, (40,), generator=torch.Generator().manual_seed(1))
+    rows = hash.row_ids(ids)[:, 0]
+    table.stage(rows, graph=False)  # larger than the initial eager staging: grows
+    table.flush(signal=False)
+    assert torch.equal(table.lookup(40), oracle(rows))
+    # the graph staging + flag protocol: capture a lookup, fill after the launch, replay
+    from freetoken.kernel.row_store import probe_wait_sync
+
+    table.wait_sync = probe_wait_sync("auto", torch.device("cuda"))
+    dec_rows = hash.row_ids(ids[40 - 1 :], context=ids[: 40 - 1])[:, 0]
+    out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda")
+    graph = torch.cuda.CUDAGraph()
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        if not table.wait_sync:
+            table.stage(dec_rows, graph=True)
+            table.flush(signal=False)
+        with torch.cuda.graph(graph, stream=stream):
+            out.copy_(table.lookup(1))
+    torch.cuda.synchronize()
+    graph.replay()  # under wait-sync the replay blocks until the host signals
+    table.stage(dec_rows, graph=True)
+    table.flush(signal=table.wait_sync)
+    torch.cuda.synchronize()
+    assert torch.equal(out, oracle(dec_rows))
+
+
+@requires_cuda
+def test_eager_staging_does_not_overwrite_an_in_flight_lookup(tmp_path):
+    """The overlap scheduler stages batch k+1 while batch k's lookup copy is still queued behind
+    the layers before the Engram layer (a sleep kernel here): batch k must read its own rows."""
+    from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source
+
+    tensors = write_tiny_checkpoint(str(tmp_path))
+    hash = EngramHash(_args(), list(range(VOCAB)), VOCAB)
+    table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=64)
+    fp8 = tensors["layers.1.engram.embed.weight"]
+    scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0)
+
+    def oracle(rows):
+        vals = fp8[rows.flatten()].float().view(-1, 1, 32) * scale[rows.flatten()].unsqueeze(-1)
+        return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda()
+
+    gen = torch.Generator().manual_seed(0)
+    batches = [hash.row_ids(torch.randint(3, VOCAB, (40,), generator=gen))[:, 0] for _ in range(4)]
+    got = []
+    for rows in batches:
+        table.stage(rows, graph=False)
+        table.flush(signal=False)
+        torch.cuda._sleep(50_000_000)
+        got.append(table.lookup(rows.shape[0]))
+    torch.cuda.synchronize()
+    for i, (rows, out) in enumerate(zip(batches, got)):
+        assert torch.equal(out, oracle(rows)), f"batch {i} read another batch's rows"
+
+
+@requires_cuda
+@torch.inference_mode()
+def test_host_fills_a_graph_decode_step_from_a_worker_thread(tmp_path):
+    """The deferred graph-decode fill must not wait for the launch call to return: a 40-layer graph
+    blocks cuGraphLaunch on the host while the GPU sits on the WAIT, so the host coordinator reads the
+    tokens back and signals from a worker thread. The worker hashes against the DISPATCH snapshot of
+    the request state: advancing the request right after dispatch (what the engine does) must not
+    change which rows are fetched."""
+    from types import SimpleNamespace
+
+    from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source
+
+    tensors = write_tiny_checkpoint(str(tmp_path))
+    args = _args()
+    hash = EngramHash(args, list(range(VOCAB)), VOCAB)
+    table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16)
+    host = EngramHost(hash, [table], torch.device("cuda"))
+    if not host.wait_sync:
+        pytest.skip("stream memops unavailable: launch-gating mode needs no worker")
+    ids = torch.randint(3, VOCAB, (20,), generator=torch.Generator().manual_seed(7))
+    req = SimpleNamespace(input_ids=ids.clone(), device_len=20, cached_len=19)
+    batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[19:20].to(torch.int32).cuda())
+    out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda")
+    stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            out.copy_(table.lookup(1))
+    torch.cuda.synchronize()
+    with torch.cuda.stream(stream):
+        # slow the readback so the request advances (complete_one + append) before the worker runs
+        torch.cuda._sleep(200_000_000)
+        with host.forward_host_ctx(batch, use_graph=True):
+            graph.replay()
+            req.device_len += 1
+            req.input_ids = torch.cat([req.input_ids, torch.tensor([5], dtype=ids.dtype)])
+    torch.cuda.synchronize()
+    fp8 = tensors["layers.1.engram.embed.weight"]
+    rows = hash.row_ids(ids[19:20], context=ids[:19])[:, 0].flatten()
+    want = fp8[rows].float().view(1, -1).to(torch.bfloat16).cuda()  # scales are all 1.0 in the synthetic table
+    assert torch.equal(out, want)
+    host.close()
+
+
+@requires_cuda
+def test_host_surfaces_a_failed_deferred_fill_at_dispatch_exit(tmp_path):
+    from types import SimpleNamespace
+
+    from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source
+
+    write_tiny_checkpoint(str(tmp_path))
+    args = _args()
+    hash = EngramHash(args, list(range(VOCAB)), VOCAB)
+    table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16)
+    host = EngramHost(hash, [table], torch.device("cuda"))
+    if not host.wait_sync:
+        pytest.skip("stream memops unavailable")
+    ids = torch.randint(3, VOCAB, (8,), generator=torch.Generator().manual_seed(8))
+    req = SimpleNamespace(input_ids=ids, device_len=8, cached_len=7)
+    batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[7:8].to(torch.int32).cuda())
+    stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda")
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            out.copy_(table.lookup(1))
+    torch.cuda.synchronize()
+    table.stage = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("disk gone"))  # the fill fails
+    with pytest.raises(RuntimeError, match="disk gone"):
+        with torch.cuda.stream(stream):
+            with host.forward_host_ctx(batch, use_graph=True):
+                graph.replay()
+    torch.cuda.synchronize()  # the failing worker released the flag: the stream is not stuck
+    host.close()
+
+
+def test_ftw_side_files_carry_the_tables_and_reopen(tmp_path):
+    """An FTW conversion copies the shards holding the Engram tables next to it (the converter's
+    ``ftw_side_files`` hook, which skips the HF index); the row source resolves the copies through
+    their headers exactly like the HF shards, byte for byte."""
+    from freetoken.kernel.row_store import RowStore
+    from freetoken.models.deepseek_v41.engram_table import engram_row_source
+    from freetoken.models.deepseek_v41.weight import ftw_side_files
+
+    src, out = tmp_path / "hf", tmp_path / "ftw"
+    out.mkdir()
+    tensors = write_tiny_checkpoint(str(src))
+    (out / "config.json").write_bytes((src / "config.json").read_bytes())  # the converter copies the metadata
+    shard = "model-00001-of-00001.safetensors"  # the tiny checkpoint's single shard holds the table
+    assert ftw_side_files(str(src), str(out)) == [shard]
+    assert (out / shard).read_bytes() == (src / shard).read_bytes() and not (out / "model.safetensors.index.json").exists()
+    source = engram_row_source(str(out), 1)
+    assert source.path == str(out / shard) and source.num_rows == 4096 and source.head_dim == 32
+    assert torch.equal(source.scales, tensors["layers.1.engram.embed.scale"].view(torch.uint8))
+    store = RowStore(paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows,
+                     row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=False)
+    ids = torch.tensor([0, 7, 4095], dtype=torch.int64)
+    buf = torch.zeros(3 * 32, dtype=torch.uint8)
+    store.stage_rows(ids.data_ptr(), 3, buf.data_ptr(), 0)
+    store.flush(0)
+    assert torch.equal(buf.view(3, 32), tensors["layers.1.engram.embed.weight"][ids].view(torch.uint8))
+
+
+def test_layer_gate_matches_the_reference():
+    """``EngramLayer.forward`` against the reference ``Engram.forward`` (bf16 wkv, fp32 gate math)."""
+    from freetoken.distributed.info import set_tp_info, try_get_tp_info
+    from freetoken.models.deepseek_v41.engram import EngramLayer
+
+    if try_get_tp_info() is None:
+        set_tp_info(0, 1)
+    args = _args()
+    torch.manual_seed(0)
+    layer = EngramLayer(args, 1)
+    T, hc, dim = 5, args.hc_mult, args.dim
+    wkv = torch.randn(dim * (hc + 1), layer.width, dtype=torch.bfloat16) * 0.05
+    layer.wkv.weight = wkv
+    layer.q_weight = torch.rand(hc, dim) + 0.5
+    layer.k_weight = torch.rand(hc, dim) + 0.5
+    x = torch.randn(T, hc, dim, dtype=torch.bfloat16)
+    rows = torch.randn(T, layer.width, dtype=torch.bfloat16)
+
+    kv = torch.nn.functional.linear(rows, wkv)
+    key, value = kv.split([hc * dim, dim], dim=-1)
+    key = key.float().view(T, hc, dim)
+    weight = layer.q_weight * layer.k_weight
+    h = x.float()
+    rstd = torch.rsqrt(h.square().mean(-1) + args.norm_eps) * torch.rsqrt(key.square().mean(-1) + args.norm_eps)
+    dot = (h * weight * key).sum(-1) * rstd * dim**-0.5
+    gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot))
+    want = (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)
+    assert torch.equal(layer.forward(x, rows), want)
diff --git a/tests/models/deepseek_v41/test_forward_smoke.py b/tests/models/deepseek_v41/test_forward_smoke.py
new file mode 100644
index 0000000000..e0861e8874
--- /dev/null
+++ b/tests/models/deepseek_v41/test_forward_smoke.py
@@ -0,0 +1,184 @@
+"""End-to-end forward smoke on the tiny model: cold prefill, batched decode, a second request, chunked
+prefill and bounded replay all produce finite logits of the right shape (numerics: test_reference_parity)."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from .common import VOCAB, requires_cuda, write_tiny_checkpoint
+
+pytestmark = requires_cuda
+
+
+@pytest.fixture(scope="module")
+def checkpoint(tmp_path_factory):
+    folder = tmp_path_factory.mktemp("dsv41-tiny")
+    write_tiny_checkpoint(str(folder))
+    return str(folder)
+
+
+def _tokens(n: int, seed: int) -> list[int]:
+    g = torch.Generator().manual_seed(seed)
+    return torch.randint(3, VOCAB, (n,), generator=g).tolist()
+
+
+def test_prefill_and_decode(checkpoint):
+    from .harness import TinyEngine
+
+    eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact")
+    r0 = eng.new_request(0, _tokens(300, 1))
+    logits = eng.prefill([r0])
+    assert logits.shape == (1, VOCAB) and torch.isfinite(logits).all()
+    eng.finish_prefill([r0])
+    r1 = eng.new_request(1, _tokens(50, 2))
+    assert torch.isfinite(eng.prefill([r1])).all()
+    eng.finish_prefill([r1])
+    for step in range(5):  # batched decode over two requests at different positions
+        out = eng.decode([r0, r1], [10 + step, 20 + step])
+        assert out.shape == (2, VOCAB) and torch.isfinite(out).all()
+
+
+def test_chunked_prefill_matches_single_shot(checkpoint):
+    """Two 128-aligned chunks must reproduce the single-shot prefill's last-token logits: the
+    compressor carry, the window ring and the compressed rows are all resumed through the pool. The
+    chunks run their projections at other row counts, so the match is to rounding, not bitwise."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact")
+    toks = _tokens(300, 3)
+    whole = eng.new_request(0, toks)
+    ref = eng.prefill([whole])
+    chunked = eng.new_request(1, toks)
+    chunked.device_len = 256  # first chunk [0, 256)
+    eng.prefill([chunked])
+    chunked.cached_len, chunked.device_len = 256, 300
+    got = eng.prefill([chunked])
+    _compare("chunked prefill", got, ref)
+
+
+def test_bounded_replay_is_exact_within_one_window(checkpoint):
+    """A prompt no longer than the window replays every token with no floor: bounded == exact."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    toks = _tokens(100, 5)
+    exact = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact")
+    want = exact.prefill([exact.new_request(0, toks)])
+    bounded = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded")
+    got = bounded.prefill([bounded.new_request(0, toks)])
+    _compare("bounded within one window", got, want)
+
+
+@pytest.mark.parametrize("cut", [128, 256])
+def test_bounded_chunked_prefill_matches_unchunked(checkpoint, cut):
+    """A chunked bounded prefill gives each chunk its part of the prompt's last window ``[172, 300)``:
+    a first chunk that ends before it (cut 128) runs no decoder row and hands the head a placeholder;
+    one that reaches into it (cut 256) leaves rows in the private ring for the last chunk, which is
+    shorter than the window (44 tokens). Both match the unchunked prefill."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded")
+    toks = _tokens(300, 6)
+    ref = eng.prefill([eng.new_request(0, toks)])
+    chunked = eng.new_request(1, toks)
+    chunked.device_len = cut
+    first = eng.prefill([chunked])
+    assert first.shape == (1, VOCAB) and torch.isfinite(first).all()
+    chunked.cached_len, chunked.device_len = cut, 300
+    _compare(f"bounded chunks cut at {cut}", eng.prefill([chunked]), ref)
+
+
+def test_bounded_batch_mixes_a_decoderless_chunk_with_a_whole_prompt(checkpoint):
+    """One batch: a chunk that ends before its prompt's last window (no decoder row, a placeholder for
+    the head) next to a whole prompt; the whole prompt's logits are those of its own prefill."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded")
+    whole = _tokens(300, 8)
+    want = eng.prefill([eng.new_request(0, whole)])
+    chunk = eng.new_request(0, _tokens(600, 9))
+    chunk.device_len = 256
+    other = eng.new_request(1, whole)
+    got = eng.prefill([chunk, other])
+    assert got.shape == (2, VOCAB) and torch.isfinite(got).all()
+    _compare("whole prompt next to a decoderless chunk", got[1:], want)
+
+
+def test_bounded_prefix_hit_matches_cold_prefill(checkpoint):
+    """A prefix hit that would leave fewer than a window of new tokens (768 cached + 20 new) is capped
+    one window before the prompt end, so the request prefills ``[640, 788)`` -- the decoder's whole
+    window among it -- and produces the cold prefill's logits."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded")
+    prefix = _tokens(768, 11)
+    donor = eng.new_request(0, prefix + _tokens(20, 12))
+    cold = eng.prefill([donor])
+    eng.finish_prefill([donor])
+    hit = eng.new_request_on_prefix(1, donor, prefix + _tokens(20, 12))  # the same prompt over the shared pages
+    assert hit.cached_len == 640 and hit.extend_len == 148
+    got = eng.prefill([hit])
+    _compare("bounded prefix hit", got, cold)
+
+
+def test_shared_prefix_replays_do_not_disturb_each_other(checkpoint):
+    """Two requests sharing a 768-token prefix with different suffix lengths: B's hit is capped at 640
+    (a window before its prompt end), its decoder writes only B's private ring and it reads the shared
+    encoder history, so A's pending decode sees the logits it would have seen without B, and neither
+    the shared pages nor A's ring change."""
+    from .harness import TinyEngine
+
+    prefix = _tokens(768, 13)
+    a_toks, b_toks = prefix + _tokens(20, 14), prefix + _tokens(40, 15)
+
+    def run(interleave: bool):
+        eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded")
+        a = eng.new_request(0, a_toks)
+        eng.prefill([a])
+        eng.finish_prefill([a])
+        dec_layer, enc_layer = eng.args.decoder_start_layer, 2
+        shared_slots = eng.pool.translate_full_to_window(eng.page_table[0, 640:768].long())
+        enc_before = eng.pool.window_pool[enc_layer][shared_slots].clone()
+        P = eng.args.window_size
+        ring_before = eng.pool.window_pool[dec_layer][0 * P : 1 * P].clone()  # A's ring row
+        if interleave:
+            b = eng.new_request_on_prefix(1, a, b_toks)
+            eng.prefill([b])
+            eng.finish_prefill([b])
+        assert torch.equal(eng.pool.window_pool[enc_layer][shared_slots], enc_before)
+        assert torch.equal(eng.pool.window_pool[dec_layer][0 * P : 1 * P], ring_before)
+        return eng.decode([a], [7])
+
+    assert torch.equal(run(True), run(False))
+
+
+def test_commit_dedup_onto_a_longer_prompts_pages_keeps_the_decoder_state(checkpoint):
+    """The radix commit may replace a request's freshly written pages with an existing node's (the
+    reviewer's donor case): D prefilled 256 tokens, so under bounded replay its decoder never computed
+    positions [0, 128); A (exactly D's first 128 tokens) is repointed onto D's pages at commit. A's
+    decoder KV lives in A's own ring, so its next decode is unchanged by the repoint (D's shared rows
+    came from a longer prefill, so to rounding)."""
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare
+
+    donor_toks = _tokens(256, 21)
+    a_toks = donor_toks[:128]
+
+    def run(repoint: bool):
+        eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded")
+        d = eng.new_request(0, donor_toks)
+        eng.prefill([d])
+        eng.finish_prefill([d])
+        a = eng.new_request(1, a_toks)
+        eng.prefill([a])
+        eng.finish_prefill([a])
+        if repoint:  # what CacheManager.cache_req does when the tree already holds the prefix
+            eng.page_table[1, :128] = eng.page_table[0, :128]
+        return eng.decode([a], [3])
+
+    _compare("decode after a commit repoint", run(True), run(False))
diff --git a/tests/models/deepseek_v41/test_reference_parity.py b/tests/models/deepseek_v41/test_reference_parity.py
new file mode 100644
index 0000000000..839cdafb82
--- /dev/null
+++ b/tests/models/deepseek_v41/test_reference_parity.py
@@ -0,0 +1,219 @@
+"""Logits parity against the vendored reference implementation on the tiny synthetic checkpoint.
+
+The reference runs bf16 or FP8/FP4 weights with its own torch quantizers, sparse attention, engram
+hash and single-pass mHC; FreeToken runs the same checkpoint through its kernels and packed pools.
+Compared: exact-mode prefill, greedy decode steps, a two-request batch, and the bounded-replay path
+against the reference's ``forward_bounded`` oracle.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, fields
+
+import pytest
+import torch
+
+from .common import VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint
+
+pytestmark = requires_cuda
+
+# The sparse-attention implementations accumulate in different orders.
+ATOL = RTOL = 1e-2
+# The fp4 experts also round the routed sum to bf16 before the shared-expert add and weight each
+# expert's down output; the reference keeps that sum fp32 and weights the intermediate before its fp8
+# round-trip (10 seeds x 16 checks: max |err| 0.019, against 0.007 in the reference's order).
+QUANT_TOL = 2e-2
+
+
+@pytest.fixture(scope="module")
+def checkpoint(tmp_path_factory):
+    folder = tmp_path_factory.mktemp("dsv41-parity")
+    tensors = write_tiny_checkpoint(str(folder), seed=11)
+    return str(folder), tensors
+
+
+def _tokens(n: int, seed: int) -> torch.Tensor:
+    return torch.randint(3, VOCAB, (n,), generator=torch.Generator().manual_seed(seed))
+
+
+class Reference:
+    """The vendored ``Transformer`` on CUDA, loaded from the synthetic checkpoint (identity token map)."""
+
+    def __init__(self, tensors: dict, text: dict, *, max_seq_len: int, max_batch_size: int, quantized: bool = False, vision_config: dict | None = None):
+        from freetoken.models.deepseek_v41.args import DeepseekV41Args
+        from types import SimpleNamespace
+
+        from . import reference
+        from .reference import engram as ref_engram
+        from .reference import model as ref_model
+
+        args = DeepseekV41Args.from_hf(SimpleNamespace(text_config=text))
+        names = {f.name for f in fields(ref_model.ModelArgs)}
+        kwargs = {k: v for k, v in asdict(args).items() if k in names}
+        kwargs.update(dtype="fp8" if quantized else "bf16", expert_dtype="fp4" if quantized else None,
+                      max_seq_len=max_seq_len, max_batch_size=max_batch_size, vision_n_layers=0, temperature=0.0)
+        if vision_config:
+            kwargs.update(vision_n_layers=vision_config["num_hidden_layers"], vision_dim=vision_config["hidden_size"],
+                          vision_n_heads=vision_config["num_attention_heads"], vision_inter_dim=vision_config["intermediate_size"],
+                          vision_patch_size=vision_config["patch_size"], vision_downsample_ratio=vision_config["downsample_ratio"])
+        self.margs = ref_model.ModelArgs(**kwargs)
+        # the synthetic checkpoint's tokenizer is the identity over VOCAB ids
+        ref_engram.build_compressed_token_map = lambda tokenizer: (list(range(VOCAB)), VOCAB)
+        torch.set_default_dtype(torch.bfloat16)
+        try:
+            with torch.device("cuda"):
+                self.model = ref_model.Transformer(self.margs, tokenizer=object())
+        finally:
+            torch.set_default_dtype(torch.float32)
+        state = {k: v for k, v in tensors.items() if not k.startswith("mtp.")
+                 and (vision_config or (not k.startswith("vision.") and not k.endswith("bias_vl")))}
+        if quantized:
+            state = {k: v.view(torch.float4_e2m1fn_x2) if ".ffn.experts." in k and k.endswith(".weight") else v
+                     for k, v in state.items()}
+            # The reference's grouped output projection holds dequantized bf16 weights.
+            for name in list(state):
+                if name.endswith(".wo_a.scale"):
+                    scale = state.pop(name).float().repeat_interleave(32, 0).repeat_interleave(32, 1)
+                    weight_name = name.removesuffix(".scale") + ".weight"
+                    state[weight_name] = (state[weight_name].float() * scale).bfloat16()
+        missing, unexpected = self.model.load_state_dict({k: v.cuda() for k, v in state.items()}, strict=False)
+        assert not unexpected, unexpected
+        assert not [m for m in missing if "freqs_cis" not in m], missing
+        self.reference = reference
+
+    # the reference builds its index tensors on the default device (generate.py sets it to cuda)
+    def prefill(self, ids: torch.Tensor) -> torch.Tensor:
+        with torch.device("cuda"):
+            return self.model(ids.view(1, -1).cuda(), 0)[1].float()
+
+    def prefill_batch(self, ids: torch.Tensor) -> torch.Tensor:
+        with torch.device("cuda"):
+            return self.model(ids.cuda(), 0)[1].float()
+
+    def prefill_bounded(self, ids: torch.Tensor) -> torch.Tensor:
+        with torch.device("cuda"):
+            return self.model.forward_bounded(ids.reshape(-1, ids.shape[-1]).cuda(), self.margs.window_size).float()
+
+    def decode(self, tokens: torch.Tensor, pos: int) -> torch.Tensor:
+        with torch.device("cuda"):
+            return self.model(tokens.view(-1, 1).cuda(), pos)[1].float()
+
+
+def _engram_table_for(engine, tensors, text):
+    """Our engine reads the Engram rows straight from the synthetic shard (the disk table), hashed
+    with the same identity token map the reference oracle uses."""
+    from freetoken.models.deepseek_v41.engram import EngramHash
+    from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source
+
+    hash = EngramHash(engine.args, list(range(VOCAB)), VOCAB)
+    tables = [
+        EngramDiskTable(engram_row_source(engine.checkpoint, layer.layer_id), hash.n_cols, engine.device, max_graph_rows=8, max_extend_tokens=64)
+        for layer in engine.model.engram_layers()
+    ]
+    for layer, table in zip(engine.model.engram_layers(), tables):
+        layer.attach_table(table)
+    host = EngramHost(hash, tables, engine.device)
+    engine.model.forward_host_ctx = host.forward_host_ctx
+    engine.engram_host = host
+
+
+def _compare(name, got, want, *, tol=ATOL):
+    err = (got - want).abs().max().item()
+    assert torch.equal(got.argmax(-1), want.argmax(-1)), f"{name}: argmax differs (max abs err {err:.4f})"
+    torch.testing.assert_close(got, want, atol=tol, rtol=tol, msg=lambda m: f"{name}: {m}")
+
+
+@pytest.mark.parametrize("mode", ["exact", "bounded"])
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_quantized_offload_prefill_and_decode_match_reference(tmp_path, mode, batch_size):
+    from .harness import TinyEngine
+
+    tensors = write_tiny_checkpoint(str(tmp_path), seed=11, quantized=True)
+    text = tiny_text_config(moe_intermediate_size=256)
+    ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3, quantized=True)
+    eng = TinyEngine(str(tmp_path), swa_decoder_replay=mode, quantized=True)
+    _engram_table_for(eng, tensors, text)
+    ids = torch.stack([_tokens(300, 1 + i) for i in range(batch_size)])
+    want = ref.prefill_batch(ids) if mode == "exact" else ref.prefill_bounded(ids)
+    reqs = [eng.new_request(i, row.tolist()) for i, row in enumerate(ids)]
+    got = eng.prefill(reqs)
+    _compare(f"quantized {mode} prefill", got, want, tol=QUANT_TOL)
+    eng.finish_prefill(reqs)
+    for step in range(3):
+        token = want.argmax(-1)
+        want = ref.decode(token, ids.shape[-1] + step)
+        got = eng.decode(reqs, token.tolist())
+        _compare(f"quantized {mode} decode {step}", got, want, tol=QUANT_TOL)
+
+
+def test_exact_prefill_and_greedy_decode_match_the_reference(checkpoint):
+    from .harness import TinyEngine
+
+    folder, tensors = checkpoint
+    text = tiny_text_config()
+    ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3)
+    eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact")
+    _engram_table_for(eng, tensors, text)
+
+    ids = _tokens(300, 1)
+    want = ref.prefill(ids)
+    req = eng.new_request(0, ids.tolist())
+    got = eng.prefill([req])
+    _compare("prefill", got, want)
+    eng.finish_prefill([req])
+    # greedy decode from the reference's own picks keeps both sides on one trajectory
+    pos = 300
+    for step in range(12):
+        nxt = int(want.argmax(-1).item())
+        want = ref.decode(torch.tensor([nxt]), pos)
+        got = eng.decode([req], [nxt])
+        _compare(f"decode step {step}", got, want)
+        pos += 1
+
+
+def test_batched_prefill_and_decode_match_the_reference(checkpoint):
+    from .harness import TinyEngine
+
+    folder, tensors = checkpoint
+    text = tiny_text_config()
+    ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3)
+    eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact")
+    _engram_table_for(eng, tensors, text)
+
+    ids = torch.stack([_tokens(200, 2), _tokens(200, 3)])
+    want = ref.prefill_batch(ids)
+    reqs = [eng.new_request(i, ids[i].tolist()) for i in range(2)]
+    got = eng.prefill(reqs)
+    _compare("batched prefill", got, want)
+    eng.finish_prefill(reqs)
+    pos = 200
+    for step in range(6):
+        nxt = want.argmax(-1).cpu()
+        want = ref.decode(nxt, pos)
+        got = eng.decode(reqs, nxt.tolist())
+        _compare(f"batched decode step {step}", got, want)
+        pos += 1
+
+
+def test_bounded_replay_matches_the_reference_oracle(checkpoint):
+    from .harness import TinyEngine
+
+    folder, tensors = checkpoint
+    text = tiny_text_config()
+    ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3)
+    eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded")
+    _engram_table_for(eng, tensors, text)
+
+    ids = _tokens(300, 4)
+    want = ref.prefill_bounded(ids)
+    req = eng.new_request(0, ids.tolist())
+    got = eng.prefill([req])
+    _compare("bounded prefill", got, want)
+    eng.finish_prefill([req])
+    pos = 300
+    for step in range(6):
+        nxt = int(want.argmax(-1).item())
+        want = ref.decode(torch.tensor([nxt]), pos)
+        got = eng.decode([req], [nxt])
+        _compare(f"decode after bounded prefill, step {step}", got, want)
+        pos += 1
diff --git a/tests/models/deepseek_v41/test_vision.py b/tests/models/deepseek_v41/test_vision.py
new file mode 100644
index 0000000000..0e1d4a4830
--- /dev/null
+++ b/tests/models/deepseek_v41/test_vision.py
@@ -0,0 +1,214 @@
+"""Image layout, official ViT/logits parity and image replay through the shared encoder cache."""
+
+from __future__ import annotations
+
+import io
+import json
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+
+from freetoken.core import Batch
+from freetoken.engine.engine import Engine
+from freetoken.mm.config import MultimodalConfig
+from freetoken.mm.encoder_cache import EncoderCache
+from freetoken.mm.processors.deepseek_v41 import DeepseekV41MMProcessor
+from freetoken.models.deepseek_v41.config import VisionConfig
+from freetoken.scheduler.mm import mm_rows_after, plan_mm_batch
+from freetoken.utils.torch_utils import torch_dtype
+
+from .common import DIM, VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint
+from .reference.vision import Aligner, ViT
+
+VC = asdict(VisionConfig(num_hidden_layers=2, hidden_size=32, num_attention_heads=4, intermediate_size=48,
+                         patch_size=2, min_pixels=0, max_image_tokens=64))
+
+
+def _args():
+    return SimpleNamespace(vision_n_layers=2, vision_dim=32, vision_n_heads=4, vision_inter_dim=48,
+                           vision_patch_size=2, vision_downsample_ratio=3, vision_rope_theta=10000.0, dim=DIM)
+
+
+def _checkpoint(path, quantized=False):
+    tensors = write_tiny_checkpoint(str(path), seed=11, quantized=quantized)
+    torch.manual_seed(12)
+    with torch_dtype(torch.bfloat16):
+        vit, aligner = ViT(_args()), Aligner(_args())
+    # Replace the unused dummy vision weight from the text-only fixture.
+    tensors = {k: v for k, v in tensors.items() if not k.startswith("vision.")}
+    for prefix, module in (("vision", vit), ("aligner", aligner)):
+        tensors.update({f"{prefix}.{k}": v.detach().bfloat16().contiguous() for k, v in module.state_dict().items()})
+    tensors.update({k: torch.randn(DIM, dtype=torch.bfloat16) for k in ("image_start", "image_end", "image_newline")})
+    save_file(tensors, path / 'model-00001-of-00001.safetensors')
+    (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {k: 'model-00001-of-00001.safetensors' for k in tensors}}))
+    config = json.loads((path / 'config.json').read_text())
+    config.update(vision_config=VC, image_token_id=VOCAB - 1)
+    (path / 'config.json').write_text(json.dumps(config))
+    return tensors
+
+
+def _processor():
+    cfg = SimpleNamespace(vision_config=SimpleNamespace(**VC), image_token_id=VOCAB - 1)
+    return DeepseekV41MMProcessor(cfg, '.', MultimodalConfig())
+
+
+def _image(color):
+    buf = io.BytesIO()
+    Image.new('RGB', (14, 10), color).save(buf, format='PNG')
+    return buf.getvalue()
+
+
+def _inputs(length=151):
+    ids = [3] * 17 + [VOCAB - 1] + [4] * (length - 19) + [VOCAB - 1]
+    return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image('red'), _image('blue')])
+
+
+def _prefill(eng, reqs):
+    if not hasattr(eng, 'encoder_cache'):
+        eng.encoder_cache = EncoderCache()
+        eng.dtype = torch.bfloat16
+        eng._mm_registered = set()
+    for req in reqs:
+        if id(req) not in eng._mm_registered:
+            for item in req.mm_items or ():  # what the scheduler claims at admission
+                eng.encoder_cache.register(item.hash, req.uid, mm_rows_after(item, req.cached_len))
+            eng._mm_registered.add(id(req))
+    batch = Batch(reqs=reqs, phase='prefill')
+    batch.padded_reqs = reqs
+    batch.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs]).cuda()
+    batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).cuda()
+    jobs, plan, rows, _ = plan_mm_batch(reqs, eng.encoder_cache)
+    batch.mm_encoder_jobs, batch.mm_gather_plan = jobs, plan
+    if plan:
+        batch.mm_rows = torch.tensor(rows, dtype=torch.long, device=eng.device)
+        Engine._run_mm_encoder(eng, batch)
+    eng._bind()
+    eng.backend.prepare_metadata(batch)
+    with eng.ctx.forward_batch(batch), eng.model.forward_host_ctx(batch, False):
+        return eng.model.forward().float()
+
+
+def test_processor_patch_order_and_all_span_rows_are_embeddings():
+    pixels = np.arange(10 * 14 * 3, dtype=np.uint8).reshape(10, 14, 3)
+    item = _processor().process([Image.fromarray(pixels)])[0]
+    assert item.grid_thw == [1, 5, 7]
+    expected = ((torch.tensor(pixels[:2, :2]).permute(2, 0, 1).float() / 255 - .5) / .5).bfloat16()
+    assert torch.equal(item.feature[0], expected)
+    assert len(_processor().prompt_replacement(item).full) == 10
+    result = _inputs()
+    assert [i.offsets for i in result.mm_items] == [[[17, 27]], [[159, 169]]]
+    assert result.mrope_positions is None
+    assert result.mm_items[0].hash != result.mm_items[1].hash
+
+
+def test_text_only_skips_every_visual_parameter(tmp_path):
+    from freetoken.models.deepseek_v41.weight import iter_weights
+    from freetoken.mm.processor import get_mm_processor
+    from .test_engine_config import _engine_config
+    _checkpoint(tmp_path)
+    enabled = _engine_config(str(tmp_path), moe_strategy='offload')
+    assert enabled.model_config.is_multimodal
+    disabled = MultimodalConfig(disabled_encoders=frozenset({'vision', 'audio'}))
+    text = _engine_config(str(tmp_path), moe_strategy='offload', mm=disabled)
+    assert not text.model_config.is_multimodal and not text.active_encoders
+    assert get_mm_processor(str(tmp_path), disabled) is None
+    loaded = dict(iter_weights(str(tmp_path), 'cpu', include_moe_experts=False, include_vision=False))
+    assert not any(k.startswith('visual.') for k in loaded)
+    assert get_mm_processor(str(tmp_path)).placeholder == [VOCAB - 1]
+
+
+@requires_cuda
+@pytest.mark.parametrize('placement', ['gpu', 'host'])
+def test_vision_encoder_matches_official_with_padded_grid(tmp_path, placement):
+    from .harness import TinyEngine
+    tensors = _checkpoint(tmp_path)
+    eng = TinyEngine(str(tmp_path))
+    eng.model.place_encoder_weights(placement)
+    item = _processor().process([Image.new('RGB', (14, 10), 'red')])[0]
+    with torch_dtype(torch.bfloat16), torch.device('cuda'):
+        vit, aligner = ViT(_args()), Aligner(_args())
+    for prefix, module in (('vision', vit), ('aligner', aligner)):
+        module.load_state_dict({k.removeprefix(prefix + '.'): v.cuda() for k, v in tensors.items() if k.startswith(prefix + '.')})
+    with torch.inference_mode(), torch.device('cuda'):
+        x = aligner(vit(item.feature.cuda(), 5, 7), 5, 7)
+    got = eng.model.encode(item)
+    torch.testing.assert_close(got[1:-1].view(2, 4, DIM)[:, :3].reshape(-1, DIM), x, atol=1e-3, rtol=1e-3)
+    assert torch.equal(got[0], tensors['image_start'].cuda())
+    assert torch.equal(got[-1], tensors['image_end'].cuda())
+    assert torch.equal(got[[4, 8]], tensors['image_newline'].cuda().expand(2, -1))
+    eng.model.place_encoder_weights('gpu')
+
+
+@requires_cuda
+@pytest.mark.parametrize('quantized', [False, True])
+def test_image_logits_and_decode_match_reference(tmp_path, quantized):
+    from .harness import TinyEngine
+    from .test_reference_parity import ATOL, QUANT_TOL, Reference, _compare, _engram_table_for
+    tol = QUANT_TOL if quantized else ATOL
+    tensors = _checkpoint(tmp_path, quantized)
+    text = tiny_text_config(moe_intermediate_size=256) if quantized else tiny_text_config()
+    ref = Reference(tensors, text, max_seq_len=512, max_batch_size=2, quantized=quantized, vision_config=VC)
+    eng = TinyEngine(str(tmp_path), max_seq_len=512, swa_decoder_replay='exact', quantized=quantized)
+    _engram_table_for(eng, tensors, text)
+    result = _inputs()
+    ids = result.input_ids.clone()
+    types = torch.full_like(ids, -1)
+    images = []
+    for item in result.mm_items:
+        lo, hi = item.offsets[0]
+        ids[lo:hi] = VOCAB - 1
+        typ = torch.tensor([0] + ([1] * 3 + [2]) * 2 + [3])
+        types[lo:hi] = typ
+        images.append(SimpleNamespace(start=lo, patches=item.feature, n_vit_h=5, n_vit_w=7, types=typ))
+    with torch.device('cuda'):
+        want = ref.model(ids.cuda()[None], images=[images], token_types=types.cuda()[None])[1]
+    req = eng.new_request(0, result.input_ids.tolist())
+    req.mm_items = result.mm_items
+    got = _prefill(eng, [req])
+    _compare('image prefill', got, want, tol=tol)
+    eng.finish_prefill([req])
+    for step in range(3):
+        token = want.argmax(-1)
+        want = ref.decode(token, len(ids) + step)
+        got = eng.decode([req], token.tolist())
+        _compare(f'image decode {step}', got, want, tol=tol)
+
+
+@requires_cuda
+@pytest.mark.parametrize('mode', ['exact', 'bounded'])
+def test_images_survive_chunking_and_prefix_replay(tmp_path, mode):
+    from .harness import TinyEngine
+    from .test_reference_parity import _compare, _engram_table_for
+    tensors = _checkpoint(tmp_path)
+    eng = TinyEngine(str(tmp_path), max_seq_len=2048, swa_decoder_replay=mode)
+    _engram_table_for(eng, tensors, tiny_text_config())
+    # bounded mode caps a hit a window before the prompt end: trailing text keeps that cap at 128
+    tail = 150 if mode == 'bounded' else 0
+    def inputs():
+        ids = [3] * 80 + [VOCAB - 1] + [4] * 37 + [VOCAB - 1] + [4] * 31 + [VOCAB - 1] + [4] * tail
+        return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image(c) for c in ('red', 'blue', 'green')])
+    whole = inputs()
+    assert [i.offsets for i in whole.mm_items] == [[[80, 90]], [[127, 137]], [[168, 178]]]
+    a = eng.new_request(0, whole.input_ids.tolist())
+    a.mm_items = whole.mm_items
+    want = _prefill(eng, [a])
+    chunk = inputs()
+    b = eng.new_request(1, chunk.input_ids.tolist())
+    b.mm_items = chunk.mm_items
+    b.device_len = 128
+    _prefill(eng, [b])
+    b.cached_len, b.device_len = 128, len(chunk.input_ids)
+    got = _prefill(eng, [b])
+    _compare('chunked image prefill', got, want)
+    prefix = inputs()
+    # The second image straddles the chunk / prefix boundary at 128 (in bounded mode the first chunk
+    # also ends before the prompt's last window and runs no decoder row).
+    hit = eng.new_request_on_prefix(1, a, prefix.input_ids.tolist())
+    assert hit.cached_len == 128
+    hit.mm_items = prefix.mm_items
+    _compare('image prefix hit', _prefill(eng, [hit]), want)
diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py
index 084a2e4aa7..bfcfb87664 100644
--- a/tests/models/test_quant_config.py
+++ b/tests/models/test_quant_config.py
@@ -1,6 +1,7 @@
 """What each checkpoint's quant config resolves to: the dialect, the method behind every layer, and agreement with the stored tensors.
 
-Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip.
+Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip. The
+stored-tensor scan also reads $FREETOKEN_TEST_MODELS_ROOT.
 """
 
 from __future__ import annotations
@@ -47,6 +48,8 @@
 from freetoken.utils.torch_utils import torch_dtype
 
 MODELS = "/mnt/nvme/models"
+# the stored-tensor scan takes every checkpoint under this root (the case table above matches by name)
+SCAN_ROOT = os.environ.get("FREETOKEN_TEST_MODELS_ROOT", MODELS)
 HF_CACHE = os.path.expanduser("~/.cache/huggingface/hub")
 
 BF16, FP8B, FP8T, MXFP8, NVFP4 = UnquantizedLinearMethod, Fp8BlockLinearMethod, Fp8TensorLinearMethod, Mxfp8LinearMethod, Nvfp4LinearMethod
@@ -435,6 +438,25 @@ def test_modelopt_nvfp4_reads_the_activation_quantizer_from_config_groups(extra,
     assert scheme.kind is QuantKind.NVFP4 and scheme.has("input_scale") is has_input_scale
 
 
+def test_fp8_dialect_reads_the_block_size_and_the_nested_expert_dtype():
+    """DeepSeek-V4.1 keeps the ``fp8`` dialect with 32x32 blocks and puts ``expert_dtype`` inside quantization_config."""
+    from freetoken.layers.quantization.scheme import fp8_block_size
+
+    q = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"}
+    quant = QuantConfig.from_hf(SimpleNamespace(quantization_config=q))
+    assert type(quant) is Fp8BlockConfig and quant.block == 32
+    dense = quant.scheme_for("layers.3.attn.wq_a")
+    assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32 and dense.weight.scale == "e8m0"
+    assert quant.storage(dense)["weight_scale_inv"].name == "scale"
+    assert quant.scheme_for("layers.3.ffn.experts.7.w1").kind is QuantKind.MXFP4
+    # the V4 128-block export still resolves to the 128 scheme
+    q128 = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"}
+    assert fp8_block_size(QuantConfig.from_hf(SimpleNamespace(quantization_config=q128)).scheme_for("layers.3.attn.wq_a")) == 128
+    for bad in ([64, 64], [32, 128], [128]):
+        with pytest.raises(NotImplementedError):
+            QuantConfig.from_hf(SimpleNamespace(quantization_config={**q, "weight_block_size": bad}))
+
+
 def test_every_dialect_names_the_tensors_behind_its_schemes():
     from freetoken.layers.quantization.registry import dialects
 
@@ -491,7 +513,7 @@ def _weight_map(p: Path) -> dict[str, str] | None:
 
 
 def _candidate_dirs() -> list[Path]:
-    roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(MODELS, "*/"))
+    roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(SCAN_ROOT, "*/"))
     return [Path(d) for d in roots if (Path(d) / "config.json").exists() and _weight_map(Path(d)) is not None]
 
 
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index f496de3142..547c94c60a 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -391,6 +391,7 @@ def test_cpu_decode_dsfp4_matches_gpu(bs):
         max_tokens=bs,
         device=dev,
         swiglu_limit=swiglu_limit,
+        act_block=128,
     )
 
     hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16)
@@ -585,7 +586,7 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4():
     ex = CpuMoeExecutor(
         cache, top_k=top_k, activation="silu",
         apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev,
-        swiglu_limit=limit,
+        swiglu_limit=limit, act_block=128,
     )
     b = cache.bank_sources
     layer_banks = [b[n][layer].to(dev) for n in (
diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py
index 6d6b49d395..11519502e2 100644
--- a/tests/moe/test_dsfp4_prequant.py
+++ b/tests/moe/test_dsfp4_prequant.py
@@ -46,7 +46,7 @@ def test_dsfp4_gpu_prequant_bit_parity():
     def run(cpu_roundtrip: bool) -> torch.Tensor:
         ex = CpuMoeExecutor(
             cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False,
-            num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0,
+            num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128,
         )
         assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant"
         if cpu_roundtrip:  # control: the original CPU-side scalar round-trip
diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py
new file mode 100644
index 0000000000..2f89cf334d
--- /dev/null
+++ b/tests/moe/test_dsfp4_reference_semantics.py
@@ -0,0 +1,170 @@
+"""ds_fp4 (DeepSeek W4A8) experts against an INDEPENDENT torch transcription over
+``fp4_gemm(act_quant(x, block), W)`` at both fp8 activation blocks (128 on V4, 32 on V4.1):
+
+    x_q  = fp8_roundtrip(x, block)                       # act_quant(..., inplace)
+    gate = bf16(x_q @ deq(W1)^T), up = bf16(x_q @ deq(W3)^T)   # fp4_gemm outputs bf16
+    h    = bf16(silu(min(gate, L)) * clamp(up, -L, L))
+    y_r  = bf16(w_r * (fp8_roundtrip(h, block) @ deq(W2)^T))   # routing weight in the down epilogue
+    y    = sum_r y_r
+
+The routing weight scales the down output (the placement main serves V4 with); the reference
+``Expert.forward`` scales the intermediate before its fp8 quant instead. The GEMV (decode), the
+grouped GEMM (prefill) and the CPU executor all follow this transcription.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+E, H, I, TOP_K, LIMIT = 16, 512, 256, 4, 7.0
+E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
+
+
+def _banks(device, seed=0):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+
+    def u8(*shape, low=0, high=256):
+        return torch.randint(low, high, shape, dtype=torch.uint8, generator=g).to(device)
+
+    return (u8(E, 2 * I, H // 2), u8(E, 2 * I, H // 32, low=120, high=130), u8(E, H, I // 2), u8(E, H, I // 32, low=120, high=130))
+
+
+def _dequant(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """[E, N, K//2] e2m1 pairs (even channel low nibble) + [E, N, K//32] e8m0 -> [E, N, K] fp32."""
+    lut = E2M1.to(packed.device)
+    lo, hi = lut[(packed & 0xF).long()], lut[(packed >> 4).long()]
+    vals = torch.stack([lo, hi], dim=-1).flatten(-2)  # [E, N, K]
+    s = torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(32, dim=-1)
+    return vals * s
+
+
+def _fp8_roundtrip(x: torch.Tensor, block: int) -> torch.Tensor:
+    g = x.float().unflatten(-1, (-1, block))
+    s = torch.exp2(torch.ceil(torch.log2(g.abs().amax(-1).clamp_min(1e-4) / 448.0)))
+    q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float()
+    return (q * s.unsqueeze(-1)).flatten(-2).to(torch.bfloat16)
+
+
+def reference(x, slots, weights, banks, block):
+    """The MoE over dequantized banks: per-route bf16 expert outputs accumulated into an fp32 ``y``,
+    returned in fp32 (the kernels round that sum to bf16)."""
+    gu_p, gu_s, dn_p, dn_s = banks
+    W13 = _dequant(gu_p, gu_s)  # [E, 2I, H]
+    W2 = _dequant(dn_p, dn_s)  # [E, H, I]
+    xq = _fp8_roundtrip(x, block).float()
+    T = x.shape[0]
+    y = torch.zeros(T, H, dtype=torch.float32, device=x.device)
+    for t in range(T):
+        for r in range(TOP_K):
+            e = int(slots[t, r])
+            gu = (xq[t] @ W13[e].T).to(torch.bfloat16).float()
+            gate, up = gu[:I].clamp(max=LIMIT), gu[I:].clamp(-LIMIT, LIMIT)
+            h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16)
+            hq = _fp8_roundtrip(h.view(1, -1), block).float().view(-1)
+            y[t] += (weights[t, r].float() * (hq @ W2[e].T)).to(torch.bfloat16).float()
+    return y
+
+
+def _inputs(T, device, seed=1):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    x = (torch.randn(T, H, dtype=torch.bfloat16, generator=g) * 0.5).to(device)
+    slots = torch.stack([torch.randperm(E, generator=g)[:TOP_K] for _ in range(T)]).to(device=device, dtype=torch.int32).contiguous()
+    w = torch.rand(T, TOP_K, generator=g)
+    w = (1.5 * w / w.sum(-1, keepdim=True)).float().to(device).contiguous()  # route_scale-style weights > 1 too
+    return x, slots, w
+
+
+@pytest.mark.parametrize("block", [32, 128])
+def test_gemv_path_matches_the_transcription(block):
+    from freetoken.moe.fused_ds_fp4 import routed_experts_fp4
+
+    banks = _banks("cuda")
+    x, slots, w = _inputs(6, "cuda")
+    got = routed_experts_fp4(x, slots, w, *banks, LIMIT, act_block=block).float()
+    want = reference(x, slots, w, banks, block)
+    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("block", [32, 128])
+def test_cpu_executor_matches_the_transcription(block):
+    """The ``ds_fp4`` CPU executor (cpu / hybrid decode) implements the same contract as the GPU path."""
+    from types import SimpleNamespace
+
+    from freetoken.kernel.pinned import alloc_pinned_tensor
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    gu_p, gu_s, dn_p, dn_s = _banks("cuda")
+    pinned = {}
+    for name, t in (("gate_up_packed", gu_p), ("gate_up_scale", gu_s), ("down_packed", dn_p), ("down_scale", dn_s)):
+        buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8)
+        buf.copy_(t.cpu())
+        pinned[name] = buf
+    cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E)
+    dev = torch.device("cuda", 0)
+    x, slots, w = _inputs(2, dev)
+    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=LIMIT, act_block=block)
+    got = ex.decode(0, x, w, slots.clone()).clone().float()
+    torch.cuda.synchronize()
+    del ex
+    want = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block)
+    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+
+
+def test_grouped_prefill_path_matches_the_transcription():
+    from freetoken.moe import fused_ds_fp4
+
+    banks = _banks("cuda")
+    x, slots, w = _inputs(64, "cuda")
+    got = fused_ds_fp4.routed_experts_fp4_prefill(x, slots, w, *banks, LIMIT, E, act_block=32).float()
+    want = reference(x, slots, w, banks, 32)
+    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("block", [32, 128])
+def test_inactive_routes_contribute_zero_without_reading_their_slots(block):
+    """Hybrid decode hands the GPU kernel slot -1 (weight 0) for the routes the CPU computes. The kernel
+    must produce exactly zero for them without touching any slot: every slot the active routes do not
+    use is filled with 0xFF (an e8m0 scale of NaN), and the result still equals the reference over the
+    active routes and the CPU executor's output for the same split."""
+    from types import SimpleNamespace
+
+    from freetoken.kernel.pinned import alloc_pinned_tensor
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+    from freetoken.moe.fused_ds_fp4 import routed_experts_fp4
+
+    gu_p, gu_s, dn_p, dn_s = (t.clone() for t in _banks("cuda"))
+    x, slots, w = _inputs(2, "cuda")
+    # routes: token 0 keeps routes 0..1 on the GPU, token 1 keeps route 3; the rest go to the CPU
+    on_gpu = torch.zeros_like(slots, dtype=torch.bool)
+    on_gpu[0, :2] = True
+    on_gpu[1, 3] = True
+    gpu_slots = torch.where(on_gpu, slots, -1)
+    gpu_w = torch.where(on_gpu, w, 0.0)
+    used = set(slots[on_gpu].tolist())
+    poison = [e for e in range(E) if e not in used]
+    for t in (gu_p, gu_s, dn_p, dn_s):
+        t[poison] = 0xFF
+    got = routed_experts_fp4(x, gpu_slots, gpu_w, gu_p, gu_s, dn_p, dn_s, LIMIT, act_block=block)
+    assert torch.isfinite(got).all()
+    clean = _banks("cuda")
+    # the poisoned-bank result equals the clean-bank result with the same holes (neither reads them),
+    # and each share matches the reference over its own routes (a zero weight drops a route)
+    got_clean = routed_experts_fp4(x, gpu_slots, gpu_w, *clean, LIMIT, act_block=block)
+    assert torch.equal(got, got_clean)
+    torch.testing.assert_close(got.float(), reference(x, slots, gpu_w, clean, block), atol=2e-2, rtol=2e-2)
+    # the CPU executor computes the complementary split from the same raw ids (ids < 0 skipped)
+    pinned = {}
+    for name, t in (("gate_up_packed", clean[0]), ("gate_up_scale", clean[1]), ("down_packed", clean[2]), ("down_scale", clean[3])):
+        buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8)
+        buf.copy_(t.cpu())
+        pinned[name] = buf
+    cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E)
+    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=torch.device("cuda", 0), swiglu_limit=LIMIT, act_block=block)
+    cpu_ids = torch.where(on_gpu, -1, slots)
+    cpu_out = ex.decode(0, x, w, cpu_ids.clone()).clone()
+    torch.cuda.synchronize()
+    del ex
+    torch.testing.assert_close(cpu_out.float(), reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block), atol=2e-2, rtol=2e-2)
diff --git a/tests/scheduler/test_dsv4_generic_manager.py b/tests/scheduler/test_dsv4_generic_manager.py
index ad69b89ec5..1bca10bc8a 100644
--- a/tests/scheduler/test_dsv4_generic_manager.py
+++ b/tests/scheduler/test_dsv4_generic_manager.py
@@ -14,8 +14,8 @@
 import torch
 
 from freetoken.core import Req, SamplingParams
-from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes
-from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache
+from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes
+from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache
 from freetoken.models.deepseek_v4.args import DeepseekV4Args
 from freetoken.scheduler.cache import CacheManager
 
@@ -191,6 +191,33 @@ def _pending(uid, n):
     cm.check_integrity()
 
 
+@pytest.mark.parametrize("replay, new_tokens, cached", [(0, 2, 256), (P, 2, 128), (P, 200, 256)])
+def test_admission_caps_the_hit_by_the_pools_prefix_replay(replay, new_tokens, cached):
+    """A pool whose model recomputes the prompt's tail (DeepSeek-V4.1's bounded decoder replay) holds
+    back ``prefix_replay_tokens`` at admission: the match ends at or before ``L - replay`` (page-aligned
+    by the radix), a hit already that far from the prompt end is untouched, and the request's
+    ``prompt_len`` is the whole prompt."""
+    from freetoken.scheduler.decode import DecodeManager
+    from freetoken.scheduler.prefill import PrefillManager
+    from freetoken.scheduler.table import TableManager
+    from freetoken.scheduler.utils import PendingReq
+
+    cm, pool, pt = _stack(num_pages=48)
+    pool.prefix_replay_tokens = replay
+    prompt = torch.arange(1, 301, dtype=torch.int32)
+    _lifecycle(cm, _req(MRR - 1, prompt, n_decode=310), total_len=310)
+    pm = PrefillManager(cm, TableManager(max_running_reqs=MRR, page_table=pt), DecodeManager(page_size=P))
+    tokens = torch.cat([prompt, torch.arange(1000, 1000 + new_tokens, dtype=torch.int32)])
+    pm.pending_list = [PendingReq(uid=1, input_ids=tokens, sampling_params=SamplingParams(max_tokens=1))]
+    batch = pm.schedule_next_batch(1024)
+    req = batch.reqs[0]
+    assert req.cached_len == cached and req.extend_len == len(tokens) - cached and req.prompt_len == len(tokens)
+    cm.allocate_paged(batch.reqs)
+    req.complete_one()
+    cm.cache_req(req, finished=True)
+    cm.check_integrity()
+
+
 def test_abort_anywhere_fuzz_conserves_and_isolates():
     """admit / extend / finish-or-abort at arbitrary points with shared prefixes + eviction
     pressure, on a deliberately small window tier. The generic conservation check (free + tree
diff --git a/tests/scheduler/test_scheduler_kv_usage.py b/tests/scheduler/test_scheduler_kv_usage.py
index 6ab949705e..bf0e825c31 100644
--- a/tests/scheduler/test_scheduler_kv_usage.py
+++ b/tests/scheduler/test_scheduler_kv_usage.py
@@ -2,6 +2,8 @@
 
 from types import SimpleNamespace
 
+from freetoken.attention import AttnType
+
 
 def test_kv_usage_pages_excludes_evictable_prefix_cache():
     # page_usage now lives on the CacheManagerLike interface (polymorphic vs DSV4); test the
@@ -53,10 +55,15 @@ def _fake_engine(swa=True, moe=True, mamba=True):
         config=SimpleNamespace(
             page_size=16,
             cache_type="swa_radix",
-            model_config=SimpleNamespace(dsv4_args=None, has_swa_attention=swa),
+            max_running_req=1,
+            model_config=SimpleNamespace(
+                dsv4_args=None, has_swa_attention=swa,
+                kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.SWA, is_swa=True, sliding_window=128)] if swa else [],
+                swa_attention_group=lambda: SimpleNamespace(sliding_window=128),
+            ),
         ),
         kv_cache=SimpleNamespace(
-            swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21)
+            window_pages=512, swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21)
         ),
         moe_offload_cache=SimpleNamespace(
             cache_size=24, num_layers=8, num_experts=16,
diff --git a/tests/server/test_effort_dialect.py b/tests/server/test_effort_dialect.py
index 6a68a4ece3..a25de8e728 100644
--- a/tests/server/test_effort_dialect.py
+++ b/tests/server/test_effort_dialect.py
@@ -159,6 +159,22 @@ def test_empty_effort_is_treated_as_absent():
     assert state.sent.chat_template_kwargs == {}
 
 
+def test_numeric_effort_rides_template_kwargs_not_the_wire_field():
+    # The wire field takes named tiers only (as vLLM and SGLang type it); a numeric budget goes
+    # through chat_template_kwargs, where the checkpoint's encoder validates it at render time.
+    from pydantic import ValidationError
+    import pytest
+
+    for effort in (37, True, 25.5):
+        with pytest.raises(ValidationError):
+            chat_request(reasoning_effort=effort)
+    ctk = {"enable_thinking": True, "reasoning_effort": 37}
+    state = FakeState(reasoning_parser="qwen3")
+    response = run(handle_chat_completion(chat_request(chat_template_kwargs=ctk), None, state, {}))
+    assert not isinstance(response, JSONResponse)
+    assert state.sent.chat_template_kwargs == ctk
+
+
 def test_foreign_thinking_shapes_stay_ignored():
     # extra="allow" swallowed any thinking shape before the field existed;
     # a bare string, a bool, or a typeless dict must keep working unchanged.
diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py
index 34c1c04d10..bc22296760 100644
--- a/tests/server/test_function_call_parser.py
+++ b/tests/server/test_function_call_parser.py
@@ -337,6 +337,34 @@ def test_dsv32_streaming_multi_param_args_prefix_stable():
     }
 
 
+V41_BLOCK = (
+    "<｜DSML｜ calls>\n"
+    '<｜DSML｜ invoke name="glob">\n'
+    '<｜DSML｜ parameter name="pattern" string="true">*.py</｜DSML｜ parameter>\n'
+    '<｜DSML｜ parameter name="path" string="true">/src</｜DSML｜ parameter>\n'
+    "</｜DSML｜ invoke>\n"
+    "</｜DSML｜ calls>"
+)
+
+
+def test_dsv41_dsml_tags_lead_with_a_space():
+    """DeepSeek-V4.1 leads every DSML tag name with a space (its encoder's ``<｜DSML｜ calls>``):
+    the V3.2 detector sees no tool call there, the V4.1 one parses it whole and streamed."""
+    text = "Let me look.\n\n" + V41_BLOCK
+    assert not FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv32").has_tool_call(text)
+    parser = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41")
+    assert parser.has_tool_call(text)
+    result = parser.parse_non_stream(text)
+    assert result.normal_text == "Let me look."
+    assert [(c.name, json.loads(c.parameters)) for c in result.calls] == [("glob", {"pattern": "*.py", "path": "/src"})]
+
+    streamed = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41")
+    texts, calls = _feed(streamed, [text[i : i + 5] for i in range(0, len(text), 5)])
+    assert "".join(texts).strip() + streamed.finish_stream() == "Let me look."
+    assert [c.name for c in calls if c.name] == ["glob"]
+    assert json.loads("".join(c.parameters for c in calls if c.name is None)) == {"pattern": "*.py", "path": "/src"}
+
+
 def test_streaming_support_flags():
     # Every registered detector is incremental-safe (buffered fallback remains as
     # the escape hatch for future formats, exercised via monkeypatch in
diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py
index 78b1663e8d..e1eb87982b 100644
--- a/tests/server/test_parser_auto_selection.py
+++ b/tests/server/test_parser_auto_selection.py
@@ -87,6 +87,13 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch():
     assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25"
 
 
+def test_deepseek_v41_is_not_shadowed_by_the_v4_branch():
+    """V4.1's DSML leads every tag name with a space; the V4 arm (``"v4"`` is a substring of
+    ``"v41"``) would hand it a tool parser that never sees a call."""
+    assert _inferred("DeepseekV41ForCausalLM") == ("deepseekv41", "deepseekv32")
+    assert _inferred("DeepseekV4ForCausalLM")[0] == "deepseekv32"
+
+
 def test_an_explicit_choice_beats_inference():
     config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"})
     with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):
diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py
index 43922021a2..eb6e3a2ed2 100644
--- a/tests/server/test_streaming_model_matrix.py
+++ b/tests/server/test_streaming_model_matrix.py
@@ -67,6 +67,11 @@
         '<｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py</｜DSML｜parameter>'
         "</｜DSML｜invoke></｜DSML｜function_calls>"
     ),
+    "deepseekv41": (
+        '<｜DSML｜ calls>\n<｜DSML｜ invoke name="read">\n'
+        '<｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py</｜DSML｜ parameter>\n'
+        "</｜DSML｜ invoke>\n</｜DSML｜ calls>"
+    ),
     "qwen25": '<tool_call>\n{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}\n</tool_call>',
     "qwen3_coder": (
         "<tool_call><function=read><parameter=filePath>/tmp/test_calc.py</parameter>"
@@ -104,6 +109,7 @@
 # Substrings that must never leak into user-visible content.
 MARKUP_MARKERS = {
     "deepseekv32": ["｜DSML｜"],
+    "deepseekv41": ["｜DSML｜"],
     "qwen25": ["<tool_call>", "</tool_call>"],
     "qwen3_coder": ["<tool_call>", "<function="],
     "glm47": ["<arg_key>", "<arg_value>", "</tool_call>"],
@@ -121,6 +127,7 @@
 # opens the block implicitly (model emits only the closing marker).
 REASONING_FAMILIES = {
     "dsv4": ("deepseekv32", "deepseekv32", "", "</think>"),
+    "dsv4.1": ("deepseekv41", "deepseekv32", "", "</think>"),
     "qwen3.5": ("qwen3_coder", "qwen3", "", "</think>"),
     "qwen": ("qwen25", "qwen3", "", "</think>"),
     "glm4.7": ("glm47", "glm", "", "</think>"),
diff --git a/tests/tokenizer/test_effort.py b/tests/tokenizer/test_effort.py
index e56df5fe68..e0cf7a81f6 100644
--- a/tests/tokenizer/test_effort.py
+++ b/tests/tokenizer/test_effort.py
@@ -1,6 +1,8 @@
 """Unit tests for the reasoning-effort dialect layer (tokenizer/effort.py)."""
 from __future__ import annotations
 
+import pytest
+
 from freetoken.tokenizer.effort import (
     EFFORT_SCALE,
     EffortProfile,
@@ -131,3 +133,37 @@ def render(kwargs, tools):
     profile = probe_effort_profile(render)
     assert profile.supported == frozenset({"xhigh", "medium", "low"})
     assert profile.consumes_effort
+
+
+def _numeric_render(kwargs, tools):
+    if not tools and not kwargs.get("enable_thinking"):
+        return "chat"
+    effort = kwargs.get("reasoning_effort", "high")
+    if isinstance(effort, str):
+        effort = {"low": 50, "high": 75, "max": 100}[effort]
+    if type(effort) is not int or not 1 <= effort <= 100:
+        raise ValueError("invalid budget")
+    return f"Reasoning Effort: {effort}"
+
+
+def test_probe_preserves_validated_numeric_budgets():
+    profile = probe_effort_profile(_numeric_render)
+    assert profile.integer_range == (1, 100)
+    assert profile.default == "high"
+    for budget in range(1, 101):
+        assert quantize_effort(budget, profile) == budget
+    for bad in (0, 101, -1, True, False, 1.0, 1.5):
+        with pytest.raises(ValueError, match="integer"):
+            quantize_effort(bad, profile)
+
+
+def test_numeric_capability_requires_validation_and_the_whole_range():
+    assert probe_effort_profile(_dsv4_render).integer_range is None
+    assert probe_effort_profile(lambda kw, tools: str(kw)).integer_range is None
+
+    def holes(kwargs, tools):
+        if kwargs.get("reasoning_effort") == 42:
+            raise ValueError("unsupported")
+        return _numeric_render(kwargs, tools)
+
+    assert probe_effort_profile(holes).integer_range is None
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 16ef6d7419..fb73fd317a 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -405,6 +405,36 @@ def encode(self, prompt, return_tensors=None, add_special_tokens=True):
         return torch.tensor([[4, 5, 6]], dtype=torch.long)
 
 
+def test_dsv41_encoder_gets_image_sources_before_and_after_fetch(tmp_path):
+    import copy
+    from freetoken.mm.media import collect_image_refs
+    encoding_dir = tmp_path / "encoding"
+    encoding_dir.mkdir()
+    (encoding_dir / "encoding_dsv4.py").write_text('''
+def encode_messages(messages, thinking_mode, reasoning_effort=None):
+    parts = messages[0]["content"]
+    assert parts[1] == {"type": "text", "text": "compare"}
+    assert parts[0].get("url") == "https://example.com/image.png" or parts[0].get("data") == b"first image bytes"
+    assert parts[2].get("data") in ("aW1hZ2U=", b"second image bytes")
+    return "two images"
+''')
+    messages = [{"role": "user", "content": [
+        {"type": "image", "freetoken_ref": {"kind": "url", "data": "https://example.com/image.png"}},
+        {"type": "text", "text": "compare"},
+        {"type": "image", "freetoken_ref": {"kind": "b64", "data": "aW1hZ2U="}},
+    ]}]
+    manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path))
+    msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams())
+    original = copy.deepcopy(messages)
+    assert manager.render_prompt(msg) == "two images"
+    assert messages == original
+    assert len(collect_image_refs(messages)) == 2
+    msg.images = [b"first image bytes", b"second image bytes"]
+    stripped = copy.deepcopy(messages)
+    assert manager.render_prompt(msg) == "two images"
+    assert messages == stripped
+
+
 def test_tokenize_manager_uses_dsv4_encoder_when_chat_template_is_missing(tmp_path):
     encoding_dir = tmp_path / "encoding"
     encoding_dir.mkdir()
@@ -761,3 +791,28 @@ def test_tokenize_survives_an_unhashable_effort():
     manager.tokenize([msg])
 
     assert "reasoning_effort" not in tokenizer.chat_template_kwargs
+
+
+@pytest.mark.parametrize("budget", [1, 25, 50, 75, 100])
+def test_numeric_encoder_effort_survives_prompt_rendering(tmp_path, budget):
+    encoding_dir = tmp_path / "encoding"
+    encoding_dir.mkdir()
+    (encoding_dir / "encoding.py").write_text('''
+def encode_messages(messages, thinking_mode, reasoning_effort=None):
+    if thinking_mode != "thinking":
+        return "chat"
+    effort = "high" if reasoning_effort is None else reasoning_effort
+    if isinstance(effort, str):
+        effort = {"low": 50, "high": 75, "max": 100}[effort]
+    if type(effort) is not int or not 1 <= effort <= 100:
+        raise ValueError("invalid budget")
+    return f"Reasoning Effort: {effort}"
+''')
+    manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path))
+    msg = TokenizeMsg(
+        uid=1, text=[{"role": "user", "content": "hello"}],
+        sampling_params=SamplingParams(),
+        chat_template_kwargs={"enable_thinking": True, "reasoning_effort": budget},
+    )
+    assert manager.render_prompt(msg) == f"Reasoning Effort: {budget}"
+    assert msg.chat_template_kwargs["reasoning_effort"] == budget

From 377a3bc4284f032e9d30777ce099828a158e7e4c Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 7 Oct 2026 13:24:45 -0700
Subject: [PATCH 071/249] ci: auto-label amd issues and prs by title (#620)

---
 .github/issue-title-labeler.yml    |  2 ++
 .github/workflows/issue-labels.yml | 18 +++++++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)
 create mode 100644 .github/issue-title-labeler.yml

diff --git a/.github/issue-title-labeler.yml b/.github/issue-title-labeler.yml
new file mode 100644
index 0000000000..48826301d5
--- /dev/null
+++ b/.github/issue-title-labeler.yml
@@ -0,0 +1,2 @@
+amd:
+  - '/\b(ROCm|HIP|RCCL|RDNA\s?\d?|CDNA\s?\d?|gfx\d{3,4}[a-z]?|MI\d{3}X?|RX\s?\d{4}|Radeon(?!\s*(\d{3}M|Graphics|Vega))|AMD[\s-]*(GPUs?|graphics|cards?|Instinct))\b|\bAMD[\s-]*显卡|\(amd\)|(?<!\w)A卡/i'
diff --git a/.github/workflows/issue-labels.yml b/.github/workflows/issue-labels.yml
index 35f325bf3b..0c11c3c404 100644
--- a/.github/workflows/issue-labels.yml
+++ b/.github/workflows/issue-labels.yml
@@ -1,23 +1,35 @@
-name: Label issues
+name: Label issues and PRs
 
 on:
   issues:
     types: [opened, edited]
+  pull_request_target:
+    types: [opened, edited]
 
 permissions:
   issues: write
+  pull-requests: write
 
 jobs:
   label:
     runs-on: ubuntu-latest
     steps:
-      - uses: github/issue-labeler@v3.4
+      - if: github.event_name == 'issues'
+        uses: github/issue-labeler@v3.4
         with:
           configuration-path: .github/issue-labeler.yml
           enable-versioned-regex: 0
           include-title: 0
           repo-token: ${{ github.token }}
-      - if: github.event.action == 'opened'
+      - if: github.event.action == 'opened' || github.event.changes.title
+        uses: github/issue-labeler@v3.4
+        with:
+          configuration-path: .github/issue-title-labeler.yml
+          enable-versioned-regex: 0
+          include-title: 1
+          include-body: 0
+          repo-token: ${{ github.token }}
+      - if: github.event_name == 'issues' && github.event.action == 'opened'
         uses: actions/github-script@v7
         with:
           script: |

From 76449847b6e7e4288f5e2247c9dcb5a040fe75fa Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 7 Oct 2026 22:42:44 -0700
Subject: [PATCH 072/249] perf(kernel): avoid layout conversions in nvfp4
 marlin decode (#631)

---
 .../kernel/triton/nvfp4_fused_moe.py          | 137 +++++++++++++-----
 python/freetoken/moe/fused_nvfp4.py           |  23 +--
 2 files changed, 115 insertions(+), 45 deletions(-)

diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
index a1064bcf7d..b1de443787 100644
--- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py
+++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
@@ -30,6 +30,7 @@
 import triton.language as tl
 
 from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32
+from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32
 
 _E2M1_VALUES = [
     0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
@@ -130,6 +131,31 @@ def _decode_nvfp4_moe_kernel(
     tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask)
 
 
+@triton.jit
+def _bf16x4_to_f32(v):
+    """Four bf16 packed in one int64 -> four fp32 (a bf16 is the high half of its fp32)."""
+    lo = v.to(tl.int32)
+    hi = (v >> 32).to(tl.int32)
+    return (
+        (lo << 16).to(tl.float32, bitcast=True),
+        (lo & -65536).to(tl.float32, bitcast=True),
+        (hi << 16).to(tl.float32, bitcast=True),
+        (hi & -65536).to(tl.float32, bitcast=True),
+    )
+
+
+@triton.jit
+def _f16x4_to_f32(v):
+    lo = v.to(tl.int32)
+    hi = (v >> 32).to(tl.int32)
+    return (
+        lo.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32),
+        (lo >> 16).to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32),
+        hi.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32),
+        (hi >> 16).to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32),
+    )
+
+
 @triton.jit
 def _decode_nvfp4_marlin_kernel(
     a_ptr,             # [M, K] activations (compute dtype)
@@ -139,7 +165,6 @@ def _decode_nvfp4_marlin_kernel(
     c_ptr,             # [M, TOP_K, N] output (compute dtype)
     topk_weights_ptr,  # [M, TOP_K] fp32
     topk_ids_ptr,      # [M, TOP_K] int32 -> cache slot
-    lut_ptr,           # [16] fp32
     total_routes,
     N,
     K,
@@ -151,9 +176,11 @@ def _decode_nvfp4_marlin_kernel(
     stride_tw_m, stride_tw_k,
     stride_tid_m, stride_tid_k,
     BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_KW: tl.constexpr,  # int32 words per K-iter (covers 8*KW k-values)
+    BLOCK_SIZE_KW: tl.constexpr,  # int32 words per tile (covers 8*KW k-values)
+    UNROLL: tl.constexpr,         # tiles per loop trip
     TOP_K: tl.constexpr,
     A_ROW_IS_ROUTE: tl.constexpr,
+    A_WIDE: tl.constexpr,         # activation rows are contiguous bf16/fp16, 16-byte aligned
     MUL_ROUTED_WEIGHT: tl.constexpr,
     compute_type: tl.constexpr,
 ):
@@ -165,13 +192,18 @@ def _decode_nvfp4_marlin_kernel(
       * Loads the FP4 codes as **int32 words** (8 codes / 4 bytes per element) so the HBM
         read issues wide coalesced transactions -- the single biggest lever for the
         gate/up GEMM, whose mem-only ceiling jumps from ~65% to ~85% of peak read BW.
-      * **Defers** the K reduction: the ``[BLOCK_KW, BLOCK_N]`` partial accumulates across
+      * **Defers** the K reduction: the ``[BLOCK_N, BLOCK_KW]`` partial accumulates across
         *all* K-iters and is reduced once at the end, removing the per-iter ``tl.sum``
         barrier so the weight loads of successive iters pipeline.
 
-    The e2m1 codes share one fp8 block scale per 16 k-values (== 2 words), applied per
-    word. Mirrors the fast kernel's route/tile mapping and epilogue so it is a drop-in
-    decode GEMV (CUDA-graph safe: fixed shapes, no host sync)."""
+    Every load in the K-loop is shaped so the compiler keeps the whole loop in the word
+    tile's register layout (no shared-memory layout conversions): the block scales are
+    loaded as a half tile and interleaved in registers (one e4m3 scale per 16 k-values
+    == 2 words), the e2m1 codes are decoded with the fp16 bit trick of
+    :func:`_nvfp4_pair_f32` (2^14 folded into the scale), and with ``A_WIDE`` the 8
+    activations of a word are read as two int64 and unpacked by shifts. Mirrors the fast
+    kernel's route/tile mapping and epilogue so it is a drop-in decode GEMV (CUDA-graph
+    safe: fixed shapes, no host sync)."""
     route_id = tl.program_id(0)
     n_block_id = tl.program_id(1)
     token_id = route_id // TOP_K
@@ -185,37 +217,70 @@ def _decode_nvfp4_marlin_kernel(
     a_base = a_ptr + a_row * stride_am
 
     offs_kw = tl.arange(0, BLOCK_SIZE_KW)
+    offs_kh = tl.arange(0, BLOCK_SIZE_KW // 2)
     K_WORDS = K // 8
-    partial = tl.zeros((BLOCK_SIZE_KW, BLOCK_SIZE_N), dtype=tl.float32)
-
-    packed_slot = packed_ptr + slot * stride_pe
-    scale_slot = scale_ptr + slot * stride_se
-    for kw_start in range(0, tl.cdiv(K_WORDS, BLOCK_SIZE_KW)):
-        widx = kw_start * BLOCK_SIZE_KW + offs_kw
-        w_mask = widx < K_WORDS
-
-        word = tl.load(
-            packed_slot + offs_n[None, :] * stride_pn + widx[:, None] * stride_pkw,
-            mask=w_mask[:, None] & n_mask[None, :], other=0,
-        )
-        # 8 codes/word fall in the same or adjacent 16-wide block -> one scale per word.
-        s_ptrs = scale_slot + offs_n[None, :] * stride_sn + (widx[:, None] // 2) * stride_sblk
-        s_mask = w_mask[:, None] & n_mask[None, :]
-        if e4m3_native_cx():
-            scale = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float32)
-        else:
-            scale = e4m3_u8_to_f32(tl.load(s_ptrs, mask=s_mask, other=0))
-
-        kbase = 8 * widx
-        acc_w = tl.zeros((BLOCK_SIZE_KW, BLOCK_SIZE_N), dtype=tl.float32)
-        for j in tl.static_range(8):
-            code = (word >> (4 * j)) & 0xF
-            b = tl.load(lut_ptr + code)
-            a_j = tl.load(a_base + (kbase + j) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
-            acc_w += a_j[:, None] * b
-        partial += acc_w * scale
-
-    accumulator = tl.sum(partial, axis=0)
+    partial = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_KW), dtype=tl.float32)
+
+    p_base = packed_ptr + slot * stride_pe + offs_n[:, None] * stride_pn
+    s_base = scale_ptr + slot * stride_se + offs_n[:, None] * stride_sn
+    for kw_start in range(0, tl.cdiv(K_WORDS, BLOCK_SIZE_KW), UNROLL):
+        for u in tl.static_range(UNROLL):
+            widx = (kw_start + u) * BLOCK_SIZE_KW + offs_kw
+            w_mask = widx < K_WORDS
+            word = tl.load(
+                p_base + widx[None, :] * stride_pkw,
+                mask=n_mask[:, None] & w_mask[None, :], other=0,
+            )
+            hidx = (kw_start + u) * (BLOCK_SIZE_KW // 2) + offs_kh
+            s_ptrs = s_base + hidx[None, :] * stride_sblk
+            s_mask = n_mask[:, None] & (hidx[None, :] < K_WORDS // 2)
+            if e4m3_native_cx():
+                s_half = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float32)
+            else:
+                s_half = e4m3_u8_to_f32(tl.load(s_ptrs, mask=s_mask, other=0))
+            s_half = s_half * 16384.0
+            scale = tl.interleave(s_half, s_half)
+
+            if A_WIDE:
+                a64 = a_base.to(tl.pointer_type(tl.int64), bitcast=True)
+                v = tl.load(
+                    a64 + widx[:, None] * 2 + tl.arange(0, 2)[None, :],
+                    mask=w_mask[:, None], other=0,
+                )
+                v0, v1 = tl.split(v)
+                if a_ptr.dtype.element_ty == tl.bfloat16:
+                    a0, a1, a2, a3 = _bf16x4_to_f32(v0)
+                    a4, a5, a6, a7 = _bf16x4_to_f32(v1)
+                else:
+                    a0, a1, a2, a3 = _f16x4_to_f32(v0)
+                    a4, a5, a6, a7 = _f16x4_to_f32(v1)
+            else:
+                kbase = 8 * widx
+                a0 = tl.load(a_base + kbase * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a1 = tl.load(a_base + (kbase + 1) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a2 = tl.load(a_base + (kbase + 2) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a3 = tl.load(a_base + (kbase + 3) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a4 = tl.load(a_base + (kbase + 4) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a5 = tl.load(a_base + (kbase + 5) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a6 = tl.load(a_base + (kbase + 6) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+                a7 = tl.load(a_base + (kbase + 7) * stride_ak, mask=w_mask, other=0.0).to(tl.float32)
+
+            acc_w = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_KW), dtype=tl.float32)
+            b0, b4 = _nvfp4_pair_f32(word)
+            b1, b5 = _nvfp4_pair_f32(word >> 4)
+            b2, b6 = _nvfp4_pair_f32(word >> 8)
+            b3, b7 = _nvfp4_pair_f32(word >> 12)
+            acc_w += a0[None, :] * b0
+            acc_w += a1[None, :] * b1
+            acc_w += a2[None, :] * b2
+            acc_w += a3[None, :] * b3
+            acc_w += a4[None, :] * b4
+            acc_w += a5[None, :] * b5
+            acc_w += a6[None, :] * b6
+            acc_w += a7[None, :] * b7
+            partial += acc_w * scale
+
+    accumulator = tl.sum(partial, axis=1)
     g = tl.load(global_ptr + slot * stride_ge + offs_n * stride_gn, mask=n_mask, other=0.0).to(tl.float32)
     accumulator = accumulator * g
 
diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py
index 7988fa9efc..375d2a547a 100644
--- a/python/freetoken/moe/fused_nvfp4.py
+++ b/python/freetoken/moe/fused_nvfp4.py
@@ -31,17 +31,15 @@
 _DECODE_BLOCK_KB = 128
 _DECODE_WARPS = 4
 
-# Marlin-style decode config (int32 wide loads + deferred reduction). Offline sweep over
-# the qwen35/qwen3moe (I=512/768) decode shapes picked BLOCK_N=16, BLOCK_KW=16 (== 128
-# k-values/iter), 4 warps -- the wide load lifts the gate/up GEMM ~43%->~51% of peak BW.
+# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps.
+# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one.
 _DECODE_MARLIN_BLOCK_N = 16
 _DECODE_MARLIN_BLOCK_KW = 16
+_DECODE_MARLIN_UNROLL = 2
 _DECODE_MARLIN_WARPS = 4
-# Deep-K variant: at K > 2048 (qwen4_exp gate_up, K=2560) a narrower N tile with the whole
-# K strip in one program iteration measures ~13% faster (18.6 vs 21.0us); short-K shapes
-# regress under it, so the split is by K, not by gemm position.
 _DECODE_MARLIN_DEEPK_BLOCK_N = 8
-_DECODE_MARLIN_DEEPK_BLOCK_KW = 128
+_DECODE_MARLIN_DEEPK_BLOCK_KW = 32
+_DECODE_MARLIN_DEEPK_UNROLL = 2
 _DECODE_MARLIN_DEEPK_THRESHOLD = 2048
 
 
@@ -111,13 +109,18 @@ def _decode_gemm_marlin(
     packed_i32 = packed.view(torch.int32)  # [S, N, K // 8]
     scale = e4m3_kernel_view(scale)
     total_routes = M * top_k
-    deep_k = K > _DECODE_MARLIN_DEEPK_THRESHOLD
+    deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD
     block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N
     block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW
+    unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL
+    # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned
+    a_wide = (
+        a.dtype in (torch.bfloat16, torch.float16)
+        and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0
+    )
     grid = (total_routes, triton.cdiv(N, block_n))
     _decode_nvfp4_marlin_kernel[grid](
         a, packed_i32, scale, glob, c, topk_weights, topk_ids,
-        _e2m1_lut(a.device.index),
         total_routes, N, K,
         a.stride(0), a.stride(1),
         packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2),
@@ -128,8 +131,10 @@ def _decode_gemm_marlin(
         topk_ids.stride(0), topk_ids.stride(1),
         BLOCK_SIZE_N=block_n,
         BLOCK_SIZE_KW=block_kw,
+        UNROLL=unroll,
         TOP_K=top_k,
         A_ROW_IS_ROUTE=a_row_is_route,
+        A_WIDE=a_wide,
         MUL_ROUTED_WEIGHT=mul_routed_weight,
         compute_type=_tl_dtype(c.dtype),
         num_warps=_DECODE_MARLIN_WARPS,

From 127b530a8ff5ae5b6f15c3b9f558c44fde774bd6 Mon Sep 17 00:00:00 2001
From: OpenClaw <openclaw@local>
Date: Thu, 8 Oct 2026 05:49:29 +0000
Subject: [PATCH 073/249] perf(moe): avoid intra-op fanout while filling NVFP4
 host banks

---
 python/freetoken/moe/expert_banks.py    | 37 +++++++++++++---
 tests/moe/test_expert_bank_threading.py | 57 +++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 7 deletions(-)
 create mode 100644 tests/moe/test_expert_bank_threading.py

diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 7500174ae4..70461f97f4 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -11,6 +11,7 @@
 import glob
 import math
 import os
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, field
 
 import torch
@@ -27,6 +28,19 @@
 _PARALLEL_READER_SUPPORTED = hasattr(os, "O_DIRECT") and hasattr(os, "preadv")
 
 
+@contextmanager
+def _single_threaded_host_copies():
+    """Keep small NVFP4 bank placements off PyTorch's intra-op worker pool."""
+    previous = torch.get_num_threads()
+    if previous > 1:
+        torch.set_num_threads(1)
+    try:
+        yield
+    finally:
+        if previous > 1:
+            torch.set_num_threads(previous)
+
+
 @dataclass(frozen=True)
 class ExpertBanks:
     """Loaded expert banks, normalized for ``OffloadMoeCache`` wiring."""
@@ -141,13 +155,22 @@ def _fill(sink) -> None:
         if missing:
             raise ValueError(f"expert banks were not filled: {len(missing)} (layer, expert) rows missing (first {missing[:4]})")
 
-    if layer_sink is not None:
-        _fill(layer_sink)
-    elif torch.cuda.is_available() and not resident:
-        with PinPipeline() as pins:
-            _fill(pins)
-    else:
-        _fill(None)
+    # Per-expert NVFP4 pieces produce thousands of small CPU copies. On high-core
+    # hosts the intra-op pool costs more than the copies and contends with the
+    # checkpoint reader. Restore its setting even if packing or pinning fails.
+    copy_threads = (
+        _single_threaded_host_copies()
+        if method.kind is QuantKind.NVFP4 and not resident
+        else nullcontext()
+    )
+    with copy_threads:
+        if layer_sink is not None:
+            _fill(layer_sink)
+        elif torch.cuda.is_available() and not resident:
+            with PinPipeline() as pins:
+                _fill(pins)
+        else:
+            _fill(None)
 
     return ExpertBanks(
         legacy_format_for(method.kind, kernel.name), banks,
diff --git a/tests/moe/test_expert_bank_threading.py b/tests/moe/test_expert_bank_threading.py
new file mode 100644
index 0000000000..da74719448
--- /dev/null
+++ b/tests/moe/test_expert_bank_threading.py
@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from freetoken.layers.quantization import QuantKind
+from freetoken.moe.expert_banks import build_expert_banks
+
+
+@pytest.mark.parametrize("fail", [False, True])
+def test_nvfp4_host_bank_copy_threads_are_restored(monkeypatch, fail):
+    threads = 24
+    changes = []
+    observed = []
+
+    def set_threads(value):
+        nonlocal threads
+        threads = value
+        changes.append(value)
+
+    monkeypatch.setattr(torch, "get_num_threads", lambda: threads)
+    monkeypatch.setattr(torch, "set_num_threads", set_threads)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+
+    class Method:
+        kind = QuantKind.NVFP4
+        kernel = SimpleNamespace(name="triton")
+        cfg = SimpleNamespace(num_experts=2)
+
+        @staticmethod
+        def layout():
+            spec = SimpleNamespace(shape=(2,), dtype=torch.uint8, resident=False)
+            return {"gate_up": spec}
+
+        @staticmethod
+        def pack(piece, out):
+            observed.append(threads)
+            if fail and piece["gate_up"][0, 0] == 2:
+                raise RuntimeError("failed during expert placement")
+            out["gate_up"].copy_(piece["gate_up"])
+            return {}
+
+    pieces = (
+        (0, expert, expert + 1, {
+            "gate_up": torch.full((1, 2), expert + 1, dtype=torch.uint8),
+        })
+        for expert in range(2)
+    )
+    if fail:
+        with pytest.raises(RuntimeError, match="failed during expert placement"):
+            build_expert_banks(Method(), 1, pieces, device=torch.device("cpu"))
+    else:
+        banks = build_expert_banks(Method(), 1, pieces, device=torch.device("cpu"))
+        assert banks.sources["gate_up"][0].tolist() == [[1, 1], [2, 2]]
+
+    assert observed == [1, 1]
+    assert changes == [1, 24]
+    assert threads == 24

From 799bec39f687f612470de604ebb8c689e0046e1c Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 7 Oct 2026 23:04:18 -0700
Subject: [PATCH 074/249] perf(kernel): read dsv41 indexer key rows as int32
 words (#634)

---
 .../freetoken/kernel/triton/dsv41/indexer.py  | 26 ++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/kernel/triton/dsv41/indexer.py b/python/freetoken/kernel/triton/dsv41/indexer.py
index 3721315837..bd3fe7ae31 100644
--- a/python/freetoken/kernel/triton/dsv41/indexer.py
+++ b/python/freetoken/kernel/triton/dsv41/indexer.py
@@ -40,6 +40,24 @@
 BLOCK_T = 64
 
 
+@triton.jit
+def _load_fp4_e8m0_rows(pool_ptr, rows, valid, D: tl.constexpr, ROW_BYTES: tl.constexpr):
+    """``load_rows(..., FMT=3)`` as bf16 for ``D == 128``, with each row read as int32 words: 16 words of 8 e2m1 codes, then one word of 4 scales.
+    The contiguous word axis pins the gather layout; a per-channel byte gather has no contiguous axis, so its layout (and speed) follows the Triton version."""
+    W: tl.constexpr = D // 8
+    base = pool_ptr.to(tl.pointer_type(tl.int32)) + tl.maximum(rows, 0).to(tl.int64) * (ROW_BYTES // 4)
+    offs_w = tl.arange(0, W)
+    x = tl.load(base[:, None] + offs_w[None, :], mask=valid[:, None], other=0)
+    sw = tl.load(base + W, mask=valid, other=0)
+    s = tl.exp2(((sw[:, None] >> (8 * (offs_w // 4))[None, :]) & 0xFF).to(tl.float32) - 127.0)
+    x = tl.join(x, x >> 16)
+    x = tl.join(x, x >> 8)
+    x = tl.join(x, x >> 4)
+    # e2m1 code -> fp32 bits: magnitude to bits 22-24, sign to bit 31 (mask 0x81C00000). That is the value times 2**-126 (code 1 is subnormal), so the multiply is exact.
+    q = (((x << 28) >> 6) & -0x7E400000).to(tl.float32, bitcast=True) * 8.507059173023462e37
+    return tl.reshape((q * s[:, :, None, None, None]).to(tl.bfloat16), (rows.shape[0], D))
+
+
 @triton.jit
 def _score_tile(
     q, w, pool_ptr, locs_row, live, pos, store_mask, out_ptrs,
@@ -52,8 +70,10 @@ def _score_tile(
     full = tl.load(locs_row + tl.maximum(full_idx, 0), mask=pos_ok & (full_idx < n_locs), other=-1)
     rows = tl.where(full >= 0, full // RATIO, -1).to(tl.int32)
     valid = pos_ok & (rows >= 0)
-    offs_d = tl.arange(0, D)
-    k = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT).to(tl.bfloat16)  # [BLOCK_T, D]
+    if FMT == 3 and D == 128:
+        k = _load_fp4_e8m0_rows(pool_ptr, rows, valid, D, ROW_BYTES)
+    else:
+        k = load_rows(pool_ptr, rows, valid, tl.arange(0, D), D, ROW_BYTES, FMT).to(tl.bfloat16)  # [BLOCK_T, D]
     score = tl.dot(q, tl.trans(k)).to(tl.bfloat16).to(tl.float32)
     score = (tl.maximum(score, 0.0) * w[:, None]).to(tl.bfloat16).to(tl.float32)
     logits = tl.sum(score, axis=0).to(tl.bfloat16).to(tl.float32)
@@ -130,7 +150,7 @@ def indexer_logits_packed(
     assert locs.ndim == 2 and locs.shape[0] == B and live.shape == (B, S), (locs.shape, live.shape, (B, S))
     assert weights.shape == (B, S, H), (weights.shape, (B, S, H))
     assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}"
-    assert k_pool.dtype == torch.uint8 and k_pool.shape[1] == fmt.row_bytes(D) and k_pool.is_contiguous()
+    assert k_pool.dtype == torch.uint8 and k_pool.shape[1] == fmt.row_bytes(D) and k_pool.is_contiguous() and k_pool.data_ptr() % 4 == 0
     assert locs.dtype in (torch.int32, torch.int64) and locs.stride(1) == 1, (locs.dtype, locs.stride())
     n_locs = locs.shape[1]
     if T is None:

From 1c3aa0aa4cd89371604be23e1405e3228a151804 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 7 Oct 2026 23:06:22 -0700
Subject: [PATCH 075/249] perf(kernel): keep 16-byte kv loads in dsv4 sparse
 attention and unroll gated pool (#633)

---
 python/freetoken/kernel/triton/dsv4/compress.py    | 5 +++--
 python/freetoken/kernel/triton/dsv4/sparse_attn.py | 6 ++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/kernel/triton/dsv4/compress.py b/python/freetoken/kernel/triton/dsv4/compress.py
index 77bacd42ce..3ab89ff0d3 100644
--- a/python/freetoken/kernel/triton/dsv4/compress.py
+++ b/python/freetoken/kernel/triton/dsv4/compress.py
@@ -31,12 +31,13 @@ def _gated_pool_kernel(
     kv_base = kv_ptr + b * stride_kb + offs * stride_kd
     sc_base = score_ptr + b * stride_sb + offs * stride_sd
     maxv = tl.full((BLOCK_D,), float("-inf"), tl.float32)
-    for r in range(R):
+    # Explicit unroll: ptxas does not unroll the plain loop Triton 3.8 emits, which serializes these latency-bound loads.
+    for r in tl.range(R, loop_unroll_factor=2):
         s = tl.load(sc_base + r * stride_sr, mask=mask, other=float("-inf")).to(tl.float32)
         maxv = tl.maximum(maxv, s)
     denom = tl.zeros((BLOCK_D,), tl.float32)
     acc = tl.zeros((BLOCK_D,), tl.float32)
-    for r in range(R):
+    for r in tl.range(R, loop_unroll_factor=2):
         s = tl.load(sc_base + r * stride_sr, mask=mask, other=float("-inf")).to(tl.float32)
         w = tl.exp(s - maxv)
         denom += w
diff --git a/python/freetoken/kernel/triton/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py
index c22f891e89..639c3e5ea5 100644
--- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py
+++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py
@@ -99,7 +99,8 @@ def _sparse_attn_paged_kernel(
         # multiple KV tiles in shared memory. Result is bit-identical to the two-pool form -- each
         # column still reads from the same pool/slot.
         is_win = offs_t < N_WINDOW
-        base = tl.where(is_win, win_ptr, cmp_ptr)  # [BLOCK_T] per-column pool base pointer
+        # The select drops the pools' 16-byte alignment (asserted in the wrapper); without this hint Triton 3.8 gathers in 2-byte loads.
+        base = tl.multiple_of(tl.where(is_win, win_ptr, cmp_ptr), 16)  # [BLOCK_T] per-column pool base pointer
         kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd
         kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32)  # [BLOCK_T, D]
 
@@ -182,7 +183,7 @@ def _sparse_attn_paged_splitk_kernel(
             idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1)
             valid = idxs >= 0
             is_win = offs_t < N_WINDOW
-            base = tl.where(is_win, win_ptr, cmp_ptr)
+            base = tl.multiple_of(tl.where(is_win, win_ptr, cmp_ptr), 16)
             kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd
             kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32)
 
@@ -314,6 +315,7 @@ def sparse_attn_paged(
     # The kernel selects one pool base per column and gathers with a single stride pair, so both
     # pools must share strides. Guaranteed here by .contiguous() on their [*, D] shape.
     assert window_pool.stride() == cmp_pool.stride(), (window_pool.stride(), cmp_pool.stride())
+    assert window_pool.data_ptr() % 16 == 0 and cmp_pool.data_ptr() % 16 == 0, (window_pool.data_ptr(), cmp_pool.data_ptr())
 
     has_counts = cmp_counts is not None
     if has_counts:

From 078132480ea4a0424f677a27eb22085624d331e6 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 8 Oct 2026 00:36:55 -0700
Subject: [PATCH 076/249] perf(kernel): skip masks and out-of-window tiles in
 triton extend attention (#632)

* perf(kernel): skip masks and out-of-window tiles in triton extend attention

* perf(kernel): drop the sm90-only register cap and k/v pipelining in triton extend attention
---
 python/freetoken/kernel/triton/attention.py | 195 +++++++++++---------
 1 file changed, 108 insertions(+), 87 deletions(-)

diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index 3f53c5fc0e..c9087e03c2 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -18,6 +18,11 @@ def _optin_smem_bytes(device_index: int) -> int:
     return int(getattr(props, "shared_memory_per_block_optin", 0))
 
 
+@functools.lru_cache(maxsize=None)
+def _num_sms(device_index: int) -> int:
+    return int(torch.cuda.get_device_properties(device_index).multi_processor_count)
+
+
 def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]:
     """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware.
 
@@ -398,10 +403,15 @@ def decode_paged_attention(
     block_h = triton.next_power_of_2(valid_block_h)
     block_d = triton.next_power_of_2(head_dim)
     block_dv = triton.next_power_of_2(head_dim)
-
-    _decode_grouped_stage1_kernel[
-        (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits)
-    ](
+    grid = (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits)
+    block_n = 32
+    if sliding_window and head_dim <= 256:
+        split_len = triton.cdiv(triton.cdiv(sliding_window, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV
+        # A window split of several tiles in a single-wave grid is latency bound; 64-key tiles halve its serial steps.
+        if split_len >= 128 and grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index):
+            block_n = 64
+
+    _decode_grouped_stage1_kernel[grid](
         q,
         k_cache,
         v_cache,
@@ -428,7 +438,7 @@ def decode_paged_attention(
         NUM_Q_HEADS=num_q_heads,
         BLOCK_D=block_d,
         BLOCK_DV=block_dv,
-        BLOCK_N=32,
+        BLOCK_N=block_n,
         BLOCK_H=block_h,
         VALID_BLOCK_H=valid_block_h,
         MIN_BLOCK_KV=_MIN_BLOCK_KV,
@@ -595,6 +605,48 @@ def _extend_attention_kernel(
     )
 
 
+@triton.jit
+def _extend_attend_tile(
+    q, m_i, l_i, acc, qk_scale, k_base, v_base, k_rows, v_rows, offs_d, offs_dv, mask_d, mask_dv, mask_n, mask,
+    MASKED: tl.constexpr,
+):
+    if MASKED:
+        k_mask = mask_n[None, :] & mask_d[:, None]
+        v_mask = mask_n[:, None] & mask_dv[None, :]
+    else:
+        k_mask = mask_d[:, None]
+        v_mask = mask_dv[None, :]
+    k = tl.load(k_base + k_rows[None, :] + offs_d[:, None], mask=k_mask, other=0.0)
+    scores = tl.dot(q.to(k.dtype), k)
+    if MASKED:
+        scores = tl.where(mask, scores, -float("inf"))
+    m_new = tl.maximum(m_i, tl.max(scores, axis=1) * qk_scale)
+    if MASKED:
+        m_new = tl.where(m_new == -float("inf"), -1e20, m_new)
+    alpha = tl.math.exp2(m_i - m_new)
+    p = tl.math.exp2(scores * qk_scale - m_new[:, None])
+    v = tl.load(v_base + v_rows[:, None] + offs_dv[None, :], mask=v_mask, other=0.0)
+    acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
+    l_i = l_i * alpha + tl.sum(p, axis=1)
+    return m_new, l_i, acc
+
+
+@triton.jit
+def _extend_tile_ranges(
+    first_row, diag, end, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr
+):
+    """Split keys [0, end) seen by rows first_row + [0, BLOCK_M) into skipped [0, lo), masked [lo, u0), unmasked [u0, u1) and masked [u1, end); keys below diag are causal for every row."""
+    if SLIDING_WINDOW > 0:
+        lo = tl.minimum(tl.maximum(first_row - SLIDING_WINDOW + 1, 0) // BLOCK_N * BLOCK_N, end)
+        full = tl.cdiv(tl.maximum(first_row + BLOCK_M - SLIDING_WINDOW, 0), BLOCK_N) * BLOCK_N
+    else:
+        lo = 0
+        full = 0
+    u1 = tl.maximum(diag, lo)
+    u0 = tl.minimum(tl.maximum(full, lo), u1)
+    return lo, u0, u1
+
+
 @triton.jit
 def _extend_attention_split_kernel(
     q_ptr,
@@ -642,7 +694,8 @@ def _extend_attention_split_kernel(
     kv_start = tl.load(kv_indptr_ptr + seq_id)
     prefix_len = tl.load(prefix_lens_ptr + seq_id)
 
-    offs_m = block_m_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    block_start = block_m_id * BLOCK_M
+    offs_m = block_start + tl.arange(0, BLOCK_M)
     offs_n = tl.arange(0, BLOCK_N)
     offs_d = tl.arange(0, BLOCK_D)
     offs_dv = tl.arange(0, BLOCK_DV)
@@ -650,6 +703,8 @@ def _extend_attention_split_kernel(
     mask_d = offs_d < D
     mask_dv = offs_dv < D
     q_abs_pos = prefix_len + offs_m
+    # Scores live in the log2 domain: exp2 lowers to ex2.approx.ftz, while tl.exp adds a denormal fixup per element.
+    qk_scale = sm_scale * 1.4426950408889634
 
     q = tl.load(
         q_ptr
@@ -661,7 +716,7 @@ def _extend_attention_split_kernel(
     )
 
     if HAS_SINKS:
-        sink = tl.load(sinks_ptr + q_head).to(tl.float32)
+        sink = tl.load(sinks_ptr + q_head).to(tl.float32) * 1.4426950408889634
         m_i = tl.full((BLOCK_M,), sink, dtype=tl.float32)
         l_i = tl.full((BLOCK_M,), 1.0, dtype=tl.float32)
     else:
@@ -669,59 +724,54 @@ def _extend_attention_split_kernel(
         l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
     acc = tl.zeros((BLOCK_M, BLOCK_DV), dtype=tl.float32)
 
-    for start_n in tl.range(0, prefix_len, BLOCK_N):
-        kv_offsets = start_n + offs_n
-        mask_n = kv_offsets < prefix_len
-        key_pos = kv_offsets
-        final_mask = mask_m[:, None] & mask_n[None, :]
-        if SLIDING_WINDOW > 0:
-            window_mask = (key_pos[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None]
-            final_mask = final_mask & window_mask
-
-        skip_tile = False
+    kc_base = k_cache_ptr + kv_head * stride_kch
+    vc_base = v_cache_ptr + kv_head * stride_vch
+    prefix_end = tl.where(block_start < q_len, prefix_len, 0)
+    lo, u0, u1 = _extend_tile_ranges(
+        prefix_len + block_start, prefix_end // BLOCK_N * BLOCK_N, prefix_end, BLOCK_M, BLOCK_N, SLIDING_WINDOW
+    )
+    for start_n in tl.range(u0, u1, BLOCK_N):
+        slots = tl.load(kv_indices_ptr + kv_start + start_n + offs_n)
+        m_i, l_i, acc = _extend_attend_tile(
+            q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs,
+            offs_d, offs_dv, mask_d, mask_dv, None, None, False,
+        )
+    n_left = (u0 - lo) // BLOCK_N
+    for i in tl.range(0, n_left + tl.cdiv(prefix_end - u1, BLOCK_N)):
+        kv_offsets = tl.where(i < n_left, lo + i * BLOCK_N, u1 + (i - n_left) * BLOCK_N) + offs_n
+        mask_n = kv_offsets < prefix_end
+        final_mask = mask_n[None, :]
         if SLIDING_WINDOW > 0:
-            skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0
-
-        if not skip_tile:
-            slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0)
-            k = tl.load(
-                k_cache_ptr
-                + slots[None, :] * stride_kcs
-                + kv_head * stride_kch
-                + offs_d[:, None],
-                mask=mask_n[None, :] & mask_d[:, None],
-                other=0.0,
-            )
-            scores = tl.dot(q.to(k.dtype), k) * sm_scale
-            scores = tl.where(final_mask, scores, -float("inf"))
-
-            row_max = tl.max(scores, axis=1)
-            row_max_fixed = tl.where(row_max == -float("inf"), -1e20, row_max)
-            m_new = tl.maximum(row_max_fixed, m_i)
-            alpha = tl.exp(m_i - m_new)
-            p = tl.exp(scores - m_new[:, None])
-
-            v = tl.load(
-                v_cache_ptr
-                + slots[:, None] * stride_vcs
-                + kv_head * stride_vch
-                + offs_dv[None, :],
-                mask=mask_n[:, None] & mask_dv[None, :],
-                other=0.0,
-            )
-            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
-            l_i = l_i * alpha + tl.sum(p, axis=1)
-            m_i = m_new
+            final_mask = final_mask & ((kv_offsets[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None])
+        slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0)
+        m_i, l_i, acc = _extend_attend_tile(
+            q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs,
+            offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True,
+        )
 
-    current_end = tl.minimum(q_len, (block_m_id + 1) * BLOCK_M)
+    ke_base = k_extend_ptr + kv_head * stride_keh
+    ve_base = v_extend_ptr + kv_head * stride_veh
+    current_end = tl.minimum(q_len, block_start + BLOCK_M)
     if HAS_BLOCKS:
         # rows inside a multimodal span also attend forward to the span's later keys: the tile loop must reach them
         block_end = tl.load(block_ends_ptr + q_start + offs_m, mask=mask_m, other=0) - prefix_len
         current_end = tl.minimum(q_len, tl.maximum(current_end, tl.max(block_end, axis=0)))
     else:
         block_end = tl.zeros((BLOCK_M,), dtype=tl.int32)
-    for start_n in tl.range(0, current_end, BLOCK_N):
-        local_kv_offsets = start_n + offs_n
+    current_end = tl.where(block_start < q_len, current_end, 0)
+    lo, u0, u1 = _extend_tile_ranges(
+        block_start, tl.minimum(block_start // BLOCK_N * BLOCK_N, current_end), current_end, BLOCK_M, BLOCK_N,
+        SLIDING_WINDOW,
+    )
+    for start_n in tl.range(u0, u1, BLOCK_N):
+        rows = q_start + start_n + offs_n
+        m_i, l_i, acc = _extend_attend_tile(
+            q, m_i, l_i, acc, qk_scale, ke_base, ve_base, rows * stride_ket, rows * stride_vet,
+            offs_d, offs_dv, mask_d, mask_dv, None, None, False,
+        )
+    n_left = (u0 - lo) // BLOCK_N
+    for i in tl.range(0, n_left + tl.cdiv(current_end - u1, BLOCK_N)):
+        local_kv_offsets = tl.where(i < n_left, lo + i * BLOCK_N, u1 + (i - n_left) * BLOCK_N) + offs_n
         mask_n = local_kv_offsets < current_end
         local_q_pos = offs_m
         causal_mask = local_kv_offsets[None, :] <= local_q_pos[:, None]
@@ -731,41 +781,12 @@ def _extend_attention_split_kernel(
             causal_mask = causal_mask & (
                 (local_kv_offsets[None, :] + SLIDING_WINDOW) > local_q_pos[:, None]
             )
-        final_mask = mask_m[:, None] & mask_n[None, :] & causal_mask
-
-        skip_tile = False
-        if SLIDING_WINDOW > 0:
-            skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0
-
-        if not skip_tile:
-            k = tl.load(
-                k_extend_ptr
-                + (q_start + local_kv_offsets[None, :]) * stride_ket
-                + kv_head * stride_keh
-                + offs_d[:, None],
-                mask=mask_n[None, :] & mask_d[:, None],
-                other=0.0,
-            )
-            scores = tl.dot(q.to(k.dtype), k) * sm_scale
-            scores = tl.where(final_mask, scores, -float("inf"))
-
-            row_max = tl.max(scores, axis=1)
-            row_max_fixed = tl.where(row_max == -float("inf"), -1e20, row_max)
-            m_new = tl.maximum(row_max_fixed, m_i)
-            alpha = tl.exp(m_i - m_new)
-            p = tl.exp(scores - m_new[:, None])
-
-            v = tl.load(
-                v_extend_ptr
-                + (q_start + local_kv_offsets[:, None]) * stride_vet
-                + kv_head * stride_veh
-                + offs_dv[None, :],
-                mask=mask_n[:, None] & mask_dv[None, :],
-                other=0.0,
-            )
-            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
-            l_i = l_i * alpha + tl.sum(p, axis=1)
-            m_i = m_new
+        final_mask = mask_n[None, :] & causal_mask
+        rows = q_start + local_kv_offsets
+        m_i, l_i, acc = _extend_attend_tile(
+            q, m_i, l_i, acc, qk_scale, ke_base, ve_base, rows * stride_ket, rows * stride_vet,
+            offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True,
+        )
 
     out = tl.where(l_i[:, None] == 0.0, 0.0, acc / l_i[:, None])
     tl.store(

From 4f4832839719cb814f928821cc40de3336fd4592 Mon Sep 17 00:00:00 2001
From: Mukkesh Ganesh <mukmckenzie@gmail.com>
Date: Thu, 8 Oct 2026 12:32:04 -0700
Subject: [PATCH 077/249] fix(engine): apply shrinking cache resizes before
 growing ones in live rebuild

---
 python/freetoken/engine/engine.py             |  59 ++++--
 python/freetoken/kvcache/base.py              |  18 ++
 python/freetoken/kvcache/dsv4/v4_pool.py      |  10 +
 .../test_rebuild_runtime_cache_ordering.py    | 199 ++++++++++++++++++
 tests/kvcache/test_kv_cache_rebuild.py        |  10 +-
 5 files changed, 277 insertions(+), 19 deletions(-)
 create mode 100644 tests/engine/test_rebuild_runtime_cache_ordering.py

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index bf3960ab8d..0f3cbdb0fc 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -996,29 +996,54 @@ def rebuild_runtime_cache(
         # 1. Tear down CUDA graphs + backend capture scratch (free-before-alloc).
         self.attn_backend.reset_capture()
         self.graph_runner.destroy_cuda_graphs()
-        # 2. Resize caches in place (each frees its old GPU tensors before allocating).
+        # 2. Resize caches in place (each frees its old GPU tensors before allocating),
+        # shrinking resizes before growing ones: a pool's grow must not run while a
+        # still-large sibling is resident, or a target that passed the fit check above can
+        # OOM mid-rebuild (#643). With shrinks first the transient peak never exceeds
+        # max(current, target) totals. Footprints are compared BEFORE the window-override
+        # write below (the old side must price the pre-override config).
+        kv_old_bytes = type(self.kv_cache).rebuild_footprint(config, self.num_pages, None)
+        kv_new_bytes = type(self.kv_cache).rebuild_footprint(
+            config, num_pages if num_pages is not None else self.num_pages, num_swa_pages
+        )
+
+        def _resize_moe() -> None:
+            assert self.moe_offload_cache is not None, "no MoE offload cache to resize"
+            self.moe_offload_cache.rebuild(moe_cache_size)
+
+        def _resize_kv() -> None:
+            # sets self.num_pages (rebuilds KV + window); a window-only change passes the
+            # CURRENT page count so only the window pool is re-derived
+            self._resize_kv_pool(
+                config, num_pages if num_pages is not None else self.num_pages, num_swa_pages
+            )
+
+        def _resize_mamba() -> None:
+            # Like every resize here it must sit between graph teardown and re-capture so
+            # the recaptured graphs bind the new state tensors. +1 for the reserved
+            # padding sink: num_mamba_slots is the usable count.
+            self.linear_state_pool.rebuild(num_mamba_slots + 1)
+
+        resize_ops = [
+            (moe_cache_size is not None,
+             moe_cache_size is None or moe_cache_size <= self.moe_offload_cache.cache_size,
+             _resize_moe),
+            (num_pages is not None or num_swa_pages is not None,
+             kv_new_bytes <= kv_old_bytes, _resize_kv),
+            (num_mamba_slots is not None,
+             num_mamba_slots is None or num_mamba_slots + 1 <= self.linear_state_pool.num_slots,
+             _resize_mamba),
+        ]
         # Pin the new window first (validated above) so any KV-pool rebuild below sizes the window
         # to it (_dsv4_pool_sizes / _swa_paged_num_tokens read config.swa_num_pages_override).
         # frozen EngineConfig — mutate in place like the moe_cache_size path; `config.x = y` raises
         # FrozenInstanceError, which here aborts the rebuild after the CUDA graphs are gone (→ 503).
         if num_swa_pages is not None:
             object.__setattr__(config, "swa_num_pages_override", num_swa_pages)
-        if moe_cache_size is not None:
-            assert self.moe_offload_cache is not None, "no MoE offload cache to resize"
-            self.moe_offload_cache.rebuild(moe_cache_size)
-        if num_pages is not None:
-            # sets self.num_pages (rebuilds KV + window)
-            self._resize_kv_pool(config, num_pages, num_swa_pages)
-        elif num_swa_pages is not None:
-            # Window-only change: no page-count change, but re-derive the window pool at the new
-            # pin against the CURRENT page count. This re-allocs the same-size full pool and
-            # the resized window, both inside the pool's own rebuild_from_config.
-            self._resize_kv_pool(config, self.num_pages, num_swa_pages)
-        if num_mamba_slots is not None:
-            # Reallocate the GDN state pool (frees old tensors first). Must sit between graph
-            # teardown and re-capture so the recaptured graphs bind the new state tensors.
-            # +1 for the reserved padding sink: num_mamba_slots is the usable count.
-            self.linear_state_pool.rebuild(num_mamba_slots + 1)
+        for shrinking_phase in (True, False):
+            for requested, shrinks, op in resize_ops:
+                if requested and shrinks == shrinking_phase:
+                    op()
         # 3. Refresh max_seq_len (+ generic page table) for the new token budget.
         self._refresh_seq_state(config)
         aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size)
diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py
index 569da1cb6c..d7bc224ce0 100644
--- a/python/freetoken/kvcache/base.py
+++ b/python/freetoken/kvcache/base.py
@@ -136,6 +136,24 @@ def validate_rebuild(
                 f"needs {mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving"
             )
 
+    @classmethod
+    def rebuild_footprint(
+        cls, config, num_pages: int, num_swa_pages: int | None
+    ) -> int:
+        """Whole-family GPU bytes at ``(num_pages, num_swa_pages)`` -- the pool's own
+        tiers priced together, same kv_cost plumbing validate_rebuild uses. The engine
+        compares old-vs-new footprints to order a live rebuild's resizes
+        shrinking-first (#643); ``num_swa_pages=None`` prices the config's current
+        window pin."""
+        import inspect
+
+        cost_params = inspect.signature(cls.kv_cost).parameters
+        cost_kwargs = (
+            {"num_swa_pages": num_swa_pages} if "num_swa_pages" in cost_params else {}
+        )
+        cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config, **cost_kwargs)
+        return num_pages * cache_per_page + fixed_cache_size
+
     @abstractmethod
     def rebuild_from_config(
         self, config, num_pages: int, *, num_swa_pages: int | None = None
diff --git a/python/freetoken/kvcache/dsv4/v4_pool.py b/python/freetoken/kvcache/dsv4/v4_pool.py
index d5eae6d080..484dcbbb29 100644
--- a/python/freetoken/kvcache/dsv4/v4_pool.py
+++ b/python/freetoken/kvcache/dsv4/v4_pool.py
@@ -279,6 +279,16 @@ def validate_rebuild(
                 f"{mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving"
             )
 
+    @classmethod
+    def rebuild_footprint(cls, config, num_pages: int, num_swa_pages: int | None) -> int:
+        # The DSV4 kv_cost does not take num_swa_pages, so price the family with the same
+        # byte model validate_rebuild uses (num_pages is USABLE; +1 for the dummy page).
+        from .v4_cost_model import _dsv4_pool_sizes, dsv4_pool_bytes
+
+        dsv4_args = config.model_config.dsv4_args
+        sizes = _dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages)
+        return dsv4_pool_bytes(sizes, dsv4_args, config.max_running_req + 1)
+
     def rebuild_from_config(
         self, config, num_pages: int, *, num_swa_pages: int | None = None
     ) -> None:
diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py
new file mode 100644
index 0000000000..229431b7a4
--- /dev/null
+++ b/tests/engine/test_rebuild_runtime_cache_ordering.py
@@ -0,0 +1,199 @@
+"""Engine.rebuild_runtime_cache resize ordering (#643): shrinking resizes must run
+before growing ones. Each pool frees only its own old tensors before reallocating, so a
+grow that runs while a still-large sibling pool is resident can OOM even though the
+TARGET geometry passed validate_rebuild. Drives the real Engine method unbound on a
+SimpleNamespace engine shell (the idiom of test_kv_cache_rebuild.py's page-table test),
+with fake pools mirroring each real pool's free-own-then-alloc sequence under a shared
+accounting ledger, and the real BaseKVCachePool.validate_rebuild as the fit check."""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+import torch
+
+from freetoken.engine.engine import Engine
+from freetoken.kvcache.base import BaseKVCachePool
+
+
+class _Ledger:
+    """Shared GPU-memory stand-in: cap in units (1 unit = 1 MoE slot / KV page)."""
+
+    def __init__(self, cap: int):
+        self.cap = cap
+        self.used = 0
+        self.peak = 0
+        self.events: list[tuple] = []
+
+    def alloc(self, tag: str, n: int) -> None:
+        if self.used + n > self.cap:
+            self.events.append(("OOM", tag, n, self.used))
+            raise torch.OutOfMemoryError(
+                f"alloc {tag}={n} would peak at {self.used + n} > cap {self.cap}"
+            )
+        self.used += n
+        self.peak = max(self.peak, self.used)
+        self.events.append(("alloc", tag, n, self.used))
+
+    def free(self, tag: str, n: int) -> None:
+        self.used -= n
+        self.events.append(("free", tag, n, self.used))
+
+
+class _FakeMoeCache:
+    """Mirrors OffloadMoeCache.rebuild: validate -> free own tensors -> cache_size
+    mutated (offload_cache.py:491, before the alloc loop) -> reallocate."""
+
+    def __init__(self, ledger: _Ledger, size: int):
+        self.ledger = ledger
+        self.num_experts = 8
+        self.max_slots = None
+        self.layout = None
+        self.quant_format = "none"
+        # one uint8 per expert row -> expert_bytes_per_slot == 1 unit
+        self.bank_sources = {"w13": [torch.zeros((self.num_experts, 1), dtype=torch.uint8)]}
+        self.cache_size = size
+        self._held = size
+        ledger.alloc("moe", size)
+
+    def validate_rebuild(self, cache_size: int) -> None:
+        if cache_size < self.num_experts:
+            raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}")
+
+    def rebuild(self, cache_size: int) -> None:
+        self.validate_rebuild(cache_size)
+        self.ledger.free("moe", self._held)
+        self._held = 0
+        self.cache_size = cache_size
+        self.ledger.alloc("moe", cache_size)
+        self._held = cache_size
+
+
+class _FakeKvPool(BaseKVCachePool):
+    """Mirrors MHAKVCache.rebuild (free the one buffer, then realloc); reuses the REAL
+    BaseKVCachePool.validate_rebuild with 1 unit per page and no fixed term."""
+
+    needs_rebind_on_rebuild = False
+
+    def __init__(self, ledger: _Ledger, num_pages: int):
+        self.ledger = ledger
+        self.pages = num_pages
+        ledger.alloc("kv", num_pages)
+
+    @classmethod
+    def kv_cost(cls, config):
+        return 1, 0, 16, 0
+
+    def rebuild_from_config(self, config, num_pages, *, num_swa_pages=None):
+        self.ledger.free("kv", self.pages)
+        self.ledger.alloc("kv", num_pages)
+        self.pages = num_pages
+
+    def attach_page_table(self, page_table):
+        pass
+
+    def unit_bytes(self):
+        return 1, 0
+
+    def k_cache(self, index): ...
+    def v_cache(self, index): ...
+    def store_kv(self, k, v, out_loc, layer_id): ...
+
+    @property
+    def device(self):
+        return torch.device("cpu")
+
+    @property
+    def dtype(self):
+        return torch.float16
+
+    @property
+    def num_layers(self):
+        return 1
+
+
+class _FakeGraphRunner:
+    def __init__(self, **kwargs):
+        self.graph_bs_list = [1, 2, 4]
+        self.destroy_cuda_graphs = lambda: None
+
+
+def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: int):
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None)
+    monkeypatch.setattr("freetoken.engine.engine.GraphRunner", _FakeGraphRunner)
+    config = types.SimpleNamespace(
+        page_size=16, max_seq_len=4096, max_running_req=4, memory_ratio=1.0,
+        cuda_graph_max_bs=4, swa_num_pages_override=None,
+        model_config=types.SimpleNamespace(vocab_size=128, model_is_mrope=False),
+    )
+    page_table = torch.zeros((config.max_running_req + 1, 256), dtype=torch.int32)
+    eng = types.SimpleNamespace(
+        config=config, moe_offload_cache=_FakeMoeCache(ledger, moe), kv_cache=_FakeKvPool(ledger, kv),
+        linear_state_pool=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None,
+        _baseline_free=budget, _weights_bytes=0, page_table=page_table, max_seq_len=4096,
+        ctx=types.SimpleNamespace(page_table=page_table),
+        dummy_req=types.SimpleNamespace(table_idx=0),
+        attn_backend=types.SimpleNamespace(reset_capture=lambda: None),
+        graph_runner=_FakeGraphRunner(),
+        _sync_get_memory=lambda: (0, 0),
+    )
+    eng._target_moe_and_expert_bytes = Engine._target_moe_and_expert_bytes.__get__(eng)
+    eng._resize_kv_pool = Engine._resize_kv_pool.__get__(eng)
+    eng._refresh_seq_state = Engine._refresh_seq_state.__get__(eng)
+    return eng
+
+
+def _alloc_index(ledger: _Ledger, tag: str, n: int) -> int:
+    return next(
+        i for i, ev in enumerate(ledger.events) if ev[0] == "alloc" and ev[1] == tag and ev[2] == n
+    )
+
+
+def test_rebuild_runtime_cache_applies_shrinks_before_grows(monkeypatch):
+    # moe 40->60 grows, kv 50->30 shrinks, target total 90 <= cap 100: the fit check
+    # passes, so the rebuild must complete live instead of OOMing on the MoE grow.
+    ledger = _Ledger(cap=100)
+    eng = _build_engine(monkeypatch, ledger, moe=40, kv=50, budget=100)
+
+    Engine.rebuild_runtime_cache(eng, moe_cache_size=60, num_pages=30)
+
+    assert eng.moe_offload_cache.cache_size == 60
+    assert eng.kv_cache.pages == 30 and eng.num_pages == 30
+    assert ledger.used == 90  # target total resident
+    assert ledger.peak == 90  # never above max(current, target) totals
+    assert _alloc_index(ledger, "kv", 30) < _alloc_index(ledger, "moe", 60)  # shrink ran first
+
+
+def test_rebuild_runtime_cache_reverse_direction_still_fits(monkeypatch):
+    # The mirror request (moe shrinks, kv grows) worked before the fix because the old
+    # fixed order happened to shrink MoE first; a naive static reorder would break it.
+    ledger = _Ledger(cap=100)
+    eng = _build_engine(monkeypatch, ledger, moe=60, kv=30, budget=100)
+
+    Engine.rebuild_runtime_cache(eng, moe_cache_size=40, num_pages=50)
+
+    assert eng.moe_offload_cache.cache_size == 40
+    assert eng.kv_cache.pages == 50 and eng.num_pages == 50
+    assert ledger.used == 90 and ledger.peak == 90
+    assert _alloc_index(ledger, "moe", 40) < _alloc_index(ledger, "kv", 50)  # shrink ran first
+
+
+def test_rebuild_runtime_cache_grow_failure_rolls_back(monkeypatch):
+    # Budget 110 prices the 105-unit target as fitting but the real cap is 100 (fragmentation
+    # / unaccounted allocations), so the grow phase OOMs after the shrink phase. The
+    # scheduler's rollback (rebuild_cache(**prior) -> another rebuild_runtime_cache) must
+    # restore the snapshot geometry. An OOM inside a KV/mamba realloc still wedges the
+    # rollback itself; that is #526, out of scope here.
+    ledger = _Ledger(cap=100)
+    eng = _build_engine(monkeypatch, ledger, moe=40, kv=50, budget=110)
+
+    with pytest.raises(torch.OutOfMemoryError):
+        Engine.rebuild_runtime_cache(eng, moe_cache_size=70, num_pages=35)
+    assert eng.rebuild_teardown_started  # the scheduler's mid-teardown signal
+
+    Engine.rebuild_runtime_cache(eng, moe_cache_size=40, num_pages=50)  # rollback to prior
+
+    assert eng.moe_offload_cache.cache_size == 40
+    assert eng.kv_cache.pages == 50 and eng.num_pages == 50
+    assert ledger.used == 90 and ledger.peak <= 100  # serving geometry restored under the cap
diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py
index 0e3b7050e8..b4596f66a0 100644
--- a/tests/kvcache/test_kv_cache_rebuild.py
+++ b/tests/kvcache/test_kv_cache_rebuild.py
@@ -269,9 +269,15 @@ def test_every_kv_pool_answers_the_sizing_surface():
     from freetoken.kvcache.mha_pool import MHAKVCache
 
     for cls in (MHAKVCache, MLAKVCache, DSAKVCache, HybridSWAKVCache, DSV4PagedKVCache):
-        for hook in ("kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild"):
+        for hook in (
+            "kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild",
+            "rebuild_footprint",
+        ):
             assert callable(getattr(cls, hook, None)), f"{cls.__name__} is missing {hook}"
-    for hook in ("kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild"):
+    for hook in (
+        "kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild",
+        "rebuild_footprint",
+    ):
         assert hook in DSV4PagedKVCache.__dict__, f"DSV4 lost its {hook} override"
 
 

From 803a3b48023b5411c27c5656c35f859d988ca36b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 8 Oct 2026 13:31:44 -0700
Subject: [PATCH 078/249] fix(kernel): correct jit tensor matcher checks and
 pdl arch guards (#638)

* fix(kernel): correct jit matcher options and pdl arch guards

- with_dtype<Ts...>(sym) and with_device<Codes...>(sym) applied no options, so their dtype and device checks accepted anything
- index and store list kDLROCM, which the ignored check let through
- emit griddepcontrol only for !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900; it failed ptxas below sm_90 and does not exist on HIP
- declare the launch attribute cache only on CUDA
- use kWarpThreads for the index and store launch geometry

Co-authored-by: zihaomu <zihaomu@outlook.com>

* fix(kernel): treat zero as a matched size in tensor matcher

- a zero dimension counted as unset, so empty inputs failed with "Size value is not set" and a zero-row tensor did not catch a mismatched one

* fix(kernel): return early on empty index and store batches

- an empty batch launched a zero-sized grid, which CUDA and HIP reject; the leftover error then failed the next GPU call

Co-authored-by: zihaomu <zihaomu@outlook.com>

---------

Co-authored-by: zihaomu <zihaomu@outlook.com>
---
 .../kernel/csrc/include/freetoken/tensor.h    | 11 +++++++--
 .../kernel/csrc/include/freetoken/utils.cuh   |  7 ++++++
 python/freetoken/kernel/csrc/jit/index.cu     | 23 ++++++++++++-------
 python/freetoken/kernel/csrc/jit/store.cu     | 15 ++++++++----
 4 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/python/freetoken/kernel/csrc/include/freetoken/tensor.h b/python/freetoken/kernel/csrc/include/freetoken/tensor.h
index 9b591a87dc..3f711a5f7b 100644
--- a/python/freetoken/kernel/csrc/include/freetoken/tensor.h
+++ b/python/freetoken/kernel/csrc/include/freetoken/tensor.h
@@ -34,7 +34,7 @@ struct DeviceRef;
 
 inline constexpr auto kAnyDeviceID = -1;
 inline constexpr auto kAnySize = static_cast<int64_t>(-1);
-inline constexpr auto kNullSize = static_cast<int64_t>(0);
+inline constexpr auto kNullSize = static_cast<int64_t>(-1);
 inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
 inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
 
@@ -147,7 +147,7 @@ struct SymbolicSize {
     RuntimeCheck(!this->has_value(), "Size value already set");
     m_value = value;
   }
-  auto has_value() const -> bool { return m_value != 0; }
+  auto has_value() const -> bool { return m_value != details::kNullSize; }
   auto get_value() const -> std::optional<int64_t> {
     return this->has_value() ? std::optional{m_value} : std::nullopt;
   }
@@ -385,6 +385,10 @@ struct TensorMatcher {
   auto with_dtype(DTypeRef &&dtype) && -> TensorMatcher && {
     m_init_dtype();
     m_dtype.rebind(*dtype);
+    // no options given: keep the ones already set on the shared symbol
+    if constexpr (sizeof...(Ts) > 0) {
+      m_dtype->set_options<Ts...>();
+    }
     return std::move(*this);
   }
 
@@ -400,6 +404,9 @@ struct TensorMatcher {
   auto with_device(DeviceRef &&device) && -> TensorMatcher && {
     m_init_device();
     m_device.rebind(*device);
+    if constexpr (sizeof...(Codes) > 0) {
+      m_device->set_options<Codes...>();
+    }
     return std::move(*this);
   }
 
diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
index 72a4959698..091642a8fd 100644
--- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
+++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh
@@ -43,16 +43,21 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const
 
 namespace PDL {
 
+// griddepcontrol needs sm_90+; older archs and HIP keep normal stream ordering.
 template <bool kUsePDL> __always_inline __device__ void wait() {
+#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900
   if constexpr (kUsePDL) {
     asm volatile("griddepcontrol.wait;" ::: "memory");
   }
+#endif
 }
 
 template <bool kUsePDL> __always_inline __device__ void launch() {
+#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900
   if constexpr (kUsePDL) {
     asm volatile("griddepcontrol.launch_dependents;" :::);
   }
+#endif
 }
 
 } // namespace PDL
@@ -144,7 +149,9 @@ private:
     return config;
   }
   cudaLaunchConfig_t m_config;
+#if !FREETOKEN_USE_ROCM
   cudaLaunchAttribute m_attr_cache;
+#endif
 };
 
 } // namespace host
diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu
index ca0e1db26e..4832e0a460 100644
--- a/python/freetoken/kernel/csrc/jit/index.cu
+++ b/python/freetoken/kernel/csrc/jit/index.cu
@@ -35,9 +35,10 @@ __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void //
   using namespace device;
   constexpr auto kSize = kElementSize;
   constexpr auto kSizePerWarp = kSize / kNumSplits;
-  constexpr auto kWarpPerBlock = static_cast<unsigned>(kNumThreads / 32);
+  constexpr auto kWarpPerBlock =
+      static_cast<unsigned>(kNumThreads / kWarpThreads);
 
-  static_assert(kNumThreads % 32 == 0);
+  static_assert(kNumThreads % kWarpThreads == 0);
   static_assert(std::has_single_bit(kNumSplits));
   static_assert(kElementSize % kNumSplits == 0);
 
@@ -66,9 +67,10 @@ __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void //
   using namespace device;
   constexpr auto kSize = kElementSize;
   constexpr auto kSizePerWarp = kSize / kNumSplits;
-  constexpr auto kWarpPerBlock = static_cast<unsigned>(kNumThreads / 32);
+  constexpr auto kWarpPerBlock =
+      static_cast<unsigned>(kNumThreads / kWarpThreads);
 
-  static_assert(kNumThreads % 32 == 0);
+  static_assert(kNumThreads % kWarpThreads == 0);
   static_assert(std::has_single_bit(kNumSplits));
   static_assert(kElementSize % kNumSplits == 0);
 
@@ -114,15 +116,15 @@ struct IndexKernel {
 
     TensorMatcher({-1, D}) //
         .with_dtype(weights_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(weights);
     TensorMatcher({L, D}) //
         .with_dtype(weights_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(output);
     TensorMatcher({L}) //
         .with_dtype<int32_t, int64_t>(indices_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(indices);
 
     const auto device = device_.unwrap();
@@ -132,7 +134,12 @@ struct IndexKernel {
     RuntimeCheck(entry_size == element_size,
                  "IndexKernel: element_size mismatch.");
 
-    constexpr auto kWarpPerBlock = num_threads / 32;
+    // Empty batches are valid no-ops; do not launch a zero-sized grid.
+    if (num_indices == 0) {
+      return;
+    }
+
+    constexpr auto kWarpPerBlock = num_threads / device::kWarpThreads;
     const auto num_warps = num_splits * num_indices;
     const auto num_blocks = div_ceil(num_warps, kWarpPerBlock);
     const auto params = IndexKernelParams{
diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu
index 8d84d76ef1..18fd45346a 100644
--- a/python/freetoken/kernel/csrc/jit/store.cu
+++ b/python/freetoken/kernel/csrc/jit/store.cu
@@ -72,18 +72,18 @@ struct StoreKernel {
 
     TensorMatcher({-1, D}) //
         .with_strides({X, 1})
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype(dtype_)
         .verify(k_cache)
         .verify(v_cache);
     TensorMatcher({L, D}) //
         .with_strides({Y, 1})
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype(dtype_)
         .verify(k)
         .verify(v);
     TensorMatcher({L}) //
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype<int32_t, int64_t>(indices_dtype_)
         .verify(indices);
 
@@ -96,6 +96,11 @@ struct StoreKernel {
     const auto kv_cache_stride = X.unwrap() * dtype_size;
     const auto kv_input_stride = Y.unwrap() * dtype_size;
 
+    // Empty batches are valid no-ops; do not launch a zero-sized grid.
+    if (length == 0) {
+      return;
+    }
+
     const auto params = StoreKernelParams{
         .k_cache = k_cache.data_ptr(),
         .v_cache = v_cache.data_ptr(),
@@ -107,8 +112,8 @@ struct StoreKernel {
         .length = length,
     };
 
-    constexpr auto kWarpPerBlock = num_threads / 32;
-    static_assert(num_threads % 32 == 0);
+    constexpr auto kWarpPerBlock = num_threads / device::kWarpThreads;
+    static_assert(num_threads % device::kWarpThreads == 0);
     const auto num_blocks = div_ceil(length, kWarpPerBlock);
     const auto kernel = use_int32
                             ? store_kv_cache<num_threads, max_concurrency,

From 23fc922bf8df728a630674b1322d55626d0261da Mon Sep 17 00:00:00 2001
From: Mukkesh Ganesh <mukmckenzie@gmail.com>
Date: Thu, 8 Oct 2026 14:42:39 -0700
Subject: [PATCH 079/249] feat(deepseek_v4): serve image input on the
 V4-Flash-Vision release (#502)

* feat(deepseek_v4): serve image input on the V4-Flash-Vision release

The block layout depends on where the image lands: lead pads align its first row token to
a 4-token compression stride, so the item hash carries the layout and two alignments never
share a radix key.

The block's tokens all carry one content pad id (the shared pad scheme); the sparse
attention reads its spans from batch.mm_spans, which the scheduler recovers from the flat
block ends its existing plan already reports.

* fix(mm): keep the shared mmprocessor signature and own the dsv4 offset in its apply

* fix(models): build dsv4 image spans in the model and start them at image_start

* fix(models): mount the dsv4 vision tower under visual and keep its weights unquantized

* fix(models): keep dsv4 image blocks whole across chunked prefill
---
 docs/models.md                                |   3 +-
 python/freetoken/mm/processor.py              |   3 +-
 python/freetoken/mm/processors/deepseek_v4.py | 129 +++++++
 python/freetoken/models/config.py             |   3 +
 .../freetoken/models/deepseek_v4/__init__.py  |  11 +-
 python/freetoken/models/deepseek_v4/args.py   |  12 +
 .../freetoken/models/deepseek_v4/attention.py | 106 +++++-
 python/freetoken/models/deepseek_v4/config.py |  59 ++-
 .../models/deepseek_v4/image_processor.py     |  50 +++
 python/freetoken/models/deepseek_v4/model.py  |  50 ++-
 python/freetoken/models/deepseek_v4/moe.py    |  29 +-
 python/freetoken/models/deepseek_v4/vision.py | 336 ++++++++++++++++++
 python/freetoken/models/deepseek_v4/weight.py |  66 +++-
 python/freetoken/models/register.py           |  13 +-
 scripts/ftw_hotfix.py                         |   2 +-
 tests/README.md                               |   2 +
 tests/checkpoint/test_ftw_hotfix_vision.py    |   1 +
 tests/mm/test_processor.py                    | 103 +++++-
 tests/models/test_dsv4_vision.py              | 335 +++++++++++++++++
 tests/scheduler/test_mm.py                    |   8 +-
 tests/tokenizer/test_mm_tokenize.py           |  55 ++-
 21 files changed, 1342 insertions(+), 34 deletions(-)
 create mode 100644 python/freetoken/mm/processors/deepseek_v4.py
 create mode 100644 python/freetoken/models/deepseek_v4/image_processor.py
 create mode 100644 python/freetoken/models/deepseek_v4/vision.py
 create mode 100644 tests/models/test_dsv4_vision.py

diff --git a/docs/models.md b/docs/models.md
index b999d29e41..cdadc2dcb1 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -6,7 +6,7 @@ for them; other checkpoints of the same architectures work too.
 | Model | HF checkpoints |
 |---|---|
 | DeepSeek-V4.1-Flash | [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) |
-| DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) |
+| DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731), [deepseek-ai/DeepSeek-V4-Flash-Vision-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp) |
 | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) |
 | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) |
 | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) |
@@ -35,6 +35,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 | GLM-5.3-Flash (`glm5_next`: ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution on a canvas zero-padded to a 28-multiple | token counts, passed through as the processor's `min_image_tokens` / `max_image_tokens`; checkpoint defaults 16 to 8000 tokens | `{"max_image_tokens": 2048}` |
 | Muse-Glimmer-30B (windowed ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, aspect ratio kept under a token cap; checkpoint default 4096 tokens | the maximum is the cap (`max_image_tokens`); the minimum has no effect | `{"max_image_tokens": 1024}` |
 | MiniMax-M3 (`minimax_m3`: CLIP-style ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 4 to 576 tokens | `{"size": {"longest_edge": 1048576}}` |
+| DeepSeek-V4-Flash-Vision (`deepseek_v4`: native-resolution ViT tower, streamed under `--mm-encoder-weights host`) | one token per 3x3 grid of 14x14 patches, against a per-image cap; the block carries sentinel tokens and lead pads that align it to a 4-token compression stride, so its length depends on where it lands in the prompt | the maximum is a cap on the block (`vision_max_n_token`, checkpoint default 384) and lowers it, the minimum is read as the checkpoint's upscale floor and has no effect | not read |
 
 ## MoE strategies
 
diff --git a/python/freetoken/mm/processor.py b/python/freetoken/mm/processor.py
index 48211a5b80..84463bc2eb 100644
--- a/python/freetoken/mm/processor.py
+++ b/python/freetoken/mm/processor.py
@@ -200,7 +200,8 @@ def get_mm_processor(model_path: str, mm: MultimodalConfig | None = None) -> MMP
     served = [e for e in spec.encoders if getattr(config, e.config_key, None) is not None and e.kind not in mm.disabled_encoders]
     if spec.mm_processor is None or not served:
         return None
-    check_mm_pad_shift(config.text_config.vocab_size)
+    # families with a flat checkpoint config (deepseek_v4) carry vocab_size at the top level
+    check_mm_pad_shift(getattr(config, "text_config", config).vocab_size)
     module, _, cls = spec.mm_processor.partition(":")
     return getattr(importlib.import_module(module), cls)(config, model_path, mm)
 
diff --git a/python/freetoken/mm/processors/deepseek_v4.py b/python/freetoken/mm/processors/deepseek_v4.py
new file mode 100644
index 0000000000..5cea991994
--- /dev/null
+++ b/python/freetoken/mm/processors/deepseek_v4.py
@@ -0,0 +1,129 @@
+"""DeepSeek-V4 image processor: each ``<｜deepseek_image｜>`` expands into the image's block of tokens.
+
+Where the reference streams ``vocab_size + type`` pseudo ids, the whole span carries one
+content pad id and the block's embeddings -- sentinels included -- come from the encoder.
+"""
+
+from __future__ import annotations
+
+import io
+import struct
+from typing import Any
+
+import torch
+
+from freetoken.message import MMItem
+from freetoken.mm import mm_pad_value
+from freetoken.mm.config import MultimodalConfig
+from freetoken.mm.processor import MMResult, MMProcessor, PromptReplacement, _find_all, content_hash
+from freetoken.models.deepseek_v4.config import IMAGE_TOKEN_ID, DSV4VisionConfig, parse_vision_config
+from freetoken.models.deepseek_v4.image_processor import load_image
+from freetoken.models.deepseek_v4.vision import COMPRESS_PAD_TO, build_image_block
+
+
+class DSV4MMProcessor(MMProcessor):
+    def __init__(self, hf_config: Any, model_path: str, mm: MultimodalConfig) -> None:
+        super().__init__(model_path, mm)
+        vc = parse_vision_config(hf_config)
+        if vc is None:
+            raise ValueError(f"{model_path} carries no vision tower (vision_n_layers <= 0)")
+        self.vc: DSV4VisionConfig = vc
+        self.image_token_id = IMAGE_TOKEN_ID
+        self.placeholder = [self.image_token_id]
+        # --image-max-tokens lowers the checkpoint's per-image block cap
+        self.max_n_token = min(vc.vision_max_n_token, mm.image_max_tokens) if mm.image_max_tokens else vc.vision_max_n_token
+
+    def process(self, images: list[Any]) -> list[MMItem]:
+        items: list[MMItem] = []
+        for image in images:
+            patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(image, self.vc, self.max_n_token)
+            # layout, hash and token count are only known at the insertion offset
+            items.append(
+                MMItem(
+                    modality="image",
+                    hash=0,
+                    pad_value=0,
+                    offsets=[],
+                    feature=patches,
+                    model_specific_data={
+                        "n_vit_h": n_vit_h,
+                        "n_vit_w": n_vit_w,
+                        "n_llm_h": n_llm_h,
+                        "n_llm_w": n_llm_w,
+                    },
+                )
+            )
+        return items
+
+    def prompt_replacement(self, item: MMItem, start: int) -> PromptReplacement:
+        """The DSV4-only ``start`` is where the block lands in input_ids; only our apply below calls this."""
+        n_llm_h, n_llm_w = item.n_llm_h, item.n_llm_w
+        types, _perm = build_image_block(n_llm_h, n_llm_w, start)
+        # the lead pads align the block, so the same image at another offset is another
+        # sequence: the layout belongs in the hash or the two would share a radix key
+        compress_pad = COMPRESS_PAD_TO - 1 - start % COMPRESS_PAD_TO
+        item.hash = content_hash(item.feature, struct.pack("<3i", n_llm_h, n_llm_w, compress_pad))
+        item.pad_value = mm_pad_value(item.hash)
+        item.model_specific_data["start"] = start
+        return PromptReplacement([self.image_token_id] * types.numel())
+
+    def apply(self, input_ids: torch.Tensor, images: list[bytes]) -> MMResult:
+        """MMProcessor.apply, save that each replacement is built knowing the offset it lands at."""
+        from PIL import Image
+
+        ids = input_ids.tolist()
+        target = self.placeholder
+        slots = _find_all(ids, target)
+        if len(slots) != len(images):
+            raise ValueError(
+                f"prompt renders {len(slots)} image placeholders but the request "
+                f"carries {len(images)} images"
+            )
+
+        pils = [Image.open(io.BytesIO(raw)).convert("RGB") for raw in images]
+        items = self.process(pils)
+        out: list[int] = []
+        cursor = 0
+        for slot, item in zip(slots, items):
+            out.extend(ids[cursor:slot])
+            base = len(out)
+            repl = self.prompt_replacement(item, base)
+            full = list(repl.full)
+            spans = repl.embed_spans()
+            # embedding slots carry the content pad id so radix keys and the model's scatter mask see the image
+            for lo, hi in spans:
+                full[lo:hi] = [item.pad_value] * (hi - lo)
+            out.extend(full)
+            item.offsets = [[base + lo, base + hi] for lo, hi in spans]
+            item.validate()
+            cursor = slot + len(target)
+        out.extend(ids[cursor:])
+        new_ids = torch.tensor(out, dtype=input_ids.dtype)
+
+        positions, delta = self.positions(len(out), items) or (None, 0)
+        return MMResult(new_ids, items, positions, delta)
+
+    def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]:
+        # The smallest block the tower can run: one 2x2 merge grid out of a 6x6 patch grid.
+        r, p = self.vc.vision_downsample_ratio, self.vc.vision_patch_size
+        n_vit, n_llm = 2 * r, 2
+        types, _perm = build_image_block(n_llm, n_llm, 0)
+        return [
+            MMItem(
+                modality="image",
+                hash=0,
+                pad_value=0,
+                offsets=[[0, types.numel()]],
+                feature=torch.zeros(n_vit * n_vit, 3, p, p, dtype=dtype, device=device),
+                model_specific_data={
+                    "n_vit_h": n_vit,
+                    "n_vit_w": n_vit,
+                    "n_llm_h": n_llm,
+                    "n_llm_w": n_llm,
+                    "start": 0,
+                },
+            )
+        ]
+
+
+__all__ = ["DSV4MMProcessor"]
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index 55b90bb58c..25bc73f5e4 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -216,6 +216,9 @@ class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attenti
     sliding_window: int  # the P-token window page
     # "v41": DeepSeek-V4.1's cross-layer-shared packed tiers (AttnType.DSV41, its own pool and backend)
     variant: Literal["v4", "v41"] = "v4"
+    # image token spans attend to each other in both directions, so prefill chunks end
+    # before a block instead of splitting it (the scheduler's keep_images_whole gate)
+    bidirectional_mm_blocks: bool = False
 
 
 AttentionGroupConfig: TypeAlias = (
diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py
index e19438d1ba..5bbb9b8e72 100644
--- a/python/freetoken/models/deepseek_v4/__init__.py
+++ b/python/freetoken/models/deepseek_v4/__init__.py
@@ -14,15 +14,22 @@
 """
 
 from .args import DeepseekV4Args, load_args
-from .config import parse_config
+from .config import DSV4VisionConfig, parse_config, parse_vision_config
 from .model import DeepseekV4ForCausalLM
-from .weight import iter_expert_pieces, iter_weights
+from .vision import DSV4Aligner, DSV4Vision, DSV4VisionTower
+from .weight import iter_expert_pieces, iter_vision_weights, iter_weights
 
 __all__ = [
     "DeepseekV4Args",
+    "DSV4Aligner",
+    "DSV4Vision",
+    "DSV4VisionConfig",
+    "DSV4VisionTower",
     "load_args",
     "parse_config",
+    "parse_vision_config",
     "DeepseekV4ForCausalLM",
     "iter_weights",
     "iter_expert_pieces",
+    "iter_vision_weights",
 ]
diff --git a/python/freetoken/models/deepseek_v4/args.py b/python/freetoken/models/deepseek_v4/args.py
index b75d2800df..d2646a2b3b 100644
--- a/python/freetoken/models/deepseek_v4/args.py
+++ b/python/freetoken/models/deepseek_v4/args.py
@@ -69,6 +69,18 @@ class DeepseekV4Args:
     hc_sinkhorn_iters: int = 20
     hc_eps: float = 1e-6
 
+    # ----- vision tower (DeepSeek-V4-Flash-Vision-Exp; 0 on the text-only releases) -----
+    vision_n_layers: int = 0
+    vision_dim: int = 1024
+    vision_n_heads: int = 16
+    vision_inter_dim: int = 2816
+    vision_patch_size: int = 14
+    vision_rope_theta: float = 10000.0
+    vision_downsample_ratio: int = 3
+    vision_max_n_token: int = 384
+    vision_min_pixels: int = 147456
+    vision_max_wh_ratio: int = 8
+
     def __post_init__(self) -> None:
         # JSON lists -> tuple so the dataclass stays hashable / immutable-ish.
         if isinstance(self.compress_ratios, list):
diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py
index cb0dd3150b..91bc8ac095 100644
--- a/python/freetoken/models/deepseek_v4/attention.py
+++ b/python/freetoken/models/deepseek_v4/attention.py
@@ -38,6 +38,7 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, *, quant_config=None, pr
         self.window_size = args.window_size
         self.compress_ratio = args.compress_ratios[layer_id]
         self.eps = args.norm_eps
+        self.max_image_tokens = args.vision_max_n_token
 
         self.attn_sink = torch.empty(self.n_heads, dtype=torch.float32)
         # the latent projections are replicated; wq_b shards over heads, wo_b over the output groups
@@ -104,7 +105,53 @@ def _wo(self, o: torch.Tensor, bsz: int, seqlen: int) -> torch.Tensor:
         o = torch.einsum("bsgd,grd->bsgr", o, wo_a).flatten(2)
         return self.wo_b.forward(o)
 
-    def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int):
+    @staticmethod
+    def visible_window_cols(
+        spans, start_pos: int, n: int, win: int, max_img: int, device: torch.device,
+        addr_lo: int = 0, width: int | None = None,
+    ) -> torch.Tensor:
+        """Window candidate columns (``-1`` padded) for a segment containing image spans.
+
+        Port of the reference ``get_image_visible`` + ``get_window_topk_idxs_visible``
+        pair: span tokens see their whole span bidirectionally (capped at ``max_img``)
+        instead of the 128-window, while text tokens -- and span tokens near the span
+        start -- keep the ordinary window into history. Returns ``[n, width]`` in
+        request-absolute coordinates with live entries in ``[addr_lo, span_end)``; the
+        caller translates to slot-lut indices (cold: as-is; extend: minus ``w_lo``).
+        """
+        if width is None:
+            width = win + max_img
+        left, right = Attention._span_left_right(spans, start_pos, n, max_img, device)
+        idx = start_pos + torch.arange(n, device=device)
+        left_add = (left - (win - 1)).clamp(min=0)
+        starts = (idx - (win - 1) - left_add).clamp(min=addr_lo)
+        cols = starts.unsqueeze(1) + torch.arange(width, device=device)
+        return torch.where(cols > (idx + right).unsqueeze(1), -1, cols)
+
+    @staticmethod
+    def _span_left_right(
+        spans, start_pos: int, n: int, max_img: int, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Per-token visible counts to the left/right within each image span.
+
+        Metadata twin of the reference ``get_image_visible`` (which derives spans from
+        virtual token ids): ``left[q] = q - span_start`` clamped to ``max_img - 1``,
+        ``right[q] = span_end - q`` clamped to ``max_img``, zero outside spans. ``spans``
+        are request-absolute positions inside ``[start_pos, start_pos + n)``.
+        """
+        left = torch.zeros(n, dtype=torch.int64, device=device)
+        right = torch.zeros(n, dtype=torch.int64, device=device)
+        for s, e in spans:
+            lo, hi = s - start_pos, e - start_pos
+            idx = torch.arange(lo, hi, device=device)
+            left[lo:hi] = (idx - lo).clamp(max=max_img - 1)
+            # the END sentinel sits at e-1: visibility reaches the END TOKEN, not past it
+            right[lo:hi] = ((hi - 1) - idx).clamp(max=max_img)
+        return left, right
+
+    def _prefill_segment(
+        self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int, spans=None
+    ):
         """One request's prefill work inside a (possibly ragged) batch: persist its window KV,
         advance its compressor/indexer carry, and return its per-query candidate lists.
 
@@ -113,6 +160,16 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in
         ``(win_global [1, n, w], cmp_global [1, n, c] | None)`` at natural widths (cold: w =
         min(n, win); extend: w = win); the caller pads BOTH halves to the batch-uniform widths
         and concatenates.
+
+        ``spans`` (VL requests) are the request's image-block ranges ``[(start, end)]`` inside
+        ``[start_pos, start_pos + n)``. Vision checkpoints set ``bidirectional_mm_blocks`` on the
+        attention group, which makes prefill end a chunk before a block instead of splitting it;
+        only a block longer than one chunk budget is still cut. Inside a span, every token
+        sees the whole span (bidirectional, capped at ``vision_max_n_token``) instead of only its
+        128-window, matching the reference's ``get_window_topk_idxs_visible``. All span-visible
+        candidates are intra-segment, so the live full-page map covers them (the 128-slot ring
+        never needs to serve >128-back reads). ``spans=None`` keeps the historical path
+        bit-identical.
         """
         win, ratio, device = self.window_size, self.compress_ratio, x_seg.device
         end = start_pos + n
@@ -120,23 +177,38 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in
         self.attn.store_window(kv_seg, self.layer_id, slots)
 
         if start_pos == 0:
-            win_cols = get_window_topk_idxs(win, 1, n, 0).to(device)
-            # natural width min(n, win); the caller pads to the batch-uniform width
+            if spans:
+                max_img = self.max_image_tokens
+                cols = Attention.visible_window_cols(
+                    spans, 0, n, win, max_img, device, addr_lo=0,
+                    width=min(n, win + max_img),
+                )
+                win_cols = cols.unsqueeze(0)
+            else:
+                win_cols = get_window_topk_idxs(win, 1, n, 0).to(device)
+            # natural width min(n, win) -- win+max_img for span segments; the caller pads
             win_global = self.attn.win_cols_to_global(win_cols, slots)
         else:
             # Candidates span [w_lo, end): each new query's 128-sliding window over the retained
             # prefix plus the new tokens, causal-masked to -1 past its own position.
             w_lo = max(0, start_pos - win + 1)
             ws_pool = self.attn.window_slots_of(ti, w_lo, end)
-            abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1)
-            cand = (abs_p - win + 1).clamp(min=w_lo) + torch.arange(win, device=device)
-            win_cols = torch.where(cand > abs_p, -1, cand - w_lo).unsqueeze(0)
+            if spans:
+                max_img = self.max_image_tokens
+                cols = Attention.visible_window_cols(
+                    spans, start_pos, n, win, max_img, device, addr_lo=w_lo,
+                )
+                win_cols = torch.where(cols < 0, cols, cols - w_lo).unsqueeze(0)
+            else:
+                abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1)
+                cand = (abs_p - win + 1).clamp(min=w_lo) + torch.arange(win, device=device)
+                win_cols = torch.where(cand > abs_p, -1, cand - w_lo).unsqueeze(0)
             win_global = self.attn.win_cols_to_global(win_cols, ws_pool)
 
         if not ratio:
             return win_global, None
-        # Only the compressor/indexer read the matched tail page's slot; resolving it costs a
-        # host sync (.item()), so do it after the ratio-0 early-out.
+        # Only the compressor/indexer read the matched tail page's slot; resolving it costs
+        # a host sync (.item()), so do it after the ratio-0 early-out.
         tail_ws = (
             int(self.attn.window_slots_of(ti, start_pos - 1, start_pos).item())
             if start_pos > 0 else None
@@ -157,7 +229,7 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in
             self.compressor.forward(x_seg, start_pos, slots, tail_window_slot=tail_ws, ti=ti)
         return win_global, self.attn.blocks_to_global(blocks, ratio, ti=ti)
 
-    def forward_ragged(self, x, segments, flat_positions):
+    def forward_ragged(self, x, segments, flat_positions, spans_by_ti=None):
         """Ragged batched prefill (cu_seqlens). ``x`` is [1, T, dim] -- the requests' NEW token
         streams concatenated (NO padding); ``segments`` is [(offset, n, table_idx, start_pos)]
         tiling [0, T); ``flat_positions`` [T] is each token's ABSOLUTE position, so a radix-hit
@@ -197,9 +269,18 @@ def forward_ragged(self, x, segments, flat_positions):
         win_parts: list[torch.Tensor] = []
         cmp_parts: list[torch.Tensor | None] = []
         max_c = 0
+        any_spans = False
         for off, n, ti, start_pos in segments:
+            spans = None
+            if spans_by_ti is not None:
+                req_spans = spans_by_ti.get(ti)
+                if req_spans:
+                    end = start_pos + n
+                    spans = [(s, e) for (s, e) in req_spans if s >= start_pos and e <= end]
+                    any_spans = any_spans or bool(spans)
             win_global, cmp_global = self._prefill_segment(
-                x[:, off:off + n], qr[:, off:off + n], kv[0, off:off + n], ti, start_pos, n
+                x[:, off:off + n], qr[:, off:off + n], kv[0, off:off + n], ti, start_pos, n,
+                spans=spans,
             )
             win_parts.append(win_global)
             cmp_parts.append(cmp_global)
@@ -213,6 +294,11 @@ def forward_ragged(self, x, segments, flat_positions):
         # is enough to flip a greedy near-tie. bs>1 pads every part to win (the historical
         # batched behavior).
         n_window = win if len(segments) > 1 else win_parts[0].shape[-1]
+        if any_spans:
+            # Span segments carry win+max_img-wide candidates; pad every part to the widest
+            # so the single flat launch stays rectangular. Text-only batches (any_spans
+            # False) keep the historical widths bit-identically.
+            n_window = max(n_window, *(w.shape[-1] for w in win_parts))
         flat = []
         for i, (off, n, ti, _start) in enumerate(segments):
             wg = win_parts[i]
diff --git a/python/freetoken/models/deepseek_v4/config.py b/python/freetoken/models/deepseek_v4/config.py
index 1005b8a3b3..9671c68290 100644
--- a/python/freetoken/models/deepseek_v4/config.py
+++ b/python/freetoken/models/deepseek_v4/config.py
@@ -13,12 +13,38 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig
 
 from .args import load_args
 
+# ``<｜deepseek_image｜>``: the only vocab-level image token. The reference expands each
+# occurrence into a block of ``vocab_size + type`` pseudo ids; here the block keeps this
+# one id and the image embeddings are scattered over the span it occupies.
+IMAGE_TOKEN_ID = 129264
+
+
+@dataclass(frozen=True)
+class DSV4VisionConfig:
+    """The tower's dims, split out of the checkpoint's flat inference config so the engine's
+    section nulling and the loader treat DSV4 like every other family with a vision tower.
+    Field names follow the reference (``inference/model.py`` ModelArgs) one to one."""
+
+    vision_n_layers: int
+    vision_dim: int
+    vision_n_heads: int
+    vision_inter_dim: int
+    vision_patch_size: int
+    vision_rope_theta: float
+    vision_downsample_ratio: int
+    vision_max_n_token: int
+    vision_min_pixels: int
+    vision_max_wh_ratio: int
+    # the language model's hidden size the aligner projects into (ModelArgs.dim)
+    text_dim: int
+
 
 def parse_config(hf_config: Any) -> ModelConfig:
     model_path = getattr(hf_config, "_name_or_path", None) or getattr(
@@ -46,6 +72,8 @@ def parse_config(hf_config: Any) -> ModelConfig:
     if max_position <= 0:
         max_position = int(args.original_seq_len * args.rope_factor)
 
+    vision_config = parse_vision_config(hf_config)
+
     return ModelConfig(
         num_layers=args.n_layers,
         num_qo_heads=args.n_heads,
@@ -88,9 +116,38 @@ def parse_config(hf_config: Any) -> ModelConfig:
                 num_kv_heads=1,  # MLA-style shared latent (K == V)
                 head_dim=args.head_dim,
                 sliding_window=args.window_size,
+                bidirectional_mm_blocks=vision_config is not None,
             ),
         ),
+        vision_config=vision_config,
+        image_token_id=IMAGE_TOKEN_ID if vision_config is not None else None,
+    )
+
+
+def parse_vision_config(hf_config: Any) -> DSV4VisionConfig | None:
+    """The tower's dims from the checkpoint config, or None when it carries no vision stack.
+
+    The vision dims sit at the top level of the checkpoint config, next to the language
+    model's, and the engine nulls ``vision_n_layers`` on the copy it hands the parser when
+    this process serves text-only (--text-model-only / --mm-disable). That null is what
+    decides, so the built tower always matches the weights the loader is told to read.
+    """
+    n_layers = int(getattr(hf_config, "vision_n_layers", 0) or 0)
+    if n_layers <= 0:
+        return None
+    return DSV4VisionConfig(
+        vision_n_layers=n_layers,
+        vision_dim=int(hf_config.vision_dim),
+        vision_n_heads=int(hf_config.vision_n_heads),
+        vision_inter_dim=int(hf_config.vision_inter_dim),
+        vision_patch_size=int(hf_config.vision_patch_size),
+        vision_rope_theta=float(hf_config.vision_rope_theta),
+        vision_downsample_ratio=int(hf_config.vision_downsample_ratio),
+        vision_max_n_token=int(hf_config.vision_max_n_token),
+        vision_min_pixels=int(hf_config.vision_min_pixels),
+        vision_max_wh_ratio=int(hf_config.vision_max_wh_ratio),
+        text_dim=int(hf_config.hidden_size),
     )
 
 
-__all__ = ["parse_config"]
+__all__ = ["IMAGE_TOKEN_ID", "DSV4VisionConfig", "parse_config", "parse_vision_config"]
diff --git a/python/freetoken/models/deepseek_v4/image_processor.py b/python/freetoken/models/deepseek_v4/image_processor.py
new file mode 100644
index 0000000000..16742b1e3e
--- /dev/null
+++ b/python/freetoken/models/deepseek_v4/image_processor.py
@@ -0,0 +1,50 @@
+"""DeepSeek-V4 image preprocessing: one image -> ViT patches plus its grid metadata.
+
+Port of the reference ``inference/image_processor.py:load_image``. The PIL work stays on the
+CPU and must be bit-identical across runs: an image's radix key rides on these patches.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+from .config import DSV4VisionConfig
+from .vision import safe_resize
+
+
+def load_image(image: Image.Image, vc: DSV4VisionConfig, max_n_token: int | None = None) -> tuple[torch.Tensor, int, int, int, int]:
+    """One image -> ``(patches bf16 [n_vit_h*n_vit_w, 3, p, p], n_vit_h, n_vit_w, n_llm_h, n_llm_w)``.
+
+    ``max_n_token`` lowers the checkpoint's per-image block cap (``--image-max-tokens``).
+    """
+    image = image.convert("RGB")
+    p = vc.vision_patch_size
+    width, height = image.size
+    if vc.vision_max_wh_ratio is not None and width > height * vc.vision_max_wh_ratio:
+        width = height * vc.vision_max_wh_ratio
+    if 0 < width * height < vc.vision_min_pixels:
+        ratio = (vc.vision_min_pixels / (width * height)) ** 0.5
+        width = int(width * ratio)
+        height = int(height * ratio)
+    best_width = math.ceil(width / p) * p
+    best_height = math.ceil(height / p) * p
+    n_llm_h, n_llm_w, best_height, best_width = safe_resize(
+        height, width, best_height, best_width, p, vc.vision_downsample_ratio,
+        vc.vision_max_n_token if max_n_token is None else min(vc.vision_max_n_token, max_n_token),
+    )
+    n_vit_h, n_vit_w = best_height // p, best_width // p
+    if vc.vision_max_wh_ratio is not None and image.width >= vc.vision_max_wh_ratio * image.height:
+        image = image.resize((best_width, best_height))
+    else:
+        image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
+    x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
+    x = ((x - 0.5) / 0.5).to(torch.bfloat16)
+    patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p)
+    return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
+
+
+__all__ = ["load_image"]
diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py
index 9faefb83fe..0c3441b951 100644
--- a/python/freetoken/models/deepseek_v4/model.py
+++ b/python/freetoken/models/deepseek_v4/model.py
@@ -29,11 +29,12 @@
 from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine
 from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn
 from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding
-from freetoken.models.blocks import BaseLLMModel
+from freetoken.models.blocks import BaseLLMModel, embed_input_ids
 
 from .args import DeepseekV4Args
 from .attention import Attention
 from .moe import MoE
+from .vision import COMPRESS_PAD_TO, DSV4Vision
 
 # Re-exports: keep every class/helper previously defined here importable from .model
 # (external import stability; the moved definitions live in their own modules).
@@ -85,7 +86,7 @@ def hc_post(self, x, residual, post, comb):
         )
         return y.view(shape)
 
-    def prefill_batched(self, x, input_ids, segments, flat_positions):
+    def prefill_batched(self, x, input_ids, segments, flat_positions, spans_by_ti=None):
         # Ragged batched prefill (cu_seqlens, no padding; bs >= 1, cold and radix-hit segments
         # mixed freely). ``x`` is [1, T, hc_mult, dim] -- the requests' token streams
         # concatenated. Per-token ops (HC / norm / MoE) run batched over ALL T tokens (the
@@ -93,10 +94,12 @@ def prefill_batched(self, x, input_ids, segments, flat_positions):
         # T queries (Attention.forward_ragged), with the stateful compressor/indexer looped per
         # request. ``segments`` = [(offset, extend_len, table_idx, start_pos)] off the attention
         # metadata; ``flat_positions`` [T] = per-token ABSOLUTE position (batch.positions).
+        # ``spans_by_ti`` maps a page-table row to its image spans: the sparse attention needs
+        # the pairs, the block-tiled backends read the flat ends (batch.mm_block_ends) instead.
         residual = x
         x, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
         x = self.attn_norm.forward(x)
-        x = self.attn.forward_ragged(x, segments, flat_positions)
+        x = self.attn.forward_ragged(x, segments, flat_positions, spans_by_ti=spans_by_ti)
         x = self.hc_post(x, residual, post, comb)
 
         residual = x
@@ -152,6 +155,7 @@ def hc_head(self, x):
 
     def prefill_batched(
         self, input_ids: torch.Tensor, segments, flat_positions: torch.Tensor,
+        spans_by_ti: dict | None = None,
     ) -> torch.Tensor:
         # Ragged batched prefill (bs >= 1). ``input_ids`` is [1, T] -- the requests' NEW tokens
         # concatenated (cu_seqlens, no padding); each request starts at its own cached_len
@@ -163,10 +167,13 @@ def prefill_batched(
         # metadata; ``flat_positions`` [T] is the scheduler-staged batch.positions (per-token
         # ABSOLUTE position); the head picks each request's final token off the attention
         # metadata -> its next-token logits row.
-        h = self.embed.forward(input_ids.view(-1)).view(1, -1, self.args.dim)
+        batch = get_global_ctx().batch
+        # on a chunk that carries image rows, batch.mm_embeds' leading columns replace the
+        # embedding of every row in batch.mm_rows -- the whole image span, sentinels included
+        h = embed_input_ids(self.embed, input_ids.view(-1), batch).view(1, -1, self.args.dim)
         h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
         for layer in self.layers.op_list:
-            h = layer.prefill_batched(h, input_ids, segments, flat_positions)
+            h = layer.prefill_batched(h, input_ids, segments, flat_positions, spans_by_ti)
         h = self.hc_head(h)
         h = self.norm.forward(h)
         return self.head.forward(h[0])  # [B, vocab]
@@ -200,6 +207,26 @@ def decode(
         return self.head.forward(h[:, -1])
 
 
+def _image_spans_by_ti(reqs) -> dict | None:
+    """The batch's image spans as ``(IMAGE_START position, block end)`` pairs per page-table row.
+
+    Only this model reads the pairs (built here, not staged on the batch): its sparse
+    attention builds candidate lists per span. The reference ``get_image_visible`` counts
+    visibility within [IMAGE_START, IMAGE_END], so the span starts at the START sentinel
+    (after ``build_image_block``'s alignment lead pads), not at the block's first token.
+    """
+    spans: dict[int, list[tuple[int, int]]] = {}
+    for req in reqs:
+        req_spans = [
+            (lo + COMPRESS_PAD_TO - 1 - lo % COMPRESS_PAD_TO, hi)
+            for item in req.mm_items or ()
+            for lo, hi in item.offsets
+        ]
+        if req_spans:
+            spans[req.table_idx] = req_spans
+    return spans or None
+
+
 class DeepseekV4ForCausalLM(BaseLLMModel):
     """Engine adapter: a registered :class:`BaseLLMModel` wrapping the DSV4 transformer.
 
@@ -211,6 +238,12 @@ def __init__(self, config):
         self._args: DeepseekV4Args = config.dsv4_args
         self.model = Transformer(self._args, config.quant, strategy=config.moe_strategy, decode_target=config.decode_target, prefix="model")
         self._bound = False
+        # Vision (DeepSeek-V4-Flash-Vision-Exp). config.is_multimodal is the engine's
+        # resolved gate, so a text-only process never builds weights the loader skips.
+        # The whole stack sits under one ``visual.`` mount, so the shared vision-key
+        # filter (VISION_KEY_PREFIXES) covers every image-only parameter.
+        if config.is_multimodal:
+            self.visual = DSV4Vision(config.vision_config, config.hidden_size)
 
     def _ensure_bound(self) -> None:
         if self._bound:
@@ -226,6 +259,12 @@ def mark_for_rebind(self) -> None:
         block count, freqs) depends on the new pool's geometry, so re-derive it via _ensure_bound."""
         self._bound = False
 
+    def place_encoder_weights(self, mode: str) -> None:
+        self.visual.place_weights(mode)
+
+    def encode(self, item):
+        return self.visual.forward(item)
+
     def forward(self) -> torch.Tensor:
         self._ensure_bound()
         batch = get_global_ctx().batch
@@ -240,6 +279,7 @@ def forward(self) -> torch.Tensor:
             # cross requests.
             return self.model.prefill_batched(
                 input_ids.view(1, -1), md.segments, batch.positions.long(),
+                spans_by_ti=_image_spans_by_ti(batch.reqs),
             )
         # DECODE (bs>=1): per-row position (GPU int tensor -> no host syncs / graph safe). The
         # compressed staging cap is the max position any row reaches (eager); a static max_seq-1
diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py
index 6b6245674d..35a316f103 100644
--- a/python/freetoken/models/deepseek_v4/moe.py
+++ b/python/freetoken/models/deepseek_v4/moe.py
@@ -9,24 +9,34 @@
 from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32
 from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu
 from freetoken.layers import BaseOP, LinearColParallelMerged, LinearRowParallel, OffloadMoELayer
+from freetoken.mm import MM_PAD_SHIFT_VALUE
 
 from .args import DeepseekV4Args
 
 
 class Gate(BaseOP):
-    """MoE router: sqrtsoftplus scoring + hash routing (first ``n_hash_layers``)."""
+    """MoE router: sqrtsoftplus scoring + hash routing (first ``n_hash_layers``).
+
+    A vision checkpoint gives every layer a ``bias_vl`` too: the expert-selection bias for
+    image rows, which never touches the routing weights. Hash layers keep routing text rows
+    through ``tid2eid`` and pick image rows by ``(scores + bias_vl).topk``, mirroring the
+    reference ``Gate.forward``.
+    """
 
     def __init__(self, layer_id: int, args: DeepseekV4Args):
         self.topk = args.n_activated_experts
         self.score_func = args.score_func
         self.route_scale = args.route_scale
         self.hash = layer_id < args.n_hash_layers
+        vl = args.vision_n_layers > 0
         self.weight = torch.empty(args.n_routed_experts, args.dim, dtype=torch.bfloat16)
         if self.hash:
             self.tid2eid = torch.empty(args.vocab_size, args.n_activated_experts, dtype=torch.int64)
-            self.bias = None
+            # present on hash layers of a vision checkpoint, loaded but never applied
+            self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32) if vl else None
         else:
             self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32)
+        self.bias_vl = torch.empty(args.n_routed_experts, dtype=torch.float32) if vl else None
 
     def forward(self, x: torch.Tensor, input_ids: torch.Tensor):
         scores = bf16_linear_fp32(x, self.weight)
@@ -37,11 +47,20 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor):
         else:
             scores = F.softplus(scores).sqrt()
         original_scores = scores
-        if self.bias is not None:
-            scores = scores + self.bias
+        # image rows are the content pad ids of an expanded image span
+        image_mask = input_ids >= MM_PAD_SHIFT_VALUE if self.bias_vl is not None else None
         if self.hash:
-            indices = self.tid2eid[input_ids]
+            if image_mask is None:
+                indices = self.tid2eid[input_ids]
+            else:
+                indices = self.tid2eid[torch.where(image_mask, 0, input_ids)]
+                vl_indices = (scores + self.bias_vl).topk(self.topk, dim=-1)[1]
+                indices = torch.where(image_mask.unsqueeze(-1), vl_indices.to(indices.dtype), indices)
         else:
+            if image_mask is None:
+                scores = scores + self.bias
+            else:
+                scores = scores + torch.where(image_mask.unsqueeze(-1), self.bias_vl, self.bias)
             indices = scores.topk(self.topk, dim=-1)[1]
         weights = original_scores.gather(1, indices)
         if self.score_func != "softmax":
diff --git a/python/freetoken/models/deepseek_v4/vision.py b/python/freetoken/models/deepseek_v4/vision.py
new file mode 100644
index 0000000000..44549460ef
--- /dev/null
+++ b/python/freetoken/models/deepseek_v4/vision.py
@@ -0,0 +1,336 @@
+"""DeepSeek-V4 vision tower, aligner and image-block layout.
+
+Port of the reference ``inference/vision.py`` / ``inference/image_processor.py`` from
+deepseek-ai/DeepSeek-V4-Flash-Vision-Exp. The tower, the aligner and the four sentinel
+vectors hang under one :class:`DSV4Vision` container the wrapper mounts as ``visual``,
+so their state-dict keys sit behind the ``visual.`` prefix every VLM family filters on
+(the checkpoint's own unprefixed names, ``vision.patch_embed.proj.weight``,
+``aligner.w1.weight``, ``image_start``, ..., are renamed on load); every tensor is bf16
+exactly as stored.
+
+Image tokens occupy plain sequential positions: the 2-D structure lives inside the
+tower (2-D RoPE over the patch grid) and in the N-layout of the token block the
+:func:`build_image_block` / :func:`assemble_block` pair renders. All block embeddings,
+sentinels included, come from the vision side, so the text embedding row of the
+placeholder id is never read -- the span is overwritten before the first layer.
+"""
+
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+
+import torch
+import torch.nn.functional as F
+
+from freetoken.layers import BaseOP, OPList
+from freetoken.models.weight_stream import BlockWeightStreamer
+
+from .config import DSV4VisionConfig
+
+# Sentinel token types inside an image block (reference inference/image_processor.py).
+IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5)
+# An image block is padded so its first row token lands on a compressor stride boundary.
+COMPRESS_PAD_TO = 4
+
+
+@lru_cache(8)
+def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
+    wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
+    freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq
+    freqs = freqs.flatten(1)
+    return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
+
+
+def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    dtype = x.dtype
+    x1, x2 = x.float().chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
+
+
+class _Linear(BaseOP):
+    """Linear in checkpoint-native layout: ``weight`` [out, in] plus an optional ``bias``."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        self.weight = torch.empty(out_features, in_features, dtype=torch.bfloat16)
+        self.bias = torch.empty(out_features, dtype=torch.bfloat16) if bias else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weight, self.bias)
+
+
+class RMSNorm(BaseOP):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        self.eps = eps
+        self.weight = torch.empty(dim, dtype=torch.bfloat16)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class PatchEmbed(BaseOP):
+    def __init__(self, patch_size: int, dim: int):
+        self.proj = _Linear(3 * patch_size**2, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj.forward(x.flatten(1))
+
+
+class VisionAttention(BaseOP):
+    def __init__(self, dim: int, n_heads: int):
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.wqkv = _Linear(dim, 3 * dim)
+        self.wo = _Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        n = x.size(0)
+        q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv.forward(x).chunk(3, dim=-1))
+        q = apply_rotary(q, cos, sin)
+        k = apply_rotary(k, cos, sin)
+        o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1))
+        return self.wo.forward(o.transpose(0, 1).reshape(n, -1))
+
+
+class VisionMLP(BaseOP):
+    def __init__(self, dim: int, inter_dim: int):
+        self.w1 = _Linear(dim, 2 * inter_dim, bias=False)
+        self.w2 = _Linear(inter_dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.w1.forward(x).chunk(2, dim=-1)
+        return self.w2.forward(F.silu(gate) * up)
+
+
+class VisionBlock(BaseOP):
+    def __init__(self, dim: int, n_heads: int, inter_dim: int, eps: float):
+        self.norm1 = RMSNorm(dim, eps)
+        self.attn = VisionAttention(dim, n_heads)
+        self.norm2 = RMSNorm(dim, eps)
+        self.mlp = VisionMLP(dim, inter_dim)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn.forward(self.norm1.forward(x), cos, sin)
+        return x + self.mlp.forward(self.norm2.forward(x))
+
+
+class DSV4VisionTower(BaseOP):
+    """Native-resolution ViT: bidirectional attention inside one image, 2-D RoPE over its patch grid."""
+
+    def __init__(self, vc: DSV4VisionConfig):
+        self.rope_dim = vc.vision_dim // vc.vision_n_heads // 2
+        self.rope_theta = vc.vision_rope_theta
+        self.patch_embed = PatchEmbed(vc.vision_patch_size, vc.vision_dim)
+        self.blocks = OPList(
+            [
+                VisionBlock(vc.vision_dim, vc.vision_n_heads, vc.vision_inter_dim, eps=1e-6)
+                for _ in range(vc.vision_n_layers)
+            ]
+        )
+        self.norm = RMSNorm(vc.vision_dim, 1e-6)
+        self._streamer: BlockWeightStreamer | None = None
+
+    def place_weights(self, mode: str) -> None:
+        """gpu: every tensor resident; host: the block stack in pinned banks, two blocks streamed at a time (patch embed and final norm stay resident)."""
+        if mode == "host" and self._streamer is None:
+            self._streamer = BlockWeightStreamer(self.blocks.op_list, self.patch_embed.proj.weight.device)
+        elif mode == "gpu" and self._streamer is not None:
+            self._streamer.unstream()
+            self._streamer = None
+        elif mode not in ("gpu", "host"):
+            raise ValueError(f"unknown vision weight placement {mode!r}")
+
+    def _blocks(self):
+        if self._streamer is None:
+            return enumerate(self.blocks.op_list)
+        return self._streamer.blocks(self.blocks.op_list)
+
+    def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        patches = patches.to(self.patch_embed.proj.weight.dtype)
+        x = self.patch_embed.forward(patches)
+        cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
+        cos, sin = cos.to(x.device), sin.to(x.device)
+        for _, block in self._blocks():
+            x = block.forward(x, cos, sin)
+        return self.norm.forward(x)
+
+
+class DSV4Aligner(BaseOP):
+    """3x3/stride-3 merge of every 3x3 patch grid into one token, then a 2-layer erf-GELU MLP."""
+
+    def __init__(self, vc: DSV4VisionConfig):
+        self.downsample_ratio = vc.vision_downsample_ratio
+        in_dim = vc.vision_dim * self.downsample_ratio**2
+        self.w1 = _Linear(in_dim, vc.text_dim)
+        self.w2 = _Linear(vc.text_dim, vc.text_dim)
+
+    def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        dtype = self.w1.weight.dtype
+        x = x.to(dtype)
+        r = self.downsample_ratio
+        x = x.view(n_h, n_w, -1).permute(2, 0, 1)
+        x = F.pad(x, (0, -n_w % r, 0, -n_h % r))
+        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
+        return self.w2.forward(F.gelu(self.w1.forward(x)))
+
+
+class DSV4Vision(BaseOP):
+    """The wrapper's whole vision stack under one ``visual.`` mount, like the V4.1 sibling."""
+
+    def __init__(self, vc: DSV4VisionConfig, text_dim: int):
+        self.vision = DSV4VisionTower(vc)
+        self.aligner = DSV4Aligner(vc)
+        for name in ("image_start", "image_end", "image_newline", "image_pad"):
+            setattr(self, name, torch.empty(text_dim, dtype=torch.bfloat16))
+
+    def place_weights(self, mode: str) -> None:
+        self.vision.place_weights(mode)
+
+    @torch.inference_mode()
+    def forward(self, item):
+        """One image -> the embeddings of its whole token block, sentinels included.
+
+        The block's layout is a function of where it lands in the prompt, so the item
+        carries the grid and the offset; ``build_image_block`` re-derives the same
+        ``(types, perm)`` the processor tokenized with.
+        """
+        device = self.vision.patch_embed.proj.weight.device
+        patches = item.feature.to(device, non_blocking=True)
+        aligned = self.aligner.forward(
+            self.vision.forward(patches, item.n_vit_h, item.n_vit_w), item.n_vit_h, item.n_vit_w
+        )
+        types, perm = build_image_block(item.n_llm_h, item.n_llm_w, item.start)
+        sentinels = torch.stack(
+            [self.image_start, self.image_pad, self.image_pad, self.image_newline, self.image_end]
+        )
+        return assemble_block(aligned, types, perm, sentinels)
+
+
+def grid_tokens(best_height: int, best_width: int, patch_size: int, downsample_ratio: int):
+    """LLM-token footprint of a resized image: the N-layout row/align padding included."""
+    n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio)
+    n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio)
+    num_tokens = n_llm_h * (n_llm_w + 1) + 2
+    if n_llm_h % 2 == 1:
+        num_tokens += n_llm_w + 1
+    num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
+    return n_llm_h, n_llm_w, num_tokens
+
+
+def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
+    r = height / width
+    max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
+    max_h_float = max_w_float * r
+    if max_w_float < 1.0:
+        max_w = 1
+        max_h = (max_n_token - 2) // (max_w + 1)
+        if max_h % 2 == 1:
+            max_h -= 1
+        best_width = max_w * patch_size * downsample_ratio
+        best_height = max_h * patch_size * downsample_ratio
+    elif max_h_float < 2.0:
+        max_h = 2
+        max_w = ((max_n_token - 2) // max_h) - 1
+        assert max_w > 1
+        best_width = max_w * patch_size * downsample_ratio
+        best_height = max_h * patch_size * downsample_ratio
+    else:
+        max_w = math.floor(max_w_float)
+        max_h = math.floor(max_h_float)
+        if max_h % 2 == 1:
+            max_h -= 1
+        beta = min(max_w * patch_size * downsample_ratio / width, max_h * patch_size * downsample_ratio / height)
+        best_width = math.floor(width * beta / patch_size) * patch_size
+        best_height = math.floor(height * beta / patch_size) * patch_size
+    n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
+    return n_llm_h, n_llm_w, best_height, best_width, num_tokens
+
+
+def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):
+    """Largest grid whose block still fits ``max_n_token`` once the alignment pads are added."""
+    max_n_token -= COMPRESS_PAD_TO - 1
+    n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
+    budget = max_n_token
+    while num_tokens > max_n_token:
+        n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio(
+            height, width, patch_size, downsample_ratio, budget
+        )
+        budget -= 1
+    return n_llm_h, n_llm_w, best_height, best_width
+
+
+def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int):
+    """The block's token types in their final order, plus the aligner-row order of its IMAGE slots.
+
+    ``types`` [n_out] places every token of the block (lead pads, START, the row-interleaved
+    IMAGE tokens, trailing pads, END). The reference streams them as ``vocab_size + types``;
+    here the span carries the image's content pad id and the assembled embeddings are
+    scattered over it. ``perm`` [G] maps aligner outputs (row-major grid) into the IMAGE slots.
+
+    ``start_pos`` is the block's absolute position in the prompt: the lead-pad count aligns
+    the first row token to a compressor stride boundary, so the layout -- and the block's
+    length -- depend on where the image lands.
+    """
+    compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO
+    pad_h = n_llm_h % 2
+    rows = n_llm_h + pad_h
+    row_len = n_llm_w + 1
+    pad_last = rows // 2 * row_len % 2 * 2
+    types = torch.tensor(
+        ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h),
+        dtype=torch.int64,
+    )
+    order = torch.arange(rows * row_len).view(rows // 2, 2, row_len).transpose(1, 2).reshape(-1)
+    image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64)
+    image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(n_llm_h, n_llm_w)
+    perm = image_idx[order]
+    perm = perm[perm >= 0]
+    types = torch.cat(
+        [
+            torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64),
+            torch.tensor([IMAGE_START]),
+            types[order],
+            torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64),
+            torch.tensor([IMAGE_END]),
+        ]
+    )
+    return types, perm
+
+
+def assemble_block(
+    aligned: torch.Tensor,
+    types: torch.Tensor,
+    perm: torch.Tensor,
+    sentinels: torch.Tensor,
+) -> torch.Tensor:
+    """Block embeddings [n_out, dim] from the aligner rows and the learned sentinels.
+
+    ``sentinels`` [5, dim] is ``(start, pad, pad, newline, end)``, indexable by the type
+    enum exactly like the reference ``merge_image_embeddings``.
+    """
+    block = sentinels[types.to(sentinels.device)]
+    block[types == IMAGE] = aligned[perm.to(aligned.device)]
+    return block
+
+
+__all__ = [
+    "COMPRESS_PAD_TO",
+    "DSV4Aligner",
+    "DSV4Vision",
+    "DSV4VisionTower",
+    "IMAGE",
+    "IMAGE_END",
+    "IMAGE_NEW_LINE",
+    "IMAGE_PAD",
+    "IMAGE_START",
+    "assemble_block",
+    "build_image_block",
+    "grid_tokens",
+    "safe_resize",
+    "solve_resize_ratio",
+]
diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py
index 4728f7f9c3..fbcddb1c06 100644
--- a/python/freetoken/models/deepseek_v4/weight.py
+++ b/python/freetoken/models/deepseek_v4/weight.py
@@ -78,12 +78,16 @@ def iter_weights(
     *,
     include_moe_experts: bool = True,
     include_non_moe: bool = True,
+    include_vision: bool = True,
 ):
     """Stream resident (non-expert) weights as ``(name, tensor)`` keyed to engine params.
 
     Routed MXFP4 experts come from the offload cache, so ``include_moe_experts`` must be
     False (DeepSeek-V4 only runs ``--moe-strategy offload``). Tensors yielded in checkpoint
     dtype (fp8 + e8m0 preserved); ``wo_a`` dequantized to bf16 to match the reference einsum.
+    ``include_vision`` (the engine's active-encoder set) drops the ``visual.*`` tower
+    tensors; the router biases stay in either way, since the gate builds
+    them on any vision checkpoint.
     """
     if include_moe_experts:
         raise ValueError(
@@ -112,6 +116,9 @@ def linear(src: str, dst: str):
         for nm in ("hc_head_fn", "hc_head_base", "hc_head_scale"):
             yield f"model.{nm}", get(nm)
 
+        if include_vision:
+            yield from _iter_vision(reader, args)
+
         for L in range(args.n_layers):
             a = f"layers.{L}.attn"
             m = f"model.{a}"
@@ -147,8 +154,13 @@ def linear(src: str, dst: str):
             yield f"model.{g}.weight", get(f"{g}.weight")
             if L < args.n_hash_layers:
                 yield f"model.{g}.tid2eid", get(f"{g}.tid2eid")
+                if args.vision_n_layers > 0:
+                    # a vision checkpoint gives its hash layers a text bias too
+                    yield f"model.{g}.bias", get(f"{g}.bias")
             else:
                 yield f"model.{g}.bias", get(f"{g}.bias")
+            if args.vision_n_layers > 0:
+                yield f"model.{g}.bias_vl", get(f"{g}.bias_vl")
             for proj in ("w1", "w2", "w3"):
                 src = f"layers.{L}.ffn.shared_experts.{proj}"
                 yield from linear(src, f"model.{src}")
@@ -162,6 +174,58 @@ def linear(src: str, dst: str):
         reader.close()
 
 
+# --------------------------------------------------------------------------------------
+# Vision tower.
+# --------------------------------------------------------------------------------------
+# The wrapper mounts the tower, the aligner and the sentinels under one ``visual``
+# container, so every engine-side name carries the ``visual.`` prefix the shared
+# VISION_KEY_PREFIXES filter matches; the checkpoint's unprefixed names
+# (``vision.*``, ``aligner.*``, ``image_*``) are renamed here on load.
+_VISION_BLOCK_SUFFIXES = (
+    "norm1.weight",
+    "attn.wqkv.weight",
+    "attn.wqkv.bias",
+    "attn.wo.weight",
+    "attn.wo.bias",
+    "norm2.weight",
+    "mlp.w1.weight",
+    "mlp.w2.weight",
+)
+_VISION_TAIL = (
+    "vision.patch_embed.proj.weight",
+    "vision.patch_embed.proj.bias",
+    "vision.norm.weight",
+    "aligner.w1.weight",
+    "aligner.w1.bias",
+    "aligner.w2.weight",
+    "aligner.w2.bias",
+    "image_start",
+    "image_end",
+    "image_newline",
+    "image_pad",
+)
+
+
+def _iter_vision(reader, args: DeepseekV4Args):
+    if args.vision_n_layers <= 0:
+        return
+    for L in range(args.vision_n_layers):
+        for suffix in _VISION_BLOCK_SUFFIXES:
+            yield f"visual.vision.blocks.{L}.{suffix}", reader.get(f"vision.blocks.{L}.{suffix}")
+    for name in _VISION_TAIL:
+        yield f"visual.{name}", reader.get(name)
+
+
+def iter_vision_weights(model_path: str, device):
+    """The vision tower alone, for the engine's encoder-only load and the FTW converter."""
+    args = load_args(model_path, max_batch_size=1)
+    reader = _ShardReader(model_path, _weight_map(model_path), device)
+    try:
+        yield from _iter_vision(reader, args)
+    finally:
+        reader.close()
+
+
 # --------------------------------------------------------------------------------------
 # Routed MXFP4 expert pieces.
 # --------------------------------------------------------------------------------------
@@ -212,4 +276,4 @@ def _serial():
     return per_expert_pieces(_serial(), locate, tensors_per_expert=6)
 
 
-__all__ = ["iter_weights", "iter_expert_pieces"]
+__all__ = ["iter_weights", "iter_expert_pieces", "iter_vision_weights"]
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index fe150ddc29..678424ff05 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -84,6 +84,10 @@ class ModelSpec:
 ) + _EXPERTS_W123_PACKED
 _MINIMAX_M3_PROCESSOR = "freetoken.mm.processors.minimax_m3:MiniMaxM3MMProcessor"
 _MINIMAX_M3_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
+# DeepSeek-V4's vision dims sit at the top level of the checkpoint config, next to the
+# language model's rather than under a section of their own.
+_DSV4_PROCESSOR = "freetoken.mm.processors.deepseek_v4:DSV4MMProcessor"
+_DSV4_ENCODERS = (EncoderSpec("vision", "vision_n_layers", ("image",)),)
 
 _MODEL_REGISTRY: dict[str, ModelSpec] = {
     "LlamaForCausalLM": ModelSpec(
@@ -153,8 +157,13 @@ class ModelSpec:
         # the checkpoint has no ``model.`` root
         checkpoint_roots=(("model.layers", "layers"), ("model.head", "head")),
         packed_modules_mapping=_EXPERTS_W123_PACKED,
-        # the head, the KV compressors and the indexer's scorer ship bf16; the fp8 config has no modules_to_not_convert
-        unquantized_modules=("head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj"),
+        # the head, the KV compressors, the indexer's scorer and the vision tower + aligner ship bf16;
+        # the fp8 config has no modules_to_not_convert (patterns match checkpoint names, like V4.1)
+        unquantized_modules=(
+            "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj", "vision.*", "aligner.*",
+        ),
+        mm_processor=_DSV4_PROCESSOR,
+        encoders=_DSV4_ENCODERS,
     ),
     # DeepSeek-V4.1-Flash (model_type deepseek_v41): multimodal wrapper config (text tower in
     # text_config). The checkpoint has no ``model.`` root; the compressors, the indexer's key projection /
diff --git a/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py
index 9640c7b7f7..12574651e2 100644
--- a/scripts/ftw_hotfix.py
+++ b/scripts/ftw_hotfix.py
@@ -264,7 +264,7 @@ def is_checkpoint_tower_name(name: str) -> bool:
     if head[0] == "model" and len(head) > 1:
         head = head[1:]
     return ("vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS
-            or head[0] in ("aligner", "image_start", "image_end", "image_newline")
+            or head[0] in ("aligner", "image_start", "image_end", "image_newline", "image_pad")
             or name.endswith(".ffn.gate.bias_vl"))
 
 
diff --git a/tests/README.md b/tests/README.md
index 76b2b1fb03..adc87035d9 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -63,6 +63,8 @@ checkpoint is set:
 | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) |
 | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file |
 | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) |
+| `FREETOKEN_DSV4_VL_MODEL` | `tokenizer/test_mm_tokenize.py` - local DeepSeek-V4 vision checkpoint |
+| `FREETOKEN_DSV4_VISION_CKPT` | `models/test_dsv4_vision.py` - the same checkpoint, for the parity tests against its `inference/` reference code |
 
 `test_aime.py` takes its sampling protocol from the checkpoint's own
 `generation_config.json` (pass@N at the recommended temperature, or a single greedy
diff --git a/tests/checkpoint/test_ftw_hotfix_vision.py b/tests/checkpoint/test_ftw_hotfix_vision.py
index eb0c592f35..ef9225a926 100644
--- a/tests/checkpoint/test_ftw_hotfix_vision.py
+++ b/tests/checkpoint/test_ftw_hotfix_vision.py
@@ -37,6 +37,7 @@ def hotfix():
     ("vision.blocks.0.norm1.weight", True),
     ("aligner.w1.weight", True),
     ("image_newline", True),
+    ("image_pad", True),
     ("layers.3.ffn.gate.bias_vl", True),
     ("model.embed_audio.embedding_projection.weight", False),
     ("model.language_model.layers.0.self_attn.q_proj.weight", False),
diff --git a/tests/mm/test_processor.py b/tests/mm/test_processor.py
index a9a50e3c2b..9562b366f4 100644
--- a/tests/mm/test_processor.py
+++ b/tests/mm/test_processor.py
@@ -9,7 +9,7 @@
 import torch
 
 from freetoken.message import MMItem
-from freetoken.mm import MM_PAD_SHIFT_VALUE
+from freetoken.mm import MM_PAD_SHIFT_VALUE, mm_pad_value
 from freetoken.mm.config import MultimodalConfig
 from freetoken.mm.processor import MMProcessor, PromptReplacement, image_positions
 from freetoken.mm.processors.gemma4 import Gemma4MMProcessor, Gemma4UnifiedMMProcessor
@@ -17,6 +17,7 @@
 from freetoken.mm.processors.muse_glimmer import MuseGlimmerMMProcessor
 from freetoken.mm.processors.minimax_m3 import IMAGE_END_ID, IMAGE_START_ID, MiniMaxM3MMProcessor
 from freetoken.mm.processors.qwen_vl import QwenVLMMProcessor
+from freetoken.mm.processors.deepseek_v4 import DSV4MMProcessor
 
 PLACEHOLDER = 7
 BOI, SOFT, EOI = 90, 91, 92
@@ -415,3 +416,103 @@ def test_registry_resolves_by_architecture(monkeypatch):
     assert get_mm_processor("/text-only") is None
     assert get_mm_processor("/unknown-vlm") is None
     assert get_mm_processor("/missing") is None  # a config that fails to load means no vision
+
+
+def _dsv4_config(**over):
+    """The vision dims sit at the top level of a deepseek_v4 checkpoint config."""
+    base = dict(
+        vision_n_layers=32,
+        vision_dim=1024,
+        vision_n_heads=16,
+        vision_inter_dim=2816,
+        vision_patch_size=14,
+        vision_rope_theta=10000.0,
+        vision_downsample_ratio=3,
+        vision_max_n_token=384,
+        vision_min_pixels=147456,
+        vision_max_wh_ratio=8,
+        hidden_size=4096,
+    )
+    return SimpleNamespace(**{**base, **over})
+
+
+def _dsv4_image(width, height):
+    from PIL import Image
+
+    return Image.new("RGB", (width, height))
+
+
+def test_dsv4_block_length_and_hash_follow_the_insertion_offset():
+    from freetoken.models.deepseek_v4.config import IMAGE_TOKEN_ID
+    from freetoken.models.deepseek_v4.vision import build_image_block
+
+    proc = DSV4MMProcessor(_dsv4_config(), "/nonexistent", MultimodalConfig())
+    assert proc.placeholder == [IMAGE_TOKEN_ID]
+    (item,) = proc.process([_dsv4_image(640, 480)])
+    assert (item.n_vit_h, item.n_vit_w, item.n_llm_h, item.n_llm_w) == (35, 46, 12, 16)
+
+    repl = proc.prompt_replacement(item, 0)
+    at_zero, hash_zero = len(repl.full), item.hash
+    assert at_zero == len(build_image_block(12, 16, 0)[0]) == 209
+    assert repl.full == [IMAGE_TOKEN_ID] * at_zero
+    # the block's sentinels are embeddings too, so the whole span takes image rows
+    assert repl.embed_spans() == [[0, at_zero]]
+    assert item.pad_value == mm_pad_value(item.hash) and item.pad_value >= MM_PAD_SHIFT_VALUE
+
+    # the lead pads align the block to a compression stride, so a later offset is a shorter
+    # block AND a different content hash: one alignment's rows must never serve another's
+    repl = proc.prompt_replacement(item, 2)
+    assert len(repl.full) == at_zero - 2 and item.hash != hash_zero
+    assert item.pad_value == mm_pad_value(item.hash) >= MM_PAD_SHIFT_VALUE
+
+
+def test_dsv4_dummy_item_is_the_smallest_encodable_block():
+    from freetoken.models.deepseek_v4.vision import build_image_block
+
+    proc = DSV4MMProcessor(_dsv4_config(), "/nonexistent", MultimodalConfig())
+    (dummy,) = proc.dummy_items(torch.bfloat16, torch.device("cpu"))
+    dummy.validate()
+    assert dummy.feature.shape == (36, 3, 14, 14)  # a 6x6 patch grid, one 2x2 merge grid
+    assert dummy.num_tokens == len(build_image_block(2, 2, 0)[0]) == 13
+
+
+def test_a_text_only_checkpoint_config_has_no_dsv4_processor():
+    proc = DSV4MMProcessor
+    with pytest.raises(ValueError):
+        proc(_dsv4_config(vision_n_layers=0), "/nonexistent", MultimodalConfig())
+
+
+def test_dsv4_apply_passes_each_block_the_offset_it_lands_at():
+    """Only the DSV4 block depends on its position, so its own apply computes the offset."""
+    from freetoken.models.deepseek_v4.config import IMAGE_TOKEN_ID
+    from freetoken.models.deepseek_v4.vision import build_image_block
+
+    proc = DSV4MMProcessor(_dsv4_config(), "/nonexistent", MultimodalConfig())
+    buf = io.BytesIO()
+    _dsv4_image(640, 480).save(buf, format="PNG")
+    raw = buf.getvalue()
+    ids = torch.tensor([1, IMAGE_TOKEN_ID, 2, IMAGE_TOKEN_ID, 3], dtype=torch.int32)
+    out = proc.apply(ids, [raw, raw])
+    first, second = out.mm_items
+    assert (first.n_llm_h, first.n_llm_w) == (12, 16)
+    # the first placeholder sits at 1, the second after the text between plus the first block
+    assert first.start == first.offsets[0][0] == 1
+    assert second.start == second.offsets[0][0]
+    assert first.num_tokens == len(build_image_block(12, 16, first.start)[0])
+    assert second.num_tokens == len(build_image_block(12, 16, second.start)[0])
+    # apply fills the embedding slots with each block's content pad id, not the placeholder
+    for item in (first, second):
+        lo, hi = item.offsets[0]
+        assert set(out.input_ids[lo:hi].tolist()) == {item.pad_value}
+    assert IMAGE_TOKEN_ID not in out.input_ids.tolist()
+
+
+def test_dsv4_image_max_tokens_lowers_the_checkpoints_block_cap():
+    from freetoken.models.deepseek_v4.vision import build_image_block
+
+    big = DSV4MMProcessor(_dsv4_config(), "/nonexistent", MultimodalConfig())
+    capped = DSV4MMProcessor(_dsv4_config(), "/nonexistent", MultimodalConfig(image_max_tokens=64))
+    (a,) = big.process([_dsv4_image(1600, 1200)])
+    (b,) = capped.process([_dsv4_image(1600, 1200)])
+    assert b.n_llm_h * b.n_llm_w < a.n_llm_h * a.n_llm_w
+    assert len(build_image_block(b.n_llm_h, b.n_llm_w, 0)[0]) <= 64
diff --git a/tests/models/test_dsv4_vision.py b/tests/models/test_dsv4_vision.py
new file mode 100644
index 0000000000..68a3f0ab19
--- /dev/null
+++ b/tests/models/test_dsv4_vision.py
@@ -0,0 +1,335 @@
+"""DeepSeek-V4 image preprocessing, block layout and the tiny vision tower (no checkpoint).
+
+The block layout is what makes an image's token count depend on where it lands in the
+prompt, so it is pinned here against hand-written expectations; the checkpoint-gated test
+at the bottom checks the same functions against the reference ``inference/`` code.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from freetoken.models.deepseek_v4.config import DSV4VisionConfig, parse_config, parse_vision_config
+from freetoken.models.deepseek_v4.image_processor import load_image
+from freetoken.models.deepseek_v4.vision import (
+    COMPRESS_PAD_TO,
+    DSV4Aligner,
+    DSV4VisionTower,
+    IMAGE,
+    IMAGE_END,
+    IMAGE_NEW_LINE,
+    IMAGE_PAD,
+    IMAGE_START,
+    assemble_block,
+    build_image_block,
+)
+
+checkpoint_path = os.environ.get("FREETOKEN_DSV4_VISION_CKPT")
+needs_checkpoint = pytest.mark.skipif(
+    not checkpoint_path, reason="needs FREETOKEN_DSV4_VISION_CKPT pointing at the vision checkpoint"
+)
+
+_HF_CONFIG = {
+    "vision_n_layers": 32,
+    "vision_dim": 1024,
+    "vision_n_heads": 16,
+    "vision_inter_dim": 2816,
+    "vision_patch_size": 14,
+    "vision_rope_theta": 10000.0,
+    "vision_downsample_ratio": 3,
+    "vision_max_n_token": 384,
+    "vision_min_pixels": 147456,
+    "vision_max_wh_ratio": 8,
+    "hidden_size": 4096,
+}
+
+
+def _hf_config(**over):
+    from types import SimpleNamespace
+
+    return SimpleNamespace(**{**_HF_CONFIG, **over})
+
+
+def _vc(**over) -> DSV4VisionConfig:
+    """A tower small enough to run: 2 patch-2 blocks over a 2x2 merge grid."""
+    base = dict(
+        vision_n_layers=2,
+        vision_dim=16,
+        vision_n_heads=2,
+        vision_inter_dim=32,
+        vision_patch_size=4,
+        vision_rope_theta=10000.0,
+        vision_downsample_ratio=2,
+        vision_max_n_token=384,
+        vision_min_pixels=147456,
+        vision_max_wh_ratio=8,
+        text_dim=64,
+    )
+    return DSV4VisionConfig(**{**base, **over})
+
+
+def _png(width, height, seed=0):
+    rng = np.random.default_rng(seed)
+    array = (rng.random((height, width, 3)) * 255).astype(np.uint8)
+    buf = io.BytesIO()
+    Image.fromarray(array, "RGB").save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_a_vision_checkpoint_config_becomes_the_tower_dims_and_a_text_one_does_not():
+    vc = parse_vision_config(_hf_config())
+    assert (vc.vision_n_layers, vc.vision_dim, vc.text_dim) == (32, 1024, 4096)
+    # the engine nulls the section on the copy it hands the parser when it serves text-only
+    assert parse_vision_config(_hf_config(vision_n_layers=None)) is None
+    assert parse_vision_config(_hf_config(vision_n_layers=0)) is None
+
+
+def test_block_layout_interleaves_rows_and_pads_to_the_compression_stride():
+    types, perm = build_image_block(2, 2, 0)
+    # 3 lead pads put the first row token on a COMPRESS_PAD_TO (4) boundary, then START,
+    # the N-layout rows (column pairs: 2 rows of 2 IMAGE tokens plus their row markers),
+    # the trailing parity pad, END.
+    assert types.tolist() == (
+        [IMAGE_PAD] * 3
+        + [IMAGE_START]
+        + [IMAGE, IMAGE, IMAGE, IMAGE, IMAGE_NEW_LINE, IMAGE_NEW_LINE]
+        + [IMAGE_PAD] * 2
+        + [IMAGE_END]
+    )
+    # row-major aligner outputs 0..3 land in the IMAGE slots in N-layout order
+    assert perm.tolist() == [0, 2, 1, 3]
+
+
+def test_the_block_shrinks_by_one_token_per_lead_pad_the_offset_absorbs():
+    # the lead pads are what aligns the block, so the same image at a later offset is a
+    # SHORTER span: the count cannot be a property of the image alone
+    lengths = [len(build_image_block(2, 2, start)[0]) for start in range(4)]
+    assert lengths == [13, 12, 11, 10]
+    # an odd row count adds a marker row, which is not pad bytes
+    assert [len(build_image_block(3, 2, start)[0]) for start in range(4)] == [17, 16, 15, 14]
+
+
+def test_assemble_block_puts_the_sentinels_and_the_aligner_rows_where_the_types_say():
+    types, perm = build_image_block(2, 2, 0)
+    sentinels = torch.stack([torch.full((4,), float(i)) for i in range(5)])
+    aligned = torch.arange(4 * 4, dtype=torch.float32).view(4, 4) * 10
+    block = assemble_block(aligned, types, perm, sentinels)
+    assert block.shape == (len(types), 4)
+    # sentinels index the type enum, so START/PAD/END rows carry 0/1/4 and the IMAGE rows
+    # carry the aligner output perm says they take
+    assert block[0].tolist() == [1.0] * 4  # lead pad
+    assert block[3].tolist() == [0.0] * 4  # IMAGE_START
+    assert block[-1].tolist() == [4.0] * 4  # IMAGE_END
+    image_rows = block[types == IMAGE]
+    assert image_rows.tolist() == aligned[perm].tolist()
+
+
+def test_preprocessing_resizes_to_the_merge_grid_and_normalizes_like_the_reference():
+    # min_pixels 0 keeps the input size: the reference upscales below it
+    vc = _vc(vision_min_pixels=0)
+    patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(
+        Image.open(io.BytesIO(_png(80, 64))), vc
+    )
+    assert (n_vit_h, n_vit_w, n_llm_h, n_llm_w) == (16, 20, 8, 10)
+    assert patches.shape == (n_vit_h * n_vit_w, 3, 4, 4)
+    assert patches.dtype == torch.bfloat16
+    assert float(patches.min()) >= -1.0 and float(patches.max()) <= 1.0
+    # identical bytes, identical patches: the content hash and the radix key ride on these
+    again, *_ = load_image(Image.open(io.BytesIO(_png(80, 64))), vc)
+    assert torch.equal(patches, again)
+
+
+def test_the_token_budget_caps_a_tall_image():
+    vc = _vc(vision_max_n_token=64, vision_min_pixels=0)
+    _patches, _n_vit_h, _n_vit_w, n_llm_h, n_llm_w = load_image(
+        Image.open(io.BytesIO(_png(64, 1024))), vc
+    )
+    types, _perm = build_image_block(n_llm_h, n_llm_w, 0)
+    assert len(types) <= vc.vision_max_n_token
+
+
+def test_tiny_tower_runs_and_its_aligner_merges_every_merge_grid():
+    vc = _vc()
+    tower = DSV4VisionTower(vc)
+    aligner = DSV4Aligner(vc)
+    patches = torch.zeros(20 * 24, 3, 4, 4, dtype=torch.bfloat16)
+    out = tower.forward(patches, 20, 24)
+    assert out.shape == (20 * 24, vc.vision_dim)
+    assert aligner.forward(out, 20, 24).shape == (10 * 12, vc.text_dim)
+
+
+def _reference_modules(path):
+    """Import the checkpoint's own ``inference/image_processor.py`` (PIL + torch only)."""
+    import importlib.util
+
+    module_path = os.path.join(path, "inference", "image_processor.py")
+    if not os.path.exists(module_path):
+        pytest.skip(f"{module_path} not found")
+    spec = importlib.util.spec_from_file_location("_dsv4_reference_image_processor", module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@needs_checkpoint
+def test_preprocessing_and_block_layout_match_the_reference():
+    from types import SimpleNamespace
+
+    reference = _reference_modules(checkpoint_path)
+    vc = parse_vision_config(_hf_config())
+    args = SimpleNamespace(
+        vision_patch_size=vc.vision_patch_size,
+        vision_downsample_ratio=vc.vision_downsample_ratio,
+        vision_max_n_token=vc.vision_max_n_token,
+        vision_min_pixels=vc.vision_min_pixels,
+        vision_max_wh_ratio=vc.vision_max_wh_ratio,
+    )
+    for width, height in ((640, 480), (100, 100), (3000, 400), (1600, 1200)):
+        raw = _png(width, height)
+        mine = load_image(Image.open(io.BytesIO(raw)), vc)
+        theirs = reference.load_image({"data": raw}, args)
+        assert torch.equal(mine[0], theirs[0])
+        assert mine[1:] == theirs[1:]
+        for start in (0, 1, 3, 7):
+            types, perm = build_image_block(mine[3], mine[4], start)
+            ref_types, ref_perm = reference.build_image_block(theirs[3], theirs[4], start)
+            assert torch.equal(types, ref_types) and torch.equal(perm, ref_perm)
+
+
+_WINDOW, _MAX_IMG = 128, 384
+
+
+def _ref_get_image_visible(input_ids: torch.Tensor, vocab_size: int, max_image_tokens: int):
+    """Verbatim port of the checkpoint's ``inference/model.py`` pair (not vendored in this
+    repo): ``get_image_visible`` derives spans from the virtual ids ``vocab + type``, so the
+    lead pads before IMAGE_START stay outside; only the START/END comparisons matter, our
+    type enum values index the same way."""
+    seqlen = input_ids.size(1)
+    idx = torch.arange(seqlen, dtype=torch.int32).unsqueeze(0)
+    is_start = input_ids == vocab_size + IMAGE_START
+    is_end = input_ids == vocab_size + IMAGE_END
+    valid = (is_start.cumsum(1) > is_end.cumsum(1)) | is_end
+    starts = torch.where(is_start, idx, 0).cummax(1)[0]
+    left = (idx - starts) * valid
+    ends = torch.where(is_end, idx, seqlen).flip(1).cummin(1)[0].flip(1)
+    right = (ends - idx) * valid
+    return left.clamp(max=max_image_tokens - 1), right.clamp(max=max_image_tokens)
+
+
+def _ref_get_window_topk_idxs_visible(window_size, seqlen, left, right, max_image_tokens):
+    width = min(seqlen, window_size + max_image_tokens)
+    idx = torch.arange(seqlen).unsqueeze(0)
+    left_add = (left - (window_size - 1)).clamp(min=0)
+    starts = (idx - (window_size - 1) - left_add).clamp(min=0)
+    matrix = starts.unsqueeze(-1) + torch.arange(width)
+    return torch.where(matrix > (idx + right).unsqueeze(-1), -1, matrix).int().contiguous()
+
+
+def _reference_candidates(prefix: int, n_llm_h: int, n_llm_w: int, vocab: int = 100_000):
+    """The reference pair's window candidates for a text prefix followed by one image block."""
+    types, _perm = build_image_block(n_llm_h, n_llm_w, prefix)
+    ids = torch.full((1, prefix + len(types)), vocab + 50, dtype=torch.int64)
+    ids[0, prefix:] = vocab + types
+    n = ids.size(1)
+    left, right = _ref_get_image_visible(ids, vocab, _MAX_IMG)
+    block = (prefix, prefix + len(types))
+    return block, types, _ref_get_window_topk_idxs_visible(_WINDOW, n, left, right, _MAX_IMG)[0]
+
+
+def test_model_spans_match_the_reference_visibility_pair():
+    from types import SimpleNamespace
+
+    from freetoken.models.deepseek_v4.attention import Attention
+    from freetoken.models.deepseek_v4.model import _image_spans_by_ti
+
+    for prefix, h, w in ((0, 3, 2), (1, 3, 2), (2, 3, 2), (3, 3, 2), (5, 4, 3), (201, 20, 16), (402, 2, 2)):
+        (lo, hi), types, ref = _reference_candidates(prefix, h, w)
+        req = SimpleNamespace(table_idx=7, mm_items=[SimpleNamespace(offsets=[[lo, hi]])])
+        spans = _image_spans_by_ti([req])
+        assert spans == {7: [(lo + types.tolist().index(IMAGE_START), hi)]}
+        cols = Attention.visible_window_cols(
+            spans[7], 0, hi, _WINDOW, _MAX_IMG, torch.device("cpu"),
+            addr_lo=0, width=min(hi, _WINDOW + _MAX_IMG),
+        )
+        assert torch.equal(cols, ref)
+    text = SimpleNamespace(table_idx=0, mm_items=None)
+    assert _image_spans_by_ti([text]) is None
+
+
+def test_a_span_started_at_the_first_lead_pad_breaks_reference_parity():
+    from freetoken.models.deepseek_v4.attention import Attention
+
+    # the maintainer's case: a 201-token prefix puts 2 lead pads in front of a 20x16 grid,
+    # and a span covering the whole block diverges from the reference on 218 rows
+    (lo, hi), _types, ref = _reference_candidates(201, 20, 16)
+    assert lo + COMPRESS_PAD_TO - 1 - lo % COMPRESS_PAD_TO - lo == 2  # the two lead pads
+    cols = Attention.visible_window_cols(
+        [(lo, hi)], 0, hi, _WINDOW, _MAX_IMG, torch.device("cpu"),
+        addr_lo=0, width=min(hi, _WINDOW + _MAX_IMG),
+    )
+    assert int((cols != ref).any(dim=1).sum()) == 218
+
+
+def test_a_vision_checkpoint_keeps_image_blocks_whole_across_chunk_boundaries(tmp_path):
+    """Image spans attend bidirectionally, so a vision checkpoint sets
+    bidirectional_mm_blocks on its attention group and the scheduler ends a chunk before a
+    block instead of splitting it; text-only defaults to False and keeps plain chunking."""
+    from types import SimpleNamespace
+
+    from freetoken.core import Context, SamplingParams, get_global_ctx, set_global_ctx
+    from freetoken.message import MMItem
+    from freetoken.scheduler.cache import CacheManager
+    from freetoken.scheduler.decode import DecodeManager
+    from freetoken.scheduler.mm import cut_image_spans
+    from freetoken.scheduler.prefill import PrefillManager
+    from freetoken.scheduler.table import TableManager
+    from freetoken.scheduler.utils import PendingReq
+
+    # DeepseekV4Args defaults fill an empty inference/config.json; the hf namespace carries the tower
+    (tmp_path / "inference").mkdir()
+    (tmp_path / "inference" / "config.json").write_text("{}")
+    vision = parse_config(SimpleNamespace(_name_or_path=str(tmp_path), **_HF_CONFIG))
+    text = parse_config(SimpleNamespace(_name_or_path=str(tmp_path), **{**_HF_CONFIG, "vision_n_layers": None}))
+    assert vision.attention_groups[0].bidirectional_mm_blocks is True
+    assert text.attention_groups[0].bidirectional_mm_blocks is False
+
+    prefix = 10
+    types, _perm = build_image_block(2, 2, prefix)
+    lo, hi = prefix, prefix + len(types)
+    image = MMItem(modality="image", hash=1, pad_value=1, offsets=[[lo, hi]], feature=torch.zeros(1))
+    assert lo < 15 < hi  # a plain 15-token chunk budget ends inside the block
+
+    def chunk_lens(keep_images_whole):
+        try:
+            get_global_ctx()
+        except AssertionError:
+            set_global_ctx(Context(page_size=1))
+        pt = torch.zeros((5, 64), dtype=torch.int32)
+        cm = CacheManager(num_pages=64, page_size=1, page_table=pt, type="radix")
+        pm = PrefillManager(
+            cm, TableManager(max_running_reqs=4, page_table=pt), DecodeManager(1),
+            keep_images_whole=keep_images_whole,
+        )
+        pm.pending_list = [PendingReq(uid=7, input_ids=torch.arange(30, dtype=torch.int32),
+                                      sampling_params=SamplingParams(max_tokens=4), mm_items=[image])]
+        lens, cuts = [], []
+        while pm.runnable:
+            batch = pm.schedule_next_batch(15)
+            assert batch is not None
+            cm.allocate_paged(batch.reqs)
+            cuts.extend(cut_image_spans(batch.reqs))
+            for r in batch.reqs:
+                lens.append(r.extend_len)
+                r.complete_one()
+        return lens, cuts
+
+    lens, cuts = chunk_lens(vision.attention_groups[0].bidirectional_mm_blocks)
+    assert lens == [lo, 15, 30 - lo - 15] and cuts == []  # first chunk stops at the block, the next holds it whole
+    lens, cuts = chunk_lens(text.attention_groups[0].bidirectional_mm_blocks)
+    assert lens == [15, 15] and cuts == [(lo, hi)]  # the block is split into two spans; the cut warning would fire
diff --git a/tests/scheduler/test_mm.py b/tests/scheduler/test_mm.py
index 1344763471..72db7e3db8 100644
--- a/tests/scheduler/test_mm.py
+++ b/tests/scheduler/test_mm.py
@@ -6,7 +6,13 @@
 
 from freetoken.message import MMItem
 from freetoken.mm.encoder_cache import EncoderCache
-from freetoken.scheduler.mm import cut_image_spans, mm_chunk_end, mm_rows_after, plan_mm_batch, plan_mm_chunk
+from freetoken.scheduler.mm import (
+    cut_image_spans,
+    mm_chunk_end,
+    mm_rows_after,
+    plan_mm_batch,
+    plan_mm_chunk,
+)
 
 CPU = torch.device("cpu")
 
diff --git a/tests/tokenizer/test_mm_tokenize.py b/tests/tokenizer/test_mm_tokenize.py
index 8f2e021128..5adac8305c 100644
--- a/tests/tokenizer/test_mm_tokenize.py
+++ b/tests/tokenizer/test_mm_tokenize.py
@@ -1,4 +1,4 @@
-"""Image tokenization against real Qwen VL, Gemma-4, GLM-5.3, Muse-Glimmer and MiniMax-M3 checkpoints (skipped when absent)."""
+"""Image tokenization against real Qwen VL, Gemma-4, GLM-5.3, Muse-Glimmer, MiniMax-M3 and DeepSeek-V4-Vision checkpoints (skipped when absent)."""
 
 from __future__ import annotations
 
@@ -20,12 +20,13 @@
 GLM = os.environ.get("FREETOKEN_GLM53_MODEL", "")
 MUSE = os.environ.get("FREETOKEN_MUSE_MODEL", "")
 MINIMAX = os.environ.get("FREETOKEN_MINIMAX_M3_MODEL", "")
+DSV4 = os.environ.get("FREETOKEN_DSV4_VL_MODEL", "")
 
 pytestmark = [
     pytest.mark.needs_weights,
     pytest.mark.skipif(
-        not MODELS and not any(os.path.exists(os.path.join(p, "config.json")) for p in (GEMMA, GLM, MUSE, MINIMAX)),
-        reason="no FREETOKEN_QWEN36_MODEL / FREETOKEN_QWEN3VL_MODEL / FREETOKEN_GEMMA4_MODEL / FREETOKEN_GLM53_MODEL / FREETOKEN_MUSE_MODEL / FREETOKEN_MINIMAX_M3_MODEL checkpoint",
+        not MODELS and not any(os.path.exists(os.path.join(p, "config.json")) for p in (GEMMA, GLM, MUSE, MINIMAX, DSV4)),
+        reason="no FREETOKEN_QWEN36_MODEL / FREETOKEN_QWEN3VL_MODEL / FREETOKEN_GEMMA4_MODEL / FREETOKEN_GLM53_MODEL / FREETOKEN_MUSE_MODEL / FREETOKEN_MINIMAX_M3_MODEL / FREETOKEN_DSV4_VL_MODEL checkpoint",
     ),
 ]
 
@@ -214,3 +215,51 @@ def test_minimax_wraps_the_pad_span_in_start_end_tokens():
     assert bool((ids[start:end] == item.pad_value).all())
     budgeted = TokenizeManager(tokenizer, get_mm_processor(MINIMAX, MultimodalConfig(image_max_tokens=64))).tokenize([_msg([_png(640, 400)])])[0]
     assert budgeted.mm_items[0].num_tokens <= 64
+
+
+@pytest.mark.skipif(
+    not os.path.exists(os.path.join(DSV4, "config.json")),
+    reason="no FREETOKEN_DSV4_VL_MODEL checkpoint",
+)
+def test_dsv4_expands_the_placeholder_into_the_reference_block():
+    """The block carries one content pad id, and its length follows the reference's grid."""
+    from freetoken.mm.processor import get_mm_processor
+    from freetoken.tokenizer.tokenize import TokenizeManager
+    from freetoken.utils.hf import load_tokenizer
+
+    manager = TokenizeManager(load_tokenizer(DSV4), get_mm_processor(DSV4))
+    # 640x480 through the reference: a 12x16 merge grid, so 206 tokens of rows and markers
+    # plus the 3 lead pads that align the block at this offset
+    out = manager.tokenize([_msg([_png(640, 480)])])[0]
+    ids = out.input_ids.tolist()
+    (item,) = out.mm_items
+    assert (item.n_llm_h, item.n_llm_w) == (12, 16)
+    assert item.num_tokens == 209
+    lo, hi = item.offsets[0]
+    assert (lo, hi) == (item.start, item.start + 209)
+    assert set(ids[lo:hi]) == {item.pad_value}
+    assert item.pad_value >= MM_PAD_SHIFT_VALUE and item.pad_value == mm_pad_value(item.hash)
+    assert 129264 not in ids  # the raw <|deepseek_image|> placeholder is gone
+    assert item.feature.shape == (35 * 46, 3, 14, 14) and item.feature.dtype == torch.bfloat16
+
+
+@pytest.mark.skipif(
+    not os.path.exists(os.path.join(DSV4, "config.json")),
+    reason="no FREETOKEN_DSV4_VL_MODEL checkpoint",
+)
+def test_dsv4_two_images_take_two_blocks_and_the_second_is_shorter():
+    """The second image sits at a different offset, so its block is a different length."""
+    from freetoken.mm.processor import get_mm_processor
+    from freetoken.tokenizer.tokenize import TokenizeManager
+    from freetoken.utils.hf import load_tokenizer
+
+    manager = TokenizeManager(load_tokenizer(DSV4), get_mm_processor(DSV4))
+    raw = _png(320, 240)
+    out = manager.tokenize([_msg([raw, raw], n_parts=2)])[0]
+    first, second = out.mm_items
+    assert (first.start, second.start) == (first.offsets[0][0], second.offsets[0][0])
+    # the lead pads take up to three tokens of slack, so the blocks nearly match
+    assert abs(second.num_tokens - first.num_tokens) <= 3
+    # two images share a pad id only when their blocks are laid out identically, which is
+    # exactly when their offsets are congruent: otherwise one must not serve the other's key
+    assert (first.pad_value == second.pad_value) == ((second.start - first.start) % 4 == 0)

From 9b585b72d0838e0e977e029c427a827029397a8d Mon Sep 17 00:00:00 2001
From: ZihaoMu <zihaomu@outlook.com>
Date: Fri, 9 Oct 2026 10:55:57 +0800
Subject: [PATCH 080/249] fix(rocm): gate CUDA-only optional backends on ROCm
 (#134)

* fix(rocm): gate cuda-only optional kernel integrations

* test(rocm): narrow optional-backend regression coverage
---
 docs/install_amd.md                           |  8 ++++
 python/freetoken/engine/engine.py             | 15 ++++++++
 python/freetoken/kernel/backend.py            | 20 +++++-----
 .../layers/quantization/linear/nvfp4.py       |  2 +
 .../layers/quantization/moe/nvfp4.py          |  4 ++
 tests/engine/test_attention_backend_arch.py   | 38 +++++++++++++++++++
 6 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/docs/install_amd.md b/docs/install_amd.md
index d8b0c1bcb1..81d1ad5009 100644
--- a/docs/install_amd.md
+++ b/docs/install_amd.md
@@ -37,3 +37,11 @@ python -m pip install --no-build-isolation -e .
 
 Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the
 actual target reported by `rocminfo`.
+
+## Optional kernel backends
+
+FreeToken's FlashInfer, `sgl_kernel`, and vLLM kernel integrations are CUDA-only.
+On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`,
+`[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`,
+or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or
+`--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error.
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index bf3960ab8d..3988caf602 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -58,6 +58,10 @@ def _flashinfer_available() -> bool:
 
 
 def _sgl_flash_attn_available() -> bool:
+    from freetoken.kernel.backend import is_rocm
+
+    if is_rocm():
+        return False
     try:
         from sgl_kernel.flash_attn import flash_attn_with_kvcache  # noqa: F401
     except Exception as exc:
@@ -167,6 +171,7 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
     where a DSV4 or MLA checkpoint rejects a generic backend before weights load,
     and where a generic model rejects dsa/dsv4_sparse."""
     from freetoken.attention import validate_attn_backend
+    from freetoken.kernel.backend import is_rocm
 
     # Name membership first (ArgumentTypeError listing the supported names): the CLI already
     # ran this, but the programmatic EngineConfig path reaches here unvalidated and would
@@ -207,12 +212,22 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
     for part in backend_parts:
         info = attention_backend_info(part)
         if info.requires_flashinfer and not _flashinfer_available():
+            if is_rocm():
+                raise RuntimeError(
+                    f"Attention backend {part!r} uses CUDA-only FlashInfer kernels "
+                    "unavailable on ROCm; use --attention-backend auto or triton."
+                )
             raise RuntimeError(
                 f"Attention backend {config.attention_backend!r} requires flashinfer, which is "
                 "not installed. Install it with `pip install 'freetoken[fi]'` (or "
                 "'freetoken[accel]'), or use --attention-backend triton."
             )
         if info.requires_sgl_kernel and not _sgl_flash_attn_available():
+            if is_rocm():
+                raise RuntimeError(
+                    f"Attention backend {part!r} uses CUDA-only sgl_kernel kernels "
+                    "unavailable on ROCm; use --attention-backend auto or triton."
+                )
             raise RuntimeError(
                 f"Attention backend {config.attention_backend!r} requires sgl_kernel, which is "
                 "not installed. Install it with `pip install 'freetoken[sgl]'` (or "
diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py
index 6d1b4a9f8e..73a7dc6f36 100644
--- a/python/freetoken/kernel/backend.py
+++ b/python/freetoken/kernel/backend.py
@@ -4,6 +4,7 @@
 ops; otherwise they fall back to the pure-Triton kernels in
 ``freetoken.kernel.triton``. ``find_spec`` only checks that the package is
 importable (no import side effects), and the result is cached.
+On ROCm, the optional-package probes return False so callers use the fallbacks.
 """
 from __future__ import annotations
 
@@ -25,16 +26,22 @@ def _importable(name: str) -> bool:
 
 @functools.cache
 def is_flashinfer_installed() -> bool:
+    if is_rocm():
+        return False
     return _importable("flashinfer")
 
 
 @functools.cache
 def is_sgl_kernel_installed() -> bool:
+    if is_rocm():
+        return False
     return _importable("sgl_kernel")
 
 
 @functools.cache
 def is_vllm_installed() -> bool:
+    if is_rocm():
+        return False
     return _importable("vllm")
 
 
@@ -54,17 +61,6 @@ def is_rocm() -> bool:
     return getattr(torch.version, "hip", None) is not None
 
 
-@functools.cache
-def driver_hip_version() -> int | None:
-    """ROCm driver version, or None if undetermined."""
-    # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used.
-    try:
-        from freetoken.kernel.pinned import _load_pinned_extension
-        return int(_load_pinned_extension().driver_cuda_version()) or None
-    except Exception:
-        return None
-
-
 @functools.cache
 def driver_cuda_version() -> int | None:
     """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0),
@@ -73,6 +69,8 @@ def driver_cuda_version() -> int | None:
     toolkit version. Resolved through the ``_pinned_tensor`` extension's link-time
     cudart, so it works wherever the extension builds (including Windows) -- no dlopen
     by soname."""
+    if is_rocm():
+        return None
     try:
         from freetoken.kernel.pinned import _load_pinned_extension
 
diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py
index 0953b82501..80f5b78d3c 100644
--- a/python/freetoken/layers/quantization/linear/nvfp4.py
+++ b/python/freetoken/layers/quantization/linear/nvfp4.py
@@ -38,6 +38,8 @@ class MarlinNvfp4LinearKernel(LinearKernel):
     name = "marlin"
 
     def unusable_reason(self, cfg: LinearConfig) -> str | None:
+        if backend.is_rocm():
+            return "Marlin is CUDA-only and unavailable on ROCm; use triton"
         if not backend.is_vllm_installed():
             return "vLLM is not installed"
         if len(cfg.output_sizes) > 1:
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index cf66e2071b..1cd1a1ac73 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -228,6 +228,8 @@ class MarlinNvfp4MoEKernel(MoEKernel):
     max_slots = MARLIN_MAX_SLOTS
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
+        if backend.is_rocm():
+            return "Marlin is CUDA-only and unavailable on ROCm; use triton"
         if not backend.is_vllm_installed():
             return "vLLM is not installed"
         reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True)
@@ -491,6 +493,8 @@ class B12xNvfp4MoEKernel(MoEKernel):
     name = "b12x"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
+        if backend.is_rocm():
+            return "b12x is CUDA-only and unavailable on ROCm; use triton"
         cc = backend.device_capability()
         if cc < (12, 0):
             return f"b12x requires sm_120+, got sm_{cc[0]}{cc[1]}"
diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py
index 9271bd4254..25eb93547c 100644
--- a/tests/engine/test_attention_backend_arch.py
+++ b/tests/engine/test_attention_backend_arch.py
@@ -116,3 +116,41 @@ def test_fa_backend_rejects_cc12(monkeypatch):
     monkeypatch.setattr(arch, "_get_torch_cuda_version", lambda: (12, 0))
     with pytest.raises(RuntimeError, match="12.x"):
         FlashAttentionBackend(SimpleNamespace(head_dim=128))
+
+
+@pytest.mark.parametrize("requested", ["auto", "fi", "fa", "trtllm", "triton,fi"])
+def test_rocm_attention_selection_avoids_cuda_packages(monkeypatch, requested):
+    import builtins
+
+    from freetoken.engine import engine
+    from freetoken.kernel import backend
+
+    real_import = builtins.__import__
+
+    def guarded_import(name, *args, **kwargs):
+        if name.split(".")[0] in {"flashinfer", "sgl_kernel"}:
+            pytest.fail(f"ROCm tried to import {name}")
+        return real_import(name, *args, **kwargs)
+
+    def unexpected_probe(_name):
+        pytest.fail("ROCm tried to discover a CUDA package")
+
+    monkeypatch.setattr(builtins, "__import__", guarded_import)
+    monkeypatch.setattr(backend, "is_rocm", lambda: True)
+    monkeypatch.setattr(backend, "_importable", unexpected_probe)
+    monkeypatch.setattr(engine, "is_sm90_family", lambda: False)
+    monkeypatch.setattr(engine, "is_sm100_family", lambda: False)
+    backend.is_flashinfer_installed.cache_clear()
+    backend.is_sgl_kernel_installed.cache_clear()
+    try:
+        assert not engine._sgl_flash_attn_available()
+        config = _engine_config(attention_backend=requested)
+        if requested == "auto":
+            engine._adjust_config(config)
+            assert config.attention_backend == "triton"
+        else:
+            with pytest.raises(RuntimeError, match="CUDA-only.*ROCm"):
+                engine._adjust_config(config)
+    finally:
+        backend.is_flashinfer_installed.cache_clear()
+        backend.is_sgl_kernel_installed.cache_clear()

From 6c36af01108bfbe0c573394499a25ff2a7533f9f Mon Sep 17 00:00:00 2001
From: ZihaoMu <zihaomu@outlook.com>
Date: Fri, 9 Oct 2026 14:51:55 +0800
Subject: [PATCH 081/249] test(kernel): cover Triton activation and norm
 fallbacks (#652)

* test(kernel): cover triton activation and norm fallbacks

* test(kernel): allow fp32 activation approximation error
---
 tests/kernels/test_triton_fallbacks.py | 57 ++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 tests/kernels/test_triton_fallbacks.py

diff --git a/tests/kernels/test_triton_fallbacks.py b/tests/kernels/test_triton_fallbacks.py
new file mode 100644
index 0000000000..45234e66d4
--- /dev/null
+++ b/tests/kernels/test_triton_fallbacks.py
@@ -0,0 +1,57 @@
+"""Numerical checks for the built-in activation and norm fallback paths."""
+
+import pytest
+import torch
+
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="a CUDA or ROCm GPU is required"
+)
+
+
+@pytest.fixture(autouse=True)
+def _use_triton_fallbacks(monkeypatch):
+    from freetoken.kernel import backend
+
+    monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False)
+    monkeypatch.setattr(backend, "is_sgl_kernel_installed", lambda: False)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("activation", ["silu", "gelu", "gelu_tanh"])
+def test_activation_fallback_matches_torch(dtype, activation):
+    from freetoken.layers.activation import gated_act_and_mul
+
+    torch.manual_seed(0)
+    x = torch.randn((7, 256), device="cuda", dtype=dtype)
+    out = torch.full((7, 128), float("nan"), device="cuda", dtype=dtype)
+    gate, up = x.float().chunk(2, dim=-1)
+    if activation == "silu":
+        expected = torch.nn.functional.silu(gate) * up
+    else:
+        approx = "tanh" if activation == "gelu_tanh" else "none"
+        expected = torch.nn.functional.gelu(gate, approximate=approx) * up
+
+    gated_act_and_mul(activation, x, out)
+
+    # CUDA intentionally uses approximate math (notably tanh.approx.f32),
+    # while ROCm uses libdevice; allow for both against the Torch reference.
+    rtol, atol = (1e-4, 1e-4) if dtype == torch.float32 else (1e-2, 1e-3)
+    torch.testing.assert_close(out, expected.to(dtype), rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("kind", ["rmsnorm", "gemma"])
+def test_norm_fallback_matches_torch(dtype, kind):
+    from freetoken.layers.norm import GemmaRMSNorm, RMSNorm
+
+    torch.manual_seed(0)
+    norm = (RMSNorm if kind == "rmsnorm" else GemmaRMSNorm)(128, 1e-6)
+    norm.weight = torch.randn(128, device="cuda", dtype=dtype)
+    x = torch.randn((7, 128), device="cuda", dtype=dtype)
+    xf = x.float()
+    expected = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + norm.eps)
+    expected *= norm.weight.float()
+
+    rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-3)
+    torch.testing.assert_close(norm.forward(x), expected.to(dtype), rtol=rtol, atol=atol)

From 2d8b9cc1eac2e8f2ae63eb0ae90fe419252d9dc6 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 9 Oct 2026 00:15:07 -0700
Subject: [PATCH 082/249] perf(kernel): avoid the activation layout conversion
 in fp8 moe decode and read mxfp4 gemv weights as column pairs (#653)

---
 python/freetoken/kernel/moe_impl.py           |  9 ++++---
 .../kernel/triton/fp8_blockscale_moe.py       | 15 ++++++++----
 python/freetoken/kernel/triton/mxfp4_moe.py   | 24 ++++++++++++-------
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/python/freetoken/kernel/moe_impl.py b/python/freetoken/kernel/moe_impl.py
index ef569fa2e6..17fd77c4ed 100644
--- a/python/freetoken/kernel/moe_impl.py
+++ b/python/freetoken/kernel/moe_impl.py
@@ -484,12 +484,15 @@ def mxfp4_splitk_gemv_triton(
     k_groups = K // 32
     kgps = triton.cdiv(k_groups, num_splits)
     lut = get_fp4_lut(x.device)
+    assert N % 2 == 0 and block_n % 2 == 0, "the kernel reads byte pairs of adjacent columns"
+    w16 = w_blocks_t.view(torch.int16)
+    s16 = w_scales_t.view(torch.int16)
 
     grid = (triton.cdiv(N, block_n), routes * num_splits)
     mxfp4_splitk_gemv_kernel[grid](
-        x, w_blocks_t, w_scales_t, bias_arg, expert_ids, partial, lut, N, K,
-        stride_xe, w_blocks_t.stride(0), w_blocks_t.stride(1),
-        w_scales_t.stride(0), w_scales_t.stride(1), bias_stride, N,
+        x, w16, s16, bias_arg, expert_ids, partial, lut, N, K,
+        stride_xe, w16.stride(0), w16.stride(1),
+        s16.stride(0), s16.stride(1), bias_stride, N,
         HAS_BIAS=has_bias, BLOCK_N=block_n,  # type: ignore
         NUM_K_SPLITS=num_splits, K_GROUPS_PER_SPLIT=kgps,  # type: ignore
         num_warps=num_warps,
diff --git a/python/freetoken/kernel/triton/fp8_blockscale_moe.py b/python/freetoken/kernel/triton/fp8_blockscale_moe.py
index 747605168a..1db54a8072 100644
--- a/python/freetoken/kernel/triton/fp8_blockscale_moe.py
+++ b/python/freetoken/kernel/triton/fp8_blockscale_moe.py
@@ -35,7 +35,7 @@ def _decode_fp8_moe_kernel(
     stride_twm, stride_twk, stride_tidm, stride_tidk,
     BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,  # BLOCK_K == 128 (one weight scale block)
     TOP_K: tl.constexpr, A_ROW_IS_ROUTE: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr,
-    compute_type: tl.constexpr,
+    compute_type: tl.constexpr, EVEN_K: tl.constexpr,
 ):
     route_id = tl.program_id(0)
     n_block = tl.program_id(1)
@@ -54,7 +54,10 @@ def _decode_fp8_moe_kernel(
     acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
     for kb in range(tl.cdiv(K, BLOCK_K)):
         offs_k = kb * BLOCK_K + tl.arange(0, BLOCK_K)
-        k_mask = offs_k < K
+        if EVEN_K:
+            k_mask = tl.full((BLOCK_K,), True, tl.int1)
+        else:
+            k_mask = offs_k < K
         if e4m3_native_cx():
             w = tl.load(
                 w_slot + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk,
@@ -66,8 +69,10 @@ def _decode_fp8_moe_kernel(
                 mask=n_mask[:, None] & k_mask[None, :], other=0,
             ))
         sc = tl.load(s_slot + sn * stride_sn + kb * stride_sk, mask=n_mask, other=0.0).to(tl.float32)
-        a = tl.load(a_base + offs_k * stride_ak, mask=k_mask, other=0.0).to(tl.float32)
-        acc += tl.sum(w * a[None, :], axis=1) * sc
+        # load the activation row at the weight tile shape so it gets the weight tile layout (no shared-memory layout conversion)
+        a_ptrs = tl.broadcast_to(a_base + offs_k[None, :] * stride_ak, (BLOCK_N, BLOCK_K))
+        a = tl.load(a_ptrs, mask=k_mask[None, :], other=0.0).to(tl.float32)
+        acc += tl.sum(w * a, axis=1) * sc
 
     if MUL_ROUTED_WEIGHT:
         acc *= tl.load(topk_weights_ptr + token_id * stride_twm + route_k * stride_twk)
@@ -92,7 +97,7 @@ def _decode_gemm(a, w, s, c, topk_weights, topk_ids, mul_routed_weight, a_row_is
         topk_weights.stride(0), topk_weights.stride(1), topk_ids.stride(0), topk_ids.stride(1),
         BLOCK_N=BLOCK_N, BLOCK_K=128, TOP_K=top_k,
         A_ROW_IS_ROUTE=a_row_is_route, MUL_ROUTED_WEIGHT=mul_routed_weight,
-        compute_type=_TL.get(c.dtype, tl.bfloat16), num_warps=4,
+        compute_type=_TL.get(c.dtype, tl.bfloat16), EVEN_K=w.shape[2] % 128 == 0, num_warps=4,
     )
 
 
diff --git a/python/freetoken/kernel/triton/mxfp4_moe.py b/python/freetoken/kernel/triton/mxfp4_moe.py
index 3981d2a391..c0237f25a1 100644
--- a/python/freetoken/kernel/triton/mxfp4_moe.py
+++ b/python/freetoken/kernel/triton/mxfp4_moe.py
@@ -53,30 +53,36 @@ def mxfp4_splitk_gemv_kernel(
     expert_id = tl.load(expert_ids_ptr + pid_e)
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     mask_n = offs_n < N
+    # w_ptr / s_ptr are int16 views: one element holds the bytes of columns 2p (low) and 2p+1 (high)
+    offs_p = pid_n * (BLOCK_N // 2) + tl.arange(0, BLOCK_N // 2)
+    mask_p = offs_p < N // 2
 
     k_groups = K // 32
     kg0 = pid_k * K_GROUPS_PER_SPLIT
     kg1 = tl.minimum(kg0 + K_GROUPS_PER_SPLIT, k_groups)
 
-    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
+    acc_even = tl.zeros([BLOCK_N // 2], dtype=tl.float32)
+    acc_odd = tl.zeros([BLOCK_N // 2], dtype=tl.float32)
     w_base = expert_id * stride_we
     s_base = expert_id * stride_se
     x_base = pid_e * stride_xe
 
     for kg in range(kg0, kg1):
-        s_val = tl.load(s_ptr + s_base + kg * stride_sk + offs_n, mask=mask_n, other=0)
-        scale_f = _e8m0_scale(s_val)
+        s_pair = tl.load(s_ptr + s_base + kg * stride_sk + offs_p, mask=mask_p, other=0).to(tl.int32)
+        scale_even = _e8m0_scale(s_pair & 0xFF)
+        scale_odd = _e8m0_scale((s_pair >> 8) & 0xFF)
         for kk in tl.static_range(16):
             k_packed = kg * 16 + kk
-            w_byte = tl.load(
-                w_ptr + w_base + k_packed * stride_wk + offs_n, mask=mask_n, other=0
+            w_pair = tl.load(
+                w_ptr + w_base + k_packed * stride_wk + offs_p, mask=mask_p, other=0
             ).to(tl.int32)
-            lo = w_byte & 0x0F
-            hi = (w_byte >> 4) & 0x0F
             x_lo = tl.load(x_ptr + x_base + kg * 32 + kk * 2).to(tl.float32)
             x_hi = tl.load(x_ptr + x_base + kg * 32 + kk * 2 + 1).to(tl.float32)
-            acc += _fp4_table_lut(lo, lut_ptr) * scale_f * x_lo
-            acc += _fp4_table_lut(hi, lut_ptr) * scale_f * x_hi
+            acc_even += _fp4_table_lut(w_pair & 0x0F, lut_ptr) * scale_even * x_lo
+            acc_even += _fp4_table_lut((w_pair >> 4) & 0x0F, lut_ptr) * scale_even * x_hi
+            acc_odd += _fp4_table_lut((w_pair >> 8) & 0x0F, lut_ptr) * scale_odd * x_lo
+            acc_odd += _fp4_table_lut((w_pair >> 12) & 0x0F, lut_ptr) * scale_odd * x_hi
+    acc = tl.reshape(tl.join(acc_even, acc_odd), (BLOCK_N,))
 
     if HAS_BIAS and pid_k == 0:
         acc += tl.load(

From 95281904548d01b2913cf7045c167cda4d025efb Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 10 Oct 2026 00:36:04 -0700
Subject: [PATCH 083/249] feat(server): add --api-key authentication (#639)

* feat(server): require an API key on every route but /health

* feat(shell): send the server's API key

* feat(ctl): send the server's API key

* feat(launch): hand the server's API key to the agent

* feat(daemon): pass the serve's API key as apiKey

The daemon reaches the serve's /v1/stats and prepare-stop itself, so the key
travels in the start request, reaches the serve through its environment and is
kept in the state file for re-adoption. The daemon's own FREETOKEN_API_KEY is no
longer passed down to the serve.

* fix(launch): keep the API key out of the variable dsh's web search reads

dsh's web search reads DEEPSEEK_API_KEY and sends it to api.deepseek.com, so
dsh now reads the key from FREETOKEN_API_KEY through the apiKeyEnv of its
llm-deepseek section, and DEEPSEEK_API_KEY stays the dummy.

* fix(server): check the API key on OPTIONS requests

* fix(launch): write agent configs that hold the API key owner-only

* chore: drop comments that restate the code
---
 docs/cli.md                               |   7 +-
 python/freetoken/control_cli.py           |  37 +++++-
 python/freetoken/daemon/README.md         |  11 +-
 python/freetoken/daemon/app.py            |  22 +++-
 python/freetoken/daemon/client.py         |   8 ++
 python/freetoken/daemon/pidfile.py        |  16 ++-
 python/freetoken/daemon/proxy.py          |  18 ++-
 python/freetoken/daemon/serve_manager.py  |  68 +++++++---
 python/freetoken/daemon/server.py         |   3 +-
 python/freetoken/launch.py                | 109 ++++++++++------
 python/freetoken/server/api_server.py     |  43 ++++++-
 python/freetoken/server/args.py           |  20 ++-
 python/freetoken/shell/__init__.py        |   8 +-
 python/freetoken/shell/client.py          |   9 +-
 python/freetoken/shell/tui.py             |   6 +-
 tests/daemon/test_daemon_serve_manager.py | 137 +++++++++++++++++++-
 tests/server/test_api_key.py              | 146 ++++++++++++++++++++++
 17 files changed, 582 insertions(+), 86 deletions(-)
 create mode 100644 tests/server/test_api_key.py

diff --git a/docs/cli.md b/docs/cli.md
index b9be878ce7..a15a34b71c 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -41,6 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--api-key` | `$FREETOKEN_API_KEY`, else off | Require the key on every route except `/health`, as `Authorization: Bearer <key>` or `x-api-key: <key>`; other requests get a 401 |
 | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |
@@ -131,14 +132,17 @@ ft shell --model ~/models/Qwen3.6-35B-A3B   # serve + chat in one process
 ```
 
 - Attach mode talks to `--server URL` (default `http://127.0.0.1:1919`)
+- `--api-key KEY` (default `$FREETOKEN_API_KEY`) attaches to a server started with `--api-key`; with `--model` the shell uses that server's own key
 - `/help` inside the shell lists the commands (`/think`, `/cache`, `/reset`).
 
 ## ft ctl
 
 ```bash
-ft ctl [--base-url http://127.0.0.1:1919] [--timeout 10] [--json] <subcommand>
+ft ctl [--base-url http://127.0.0.1:1919] [--api-key KEY] [--timeout 10] [--json] <subcommand>
 ```
 
+`--api-key` defaults to `$FREETOKEN_API_KEY`.
+
 | Subcommand | Endpoint | Purpose |
 |---|---|---|
 | `health` | `GET /health` | Server status, model, load progress |
@@ -166,6 +170,7 @@ Code and Hermes need no declaration.
 | Flag | Meaning |
 |---|---|
 | `--server URL` | Server to point the agent at (default `http://127.0.0.1:1919`) |
+| `--api-key KEY` | Key of a server started with `--api-key` (default `$FREETOKEN_API_KEY`). The agent gets it in its environment; Hermes and OpenClaw keep it in their config files |
 | `--dry-run` | Print the planned config changes and command, touch nothing |
 | `-y`, `--yes` | Approve install/config prompts |
 | `--config` | Configure without launching |
diff --git a/python/freetoken/control_cli.py b/python/freetoken/control_cli.py
index 50a32aae6e..e33eb4acb8 100644
--- a/python/freetoken/control_cli.py
+++ b/python/freetoken/control_cli.py
@@ -10,6 +10,8 @@
 from collections.abc import Sequence
 from typing import Any
 
+from freetoken.launch import resolve_api_key
+
 
 DEFAULT_BASE_URL = "http://127.0.0.1:1919"
 # (CLI dest, pool, rebuild body key). moe/mamba are slot counts and pass straight through;
@@ -74,6 +76,7 @@ def _request_json(
     body: dict[str, Any] | None = None,
     query: dict[str, Any] | None = None,
     timeout: float = 10.0,
+    api_key: str | None,
 ) -> dict[str, Any]:
     with _open_request(
         method,
@@ -83,6 +86,7 @@ def _request_json(
         query=query,
         accept="application/json",
         timeout=timeout,
+        api_key=api_key,
     ) as response:
         raw = response.read()
 
@@ -104,6 +108,7 @@ def _request_sse_text(
     *,
     body: dict[str, Any] | None = None,
     timeout: float = 10.0,
+    api_key: str | None,
 ) -> str:
     with _open_request(
         method,
@@ -112,6 +117,7 @@ def _request_sse_text(
         body=body,
         accept="text/event-stream",
         timeout=timeout,
+        api_key=api_key,
     ) as response:
         return _read_sse_text(response)
 
@@ -125,9 +131,12 @@ def _open_request(
     query: dict[str, Any] | None = None,
     accept: str,
     timeout: float,
+    api_key: str | None,
 ) -> Any:
     data = None
     headers = {"Accept": accept}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
     if body is not None:
         data = json.dumps(body).encode("utf-8")
         headers["Content-Type"] = "application/json"
@@ -284,6 +293,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser:
         default=DEFAULT_BASE_URL,
         help=f"FreeToken server URL (default: {DEFAULT_BASE_URL})",
     )
+    parser.add_argument(
+        "--api-key",
+        help="API key of a server started with --api-key (default: $FREETOKEN_API_KEY)",
+    )
     parser.add_argument("--timeout", type=float, default=10.0, help="HTTP timeout in seconds")
     parser.add_argument("--json", action="store_true", help="Print raw JSON response")
 
@@ -355,14 +368,19 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int:
         args = parser.parse_args(list(argv) if argv is not None else None)
     except SystemExit as exc:
         return int(exc.code) if isinstance(exc.code, int) else 2
+    args.api_key = resolve_api_key(args.api_key)
 
     try:
         if args.command == "health":
-            doc = _request_json("GET", args.base_url, "/health", timeout=args.timeout)
+            doc = _request_json(
+                "GET", args.base_url, "/health", timeout=args.timeout, api_key=args.api_key
+            )
             _print_doc(doc, _format_health, raw_json=args.json)
             return 0
         if args.command == "stats":
-            doc = _request_json("GET", args.base_url, "/v1/stats", timeout=args.timeout)
+            doc = _request_json(
+                "GET", args.base_url, "/v1/stats", timeout=args.timeout, api_key=args.api_key
+            )
             _print_doc(doc, _format_stats, raw_json=args.json)
             return 0
         if args.command == "generate":
@@ -376,6 +394,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int:
                     "ignore_eos": args.ignore_eos,
                 },
                 timeout=args.timeout,
+                api_key=args.api_key,
             )
             if args.json:
                 print(json.dumps({"text": text}, ensure_ascii=False, indent=2, sort_keys=True))
@@ -390,13 +409,16 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int:
                 "/v1/requests",
                 query=query,
                 timeout=args.timeout,
+                api_key=args.api_key,
             )
             _print_doc(doc, _format_requests, raw_json=args.json)
             return 0
         if args.command == "cache" and args.cache_command in (None, "status"):
             if args.cache_command is None and _cache_targets(args):
                 return _run_cache_rebuild(args)
-            doc = _request_json("GET", args.base_url, "/v1/cache/status", timeout=args.timeout)
+            doc = _request_json(
+                "GET", args.base_url, "/v1/cache/status", timeout=args.timeout, api_key=args.api_key
+            )
             _print_doc(doc, _format_cache_status, raw_json=args.json)
             return 0
         if args.command == "cache" and args.cache_command == "rebuild":
@@ -444,7 +466,9 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int:
         print("error: cache rebuild requires at least one cache target", file=sys.stderr)
         return 2
     # The token targets can only be converted against the live geometry, so read it first.
-    status = _request_json("GET", args.base_url, "/v1/cache/status", timeout=args.timeout)
+    status = _request_json(
+        "GET", args.base_url, "/v1/cache/status", timeout=args.timeout, api_key=args.api_key
+    )
     body = _rebuild_body(targets, status.get("geometry") or {})
     body["timeout"] = args.wait
     doc = _request_json(
@@ -453,6 +477,7 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int:
         "/v1/cache/rebuild",
         body=body,
         timeout=args.wait + args.timeout,
+        api_key=args.api_key,
     )
     if args.json or doc.get("status") != "ok":
         _print_doc(doc, _format_rebuild, raw_json=args.json)
@@ -462,7 +487,9 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int:
     # /cache prints. Best-effort -- the rebuild already succeeded.
     print("status=ok")
     try:
-        print(_format_cache_status(_request_json("GET", args.base_url, "/v1/cache/status")))
+        print(_format_cache_status(
+            _request_json("GET", args.base_url, "/v1/cache/status", api_key=args.api_key)
+        ))
     except ControlCliError:
         pass
     return 0
diff --git a/python/freetoken/daemon/README.md b/python/freetoken/daemon/README.md
index ea3b016eaf..ebe4a17008 100644
--- a/python/freetoken/daemon/README.md
+++ b/python/freetoken/daemon/README.md
@@ -60,9 +60,9 @@ Target a non-default daemon with `--url http://host:1900` (or `$FREETOKEN_DAEMON
 | Method / path | Notes |
 | --- | --- |
 | `GET /health` | Daemon self-health; always answers, never gated by `--token`. |
-| `POST /engine/start` `{model,port,args[]}` | Idempotent on the full `(model,port,args)`; a differing config on the same port → `409`. |
+| `POST /engine/start` `{model,port,args[],apiKey?}` | Idempotent on the full `(model,port,args,apiKey)`; a differing config on the same port → `409`. |
 | `POST /engine/stop` `{force?:false}` | Close admission, drain/abort, durably enqueue the final-accounting receipt, then `SIGTERM`→grace→`SIGKILL`. A prepare/outbox failure preserves the engine. |
-| `POST /engine/switch` `{model,port,args[],force?:false}` | One serialized stop-accounting-start transaction. |
+| `POST /engine/switch` `{model,port,args[],apiKey?,force?:false}` | One serialized stop-accounting-start transaction. |
 | `GET /engine/status` | `{running,pid,model,port,uptimeS,lastExitCode,…}`; outlives any single serve. |
 | `GET /engine/logs?since=` | SSE, ANSI-stripped, tqdm-`\r` collapsed, ring replay, `id:<seq>`, `Last-Event-ID` resume. |
 | `GET /engine/metrics` | `{ramBytes,vramBytes}` — the serve tree's own footprint only. |
@@ -75,6 +75,13 @@ Target a non-default daemon with `--url http://host:1900` (or `$FREETOKEN_DAEMON
 Set `--token` (or `$FREETOKEN_DAEMON_TOKEN`) to require an `X-FT-Token` header on everything
 except `/health`.
 
+`apiKey` is the serve's own `--api-key`. The daemon hands it to the serve as `FREETOKEN_API_KEY` in
+its environment (never on its command line), sends it on its own calls to the serve, and keeps it in
+the owner-only state file so a re-adopted serve stays reachable. The daemon's own
+`FREETOKEN_API_KEY` is never passed down: a serve requires a key only when its start request carries
+one, and an `--api-key` inside `args` is rejected with a 400 because the daemon would not know it.
+`ft daemon start|switch --api-key KEY` (default `$FREETOKEN_API_KEY`) sets it.
+
 The daemon reaches the serve's destructive `POST /v1/admin/prepare-stop` endpoint only over
 loopback; the serve rejects non-loopback callers even when its inference API is bound to
 `0.0.0.0`. `force` is never implicit: it is an explicit recovery choice that may record null
diff --git a/python/freetoken/daemon/app.py b/python/freetoken/daemon/app.py
index d7a0a53c60..dd506f78fc 100644
--- a/python/freetoken/daemon/app.py
+++ b/python/freetoken/daemon/app.py
@@ -30,6 +30,7 @@ class StartBody(BaseModel):
     model: str
     port: int | None = None
     args: list[str] = []
+    apiKey: str | None = None
 
 
 class StopBody(BaseModel):
@@ -161,6 +162,15 @@ def resolve_port(explicit: int | None) -> int:
         st = manager.status()
         return st.get("port") or default_serve_port
 
+    def reject_api_key_arg(args: list[str]) -> None:
+        # the daemon must know the key to reach the serve, so the key only travels as apiKey
+        # argparse also accepts any unambiguous prefix, and no other serve flag starts with --ap
+        flags = (a.split("=", 1)[0] for a in args)
+        if any(len(f) >= 4 and "--api-key".startswith(f) for f in flags):
+            raise HTTPException(
+                status_code=400, detail="pass the serve's API key as apiKey, not in args"
+            )
+
     def accounting_error(exc: Exception) -> JSONResponse:
         code = (
             "accounting_outbox_failed"
@@ -192,9 +202,17 @@ async def health():
 
     @app.post("/engine/start", dependencies=auth)
     async def engine_start(body: StartBody):
+        reject_api_key_arg(body.args)
         port = resolve_port(body.port)
         try:
-            return await run(lifecycle_pool, manager.start, body.model, port, list(body.args))
+            return await run(
+                lifecycle_pool,
+                manager.start,
+                body.model,
+                port,
+                list(body.args),
+                body.apiKey or None,
+            )
         except Conflict as exc:
             st = manager.status()
             return JSONResponse(
@@ -236,6 +254,7 @@ async def shutdown_daemon(request: Request, body: StopBody | None = None):
 
     @app.post("/engine/switch", dependencies=auth)
     async def engine_switch(body: SwitchBody):
+        reject_api_key_arg(body.args)
         port = resolve_port(body.port)
         try:
             return await run(
@@ -245,6 +264,7 @@ async def engine_switch(body: SwitchBody):
                 port,
                 list(body.args),
                 body.force,
+                body.apiKey or None,
             )
         except (AccountingPrepareError, AccountingOutboxError) as exc:
             return accounting_error(exc)
diff --git a/python/freetoken/daemon/client.py b/python/freetoken/daemon/client.py
index 6616778e72..9ded96276e 100644
--- a/python/freetoken/daemon/client.py
+++ b/python/freetoken/daemon/client.py
@@ -16,6 +16,8 @@
 from collections.abc import Sequence
 from typing import Any
 
+from freetoken.launch import resolve_api_key
+
 DEFAULT_URL = "http://127.0.0.1:1900"
 DEFAULT_TIMEOUT = 10.0
 # prepare-stop (15s transport budget) + default SIGTERM grace (10s) + reap wait (10s),
@@ -144,6 +146,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser:
         sp = sub.add_parser(name, parents=[common], help=f"POST /engine/{name}")
         sp.add_argument("model", help="Model path/id")
         sp.add_argument("--port", type=int, default=None, help="Serve port")
+        sp.add_argument(
+            "--api-key",
+            help="API key the serve requires, set in its environment (default: $FREETOKEN_API_KEY)",
+        )
         if name == "switch":
             sp.add_argument(
                 "--force",
@@ -181,6 +187,8 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft daemon") -> int:
             body: dict[str, Any] = {"model": args.model, "args": list(args.serve_args)}
             if args.port is not None:
                 body["port"] = args.port
+            if api_key := resolve_api_key(args.api_key):
+                body["apiKey"] = api_key
             if args.verb == "switch" and args.force:
                 body["force"] = True
             method, path = "POST", f"/engine/{args.verb}"
diff --git a/python/freetoken/daemon/pidfile.py b/python/freetoken/daemon/pidfile.py
index 89ffc2b4df..b9865a83d3 100644
--- a/python/freetoken/daemon/pidfile.py
+++ b/python/freetoken/daemon/pidfile.py
@@ -4,15 +4,16 @@
 open for the daemon's whole life. flock (not a bare port-bind) avoids TIME_WAIT races and is
 released automatically if the daemon dies, so a crashed daemon never wedges its own restart.
 
-``ServeStateStore`` persists ``{model, port, pid, args, starttime, logPath}`` as JSON on every
+``ServeStateStore`` persists ``{model, port, pid, args, starttime, log_path, api_key}`` as JSON on every
 lifecycle change, so a restarted daemon can re-adopt a still-running serve. ``starttime`` +
-``args`` are what make re-adoption PID-reuse-safe and config-exact."""
+``args`` are what make re-adoption PID-reuse-safe and config-exact.
+The file is owner-only because it holds the serve's API key."""
 
 from __future__ import annotations
 
 import json
 import os
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 
 
 class AlreadyRunning(RuntimeError):
@@ -63,6 +64,7 @@ class ServeState:
     args: list[str]
     starttime: int | None = None
     log_path: str | None = None
+    api_key: str | None = field(default=None, repr=False)
 
 
 class ServeStateStore:
@@ -72,7 +74,12 @@ def __init__(self, path: str) -> None:
     def save(self, state: ServeState) -> None:
         os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
         tmp = f"{self.path}.tmp"
-        with open(tmp, "w") as fh:
+        # unlink first: O_CREAT applies the 0o600 mode only to a file it creates
+        try:
+            os.remove(tmp)
+        except FileNotFoundError:
+            pass
+        with os.fdopen(os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w") as fh:
             json.dump(asdict(state), fh)
             fh.flush()
             os.fsync(fh.fileno())
@@ -102,6 +109,7 @@ def load(self) -> ServeState | None:
                 args=list(doc.get("args") or []),
                 starttime=doc.get("starttime"),
                 log_path=doc.get("log_path"),
+                api_key=doc.get("api_key"),
             )
         except (KeyError, TypeError, ValueError):
             return None
diff --git a/python/freetoken/daemon/proxy.py b/python/freetoken/daemon/proxy.py
index 9e382a4d72..c4f0ba2864 100644
--- a/python/freetoken/daemon/proxy.py
+++ b/python/freetoken/daemon/proxy.py
@@ -49,6 +49,7 @@ def __init__(
         now: Callable[[], float] = time.monotonic,
         opener: Callable[[str, float], dict] | None = None,
         prepare_opener: Callable[[str, float], dict] | None = None,
+        api_key: Callable[[], str | None] = lambda: None,
     ) -> None:
         self._host = host
         self._ttl = ttl_s
@@ -57,6 +58,7 @@ def __init__(
         self._now = now
         self._opener = opener or self._urlopen
         self._prepare_opener = prepare_opener or self._urlopen_prepare
+        self._api_key = api_key
         self._lock = threading.Lock()
         self._cache: dict[tuple[str, int], tuple[float, dict]] = {}
 
@@ -109,19 +111,23 @@ def _fetch(self, path: str, port: int) -> dict:
         result["reachable"] = True
         return result
 
-    @staticmethod
-    def _urlopen(url: str, timeout: float) -> dict:
-        req = urllib.request.Request(url, headers={"Accept": "application/json"}, method="GET")
+    def _headers(self) -> dict[str, str]:
+        headers = {"Accept": "application/json"}
+        if key := self._api_key():
+            headers["Authorization"] = f"Bearer {key}"
+        return headers
+
+    def _urlopen(self, url: str, timeout: float) -> dict:
+        req = urllib.request.Request(url, headers=self._headers(), method="GET")
         with urllib.request.urlopen(req, timeout=timeout) as resp:
             raw = resp.read()
         return json.loads(raw.decode("utf-8"))
 
-    @staticmethod
-    def _urlopen_prepare(url: str, timeout: float) -> dict:
+    def _urlopen_prepare(self, url: str, timeout: float) -> dict:
         req = urllib.request.Request(
             url,
             data=b"{}",
-            headers={"Accept": "application/json", "Content-Type": "application/json"},
+            headers={**self._headers(), "Content-Type": "application/json"},
             method="POST",
         )
         with urllib.request.urlopen(req, timeout=timeout) as resp:
diff --git a/python/freetoken/daemon/serve_manager.py b/python/freetoken/daemon/serve_manager.py
index 681bdef497..9d5e51fe71 100644
--- a/python/freetoken/daemon/serve_manager.py
+++ b/python/freetoken/daemon/serve_manager.py
@@ -26,6 +26,8 @@
 from dataclasses import dataclass
 from typing import Any, Callable
 
+from freetoken.launch import API_KEY_ENV
+
 from . import osproc
 from .accounting import (
     AccountingOutbox,
@@ -37,7 +39,7 @@
 
 
 class Conflict(RuntimeError):
-    """A different serve (model/port/args) is already running; the client should switch()."""
+    """A different serve (model/port/args/key) is already running; the client should switch()."""
 
 
 @dataclass
@@ -128,7 +130,14 @@ def build_serve_command(
     return argv, log_path
 
 
-def spawn_serve(argv: list[str], log_path: str) -> PopenChild:
+def serve_env(api_key: str | None) -> dict[str, str]:
+    env = {k: v for k, v in os.environ.items() if k != API_KEY_ENV}
+    if api_key:
+        env[API_KEY_ENV] = api_key
+    return env
+
+
+def spawn_serve(argv: list[str], log_path: str, api_key: str | None = None) -> PopenChild:
     """Spawn the serve as a session leader with its stdout+stderr going to a real logfile fd (never
     a PIPE): the file survives daemon death, so a re-adopting daemon just resumes tailing it and
     the serve never writes to a dead pipe. ``start_new_session`` makes it a process
@@ -138,6 +147,7 @@ def spawn_serve(argv: list[str], log_path: str) -> PopenChild:
     try:
         proc = subprocess.Popen(
             argv,
+            env=serve_env(api_key),
             stdout=logf,
             stderr=subprocess.STDOUT,
             stdin=subprocess.DEVNULL,
@@ -158,7 +168,7 @@ def __init__(
         ring,
         state_store,
         *,
-        spawn_fn: Callable[[str, int, list[str]], object] | None = None,
+        spawn_fn: Callable[[str, int, list[str], str | None], object] | None = None,
         adopt_fn: Callable[[object], object] | None = None,
         tailer_factory: Callable[[object], object] | None = None,
         signal_fn: Callable[[int, int], None] | None = None,
@@ -217,6 +227,7 @@ def __init__(
         self._model: str | None = None
         self._port: int | None = None
         self._args: list[str] = []
+        self._api_key: str | None = None
         self._started_at: float | None = None
         self._last_exit: ExitInfo | None = None
         self._adopted = False
@@ -232,11 +243,11 @@ def __init__(
 
     # ---- default factories (overridable for tests) ----
 
-    def _default_spawn(self, model: str, port: int, args: list[str]):
+    def _default_spawn(self, model: str, port: int, args: list[str], api_key: str | None):
         argv, log_path = build_serve_command(
             model, port, args, python=self._python, log_dir=self._log_dir
         )
-        return spawn_serve(argv, log_path)
+        return spawn_serve(argv, log_path, api_key)
 
     def _default_adopt(self, state):
         if not osproc.is_ft_serve_on_port(state.pid, state.port, starttime=state.starttime):
@@ -256,13 +267,25 @@ def _default_adopt(self, state):
     # ---- lifecycle ----
 
     def start(
-        self, model: str, port: int, args: list[str] | None = None, *, _auto: bool = False
+        self,
+        model: str,
+        port: int,
+        args: list[str] | None = None,
+        api_key: str | None = None,
+        *,
+        _auto: bool = False,
     ) -> dict:
         with self._lifecycle:
-            return self._start(model, port, args, _auto=_auto)
+            return self._start(model, port, args, api_key, _auto=_auto)
 
     def _start(
-        self, model: str, port: int, args: list[str] | None = None, *, _auto: bool = False
+        self,
+        model: str,
+        port: int,
+        args: list[str] | None = None,
+        api_key: str | None = None,
+        *,
+        _auto: bool = False,
     ) -> dict:
         args = list(args or [])
         with self._cond:
@@ -278,10 +301,13 @@ def _start(
                 self._stop_requested = False  # an explicit client start clears a prior stop intent
             while True:
                 if self._child is not None and not self._stopping:
-                    if (self._model, self._port, self._args) == (model, port, args):
+                    if (self._model, self._port, self._args, self._api_key) == (
+                        model, port, args, api_key
+                    ):
                         return {"pid": self._child.pid, "idempotent": True}
                     raise Conflict(
-                        f"serve already running (model={self._model!r} port={self._port}); "
+                        f"serve already running with a different config "
+                        f"(model={self._model!r} port={self._port}); "
                         f"use switch to replace it"
                     )
                 if self._starting:
@@ -299,7 +325,7 @@ def _start(
 
         child = None
         try:
-            child = self._spawn_fn(model, port, args)
+            child = self._spawn_fn(model, port, args, api_key)
         finally:
             if child is None:
                 with self._cond:
@@ -311,6 +337,7 @@ def _start(
             self._model = model
             self._port = port
             self._args = args
+            self._api_key = api_key
             self._started_at = self._now()
             self._last_exit = None
             self._adopted = getattr(child, "adopted", False)
@@ -318,7 +345,7 @@ def _start(
             self._starting = False
             self._cond.notify_all()
 
-        self._persist(child, model, port, args)
+        self._persist(child, model, port, args, api_key)
         self._maybe_apply_oom(child.pid)
         self._begin_watch(child)
         self._emit(f"serve started (pid={child.pid} model={model} port={port})")
@@ -407,10 +434,11 @@ def switch(
         port: int,
         args: list[str] | None = None,
         force: bool = False,
+        api_key: str | None = None,
     ) -> dict:
         with self._lifecycle:
             stopped = self._stop(force=force)
-            started = self._start(model, port, args)
+            started = self._start(model, port, args, api_key)
             return {**started, "accounting": stopped["accounting"]}
 
     def pending_accounting(self) -> list[dict[str, Any]]:
@@ -734,6 +762,7 @@ def readopt(self) -> bool:
             self._model = state.model
             self._port = state.port
             self._args = list(state.args)
+            self._api_key = state.api_key
             self._started_at = self._now()  # unknown true start; report uptime since adoption
             self._last_exit = None
             self._adopted = True
@@ -824,13 +853,13 @@ def _exit_line(self, child, info: ExitInfo, was_stopping: bool) -> str:
         return f"serve exited with code {info.code} (pid={child.pid})"
 
     def _restart_async(self, dead_child) -> None:
-        model, port, args = self._model, self._port, list(self._args)
+        model, port, args, api_key = self._model, self._port, list(self._args), self._api_key
         if model is None or port is None:
             return
 
         def _run():
             try:
-                self.start(model, port, args, _auto=True)
+                self.start(model, port, args, api_key, _auto=True)
             except Exception as exc:  # noqa: BLE001
                 self._emit(f"auto-restart failed: {exc}")
 
@@ -867,6 +896,10 @@ def serve_args(self) -> list[str]:
         with self._cond:
             return list(self._args)
 
+    def serve_api_key(self) -> str | None:
+        with self._cond:
+            return self._api_key
+
     def status(self) -> dict:
         with self._cond:
             child = self._child
@@ -892,7 +925,9 @@ def status(self) -> dict:
 
     # ---- helpers ----
 
-    def _persist(self, child, model: str, port: int, args: list[str]) -> None:
+    def _persist(
+        self, child, model: str, port: int, args: list[str], api_key: str | None
+    ) -> None:
         from .pidfile import ServeState
 
         self._store.save(
@@ -903,6 +938,7 @@ def _persist(self, child, model: str, port: int, args: list[str]) -> None:
                 args=list(args),
                 starttime=getattr(child, "starttime", None),
                 log_path=getattr(child, "log_path", None),
+                api_key=api_key,
             )
         )
 
diff --git a/python/freetoken/daemon/server.py b/python/freetoken/daemon/server.py
index d4e742b0be..76e7e789c9 100644
--- a/python/freetoken/daemon/server.py
+++ b/python/freetoken/daemon/server.py
@@ -132,7 +132,8 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft daemon") -> int:
 
     ring = LogRing(capacity=args.log_capacity)
     store = ServeStateStore(os.path.join(state_dir, "serve.json"))
-    probe = ServeProbe()
+    # read at call time: the running serve's key, from the manager created below
+    probe = ServeProbe(api_key=lambda: manager.serve_api_key())
 
     def tailer_factory(child):
         log_path = getattr(child, "log_path", None)
diff --git a/python/freetoken/launch.py b/python/freetoken/launch.py
index 5aac9389fe..e3c42154ab 100644
--- a/python/freetoken/launch.py
+++ b/python/freetoken/launch.py
@@ -5,6 +5,7 @@
 import os
 import shlex
 import shutil
+import stat
 import subprocess
 import sys
 import time
@@ -20,7 +21,7 @@
 CODEX_PROFILE = "freetoken-launch"
 CODEX_PROVIDER_NAME = "FreeToken"
 CODEX_CATALOG_NAME = "freetoken-model.json"
-CODEX_PROVIDER_API_KEY_ENV = "FREETOKEN_API_KEY"
+API_KEY_ENV = "FREETOKEN_API_KEY"
 # Used when the server reports no context length. Guessing low only costs earlier compaction.
 FALLBACK_CONTEXT_WINDOW = 128_000
 MAX_OUTPUT_TOKENS_CAP = 32_768
@@ -71,6 +72,7 @@ class LaunchContext:
     extra_args: list[str]
     dry_run: bool
     assume_yes: bool = False
+    api_key: str | None = None
 
 
 @dataclass(frozen=True)
@@ -167,9 +169,15 @@ def resolve_server_url(server: str | None) -> ServerURL:
     return ServerURL(origin=origin, openai_base_url=f"{origin}/v1")
 
 
-def _get_json(url: str) -> object:
-    request = Request(url, headers={"Accept": "application/json"})
-    with urlopen(request, timeout=5) as response:
+def resolve_api_key(api_key: str | None) -> str | None:
+    return api_key or os.environ.get(API_KEY_ENV) or None
+
+
+def _get_json(url: str, api_key: str | None) -> object:
+    headers = {"Accept": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    with urlopen(Request(url, headers=headers), timeout=5) as response:
         return json.loads(response.read().decode("utf-8"))
 
 
@@ -177,10 +185,10 @@ def _positive_int(value: object) -> int | None:
     return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None
 
 
-def _stats_model(server: ServerURL) -> dict[str, object]:
+def _stats_model(server: ServerURL, api_key: str | None) -> dict[str, object]:
     """The /v1/stats model card; empty when the server predates it."""
     try:
-        payload = _get_json(f"{server.origin}/v1/stats")
+        payload = _get_json(f"{server.origin}/v1/stats", api_key)
     except (HTTPError, URLError, HTTPException, OSError, TimeoutError, ValueError):
         return {}
     model = payload.get("model") if isinstance(payload, dict) else None
@@ -193,11 +201,16 @@ def _input_modalities(stats_model: dict[str, object]) -> tuple[str, ...]:
     return ("text", *extra)
 
 
-def discover_server_model(server: ServerURL) -> ServedModel:
+def discover_server_model(server: ServerURL, api_key: str | None) -> ServedModel:
     try:
-        _get_json(server.openai_base_url)
-        payload = _get_json(f"{server.openai_base_url}/models")
+        _get_json(server.openai_base_url, api_key)
+        payload = _get_json(f"{server.openai_base_url}/models", api_key)
     except (HTTPError, URLError, HTTPException, OSError, TimeoutError) as exc:
+        if isinstance(exc, HTTPError) and exc.code == 401:
+            raise RuntimeError(
+                f"FreeToken server at {server.origin} needs a valid API key: "
+                f"pass --api-key or set {API_KEY_ENV}"
+            ) from exc
         raise RuntimeError(f"Cannot connect to FreeToken server at {server.origin}: {exc}") from exc
 
     if not isinstance(payload, dict):
@@ -220,7 +233,7 @@ def discover_server_model(server: ServerURL) -> ServedModel:
     if not models:
         raise RuntimeError(f"FreeToken server at {server.origin} reported no models")
 
-    stats_model = _stats_model(server)
+    stats_model = _stats_model(server, api_key)
     if context_length is None:
         context_length = _positive_int(stats_model.get("ctx"))
 
@@ -236,15 +249,22 @@ def _backup_path(path: Path) -> Path:
     return path.with_name(f"{path.name}.{time.time_ns()}.bak")
 
 
-def _write_text_with_backup(path: Path, text: str) -> None:
+def _write_text_with_backup(path: Path, text: str, *, private: bool = False) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     if path.exists() and path.read_text() != text:
-        shutil.copy2(path, _backup_path(path))
+        backup = _backup_path(path)
+        # copy2 sets the mode only after it copies the content, so create the backup with the source mode first
+        backup.touch(mode=stat.S_IMODE(path.stat().st_mode))
+        shutil.copy2(path, backup)
+    if private:
+        # before the write, so the key never sits in a file others can read
+        path.touch(mode=0o600)
+        path.chmod(0o600)
     path.write_text(text)
 
 
-def _write_json_with_backup(path: Path, value: object) -> None:
-    _write_text_with_backup(path, json.dumps(value, indent=2) + "\n")
+def _write_json_with_backup(path: Path, value: object, *, private: bool = False) -> None:
+    _write_text_with_backup(path, json.dumps(value, indent=2) + "\n", private=private)
 
 
 def _toml_string(value: str) -> str:
@@ -371,7 +391,7 @@ def _codex_profile_text(ctx: LaunchContext, catalog_path: Path) -> str:
         + f"name = {_toml_string(CODEX_PROVIDER_NAME)}\n"
         + f"base_url = {_toml_string(ctx.server.openai_base_url)}\n"
         + 'wire_api = "responses"\n'
-        + f"env_key = {_toml_string(CODEX_PROVIDER_API_KEY_ENV)}\n"
+        + f"env_key = {_toml_string(API_KEY_ENV)}\n"
     )
 
 
@@ -408,7 +428,7 @@ def prepare_codex(ctx: LaunchContext) -> CommandSpec:
     ]
     return CommandSpec(
         argv=argv,
-        env={CODEX_PROVIDER_API_KEY_ENV: "freetoken"},
+        env={API_KEY_ENV: ctx.api_key or "freetoken"},
         unset_env=CODEX_CLEAR_ENV,
     )
 
@@ -423,7 +443,7 @@ def prepare_claude(ctx: LaunchContext) -> CommandSpec:
             "CLAUDE_CODE_MAX_OUTPUT_TOKENS": str(_max_output_tokens(ctx)),
             "ANTHROPIC_BASE_URL": ctx.server.origin,
             "ANTHROPIC_API_KEY": "",
-            "ANTHROPIC_AUTH_TOKEN": "freetoken",
+            "ANTHROPIC_AUTH_TOKEN": ctx.api_key or "freetoken",
             "ANTHROPIC_MODEL": model_id,
             "ANTHROPIC_DEFAULT_OPUS_MODEL": model_id,
             "ANTHROPIC_DEFAULT_SONNET_MODEL": model_id,
@@ -456,13 +476,16 @@ def _opencode_model_entries(
 
 def _opencode_config(ctx: LaunchContext) -> str:
     model_ids = _ordered_model_ids(ctx)
+    options = {"baseURL": ctx.server.openai_base_url}
+    if ctx.api_key:
+        options["apiKey"] = ctx.api_key
     config = {
         "$schema": "https://opencode.ai/config.json",
         "provider": {
             OPENCODE_PROVIDER: {
                 "npm": "@ai-sdk/openai-compatible",
                 "name": OPENCODE_PROVIDER_NAME,
-                "options": {"baseURL": ctx.server.openai_base_url},
+                "options": options,
                 "models": _opencode_model_entries(
                     model_ids, _context_window(ctx), _max_output_tokens(ctx), _accepts_images(ctx)
                 ),
@@ -614,7 +637,7 @@ def _patch_openclaw_config(
         new_models.append(entry)
 
     provider["baseUrl"] = ctx.server.openai_base_url
-    provider["apiKey"] = "freetoken-local"
+    provider["apiKey"] = ctx.api_key or "freetoken-local"
     provider["api"] = "openai-completions"
     provider["models"] = new_models
     providers[OPENCLAW_PROVIDER] = provider
@@ -661,7 +684,9 @@ def prepare_openclaw(ctx: LaunchContext) -> CommandSpec:
         if not _openclaw_has_freetoken_provider(config):
             if not ctx.assume_yes and not _confirm_openclaw_first_patch(ctx, config_path):
                 raise RuntimeError("OpenClaw config patch cancelled")
-        _write_json_with_backup(config_path, _patch_openclaw_config(ctx, config))
+        _write_json_with_backup(
+            config_path, _patch_openclaw_config(ctx, config), private=bool(ctx.api_key)
+        )
         _clear_openclaw_session_overrides(ctx.model.model_id)
 
     argv = ["openclaw", *ctx.extra_args] if ctx.extra_args else ["openclaw", "chat"]
@@ -697,8 +722,8 @@ def prepare_hermes(ctx: LaunchContext) -> CommandSpec:
         model_section["default"] = ctx.model.model_id
         model_section["provider"] = "custom"
         model_section["base_url"] = ctx.server.openai_base_url
-        # A non-empty dummy: FreeToken is unauthenticated, but some clients reject an empty key.
-        model_section["api_key"] = HERMES_API_KEY
+        # Dummy for a server without --api-key; some clients reject an empty key.
+        model_section["api_key"] = ctx.api_key or HERMES_API_KEY
         # Hermes' name for the window (prompt + generation). Left unset it auto-detects, and a
         # miss lands on a 32k default that trips the floor below.
         window = _context_window(ctx)
@@ -712,7 +737,9 @@ def prepare_hermes(ctx: LaunchContext) -> CommandSpec:
                 file=sys.stderr,
             )
 
-        _write_text_with_backup(config_path, yaml.safe_dump(config, sort_keys=False))
+        _write_text_with_backup(
+            config_path, yaml.safe_dump(config, sort_keys=False), private=bool(ctx.api_key)
+        )
 
     argv = ["hermes", *ctx.extra_args] if ctx.extra_args else ["hermes", "chat"]
     return CommandSpec(argv=argv, env={})
@@ -753,8 +780,7 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec:
     ``reasoning_content``, while llm-pi-ai's JSON round-trip can change argument
     values and key order — drift that survives render canonicalization, cuts
     the prefix cache, and shows the model a rewrite of its own output. The
-    dummy DEEPSEEK_API_KEY satisfies dsh's non-empty key requirement; FreeToken
-    itself is unauthenticated."""
+    dummy DEEPSEEK_API_KEY satisfies dsh's non-empty key requirement."""
     settings_path = _dsh_home() / DSH_LAUNCH_SETTINGS_NAME
     patch_path = _dsh_home() / DSH_LAUNCH_PATCH_NAME
     if not ctx.dry_run:
@@ -774,7 +800,7 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec:
 
         window, output = _context_window(ctx), _max_output_tokens(ctx)
         modalities = ["text", "image"] if _accepts_images(ctx) else ["text"]
-        config["llm-deepseek"] = {
+        section: dict[str, object] = {
             "baseURL": ctx.server.openai_base_url,
             "models": [
                 {
@@ -787,6 +813,10 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec:
                 for model_id in _ordered_model_ids(ctx)
             ],
         }
+        if ctx.api_key:
+            # not DEEPSEEK_API_KEY: dsh's web search also reads it and sends it to api.deepseek.com
+            section["apiKeyEnv"] = API_KEY_ENV
+        config["llm-deepseek"] = section
         config["agent-default-model"] = {
             "provider": "deepseek-official",
             "model": ctx.model.model_id,
@@ -804,15 +834,14 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec:
     elif len(app_args) >= 2 and app_args[0] == "--profile":
         profile, app_args = app_args[1], app_args[2:]
     argv = ["dsh", "--profile", profile, "--patch", str(patch_path), *app_args]
-    return CommandSpec(
-        argv=argv,
-        env={
-            "DEEPSEEK_BASE_URL": ctx.server.openai_base_url,
-            "DEEPSEEK_API_KEY": DSH_API_KEY,
-            "DSH_TELEMETRY_DISABLED": "1",
-        },
-        unset_env=CLOUD_PROVIDER_API_KEY_ENV,
-    )
+    env = {
+        "DEEPSEEK_BASE_URL": ctx.server.openai_base_url,
+        "DEEPSEEK_API_KEY": DSH_API_KEY,
+        "DSH_TELEMETRY_DISABLED": "1",
+    }
+    if ctx.api_key:
+        env[API_KEY_ENV] = ctx.api_key
+    return CommandSpec(argv=argv, env=env, unset_env=CLOUD_PROVIDER_API_KEY_ENV)
 
 
 PREPARERS = {
@@ -840,6 +869,10 @@ def parse_argv(argv: list[str], prog: str = "python -m freetoken.launch") -> arg
     )
     parser.add_argument("agent", choices=sorted(PREPARERS))
     parser.add_argument("--server", help="FreeToken server origin or /v1 URL")
+    parser.add_argument(
+        "--api-key",
+        help=f"API key of a server started with --api-key (default: ${API_KEY_ENV})",
+    )
     parser.add_argument(
         "--dry-run",
         action="store_true",
@@ -1028,13 +1061,16 @@ def main(argv: list[str] | None = None, prog: str = "python -m freetoken.launch"
             return 0
 
         server = resolve_server_url(args.server)
-        model = discover_server_model(server)
+        api_key = resolve_api_key(args.api_key)
+        model = discover_server_model(server, api_key)
         plan_ctx = LaunchContext(
             server=server,
             model=model,
             extra_args=args.extra_args,
             dry_run=True,
             assume_yes=args.yes,
+            # the plan is printed by --dry-run, so it never holds the real key
+            api_key="<api-key>" if api_key else None,
         )
         spec = PREPARERS[args.agent](plan_ctx)
         if args.dry_run:
@@ -1052,6 +1088,7 @@ def main(argv: list[str] | None = None, prog: str = "python -m freetoken.launch"
             extra_args=args.extra_args,
             dry_run=False,
             assume_yes=args.yes,
+            api_key=api_key,
         )
         spec = PREPARERS[args.agent](ctx)
         if binary_path != spec.argv[0]:
diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 34221cfc70..a9b16adfd7 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -2,6 +2,8 @@
 
 import asyncio
 import contextlib
+import hashlib
+import hmac
 import json
 import os
 import signal
@@ -34,6 +36,7 @@
     load_generation_sampling,
 )
 from pydantic import BaseModel
+from starlette.datastructures import Headers
 
 from .args import ServerArgs
 from .anthropic_api import register_anthropic_routes
@@ -417,6 +420,39 @@ def install_cors(app: FastAPI, origins_csv: str) -> None:
     )
 
 
+class _ApiKeyMiddleware:
+    def __init__(self, app, api_key: str) -> None:
+        self.app = app
+        self._digest = hashlib.sha256(api_key.encode()).digest()
+
+    def _matches(self, value: str | None) -> bool:
+        # compare digests so the time taken does not depend on the key length
+        return value is not None and hmac.compare_digest(
+            hashlib.sha256(value.encode()).digest(), self._digest
+        )
+
+    async def __call__(self, scope, receive, send) -> None:
+        if scope["type"] not in ("http", "websocket") or scope["path"] == "/health":
+            return await self.app(scope, receive, send)
+        headers = Headers(scope=scope)
+        scheme, _, token = headers.get("authorization", "").partition(" ")
+        if (scheme.lower() == "bearer" and self._matches(token)) or self._matches(
+            headers.get("x-api-key")
+        ):
+            return await self.app(scope, receive, send)
+        await JSONResponse({"error": "Unauthorized"}, status_code=401)(scope, receive, send)
+
+
+def install_api_key(app: FastAPI, api_key: str | None) -> None:
+    """Require ``api_key`` on every route but /health.
+
+    Run it before install_cors, so CORS answers preflights and adds its headers to 401s."""
+    if api_key is None:
+        return
+    app.add_middleware(_ApiKeyMiddleware, api_key=api_key)
+    logger.info("API key required on every route except /health")
+
+
 app = FastAPI(title="FreeToken API Server", version=__version__, lifespan=lifespan)
 register_openai_routes(app, get_global_state, lambda: _MODEL_SAMPLING)
 register_anthropic_routes(app, get_global_state, lambda: _MODEL_SAMPLING)
@@ -883,7 +919,7 @@ def _flag_shutdown(signum, frame) -> None:
         signal.signal(sig, _flag_shutdown)
 
 
-def _serve_and_run_shell(host: str, port: int) -> None:
+def _serve_and_run_shell(host: str, port: int, api_key: str | None) -> None:
     """Shell mode: serve the API here, and attach the terminal client to it over the loopback.
 
     The shell is an ordinary API client (see ``freetoken.shell``), so shell mode is just
@@ -916,7 +952,7 @@ def _serve_and_run_shell(host: str, port: int) -> None:
         # /health and echoes the same load progress the desktop app polls for. A ^C during that
         # wait is a stop, not a crash -- exit through the teardown below, not a traceback.
         with contextlib.suppress(KeyboardInterrupt):
-            asyncio.run(run_shell(origin, connect_grace=30.0))
+            asyncio.run(run_shell(origin, connect_grace=30.0, api_key=api_key))
     finally:
         server.should_exit = True
         thread.join(timeout=15)
@@ -958,6 +994,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s
     host = config.server_host
     port = config.server_port
 
+    install_api_key(app, config.api_key)
     # Create/validate FREETOKEN_API_LOG_DIR and start the writer thread up front, so a
     # bad path is reported at boot rather than silently on the first request.
     install_cors(app, config.cors_origins)
@@ -1044,7 +1081,7 @@ def _on_meta(meta: dict) -> None:
     ).start()
 
     if run_shell:
-        _serve_and_run_shell(host, port)
+        _serve_and_run_shell(host, port, config.api_key)
         return
     # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group.
     uvicorn.run(app, host=host, port=port)
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index b8d06018d8..5e6caa759a 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -3,7 +3,7 @@
 import argparse
 import json
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import List, Tuple
 
 import torch
@@ -38,6 +38,8 @@ def _nvfp4_entry(value: str) -> str:
 class ServerArgs(SchedulerConfig):
     server_host: str = "127.0.0.1"
     server_port: int = 1919
+    # kept out of repr: parse_args logs the whole ServerArgs
+    api_key: str | None = field(default=None, repr=False)
     num_tokenizer: int = 0
     silent_output: bool = False
     # The terminal shell is attached to this server (ft shell --model / ft serve --shell-mode).
@@ -368,6 +370,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="The port number for the server to listen on.",
     )
 
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=ServerArgs.api_key,
+        help="Require this key on every route except /health, sent as "
+        "'Authorization: Bearer <key>' or 'x-api-key: <key>'. Defaults to $FREETOKEN_API_KEY; "
+        "unset serves without authentication.",
+    )
+
     parser.add_argument(
         "--cuda-graph-max-bs",
         "--graph",
@@ -842,6 +853,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         if entry:
             kwargs["quant_backend"] = entry
 
+    # an empty flag usually means an unset shell variable; refuse it rather than serve open
+    if kwargs["api_key"] == "":
+        parser.error("--api-key is empty; omit it to serve without authentication")
+    from freetoken.launch import resolve_api_key
+
+    kwargs["api_key"] = resolve_api_key(kwargs["api_key"])
+
     if kwargs["model_path"].startswith("~"):
         kwargs["model_path"] = os.path.expanduser(kwargs["model_path"])
 
diff --git a/python/freetoken/shell/__init__.py b/python/freetoken/shell/__init__.py
index e429f487a0..d31eed4cd5 100644
--- a/python/freetoken/shell/__init__.py
+++ b/python/freetoken/shell/__init__.py
@@ -45,6 +45,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser:
         default=None,
         help="FreeToken server URL (default: $FREETOKEN_HOST, else http://127.0.0.1:1919)",
     )
+    parser.add_argument(
+        "--api-key",
+        help="API key of a server started with --api-key (default: $FREETOKEN_API_KEY)",
+    )
     return parser
 
 
@@ -63,7 +67,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft shell") -> int:
     except SystemExit as exc:
         return int(exc.code) if isinstance(exc.code, int) else 2
 
-    from freetoken.launch import resolve_server_url
+    from freetoken.launch import resolve_api_key, resolve_server_url
 
     try:
         server = resolve_server_url(parsed.server)
@@ -74,7 +78,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft shell") -> int:
     from .tui import run_shell
 
     try:
-        return asyncio.run(run_shell(server.origin))
+        return asyncio.run(run_shell(server.origin, api_key=resolve_api_key(parsed.api_key)))
     except KeyboardInterrupt:
         return 130
 
diff --git a/python/freetoken/shell/client.py b/python/freetoken/shell/client.py
index 183fa4dff2..d111536294 100644
--- a/python/freetoken/shell/client.py
+++ b/python/freetoken/shell/client.py
@@ -27,7 +27,7 @@
 import openai
 from openai import AsyncOpenAI
 
-# Any string is accepted by the server; sent so the SDK doesn't refuse to build a request.
+# Sent when no key is given: a server without --api-key accepts any string, and the SDK needs one.
 LOCAL_API_KEY = "freetoken-local"
 CONTROL_TIMEOUT = 10.0
 # Per-read timeout on the chat stream. Generous: it has to cover the gap before the first
@@ -97,14 +97,15 @@ def __init__(
         self,
         origin: str,
         *,
-        api_key: str = LOCAL_API_KEY,
+        api_key: str | None = None,
         timeout: float = CONTROL_TIMEOUT,
     ) -> None:
         self.origin = origin.rstrip("/")
         self.timeout = timeout
+        self._api_key = api_key or LOCAL_API_KEY
         self._openai = AsyncOpenAI(
             base_url=f"{self.origin}/v1",
-            api_key=api_key,
+            api_key=self._api_key,
             max_retries=0,  # a retried chat request would generate the turn twice
             timeout=openai.Timeout(CHAT_READ_TIMEOUT, connect=10.0),
         )
@@ -118,7 +119,7 @@ def _request_json_blocking(
         self, method: str, path: str, body: dict[str, Any] | None, timeout: float
     ) -> dict[str, Any]:
         data = None
-        headers = {"Accept": "application/json"}
+        headers = {"Accept": "application/json", "Authorization": f"Bearer {self._api_key}"}
         if body is not None:
             data = json.dumps(body).encode("utf-8")
             headers["Content-Type"] = "application/json"
diff --git a/python/freetoken/shell/tui.py b/python/freetoken/shell/tui.py
index 93bd18cc33..72d7ebe0e0 100644
--- a/python/freetoken/shell/tui.py
+++ b/python/freetoken/shell/tui.py
@@ -515,13 +515,15 @@ def _format_load_progress(doc: dict) -> str:
     return f"loading ({phase})..."
 
 
-async def run_shell(origin: str, *, connect_grace: float = 0.0) -> int:
+async def run_shell(
+    origin: str, *, connect_grace: float = 0.0, api_key: str | None = None
+) -> int:
     """Attach to the FreeToken server at ``origin`` and run the terminal chat.
 
     ``connect_grace`` is how long to keep retrying a refused connection before giving up --
     left at 0 when attaching to a server the user says is already running, raised when the
     caller just started one in this process (see ``server/api_server.py``)."""
-    client = ShellClient(origin)
+    client = ShellClient(origin, api_key=api_key)
     try:
         return await _run_shell(client, origin, connect_grace=connect_grace)
     finally:
diff --git a/tests/daemon/test_daemon_serve_manager.py b/tests/daemon/test_daemon_serve_manager.py
index fd9956ebba..cfd84ae6e6 100644
--- a/tests/daemon/test_daemon_serve_manager.py
+++ b/tests/daemon/test_daemon_serve_manager.py
@@ -1,19 +1,25 @@
 from __future__ import annotations
 
+import os
 import signal
+import sys
 import threading
 import time
+from concurrent.futures import ThreadPoolExecutor
 
 import pytest
+from fastapi.testclient import TestClient
 
+from freetoken.daemon import proxy
 from freetoken.daemon.accounting import (
     AccountingOutbox,
     AccountingOutboxError,
     AccountingPrepareError,
 )
+from freetoken.daemon.app import build_app
 from freetoken.daemon.logring import LogRing
 from freetoken.daemon.pidfile import ServeState, ServeStateStore
-from freetoken.daemon.serve_manager import Conflict, ExitInfo, ServeManager
+from freetoken.daemon.serve_manager import Conflict, ExitInfo, ServeManager, spawn_serve
 
 
 # --------------------------------------------------------------------------- test doubles
@@ -50,11 +56,13 @@ class Spawner:
     def __init__(self):
         self.children: list[FakeChild] = []
         self.calls: list[tuple] = []
+        self.api_keys: list[str | None] = []
         self._pid = 1000
         self.gate: threading.Event | None = None
 
-    def __call__(self, model, port, args) -> FakeChild:
+    def __call__(self, model, port, args, api_key=None) -> FakeChild:
         self.calls.append((model, port, list(args)))
+        self.api_keys.append(api_key)
         if self.gate is not None:
             self.gate.wait()
         self._pid += 1
@@ -743,6 +751,131 @@ def test_readopt_no_state_is_noop(tmp_path):
     assert mgr.readopt() is False
 
 
+# --------------------------------------------------------------------------- API key
+
+
+def test_api_key_is_serve_config_kept_owner_only_and_readopted(tmp_path):
+    sp = Spawner()
+    mgr, store, _ = make_manager(tmp_path, sp)
+    # an older daemon left a world-readable state file and a stale tmp behind
+    for path in (store.path, f"{store.path}.tmp"):
+        with open(path, "w") as fh:
+            fh.write("{}")
+        os.chmod(path, 0o644)
+    mgr.start("m", 1919, [], "k1")
+    assert sp.api_keys == ["k1"] and mgr.serve_api_key() == "k1"
+    assert mgr.start("m", 1919, [], "k1")["idempotent"] is True
+    with pytest.raises(Conflict):
+        mgr.start("m", 1919, [], "k2")
+    assert store.load().api_key == "k1"
+    assert os.stat(store.path).st_mode & 0o777 == 0o600
+
+    adopted = FakeChild(4242, adopted=True)
+    fresh = ServeManager(
+        LogRing(), store, spawn_fn=Spawner(), adopt_fn=lambda state: adopted, tailer_factory=None,
+        signal_fn=lambda p, s: None, apply_oom=False,
+    )
+    assert fresh.readopt() is True and fresh.serve_api_key() == "k1"
+    adopted.die()
+    sp.children[0].die()
+
+
+def test_switch_replaces_the_api_key(tmp_path):
+    sp = Spawner()
+
+    def sig(pid, s):
+        if s == signal.SIGTERM:
+            sp.by_pid(pid).die(0)
+
+    mgr, store, _ = make_manager(tmp_path, sp, signal_fn=sig)
+    mgr.start("m", 1919, [], "k1")
+    mgr.switch("m", 1919, [], False, "k2")
+    assert sp.api_keys == ["k1", "k2"]
+    assert mgr.serve_api_key() == "k2" and store.load().api_key == "k2"
+    sp.children[1].die()
+
+
+def test_auto_restart_keeps_the_api_key(tmp_path):
+    sp = Spawner()
+    mgr, _, _ = make_manager(tmp_path, sp, auto_restart=True)
+    mgr.start("m", 1919, [], "k1")
+    sp.children[0].die(1)
+    assert wait_until(lambda: len(sp.children) == 2)
+    assert sp.api_keys == ["k1", "k1"]
+    sp.children[1].die()
+
+
+def test_spawned_serve_sees_only_the_requested_key(tmp_path, monkeypatch):
+    monkeypatch.setenv("FREETOKEN_API_KEY", "from-daemon-env")
+    code = "import os; print(os.environ.get('FREETOKEN_API_KEY'))"
+    for api_key in (None, "k1"):
+        log = tmp_path / f"{api_key}.log"
+        child = spawn_serve([sys.executable, "-c", code], str(log), api_key)
+        child.wait()
+        child.close()
+        assert log.read_text().strip() == str(api_key)
+
+
+def test_probe_sends_the_running_serves_key(monkeypatch):
+    seen = []
+
+    class Response:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+        def read(self):
+            return b"{}"
+
+    def fake_urlopen(req, timeout):
+        seen.append(req.get_header("Authorization"))
+        return Response()
+
+    monkeypatch.setattr(proxy.urllib.request, "urlopen", fake_urlopen)
+    key = None
+    probe = proxy.ServeProbe(api_key=lambda: key)
+    probe.fresh_stats(1919)
+    key = "k1"
+    probe.fresh_stats(1919)
+    probe.prepare_stop(1919)
+    assert seen == [None, "Bearer k1", "Bearer k1"]
+
+
+def test_start_and_switch_routes_hand_the_api_key_to_the_manager():
+    calls = []
+
+    class Manager:
+        def status(self):
+            return {"running": False}
+
+        def start(self, *args):
+            calls.append(args)
+            return {"pid": 1, "idempotent": False}
+
+        def switch(self, *args):
+            calls.append(args)
+            return {"pid": 1, "idempotent": False, "accounting": None}
+
+    with ThreadPoolExecutor(1) as lifecycle_pool, ThreadPoolExecutor(1) as proxy_pool:
+        app = build_app(
+            manager=Manager(), ring=LogRing(), probe=None, footprint_fn=lambda pid: {},
+            lifecycle_pool=lifecycle_pool, proxy_pool=proxy_pool,
+        )
+        with TestClient(app) as client:
+            client.post("/engine/start", json={"model": "m", "port": 1919, "apiKey": "k1"})
+            client.post("/engine/start", json={"model": "m", "port": 1919, "apiKey": ""})
+            client.post("/engine/start", json={"model": "m", "port": 1919})
+            client.post("/engine/switch", json={"model": "m", "port": 1919, "apiKey": "k2"})
+            for args in (["--api-key", "k3"], ["--api-key=k3"], ["--api", "k3"], ["--ap=k3"]):
+                rejected = client.post("/engine/start", json={"model": "m", "args": args})
+                assert rejected.status_code == 400
+            allowed = client.post("/engine/start", json={"model": "m", "args": ["--attn", "fa"]})
+            assert allowed.status_code == 200
+    assert [args[-1] for args in calls] == ["k1", None, None, "k2", None]
+
+
 # --------------------------------------------------------------------------- resource leak (L2)
 
 
diff --git a/tests/server/test_api_key.py b/tests/server/test_api_key.py
new file mode 100644
index 0000000000..a655c68831
--- /dev/null
+++ b/tests/server/test_api_key.py
@@ -0,0 +1,146 @@
+"""--api-key: the gate in front of every route but /health, and where the key comes from."""
+
+from __future__ import annotations
+
+import json
+import os
+import stat
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import yaml
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from freetoken import launch
+from freetoken.server.api_server import install_api_key, install_cors
+from freetoken.server.args import parse_args
+
+KEY = "sk-test"
+ORIGIN = "http://localhost:1420"
+
+
+def _client(api_key: str | None) -> TestClient:
+    app = FastAPI()
+    app.get("/health")(lambda: {"status": "ok"})
+    app.get("/v1/models")(lambda: {"data": []})
+    install_api_key(app, api_key)
+    install_cors(app, ORIGIN)
+    return TestClient(app)
+
+
+def _parse(*argv: str):
+    hf = SimpleNamespace(to_dict=lambda: {"architectures": ["Qwen3ForCausalLM"], "torch_dtype": "bfloat16"})
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: hf):
+        return parse_args(["--model", "/models/anon", *argv])[0]
+
+
+@pytest.mark.parametrize(
+    "headers, status",
+    [
+        ({}, 401),
+        ({"Authorization": "Bearer wrong"}, 401),
+        ({"Authorization": f"Basic {KEY}"}, 401),
+        ({"Authorization": KEY}, 401),
+        ({"Authorization": b"Bearer caf\xe9"}, 401),
+        ({"x-api-key": "wrong"}, 401),
+        ({"Authorization": f"Bearer {KEY}"}, 200),
+        ({"Authorization": f"bearer {KEY}"}, 200),
+        ({"x-api-key": KEY}, 200),
+    ],
+)
+def test_key_gates_every_route_but_health(headers, status):
+    client = _client(KEY)
+    response = client.get("/v1/models", headers=headers)
+    assert response.status_code == status
+    if status == 401:
+        assert response.json() == {"error": "Unauthorized"}
+    assert client.get("/health").status_code == 200
+
+
+def test_no_key_serves_open():
+    assert _client(None).get("/v1/models").status_code == 200
+
+
+def test_cors_preflight_passes_and_401_carries_cors_headers():
+    client = _client(KEY)
+    preflight = client.options(
+        "/v1/models",
+        headers={"Origin": ORIGIN, "Access-Control-Request-Method": "GET",
+                 "Access-Control-Request-Headers": "authorization"},
+    )
+    assert preflight.status_code == 200
+    denied = client.get("/v1/models", headers={"Origin": ORIGIN})
+    assert denied.status_code == 401
+    assert denied.headers["access-control-allow-origin"] == ORIGIN
+
+
+def test_options_outside_a_cors_preflight_needs_the_key():
+    assert _client(KEY).options("/v1/models").status_code == 401
+
+
+def test_flag_wins_over_env_and_empty_values_do_not_enable_auth(monkeypatch):
+    monkeypatch.delenv("FREETOKEN_API_KEY", raising=False)
+    assert _parse().api_key is None
+    monkeypatch.setenv("FREETOKEN_API_KEY", "from-env")
+    assert _parse().api_key == "from-env"
+    assert _parse("--api-key", KEY).api_key == KEY
+    monkeypatch.setenv("FREETOKEN_API_KEY", "")
+    assert _parse().api_key is None
+    with pytest.raises(SystemExit):
+        _parse("--api-key", "")
+
+
+def test_key_stays_out_of_the_logged_config(monkeypatch):
+    monkeypatch.delenv("FREETOKEN_API_KEY", raising=False)
+    assert KEY not in str(_parse("--api-key", KEY))
+
+
+@pytest.mark.parametrize("agent", sorted(launch.PREPARERS))
+def test_dry_run_prints_a_placeholder_instead_of_the_key(agent, monkeypatch, capsys):
+    key = 'p"w\\d\xe9'
+    model = launch.ServedModel(model_id="m", models=["m"], context_length=8192)
+    monkeypatch.setattr(launch, "discover_server_model", lambda _server, _key: model)
+    assert launch.main([agent, "--dry-run", "--api-key", key]) == 0
+    out = capsys.readouterr().out
+    assert key not in out and json.dumps(key)[1:-1] not in out
+    if agent in ("claude", "codex", "dsh", "opencode"):
+        assert "<api-key>" in out
+
+
+def test_dsh_keeps_the_key_out_of_the_variable_its_web_search_reads(monkeypatch, tmp_path):
+    monkeypatch.setenv("DSH_HOME", str(tmp_path))
+    monkeypatch.setattr(launch, "_warn_dsh_node_version", lambda: None)
+    ctx = launch.LaunchContext(
+        server=launch.resolve_server_url("http://127.0.0.1:1919"),
+        model=launch.ServedModel(model_id="m", models=["m"], context_length=8192),
+        extra_args=[],
+        dry_run=False,
+        api_key=KEY,
+    )
+    env = launch.prepare_dsh(ctx).env
+    settings = yaml.safe_load((tmp_path / launch.DSH_LAUNCH_SETTINGS_NAME).read_text())
+    assert env["DEEPSEEK_API_KEY"] != KEY
+    assert env[settings["llm-deepseek"]["apiKeyEnv"]] == KEY
+
+
+@pytest.mark.parametrize("agent", ["hermes", "openclaw"])
+def test_config_files_holding_the_key_are_owner_only(agent, monkeypatch, tmp_path):
+    monkeypatch.setenv("HOME", str(tmp_path))
+    old_umask = os.umask(0o022)
+    try:
+        for key in ("sk-first", "sk-second"):
+            ctx = launch.LaunchContext(
+                server=launch.resolve_server_url("http://127.0.0.1:1919"),
+                model=launch.ServedModel(model_id="m", models=["m"], context_length=65536),
+                extra_args=[],
+                dry_run=False,
+                assume_yes=True,
+                api_key=key,
+            )
+            launch.PREPARERS[agent](ctx)
+    finally:
+        os.umask(old_umask)
+    written = [p for p in tmp_path.rglob("*") if p.is_file() and "sk-" in p.read_text()]
+    assert len(written) == 2
+    assert all(stat.S_IMODE(p.stat().st_mode) == 0o600 for p in written)

From ee9789ff6cdcc2e9c5e9467aba630d38cd230fa6 Mon Sep 17 00:00:00 2001
From: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
Date: Sat, 10 Oct 2026 04:17:57 -0400
Subject: [PATCH 084/249] fix(kernel): avoid nan in triton sampling softmax
 when a chunk starts with -inf (#654)

---
 python/freetoken/kernel/triton/sampling.py |  4 +++-
 tests/kernels/test_triton_fallbacks.py     | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py
index a970714406..38bc253386 100644
--- a/python/freetoken/kernel/triton/sampling.py
+++ b/python/freetoken/kernel/triton/sampling.py
@@ -92,7 +92,9 @@ def _sm_partial(
         x = tl.load(logits_ptr + base + offs, mask=mask, other=-float("inf")).to(tl.float32) * inv_t
         blk_max = tl.max(x, 0)
         new_m = tl.maximum(m, blk_max)
-        d = d * tl.exp(m - new_m) + tl.sum(tl.exp(x - new_m), 0)
+        # While every logit so far is -inf, shift by 0: exp(-inf - -inf) would be NaN.
+        shift = tl.where(new_m == -float("inf"), 0.0, new_m)
+        d = d * tl.exp(m - shift) + tl.sum(tl.exp(x - shift), 0)
         m = new_m
     tl.store(pm_ptr + pid, m)
     tl.store(pl_ptr + pid, d)
diff --git a/tests/kernels/test_triton_fallbacks.py b/tests/kernels/test_triton_fallbacks.py
index 45234e66d4..2306a40044 100644
--- a/tests/kernels/test_triton_fallbacks.py
+++ b/tests/kernels/test_triton_fallbacks.py
@@ -55,3 +55,22 @@ def test_norm_fallback_matches_torch(dtype, kind):
 
     rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-3)
     torch.testing.assert_close(norm.forward(x), expected.to(dtype), rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("mask", ["prefix", "allowed_ids"])
+def test_sampling_softmax_fallback_with_masked_logits(mask):
+    from freetoken.kernel.triton.sampling import softmax
+
+    torch.manual_seed(0)
+    logits = torch.randn((3, 151936), device="cuda") * 4
+    if mask == "prefix":
+        # Long enough that some column chunk starts with an all -inf block.
+        logits[:2, :16384] = float("-inf")
+    else:
+        allowed = torch.tensor([5, 70000, 151000], device="cuda")
+        logits[:2] = torch.full_like(logits[:2], float("-inf")).index_copy(
+            1, allowed, logits[:2, allowed])
+    temperature = torch.tensor([1.0, 0.7, 1.3], device="cuda")
+    expected = torch.softmax(logits / temperature[:, None], dim=-1)
+
+    torch.testing.assert_close(softmax(logits, temperature), expected, rtol=1e-5, atol=1e-7)

From 74a387777708b9650a027494bcb37d2ab9cf2443 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 10 Oct 2026 13:30:22 -0700
Subject: [PATCH 085/249] build(deps): move to torch 2.14.1, triton 3.8 and
 sglang-kernel 0.4.9 (#630)

---
 freetoken-kernel-cache/pyproject.toml |  4 ++--
 install.sh                            |  2 +-
 pyproject.toml                        | 18 +++++++++---------
 scripts/ci/manylinux-build.sh         |  4 ++--
 setup.py                              |  6 +++---
 5 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/freetoken-kernel-cache/pyproject.toml b/freetoken-kernel-cache/pyproject.toml
index ec37efbfa6..bc75e1b2f7 100644
--- a/freetoken-kernel-cache/pyproject.toml
+++ b/freetoken-kernel-cache/pyproject.toml
@@ -8,7 +8,7 @@ requires = [
     # 77 is the first release that understands the PEP 639 `license` SPDX string below.
     "setuptools>=77",
     "wheel",
-    "torch>=2.11,<2.12",
+    "torch>=2.14,<2.15",
     "apache-tvm-ffi==0.1.13.post3",
 ]
 build-backend = "build_backend"
@@ -22,7 +22,7 @@ requires-python = ">=3.10"
 license = "Apache-2.0"
 dependencies = ["apache-tvm-ffi==0.1.13.post3"]
 
-# Provenance pin: PyPI's torch 2.11.0 is itself the cu130 build, but this index serves
+# Provenance pin: PyPI's torch 2.14.1 is itself the cu130 build, but this index serves
 # ONLY cu130 wheels, so an isolated build can never resolve a different-CUDA torch and
 # stamp the wheel +cu126/+cu128. uv applies these to build-system.requires when this
 # directory is the build root. Mirrors ../pyproject.toml.
diff --git a/install.sh b/install.sh
index f623ffece7..68738cc695 100755
--- a/install.sh
+++ b/install.sh
@@ -206,7 +206,7 @@ mkdir -p "$FT_HOME"
 # re-install can't inherit a stale/mismatched torch (e.g. an old cu128 venv after a cu130 bump).
 "$UV" venv "$VENV" --python "$PY_VERSION" --clear
 
-# PyPI's torch 2.11.0 is the same cu130 build the pytorch index
+# PyPI's torch 2.14.1 is the same cu130 build the pytorch index
 # serves; the explicit index pins provenance to the cu130 channel. `unsafe-best-match`
 # is needed because the pytorch index also mirrors stale copies of common deps (e.g.
 # packaging<=24.1) that would shadow PyPI under uv's first-index strategy; all indexes
diff --git a/pyproject.toml b/pyproject.toml
index 8d00d2bbc6..95b872f5b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@
 # own, and a mismatch links the C++ extensions against the wrong libtorch.
 # setuptools floor: 77 is the first release that understands the PEP 639 `license`
 # SPDX string and `license-files` below.
-requires = ["setuptools>=77", "torch>=2.11,<2.12", "wheel"]
+requires = ["setuptools>=77", "torch>=2.14,<2.15", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -53,17 +53,17 @@ dependencies = [
     "pydantic>=2.9,<3",
     "pyzmq>=27,<28",
     "safetensors>=0.6,<1",
-    # floor+ceiling: sglang-kernel 0.4.5 links libtorch symbols only 2.11 has.
-    # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves
+    # floor+ceiling: sglang-kernel 0.4.9 requires torch==2.14.1.
+    # PyPI's torch 2.14.1 wheel is itself the cu130 build, so plain pip resolves
     # correctly from PyPI alone; uv additionally pins the index below.
-    "torch>=2.11,<2.12",
-    "torchvision>=0.26,<0.27",
+    "torch>=2.14,<2.15",
+    "torchvision>=0.29,<0.30",
     "tqdm>=4.66,<5",
     "transformers>=5.16,<5.17",
-    # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its
+    # CUDA torch 2.14 resolves Triton 3.8; AMD's ROCm 7.14 image supplies its
     # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
     # runtime-specific wheel selected by the PyTorch distribution.
-    "triton>=3.6,<3.8; platform_system == 'Linux'",
+    "triton>=3.6,<3.9; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 
@@ -82,7 +82,7 @@ dev = ["pytest>=6.0"]
 # `freetoken[accel]` for the full native fast path.
 fi = ["flashinfer-python[cu13]==0.6.18.post1"]
 # renamed from sgl-kernel at 0.4; still imports as `sgl_kernel`, so never co-install both
-sgl = ["sglang-kernel==0.4.5"]
+sgl = ["sglang-kernel==0.4.9"]
 accel = ["freetoken[fi,sgl]"]
 # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel
 # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the
@@ -91,7 +91,7 @@ accel = ["freetoken[fi,sgl]"]
 # in a dedicated environment when you specifically need the Marlin path.
 
 # uv-only provenance pins — pip ignores this section and resolves from PyPI, which
-# works: PyPI's torch 2.11.0 is the same cu130 build this index serves.
+# works: PyPI's torch 2.14.1 is the same cu130 build this index serves.
 # `explicit = true` scopes the index to the one package that needs it
 # (the torch index also mirrors stale copies of common deps, e.g. packaging<=24.1,
 # which would otherwise shadow PyPI under uv's first-index strategy).
diff --git a/scripts/ci/manylinux-build.sh b/scripts/ci/manylinux-build.sh
index d2dd3ca2f3..5fe0487fed 100755
--- a/scripts/ci/manylinux-build.sh
+++ b/scripts/ci/manylinux-build.sh
@@ -97,11 +97,11 @@ for cptag in $MATRIX; do
   VENV="/tmp/build-venv-$cptag"
   say "creating build venv ($cptag)"
   uv venv --quiet --python "$PYBIN" "$VENV"
-  # Provenance pin: PyPI's torch 2.11.0 is itself the cu130 build, but this index
+  # Provenance pin: PyPI's torch 2.14.1 is itself the cu130 build, but this index
   # serves ONLY cu130 wheels, so the resolve can never pick a different-CUDA torch
   # and tag the kernel-cache wheel wrong; everything else is plain PyPI.
   uv pip install --quiet --python "$VENV/bin/python" \
-    --index-url https://download.pytorch.org/whl/cu130 "torch>=2.11,<2.12"
+    --index-url https://download.pytorch.org/whl/cu130 "torch>=2.14,<2.15"
   uv pip install --quiet --python "$VENV/bin/python" \
     "setuptools>=77" wheel ninja "apache-tvm-ffi==0.1.13.post3"
 
diff --git a/setup.py b/setup.py
index 997745f1a8..b8e4693c89 100644
--- a/setup.py
+++ b/setup.py
@@ -80,12 +80,12 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
     # These extensions contain host code only. BuildExtension supplies the ROCm
     # platform defines to the C++ compiler; offload architecture flags belong on
     # HIP device sources and would be rejected by the host compiler here.
-    extra_compile = ["-O3", "-std=c++17"]
+    extra_compile = ["-O3", "-std=c++20"]
 else:
     runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths()
     runtime_lib = "cudart"
     runtime_link_args = []
-    extra_compile = ["-O3", "-std=c++17"]
+    extra_compile = ["-O3", "-std=c++20"]
 
 _check_toolchain()
 
@@ -126,7 +126,7 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
                 sources=[
                     "python/freetoken/kernel/csrc/row_store/row_store_ext.cpp",
                 ],
-                extra_compile_args=["-O3", "-std=c++17"],
+                extra_compile_args=["-O3", "-std=c++20"],
             )
         ] if sys.platform == "linux" else []),
     ],

From c8186396f5834994628f7750b11da7300e1d4714 Mon Sep 17 00:00:00 2001
From: cherry77-cloud <1615405@qq.com>
Date: Mon, 12 Oct 2026 03:41:07 +0800
Subject: [PATCH 086/249] fix(sampling): apply presence and frequency penalties
 (#527)

* fix(sampling): apply presence and frequency penalties

* fix(server): reject non-finite sampling penalties

* fix(server): bound presence and frequency penalties
---
 python/freetoken/core.py              |  3 ++
 python/freetoken/engine/sample.py     | 46 +++++++++++++++++++++------
 python/freetoken/server/api_models.py |  8 ++---
 python/freetoken/server/generation.py |  4 +++
 python/freetoken/server/openai_api.py |  4 +++
 5 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 45a763c0fe..812f9a4349 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -24,6 +24,8 @@ class SamplingParams:
     # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one
     # appears in the decoded output; the matched substring (and anything after) is trimmed.
     stop_strs: list[str] = field(default_factory=list)
+    presence_penalty: float = 0.0
+    frequency_penalty: float = 0.0
 
     @property
     def is_greedy(self) -> bool:
@@ -66,6 +68,7 @@ class Req:
     # handler must not free resources under an in-flight forward; it sets this flag and
     # _process_last_data frees the request when the batch drains (after copy_done.synchronize).
     aborted: bool = False
+    output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         assert self.input_ids.is_cpu
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 3bdae76957..c6fb2f40ee 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, List
 
 import torch
@@ -16,6 +16,7 @@ class BatchSamplingArgs:
     top_k: torch.Tensor | None = None
     top_p: torch.Tensor | None = None
     greedy_mask: torch.Tensor | None = None
+    penalties: list[tuple[int, torch.Tensor, float, float]] = field(default_factory=list)
 
 
 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
@@ -59,8 +60,20 @@ class Sampler:
     def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
         is_greedy = [p.is_greedy for p in params]
+        penalties = []
+        for row, req in enumerate(batch.reqs):
+            p = req.sampling_params
+            if not (p.presence_penalty or p.frequency_penalty) or not req.can_decode:
+                continue
+            if req.output_token_counts is None:
+                req.output_token_counts = torch.zeros(
+                    self.vocab_size, dtype=torch.int32, device=self.device
+                )
+            penalties.append(
+                (row, req.output_token_counts, p.presence_penalty, p.frequency_penalty)
+            )
         if all(is_greedy):
-            return BatchSamplingArgs(temperatures=None)
+            return BatchSamplingArgs(temperatures=None, penalties=penalties)
 
         MIN_P = MIN_T = 1e-6
         # Greedy outputs are selected explicitly in sample(); use neutral sampling
@@ -83,17 +96,30 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
         greedy_mask = (
             make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None
         )
-        return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask)
+        return BatchSamplingArgs(
+            temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, penalties=penalties
+        )
 
     @nvtx_annotate("Sampler")
     def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
         with torch.cuda.nvtx.range("Sampler"):
+            if args.penalties:
+                logits = logits.float().clone()
+                for row, counts, presence, frequency in args.penalties:
+                    logits[row] -= frequency * counts + presence * (counts > 0)
             if args.temperatures is None:  # greedy sampling
-                return torch.argmax(logits, dim=-1)
-            tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
-            if args.greedy_mask is not None:
-                # Mixed batches still run probability sampling for all rows, but
-                # greedy rows must follow argmax's deterministic tie-breaking.
-                greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
-                tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
+                tokens = torch.argmax(logits, dim=-1)
+            else:
+                tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
+                if args.greedy_mask is not None:
+                    # Mixed batches still run probability sampling for all rows, but
+                    # greedy rows must follow argmax's deterministic tie-breaking.
+                    greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
+                    tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
+            # Update on the sampling stream: overlapped scheduling can prepare the next
+            # batch before the previous token reaches Req.input_ids on the CPU.
+            for row, counts, _, _ in args.penalties:
+                counts.scatter_add_(
+                    0, tokens[row : row + 1].long(), counts.new_ones(1)
+                )
             return tokens
diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index ffd7172802..6c3ade0924 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -74,8 +74,8 @@ class ChatCompletionRequest(BaseModel):
     stream: bool = False
     stream_options: StreamOptions | None = None
     stop: str | list[str] | None = None
-    presence_penalty: float = 0.0
-    frequency_penalty: float = 0.0
+    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
+    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)
     reasoning_effort: str | None = None
     # DeepSeek-wire thinking toggle ({"type": "enabled"|"disabled"}). Any so a
@@ -112,8 +112,8 @@ class CompletionRequest(BaseModel):
     stream: bool = False
     stream_options: StreamOptions | None = None
     stop: str | list[str] | None = None
-    presence_penalty: float = 0.0
-    frequency_penalty: float = 0.0
+    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
+    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     ignore_eos: bool = False
     logprobs: int | None = None
     echo: bool = False
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index ac988eddc7..2e342ff899 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -169,6 +169,8 @@ def resolve_sampling(
     model_sampling: dict[str, Any],
     stop: str | list[str] | None = None,
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    presence_penalty: float = 0.0,
+    frequency_penalty: float = 0.0,
 ) -> SamplingParams:
     """Map a protocol's sampling fields onto the engine's neutral SamplingParams,
     filling unspecified fields from the checkpoint's recommended defaults."""
@@ -190,6 +192,8 @@ def pick(value, key, framework):
         top_k=pick(top_k, "top_k", -1),
         top_p=pick(top_p, "top_p", 1.0),
         stop_strs=[s for s in stop_list if s],  # drop empty strings (would match everything)
+        presence_penalty=presence_penalty,
+        frequency_penalty=frequency_penalty,
     )
 
 
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index dc2f73a979..99a4283549 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -79,6 +79,8 @@ def chat_request_to_genspec(
             model_sampling=model_sampling,
             stop=req.stop,
             default_max_tokens=default_max_tokens,
+            presence_penalty=req.presence_penalty,
+            frequency_penalty=req.frequency_penalty,
         ),
         chat_template_kwargs=ctk,
         template_tools=_tools_for_template(req),
@@ -550,6 +552,8 @@ def _resolve_sampling(
         model_sampling=model_sampling,
         stop=req.stop,
         default_max_tokens=default_max_tokens,
+        presence_penalty=req.presence_penalty,
+        frequency_penalty=req.frequency_penalty,
     )
 
 

From 95ac86f21567bc244c7a9d91972838fbf049b4ed Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 11 Oct 2026 22:12:41 -0700
Subject: [PATCH 087/249] feat(cpu-moe): serve block-fp8 experts on the cpu and
 hybrid backends (#662)

* feat(cpu-moe): block-FP8 expert GEMV (W8A16) for the cpu/hybrid backends

Block-fp8 checkpoints (Qwen3.5/3.6-FP8, DeepSeek-V3-style, GLM) had no CPU expert
path at all:

    NotImplementedError: --moe-backend cpu/hybrid ... supports
    ['bf16', 'ds_fp4', 'mxfp4_triton', 'nvfp4', 'q4_0'] formats,
    but this checkpoint's experts are 'fp8_block'

so every expert miss had to stream over PCIe at ~25 GB/s while the CPU sat idle.

The bank is fp8-e4m3 row-major [rows, K] with one bf16 scale per 128x128 block --
K contiguous per output row, unlike mxfp4's transposed bank, so it maps straight
onto a bf16 dot product. e4m3 -> bf16 is *exact* (3 mantissa bits into 7, and
bf16's exponent range covers e4m3's whole 2^-9..448 span), so the AVX-512 path
widens weights in-register and feeds `_mm512_dpbf16_ps` against the bf16
activations: identical products to the reference's fp32 math, differing only in
summation order -- the same latitude the bf16 dot already takes.

The normal range rebases the exponent with a shift and a constant +120; exp == 0
is subnormal (m * 2^-9) and does not obey that, so those lanes blend in from a
small table. AVX2 does the same in fp32; scalar goes through the existing 256-entry
e4m3 LUT.

`ft bench bw` gains an fp8_block CPU measurement (it could already build the banks).
Measured on main with this commit alone, `--dtype fp8_block --reps 5`, 2x Xeon Gold
6526Y:

                        CPU MoE      PCIe    ratio   backend picked
    main             not measurable  25.1     --     offload
    main + this      53.7 GB/s       25.1    2.14x   hybrid

Before this commit the format reports no CPU number, so the backend picker has
nothing to weigh and selects offload. The ratio is what decides the backend, so
adding the kernel is what lets a block-fp8 deployment use the CPU at all.

End to end on Qwen/Qwen3.6-35B-A3B-FP8, median of 3 generations at matched GPU
expert residency:

    residency   offload   hybrid    delta
    10%           43.94    50.94    +15.9%
    25%           82.65    55.12    -33%

The win is real but conditional: hybrid pays off once the miss rate is high enough
that PCIe is the bottleneck, and costs throughput when enough experts are resident
that the CPU becomes the limiter. See the PR discussion -- the ratio makes
`--moe-backend auto` select hybrid for every block-fp8 checkpoint, which is right
for the oversubscribed case this project targets and wrong for a model that nearly
fits, and the bench models bandwidth but not miss rate.

Correctness is checked against the Triton block-fp8 decode kernel on identical
banks and routing (bs 1/2/5), plus an exactness check on the e4m3 -> bf16 widening
over all 256 codes. A weight-format ABI probe (`max_weight_format_id`) makes a
stale prebuilt `.so` fail loudly instead of falling through to the wrong dequant
branch and returning silently wrong numbers.

Assisted-by: Claude

* fix(cpu-moe): honor padded fp8 scale strides

The scale banks' trailing dimension is padded so per-expert rows land on
a 16-byte boundary for the fused copy (kernel/aot_models.py:
fp8_block_scale_pad), so the row stride is not ceil(K/128). Recomputing
it from the unpadded block count reads scales at the wrong offset.

Qwen3.8 is the case in the tree: models/qwen3_5_moe/weight.py records
down_scale as 20x5 bf16 = 200 B, and 200 % 16 != 0 pads cols 5 -> 6.
The benchbw workload (H=2048, I=512) and DSV4-Flash (H=7168, I=2048)
both happen to give rows*cols divisible by 8, so neither trips it.

Read the stride off the tensor and pass it through to the kernel, and
let the shape check assert the unpadded extent is present rather than an
exact shape. The fixture now allocates padded banks, so the existing GPU
cross-check covers this too; the GPU kernel was already stride-aware.

* test(cpu-moe): add float64 parity reference and an avx512f fp8 tier

The existing GPU-vs-CPU check compares two fp32 accumulations at a 2e-2
tolerance, which cannot see accumulator precision loss. Add a float64
reference over mixed-sign, full-range e4m3 weights and zero-mean
activations, and force each ISA tier explicitly so the sweep is real.

Also add an avx512f tier between avx512bf16 and avx2, so CPUs with
AVX-512 but no BF16 extension (Skylake-SP through Ice Lake-SP) stop
falling back to 8-wide fp32.

Measured on Xeon Platinum 8368 (Ice Lake-SP, no avx512_bf16), three
repetitions each: +19% single-thread cache-resident (3.31 -> 3.94 GB/s,
disjoint intervals), but only +4% under `ft bench bw --dtype fp8` at 75
threads (60.4 -> 62.7 GB/s, overlapping intervals). At that thread count
the kernel is memory-bound, so the project's own benchmark cannot
resolve the tier; the single-thread gain is the honest figure.

* test(cpu-moe): cover fp8 dispatch and cpu-only regressions

The new tier shipped without test coverage: deleting the avx512f branch
from select_fp8dot left every test green, because isa_name() reports
select_dot() (the bf16 chain) rather than the fp8 dispatch. Injecting a
ragged-tail bug into dot_fp8_block_avx2 was invisible for the same
reason -- no test ran numerics under a forced tier.

Force each available fp8 tier through the float64 parity GEMV, expose
the selected fp8 tier so the assertion can see it, and make the padded
bank test run a GEMV instead of only constructing an executor.

Also move the module-level CUDA guard onto the tests that actually need
a GPU. A CPU MoE kernel's most relevant hosts are the ones without a
usable GPU, and all ten tests were skipping there.

* review: force each ISA tier in the fp8_block test, and stop the docstring overclaiming

The comparison test only varied batch size, so all three cases ran
whatever tier this machine auto-selects and the scalar / AVX2 /
AVX-512-BF16 fp8_block kernels were never compared against the GPU at
all.

Parametrise FREETOKEN_CPU_MOE_ISA over the four tiers. Setting the
variable alone would not be enough: pick_isa() caps *downward*, so asking
for a tier the CPU or build lacks silently returns a lower one and the
test would pass four times over the same kernel. The executor latches its
choice into isa_name() at construction, so read it back and skip when it
does not match what was asked for. On a machine with all four, that is 4
x 3 cases against a 8-expert, 256-hidden problem.

test_subnormal_and_sign_decode_exactly said it validated that "every e4m3
byte must widen to the value the reference LUT gives", but it never calls
the CPU kernel -- it checks that bf16 round-trips e4m3 losslessly and
that subnormals are finite and small. That is the premise the widening
rests on, which is worth pinning, but it is not the widening. Renamed to
test_bf16_carries_every_e4m3_byte and the docstring now says what the
body asserts.

Not run here: this file is skipif(not torch.cuda.is_available()), and
this machine has no GPU and no torch. Compiles clean and ruff reports the
same two pre-existing E741 findings before and after.

Assisted-by: Claude

* fix(moe): let fp8_block experts decode on the cpu executor

* fix(cpu-moe): build the bf16 fp8 tier only when the compiler has avx512bf16

* fix(bench): version the bw profile and ignore older ones

* perf(cpu-moe): read fp8 weights as fp16 bits in the cpu dot kernels

* refactor(cpu-moe): write the fp8_block bank checks like the other formats

* refactor(cpu-moe): name the fp8 cpu symbols fp8_block

* docs(server): list fp8_block in the --moe-cpu-layers help

---------

Co-authored-by: Gabriel A. Devenyi <gdevenyi@gmail.com>
Co-authored-by: ChenyuHe <hechenyu@zju.edu.cn>
---
 docs/cli.md                                   |   2 +-
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       | 188 ++++++++++--
 .../layers/quantization/moe/fp8_block.py      |   3 +-
 python/freetoken/moe/bench_profile.py         |   9 +
 python/freetoken/moe/benchbw.py               |  23 +-
 python/freetoken/moe/cpu_executor.py          |  55 +++-
 python/freetoken/server/args.py               |   2 +-
 tests/models/test_quant_config.py             |   3 +-
 tests/moe/test_cpu_moe_fp8_block.py           | 268 ++++++++++++++++++
 tests/moe/test_hybrid_fetch.py                |   8 +-
 10 files changed, 526 insertions(+), 35 deletions(-)
 create mode 100644 tests/moe/test_cpu_moe_fp8_block.py

diff --git a/docs/cli.md b/docs/cli.md
index a15a34b71c..2166e9e86e 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -87,7 +87,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) |
 | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts |
 | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor |
-| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, nvfp4, mxfp4), so fp8 experts cannot use it |
+| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) |
 | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU |
 | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) |
 | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap |
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index 86765013e3..6f85971e9c 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -10,7 +10,7 @@
 // sync() blocks the host-func thread until the pool drains. The heavy GEMV runs
 // on the persistent worker threads, not the host-func thread.
 //
-// Weight formats: bf16, NVFP4, MXFP4, ds_fp4 and Q4_0 expert banks (see WFmt and
+// Weight formats: bf16, NVFP4, MXFP4, ds_fp4, Q4_0 and block-FP8 expert banks (see WFmt and
 // the per-format bank schemas). Compute is FP32-accumulate; the intermediate is
 // stored bf16 to match the GPU decode path. ISA is chosen once at construction
 // (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar).
@@ -271,6 +271,108 @@ inline float e4m3_decode(uint8_t v) {
   return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7);
 }
 
+// ------------------------ block-FP8 (W8A16) dequant ------------------------
+// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py.
+// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x.
+// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does.
+using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*,
+                                   const bf16_t*, int, const float*);
+
+constexpr int FP8_BLK = 128;  // K elements covered by one weight scale
+
+inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) {
+  float blk = 0.0f;
+  for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]);
+  return blk;
+}
+
+float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                           const bf16_t* s, int K, const float* e4m3) {
+  float acc = 0.0f;
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    // 8 independent sums: a single running sum is bound by the add latency
+    float a[8] = {};
+    const uint8_t* wb = w + k0;
+    const float* pe = xe + k0 / 2;
+    const float* po = xo + k0 / 2;
+    for (int j = 0; j < FP8_BLK / 2; j += 4)
+      for (int l = 0; l < 4; ++l) {
+        a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l];
+        a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l];
+      }
+    acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) *
+           bf16_to_f32(s[b]);
+  }
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+
+#if CPU_MOE_X86
+// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256.
+// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C.
+__attribute__((target("avx2,fma,f16c")))
+float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                         const bf16_t* s, int K, const float* e4m3) {
+  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
+  __m256 total = _mm256_setzero_ps();
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
+    __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
+    __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps();
+    for (int j = 0; j < FP8_BLK; j += 32) {
+      const __m256i v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
+      const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep);
+      const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep);
+      const float* pe = xe + (k0 + j) / 2;
+      const float* po = xo + (k0 + j) / 2;
+      a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0);
+      a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1);
+      a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2);
+      a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3);
+    }
+    const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3));
+    total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
+  }
+  float acc = hsum256(total);
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+
+__attribute__((target("avx512f,avx2,fma,f16c")))
+float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                            const bf16_t* s, int K, const float* e4m3) {
+  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
+  __m512 total = _mm512_setzero_ps();
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
+    __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
+    __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
+    for (int j = 0; j < FP8_BLK; j += 64) {
+      const __m256i v0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
+      const __m256i v1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j + 32));
+      const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep);
+      const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep);
+      const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep);
+      const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep);
+      const float* pe = xe + (k0 + j) / 2;
+      const float* po = xo + (k0 + j) / 2;
+      a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0);
+      a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1);
+      a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2);
+      a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3);
+    }
+    const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3));
+    total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
+  }
+  float acc = _mm512_reduce_add_ps(total);
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+#endif  // CPU_MOE_X86
+
 // Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the
 // ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights).
 using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*,
@@ -677,7 +779,7 @@ struct DotChoice {
 };
 
 // SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by
-// pick_isa() that it implements (fp4 formats have no bf16-specific tier, so the
+// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the
 // avx512bf16 tier maps to their avx512 kernel).
 enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 };
 
@@ -1017,6 +1119,17 @@ void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const b
 }
 #endif
 
+fp8_block_dot_fn select_fp8_block_dot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16
+  if (t >= ISA_AVX512) return dot_fp8_block_avx512f;
+  if (t >= ISA_AVX2) return dot_fp8_block_avx2;
+#endif
+  (void)t;
+  return dot_fp8_block_scalar;
+}
+
 mxgemv_fn select_mxgemv() {
   const IsaTier t = pick_isa();
 #if CPU_MOE_X86
@@ -1229,7 +1342,8 @@ q4dot_fn select_q4dot() {
   return q4_0_dot_i8_scalar;
 }
 
-enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 };
+enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4,
+             WF_FP8_BLOCK = 5 };
 
 // Each ctor pointer arg is the address of a CPU int64 array of length
 // num_layers (one base address per layer, built by cpu_executor.py's
@@ -1246,7 +1360,7 @@ struct CpuMoeExecutor {
   int H, I;
   int act, apply_on_input;
   int fmt;                // WFmt
-  bool needs_di = false;  // pre-deinterleave activations to fp32 (nvfp4/ds_fp4)
+  bool needs_di = false;  // pre-deinterleave activations to fp32 (nvfp4/ds_fp4/fp8_block)
   // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl
   // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up
   // table (down_tbl likewise for down); which reinterpretation applies is picked by
@@ -1269,6 +1383,10 @@ struct CpuMoeExecutor {
   dsdot_fn dsdot;
   mxgemv_fn mxgemv;
   q4dot_fn q4dot;
+  fp8_block_dot_fn fp8_block_dot;
+  // fp8_block: bf16 scale per 128x128 weight block. These are the scale tensor's
+  // trailing strides, including any padding used by the host bank layout.
+  int fp8_block_gu_scale_row = 0, fp8_block_dn_scale_row = 0;
   // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU
   // (same reference grid), so submit() must not repeat it on the host-callback
   // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL
@@ -1287,7 +1405,7 @@ struct CpuMoeExecutor {
 
   std::vector<bf16_t> g_scratch;   // [max_tokens * top_k * I] intermediate
   std::vector<bf16_t> xq_scratch;  // [max_tokens * H] ds_fp4 fp8-roundtripped input
-  // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot.
+  // nvfp4/ds_fp4/fp8_block activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot.
   std::vector<float> xe_scratch, xo_scratch;  // [max_tokens * H/2]   (input)
   std::vector<float> ge_scratch, go_scratch;  // [max_tokens*top_k*I/2] (intermediate)
   // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale.
@@ -1315,7 +1433,7 @@ struct CpuMoeExecutor {
 
   alignas(64) std::atomic<int64_t> p1_next{0};
   std::atomic<int64_t> p2_next{0};
-  std::atomic<int64_t> prt_next{0};  // ds_fp4 intermediate fp8 round-trip phase
+  std::atomic<int64_t> prt_next{0};  // intermediate-row prep phase (prep_g_row)
   int64_t p1_total = 0, p2_total = 0, prt_total = 0;
   int n_iblk = 0, n_hblk = 0;
   std::atomic<int> done_count{0};
@@ -1373,7 +1491,8 @@ struct CpuMoeExecutor {
                  uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr,
                  uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr,
                  uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_,
-                 std::vector<int> core_ids_)
+                 std::vector<int> core_ids_, int fp8_block_gu_scale_stride_,
+                 int fp8_block_dn_scale_stride_)
       : num_threads(num_threads_ > 0 ? num_threads_ : 1),
         num_layers(num_layers_),
         num_experts(num_experts_),
@@ -1393,6 +1512,8 @@ struct CpuMoeExecutor {
         dn_bias_tbl(reinterpret_cast<const uint64_t*>(down_bias_ptr)),
         swiglu_alpha(static_cast<float>(swiglu_alpha_)),
         swiglu_limit(static_cast<float>(swiglu_limit_)),
+        fp8_block_gu_scale_row(fp8_block_gu_scale_stride_),
+        fp8_block_dn_scale_row(fp8_block_dn_scale_stride_),
         core_ids(std::move(core_ids_)) {
     DotChoice c = select_dot();
     dot = c.fn;
@@ -1400,6 +1521,7 @@ struct CpuMoeExecutor {
     dsdot = select_dsdot();
     mxgemv = select_mxgemv();
     q4dot = select_q4dot();
+    fp8_block_dot = select_fp8_block_dot();
     if (weight_format == WF_Q4_0) {
       if (H % 32 != 0 || I % 32 != 0)
         throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32");
@@ -1423,8 +1545,8 @@ struct CpuMoeExecutor {
     // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254].
     for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127);
     g_scratch.assign(static_cast<size_t>(max_tokens) * top_k * I, 0);
-    // Row-major fp4 (nvfp4/ds_fp4) pre-deinterleaves activations to fp32 even/odd.
-    needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4);
+    // Row-major fp4 (nvfp4/ds_fp4) and fp8_block pre-deinterleave activations to fp32 even/odd.
+    needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4 || fmt == WF_FP8_BLOCK);
     if (needs_di) {
       if (fmt == WF_DSFP4) xq_scratch.assign(static_cast<size_t>(max_tokens) * H, 0);
       xe_scratch.assign(static_cast<size_t>(max_tokens) * (H / 2), 0);
@@ -1498,6 +1620,7 @@ struct CpuMoeExecutor {
   // once per task/pass by the caller via tbl_at) pick the layer's own tensors.
   // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves
   // (or, with AVX-VNNI, the per-16-block int8 activations).
+  // fp8_block uses the halves for whole 128-blocks and the bf16 row for a ragged tail.
   inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l,
                          const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row,
                          const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8,
@@ -1511,6 +1634,13 @@ struct CpuMoeExecutor {
           gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes;
       return q4dot(w, xi8, xas, H);  // W4A8: int8 activations (Q8_0), scale in xas
     }
+    if (fmt == WF_FP8_BLOCK) {
+      const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)H;
+      const bf16_t* sc = reinterpret_cast<const bf16_t*>(gu_scale_l) +
+                         ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) *
+                             (size_t)fp8_block_gu_scale_row;
+      return fp8_block_dot(w, xe, xo, x, sc, H, e4m3_lut);
+    }
     const size_t r = (size_t)e * (2 * I) + row;
     if (use_vnni)
       return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16),
@@ -1533,6 +1663,13 @@ struct CpuMoeExecutor {
       const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes;
       return q4dot(w, gi8, gas, I);  // W4A8: int8 activations (Q8_0), scale in gas
     }
+    if (fmt == WF_FP8_BLOCK) {
+      const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)I;
+      const bf16_t* sc = reinterpret_cast<const bf16_t*>(dn_scale_l) +
+                         ((size_t)e * ((H + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) *
+                             (size_t)fp8_block_dn_scale_row;
+      return fp8_block_dot(w, ge, go, g, sc, I, e4m3_lut);
+    }
     const size_t r = (size_t)e * H + row;
     if (use_vnni)
       return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16),
@@ -1582,6 +1719,14 @@ struct CpuMoeExecutor {
 
   const char* isa_name() const { return isa; }
 
+  const char* fp8_block_isa_name() const {
+#if CPU_MOE_X86
+    if (fp8_block_dot == dot_fp8_block_avx512f) return "avx512f";
+    if (fp8_block_dot == dot_fp8_block_avx2) return "avx2";
+#endif
+    return "scalar";
+  }
+
   void barrier(int& local_sense) {
     local_sense ^= 1;
     if (bar_count.fetch_add(1) + 1 == num_threads) {
@@ -1810,9 +1955,8 @@ struct CpuMoeExecutor {
     }
   }
 
-  // Prepare one intermediate row (token,route) for the down GEMV: ds_fp4 first FP8
-  // round-trips it (DSV4 act_quant), then both formats deinterleave to fp32 even/odd
-  // (reused across every down output row).
+  // Prepare one intermediate row (token,route) for the down GEMV, reused across every down output row.
+  // ds_fp4 first FP8 round-trips it (DSV4 act_quant); nvfp4/ds_fp4/fp8_block deinterleave it to fp32 even/odd; q4_0 Q8_0-quantizes it.
   void prep_g_row(int64_t r) {
     bf16_t* g = g_scratch.data() + (size_t)r * I;
     if (use_q4a8) {  // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV.
@@ -1865,9 +2009,8 @@ struct CpuMoeExecutor {
       do_pass1(t, p);
     }
     barrier(local_sense);
-    // Row-major fp4: prepare the intermediate rows (per token,route) before the down
-    // GEMV -- ds_fp4 FP8 round-trips (DSV4 act_quant), both deinterleave to fp32; q4_0
-    // W4A8 Q8_0-quantizes. Needs all of pass1 done (a full row spans every iblk).
+    // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_g_row.
+    // Needs all of pass1 done (a full row spans every iblk).
     if (needs_di || use_q4a8) {
       for (;;) {
         int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed);
@@ -1938,8 +2081,8 @@ struct CpuMoeExecutor {
     done_count.store(0, std::memory_order_relaxed);
     bar_count.store(0, std::memory_order_relaxed);
     bar_sense.store(0, std::memory_order_relaxed);
-    // ds_fp4: FP8 round-trip the per-token input once, up front (single-threaded;
-    // tiny for decode, and done before the workers are woken below).
+    // Deinterleave the per-token input once, up front (ds_fp4 FP8-round-trips it first).
+    // Single-threaded: tiny for decode, and done before the workers are woken below.
     if (needs_di) {
       const size_t xn = static_cast<size_t>(t->num_tokens) * H;
       if (xn / 2 > xe_scratch.size()) {
@@ -2160,7 +2303,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   py::class_<CpuMoeExecutor>(m, "CpuMoeExecutor")
       .def(py::init<int, int, int, int, int, int, int, int, int, int, uintptr_t, uintptr_t,
                     uintptr_t, uintptr_t, uintptr_t, uintptr_t, uintptr_t, uintptr_t,
-                    double, double, std::vector<int>>(),
+                    double, double, std::vector<int>, int, int>(),
            py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"),
            py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"),
            py::arg("max_tokens"), py::arg("activation_id"),
@@ -2169,7 +2312,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"),
            py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"),
            py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"),
-           py::arg("core_ids"))
+           py::arg("core_ids"), py::arg("fp8_block_gu_scale_stride") = 0,
+           py::arg("fp8_block_dn_scale_stride") = 0)
       .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"),
            py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"),
            py::arg("y_ptr"))
@@ -2192,6 +2336,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            [](CpuMoeExecutor& e, int v) { e.act_block = v; },
            py::arg("value"))
       .def("isa_name", &CpuMoeExecutor::isa_name)
+      .def("fp8_block_isa_name", &CpuMoeExecutor::fp8_block_isa_name)
       .def("set_spin_wait",
            [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); },
            py::arg("enabled"))
@@ -2208,4 +2353,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   // (act_apply falls through to gelu_tanh); the probe turns a stale extension
   // into a loud rebuild instruction instead of wrong model outputs.
   m.def("max_generic_act_id", []() { return static_cast<int>(ACT_SWIGLU_CLAMP); });
+  // Weight-format ABI probe. A stale prebuilt .so accepts a newer WFmt id without
+  // complaint and then falls through to the wrong dequant branch -- silently wrong
+  // numbers, not a crash. The Python executor refuses to build against an extension
+  // that predates the format it was asked for.
+  m.def("max_weight_format_id", []() { return static_cast<int>(WF_FP8_BLOCK); });
 }
diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py
index 7552978d3f..6cda6b1ecd 100644
--- a/python/freetoken/layers/quantization/moe/fp8_block.py
+++ b/python/freetoken/layers/quantization/moe/fp8_block.py
@@ -19,9 +19,10 @@ def _pad_scale(rows: int, cols: int) -> int:
 
 class TritonFp8BlockMoEKernel(MoEKernel):
     name = "triton"
+    cpu_format = "fp8_block"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False)
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
         if reason:
             return reason
         if cfg.scheme is not None and fp8_block_size(cfg.scheme) != BLOCK:
diff --git a/python/freetoken/moe/bench_profile.py b/python/freetoken/moe/bench_profile.py
index d7b720738d..5fa38b6f9d 100644
--- a/python/freetoken/moe/bench_profile.py
+++ b/python/freetoken/moe/bench_profile.py
@@ -18,6 +18,9 @@
 
 logger = init_logger(__name__)
 
+# Bump when a format gains or loses a CPU MoE path: verdicts and fetch fractions from an older profile are stale.
+PROFILE_VERSION = 5
+
 # Engine ``expert_quant`` (models/config.py) -> benchbw format key (offload_cache._BANK_SCHEMAS
 # / benchbw._offload_bank_specs). Only the offload-family formats with a CPU MoE weight path can
 # ever resolve to hybrid; anything not listed falls through unmapped and finds no profile entry
@@ -101,6 +104,12 @@ def _usable_profile(
             return None
     if not isinstance(prof, dict):
         return None
+    if not isinstance(prof.get("version"), int) or prof["version"] < PROFILE_VERSION:
+        logger.warning(
+            f"benchbw profile {src} predates this engine (version {prof.get('version')}, "
+            f"need {PROFILE_VERSION}); ignoring it, rerun `ft bench bw`"
+        )
+        return None
     prof_gpu = (prof.get("gpu") or {}).get("name")
     if gpu_name and prof_gpu and prof_gpu != gpu_name:
         logger.warning(
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index e9d6ffe54e..e70cad31d3 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -56,6 +56,7 @@
     single_gpu_arg,
 )
 from freetoken.kernel.pinned import alloc_pinned_tensor
+from freetoken.moe.bench_profile import PROFILE_VERSION
 from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity
 from freetoken.utils import init_logger
 
@@ -63,7 +64,7 @@
 
 # Formats the CPU MoE C++ kernel can compute AND this bench can build banks for; anything
 # else is offload-only here. (The kernel also does q4_0, but this bench has no q4_0 banks.)
-_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4"})
+_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4", "fp8_block"})
 # Formats this bench can build synthetic (correctly-sized) banks for.
 _BUILDABLE_FORMATS = frozenset({"bf16", "nvfp4", "fp8_block", "mxfp4_triton", "ds_fp4"})
 # Friendlier CLI/display aliases for the internal quant_format strings.
@@ -342,6 +343,22 @@ def pin(*shape, dtype):
         gate_up.fill_(0.02)  # uninitialized bf16 can be denormal -> x86 FP slowdown
         down.fill_(0.02)
         return {"gate_up": gate_up, "down": down}
+    if fmt == "fp8_block":
+        def nb(n):
+            return (n + 127) // 128
+        b = {
+            "gate_up": pin(E, 2 * I, H, dtype=torch.float8_e4m3fn),
+            "gate_up_scale": pin(E, nb(2 * I), nb(H), dtype=torch.bfloat16),
+            "down": pin(E, H, I, dtype=torch.float8_e4m3fn),
+            "down_scale": pin(E, nb(H), nb(I), dtype=torch.bfloat16),
+        }
+        # e4m3 0x00 is a clean zero and the scales are unit-ish, so nothing lands in
+        # the fp32 denormal range (which would slow the GEMV and skew the timing).
+        b["gate_up"].view(torch.uint8).fill_(0x38)  # 1.0
+        b["down"].view(torch.uint8).fill_(0x38)
+        b["gate_up_scale"].fill_(1.0)
+        b["down_scale"].fill_(1.0)
+        return b
     if fmt == "nvfp4":
         b = {
             "gate_up_packed": pin(E, 2 * I, H // 2, dtype=torch.uint8),
@@ -649,7 +666,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
                 if fmt == "nvfp4":  # WF_NVFP4 is the only use_vnni format
                     _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI "
                                  "regardless of the forced tier")
-                elif fmt in ("mxfp4_triton", "ds_fp4"):
+                elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"):
                     _note(entry, "isa tier labels are nominal: avx512bf16 == avx512f here "
                                  "(3 real kernels: scalar/avx2/avx512, no bf16/VNNI variant)")
             entry["expert_bytes"] = entry["expert_bytes"] or c["expert_bytes"]
@@ -781,7 +798,7 @@ def _prog(label: str) -> None:
     _prog("done")
 
     result = {
-        "version": 4,
+        "version": PROFILE_VERSION,
         "timestamp": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"),
         "epoch": int(time.time()),
         "host": socket.gethostname(),
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index c4809075c8..f519687ed6 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -71,7 +71,8 @@
 }
 
 # Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp.
-_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4}
+_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4,
+             "fp8_block": 5}
 
 
 def compiled_extension_supports(activation: str) -> bool:
@@ -156,7 +157,7 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]:
 
 class CpuMoeExecutor:
     """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks
-    (bf16, nvfp4, mxfp4_triton, ds_fp4 or q4_0 — see ``_WFMT_IDS`` / ``_resolve_banks``)."""
+    (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0 or fp8_block — see ``_WFMT_IDS`` / ``_resolve_banks``)."""
 
     def __init__(
         self,
@@ -189,6 +190,14 @@ def __init__(
         # error and silently computes the wrong activation in the generic
         # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4
         # handles its act inside the kernel and predates the marker.)
+        # A prebuilt .so that predates this format would accept the id and then take
+        # the wrong dequant branch -- silently wrong output rather than a crash.
+        if _WFMT_IDS[fmt] > getattr(_cpu_moe, "max_weight_format_id", lambda: 4)():
+            raise RuntimeError(
+                f"the compiled _cpu_moe extension predates the {fmt!r} expert format; "
+                "rebuild it with `python setup.py build_ext --inplace` (or reinstall "
+                "the wheel) before serving this checkpoint on the cpu/hybrid backend."
+            )
         if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton":
             supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)()
             if _ACT_IDS[activation] > supported:
@@ -359,6 +368,32 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor:
         self._banks.extend(layers)
         return table
 
+    def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
+        """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride."""
+        gu, gus = banks["gate_up"], banks["gate_up_scale"]
+        dn, dns = banks["down"], banks["down_scale"]
+        assert gu[0].element_size() == 1 and dn[0].element_size() == 1, (gu[0].dtype, dn[0].dtype)
+        assert gus[0].dtype == torch.bfloat16 and dns[0].dtype == torch.bfloat16, (gus[0].dtype, dns[0].dtype)
+        I = int(gu[0].shape[1] // 2)
+        H = int(gu[0].shape[2])
+        assert gu[0].shape[1] == 2 * I
+        assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I)
+        assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I)
+        assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I)
+        ptrs = dict(
+            gate_up_ptr=self._make_table(gu).data_ptr(),
+            down_ptr=self._make_table(dn).data_ptr(),
+            gate_up_scale_ptr=self._make_table(gus).data_ptr(),
+            gate_up_global_ptr=0,
+            down_scale_ptr=self._make_table(dns).data_ptr(),
+            down_global_ptr=0,
+            gate_up_bias_ptr=0,
+            down_bias_ptr=0,
+            fp8_block_gu_scale_stride=int(gus[0].stride(1)),
+            fp8_block_dn_scale_stride=int(dns[0].stride(1)),
+        )
+        return ptrs, (H, I)
+
     def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]:
         """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format.
 
@@ -400,6 +435,9 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]:
         if fmt == "ds_fp4":
             return self._resolve_dsfp4_banks(banks)
 
+        if fmt == "fp8_block":
+            return self._resolve_fp8_block_banks(banks)
+
         # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals.
         gup, gus, gug = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_global"]
         dnp, dns, dng = banks["down"], banks["down_scale"], banks["down_global"]
@@ -515,11 +553,16 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
     def _io_for(self, bs: int) -> dict[str, torch.Tensor]:
         io = self._io.get(bs)
         if io is None:
+            alloc_io = (
+                (lambda *shape, dtype: torch.empty(*shape, dtype=dtype))
+                if self.device.type == "cpu"
+                else alloc_pinned_tensor
+            )
             io = {
-                "x": alloc_pinned_tensor(bs, self.H, dtype=torch.bfloat16),
-                "ids": alloc_pinned_tensor(bs, self.top_k, dtype=torch.int32),
-                "w": alloc_pinned_tensor(bs, self.top_k, dtype=torch.float32),
-                "y": alloc_pinned_tensor(bs, self.H, dtype=torch.bfloat16),
+                "x": alloc_io(bs, self.H, dtype=torch.bfloat16),
+                "ids": alloc_io(bs, self.top_k, dtype=torch.int32),
+                "w": alloc_io(bs, self.top_k, dtype=torch.float32),
+                "y": alloc_io(bs, self.H, dtype=torch.bfloat16),
             }
             self._io[bs] = io
         return io
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 5e6caa759a..7c0ed9b047 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -751,7 +751,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
             "only, where CUDA pinned memory is capped: it locks just enough head+tail layers "
             "for the banks over the pin budget. Any value, 'auto' included, commits to CPU "
             "decode before the model is built, so the expert format must have a CPU executor "
-            "path (bf16, nvfp4, mxfp4); do not pass it on Linux. Unset = every layer on the "
+            "path (bf16, fp8_block, nvfp4, mxfp4); do not pass it on Linux. Unset = every layer on the "
             "GPU; a boot whose banks exceed a known pin budget stops and asks for this flag."
         ),
     )
diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py
index bfcfb87664..0f31d6dee7 100644
--- a/tests/models/test_quant_config.py
+++ b/tests/models/test_quant_config.py
@@ -26,7 +26,6 @@
 from freetoken.layers.quantization import (
     CompressedTensorsConfig,
     Fp8BlockConfig,
-    KernelSelectionError,
     ModelOptConfig,
     Mxfp4Config,
     NoQuantConfig,
@@ -352,7 +351,7 @@ def _gemma_tied_head(model, idx):
     # hybrid decodes on the CPU executor too, so the expert kernel must have a CPU format
     Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {"model.layers.{lin}.mlp.experts": (Nvfp4MoEMethod, "triton")}, strategy="hybrid"),
     Case("Qwen3.6-35B-A3B", NoQuantConfig, {"model.layers.{lin}.mlp.experts": (UnquantizedMoEMethod, "fused")}, strategy="hybrid"),
-    Case("Qwen3.6-35B-A3B-FP8", Fp8BlockConfig, {}, strategy="hybrid", raises=KernelSelectionError),
+    Case("Qwen3.6-35B-A3B-FP8", Fp8BlockConfig, {"model.layers.{lin}.mlp.experts": (Fp8BlockMoEMethod, "triton")}, strategy="hybrid"),
 ]
 
 
diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py
new file mode 100644
index 0000000000..06f954d8d9
--- /dev/null
+++ b/tests/moe/test_cpu_moe_fp8_block.py
@@ -0,0 +1,268 @@
+"""Block-FP8 CPU MoE GEMV vs a float64 reference and FreeToken's Triton block-fp8 decode kernel.
+
+Two things this is really guarding:
+  * the scale is indexed [row // 128][k // 128] and *multiplies*, despite upstream naming it ``weight_scale_inv``. Inverting it, or transposing the two axes, still produces plausible-looking output.
+  * the SIMD kernels decode e4m3 by reading each byte as fp16 bits. Subnormals carry too little weight in a full-range GEMV to show a decode error, so they get their own case.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+BLK = 128
+
+
+def _nb(n: int) -> int:
+    return (n + BLK - 1) // BLK
+
+
+def _make_fp8_block_cache(L, E, H, I, seed=0):
+    from freetoken.kernel.aot_models import fp8_block_scale_pad
+
+    if torch.cuda.is_available():
+        from freetoken.kernel.pinned import alloc_pinned_tensor
+    else:
+        def alloc_pinned_tensor(*shape, dtype):
+            return torch.empty(*shape, dtype=dtype)
+
+    torch.manual_seed(seed)
+    S = L * E
+
+    def rows(OUT, IN):
+        w = alloc_pinned_tensor(S, OUT, IN, dtype=torch.float8_e4m3fn)
+        w.copy_((torch.randn(S, OUT, IN) * 6.0).to(torch.float8_e4m3fn))
+        s = alloc_pinned_tensor(
+            S, _nb(OUT), fp8_block_scale_pad(_nb(OUT), _nb(IN)), dtype=torch.bfloat16
+        )
+        s.copy_((0.01 + 0.02 * torch.rand_like(s)).to(torch.bfloat16))
+        return w, s
+
+    gu, gus = rows(2 * I, H)
+    dn, dns = rows(H, I)
+    return SimpleNamespace(
+        quant_format="fp8_block",
+        bank_sources={
+            "gate_up": list(gu.split(E)), "gate_up_scale": list(gus.split(E)),
+            "down": list(dn.split(E)), "down_scale": list(dns.split(E)),
+        },
+        num_layers=L,
+        num_experts=E,
+        decode_target="cpu",
+        cpu_executor=None,
+    )
+
+
+def test_cpu_fp8_block_accepts_qwen38_padded_scale_bank():
+    from freetoken.kernel.aot_models import fp8_block_scale_pad
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    H, I = 2560, 640
+    cache = _make_fp8_block_cache(1, 1, H, I)
+    assert cache.bank_sources["down_scale"][0].shape[1:] == (
+        _nb(H), fp8_block_scale_pad(_nb(H), _nb(I))
+    ) == (20, 6)
+    ex = CpuMoeExecutor(
+        cache,
+        top_k=1,
+        activation="silu",
+        apply_router_weight_on_input=False,
+        num_threads=1,
+        max_tokens=1,
+        device=torch.device("cpu"),
+    )
+    hidden = torch.randn(1, H, dtype=torch.bfloat16)
+    io = ex._io_for(1)
+    io["x"].copy_(hidden)
+    io["ids"].zero_()
+    io["w"].fill_(1.0)
+    ex._ext.run_task(ex._task_for(0, 1))
+    cpu_out = io["y"][0].float().to(torch.float64)
+    banks = cache.bank_sources
+    gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0])
+    intermediate = torch.nn.functional.silu(gate_up[:I]) * gate_up[I:]
+    expected = _reference_fp8_block_gemv(banks["down"][0], banks["down_scale"][0], intermediate)
+    rel = (cpu_out - expected).abs().max() / (expected.abs().max() + 1e-6)
+    assert rel < 5e-3, f"max relative error {rel.item()}"
+
+
+# every finite e4m3 code, and the subnormal ones (exp == 0) with both signs
+_FINITE_CODES = [c for c in range(256) if c not in (0x7F, 0xFF)]
+_SUBNORMAL_CODES = [c for c in range(256) if c & 0x78 == 0]
+
+
+def _make_coded_fp8_block_cache(H, I, codes, seed=0):
+    from freetoken.kernel.aot_models import fp8_block_scale_pad
+
+    if torch.cuda.is_available():
+        from freetoken.kernel.pinned import alloc_pinned_tensor
+    else:
+        def alloc_pinned_tensor(*shape, dtype):
+            return torch.empty(*shape, dtype=dtype)
+
+    torch.manual_seed(seed)
+    pool = torch.tensor(codes, dtype=torch.uint8)
+
+    def rows(OUT, IN):
+        count = OUT * IN
+        picked = pool.repeat((count + pool.numel() - 1) // pool.numel())[:count]
+        picked = picked[torch.randperm(count)].reshape(1, OUT, IN)
+        weights = alloc_pinned_tensor(1, OUT, IN, dtype=torch.float8_e4m3fn)
+        weights.copy_(picked.view(torch.float8_e4m3fn))
+        scales = alloc_pinned_tensor(
+            1, _nb(OUT), fp8_block_scale_pad(_nb(OUT), _nb(IN)), dtype=torch.bfloat16
+        )
+        scales.copy_((0.01 + 0.09 * torch.rand_like(scales)).to(torch.bfloat16))
+        return weights, scales
+
+    gu, gus = rows(2 * I, H)
+    dn, dns = rows(H, I)
+    return SimpleNamespace(
+        quant_format="fp8_block",
+        bank_sources={
+            "gate_up": [gu], "gate_up_scale": [gus],
+            "down": [dn], "down_scale": [dns],
+        },
+        num_layers=1,
+        num_experts=1,
+        decode_target="cpu",
+        cpu_executor=None,
+    )
+
+
+def _reference_fp8_block_gemv(weights, scales, x):
+    rows, cols = weights.shape[-2:]
+    scale = scales[0, :_nb(rows), :_nb(cols)].to(torch.float64)
+    scale = scale.repeat_interleave(BLK, 0).repeat_interleave(BLK, 1)[:rows, :cols]
+    return (weights[0].to(torch.float64) * scale) @ x.to(torch.float64)
+
+
+@pytest.mark.parametrize("codes", ["finite", "subnormal"])
+@pytest.mark.parametrize(
+    "isa, expected_isa",
+    [("scalar", "scalar"), ("avx2", "avx2"), ("avx512", "avx512f"), ("avx512bf16", "avx512f")],
+)
+def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, codes):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    H, I = 259, 257  # ragged K for both GEMVs; 257 down rows exercise many reductions
+    cache = _make_coded_fp8_block_cache(H, I, _FINITE_CODES if codes == "finite" else _SUBNORMAL_CODES, seed=20260904)
+    monkeypatch.setenv("FREETOKEN_CPU_MOE_ISA", isa)
+    ex = CpuMoeExecutor(
+        cache,
+        top_k=1,
+        activation="silu",
+        apply_router_weight_on_input=False,
+        num_threads=1,
+        max_tokens=1,
+        device=torch.device("cpu"),
+    )
+    tiers = {"scalar": 0, "avx2": 1, "avx512": 2, "avx512bf16": 3}
+    if tiers[isa] > ("scalar", "avx2", "avx512f", "avx512bf16").index(ex.isa):
+        pytest.skip(f"requested {isa}, host/build selected {ex.isa}")
+    assert ex._ext.fp8_block_isa_name() == expected_isa
+
+    hidden = torch.randn(1, H, dtype=torch.bfloat16)
+    io = ex._io_for(1)
+    io["x"].copy_(hidden)
+    io["ids"].zero_()
+    io["w"].fill_(1.0)
+    ex._ext.run_task(ex._task_for(0, 1))
+    cpu_out = io["y"][0].float().to(torch.float64)
+
+    banks = cache.bank_sources
+    gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0])
+    intermediate = (torch.nn.functional.silu(gate_up[:I]) * gate_up[I:])
+    expected = _reference_fp8_block_gemv(banks["down"][0], banks["down_scale"][0], intermediate)
+    rel = (cpu_out - expected).abs().max() / (expected.abs().max() + 1e-6)
+    # The executor stores bf16; measured max relative error is 3.3e-3 (finite codes) and 1.2e-3 (subnormal codes).
+    # The bound leaves room for reduction order, not for a narrow accumulator.
+    assert rel < 5e-3, f"max relative error {rel.item()}"
+
+
+@pytest.mark.parametrize("isa", ["scalar", "avx2", "avx512", "avx512bf16"])
+def test_cpu_fp8_block_isa_override(monkeypatch, isa):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    cache = _make_fp8_block_cache(1, 1, 256, 128, seed=17)
+    monkeypatch.delenv("FREETOKEN_CPU_MOE_ISA", raising=False)
+    auto = CpuMoeExecutor(
+        cache, top_k=1, activation="silu", apply_router_weight_on_input=False,
+        num_threads=1, max_tokens=1, device=torch.device("cpu"),
+    )
+    names = ("scalar", "avx2", "avx512f", "avx512bf16")
+    auto_rank = names.index(auto.isa)
+    monkeypatch.setenv("FREETOKEN_CPU_MOE_ISA", isa)
+    forced = CpuMoeExecutor(
+        cache, top_k=1, activation="silu", apply_router_weight_on_input=False,
+        num_threads=1, max_tokens=1, device=torch.device("cpu"),
+    )
+    assert names.index(forced.isa) == min(names.index({
+        "scalar": "scalar", "avx2": "avx2", "avx512": "avx512f",
+        "avx512bf16": "avx512bf16",
+    }[isa]), auto_rank)
+
+
+# FREETOKEN_CPU_MOE_ISA tier -> the fp8_block kernel it selects; the avx512bf16 tier selects the avx512f kernel, so it is not repeated here.
+# The env var caps *down* only, so a tier the CPU or the build lacks selects a lower kernel, which turns into a skip instead of testing the same kernel twice.
+_ISA_TIERS = {"scalar": "scalar", "avx2": "avx2", "avx512": "avx512f"}
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("isa", list(_ISA_TIERS))
+@pytest.mark.parametrize("bs", [1, 2, 5])
+def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+    from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block
+
+    L, E, H, I, top_k = 2, 8, 256, 128, 4
+    layer = 1
+    dev = torch.device("cuda")
+    cache = _make_fp8_block_cache(L, E, H, I, seed=bs)
+
+    # Set before construction: the executor latches its dot kernel there.
+    monkeypatch.setenv("FREETOKEN_CPU_MOE_ISA", isa)
+    ex = CpuMoeExecutor(
+        cache,
+        top_k=top_k,
+        activation="silu",
+        apply_router_weight_on_input=False,
+        num_threads=0,
+        max_tokens=bs,
+        device=dev,
+    )
+    got = ex._ext.fp8_block_isa_name()
+    if got != _ISA_TIERS[isa]:
+        pytest.skip(f"asked for {isa}, kernel selected {got}: not supported here")
+
+    hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16)
+    ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32)
+    w = torch.rand(bs, top_k, device=dev, dtype=torch.float32)
+
+    cpu_out = ex.decode(layer, hidden, w, ids).float()
+    torch.cuda.synchronize()
+
+    gpu_out = fused_experts_decode_fp8_block(
+        hidden,
+        cache.bank_sources["gate_up"][layer].to(dev),
+        cache.bank_sources["gate_up_scale"][layer].to(dev),
+        cache.bank_sources["down"][layer].to(dev),
+        cache.bank_sources["down_scale"][layer].to(dev),
+        w, ids.clone(), "silu", False,
+    ).float()
+
+    rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)
+    assert rel < 2e-2, f"bs={bs} isa={got} rel err {rel.item()}"
+
+
+def test_fp16_bits_carry_every_e4m3_byte():
+    """A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, which the SIMD kernels rely on."""
+    codes = torch.arange(256, dtype=torch.int32)
+    ref = codes.to(torch.uint8).view(torch.float8_e4m3fn).float()
+    bits = ((codes & 0x80) << 8) | ((codes & 0x7F) << 7)
+    got = torch.where(bits >= 0x8000, bits - 0x10000, bits).to(torch.int16).view(torch.float16).float() * 256
+    finite = torch.isfinite(ref)
+    assert torch.equal(got[finite], ref[finite])
+    assert torch.equal(got[~finite].abs(), torch.full((2,), 480.0))  # 0x7F / 0xFF, as e4m3_decode
diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py
index 169f8e2b6b..a26630b894 100644
--- a/tests/moe/test_hybrid_fetch.py
+++ b/tests/moe/test_hybrid_fetch.py
@@ -11,7 +11,7 @@
 import pytest
 import torch
 
-from freetoken.moe.bench_profile import default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction
+from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction
 from freetoken.moe.offload_cache import OffloadMoeCache
 
 Q = 1 << 16
@@ -42,6 +42,7 @@ def test_balanced_fetch_tracks_fraction():
 
 def test_load_hybrid_fetch_fraction(tmp_path):
     prof = {
+        "version": PROFILE_VERSION,
         "gpu": {"name": "FAKE GPU"},
         "dtype_kernels": {
             "bf16": {"cpu_moe_gbs": 100.0, "pcie_gather_gbs": 40.0},
@@ -64,6 +65,9 @@ def test_load_hybrid_fetch_fraction(tmp_path):
     assert load_hybrid_fetch_fraction("nvfp4", path=str(path)) is None
     # a profile from different hardware is ignored
     assert load_hybrid_fetch_fraction("bf16", gpu_name="OTHER", path=str(path)) is None
+    # so is one written before the current set of CPU formats
+    path.write_text(json.dumps({**prof, "version": PROFILE_VERSION - 1}))
+    assert load_hybrid_fetch_fraction("bf16", path=str(path)) is None
 
 
 def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
@@ -74,7 +78,7 @@ def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
     def write(path, name, verdict):
         os.makedirs(os.path.dirname(path), exist_ok=True)
         with open(path, "w") as f:
-            json.dump({"gpu": {"name": name}, "dtypes": {"bf16": verdict}}, f)
+            json.dump({"version": PROFILE_VERSION, "gpu": {"name": name}, "dtypes": {"bf16": verdict}}, f)
 
     # legacy single file only: used when the name matches, ignored otherwise
     write(default_profile_path(), "FAKE GPU", "hybrid")

From 7080dad4c200e078e08dacf05810faeec3c32896 Mon Sep 17 00:00:00 2001
From: Mukkesh Ganesh <mukmckenzie@gmail.com>
Date: Thu, 8 Oct 2026 22:55:37 -0700
Subject: [PATCH 088/249] test(moe): grant one bf16 ulp only at rounding ties
 in the dsfp4 transcription check

---
 tests/moe/test_dsfp4_reference_semantics.py | 44 +++++++++++++++------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py
index 2f89cf334d..8aef4e661e 100644
--- a/tests/moe/test_dsfp4_reference_semantics.py
+++ b/tests/moe/test_dsfp4_reference_semantics.py
@@ -50,13 +50,16 @@ def _fp8_roundtrip(x: torch.Tensor, block: int) -> torch.Tensor:
 
 def reference(x, slots, weights, banks, block):
     """The MoE over dequantized banks: per-route bf16 expert outputs accumulated into an fp32 ``y``,
-    returned in fp32 (the kernels round that sum to bf16)."""
+    returned in fp32 (the kernels round that sum to bf16). ``tie`` is per element the largest bf16
+    ulp among routes sitting on a rounding tie, where a different fp32 accumulation order can flip
+    the rounded result."""
     gu_p, gu_s, dn_p, dn_s = banks
     W13 = _dequant(gu_p, gu_s)  # [E, 2I, H]
     W2 = _dequant(dn_p, dn_s)  # [E, H, I]
     xq = _fp8_roundtrip(x, block).float()
     T = x.shape[0]
     y = torch.zeros(T, H, dtype=torch.float32, device=x.device)
+    tie = torch.zeros_like(y)
     for t in range(T):
         for r in range(TOP_K):
             e = int(slots[t, r])
@@ -64,8 +67,25 @@ def reference(x, slots, weights, banks, block):
             gate, up = gu[:I].clamp(max=LIMIT), gu[I:].clamp(-LIMIT, LIMIT)
             h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16)
             hq = _fp8_roundtrip(h.view(1, -1), block).float().view(-1)
-            y[t] += (weights[t, r].float() * (hq @ W2[e].T)).to(torch.bfloat16).float()
-    return y
+            raw = weights[t, r].float() * (hq @ W2[e].T)
+            out = raw.to(torch.bfloat16).float()
+            y[t] += out
+            ulp = torch.exp2(torch.floor(torch.log2(out.abs())) - 7.0)
+            # a differing fp32 sum order (~1e-5 relative) only flips a value within noise of a
+            # rounding tie: distance to the rounded value at ulp/2; 10% of an ulp is generous
+            flippable = 0.5 * ulp - (raw - out).abs() < 0.1 * ulp
+            tie[t] = torch.maximum(tie[t], ulp * flippable)
+    return y, tie
+
+
+def _assert_close(got, want, tie, atol=2e-2, rtol=2e-2):
+    # torch.testing.assert_close(atol, rtol) widened per element by ``tie``: a different fp32
+    # accumulation order may flip one per-route bf16 rounding tie, moving y by one ulp
+    bad = (got - want).abs() > atol + rtol * want.abs() + tie
+    assert not bad.any(), (
+        f"Mismatched elements: {int(bad.sum())} / {want.numel()}, "
+        f"greatest absolute difference: {(got - want).abs().max():.6f}"
+    )
 
 
 def _inputs(T, device, seed=1):
@@ -84,8 +104,8 @@ def test_gemv_path_matches_the_transcription(block):
     banks = _banks("cuda")
     x, slots, w = _inputs(6, "cuda")
     got = routed_experts_fp4(x, slots, w, *banks, LIMIT, act_block=block).float()
-    want = reference(x, slots, w, banks, block)
-    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+    want, tie = reference(x, slots, w, banks, block)
+    _assert_close(got, want, tie)
 
 
 @pytest.mark.parametrize("block", [32, 128])
@@ -109,8 +129,8 @@ def test_cpu_executor_matches_the_transcription(block):
     got = ex.decode(0, x, w, slots.clone()).clone().float()
     torch.cuda.synchronize()
     del ex
-    want = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block)
-    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+    want, tie = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block)
+    _assert_close(got, want, tie)
 
 
 def test_grouped_prefill_path_matches_the_transcription():
@@ -119,8 +139,8 @@ def test_grouped_prefill_path_matches_the_transcription():
     banks = _banks("cuda")
     x, slots, w = _inputs(64, "cuda")
     got = fused_ds_fp4.routed_experts_fp4_prefill(x, slots, w, *banks, LIMIT, E, act_block=32).float()
-    want = reference(x, slots, w, banks, 32)
-    torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2)
+    want, tie = reference(x, slots, w, banks, 32)
+    _assert_close(got, want, tie)
 
 
 @pytest.mark.parametrize("block", [32, 128])
@@ -154,7 +174,8 @@ def test_inactive_routes_contribute_zero_without_reading_their_slots(block):
     # and each share matches the reference over its own routes (a zero weight drops a route)
     got_clean = routed_experts_fp4(x, gpu_slots, gpu_w, *clean, LIMIT, act_block=block)
     assert torch.equal(got, got_clean)
-    torch.testing.assert_close(got.float(), reference(x, slots, gpu_w, clean, block), atol=2e-2, rtol=2e-2)
+    want, tie = reference(x, slots, gpu_w, clean, block)
+    _assert_close(got.float(), want, tie)
     # the CPU executor computes the complementary split from the same raw ids (ids < 0 skipped)
     pinned = {}
     for name, t in (("gate_up_packed", clean[0]), ("gate_up_scale", clean[1]), ("down_packed", clean[2]), ("down_scale", clean[3])):
@@ -167,4 +188,5 @@ def test_inactive_routes_contribute_zero_without_reading_their_slots(block):
     cpu_out = ex.decode(0, x, w, cpu_ids.clone()).clone()
     torch.cuda.synchronize()
     del ex
-    torch.testing.assert_close(cpu_out.float(), reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block), atol=2e-2, rtol=2e-2)
+    want, tie = reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block)
+    _assert_close(cpu_out.float(), want, tie)

From 9c635a157798da4aa43e6d051e92620521ed1858 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 13 Oct 2026 01:24:13 -0700
Subject: [PATCH 089/249] build(deps): move to flashlib 0.4.0 and drop the
 numpy<2.5 cap (#669)

---
 pyproject.toml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 95b872f5b8..b308045f6c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,13 +39,12 @@ dependencies = [
     "einops>=0.8,<1",
     "fastapi>=0.115,<1",
     # slot_cache: the device-side LRU admission kernel behind the MoE expert cache.
-    "flashlib==0.3.0", # pin it
+    "flashlib==0.4.0", # pin it
     "gguf>=0.19,<1",
     "huggingface_hub>=1.5,<2",
     "msgpack>=1.1,<2",
     "modelscope>=1.37,<2",
-    # ceiling: numba (via flashlib) needs numpy<2.5
-    "numpy>=2.0,<2.5",
+    "numpy>=2.0,<3",
     "openai>=2.0,<3",
     "partial-json-parser>=0.2,<1",
     "pillow>=10,<13",

From aae638d877f23f3e1f954cc8c9b7b2703bf85478 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 14:07:09 +0300
Subject: [PATCH 090/249] test(moe): a short prefill chunk loads its experts
 through the slot cache, bit-identical to the stream

---
 tests/moe/test_nvfp4_backends.py | 35 ++++++++++++++++++++++++++
 tests/moe/test_offload.py        | 43 ++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 0546fbcd6f..91abeee1a8 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -325,6 +325,41 @@ def test_triton_overlap_prefill_matches_dequant_reference():
         cache.release_prefill_layer(layer_id)
 
 
+@cuda
+def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit():
+    """A short prefill chunk loads only its routed experts, through the slot cache, and runs the
+    same grouped GEMM over the slots: its output must equal the whole-layer stream's exactly, and
+    its admissions must not count as decode misses."""
+    from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4
+
+    device = torch.device("cuda")
+    torch.manual_seed(15)
+    M = 12
+    hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4
+    topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous()
+    topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device)
+
+    streamed_cache, _ = _triton_cache(device)
+    streamed_cache.materialize_layer(1)
+    streamed_cache.copy_missing()
+    streamed = fused_experts_nvfp4(
+        hidden, *streamed_cache.bank_views(E), topk_weights, topk_ids.clone(), E, "silu", False
+    )
+
+    demand_cache, _ = _triton_cache(device, cache_size=2 * S)
+    demand_cache.collect_stats = True
+    slots = topk_ids.clone()
+    demand_cache.ensure_experts(1, slots, counts_as_decode=False)
+    demand_cache.copy_missing()
+    demanded = fused_experts_nvfp4(
+        hidden, *demand_cache.bank_views(), topk_weights, slots, demand_cache.cache_size, "silu", False
+    )
+
+    assert not torch.equal(slots, topk_ids), "the admission must have rewritten expert ids to slots"
+    assert torch.equal(demanded, streamed)
+    assert int(demand_cache.lru_stats.sum()) == 0
+
+
 @cuda
 def test_triton_swigluoai_matches_dequant_reference():
     """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index cd5f5dc773..5e3a20c8fb 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -125,6 +125,49 @@ def fake_fused(
     assert calls["topk_ids"].tolist() == [[2, 1]]
 
 
+def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch):
+    layer, cache = _make_layer_and_cache()
+    cache.prefill_demand_tokens = 1
+    topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32)
+    hidden_states = torch.randn(1, 8)
+    calls = {}
+
+    monkeypatch.setattr(
+        "freetoken.layers.moe.fused_topk",
+        lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, torch.tensor([[2, 1]], dtype=torch.int32)),
+    )
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id))
+
+    def ensure(layer_id, expert_ids, *, counts_as_decode=True):
+        calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode)
+        expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32))
+
+    monkeypatch.setattr(cache, "ensure_experts", ensure)
+    monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True))
+
+    def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation, apply_router_weight_on_input,
+                   act_alpha=1.0, act_limit=float("inf")):
+        calls["w1"] = w1
+        calls["topk_ids"] = got_topk_ids.clone()
+        return hidden_states
+
+    monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused)
+
+    out = layer.prefill_forward(hidden_states, torch.randn(1, 4))
+
+    assert out is hidden_states
+    assert "materialized" not in calls
+    assert calls["ensured"] == (0, [[2, 1]], False)
+    assert calls["copied"] is True
+    # the grouped GEMM reads the whole slot cache at the slots the admission wrote
+    assert calls["w1"].shape[0] == cache.cache_size
+    assert calls["topk_ids"].tolist() == [[5, 3]]
+
+    calls.clear()
+    layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4))
+    assert calls["materialized"] == 0 and "ensured" not in calls
+
+
 def test_offload_moe_layer_prefill_overlap_prefetches_layers_into_two_buffers(monkeypatch):
     from freetoken.moe.offload_cache import OffloadMoeCache
 

From 689f3d137e23aefbb1eecb961e59f4afd72e2a0c Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 14:08:19 +0300
Subject: [PATCH 091/249] feat(moe): load a short prefill chunk's experts on
 demand through the slot cache

---
 docs/cli.md                             |  1 +
 python/freetoken/engine/config.py       |  1 +
 python/freetoken/engine/engine.py       |  2 ++
 python/freetoken/layers/moe.py          | 27 +++++++++++++++++++------
 python/freetoken/moe/offload_cache.py   | 16 ++++++++++++---
 python/freetoken/moe/offload_kernels.py |  4 ++--
 python/freetoken/server/args.py         | 11 ++++++++++
 7 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 2166e9e86e..b4c318e424 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -89,6 +89,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor |
 | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) |
 | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU |
+| `--moe-prefill-demand-tokens` | 0 | Prefill: a chunk of at most this many tokens loads only the experts it routes to, through the expert cache as decode does (GPU decode only); 0 streams whole layers |
 | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) |
 | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap |
 
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index b6fa8b8949..4b30fe7fda 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -47,6 +47,7 @@ class EngineConfig:
     moe_cache_auto: bool = False
     kv_reserve_tokens: int = 8192  # KV floor for --moe-cache-auto; small by design (MoE-priority)
     moe_cache_policy: str = "lru"
+    moe_prefill_demand_tokens: int = 0
     moe_prefill_overlap: bool = True
     # Prefill hit/miss split: serve cache-resident experts D2D during prefill
     # prefetch instead of re-streaming the full layer over PCIe. Needs CUDA >= 12.8
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 3988caf602..5e02fb2ed8 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -753,6 +753,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             cache_size=config.moe_cache_size,
             device=self.device,
             cache_policy=config.moe_cache_policy,
+            prefill_demand_tokens=config.moe_prefill_demand_tokens,
             prefill_overlap=config.moe_prefill_overlap,
             prefill_hit_d2d=config.moe_prefill_hit_d2d,
             quant_format=banks.quant_format,
@@ -1459,6 +1460,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int
     "moe_cpu_layers": None,
     "moe_cpu_threads": 0,
     "moe_hybrid_max_fetch": -1,
+    "moe_prefill_demand_tokens": 0,
     "moe_prefill_overlap": True,
     "moe_prefill_hit_d2d": False,
     "expert_load": "auto",
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 4918a032fe..710a44c79a 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -253,8 +253,9 @@ def prefill_forward(
     # ------------------------------------------------------------------
     # Data movement -- one decision tree for every quant format (the banks
     # registry makes the cache machinery bank-count agnostic). Decode loads
-    # on demand; prefill streams whole layers, double-buffered when overlap
-    # is enabled. The kernels only ever see bank views plus row indices;
+    # on demand, and so does a short enough prefill chunk; a longer one
+    # streams whole layers, double-buffered when overlap is enabled. The
+    # kernels only ever see bank views plus row indices;
     # which kernel runs is decided afterwards, in ``_expert_gemm``.
     # ------------------------------------------------------------------
 
@@ -355,12 +356,26 @@ def _prefill_routed(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> torch.Tensor:
-        """Prefill movement: stream whole layers -- double-buffered behind the
-        previous layer's GEMMs when ``prefill_overlap`` is on, else a synchronous
-        ``materialize_layer``. In both, position == expert id, so the routing ids
-        pass through unmapped."""
+        """Prefill movement. A chunk within ``prefill_demand_tokens`` loads its routed experts
+        into the slot cache as decode does and runs the same grouped GEMM over the slots. A
+        longer one streams whole layers -- double-buffered behind the previous layer's GEMMs
+        when ``prefill_overlap`` is on, else a synchronous ``materialize_layer`` -- where
+        position == expert id, so the routing ids pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
+        if cache.loads_prefill_on_demand(hidden_states.shape[0]):
+            cache.ensure_experts(self.layer_id, topk_ids, counts_as_decode=False)
+            cache.copy_missing()
+            return self._expert_gemm(
+                cache,
+                hidden_states,
+                topk_weights,
+                topk_ids,
+                views=cache.bank_views(),
+                n=cache.cache_size,
+                alphas=cache.alphas_for_slots(self.layer_id),
+                is_prefill=True,
+            )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 0969f73829..2afbb1a79f 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -108,6 +108,9 @@ class OffloadMoeCache:
     cache_size: int
     device: torch.device
     cache_policy: str = "lru"
+    # A prefill chunk of at most this many tokens admits the experts it routes to into the slot
+    # cache, as decode does, instead of streaming every expert of each layer; 0 streams them all.
+    prefill_demand_tokens: int = 0
     prefill_overlap: bool = False
     # Prefill hit/miss split: experts already resident in the slot cache (slots
     # >= 2 * num_experts) are gathered device-side into the double buffer instead
@@ -150,6 +153,10 @@ def __post_init__(self) -> None:
         policy_ids = {"lru": 0}
         assert self.cache_policy in policy_ids
         assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target
+        if self.prefill_demand_tokens and self.decode_target != "gpu":
+            raise ValueError(
+                f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}"
+            )
         if self.layout is None:
             assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}"
         # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None
@@ -843,17 +850,20 @@ def release_prefill_layer(self, layer_id: int) -> None:
             self._prefill_buffer_has_release_event[buffer_id] = True
         self._prefill_buffer_released[buffer_id] = True
 
-    def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
+    def loads_prefill_on_demand(self, num_tokens: int) -> bool:
+        return num_tokens <= self.prefill_demand_tokens
+
+    def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor, *, counts_as_decode: bool = True) -> None:
         from freetoken.moe.offload_kernels import ensure_experts
 
-        if self.collect_decode_freq:
+        if self.collect_decode_freq and counts_as_decode:
             # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to
             # slot ids in place), so snapshot the routing histogram before that happens.
             ids = expert_ids.reshape(-1).long()
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        ensure_experts(self, layer_id, expert_ids)
+        ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode)
 
     def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         """Capped-fetch LRU for the hybrid backend.
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index cf513f52d1..5df2ffec68 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -16,7 +16,7 @@
 )
 
 
-def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
+def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None:
     """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids.
 
     Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the
@@ -35,7 +35,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
         cache.src_indices,
         cache.evict_slots,
         cache.num_indices,
-        stats=cache.lru_stats[layer_id] if cache.collect_stats else None,
+        stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None,
         id_base=layer_id * cache.num_experts,
     )
 
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 7c0ed9b047..519447081b 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -781,6 +781,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--moe-prefill-demand-tokens",
+        type=int,
+        default=ServerArgs.moe_prefill_demand_tokens,
+        help=(
+            "A prefill chunk of at most this many tokens loads only the experts it routes to, "
+            "through the expert cache as decode does, instead of every expert of each layer. "
+            "0 streams whole layers for every chunk."
+        ),
+    )
+
     parser.add_argument(
         "--enable-special-token-ckpt",
         action="store_true",

From 96c731085aca69b74f5fae066bcdc806df540948 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 17:38:27 +0300
Subject: [PATCH 092/249] test(moe): a dense model drops the demand-prefill
 threshold

---
 tests/moe/test_offload.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 5e3a20c8fb..665cd7876e 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -498,6 +498,37 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch):
     assert is_offload_moe_strategy(config.moe_strategy)
 
 
+def test_adjust_config_drops_the_demand_prefill_threshold_for_a_dense_model():
+    from types import SimpleNamespace
+
+    from freetoken.distributed import DistributedInfo
+    from freetoken.engine.config import EngineConfig
+    from freetoken.engine.engine import _adjust_config
+
+    config = EngineConfig(
+        model_path="/tmp/freetoken-test-model",
+        tp_info=DistributedInfo(rank=0, size=1),
+        dtype=torch.float16,
+        attention_backend="triton",
+        moe_prefill_demand_tokens=512,
+    )
+    object.__setattr__(
+        config,
+        "model_config",
+        SimpleNamespace(
+            has_swa_attention=False,
+            has_linear_attention=False,
+            is_moe=False,
+            num_layers=10,
+            expert_quant="none",
+        ),
+    )
+
+    _adjust_config(config)
+
+    assert config.moe_prefill_demand_tokens == 0
+
+
 def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch):
     import freetoken.core as core
     from freetoken.core import Context, Req, get_global_ctx

From a8f24b1764ccc54a7a1509be7f9a923ca23f1a24 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Thu, 1 Oct 2026 06:27:12 +0300
Subject: [PATCH 093/249] test(moe): the demand-prefill threshold is refused
 off GPU decode

OffloadMoeCache raises when --moe-prefill-demand-tokens is set for cpu or hybrid decode, whose experts are not read from the slot cache a demand-loaded chunk fills; the test drives both targets.
---
 tests/moe/test_offload.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 665cd7876e..2cbd29eb29 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -283,6 +283,22 @@ def test_offload_moe_cache_marlin_rejects_slot_count_beyond_kernel_limit():
         )
 
 
+@pytest.mark.parametrize("decode_target", ["cpu", "hybrid"])
+def test_offload_moe_cache_rejects_demand_prefill_without_gpu_decode(decode_target):
+    # a demand-loaded chunk admits its experts into GPU decode's slot cache, which cpu and hybrid decode do not read
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    with pytest.raises(ValueError, match="needs GPU decode"):
+        OffloadMoeCache(
+            num_layers=1,
+            num_experts=4,
+            cache_size=6,
+            device=torch.device("cpu"),
+            decode_target=decode_target,
+            prefill_demand_tokens=64,
+        )
+
+
 def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots():
     from freetoken.moe.offload_cache import OffloadMoeCache
 

From 2010271266f1a0253403a20f897d39d85e70cfe2 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 13 Oct 2026 20:12:58 -0700
Subject: [PATCH 094/249] refactor(cpu-moe): split the executor into weight
 formats and gpu sync backends (#670)

* refactor(cpu-moe): split the executor into weight formats and gpu sync backends

- move each expert weight format into kernels/<fmt>.cpp behind the MoeFormat
  interface in format.h; the executor keeps the threads, the work split and
  the pass 2 route sum
- pass each format its banks and integer options by name; MoeParams keeps only
  the shared shape and activation fields
- key the formats by name in one table that weight_formats() exports, replacing
  the integer format ids and the fp8_block_isa_name hook
- replace the per-(layer, batch size) flag slots, pinned IO buffers and task
  table with one doorbell word MoeTask{layer, num_tokens} and one shared pinned
  IO block of max_tokens rows
- move the GPU/CPU handshake behind GpuSync with a CUDA and a HIP backend picked
  by setup.py; HIP waits on signal memory and adds explicit batch-memop graph
  nodes where the runtime does not record the native memops under capture
- raise on graph capture when the handshake is off and host functions are not
  replay-safe
- allocate format scratch on 64-byte boundaries
- size max_tokens by the explicit cuda_graph_bs as well

Co-authored-by: zihaomu <zihaomu@outlook.com>

* perf(cpu-moe): skip the worker pool for tasks without routed experts

* perf(cpu-moe): publish and combine the cpu handshake with kernels

* fix(cpu-moe): clean up failed hip capability probes

---------

Co-authored-by: zihaomu <zihaomu@outlook.com>
---
 docs/install_amd.md                           |    8 +
 python/freetoken/engine/engine.py             |   13 +-
 python/freetoken/kernel/csrc/cpu_moe/common.h |  235 ++
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       | 2307 ++---------------
 python/freetoken/kernel/csrc/cpu_moe/format.h |  131 +
 .../freetoken/kernel/csrc/cpu_moe/gpu_sync.h  |   78 +
 .../kernel/csrc/cpu_moe/gpu_sync_cuda.cpp     |  166 ++
 .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp      |  275 ++
 .../kernel/csrc/cpu_moe/gpu_sync_io.cu        |  199 ++
 .../kernel/csrc/cpu_moe/gpu_sync_io.h         |   39 +
 .../kernel/csrc/cpu_moe/kernels/bf16.cpp      |  159 ++
 .../kernel/csrc/cpu_moe/kernels/ds_fp4.cpp    |  269 ++
 .../kernel/csrc/cpu_moe/kernels/fp8_block.cpp |  203 ++
 .../kernel/csrc/cpu_moe/kernels/mxfp4.cpp     |  247 ++
 .../kernel/csrc/cpu_moe/kernels/nvfp4.cpp     |  466 ++++
 .../kernel/csrc/cpu_moe/kernels/q4_0.cpp      |  216 ++
 .../csrc/include/freetoken/hip_compat.h       |   12 +
 python/freetoken/layers/moe.py                |    5 +-
 python/freetoken/moe/benchbw.py               |    9 +-
 python/freetoken/moe/cpu_executor.py          |  474 ++--
 setup.py                                      |   17 +-
 tests/moe/hip_probe_faults.cpp                |   75 +
 tests/moe/test_cpu_moe.py                     |  115 +-
 tests/moe/test_cpu_moe_fp8_block.py           |   25 +-
 tests/moe/test_cpu_moe_q4_0.py                |    2 +-
 tests/moe/test_dsfp4_prequant.py              |   10 +-
 26 files changed, 3372 insertions(+), 2383 deletions(-)
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/common.h
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/format.h
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
 create mode 100644 tests/moe/hip_probe_faults.cpp

diff --git a/docs/install_amd.md b/docs/install_amd.md
index 81d1ad5009..ed934ad657 100644
--- a/docs/install_amd.md
+++ b/docs/install_amd.md
@@ -45,3 +45,11 @@ On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`,
 `[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`,
 or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or
 `--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error.
+
+## CPU and hybrid MoE
+
+`--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU
+expert pool through HIP stream memory operations. ROCm 10 records them into CUDA
+graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes
+explicitly. When the startup probe of that path fails, CUDA graph capture stops
+with an error; run those strategies with `--cuda-graph-max-bs 0`.
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 3988caf602..feb33fbb94 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -813,9 +813,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         """Build the persistent CPU MoE executor (decode-time expert compute).
 
         Must run before CUDA graph capture: the worker pool has to be live for the
-        eager warmup forward, and the pinned IO buffers / host-func task pointers
-        must be stable for the captured nodes. Buffers/tasks themselves are
-        allocated lazily on the first (eager) forward at each batch size.
+        eager warmup forward, and the pinned IO block must be stable for the captured
+        nodes.
         """
         from freetoken.moe.cpu_executor import CpuMoeExecutor
 
@@ -828,7 +827,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
             )
         # Decode batches never exceed max_running_req, but CUDA-graph padding can
         # round a batch up to the largest captured size; cover both.
-        max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1)
+        max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1)
         method = sample.quant_method
         executor = CpuMoeExecutor(
             cache,
@@ -1365,10 +1364,10 @@ def _cpu_moe_executor_viable(model_config) -> bool:
     """Whether an automatic CPU-decode decision may target the CPU MoE executor.
 
     A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly."""
-    from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports
+    from freetoken.moe.cpu_executor import compiled_extension_supports
 
     try:
-        from freetoken.kernel import _cpu_moe  # noqa: F401
+        from freetoken.kernel import _cpu_moe
     except ImportError:
         return False
     act = getattr(model_config, "hidden_act", "silu")
@@ -1379,7 +1378,7 @@ def _cpu_moe_executor_viable(model_config) -> bool:
         return False
     expert_quant = getattr(model_config, "expert_quant", "none")
     fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
-    return fmt == "mxfp4" or fmt in _WFMT_IDS
+    return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats()
 
 
 def _pin_budget_bytes(reserved: int = 0) -> int | None:
diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h
new file mode 100644
index 0000000000..adf7e14584
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/common.h
@@ -0,0 +1,235 @@
+// Numeric helpers and ISA detection shared by the CPU MoE formats.
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <new>
+#include <vector>
+
+#if defined(__x86_64__) || defined(__i386__)
+#include <immintrin.h>
+#define CPU_MOE_X86 1
+#else
+#define CPU_MOE_X86 0
+#endif
+
+// Each SIMD kernel is compiled for its own ISA and picked at runtime (pick_isa), so the build sets no global -march.
+#define CPU_MOE_TARGET(...) __attribute__((target(__VA_ARGS__)))
+
+#if CPU_MOE_X86 && ((defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__))
+#define CPU_MOE_HAS_AVX512BF16 1
+#define CPU_MOE_HAS_AVX512VNNI 1
+#endif
+
+namespace freetoken::cpu_moe {
+
+// Scratch rows are read with 64-byte SIMD loads, so their base is cache-line aligned whatever the heap state; std::vector alone guarantees 16 bytes.
+template <class T>
+struct CacheLineAllocator {
+  using value_type = T;
+  CacheLineAllocator() = default;
+  template <class U>
+  CacheLineAllocator(const CacheLineAllocator<U>&) {}
+  T* allocate(size_t n) { return static_cast<T*>(::operator new(n * sizeof(T), std::align_val_t(64))); }
+  void deallocate(T* p, size_t) { ::operator delete(p, std::align_val_t(64)); }
+  template <class U>
+  bool operator==(const CacheLineAllocator<U>&) const { return true; }
+};
+
+template <class T>
+using AlignedVector = std::vector<T, CacheLineAllocator<T>>;
+
+using bf16_t = uint16_t;
+
+inline float bf16_to_f32(bf16_t v) {
+  uint32_t u = static_cast<uint32_t>(v) << 16;
+  float f;
+  std::memcpy(&f, &u, sizeof(f));
+  return f;
+}
+
+inline bf16_t f32_to_bf16(float f) {
+  uint32_t u;
+  std::memcpy(&u, &f, sizeof(u));
+  // round-to-nearest-even
+  const uint32_t lsb = (u >> 16) & 1u;
+  u += 0x7fffu + lsb;
+  return static_cast<bf16_t>(u >> 16);
+}
+
+// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" /
+// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit
+// scalars, so it is handled in the gated_rows epilogue (act_apply never sees it;
+// the mxfp4 kernel additionally fuses its own copy of the same math).
+// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT
+// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim).
+enum ActKind {
+  ACT_SILU = 0,
+  ACT_GELU = 1,
+  ACT_GELU_TANH = 2,
+  ACT_SWIGLUOAI = 3,
+  ACT_SWIGLU_CLAMP = 4,
+};
+
+inline float act_apply(int act, float x) {
+  if (act == ACT_SILU) return x / (1.0f + std::exp(-x));
+  if (act == ACT_GELU)
+    return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f));
+  // gelu_tanh
+  const float k0 = 0.79788456080286535588f;  // sqrt(2/pi)
+  const float inner = k0 * (x + 0.044715f * x * x * x);
+  return 0.5f * x * (1.0f + std::tanh(inner));
+}
+
+// Software prefetch distance (bytes) ahead of the current weight row stream. The
+// weight stream is the bandwidth bottleneck (read once, never reused); nudging the
+// HW prefetcher with a few cache lines of lookahead raises sustained throughput.
+constexpr int PF_AHEAD = 512;
+
+const float kE2M1[16] = {0.0f,  0.5f,  1.0f,  1.5f,  2.0f,  3.0f,  4.0f,  6.0f,
+                         -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f};
+
+inline float fp16_to_f32(uint16_t h) {
+  const uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
+  uint32_t exp = (h >> 10) & 0x1Fu;
+  uint32_t man = h & 0x3FFu;
+  uint32_t f;
+  if (exp == 0) {
+    if (man == 0) {
+      f = sign;
+    } else {
+      exp = 127 - 15 + 1;
+      while ((man & 0x400u) == 0) {
+        man <<= 1;
+        --exp;
+      }
+      man &= 0x3FFu;
+      f = sign | (exp << 23) | (man << 13);
+    }
+  } else if (exp == 0x1Fu) {
+    f = sign | 0x7F800000u | (man << 13);
+  } else {
+    f = sign | ((exp + (127 - 15)) << 23) | (man << 13);
+  }
+  float out;
+  std::memcpy(&out, &f, sizeof(out));
+  return out;
+}
+
+// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT.
+inline float e4m3_decode(uint8_t v) {
+  const float sign = (v & 0x80u) ? -1.0f : 1.0f;
+  const uint32_t exp = (v >> 3) & 0xFu;
+  const uint32_t man = v & 0x7u;
+  if (exp == 0) return sign * (man / 8.0f) * 0.015625f;  // 2^(1-7) = 2^-6
+  return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7);
+}
+
+#if CPU_MOE_X86
+CPU_MOE_TARGET("avx2,fma")
+inline float hsum256(__m256 v) {
+  __m128 lo = _mm256_castps256_ps128(v);
+  lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1));
+  lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo));
+  lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55));
+  return _mm_cvtss_f32(lo);
+}
+
+// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is
+// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry:
+// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of
+// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT.
+CPU_MOE_TARGET("avx2,fma")
+inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) {
+  __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7)));
+  __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28);
+  return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn));
+}
+#endif  // CPU_MOE_X86
+
+// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by
+// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the
+// avx512bf16 tier maps to their avx512 kernel).
+enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 };
+
+// Best tier the CPU+build supports, optionally capped DOWN by
+// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine
+// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias).
+inline IsaTier pick_isa() {
+#if CPU_MOE_X86
+  if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR;
+  IsaTier best = ISA_SCALAR;
+  if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2;
+  if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512;
+#ifdef CPU_MOE_HAS_AVX512BF16
+  if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16;
+#endif
+  if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) {
+    IsaTier want = best;
+    if (!std::strcmp(f, "scalar")) want = ISA_SCALAR;
+    else if (!std::strcmp(f, "avx2")) want = ISA_AVX2;
+    else if (!std::strcmp(f, "avx512")) want = ISA_AVX512;
+    else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16;
+    if (want < best) best = want;  // cap downward; never force above hw/build support
+  }
+  return best;
+#else
+  return ISA_SCALAR;
+#endif
+}
+
+inline const char* tier_name(IsaTier t) {
+#if CPU_MOE_X86
+#ifdef CPU_MOE_HAS_AVX512BF16
+  if (t >= ISA_AVX512BF16) return "avx512bf16";
+#endif
+  if (t >= ISA_AVX512) return "avx512f";
+  if (t >= ISA_AVX2) return "avx2";
+#endif
+  (void)t;
+  return "scalar";
+}
+
+// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5.
+// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path).
+inline bool cpu_has_avxvnni() {
+#if CPU_MOE_X86
+  const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI");
+  if (no && no[0] && no[0] != '0') return false;  // ignore unset/empty/"0"
+  return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni");
+#else
+  return false;
+#endif
+}
+
+// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids,
+// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the
+// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1
+// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is
+// honored here too.
+inline bool cpu_has_avx512vnni() {
+#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI)
+  const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI");
+  if (no && no[0] && no[0] != '0') return false;
+  const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI");
+  if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false;
+  return __builtin_cpu_supports("avx512vnni");
+#else
+  return false;
+#endif
+}
+
+// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]).
+// bf16->fp32 is exact, so this only reorders -- done once per token/route and
+// reused across every output row of the GEMV.
+inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) {
+  for (int m = 0; m < K / 2; ++m) {
+    even[m] = bf16_to_f32(src[2 * m]);
+    odd[m] = bf16_to_f32(src[2 * m + 1]);
+  }
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index 6f85971e9c..c8b4d1b465 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -1,32 +1,27 @@
-// CPU-compute MoE executor for the "cpu" offload backend.
+// CPU-compute MoE executor for the cpu and hybrid MoE strategies, and its Python binding.
 //
 // Decode ships activations to the CPU, computes the routed experts here (reading
-// the pinned host expert banks at full RAM bandwidth), and ships the results
-// back. To keep the whole decode path inside a single CUDA graph we expose
-// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a
-// CPU worker pool + pinned host buffers and never call any CUDA API.
+// the pinned host expert banks at full RAM bandwidth), and ships the results back.
+// One task (one MoE layer) is in flight at a time: submit() wakes the persistent worker
+// pool, the pool runs pass 1 (gate_up) and pass 2 (down), and sync() waits for it.
+// The GPU rings a doorbell to submit and waits on it to sync (GpuSync, gpu_sync.h), so the
+// whole decode step stays in one CUDA graph; cudaLaunchHostFunc nodes are the fallback.
 //
-// One task is in flight at a time (per MoE layer): submit() wakes the pool,
-// sync() blocks the host-func thread until the pool drains. The heavy GEMV runs
-// on the persistent worker threads, not the host-func thread.
-//
-// Weight formats: bf16, NVFP4, MXFP4, ds_fp4, Q4_0 and block-FP8 expert banks (see WFmt and
-// the per-format bank schemas). Compute is FP32-accumulate; the intermediate is
-// stored bf16 to match the GPU decode path. ISA is chosen once at construction
-// (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar).
+// Weight formats live in kernels/, one file each behind the MoeFormat interface (format.h).
+// Compute is FP32-accumulate; the intermediate is stored bf16 to match the GPU decode path.
 
 #include <algorithm>
 #include <atomic>
 #include <condition_variable>
-#include <cmath>
 #include <cstdint>
 #include <chrono>
-#include <cstdio>
 #include <cstring>
+#include <map>
 #include <mutex>
 #include <stdexcept>
 #include <string>
 #include <thread>
+#include <utility>
 #include <vector>
 
 #include <freetoken/hip_compat.h>
@@ -40,1379 +35,99 @@
 #define CPU_MOE_HAS_AFFINITY 0
 #endif
 
-#if defined(__x86_64__) || defined(__i386__)
-#include <immintrin.h>
-#define CPU_MOE_X86 1
-#else
-#define CPU_MOE_X86 0
-#endif
-
-namespace {
-
-inline void cpu_relax() {
-#if CPU_MOE_X86
-  _mm_pause();
-#elif defined(__aarch64__)
-  asm volatile("yield" ::: "memory");
-#endif
-}
-
-using bf16_t = uint16_t;
-
-inline float bf16_to_f32(bf16_t v) {
-  uint32_t u = static_cast<uint32_t>(v) << 16;
-  float f;
-  std::memcpy(&f, &u, sizeof(f));
-  return f;
-}
-
-inline bf16_t f32_to_bf16(float f) {
-  uint32_t u;
-  std::memcpy(&u, &f, sizeof(u));
-  // round-to-nearest-even
-  const uint32_t lsb = (u >> 16) & 1u;
-  u += 0x7fffu + lsb;
-  return static_cast<bf16_t>(u >> 16);
-}
-
-// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" /
-// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit
-// scalars, so it is handled in the do_pass1 epilogue (act_apply never sees it;
-// the mxfp4 kernel additionally fuses its own copy of the same math).
-// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT
-// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim).
-enum ActKind {
-  ACT_SILU = 0,
-  ACT_GELU = 1,
-  ACT_GELU_TANH = 2,
-  ACT_SWIGLUOAI = 3,
-  ACT_SWIGLU_CLAMP = 4,
-};
-
-inline float act_apply(int act, float x) {
-  if (act == ACT_SILU) return x / (1.0f + std::exp(-x));
-  if (act == ACT_GELU)
-    return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f));
-  // gelu_tanh
-  const float k0 = 0.79788456080286535588f;  // sqrt(2/pi)
-  const float inner = k0 * (x + 0.044715f * x * x * x);
-  return 0.5f * x * (1.0f + std::tanh(inner));
-}
-
-// ------------------------------- dot products -------------------------------
-// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function
-// pointer chosen at runtime; the per-row call overhead is negligible vs n.
-
-using dot_fn = float (*)(const bf16_t*, const bf16_t*, int);
-
-float dot_scalar(const bf16_t* w, const bf16_t* x, int n) {
-  float acc = 0.0f;
-  for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
-  return acc;
-}
-
-// Software prefetch distance (bytes) ahead of the current weight row stream. The
-// weight stream is the bandwidth bottleneck (read once, never reused); nudging the
-// HW prefetcher with a few cache lines of lookahead raises sustained throughput.
-constexpr int PF_AHEAD = 512;
-
-#if CPU_MOE_X86
-__attribute__((target("avx512f")))
-float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) {
-  // 4 independent accumulators -> more in-flight loads (memory-level parallelism),
-  // which is what lifts a bandwidth-bound GEMV toward peak.
-  __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
-  __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
-  int i = 0;
-  for (; i + 64 <= n; i += 64) {
-    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
-    for (int j = 0; j < 64; j += 16) {
-      __m256i wi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + i + j));
-      __m256i xi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + i + j));
-      __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16));
-      __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16));
-      __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3;
-      acc = _mm512_fmadd_ps(wf, xf, acc);
-    }
-  }
-  for (; i + 16 <= n; i += 16) {
-    __m256i wi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + i));
-    __m256i xi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + i));
-    __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16));
-    __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16));
-    a0 = _mm512_fmadd_ps(wf, xf, a0);
-  }
-  float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)));
-  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
-  return s;
-}
-
-#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)
-#define CPU_MOE_HAS_AVX512BF16 1
-__attribute__((target("avx512bf16,avx512f")))
-static inline __m512bh load_bh(const bf16_t* p) {
-  __m512i raw = _mm512_loadu_si512(reinterpret_cast<const void*>(p));
-  __m512bh out;
-  std::memcpy(&out, &raw, sizeof(out));
-  return out;
-}
-
-__attribute__((target("avx512bf16,avx512f")))
-float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) {
-  // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge.
-  __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
-  __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
-  int i = 0;
-  for (; i + 128 <= n; i += 128) {
-    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
-    a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i));
-    a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32));
-    a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64));
-    a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96));
-  }
-  for (; i + 32 <= n; i += 32) {
-    a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i));
-  }
-  float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)));
-  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
-  return s;
-}
-#endif  // avx512bf16 available
-
-// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ----------
-// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The
-// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the
-// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width).
-__attribute__((target("avx2,fma")))
-inline float hsum256(__m256 v) {
-  __m128 lo = _mm256_castps256_ps128(v);
-  lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1));
-  lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo));
-  lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55));
-  return _mm_cvtss_f32(lo);
-}
-
-__attribute__((target("avx2,fma")))
-float dot_avx2(const bf16_t* w, const bf16_t* x, int n) {
-  __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
-  __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps();
-  int i = 0;
-  for (; i + 32 <= n; i += 32) {
-    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
-    for (int j = 0; j < 32; j += 8) {
-      __m128i wi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(w + i + j));
-      __m128i xi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + i + j));
-      __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16));
-      __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16));
-      __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3;
-      acc = _mm256_fmadd_ps(wf, xf, acc);
-    }
-  }
-  for (; i + 8 <= n; i += 8) {
-    __m128i wi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(w + i));
-    __m128i xi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + i));
-    __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16));
-    __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16));
-    a0 = _mm256_fmadd_ps(wf, xf, a0);
-  }
-  float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)));
-  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
-  return s;
-}
-#endif  // CPU_MOE_X86
+#include "format.h"
+#include "gpu_sync.h"
+#include "gpu_sync_io.h"
 
-// --------------------------- NVFP4 (W4A16) dequant ---------------------------
-// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in
-// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels
-// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global.
-// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop.
+namespace freetoken::cpu_moe {
 
-const float kE2M1[16] = {0.0f,  0.5f,  1.0f,  1.5f,  2.0f,  3.0f,  4.0f,  6.0f,
-                         -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f};
+std::unique_ptr<MoeFormat> make_bf16_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_nvfp4_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_mxfp4_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_ds_fp4_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_q4_0_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_fp8_block_format(const MoeParams& p, const FormatArgs& a);
 
-// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI
-// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations;
-// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4.
-alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12};
-
-inline float fp16_to_f32(uint16_t h) {
-  const uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
-  uint32_t exp = (h >> 10) & 0x1Fu;
-  uint32_t man = h & 0x3FFu;
-  uint32_t f;
-  if (exp == 0) {
-    if (man == 0) {
-      f = sign;
-    } else {
-      exp = 127 - 15 + 1;
-      while ((man & 0x400u) == 0) {
-        man <<= 1;
-        --exp;
-      }
-      man &= 0x3FFu;
-      f = sign | (exp << 23) | (man << 13);
-    }
-  } else if (exp == 0x1Fu) {
-    f = sign | 0x7F800000u | (man << 13);
-  } else {
-    f = sign | ((exp + (127 - 15)) << 23) | (man << 13);
-  }
-  float out;
-  std::memcpy(&out, &f, sizeof(out));
-  return out;
-}
-
-// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT.
-inline float e4m3_decode(uint8_t v) {
-  const float sign = (v & 0x80u) ? -1.0f : 1.0f;
-  const uint32_t exp = (v >> 3) & 0xFu;
-  const uint32_t man = v & 0x7u;
-  if (exp == 0) return sign * (man / 8.0f) * 0.015625f;  // 2^(1-7) = 2^-6
-  return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7);
-}
-
-// ------------------------ block-FP8 (W8A16) dequant ------------------------
-// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py.
-// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x.
-// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does.
-using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*,
-                                   const bf16_t*, int, const float*);
-
-constexpr int FP8_BLK = 128;  // K elements covered by one weight scale
-
-inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) {
-  float blk = 0.0f;
-  for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]);
-  return blk;
-}
-
-float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
-                           const bf16_t* s, int K, const float* e4m3) {
-  float acc = 0.0f;
-  int b = 0, k0 = 0;
-  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
-    // 8 independent sums: a single running sum is bound by the add latency
-    float a[8] = {};
-    const uint8_t* wb = w + k0;
-    const float* pe = xe + k0 / 2;
-    const float* po = xo + k0 / 2;
-    for (int j = 0; j < FP8_BLK / 2; j += 4)
-      for (int l = 0; l < 4; ++l) {
-        a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l];
-        a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l];
-      }
-    acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) *
-           bf16_to_f32(s[b]);
-  }
-  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
-  return acc;
-}
-
-#if CPU_MOE_X86
-// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256.
-// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C.
-__attribute__((target("avx2,fma,f16c")))
-float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
-                         const bf16_t* s, int K, const float* e4m3) {
-  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
-  __m256 total = _mm256_setzero_ps();
-  int b = 0, k0 = 0;
-  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
-    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
-    __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
-    __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps();
-    for (int j = 0; j < FP8_BLK; j += 32) {
-      const __m256i v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
-      const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep);
-      const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep);
-      const float* pe = xe + (k0 + j) / 2;
-      const float* po = xo + (k0 + j) / 2;
-      a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0);
-      a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1);
-      a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2);
-      a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3);
-    }
-    const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3));
-    total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
-  }
-  float acc = hsum256(total);
-  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
-  return acc;
-}
-
-__attribute__((target("avx512f,avx2,fma,f16c")))
-float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
-                            const bf16_t* s, int K, const float* e4m3) {
-  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
-  __m512 total = _mm512_setzero_ps();
-  int b = 0, k0 = 0;
-  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
-    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
-    __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
-    __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
-    for (int j = 0; j < FP8_BLK; j += 64) {
-      const __m256i v0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
-      const __m256i v1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j + 32));
-      const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep);
-      const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep);
-      const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep);
-      const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep);
-      const float* pe = xe + (k0 + j) / 2;
-      const float* po = xo + (k0 + j) / 2;
-      a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0);
-      a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1);
-      a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2);
-      a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3);
-    }
-    const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3));
-    total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
-  }
-  float acc = _mm512_reduce_add_ps(total);
-  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
-  return acc;
-}
-#endif  // CPU_MOE_X86
-
-// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the
-// ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights).
-using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*,
-                           int, const float*, const float*);
-
-float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global,
-                       const float* xe, const float* xo, int K, const float* e2m1,
-                       const float* e4m3) {
-  float acc = 0.0f;
-  const int nb = K / 16;
-  for (int b = 0; b < nb; ++b) {
-    const float bs = e4m3[scale[b]];
-    const uint8_t* pk = packed + (size_t)b * 8;
-    const float* xeb = xe + (size_t)b * 8;  // 16 K -> 8 even + 8 odd
-    const float* xob = xo + (size_t)b * 8;
-    float bsum = 0.0f;
-    for (int j = 0; j < 8; ++j) {
-      const uint8_t byte = pk[j];
-      bsum += e2m1[byte & 0xF] * xeb[j];
-      bsum += e2m1[byte >> 4] * xob[j];
-    }
-    acc += bs * bsum;
-  }
-  return acc * global;
-}
-
-// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ----
-// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block
-// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale.
-using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int,
-                             const float*, const float*);
-
-[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale,
-                          float global, const int8_t* asi8, int K, const float* e4m3,
-                          const float* asb) {
-  float acc = 0.0f;
-  const int nb = K / 16;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* pk = packed + (size_t)b * 8;
-    const int8_t* ae = asi8 + (size_t)b * 16;       // even(8)
-    const int8_t* ao = ae + 8;                      // odd(8)
-    int isum = 0;
-    for (int j = 0; j < 8; ++j) {
-      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j];
-      isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
-    }
-    acc += (e4m3[scale[b]] * asb[b]) * (float)isum;
-  }
-  return acc * (0.5f * global);
-}
-
-#if CPU_MOE_X86
-// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is
-// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry:
-// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of
-// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT.
-__attribute__((target("avx2,fma")))
-inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) {
-  __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7)));
-  __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28);
-  return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn));
-}
-
-// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K,
-// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so
-// it is applied per lane (low 8 lanes = block b, high 8 = block b+1).
-__attribute__((target("avx512f")))
-inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut,
-                         __m512i loma, float s0, float s1) {
-  __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(pk)));
-  __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
-  __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
-  __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob)));
-  // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1  (pure AVX512F mask move)
-  __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1));
-  return _mm512_mul_ps(prod, scv);
-}
-
-__attribute__((target("avx512f")))
-float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global,
-                       const float* xe, const float* xo, int K, const float* e2m1,
-                       const float* e4m3) {
-  const __m512 lut = _mm512_loadu_ps(e2m1);  // 16 e2m1 values
-  const __m512i loma = _mm512_set1_epi32(0xF);
-  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
-  const int nb = K / 16;  // 8 packed bytes + one e4m3 scale per block
-  int b = 0;
-  for (; b + 4 <= nb; b += 4) {  // two blk2 calls -> 4 blocks, 2 accumulators
-    acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8,
-                                          xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]]));
-    acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8,
-                                          xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]]));
-  }
-  for (; b + 2 <= nb; b += 2)
-    acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8,
-                                          xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]]));
-  float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
-  for (; b < nb; ++b) {  // odd final 16-K block
-    const uint8_t* pk = packed + (size_t)b * 8;
-    const float* xeb = xe + (size_t)b * 8;
-    const float* xob = xo + (size_t)b * 8;
-    float bsum = 0.0f;
-    for (int j = 0; j < 8; ++j) {
-      bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j];
-    }
-    s += e4m3[scale[b]] * bsum;
-  }
-  return s * global;
-}
-
-// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes.
-__attribute__((target("avx2,fma")))
-inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob,
-                             __m256 mag8, float sc) {
-  __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk)));
-  __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
-  __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
-  __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob)));
-  return _mm256_mul_ps(prod, _mm256_set1_ps(sc));
-}
-
-__attribute__((target("avx2,fma")))
-float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global,
-                     const float* xe, const float* xo, int K, const float* e2m1,
-                     const float* e4m3) {
-  const __m256 mag8 = _mm256_loadu_ps(e2m1);  // e2m1[0..7] magnitudes
-  __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps();
-  const int nb = K / 16;
-  int b = 0;
-  for (; b + 2 <= nb; b += 2) {
-    acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8,
-                                              xo + (size_t)b * 8, mag8, e4m3[scale[b]]));
-    acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8,
-                                              xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]]));
-  }
-  for (; b < nb; ++b)
-    acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8,
-                                              xo + (size_t)b * 8, mag8, e4m3[scale[b]]));
-  return hsum256(_mm256_add_ps(acc0, acc1)) * global;
-}
-
-// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)]
-// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less
-// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout.
-__attribute__((target("avx2,avxvnni,fma")))
-inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) {
-  __m128i b = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk));   // 8 bytes
-  __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F));
-  __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F));
-  return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi));            // [lo(8),hi(8)] -> int8
-}
-
-// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a,
-// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32
-// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant
-// (W4A8) differs from the bf16 reference.
-__attribute__((target("avx2,avxvnni,fma")))
-float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global,
-                        const int8_t* asi8, int K, const float* e4m3, const float* asb) {
-  const __m128i lut = _mm_loadu_si128(reinterpret_cast<const __m128i*>(kE2M1x2));
-  __m256 accF = _mm256_setzero_ps();
-  const int nb = K / 16;
-  int b = 0;
-  for (; b + 2 <= nb; b += 2) {
-    __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut);
-    __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut);
-    __m256i w = _mm256_set_m128i(wb1, wb);                              // [blk b | blk b+1]
-    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(asi8 + (size_t)b * 16));
-    __m256i aw = _mm256_sign_epi8(w, w);                               // |w| (u8 operand)
-    __m256i sa = _mm256_sign_epi8(a, w);                               // sign(w)*a (s8 operand)
-    __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa);
-    __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]),
-                                 _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0);
-    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF);
-  }
-  float s = hsum256(accF);
-  for (; b < nb; ++b) {  // tail (odd block count)
-    const uint8_t* pk = packed + (size_t)b * 8;
-    const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8;
-    int isum = 0;
-    for (int j = 0; j < 8; ++j)
-      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
-    s += (e4m3[scale[b]] * asb[b]) * (float)isum;
-  }
-  return s * (0.5f * global);
-}
-
-#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)
-#define CPU_MOE_HAS_AVX512VNNI 1
-
-// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed
-// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then
-// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical
-// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row
-// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte
-// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that
-// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim
-// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20%
-// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it.
-// Prefetch never faults, so overshooting a row/bank tail is safe.
-static int nvfp4_pf_blocks() {
-  static const int v = [] {
-    const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS");
-    return (s && s[0]) ? atoi(s) : -1;
-  }();
-  return v;
-}
-
-// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI
-// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8
-// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so
-// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate.
-// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16.
-// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so
-// the caller can run several independent chains into separate accumulators (the decode ->
-// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound).
-__attribute__((target("avx512f,avx512bw,avx512vnni,avx2")))
-static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale,
-                                   const int8_t* asi8, const float* e4m3, const float* asb,
-                                   int b, __m512i lut, __m512i idx, __m512i mask0F,
-                                   __m512i idxsc) {
-  const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL;  // bytes 8-15 of each 128b lane -> hi nibbles
-  __m256i raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(packed + (size_t)b * 8));
-  __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw));
-  __m512i lo = _mm512_and_si512(src, mask0F);
-  __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F);
-  __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi);
-  __m512i w = _mm512_shuffle_epi8(lut, comb);  // 64 int8 weights (e2m1*2)
-  // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0)
-  __m512i a = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(asi8 + (size_t)b * 16));
-  __m512i aw = _mm512_abs_epi8(w);
-  __mmask64 neg = _mm512_movepi8_mask(w);
-  __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a);
-  __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa);  // 16 int32 (groups of 4)
-  // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales
-  // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32
-  // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls +
-  // a set_ps assembly, which otherwise dominates the per-group op count.
-  int sc_raw;
-  memcpy(&sc_raw, scale + b, 4);
-  __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw));
-  __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b));
-  __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4));
-  return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv);
-}
-
-__attribute__((target("avx512f,avx512bw,avx512vnni,avx2")))
-float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global,
-                              const int8_t* asi8, int K, const float* e4m3, const float* asb) {
-  const __m512i lut = _mm512_broadcast_i32x4(
-      _mm_loadu_si128(reinterpret_cast<const __m128i*>(kE2M1x2)));
-  const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0);  // block i -> 128b lane i
-  const __m512i mask0F = _mm512_set1_epi8(0x0F);
-  const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
-  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
-  __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps();
-  const int nb = K / 16;
-  // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set
-  // env value is used verbatim so a per-machine sweep can reach operating points past
-  // the conservative default.
-  const int pfb = nvfp4_pf_blocks();
-  const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb;
-  int b = 0;
-  // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op
-  // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core
-  // rate sat at ~60% of the core's achievable DRAM stream rate).
-  for (; b + 16 <= nb; b += 16) {
-    // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss
-    // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids).
-    if (pf > 0) {
-      _mm_prefetch(reinterpret_cast<const char*>(packed + ((size_t)b + (size_t)pf) * 8),
-                   _MM_HINT_T0);
-      _mm_prefetch(reinterpret_cast<const char*>(packed + ((size_t)b + (size_t)pf) * 8 + 64),
-                   _MM_HINT_T0);
-    }
-    acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx,
-                                             mask0F, idxsc));
-    acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx,
-                                             mask0F, idxsc));
-    acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx,
-                                             mask0F, idxsc));
-    acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx,
-                                             mask0F, idxsc));
-  }
-  for (; b + 4 <= nb; b += 4)
-    acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx,
-                                             mask0F, idxsc));
-  float s = _mm512_reduce_add_ps(
-      _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3)));
-  for (; b < nb; ++b) {  // tail (<4 remaining 16-K blocks)
-    const uint8_t* pk = packed + (size_t)b * 8;
-    const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8;
-    int isum = 0;
-    for (int j = 0; j < 8; ++j)
-      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
-    s += (e4m3[scale[b]] * asb[b]) * (float)isum;
-  }
-  return s * (0.5f * global);
-}
-#endif  // avx512vnni available
-#endif
-
-// =====================================================================================
-// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or
-// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE
-// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is
-// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays
-// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%,
-// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max
-// frequency (GEMV workers -1.5x: the reported edge regression). Availability is
-// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM,
-// vGPU, old drivers) falls back to the cudaLaunchHostFunc path.
-#if defined(_WIN32)
-#include <windows.h>
-static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); }
-static void* cumemop_dlsym(void* h, const char* n) {
-  return (void*)::GetProcAddress((HMODULE)h, n);
-}
-#else
-#include <dlfcn.h>
-static void* cumemop_dlopen() {
-  void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
-  if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL);
-  return h;
-}
-static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); }
-#endif
-
-using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value,
-                             unsigned int flags);
-static cuMemOp64_fn g_cu_write64 = nullptr;
-static cuMemOp64_fn g_cu_wait64 = nullptr;
-static constexpr unsigned int kCuWaitValueGeq = 0x0;   // CU_STREAM_WAIT_VALUE_GEQ
-static constexpr unsigned int kCuWriteDefault = 0x0;   // CU_STREAM_WRITE_VALUE_DEFAULT
-
-static bool cumemop_resolve() {
-  static bool resolved = [] {
-    void* h = cumemop_dlopen();
-    if (h == nullptr) return false;
-    // 11.7+ made the v2 entry points the default; older drivers export only the v1
-    // names with the same signature.
-    g_cu_write64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWriteValue64_v2"));
-    if (g_cu_write64 == nullptr)
-      g_cu_write64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWriteValue64"));
-    g_cu_wait64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWaitValue64_v2"));
-    if (g_cu_wait64 == nullptr)
-      g_cu_wait64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWaitValue64"));
-    return g_cu_write64 != nullptr && g_cu_wait64 != nullptr;
-  }();
-  return resolved;
-}
-
-// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync.
-// Returns true only if the whole memop path works on THIS stream/device/driver.
-static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) {
-  if (!cumemop_resolve()) return false;
-  auto* s = reinterpret_cast<void*>(stream);
-  if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false;
-  if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false;
-  return cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)) == cudaSuccess;
-}
-
-// GPU side of the flag handshake (see the block comment above): enqueued on the
-// caller's (possibly capturing) stream; the WAIT immediate is the constant 1,
-// replay-safe under CUDA graphs.
-// The startup probe validates EAGER memops; a driver could still reject them at graph
-// capture time. Those enqueue errors would otherwise be swallowed here and surface only
-// as a later EndCapture failure -- log the first CUresult so triage is one step.
-static void cumemop_check(int rc, const char* what) {
-  static std::atomic<bool> warned{false};
-  if (rc != 0 && !warned.exchange(true)) {
-    std::fprintf(stderr,
-                 "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; "
-                 "subsequent errors are not repeated). If this happened during CUDA "
-                 "graph capture, the driver lacks capture support for stream memops -- "
-                 "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n",
-                 what, rc);
-  }
-}
-
-static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr,
-                           int64_t slot) {
-  auto* s = reinterpret_cast<void*>(stream);
-  // Order matters and is preserved by the front end: reset done BEFORE raising ready,
-  // so the coordinator's completion write for THIS step can never be wiped.
-  cumemop_check(g_cu_write64(s, (unsigned long long)(done_addr + (size_t)slot * 8), 0ULL,
-                             kCuWriteDefault),
-                "cuStreamWriteValue64(done)");
-  cumemop_check(g_cu_write64(s, (unsigned long long)(ready_addr + (size_t)slot * 8), 1ULL,
-                             kCuWriteDefault),
-                "cuStreamWriteValue64(ready)");
-}
-
-static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) {
-  cumemop_check(g_cu_wait64(reinterpret_cast<void*>(stream),
-                            (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL,
-                            kCuWaitValueGeq),
-                "cuStreamWaitValue64(done)");
-}
-
-struct DotChoice {
-  dot_fn fn;
-  const char* name;
-};
-
-// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by
-// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the
-// avx512bf16 tier maps to their avx512 kernel).
-enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 };
-
-// Best tier the CPU+build supports, optionally capped DOWN by
-// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine
-// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias).
-inline IsaTier pick_isa() {
-#if CPU_MOE_X86
-  if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR;
-  IsaTier best = ISA_SCALAR;
-  if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2;
-  if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512;
-#ifdef CPU_MOE_HAS_AVX512BF16
-  if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16;
-#endif
-  if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) {
-    IsaTier want = best;
-    if (!std::strcmp(f, "scalar")) want = ISA_SCALAR;
-    else if (!std::strcmp(f, "avx2")) want = ISA_AVX2;
-    else if (!std::strcmp(f, "avx512")) want = ISA_AVX512;
-    else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16;
-    if (want < best) best = want;  // cap downward; never force above hw/build support
-  }
-  return best;
-#else
-  return ISA_SCALAR;
-#endif
-}
-
-DotChoice select_dot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-#ifdef CPU_MOE_HAS_AVX512BF16
-  if (t >= ISA_AVX512BF16) return {dot_avx512bf16, "avx512bf16"};
-#endif
-  if (t >= ISA_AVX512) return {dot_avx512f, "avx512f"};
-  if (t >= ISA_AVX2) return {dot_avx2, "avx2"};
-#endif
-  (void)t;
-  return {dot_scalar, "scalar"};
-}
-
-nvdot_fn select_nvdot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-  if (t >= ISA_AVX512) return dot_nvfp4_avx512;
-  if (t >= ISA_AVX2) return dot_nvfp4_avx2;
-#endif
-  (void)t;
-  return dot_nvfp4_scalar;
-}
-
-// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5.
-// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path).
-inline bool cpu_has_avxvnni() {
-#if CPU_MOE_X86
-  const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI");
-  if (no && no[0] && no[0] != '0') return false;  // ignore unset/empty/"0"
-  return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni");
-#else
-  return false;
-#endif
-}
+namespace {
 
-// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids,
-// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the
-// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1
-// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is
-// honored here too.
-inline bool cpu_has_avx512vnni() {
-#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI)
-  const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI");
-  if (no && no[0] && no[0] != '0') return false;
-  const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI");
-  if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false;
-  return __builtin_cpu_supports("avx512vnni");
+inline void pin_this_thread(int cpu) {
+#if CPU_MOE_HAS_AFFINITY
+  cpu_set_t set;
+  CPU_ZERO(&set);
+  CPU_SET(cpu, &set);
+  pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
 #else
-  return false;
-#endif
-}
-
-// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the
-// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference.
-nvi8dot_fn select_nvi8dot() {
-#if CPU_MOE_X86
-#if defined(CPU_MOE_HAS_AVX512VNNI)
-  if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni;
-#endif
-  if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni;
-#endif
-  return nullptr;
-}
-
-// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) ---------------------------
-// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global
-// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128
-// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe +
-// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the
-// down output.
-
-// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per
-// token/route and reused across all output rows. This drops the hot dot to one
-// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16,
-// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5
-// bound and approaches the bf16 memory-bandwidth ceiling.
-using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int,
-                           const float*, const float*);
-
-float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe,
-                       const float* xo, int K, const float* e2m1, const float* e8m0) {
-  float acc = 0.0f;
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const float sc = e8m0[scale[b]];
-    const uint8_t* pk = packed + (size_t)b * 16;  // 16 bytes = 32 codes
-    const float* xeb = xe + (size_t)b * 16;
-    const float* xob = xo + (size_t)b * 16;
-    float bsum = 0.0f;
-    for (int j = 0; j < 16; ++j) {
-      const uint8_t byte = pk[j];
-      bsum += e2m1[byte & 0xF] * xeb[j];   // low nibble  -> even-K activation
-      bsum += e2m1[byte >> 4] * xob[j];    // high nibble -> odd-K activation
-    }
-    acc += sc * bsum;
-  }
-  return acc;
-}
-
-#if CPU_MOE_X86
-// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times
-// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale.
-__attribute__((target("avx512f")))
-inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut,
-                        __m512i loma, float sc) {
-  __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(pk)));
-  __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
-  __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
-  __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob)));
-  return _mm512_mul_ps(prod, _mm512_set1_ps(sc));
-}
-
-__attribute__((target("avx512f")))
-float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe,
-                       const float* xo, int K, const float* e2m1, const float* e8m0) {
-  const __m512 lut = _mm512_loadu_ps(e2m1);
-  const __m512i loma = _mm512_set1_epi32(0xF);
-  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
-  const int nb = K / 32;  // 16 packed bytes + one e8m0 scale per block
-  int b = 0;
-  for (; b + 2 <= nb; b += 2) {  // two independent accumulators hide FMA latency
-    acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16,
-                                         xo + (size_t)b * 16, lut, loma, e8m0[scale[b]]));
-    acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16,
-                                         xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]]));
-  }
-  for (; b < nb; ++b)
-    acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16,
-                                         xo + (size_t)b * 16, lut, loma, e8m0[scale[b]]));
-  return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
-}
-
-// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each).
-__attribute__((target("avx2,fma")))
-inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) {
-  __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk)));
-  __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
-  __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
-  return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob)));
-}
-
-__attribute__((target("avx2,fma")))
-float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe,
-                     const float* xo, int K, const float* e2m1, const float* e8m0) {
-  const __m256 mag8 = _mm256_loadu_ps(e2m1);
-  __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps();
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* pk = packed + (size_t)b * 16;
-    const float* xeb = xe + (size_t)b * 16;
-    const float* xob = xo + (size_t)b * 16;
-    const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]);
-    acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0);
-    acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1);
-  }
-  return hsum256(_mm256_add_ps(acc0, acc1));
-}
-#endif
-
-dsdot_fn select_dsdot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-  if (t >= ISA_AVX512) return dot_dsfp4_avx512;
-  if (t >= ISA_AVX2) return dot_dsfp4_avx2;
-#endif
-  (void)t;
-  return dot_dsfp4_scalar;
-}
-
-// ------------------------- mxfp4 (gpt-oss) GEMV -----------------------------
-// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2]
-// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1])
-// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of
-// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache-
-// sequential) loop. Used by both gate_up (K=H) and down (K=I).
-using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int,
-                           int, const float*, const float*);
-
-void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
-                       int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
-  for (int c = 0; c < ncol; ++c) out[c] = 0.0f;
-  for (int kb = 0; kb < Kpairs; ++kb) {
-    const uint8_t* w = blk + (size_t)kb * N2;
-    const uint8_t* s = scl + (size_t)(kb >> 4) * N2;
-    const float xl = bf16_to_f32(x[2 * kb]);
-    const float xh = bf16_to_f32(x[2 * kb + 1]);
-    for (int c = 0; c < ncol; ++c) {
-      const uint8_t byte = w[c];
-      out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]];
-    }
-  }
-}
-
-#if CPU_MOE_X86
-__attribute__((target("avx512f")))
-void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
-                       int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
-  (void)e8m0;  // e8m0[c]=2^(c-127) computed via bit construction (no gather)
-  const __m512 lut = _mm512_loadu_ps(e2m1);
-  const __m512i loma = _mm512_set1_epi32(0xF);
-  // K-outer / N-inner: each kb cache line is read once and all live column chunks
-  // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once.
-  int c0 = 0;
-  for (; c0 + 16 <= ncol; c0 += 64) {
-    const int nchunk = std::min(4, (ncol - c0) / 16);
-    __m512 acc[4];
-    for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps();
-    for (int kblk = 0; kblk < Kpairs; kblk += 16) {  // 16 K-pairs = 32 K = one scale row
-      __m512 sc[4];
-      for (int ci = 0; ci < nchunk; ++ci) {
-        __m128i sraw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(
-            scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16));
-        sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23));
-      }
-      __m512 blk_acc[4];
-      for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps();
-      for (int kk = 0; kk < 16; ++kk) {
-        const int kb = kblk + kk;
-        const uint8_t* wbase = blk + (size_t)kb * N2 + c0;
-        // The transposed layout strides K by N2 bytes; prefetch ahead so the strided
-        // reads are not exposed to DRAM latency (the HW streamer misses big strides).
-        constexpr int PFD = 8;
-        if (kb + PFD < Kpairs)
-          _mm_prefetch(reinterpret_cast<const char*>(blk + (size_t)(kb + PFD) * N2 + c0),
-                       _MM_HINT_T0);
-        const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb]));
-        const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1]));
-        for (int ci = 0; ci < nchunk; ++ci) {
-          __m512i wi = _mm512_cvtepu8_epi32(
-              _mm_loadu_si128(reinterpret_cast<const __m128i*>(wbase + ci * 16)));
-          __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
-          __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
-          blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]);
-          blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]);
-        }
-      }
-      for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]);
-    }
-    for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]);
-  }
-  for (int c = c0; c < ncol; ++c) {  // tail columns (< 16)
-    float o = 0.0f;
-    for (int kb = 0; kb < Kpairs; ++kb) {
-      const uint8_t byte = blk[(size_t)kb * N2 + c];
-      uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23;
-      float sc;
-      std::memcpy(&sc, &bits, 4);
-      o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) +
-            e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc;
-    }
-    out[c] = o;
-  }
-}
-
-__attribute__((target("avx2,fma")))
-void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
-                     int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
-  (void)e8m0;  // e8m0[s]=2^(s-127) built via s<<23 (no gather)
-  const __m256 mag8 = _mm256_loadu_ps(e2m1);
-  int c0 = 0;
-  for (; c0 + 8 <= ncol; c0 += 32) {  // up to 4 chunks of 8 = 32 cols
-    const int nchunk = std::min(4, (ncol - c0) / 8);
-    __m256 acc[4];
-    for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps();
-    for (int kblk = 0; kblk < Kpairs; kblk += 16) {  // 16 K-pairs = one scale row
-      __m256 sc[4];
-      for (int ci = 0; ci < nchunk; ++ci) {
-        __m128i sraw = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(
-            scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8));
-        sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23));
-      }
-      __m256 blk_acc[4];
-      for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps();
-      for (int kk = 0; kk < 16; ++kk) {
-        const int kb = kblk + kk;
-        const uint8_t* wbase = blk + (size_t)kb * N2 + c0;
-        constexpr int PFD = 8;
-        if (kb + PFD < Kpairs)
-          _mm_prefetch(reinterpret_cast<const char*>(blk + (size_t)(kb + PFD) * N2 + c0),
-                       _MM_HINT_T0);
-        const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb]));
-        const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1]));
-        for (int ci = 0; ci < nchunk; ++ci) {
-          __m256i wi = _mm256_cvtepu8_epi32(
-              _mm_loadl_epi64(reinterpret_cast<const __m128i*>(wbase + ci * 8)));
-          __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
-          __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
-          blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]);
-          blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]);
-        }
-      }
-      for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]);
-    }
-    for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]);
-  }
-  for (int c = c0; c < ncol; ++c) {  // tail columns (< 8); none when ncol%8==0
-    float o = 0.0f;
-    for (int kb = 0; kb < Kpairs; ++kb) {
-      const uint8_t byte = blk[(size_t)kb * N2 + c];
-      uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23;
-      float sc;
-      std::memcpy(&sc, &bits, 4);
-      o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) +
-            e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc;
-    }
-    out[c] = o;
-  }
-}
-#endif
-
-fp8_block_dot_fn select_fp8_block_dot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-  // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16
-  if (t >= ISA_AVX512) return dot_fp8_block_avx512f;
-  if (t >= ISA_AVX2) return dot_fp8_block_avx2;
+  (void)cpu;
 #endif
-  (void)t;
-  return dot_fp8_block_scalar;
 }
 
-mxgemv_fn select_mxgemv() {
-  const IsaTier t = pick_isa();
+inline void cpu_relax() {
 #if CPU_MOE_X86
-  if (t >= ISA_AVX512) return mxfp4_gemv_avx512;
-  if (t >= ISA_AVX2) return mxfp4_gemv_avx2;
+  _mm_pause();
+#elif defined(__aarch64__)
+  asm volatile("yield" ::: "memory");
 #endif
-  (void)t;
-  return mxfp4_gemv_scalar;
-}
-
-// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches
-// torch.float8_e4m3fn / triton .to(float8e4nv).
-inline float e4m3_round(float x) {
-  const float sign = x < 0.0f ? -1.0f : 1.0f;
-  const float a = std::fabs(x);
-  if (a == 0.0f) return 0.0f;
-  if (a >= 448.0f) return sign * 448.0f;
-  int e;
-  std::frexp(a, &e);  // a in [2^(e-1), 2^e)
-  float step = std::ldexp(1.0f, e - 4);
-  const float min_step = std::ldexp(1.0f, -9);  // e4m3 subnormal step (2^-9)
-  if (step < min_step) step = min_step;
-  float r = std::nearbyint(a / step) * step;
-  if (r > 448.0f) r = 448.0f;
-  return sign * r;
-}
-
-// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale).
-inline int ceil_log2_pos(float v) {
-  uint32_t bits;
-  std::memcpy(&bits, &v, sizeof(bits));
-  const int exp = (int)((bits >> 23) & 0xFF);
-  const int man = (int)(bits & 0x7FFFFF);
-  return exp - 127 + (man != 0 ? 1 : 0);
-}
-
-// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]).
-// bf16->fp32 is exact, so this only reorders -- done once per token/route and
-// reused across every output row of the GEMV.
-inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) {
-  for (int m = 0; m < K / 2; ++m) {
-    even[m] = bf16_to_f32(src[2 * m]);
-    odd[m] = bf16_to_f32(src[2 * m + 1]);
-  }
-}
-
-// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's
-// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448));
-// y = round_e4m3(clamp(x/s,+-448)) * s.
-void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) {
-  for (int b0 = 0; b0 < K; b0 += block) {
-    const int b1 = std::min(K, b0 + block);
-    float amax = 1e-4f;
-    for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i])));
-    const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f)));
-    const float inv_s = 1.0f / s;
-    for (int i = b0; i < b1; ++i) {
-      float q = bf16_to_f32(src[i]) * inv_s;
-      q = std::min(448.0f, std::max(-448.0f, q));
-      dst[i] = f32_to_bf16(e4m3_round(q) * s);
-    }
-  }
 }
 
-// --------------------------------- executor ---------------------------------
-
-struct CpuMoeExecutor;
+constexpr int kMaxTopK = 64;
 
+// One MoE layer on the first num_tokens rows of the pinned IO.
+// It is also the doorbell word the GPU writes, a constant baked into each captured graph; num_tokens >= 1 keeps it nonzero, since 0 means no doorbell.
 struct MoeTask {
-  CpuMoeExecutor* exec;
-  int layer_id;
-  int num_tokens;
-  const bf16_t* x;     // [num_tokens, H]
-  const int32_t* ids;  // [num_tokens, top_k]  (raw expert ids; <0 = skip)
-  const float* w;      // [num_tokens, top_k]
-  bf16_t* y;           // [num_tokens, H]
+  uint64_t layer : 32;
+  uint64_t num_tokens : 32;
 };
+static_assert(sizeof(MoeTask) == sizeof(uint64_t), "MoeTask must pack into one 64-bit word");
 
-// Output-row tiling. Small enough to give every worker independent work even at
-// batch size 1; large enough to amortize the atomic work-grab.
-//
-// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read
-// every expert weight byte exactly once per token (each output row block is owned
-// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the
-// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe).
-// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription
-// thrashes the spin-barrier. Deferred (not worth it here / for this workload):
-//   - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would
-//     only pay off in a grouped/batched (dedup) path.
-//   - expert dedup for bs>1: read each distinct expert once and GEMM its tokens.
-//     Helps locality+bytes when bs is large; decode batches here are tiny (<=4).
-//   - NUMA: a single node is assumed. Multi-socket machines would split each
-//     expert's K dimension per node (banks are already per-row contiguous).
-constexpr int IBLK = 32;
-constexpr int HBLK = 32;
-
-// -------------------------------- Q4_0 (W4A8) --------------------------------
-// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed
-// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a
-// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches
-// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload
-// path streams.
-//
-// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0
-// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to
-// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high
-// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD
-// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then
-// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU
-// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close.
-using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int);
-
-float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  float acc = 0.0f;
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    const uint8_t* q = blk + 2;  // 16 nibble bytes
-    const int8_t* a = aq + (size_t)b * 32;
-    int isum = 0;
-    for (int j = 0; j < 16; ++j) {
-      isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j];       // elem j
-      isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j];    // elem 16+j
-    }
-    acc += fp16_to_f32(dh) * asb[b] * (float)isum;
-  }
-  return acc;
-}
-
-#if CPU_MOE_X86
-// fp16 block scale -> fp32 via HW F16C (single value in lane 0).
-__attribute__((target("f16c")))
-static inline float q4_scale(uint16_t h) {
-  return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h)));
+inline uint64_t to_word(MoeTask t) {
+  uint64_t w;
+  std::memcpy(&w, &t, sizeof(w));
+  return w;
 }
 
-// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low
-// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8).
-__attribute__((target("avx2")))
-static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) {
-  const __m128i qb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(blk + 2));
-  const __m128i lo = _mm_and_si128(qb, mask);
-  const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask);
-  return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight);
+inline MoeTask from_word(uint64_t w) {
+  MoeTask t;
+  std::memcpy(&t, &w, sizeof(t));
+  return t;
 }
 
-// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW +
-// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no
-// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the
-// avx512-tier fallback, since the block dot is 256-bit either way).
-__attribute__((target("avx2,fma,f16c")))
-float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  const __m128i mask = _mm_set1_epi8(0x0F);
-  const __m256i eight = _mm256_set1_epi8(8);
-  const __m256i ones16 = _mm256_set1_epi16(1);
-  __m256 accF = _mm256_setzero_ps();
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    __m256i wq = q4_unpack32(blk, mask, eight);
-    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
-    __m256i aw = _mm256_sign_epi8(wq, wq);              // |wq|          (unsigned operand)
-    __m256i sa = _mm256_sign_epi8(a, wq);               // sign(wq) * a  (signed operand)
-    __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16);  // 8 int32
-    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
-  }
-  return hsum256(accF);
-}
-
-// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs).
-__attribute__((target("avx2,avxvnni,fma,f16c")))
-float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  const __m128i mask = _mm_set1_epi8(0x0F);
-  const __m256i eight = _mm256_set1_epi8(8);
-  __m256 accF = _mm256_setzero_ps();
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    __m256i wq = q4_unpack32(blk, mask, eight);
-    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
-    __m256i aw = _mm256_sign_epi8(wq, wq);   // |wq|            (unsigned operand)
-    __m256i sa = _mm256_sign_epi8(a, wq);    // sign(wq) * a    (signed operand)
-    __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa);
-    // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and
-    // reduce at the end (equivalent to scale * block_total).
-    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
-  }
-  return hsum256(accF);
-}
-#endif  // CPU_MOE_X86
-
-// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to
-// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when
-// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers.
-q4dot_fn select_q4dot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-  if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni;
-  if (t >= ISA_AVX2) return q4_0_dot_i8_avx2;
-#endif
-  (void)t;
-  return q4_0_dot_i8_scalar;
-}
-
-enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4,
-             WF_FP8_BLOCK = 5 };
+// Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names.
+const std::pair<const char*, std::unique_ptr<MoeFormat> (*)(const MoeParams&, const FormatArgs&)> kFormats[] = {
+    {"bf16", make_bf16_format},     {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format},
+    {"ds_fp4", make_ds_fp4_format}, {"q4_0", make_q4_0_format},   {"fp8_block", make_fp8_block_format},
+};
 
-// Each ctor pointer arg is the address of a CPU int64 array of length
-// num_layers (one base address per layer, built by cpu_executor.py's
-// _make_table), not a single flat bank. tbl_at resolves
-// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr
-// arg 0) resolves to nullptr without dereferencing.
-inline const void* tbl_at(const uint64_t* tbl, int layer_id) {
-  return tbl ? reinterpret_cast<const void*>(tbl[layer_id]) : nullptr;
+std::unique_ptr<MoeFormat> make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) {
+  for (const auto& [n, make] : kFormats)
+    if (name == n) return make(p, a);
+  throw std::invalid_argument("unknown CPU MoE weight format " + name);
 }
 
 struct CpuMoeExecutor {
+  // A host-function callback gets one pointer, so it carries the executor with the task.
+  struct HostTask {
+    CpuMoeExecutor* exec;
+    MoeTask task;
+  };
+
+  const int num_layers, num_experts;
+  const bool apply_on_input;
+  MoeParams params;  // the format keeps a reference; declared before fmt
+  std::unique_ptr<MoeFormat> fmt;
+  bool preps_inter = false;
+  bool rounds_routes = false;
+  int n_iblk = 0, n_hblk = 0;  // kRowBlock blocks per intermediate row and per output row
+
+  // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w.
+  const bf16_t* io_x;
+  const int32_t* io_ids;
+  const float* io_w;
+  bf16_t* io_y;
+  AlignedVector<bf16_t> g_scratch;   // [max_tokens * top_k * I] intermediate
+
+  // Worker pool.
   int num_threads;
-  int num_layers, num_experts, top_k;
-  int H, I;
-  int act, apply_on_input;
-  int fmt;                // WFmt
-  bool needs_di = false;  // pre-deinterleave activations to fp32 (nvfp4/ds_fp4/fp8_block)
-  // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl
-  // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up
-  // table (down_tbl likewise for down); which reinterpretation applies is picked by
-  // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4).
-  const uint64_t* gate_up_tbl;   // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks
-  const uint64_t* down_tbl;      // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks
-  const uint64_t* gu_scale_tbl;  // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales
-  const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals
-  const uint64_t* dn_scale_tbl;  // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales
-  const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals
-  const uint64_t* gu_bias_tbl;   // mxfp4: [E,2I] bf16 biases
-  const uint64_t* dn_bias_tbl;   // mxfp4: [E,H] bf16 biases
-  float swiglu_alpha;
-  float swiglu_limit;          // +inf == no clamp
-  dot_fn dot;
-  nvdot_fn nvdot;
-  nvi8dot_fn nvi8dot = nullptr;  // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot)
-  bool use_vnni = false;         // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8)
-  bool use_q4a8 = false;       // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant
-  dsdot_fn dsdot;
-  mxgemv_fn mxgemv;
-  q4dot_fn q4dot;
-  fp8_block_dot_fn fp8_block_dot;
-  // fp8_block: bf16 scale per 128x128 weight block. These are the scale tensor's
-  // trailing strides, including any padding used by the host bank layout.
-  int fp8_block_gu_scale_row = 0, fp8_block_dn_scale_row = 0;
-  // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU
-  // (same reference grid), so submit() must not repeat it on the host-callback
-  // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL
-  // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving
-  // it to a captured GPU elementwise kernel removes it while keeping the official
-  // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py).
-  bool input_prequant = false;
-  // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block.
-  int act_block = 128;
-  // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I).
-  int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0;
-  float e2m1_lut[16];
-  float e4m3_lut[256];
-  float e8m0_lut[256];         // mxfp4 block scale: 2^(s-127), s clamped to [0,254]
-  const char* isa;
-
-  std::vector<bf16_t> g_scratch;   // [max_tokens * top_k * I] intermediate
-  std::vector<bf16_t> xq_scratch;  // [max_tokens * H] ds_fp4 fp8-roundtripped input
-  // nvfp4/ds_fp4/fp8_block activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot.
-  std::vector<float> xe_scratch, xo_scratch;  // [max_tokens * H/2]   (input)
-  std::vector<float> ge_scratch, go_scratch;  // [max_tokens*top_k*I/2] (intermediate)
-  // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale.
-  std::vector<int8_t> xi8_scratch, gi8_scratch;  // [max_tokens*H], [max_tokens*top_k*I]
-  std::vector<float> xas_scratch, gas_scratch;   // [max_tokens*H/16], [..*top_k*I/16]
-  std::string isa_str;
-
+  std::vector<int> core_ids;          // worker tid -> logical CPU to pin to (may be empty)
   std::vector<std::thread> workers;
   std::mutex task_mtx;
   std::condition_variable task_cv;
@@ -1421,7 +136,7 @@ struct CpuMoeExecutor {
 
   std::atomic<bool> stop{false};  // spinning workers read it without the lock
   uint64_t cur_gen = 0;
-  MoeTask* cur_task = nullptr;
+  MoeTask cur_task{};
   // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next.
   alignas(64) std::atomic<uint64_t> submitted{0};
   alignas(64) std::atomic<uint64_t> completed{0};
@@ -1431,264 +146,56 @@ struct CpuMoeExecutor {
   std::atomic<bool> spin_wait{false};
   bool coord_pinned = false;
 
+  // Work split of the current task.
   alignas(64) std::atomic<int64_t> p1_next{0};
   std::atomic<int64_t> p2_next{0};
-  std::atomic<int64_t> prt_next{0};  // intermediate-row prep phase (prep_g_row)
+  std::atomic<int64_t> prt_next{0};  // intermediate-row prep phase (prep_inter)
   int64_t p1_total = 0, p2_total = 0, prt_total = 0;
-  int n_iblk = 0, n_hblk = 0;
   std::atomic<int> done_count{0};
   std::atomic<int> bar_count{0};
   std::atomic<int> bar_sense{0};
 
-  std::vector<MoeTask*> owned_tasks;  // persistent task descriptors (graph-stable)
-  std::vector<int> core_ids;          // worker tid -> logical CPU to pin to (may be empty)
-
   // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ----
-  // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls
-  // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU
-  // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round
-  // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a
-  // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side
-  // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA:
-  // the same pointers are used by the GPU kernels and by this thread).
+  // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync.
+  // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model).
+  std::unique_ptr<GpuSync> gpu_sync;
+  std::unique_ptr<KernelIo> kernel_io;  // replaces the copies around the doorbell once enable_kernel_io() runs
   std::thread coord_thread;
   std::atomic<bool> coord_stop{false};
-  volatile int64_t* ready_flags = nullptr;  // GPU increments, this thread polls
-  volatile int64_t* done_flags = nullptr;   // this thread sets, GPU spin-waits
-  int coord_num_slots = 0;
-  std::vector<MoeTask*> flag_task;           // slot -> task (registered lazily)
-  std::vector<int64_t> flag_served;          // slot -> completed dispatch count (tests/debug)
-  std::mutex flag_task_mtx;
-
-  // Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that
-  // preceded the bump are visible" and "y stores are visible before done". Plain
-  // volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too
-  // (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier
-  // + TSO.)
-  static int64_t flag_load_acquire(const volatile int64_t* p) {
-#if defined(_MSC_VER)
-    const int64_t v = *p;
-    _ReadWriteBarrier();
-    return v;
-#else
-    return __atomic_load_n(const_cast<const int64_t*>(p), __ATOMIC_ACQUIRE);
-#endif
-  }
-
-  static void flag_store_release(volatile int64_t* p, int64_t v) {
-#if defined(_MSC_VER)
-    _ReadWriteBarrier();
-    *p = v;
-#else
-    __atomic_store_n(const_cast<int64_t*>(p), v, __ATOMIC_RELEASE);
-#endif
-  }
+  std::atomic<int64_t> flag_served{0};  // completed doorbells (tests, watchdog)
+  std::map<uint64_t, HostTask> host_tasks;  // host-function userData, keyed by doorbell word (graph-stable)
 
   CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_,
-                 int hidden_size, int inter_size, int max_tokens, int activation_id,
-                 int apply_router_weight_on_input, int weight_format,
-                 uintptr_t gate_up_ptr, uintptr_t down_ptr, uintptr_t gate_up_scale_ptr,
-                 uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr,
-                 uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr,
-                 uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_,
-                 std::vector<int> core_ids_, int fp8_block_gu_scale_stride_,
-                 int fp8_block_dn_scale_stride_)
-      : num_threads(num_threads_ > 0 ? num_threads_ : 1),
-        num_layers(num_layers_),
+                 int hidden_size, int inter_size, int max_tokens_, int activation_id,
+                 int apply_router_weight_on_input, const std::string& weight_format,
+                 std::map<std::string, uint64_t> banks, std::map<std::string, int64_t> options,
+                 double swiglu_alpha_, double swiglu_limit_,
+                 std::vector<int> core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr,
+                 uintptr_t y_ptr)
+      : num_layers(num_layers_),
         num_experts(num_experts_),
-        top_k(top_k_),
-        H(hidden_size),
-        I(inter_size),
-        act(activation_id),
-        apply_on_input(apply_router_weight_on_input),
-        fmt(weight_format),
-        gate_up_tbl(reinterpret_cast<const uint64_t*>(gate_up_ptr)),
-        down_tbl(reinterpret_cast<const uint64_t*>(down_ptr)),
-        gu_scale_tbl(reinterpret_cast<const uint64_t*>(gate_up_scale_ptr)),
-        gu_global_tbl(reinterpret_cast<const uint64_t*>(gate_up_global_ptr)),
-        dn_scale_tbl(reinterpret_cast<const uint64_t*>(down_scale_ptr)),
-        dn_global_tbl(reinterpret_cast<const uint64_t*>(down_global_ptr)),
-        gu_bias_tbl(reinterpret_cast<const uint64_t*>(gate_up_bias_ptr)),
-        dn_bias_tbl(reinterpret_cast<const uint64_t*>(down_bias_ptr)),
-        swiglu_alpha(static_cast<float>(swiglu_alpha_)),
-        swiglu_limit(static_cast<float>(swiglu_limit_)),
-        fp8_block_gu_scale_row(fp8_block_gu_scale_stride_),
-        fp8_block_dn_scale_row(fp8_block_dn_scale_stride_),
+        apply_on_input(apply_router_weight_on_input != 0),
+        params{top_k_, max_tokens_, hidden_size, inter_size, activation_id, static_cast<float>(swiglu_alpha_),
+               static_cast<float>(swiglu_limit_)},
+        io_x(reinterpret_cast<const bf16_t*>(x_ptr)),
+        io_ids(reinterpret_cast<const int32_t*>(ids_ptr)),
+        io_w(reinterpret_cast<const float*>(w_ptr)),
+        io_y(reinterpret_cast<bf16_t*>(y_ptr)),
+        num_threads(num_threads_ > 0 ? num_threads_ : 1),
         core_ids(std::move(core_ids_)) {
-    DotChoice c = select_dot();
-    dot = c.fn;
-    nvdot = select_nvdot();
-    dsdot = select_dsdot();
-    mxgemv = select_mxgemv();
-    q4dot = select_q4dot();
-    fp8_block_dot = select_fp8_block_dot();
-    if (weight_format == WF_Q4_0) {
-      if (H % 32 != 0 || I % 32 != 0)
-        throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32");
-      q4_gu_row_bytes = (H / 32) * 18;  // K = H (gate_up rows)
-      q4_dn_row_bytes = (I / 32) * 18;  // K = I (down rows)
-    }
-    isa = c.name;
-    // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always
-    // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW
-    // / scalar for the tier, so the tag reflects which of those q4dot resolved to.
-    nvi8dot = select_nvi8dot();
-    use_vnni = (weight_format == WF_NVFP4) && (nvi8dot != nullptr);
-    use_q4a8 = (weight_format == WF_Q4_0);
-    const char* q4tag = use_q4a8 ? (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8") : "";
-    const char* vnni_tag =
-        cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)";
-    isa_str = std::string(c.name) + (use_vnni ? vnni_tag : "") + q4tag;
-    isa = isa_str.c_str();
-    for (int i = 0; i < 16; ++i) e2m1_lut[i] = kE2M1[i];
-    for (int i = 0; i < 256; ++i) e4m3_lut[i] = e4m3_decode((uint8_t)i);
-    // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254].
-    for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127);
-    g_scratch.assign(static_cast<size_t>(max_tokens) * top_k * I, 0);
-    // Row-major fp4 (nvfp4/ds_fp4) and fp8_block pre-deinterleave activations to fp32 even/odd.
-    needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4 || fmt == WF_FP8_BLOCK);
-    if (needs_di) {
-      if (fmt == WF_DSFP4) xq_scratch.assign(static_cast<size_t>(max_tokens) * H, 0);
-      xe_scratch.assign(static_cast<size_t>(max_tokens) * (H / 2), 0);
-      xo_scratch.assign(static_cast<size_t>(max_tokens) * (H / 2), 0);
-      ge_scratch.assign(static_cast<size_t>(max_tokens) * top_k * (I / 2), 0);
-      go_scratch.assign(static_cast<size_t>(max_tokens) * top_k * (I / 2), 0);
-      if (use_vnni) {
-        xi8_scratch.assign(static_cast<size_t>(max_tokens) * H, 0);
-        xas_scratch.assign(static_cast<size_t>(max_tokens) * (H / 16), 0);
-        gi8_scratch.assign(static_cast<size_t>(max_tokens) * top_k * I, 0);
-        gas_scratch.assign(static_cast<size_t>(max_tokens) * top_k * (I / 16), 0);
-      }
-    }
-    // q4_0 W4A8: per-32-block Q8_0 activations (int8 + fp32 scale) for input + intermediate.
-    if (use_q4a8) {
-      xi8_scratch.assign(static_cast<size_t>(max_tokens) * H, 0);
-      xas_scratch.assign(static_cast<size_t>(max_tokens) * (H / 32), 0);
-      gi8_scratch.assign(static_cast<size_t>(max_tokens) * top_k * I, 0);
-      gas_scratch.assign(static_cast<size_t>(max_tokens) * top_k * (I / 32), 0);
-    }
+    if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK));
+    fmt = make_format(weight_format, params, FormatArgs{weight_format, std::move(banks), std::move(options)});
+    preps_inter = fmt->preps_inter();
+    rounds_routes = fmt->rounds_routes();
+    n_iblk = (params.I + kRowBlock - 1) / kRowBlock;
+    n_hblk = (params.H + kRowBlock - 1) / kRowBlock;
+    g_scratch.assign(static_cast<size_t>(params.max_tokens) * params.top_k * params.I, 0);
     for (int t = 0; t < num_threads; ++t)
       workers.emplace_back([this, t] { worker_loop(t); });
   }
 
-  // Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in
-  // natural order + one fp32 scale (amax/127) per block. Done once per token/route,
-  // amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0.
-  void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) {
-    const int nb = K / 32;
-    for (int b = 0; b < nb; ++b) {
-      const bf16_t* xb = x + (size_t)b * 32;
-      float xf[32], amax = 0.0f;
-      for (int j = 0; j < 32; ++j) {
-        xf[j] = bf16_to_f32(xb[j]);
-        amax = std::max(amax, std::fabs(xf[j]));
-      }
-      const float d = amax > 0.0f ? amax / 127.0f : 1.0f;
-      asb[b] = d;
-      const float inv = amax > 0.0f ? 1.0f / d : 0.0f;
-      int8_t* o = aq + (size_t)b * 32;
-      for (int j = 0; j < 32; ++j)
-        o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv)));
-    }
-  }
-
-  // Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the
-  // [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over
-  // every output row), so a scalar pass is fine relative to the GEMV.
-  void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) {
-    const int nb = K / 16;
-    for (int b = 0; b < nb; ++b) {
-      const float* xeb = xe + (size_t)b * 8;
-      const float* xob = xo + (size_t)b * 8;
-      float amax = 0.0f;
-      for (int j = 0; j < 8; ++j)
-        amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j])));
-      const float s = amax > 0.0f ? amax / 127.0f : 1.0f;
-      asb[b] = s;
-      const float inv = 1.0f / s;
-      int8_t* ae = asi8 + (size_t)b * 16;
-      for (int j = 0; j < 8; ++j) {
-        int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv);
-        ae[j] = (int8_t)std::max(-127, std::min(127, qe));
-        ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo));
-      }
-    }
-  }
-
-  // gate_up output row `row` (in [0, 2I)) dotted with activation over K = H. ``e`` is
-  // the layer-local expert row (0..num_experts); the layer bases (already resolved
-  // once per task/pass by the caller via tbl_at) pick the layer's own tensors.
-  // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves
-  // (or, with AVX-VNNI, the per-16-block int8 activations).
-  // fp8_block uses the halves for whole 128-blocks and the bf16 row for a ragged tail.
-  inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l,
-                         const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row,
-                         const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8,
-                         const float* xas) {
-    if (fmt == WF_BF16) {
-      const bf16_t* w = gate_up_l + ((size_t)e * (2 * I) + row) * H;
-      return dot(w, x, H);
-    }
-    if (fmt == WF_Q4_0) {
-      const uint8_t* w =
-          gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes;
-      return q4dot(w, xi8, xas, H);  // W4A8: int8 activations (Q8_0), scale in xas
-    }
-    if (fmt == WF_FP8_BLOCK) {
-      const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)H;
-      const bf16_t* sc = reinterpret_cast<const bf16_t*>(gu_scale_l) +
-                         ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) *
-                             (size_t)fp8_block_gu_scale_row;
-      return fp8_block_dot(w, xe, xo, x, sc, H, e4m3_lut);
-    }
-    const size_t r = (size_t)e * (2 * I) + row;
-    if (use_vnni)
-      return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16),
-                     fp16_to_f32(gu_global_l[r]), xi8, H, e4m3_lut, xas);
-    return nvdot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16),
-                 fp16_to_f32(gu_global_l[r]), xe, xo, H, e2m1_lut, e4m3_lut);
-  }
-
-  // down output row `row` (in [0, H)) dotted with the intermediate over K = I. Same
-  // layer-local-base convention as gemm1_dot.
-  inline float gemm2_dot(const bf16_t* down_l, const uint8_t* dn_packed_l,
-                         const uint8_t* dn_scale_l, const uint16_t* dn_global_l, int e, int row,
-                         const bf16_t* g, const float* ge, const float* go, const int8_t* gi8,
-                         const float* gas) {
-    if (fmt == WF_BF16) {
-      const bf16_t* w = down_l + ((size_t)e * H + row) * I;
-      return dot(w, g, I);
-    }
-    if (fmt == WF_Q4_0) {
-      const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes;
-      return q4dot(w, gi8, gas, I);  // W4A8: int8 activations (Q8_0), scale in gas
-    }
-    if (fmt == WF_FP8_BLOCK) {
-      const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)I;
-      const bf16_t* sc = reinterpret_cast<const bf16_t*>(dn_scale_l) +
-                         ((size_t)e * ((H + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) *
-                             (size_t)fp8_block_dn_scale_row;
-      return fp8_block_dot(w, ge, go, g, sc, I, e4m3_lut);
-    }
-    const size_t r = (size_t)e * H + row;
-    if (use_vnni)
-      return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16),
-                     fp16_to_f32(dn_global_l[r]), gi8, I, e4m3_lut, gas);
-    return nvdot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16),
-                 fp16_to_f32(dn_global_l[r]), ge, go, I, e2m1_lut, e4m3_lut);
-  }
-
   void pin_self(int tid) {
-#if CPU_MOE_HAS_AFFINITY
-    if (core_ids.empty()) return;
-    const int cpu = core_ids[tid % static_cast<int>(core_ids.size())];
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    CPU_SET(cpu, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
-#else
-    (void)tid;
-#endif
+    if (!core_ids.empty()) pin_this_thread(core_ids[tid % static_cast<int>(core_ids.size())]);
   }
 
   ~CpuMoeExecutor() {
@@ -1701,30 +208,23 @@ struct CpuMoeExecutor {
     task_cv.notify_all();
     for (auto& th : workers)
       if (th.joinable()) th.join();
-    for (MoeTask* t : owned_tasks) delete t;
   }
 
-  uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr,
-                        uintptr_t ids_ptr, uintptr_t w_ptr, uintptr_t y_ptr) {
-    MoeTask* t = new MoeTask{this,
-                             layer_id,
-                             num_tokens,
-                             reinterpret_cast<const bf16_t*>(x_ptr),
-                             reinterpret_cast<const int32_t*>(ids_ptr),
-                             reinterpret_cast<const float*>(w_ptr),
-                             reinterpret_cast<bf16_t*>(y_ptr)};
-    owned_tasks.push_back(t);
-    return reinterpret_cast<uintptr_t>(t);
+  MoeTask task(int layer_id, int num_tokens) const {
+    if (layer_id < 0 || layer_id >= num_layers || num_tokens < 1 || num_tokens > params.max_tokens)
+      throw std::out_of_range("CPU MoE task (layer " + std::to_string(layer_id) + ", " +
+                              std::to_string(num_tokens) + " tokens) is outside the executor's layers or IO");
+    return MoeTask{static_cast<uint64_t>(layer_id), static_cast<uint64_t>(num_tokens)};
   }
 
-  const char* isa_name() const { return isa; }
+  std::string isa_name() const { return fmt->isa_name(); }
 
-  const char* fp8_block_isa_name() const {
-#if CPU_MOE_X86
-    if (fp8_block_dot == dot_fp8_block_avx512f) return "avx512f";
-    if (fp8_block_dot == dot_fp8_block_avx2) return "avx2";
-#endif
-    return "scalar";
+  // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others.
+  bool has_routes(const MoeTask& t) const {
+    const int64_t n = static_cast<int64_t>(t.num_tokens) * params.top_k;
+    for (int64_t i = 0; i < n; ++i)
+      if (io_ids[i] >= 0 && io_ids[i] < num_experts) return true;
+    return false;
   }
 
   void barrier(int& local_sense) {
@@ -1733,275 +233,58 @@ struct CpuMoeExecutor {
       bar_count.store(0);
       bar_sense.store(local_sense);
     } else {
-      while (bar_sense.load() != local_sense) {
-#if CPU_MOE_X86
-        _mm_pause();
-#endif
-      }
+      while (bar_sense.load() != local_sense) cpu_relax();
     }
   }
 
-  void do_pass1(const MoeTask* t, int64_t p) {
-    if (fmt == WF_MXFP4) {
-      do_pass1_mxfp4(t, p);
-      return;
-    }
-    if (fmt == WF_DSFP4) {
-      do_pass1_dsfp4(t, p);
-      return;
-    }
+  // Pass 1 item p = one kRowBlock block of one (token, route)'s intermediate row.
+  void do_pass1(const MoeTask& t, int64_t p) {
+    const int top_k = params.top_k, H = params.H, I = params.I;
     const int64_t ib = p % n_iblk;
     const int64_t tk = p / n_iblk;
     const int k = static_cast<int>(tk % top_k);
     const int tok = static_cast<int>(tk / top_k);
-    const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
+    const int e = io_ids[static_cast<size_t>(tok) * top_k + k];
     if (e < 0 || e >= num_experts) return;
-    const float w_in = apply_on_input ? t->w[static_cast<size_t>(tok) * top_k + k] : 1.0f;
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const bf16_t* gate_up_l = reinterpret_cast<const bf16_t*>(tbl_at(gate_up_tbl, t->layer_id));
-    const uint8_t* gu_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(gate_up_tbl, t->layer_id));
-    const uint8_t* gu_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(gu_scale_tbl, t->layer_id));
-    const uint16_t* gu_global_l =
-        reinterpret_cast<const uint16_t*>(tbl_at(gu_global_tbl, t->layer_id));
-    const bf16_t* x_row = t->x + (size_t)tok * H;
-    const float* xe = needs_di ? xe_scratch.data() + (size_t)tok * (H / 2) : nullptr;
-    const float* xo = needs_di ? xo_scratch.data() + (size_t)tok * (H / 2) : nullptr;
-    const int8_t* xi8 =
-        (use_vnni || use_q4a8) ? xi8_scratch.data() + (size_t)tok * H : nullptr;
-    const float* xas = use_vnni ? xas_scratch.data() + (size_t)tok * (H / 16)
-                     : use_q4a8 ? xas_scratch.data() + (size_t)tok * (H / 32)
-                                  : nullptr;
-    bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I;
-    const int i0 = static_cast<int>(ib) * IBLK;
-    const int i1 = std::min(I, i0 + IBLK);
-    const bool clamped = act == ACT_SWIGLUOAI || act == ACT_SWIGLU_CLAMP;
-    const float up_bias = act == ACT_SWIGLUOAI ? 1.0f : 0.0f;
-    const float lim = swiglu_limit, alpha = swiglu_alpha;
-    for (int i = i0; i < i1; ++i) {
-      // gate = row i, up = row I+i
-      float gate =
-          gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, i, x_row, xe, xo, xi8, xas) * w_in;
-      float up = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, I + i, x_row,
-                           xe, xo, xi8, xas) * w_in;
-      if (clamped) {
-        // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias)
-        // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp
-        // (GLM-5.3) does not. lim == +inf: no clamp.
-        if (gate > lim) gate = lim;
-        if (up > lim) up = lim;
-        else if (up < -lim) up = -lim;
-        const float glu = gate / (1.0f + std::exp(-gate * alpha));
-        g_row[i] = f32_to_bf16(glu * (up + up_bias));
-      } else {
-        g_row[i] = f32_to_bf16(act_apply(act, gate) * up);
-      }
-    }
+    const float w_in = apply_on_input ? io_w[static_cast<size_t>(tok) * top_k + k] : 1.0f;
+    const int i0 = static_cast<int>(ib) * kRowBlock;
+    const int i1 = std::min(I, i0 + kRowBlock);
+    fmt->gate_up(static_cast<int>(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in,
+                 g_scratch.data() + ((size_t)tok * top_k + k) * I);
   }
 
-  void do_pass2(const MoeTask* t, int64_t p) {
-    if (fmt == WF_MXFP4) {
-      do_pass2_mxfp4(t, p);
-      return;
-    }
-    if (fmt == WF_DSFP4) {
-      do_pass2_dsfp4(t, p);
-      return;
-    }
-    const int64_t hb = p % n_hblk;
-    const int tok = static_cast<int>(p / n_hblk);
-    const int h0 = static_cast<int>(hb) * HBLK;
-    const int h1 = std::min(H, h0 + HBLK);
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const bf16_t* down_l = reinterpret_cast<const bf16_t*>(tbl_at(down_tbl, t->layer_id));
-    const uint8_t* dn_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(down_tbl, t->layer_id));
-    const uint8_t* dn_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(dn_scale_tbl, t->layer_id));
-    const uint16_t* dn_global_l =
-        reinterpret_cast<const uint16_t*>(tbl_at(dn_global_tbl, t->layer_id));
-    bf16_t* y_row = t->y + (size_t)tok * H;
-    for (int h = h0; h < h1; ++h) {
-      float acc = 0.0f;
-      for (int k = 0; k < top_k; ++k) {
-        const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
-        if (e < 0 || e >= num_experts) continue;
-        const float w_out = apply_on_input ? 1.0f : t->w[static_cast<size_t>(tok) * top_k + k];
-        const size_t gr = (size_t)tok * top_k + k;
-        const bf16_t* g_row = g_scratch.data() + gr * I;
-        const float* ge = needs_di ? ge_scratch.data() + gr * (I / 2) : nullptr;
-        const float* go = needs_di ? go_scratch.data() + gr * (I / 2) : nullptr;
-        const int8_t* gi8 = (use_vnni || use_q4a8) ? gi8_scratch.data() + gr * I : nullptr;
-        const float* gas = use_vnni ? gas_scratch.data() + gr * (I / 16)
-                         : use_q4a8 ? gas_scratch.data() + gr * (I / 32)
-                                      : nullptr;
-        acc += gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, gi8,
-                         gas) * w_out;
-      }
-      y_row[h] = f32_to_bf16(acc);
-    }
-  }
-
-  // ----------------------------- mxfp4 (gpt-oss) -----------------------------
-  // Transposed split-K layout (N innermost), so the GEMV streams K and accumulates
-  // a contiguous block of N output columns -> cache-line-efficient, no repack, no
-  // extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 fuses
-  // down(+bias) * router-weight, summed over the token's routes.
-  //
-  // Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble
-  // first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py.
-
-  void do_pass1_mxfp4(const MoeTask* t, int64_t p) {
-    const int64_t ib = p % n_iblk;
-    const int64_t tk = p / n_iblk;
-    const int k = static_cast<int>(tk % top_k);
-    const int tok = static_cast<int>(tk / top_k);
-    const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
-    if (e < 0 || e >= num_experts) return;
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const uint8_t* gu_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(gate_up_tbl, t->layer_id));
-    const uint8_t* gu_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(gu_scale_tbl, t->layer_id));
-    const bf16_t* gu_bias_l = reinterpret_cast<const bf16_t*>(tbl_at(gu_bias_tbl, t->layer_id));
-    const int N2 = 2 * I;            // gate_up output width (gate/up interleaved)
-    const int Hh = H / 2;            // packed-K rows
-    const int i0 = static_cast<int>(ib) * IBLK;
-    const int i1 = std::min(I, i0 + IBLK);
-    const int nunit = i1 - i0;       // intermediate units owned by this tile
-    const int col0 = 2 * i0;         // first gate_up column
-    const int ncol = 2 * nunit;      // gate_up columns owned by this tile
-    const bf16_t* x_row = t->x + (size_t)tok * H;
-    const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2;
-    const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2;
-    float gu[2 * IBLK];
-    mxgemv(gu, blk_e + col0, scl_e + col0, x_row, Hh, N2, ncol, e2m1_lut, e8m0_lut);
-    const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0;
-    bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I;
-    const float lim = swiglu_limit, alpha = swiglu_alpha;
-    for (int j = 0; j < nunit; ++j) {
-      float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]);
-      float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]);
-      if (gate > lim) gate = lim;
-      if (up > lim) up = lim;
-      else if (up < -lim) up = -lim;
-      const float glu = gate / (1.0f + std::exp(-gate * alpha));  // gate * sigmoid(alpha*gate)
-      g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f));
-    }
-  }
-
-  void do_pass2_mxfp4(const MoeTask* t, int64_t p) {
+  // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes.
+  void do_pass2(const MoeTask& t, int64_t p) {
+    const int top_k = params.top_k, H = params.H, I = params.I;
     const int64_t hb = p % n_hblk;
     const int tok = static_cast<int>(p / n_hblk);
-    const int Ih = I / 2;            // packed-K rows
-    const int h0 = static_cast<int>(hb) * HBLK;
-    const int h1 = std::min(H, h0 + HBLK);
-    const int nh = h1 - h0;
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const uint8_t* dn_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(down_tbl, t->layer_id));
-    const uint8_t* dn_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(dn_scale_tbl, t->layer_id));
-    const bf16_t* dn_bias_l = reinterpret_cast<const bf16_t*>(tbl_at(dn_bias_tbl, t->layer_id));
-    float acc[HBLK];
-    for (int c = 0; c < nh; ++c) acc[c] = 0.0f;
+    const int h0 = static_cast<int>(hb) * kRowBlock;
+    const int h1 = std::min(H, h0 + kRowBlock);
+    DownRoute routes[kMaxTopK];
+    float w_out[kMaxTopK];
+    int n = 0;
     for (int k = 0; k < top_k; ++k) {
-      const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
+      const int e = io_ids[static_cast<size_t>(tok) * top_k + k];
       if (e < 0 || e >= num_experts) continue;
-      const float wt = t->w[static_cast<size_t>(tok) * top_k + k];
-      const bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I;
-      const uint8_t* blk_e = dn_packed_l + (size_t)e * Ih * H;
-      const uint8_t* scl_e = dn_scale_l + (size_t)e * (size_t)(I / 32) * H;
-      float part[HBLK];
-      mxgemv(part, blk_e + h0, scl_e + h0, g_row, Ih, H, nh, e2m1_lut, e8m0_lut);
-      const bf16_t* bias_e = dn_bias_l + (size_t)e * H + h0;
-      for (int c = 0; c < nh; ++c) acc[c] += (part[c] + bf16_to_f32(bias_e[c])) * wt;
-    }
-    bf16_t* y_row = t->y + (size_t)tok * H;
-    for (int c = 0; c < nh; ++c) y_row[h0 + c] = f32_to_bf16(acc[c]);
-  }
-
-  // ----------------------------- ds_fp4 (DSV4) -------------------------------
-  // Row-major e2m1 + e8m0/32 (no global); silu-swiglu with clamp; FP8-roundtripped
-  // activations (x once in submit -> xq_scratch; the intermediate g in a dedicated
-  // round-trip phase between the two passes). Router weight applies on the down output.
-
-  void do_pass1_dsfp4(const MoeTask* t, int64_t p) {
-    const int64_t ib = p % n_iblk;
-    const int64_t tk = p / n_iblk;
-    const int k = static_cast<int>(tk % top_k);
-    const int tok = static_cast<int>(tk / top_k);
-    const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
-    if (e < 0 || e >= num_experts) return;
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const uint8_t* gu_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(gate_up_tbl, t->layer_id));
-    const uint8_t* gu_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(gu_scale_tbl, t->layer_id));
-    const int N2 = 2 * I, Hh = H / 2, Hs = H / 32;
-    // fp8-roundtripped input, pre-deinterleaved to fp32 even/odd halves.
-    const float* xe = xe_scratch.data() + (size_t)tok * (H / 2);
-    const float* xo = xo_scratch.data() + (size_t)tok * (H / 2);
-    const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh;
-    const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs;
-    bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I;
-    const int i0 = static_cast<int>(ib) * IBLK;
-    const int i1 = std::min(I, i0 + IBLK);
-    const float lim = swiglu_limit;
-    for (int i = i0; i < i1; ++i) {
-      // gate_up is stored bf16 by the reference GEMV before swiglu; round to match.
-      float gate = bf16_to_f32(f32_to_bf16(
-          dsdot(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut, e8m0_lut)));
-      float up = bf16_to_f32(f32_to_bf16(dsdot(
-          gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut, e8m0_lut)));
-      if (lim > 0.0f) {
-        if (gate > lim) gate = lim;
-        if (up > lim) up = lim;
-        else if (up < -lim) up = -lim;
-      }
-      const float glu = gate / (1.0f + std::exp(-gate));  // silu(gate)
-      g_row[i] = f32_to_bf16(glu * up);
-    }
-  }
-
-  // Prepare one intermediate row (token,route) for the down GEMV, reused across every down output row.
-  // ds_fp4 first FP8 round-trips it (DSV4 act_quant); nvfp4/ds_fp4/fp8_block deinterleave it to fp32 even/odd; q4_0 Q8_0-quantizes it.
-  void prep_g_row(int64_t r) {
-    bf16_t* g = g_scratch.data() + (size_t)r * I;
-    if (use_q4a8) {  // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV.
-      quant_q8_0(g, I, gi8_scratch.data() + (size_t)r * I,
-                 gas_scratch.data() + (size_t)r * (I / 32));
-      return;
-    }
-    if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block);
-    float* ge = ge_scratch.data() + (size_t)r * (I / 2);
-    float* go = go_scratch.data() + (size_t)r * (I / 2);
-    deinterleave_bf16_f32(g, ge, go, I);
-    if (use_vnni)
-      quant_i8_pg16(ge, go, I, gi8_scratch.data() + (size_t)r * I,
-                    gas_scratch.data() + (size_t)r * (I / 16));
-  }
-
-  void do_pass2_dsfp4(const MoeTask* t, int64_t p) {
-    const int64_t hb = p % n_hblk;
-    const int tok = static_cast<int>(p / n_hblk);
-    const int Ih = I / 2, Is = I / 32;
-    const int h0 = static_cast<int>(hb) * HBLK;
-    const int h1 = std::min(H, h0 + HBLK);
-    // Resolve this task's layer base once; row indexing below is layer-local (e).
-    const uint8_t* dn_packed_l = reinterpret_cast<const uint8_t*>(tbl_at(down_tbl, t->layer_id));
-    const uint8_t* dn_scale_l = reinterpret_cast<const uint8_t*>(tbl_at(dn_scale_tbl, t->layer_id));
-    bf16_t* y_row = t->y + (size_t)tok * H;
+      const size_t r = (size_t)tok * top_k + k;
+      routes[n] = {e, static_cast<int64_t>(r), g_scratch.data() + r * I};
+      w_out[n] = apply_on_input ? 1.0f : io_w[r];
+      ++n;
+    }
+    float part[kMaxTopK * kRowBlock];
+    if (n > 0) fmt->down(static_cast<int>(t.layer), h0, h1, routes, n, part);
+    bf16_t* y_row = io_y + (size_t)tok * H;
     for (int h = h0; h < h1; ++h) {
       float acc = 0.0f;
-      for (int k = 0; k < top_k; ++k) {
-        const int e = t->ids[static_cast<size_t>(tok) * top_k + k];
-        if (e < 0 || e >= num_experts) continue;
-        const float wt = t->w[static_cast<size_t>(tok) * top_k + k];
-        const float* ge = ge_scratch.data() + ((size_t)tok * top_k + k) * (I / 2);
-        const float* go = go_scratch.data() + ((size_t)tok * top_k + k) * (I / 2);
-        const uint8_t* dp = dn_packed_l + (size_t)e * (size_t)H * Ih + (size_t)h * Ih;
-        const uint8_t* ds = dn_scale_l + (size_t)e * (size_t)H * Is + (size_t)h * Is;
-        // The reference rounds each route's weighted down output to bf16 before the
-        // fp32 sum over routes (down [T, top_k, H] bf16 -> .sum(dim=1)).
-        acc += bf16_to_f32(f32_to_bf16(dsdot(dp, ds, ge, go, I, e2m1_lut, e8m0_lut) * wt));
+      for (int j = 0; j < n; ++j) {
+        const float v = part[j * kRowBlock + h - h0] * w_out[j];
+        acc += rounds_routes ? bf16_to_f32(f32_to_bf16(v)) : v;
       }
       y_row[h] = f32_to_bf16(acc);
     }
   }
 
-  void run_task_body(const MoeTask* t) {
+  void run_task_body(const MoeTask& t) {
     int local_sense = 0;
     for (;;) {
       int64_t p = p1_next.fetch_add(1, std::memory_order_relaxed);
@@ -2009,13 +292,13 @@ struct CpuMoeExecutor {
       do_pass1(t, p);
     }
     barrier(local_sense);
-    // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_g_row.
+    // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter.
     // Needs all of pass1 done (a full row spans every iblk).
-    if (needs_di || use_q4a8) {
+    if (preps_inter) {
       for (;;) {
         int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed);
         if (r >= prt_total) break;
-        prep_g_row(r);
+        fmt->prep_inter(r, g_scratch.data() + (size_t)r * params.I);
       }
       barrier(local_sense);
     }
@@ -2051,8 +334,7 @@ struct CpuMoeExecutor {
       }
       // submit() writes cur_task before the release store to `submitted`.
       my_gen = gen;
-      MoeTask* t = cur_task;
-      run_task_body(t);
+      run_task_body(cur_task);
       last_active = wclk::now();
       if (done_count.fetch_add(1) + 1 == num_threads) {
         completed.store(my_gen, std::memory_order_release);
@@ -2064,69 +346,20 @@ struct CpuMoeExecutor {
     }
   }
 
-  void submit(MoeTask* t) {
-    n_iblk = (I + IBLK - 1) / IBLK;
-    n_hblk = (H + HBLK - 1) / HBLK;
-    // Grow the per-token intermediate scratch if a larger batch shows up than the
-    // construction-time hint (CUDA-graph capture warms the largest bs first, so
-    // this happens at most once, before any capture, while the pool is idle).
-    const size_t need = static_cast<size_t>(t->num_tokens) * top_k * I;
-    if (need > g_scratch.size()) g_scratch.resize(need);
-    p1_total = static_cast<int64_t>(t->num_tokens) * top_k * n_iblk;
-    p2_total = static_cast<int64_t>(t->num_tokens) * n_hblk;
-    prt_total = (needs_di || use_q4a8) ? static_cast<int64_t>(t->num_tokens) * top_k : 0;
+  void submit(MoeTask t) {
+    const int num_tokens = static_cast<int>(t.num_tokens);
+    p1_total = static_cast<int64_t>(num_tokens) * params.top_k * n_iblk;
+    p2_total = static_cast<int64_t>(num_tokens) * n_hblk;
+    prt_total = preps_inter ? static_cast<int64_t>(num_tokens) * params.top_k : 0;
     p1_next.store(0, std::memory_order_relaxed);
     p2_next.store(0, std::memory_order_relaxed);
     prt_next.store(0, std::memory_order_relaxed);
     done_count.store(0, std::memory_order_relaxed);
     bar_count.store(0, std::memory_order_relaxed);
     bar_sense.store(0, std::memory_order_relaxed);
-    // Deinterleave the per-token input once, up front (ds_fp4 FP8-round-trips it first).
+    // Prepare the per-token input once, up front.
     // Single-threaded: tiny for decode, and done before the workers are woken below.
-    if (needs_di) {
-      const size_t xn = static_cast<size_t>(t->num_tokens) * H;
-      if (xn / 2 > xe_scratch.size()) {
-        xe_scratch.resize(xn / 2);
-        xo_scratch.resize(xn / 2);
-        ge_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * (I / 2));
-        go_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * (I / 2));
-        if (use_vnni) {
-          xi8_scratch.resize(xn);
-          xas_scratch.resize(xn / 16);
-          gi8_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * I);
-          gas_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * (I / 16));
-        }
-      }
-      const bool ds = (fmt == WF_DSFP4) && !input_prequant;
-      if (ds && xn > xq_scratch.size()) xq_scratch.resize(xn);
-      for (int tok = 0; tok < t->num_tokens; ++tok) {
-        const bf16_t* src = t->x + (size_t)tok * H;
-        if (ds) {  // DSV4 FP8-round-trips the input before the gate_up GEMV
-          bf16_t* xq = xq_scratch.data() + (size_t)tok * H;
-          fp8_roundtrip_bf16(src, xq, H, act_block);
-          src = xq;
-        }
-        float* xe = xe_scratch.data() + (size_t)tok * (H / 2);
-        float* xo = xo_scratch.data() + (size_t)tok * (H / 2);
-        deinterleave_bf16_f32(src, xe, xo, H);
-        if (use_vnni)
-          quant_i8_pg16(xe, xo, H, xi8_scratch.data() + (size_t)tok * H,
-                        xas_scratch.data() + (size_t)tok * (H / 16));
-      }
-    }
-    // q4_0 W4A8: Q8_0-quantize the per-token input once (single-threaded, tiny for decode).
-    if (use_q4a8) {
-      const size_t xn = static_cast<size_t>(t->num_tokens) * H;
-      if (xn > xi8_scratch.size()) {
-        xi8_scratch.resize(xn);
-        xas_scratch.resize(xn / 32);
-        gi8_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * I);
-        gas_scratch.resize(static_cast<size_t>(t->num_tokens) * top_k * (I / 32));
-      }
-      for (int tok = 0; tok < t->num_tokens; ++tok)
-        quant_q8_0(t->x + (size_t)tok * H, H, xi8_scratch.data() + (size_t)tok * H,
-                   xas_scratch.data() + (size_t)tok * (H / 32));
-    }
+    for (int tok = 0; tok < num_tokens; ++tok) fmt->prep_input(tok, io_x + (size_t)tok * params.H);
     {
       std::lock_guard<std::mutex> lk(task_mtx);
       cur_task = t;
@@ -2153,62 +386,56 @@ struct CpuMoeExecutor {
     sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; });
   }
 
-  void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) {
-    cudaLaunchHostFunc(reinterpret_cast<cudaStream_t>(stream), &CpuMoeExecutor::submit_cb,
-                       reinterpret_cast<void*>(task));
+  void submit_with_cuda_stream(uintptr_t stream, int layer_id, int num_tokens) {
+    const MoeTask t = task(layer_id, num_tokens);
+    HostTask& h = host_tasks.try_emplace(to_word(t), HostTask{this, t}).first->second;
+    cudaLaunchHostFunc(reinterpret_cast<cudaStream_t>(stream), &CpuMoeExecutor::submit_cb, &h);
   }
 
-  void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) {
-    cudaLaunchHostFunc(reinterpret_cast<cudaStream_t>(stream), &CpuMoeExecutor::sync_cb,
-                       reinterpret_cast<void*>(task));
+  void sync_with_cuda_stream(uintptr_t stream) {
+    cudaLaunchHostFunc(reinterpret_cast<cudaStream_t>(stream), &CpuMoeExecutor::sync_cb, this);
   }
 
-  // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a
-  // flag bump.
-  void register_flag_task(int slot, uintptr_t task) {
-    std::lock_guard<std::mutex> lk(flag_task_mtx);
-    if (static_cast<int>(flag_task.size()) <= slot) flag_task.resize(slot + 1, nullptr);
-    flag_task[slot] = reinterpret_cast<MoeTask*>(task);
-  }
-
-  int64_t flag_served_count(int slot) const {
-    return (slot >= 0 && slot < static_cast<int>(flag_served.size())) ? flag_served[slot] : 0;
-  }
-
-  // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0
+  // Start the busy-poll coordinator over the doorbell. ``pin_core`` >= 0
   // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its
   // polling never migrates onto / contends with a GEMV worker's core.
-  void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots,
-                              int pin_core) {
-    ready_flags = reinterpret_cast<volatile int64_t*>(ready_ptr);
-    done_flags = reinterpret_cast<volatile int64_t*>(done_ptr);
-    coord_num_slots = num_slots;
-    {
-      std::lock_guard<std::mutex> lk(flag_task_mtx);
-      if (static_cast<int>(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr);
-    }
-    flag_served.assign(num_slots, 0);
+  void start_flag_coordinator(int pin_core) {
+    gpu_sync = std::make_unique<GpuSync>();
     coord_stop.store(false);
     coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0;
     coord_thread = std::thread([this, pin_core] {
-#if CPU_MOE_HAS_AFFINITY
-      if (pin_core >= 0) {
-        cpu_set_t set;
-        CPU_ZERO(&set);
-        CPU_SET(pin_core, &set);
-        pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
-      }
-#endif
+      if (pin_core >= 0) pin_this_thread(pin_core);
       coordinator_loop();
     });
   }
 
+  void flag_signal(uintptr_t stream, int layer_id, int num_tokens) {
+    gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens)));
+  }
+  void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); }
+  void enable_kernel_io() {
+    kernel_io = std::make_unique<KernelIo>(gpu_sync->ready_word(), gpu_sync->done_word(), io_x, io_ids, io_w, io_y, params.H,
+                                           params.top_k);
+  }
+  void flag_publish(uintptr_t stream, uintptr_t x, uintptr_t ids, uintptr_t w, int layer_id, int num_tokens) {
+    kernel_io->publish(stream, reinterpret_cast<const void*>(x), reinterpret_cast<const int32_t*>(ids),
+                       reinterpret_cast<const float*>(w), num_tokens, to_word(task(layer_id, num_tokens)));
+  }
+  void flag_combine(uintptr_t stream, uintptr_t out, int num_tokens, bool accumulate) {
+    kernel_io->combine(stream, reinterpret_cast<void*>(out), num_tokens, accumulate);
+  }
+  bool flag_pending() const { return gpu_sync && gpu_sync->pending(); }
+  // Complete a doorbell the coordinator never answered so the stream can make progress.
+  void poison_flag() {
+    if (gpu_sync) gpu_sync->cpu_done();
+  }
+
   void coordinator_loop() {
     // Idle backoff (edge-device friendliness): spin hot only while decode traffic is
-    // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot
-    // count and pause cost, so some configs dozed every token). 50 ms since the last
-    // flag comfortably covers intra- and inter-token gaps -- including host-func-only
-    // stretches (prefill bursts, --moe-cpu-layers subsets) -- so steady decode never
+    // flowing, with a TIME-based hot window (an iteration count varies 2-6x with pause
+    // cost, so some configs dozed every token). 50 ms since the last
+    // flag comfortably covers intra- and inter-token gaps -- including prefill bursts and
+    // --moe-cpu-layers gaps -- so steady decode never
     // sleeps and keeps the ~us-level wakeup. Past it, i.e. the engine is actually idle,
     // escalate to timed sleeps capped at 2 ms: a dozing coordinator costs <0.1% of a
     // core instead of 100%, and the only price is a <=2 ms discovery delay on the FIRST
@@ -2222,30 +449,19 @@ struct CpuMoeExecutor {
     int64_t sleep_us = 100;
     bool dozing = false;
     while (!coord_stop.load(std::memory_order_relaxed)) {
-      bool any = false;
-      for (int L = 0; L < coord_num_slots; ++L) {
-        // Binary handshake (memop-compatible: the GPU-side WAIT compares against an
-        // immediate baked at graph capture, so the protocol resets per step instead of
-        // counting). Acquire: everything the GPU made visible before setting ready --
-        // the D2H input copies -- is visible to the worker pool after this read.
-        if (flag_load_acquire(&ready_flags[L]) != 0) {
-          flag_store_release(&ready_flags[L], 0);  // consume this step's doorbell
-          MoeTask* t;
-          {
-            std::lock_guard<std::mutex> lk(flag_task_mtx);
-            t = (L < static_cast<int>(flag_task.size())) ? flag_task[L] : nullptr;
-          }
-          if (t != nullptr) {
-            submit(t);
-            sync(coord_pinned && spin_wait.load(std::memory_order_relaxed));
-          }
-          // Release: the workers' y stores are visible before the GPU sees done.
-          flag_store_release(&done_flags[L], 1);
-          if (L < static_cast<int>(flag_served.size())) ++flag_served[L];
-          any = true;
+      const uint64_t word = gpu_sync->cpu_poll();
+      if (word != 0) {
+        const MoeTask t = from_word(word);
+        if (has_routes(t)) {
+          submit(t);
+          sync(coord_pinned && spin_wait.load(std::memory_order_relaxed));
+        } else {
+          // What the pool would write for a task without routes, minus its fork and join.
+          std::memset(io_y, 0, sizeof(bf16_t) * t.num_tokens * params.H);
         }
-      }
-      if (any) {
+        // Release: the workers' y stores are visible before the GPU sees done.
+        gpu_sync->cpu_done();
+        flag_served.fetch_add(1, std::memory_order_relaxed);
         last_active = coord_clock::now();
         empty_polls = 0;
         sleep_us = 100;
@@ -2254,15 +470,13 @@ struct CpuMoeExecutor {
       }
       // Hot phase: pause-spin, consulting the clock only every 1024 empty polls (the
       // amortization must gate the CLOCK, not the sleep -- gating the sleep decision
-      // let 1023/1024 idle iterations hot-scan all slots, ~20% of a core at 992 slots).
+      // let 1023/1024 idle iterations hot-poll).
       // Doze phase: sleep on EVERY iteration until traffic returns; one wake-up scan
       // (~0.5 us) per 2 ms sleep is ~0.03% duty.
       if (!dozing) {
         if ((++empty_polls & 1023u) != 0 ||
             coord_clock::now() - last_active < kHotWindow) {
-#if CPU_MOE_X86
-          _mm_pause();
-#endif
+          cpu_relax();
           continue;
         }
         dozing = true;
@@ -2270,92 +484,79 @@ struct CpuMoeExecutor {
       std::this_thread::sleep_for(std::chrono::microseconds(sleep_us));
       sleep_us = std::min<int64_t>(sleep_us * 2, kSleepCapUs);
     }
-    // Teardown: release any in-flight (or future) spin-wait immediately so a replay
-    // caught mid-shutdown exits its sync kernel now instead of owning the watchdog
-    // stall. Runs before the destructor's join() returns, while the flag arrays are
-    // still alive on the Python side.
-    for (int L = 0; L < coord_num_slots; ++L) {
-      flag_store_release(&done_flags[L], INT64_MAX);
-    }
+    // Teardown: complete any in-flight wait so a replay caught mid-shutdown does not stall
+    // until the watchdog fires. Runs before the destructor's join() returns, while the doorbell is still alive.
+    gpu_sync->cpu_done();
   }
 
   // Eager (non-graph) path: run one task to completion on the pool.
-  void run_task(uintptr_t task) {
-    MoeTask* t = reinterpret_cast<MoeTask*>(task);
-    submit(t);
+  void run(int layer_id, int num_tokens) {
+    submit(task(layer_id, num_tokens));
     sync();
   }
 
   static void CUDART_CB submit_cb(void* ud) {
-    MoeTask* t = reinterpret_cast<MoeTask*>(ud);
-    t->exec->submit(t);
-  }
-  static void CUDART_CB sync_cb(void* ud) {
-    MoeTask* t = reinterpret_cast<MoeTask*>(ud);
-    t->exec->sync();
+    HostTask* h = reinterpret_cast<HostTask*>(ud);
+    h->exec->submit(h->task);
   }
+  static void CUDART_CB sync_cb(void* ud) { reinterpret_cast<CpuMoeExecutor*>(ud)->sync(); }
 };
 
 }  // namespace
+}  // namespace freetoken::cpu_moe
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   namespace py = pybind11;
+  using freetoken::cpu_moe::CpuMoeExecutor;
+  using freetoken::cpu_moe::GpuSync;
   py::class_<CpuMoeExecutor>(m, "CpuMoeExecutor")
-      .def(py::init<int, int, int, int, int, int, int, int, int, int, uintptr_t, uintptr_t,
-                    uintptr_t, uintptr_t, uintptr_t, uintptr_t, uintptr_t, uintptr_t,
-                    double, double, std::vector<int>, int, int>(),
+      .def(py::init<int, int, int, int, int, int, int, int, int, std::string,
+                    std::map<std::string, uint64_t>, std::map<std::string, int64_t>,
+                    double, double, std::vector<int>, uintptr_t, uintptr_t, uintptr_t, uintptr_t>(),
            py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"),
            py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"),
            py::arg("max_tokens"), py::arg("activation_id"),
            py::arg("apply_router_weight_on_input"), py::arg("weight_format"),
-           py::arg("gate_up_ptr"), py::arg("down_ptr"), py::arg("gate_up_scale_ptr"),
-           py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"),
-           py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"),
-           py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"),
-           py::arg("core_ids"), py::arg("fp8_block_gu_scale_stride") = 0,
-           py::arg("fp8_block_dn_scale_stride") = 0)
-      .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"),
-           py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"),
+           py::arg("banks"), py::arg("options"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"),
+           py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"),
            py::arg("y_ptr"))
       .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream,
-           py::arg("stream"), py::arg("task"), py::call_guard<py::gil_scoped_release>())
+           py::arg("stream"), py::arg("layer_id"), py::arg("num_tokens"),
+           py::call_guard<py::gil_scoped_release>())
       .def("sync_with_cuda_stream", &CpuMoeExecutor::sync_with_cuda_stream,
-           py::arg("stream"), py::arg("task"), py::call_guard<py::gil_scoped_release>())
-      .def("run_task", &CpuMoeExecutor::run_task, py::arg("task"),
+           py::arg("stream"), py::call_guard<py::gil_scoped_release>())
+      .def("run", &CpuMoeExecutor::run, py::arg("layer_id"), py::arg("num_tokens"),
            py::call_guard<py::gil_scoped_release>())
-      .def("register_flag_task", &CpuMoeExecutor::register_flag_task,
-           py::arg("slot"), py::arg("task"))
-      .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot"))
-      .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator,
-           py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"),
-           py::arg("pin_core"))
-      .def("set_input_prequant",
-           [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; },
-           py::arg("value"))
-      .def("set_act_block",
-           [](CpuMoeExecutor& e, int v) { e.act_block = v; },
-           py::arg("value"))
+      .def("flag_served", [](const CpuMoeExecutor& e) { return e.flag_served.load(); })
+      .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("pin_core"))
+      .def("flag_signal", &CpuMoeExecutor::flag_signal, py::arg("stream"), py::arg("layer_id"),
+           py::arg("num_tokens"))
+      .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream"))
+      .def("flag_pending", &CpuMoeExecutor::flag_pending)
+      .def("enable_kernel_io", &CpuMoeExecutor::enable_kernel_io)
+      .def("flag_publish", &CpuMoeExecutor::flag_publish, py::arg("stream"), py::arg("x"), py::arg("ids"), py::arg("w"),
+           py::arg("layer_id"), py::arg("num_tokens"))
+      .def("flag_combine", &CpuMoeExecutor::flag_combine, py::arg("stream"), py::arg("out"), py::arg("num_tokens"),
+           py::arg("accumulate"))
+      .def("poison_flag", &CpuMoeExecutor::poison_flag)
       .def("isa_name", &CpuMoeExecutor::isa_name)
-      .def("fp8_block_isa_name", &CpuMoeExecutor::fp8_block_isa_name)
       .def("set_spin_wait",
            [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); },
            py::arg("enabled"))
       .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); });
-  m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr"));
-  m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"),
-        py::arg("ready_addr"), py::arg("slot"));
-  m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"),
-        py::arg("slot"));
+  m.def("memops_probe", &GpuSync::probe, py::arg("stream"));
+  m.def("kernel_io_probe", &freetoken::cpu_moe::KernelIo::probe, py::arg("stream"));
+  m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe);
   // ABI capability marker: the highest ActKind this build implements in the
   // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act
   // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI
   // accepts id 3 without error and silently computes the wrong activation
   // (act_apply falls through to gelu_tanh); the probe turns a stale extension
   // into a loud rebuild instruction instead of wrong model outputs.
-  m.def("max_generic_act_id", []() { return static_cast<int>(ACT_SWIGLU_CLAMP); });
-  // Weight-format ABI probe. A stale prebuilt .so accepts a newer WFmt id without
-  // complaint and then falls through to the wrong dequant branch -- silently wrong
-  // numbers, not a crash. The Python executor refuses to build against an extension
-  // that predates the format it was asked for.
-  m.def("max_weight_format_id", []() { return static_cast<int>(WF_FP8_BLOCK); });
+  m.def("max_generic_act_id", []() { return static_cast<int>(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); });
+  m.def("weight_formats", []() {
+    py::list names;
+    for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first);
+    return names;
+  });
 }
diff --git a/python/freetoken/kernel/csrc/cpu_moe/format.h b/python/freetoken/kernel/csrc/cpu_moe/format.h
new file mode 100644
index 0000000000..286ae110f7
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/format.h
@@ -0,0 +1,131 @@
+// The interface between the CPU MoE executor and one expert weight format.
+#pragma once
+
+#include <map>
+#include <memory>
+#include <stdexcept>
+#include <string>
+
+#include "common.h"
+
+namespace freetoken::cpu_moe {
+
+// Output-row tiling. Small enough to give every worker independent work even at
+// batch size 1; large enough to amortize the atomic work-grab.
+//
+// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read
+// every expert weight byte exactly once per token (each output row block is owned
+// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the
+// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe).
+// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription
+// thrashes the spin-barrier. Deferred (not worth it here / for this workload):
+//   - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would
+//     only pay off in a grouped/batched (dedup) path.
+//   - expert dedup for bs>1: read each distinct expert once and GEMM its tokens.
+//     Helps locality+bytes when bs is large; decode batches here are tiny (<=4).
+//   - NUMA: a single node is assumed. Multi-socket machines would split each
+//     expert's K dimension per node (banks are already per-row contiguous).
+constexpr int kRowBlock = 32;
+
+struct MoeParams {
+  int top_k = 0, max_tokens = 0;
+  int H = 0, I = 0;
+  int act = ACT_SILU;
+  float swiglu_alpha = 1.0f;
+  float swiglu_limit = INFINITY;  // +inf == no clamp
+};
+
+// What cpu_executor.py passes for one format, by name: banks map canonical roles (gate_up, gate_up_scale, ...) to per-layer tables, options hold small integers.
+// A format reads what it needs once, in its constructor.
+struct FormatArgs {
+  std::string format;
+  std::map<std::string, uint64_t> banks;
+  std::map<std::string, int64_t> options;
+
+  const uint64_t* bank(const std::string& role) const {
+    const auto it = banks.find(role);
+    if (it == banks.end() || it->second == 0)
+      throw std::invalid_argument("CPU MoE " + format + " experts need the " + role + " bank");
+    return reinterpret_cast<const uint64_t*>(it->second);
+  }
+  int64_t option(const std::string& name) const {
+    const auto it = options.find(name);
+    if (it == options.end()) throw std::invalid_argument("CPU MoE " + format + " experts need the " + name + " option");
+    return it->second;
+  }
+  int64_t option(const std::string& name, int64_t fallback) const {
+    const auto it = options.find(name);
+    return it == options.end() ? fallback : it->second;
+  }
+};
+
+// A bank table is the address of a CPU int64 array of length
+// num_layers (one base address per layer, built by cpu_executor.py's
+// _make_table), not a single flat bank. tbl_at resolves
+// tbl[layer_id] once per task/pass.
+inline const void* tbl_at(const uint64_t* tbl, int layer_id) {
+  return reinterpret_cast<const void*>(tbl[layer_id]);
+}
+
+// One route of a pass 2 item: expert e on intermediate row r = tok * top_k + k, whose bf16 values are g.
+struct DownRoute {
+  int e;
+  int64_t r;
+  const bf16_t* g;
+};
+
+// g[i] for i in [i0, i1) from the gate row i and the up row I + i; dot(row) is the row's GEMV.
+// ACT_SWIGLUOAI and ACT_SWIGLU_CLAMP clamp both inputs; the other kinds go through act_apply.
+template <class RowDot>
+inline void gated_rows(const MoeParams& p, int i0, int i1, float w_in, bf16_t* g, RowDot dot) {
+  const bool clamped = p.act == ACT_SWIGLUOAI || p.act == ACT_SWIGLU_CLAMP;
+  const float up_bias = p.act == ACT_SWIGLUOAI ? 1.0f : 0.0f;
+  const float lim = p.swiglu_limit, alpha = p.swiglu_alpha;
+  for (int i = i0; i < i1; ++i) {
+    // gate = row i, up = row I+i
+    float gate = dot(i) * w_in;
+    float up = dot(p.I + i) * w_in;
+    if (clamped) {
+      // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias)
+      // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp
+      // (GLM-5.3) does not. lim == +inf: no clamp.
+      if (gate > lim) gate = lim;
+      if (up > lim) up = lim;
+      else if (up < -lim) up = -lim;
+      const float glu = gate / (1.0f + std::exp(-gate * alpha));
+      g[i] = f32_to_bf16(glu * (up + up_bias));
+    } else {
+      g[i] = f32_to_bf16(act_apply(p.act, gate) * up);
+    }
+  }
+}
+
+// One expert weight format: its kernels, its activation layout and the scratch for it.
+// The executor owns the threads and the work split.
+// Scratch is sized once from max_tokens; gate_up, down and prep_inter run on the worker threads, prep_input before they wake.
+class MoeFormat {
+ public:
+  explicit MoeFormat(const MoeParams& p) : p(p) {}
+  virtual ~MoeFormat() = default;
+
+  virtual std::string isa_name() const = 0;
+  // Once per input row x[tok] before pass 1.
+  virtual void prep_input(int /*tok*/, const bf16_t* /*x*/) {}
+  // Whether pass 2 needs prep_inter on every intermediate row.
+  virtual bool preps_inter() const { return false; }
+  virtual void prep_inter(int64_t /*r*/, bf16_t* /*g*/) {}
+  // Pass 1: g[i] for i in [i0, i1) (at most kRowBlock rows), expert e of `layer` on token tok.
+  // w_in scales the gate and up dots: the router weight when it applies on the input, else 1.
+  virtual void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in,
+                       bf16_t* g) const = 0;
+  // Pass 2: out[j * kRowBlock + h - h0] for h in [h0, h1) (at most kRowBlock rows) and each of the n routes, before the router weight.
+  // Visiting the routes' rows h together keeps n weight streams in flight, which the down GEMV needs to reach DRAM bandwidth.
+  virtual void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const = 0;
+  // Round each router-weighted down output to bf16 before the sum over routes.
+  virtual bool rounds_routes() const { return false; }
+
+ protected:
+  const MoeParams& p;
+};
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h
new file mode 100644
index 0000000000..829a8f34c3
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h
@@ -0,0 +1,78 @@
+// GPU <-> CPU handshake for the CPU MoE coordinator, one implementation per GPU runtime (gpu_sync_cuda.cpp, gpu_sync_hip.cpp).
+//
+// A backend keeps two contracts, in eager mode and in graph replay:
+// 1. submit, GPU -> CPU: after gpu_submit(stream, task), cpu_poll() returns task once, and every earlier write of the stream is visible to the CPU.
+// 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done().
+// task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph.
+// One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals.
+// KernelIo (gpu_sync_io.h) may replace the copies around submit and wait: its kernel copies the inputs into the IO rows and does gpu_submit's writes itself.
+// It skips a step without a CPU route: done = 1 and no ready, so cpu_poll() never sees the step and gpu_wait() passes at once.
+#pragma once
+
+#include <cstdint>
+#include <memory>
+
+#if defined(_MSC_VER)
+#include <intrin.h>
+#endif
+
+namespace freetoken::cpu_moe {
+
+// Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that
+// preceded the bump are visible" and "y stores are visible before done". Plain
+// volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too
+// (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier
+// + TSO.)
+inline int64_t flag_load_acquire(const volatile int64_t* p) {
+#if defined(_MSC_VER)
+  const int64_t v = *p;
+  _ReadWriteBarrier();
+  return v;
+#else
+  return __atomic_load_n(const_cast<const int64_t*>(p), __ATOMIC_ACQUIRE);
+#endif
+}
+
+inline void flag_store_release(volatile int64_t* p, int64_t v) {
+#if defined(_MSC_VER)
+  _ReadWriteBarrier();
+  *p = v;
+#else
+  __atomic_store_n(const_cast<int64_t*>(p), v, __ATOMIC_RELEASE);
+#endif
+}
+
+class GpuSync {
+ public:
+  // Allocates the handshake memory; throws if this runtime cannot.
+  GpuSync();
+  ~GpuSync();
+  GpuSync(const GpuSync&) = delete;
+  GpuSync& operator=(const GpuSync&) = delete;
+
+  // Whether both contracts hold on this device and stream, including graph replay where the runtime needs a real check.
+  static bool probe(uintptr_t stream);
+  // Whether the cudaLaunchHostFunc fallback is replayed correctly from a captured graph.
+  static bool host_func_capture_safe();
+
+  // 1. submit. gpu_submit() is enqueued on `stream`, which may be capturing; cpu_poll() returns 0 when nothing is submitted.
+  void gpu_submit(uintptr_t stream, uint64_t task);
+  uint64_t cpu_poll();
+
+  // 2. wait. cpu_done() also unblocks the stream on teardown and from the watchdog.
+  void gpu_wait(uintptr_t stream);
+  void cpu_done();
+
+  // Submitted but not polled yet (watchdog only).
+  bool pending() const;
+
+  // Device addresses of the two words, for KernelIo.
+  uintptr_t ready_word() const;
+  uintptr_t done_word() const;
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp
new file mode 100644
index 0000000000..a2a089bf67
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp
@@ -0,0 +1,166 @@
+#include <atomic>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <stdexcept>
+#include <string>
+
+#if defined(_WIN32)
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+#include <cuda_runtime_api.h>
+
+#include "gpu_sync.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// =====================================================================================
+// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or
+// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE
+// (done=0 then ready=task word), sync = WAIT_VALUE(done >= 1). The wait is
+// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays
+// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%,
+// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max
+// frequency (GEMV workers -1.5x: the reported edge regression). Availability is
+// probed functionally at startup (memops_probe); anything unsupported (some driver models,
+// vGPU, old drivers) falls back to the cudaLaunchHostFunc path.
+#if defined(_WIN32)
+static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); }
+static void* cumemop_dlsym(void* h, const char* n) {
+  return (void*)::GetProcAddress((HMODULE)h, n);
+}
+#else
+static void* cumemop_dlopen() {
+  void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
+  if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL);
+  return h;
+}
+static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); }
+#endif
+
+using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value,
+                             unsigned int flags);
+static cuMemOp64_fn g_cu_write64 = nullptr;
+static cuMemOp64_fn g_cu_wait64 = nullptr;
+static constexpr unsigned int kCuWaitValueGeq = 0x0;   // CU_STREAM_WAIT_VALUE_GEQ
+static constexpr unsigned int kCuWriteDefault = 0x0;   // CU_STREAM_WRITE_VALUE_DEFAULT
+
+static bool cumemop_resolve() {
+  static bool resolved = [] {
+    void* h = cumemop_dlopen();
+    if (h == nullptr) return false;
+    // 11.7+ made the v2 entry points the default; older drivers export only the v1
+    // names with the same signature.
+    g_cu_write64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWriteValue64_v2"));
+    if (g_cu_write64 == nullptr)
+      g_cu_write64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWriteValue64"));
+    g_cu_wait64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWaitValue64_v2"));
+    if (g_cu_wait64 == nullptr)
+      g_cu_wait64 = reinterpret_cast<cuMemOp64_fn>(cumemop_dlsym(h, "cuStreamWaitValue64"));
+    return g_cu_write64 != nullptr && g_cu_wait64 != nullptr;
+  }();
+  return resolved;
+}
+
+// GPU side of the flag handshake (see the block comment above): enqueued on the
+// caller's (possibly capturing) stream; the WAIT immediate is the constant 1,
+// replay-safe under CUDA graphs.
+// The startup probe validates EAGER memops; a driver could still reject them at graph
+// capture time. Those enqueue errors would otherwise be swallowed here and surface only
+// as a later EndCapture failure -- log the first CUresult so triage is one step.
+static void cumemop_check(int rc, const char* what) {
+  static std::atomic<bool> warned{false};
+  if (rc != 0 && !warned.exchange(true)) {
+    std::fprintf(stderr,
+                 "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; "
+                 "subsequent errors are not repeated). If this happened during CUDA "
+                 "graph capture, the driver lacks capture support for stream memops -- "
+                 "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n",
+                 what, rc);
+  }
+}
+
+}  // namespace
+
+// ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses.
+struct GpuSync::Impl {
+  void* block = nullptr;
+  void* device_block = nullptr;
+  volatile int64_t* ready = nullptr;
+  volatile int64_t* done = nullptr;
+};
+
+GpuSync::GpuSync() : impl_(std::make_unique<Impl>()) {
+  if (!cumemop_resolve()) throw std::runtime_error("CPU MoE flag sync: CUDA stream memops are unavailable");
+  const cudaError_t rc = cudaHostAlloc(&impl_->block, 128, cudaHostAllocPortable | cudaHostAllocMapped);
+  if (rc != cudaSuccess)
+    throw std::runtime_error(std::string("CPU MoE flag allocation failed: ") + cudaGetErrorString(rc));
+  std::memset(impl_->block, 0, 128);
+  impl_->ready = static_cast<volatile int64_t*>(impl_->block);
+  impl_->done = impl_->ready + 8;
+  const cudaError_t dev_rc = cudaHostGetDevicePointer(&impl_->device_block, impl_->block, 0);
+  if (dev_rc != cudaSuccess) {
+    cudaFreeHost(impl_->block);
+    throw std::runtime_error(std::string("CPU MoE flag device address: ") + cudaGetErrorString(dev_rc));
+  }
+}
+
+GpuSync::~GpuSync() { cudaFreeHost(impl_->block); }
+
+// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync.
+// Returns true only if the whole memop path works on THIS stream/device/driver.
+bool GpuSync::probe(uintptr_t stream) {
+  if (!cumemop_resolve()) return false;
+  void* scratch = nullptr;
+  if (cudaHostAlloc(&scratch, sizeof(int64_t), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess)
+    return false;
+  *static_cast<int64_t*>(scratch) = 0;
+  auto* s = reinterpret_cast<void*>(stream);
+  const auto addr = reinterpret_cast<unsigned long long>(scratch);
+  const bool ok = g_cu_write64(s, addr, 7ULL, kCuWriteDefault) == 0 && g_cu_wait64(s, addr, 7ULL, kCuWaitValueGeq) == 0 &&
+                  cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)) == cudaSuccess;
+  // a failed enqueue can leave the WRITE pending on the scratch word
+  cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream));
+  cudaFreeHost(scratch);
+  return ok;
+}
+
+bool GpuSync::host_func_capture_safe() { return true; }
+
+void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) {
+  auto* s = reinterpret_cast<void*>(stream);
+  // Order matters and is preserved by the front end: reset done BEFORE raising ready,
+  // so the coordinator's completion write for THIS step can never be wiped.
+  cumemop_check(g_cu_write64(s, reinterpret_cast<unsigned long long>(impl_->done), 0ULL, kCuWriteDefault),
+                "cuStreamWriteValue64(done)");
+  cumemop_check(g_cu_write64(s, reinterpret_cast<unsigned long long>(impl_->ready), task, kCuWriteDefault),
+                "cuStreamWriteValue64(ready)");
+}
+
+void GpuSync::gpu_wait(uintptr_t stream) {
+  cumemop_check(g_cu_wait64(reinterpret_cast<void*>(stream), reinterpret_cast<unsigned long long>(impl_->done), 1ULL,
+                            kCuWaitValueGeq),
+                "cuStreamWaitValue64(done)");
+}
+
+uint64_t GpuSync::cpu_poll() {
+  // Acquire: everything the GPU made visible before writing ready -- the D2H input copies -- is visible to the worker pool after this read.
+  const int64_t task = flag_load_acquire(impl_->ready);
+  if (task == 0) return 0;
+  flag_store_release(impl_->ready, 0);  // consume this step's doorbell
+  return static_cast<uint64_t>(task);
+}
+
+void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); }
+
+bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; }
+
+uintptr_t GpuSync::ready_word() const { return reinterpret_cast<uintptr_t>(impl_->device_block); }
+
+uintptr_t GpuSync::done_word() const { return reinterpret_cast<uintptr_t>(impl_->device_block) + 8 * sizeof(int64_t); }
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
new file mode 100644
index 0000000000..31ff9be539
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
@@ -0,0 +1,275 @@
+// HIP stream memory operations for the CPU MoE flag handshake.
+//
+// HIP WAIT addresses must come from hipMallocSignalMemory, one 64-bit word per allocation.
+//
+// ROCm 10 records hipStreamWrite/WaitValue64 into the graph under stream capture, one node per call.
+// ROCm 7.14 runs them while the stream is being captured but records nothing, so there we add explicit batch-memop nodes and splice them into the stream's dependency set.
+// Which of the two applies is measured once at runtime (records_native_memops), not read from the version number.
+// The probe then does a real capture + instantiate + replay through the chosen path.
+#include <array>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <map>
+#include <stdexcept>
+#include <string>
+#include <thread>
+
+#include <freetoken/hip_compat.h>
+
+#include "gpu_sync.h"
+
+namespace freetoken::cpu_moe {
+
+namespace {
+
+void hip_check(hipError_t rc, const char* what) {
+  if (rc != hipSuccess)
+    throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc));
+}
+
+// Capability failures must not poison eager fallback or hide a caller's earlier error.
+struct ProbeErrorScope {
+  ProbeErrorScope() { hip_check(hipPeekAtLastError(), "pending error before probe"); }
+  ~ProbeErrorScope() { (void)hipGetLastError(); }
+};
+
+hipError_t alloc_signal(uint64_t** out) {
+  void* storage = nullptr;
+  const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory);
+  if (rc == hipSuccess) {
+    *out = static_cast<uint64_t*>(storage);
+    __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE);
+  }
+  return rc;
+}
+
+void free_signal(uint64_t* signal) {
+  if (signal != nullptr) (void)hipFree(signal);
+}
+
+volatile int64_t* as_flag(uint64_t* signal) { return reinterpret_cast<volatile int64_t*>(signal); }
+
+bool capturing(hipStream_t stream) {
+  hipStreamCaptureStatus status = hipStreamCaptureStatusNone;
+  hip_check(hipStreamIsCapturing(stream, &status), "capture status query");
+  return status == hipStreamCaptureStatusActive;
+}
+
+// Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node.
+bool records_native_memops() {
+  static const bool native = [] {
+    ProbeErrorScope errors;
+    uint64_t* word = nullptr;
+    hipStream_t stream = nullptr;
+    hipGraph_t graph = nullptr;
+    hipError_t rc = alloc_signal(&word);
+    if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking);
+    if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed);
+    const bool capture_started = rc == hipSuccess;
+    if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, word, 1, hipStreamWriteValueDefault);
+    if (capture_started) {
+      const hipError_t end_rc = hipStreamEndCapture(stream, &graph);
+      if (rc == hipSuccess) rc = end_rc;
+    }
+    size_t nodes = 0;
+    if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes);
+    if (graph != nullptr) (void)hipGraphDestroy(graph);
+    if (stream != nullptr) {
+      (void)hipStreamSynchronize(stream);
+      (void)hipStreamDestroy(stream);
+    }
+    free_signal(word);
+    return rc == hipSuccess && nodes == 1;
+  }();
+  return native;
+}
+
+// Appends one batch-memop node to the capture of `stream`; the node keeps a pointer to `ops`.
+hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) {
+  hipStreamCaptureStatus status = hipStreamCaptureStatusNone;
+  hipGraph_t graph = nullptr;
+  const hipGraphNode_t* dependencies = nullptr;
+  size_t dependency_count = 0;
+  hipError_t rc = hipStreamGetCaptureInfo_v2(stream, &status, nullptr, &graph, &dependencies, &dependency_count);
+  if (rc != hipSuccess) return rc;
+  if (status != hipStreamCaptureStatusActive || graph == nullptr) return hipErrorIllegalState;
+
+  hipCtx_t context = nullptr;
+#if defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+#endif
+  rc = hipCtxGetCurrent(&context);
+#if defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif
+  if (rc != hipSuccess) return rc;
+
+  hipBatchMemOpNodeParams params{};
+  params.ctx = context;
+  params.count = count;
+  params.paramArray = ops;
+  hipGraphNode_t node = nullptr;
+  rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, &params);
+  if (rc != hipSuccess) return rc;
+  rc = hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies);
+  // EndCapture rejects a node that was added but never joined to the stream.
+  if (rc != hipSuccess && node != nullptr) (void)hipGraphDestroyNode(node);
+  return rc;
+}
+
+hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) {
+  hipStreamBatchMemOpParams op{};
+  op.writeValue.operation = hipStreamMemOpWriteValue64;
+  op.writeValue.address = reinterpret_cast<hipDeviceptr_t>(address);
+  op.writeValue.value64 = value;
+  return op;
+}
+
+hipStreamBatchMemOpParams wait_op(uint64_t* address) {
+  hipStreamBatchMemOpParams op{};
+  op.waitValue.operation = hipStreamMemOpWaitValue64;
+  op.waitValue.address = reinterpret_cast<hipDeviceptr_t>(address);
+  op.waitValue.value64 = 1;
+  op.waitValue.flags = hipStreamWaitValueGte;
+  return op;
+}
+
+bool needs_explicit_nodes(hipStream_t stream) { return capturing(stream) && !records_native_memops(); }
+
+// node_op is null for a native call, else the parameter slot of an explicit node, which the node keeps a pointer to.
+// The ops of one batch node run in parallel, so each write gets its own node and consecutive writes stay in order.
+hipError_t write_value(hipStream_t stream, uint64_t* address, uint64_t value, hipStreamBatchMemOpParams* node_op) {
+  if (node_op == nullptr) return hipStreamWriteValue64(stream, address, value, hipStreamWriteValueDefault);
+  *node_op = write_op(address, value);
+  return add_capture_memop_node(stream, node_op, 1);
+}
+
+hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpParams* node_op) {
+  if (node_op == nullptr) return hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX);
+  *node_op = wait_op(done);
+  return add_capture_memop_node(stream, node_op, 1);
+}
+
+}  // namespace
+
+struct GpuSync::Impl {
+  uint64_t* ready = nullptr;  // signal memory
+  uint64_t* done = nullptr;
+  // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable).
+  std::map<uint64_t, std::array<hipStreamBatchMemOpParams, 2>> submit_ops;  // done = 0, ready = task
+  hipStreamBatchMemOpParams wait_op{};
+};
+
+GpuSync::GpuSync() : impl_(std::make_unique<Impl>()) {
+  hipError_t rc = alloc_signal(&impl_->ready);
+  if (rc == hipSuccess) rc = alloc_signal(&impl_->done);
+  if (rc != hipSuccess) {
+    free_signal(impl_->ready);
+    free_signal(impl_->done);
+    hip_check(rc, "signal allocation");
+  }
+  // Measure now: the measurement captures and allocates, which a graph capture in progress would reject.
+  (void)records_native_memops();
+}
+
+GpuSync::~GpuSync() {
+  free_signal(impl_->ready);
+  free_signal(impl_->done);
+}
+
+bool GpuSync::probe(uintptr_t /*stream*/) {
+  ProbeErrorScope errors;
+  int device = 0;
+  int can_wait = 0;
+  if (hipGetDevice(&device) != hipSuccess ||
+      hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0)
+    return false;
+  const bool explicit_nodes = !records_native_memops();
+
+  uint64_t* ready = nullptr;
+  uint64_t* done = nullptr;
+  hipStream_t stream = nullptr;
+  hipGraph_t graph = nullptr;
+  hipGraphExec_t graph_exec = nullptr;
+  std::array<hipStreamBatchMemOpParams, 2> node_ops{};
+  hipError_t rc = alloc_signal(&ready);
+  if (rc == hipSuccess) rc = alloc_signal(&done);
+  if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking);
+  // A runtime that runs the WAIT during capture instead of recording it must not block here.
+  if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE);
+  if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed);
+  const bool capture_started = rc == hipSuccess;
+  if (rc == hipSuccess) rc = write_value(stream, ready, 1, explicit_nodes ? &node_ops[0] : nullptr);
+  if (rc == hipSuccess) rc = wait_done(stream, done, explicit_nodes ? &node_ops[1] : nullptr);
+  // End the capture even on failure so its stream and graph are not left active during cleanup.
+  if (capture_started) {
+    const hipError_t end_rc = hipStreamEndCapture(stream, &graph);
+    if (rc == hipSuccess) rc = end_rc;
+  }
+  if (rc == hipSuccess) rc = hipStreamSynchronize(stream);
+  size_t nodes = 0;
+  if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes);
+  if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported;
+  if (rc == hipSuccess) rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0);
+
+  bool served = false;
+  if (rc == hipSuccess) {
+    __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE);
+    __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE);
+    std::thread responder([&] {
+      const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
+      while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline)
+        std::this_thread::yield();
+      served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1;
+      // The probe graph never resets done, so this also releases a launch that arrives after the deadline.
+      __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE);
+    });
+    rc = hipGraphLaunch(graph_exec, stream);
+    if (rc == hipSuccess) rc = hipStreamSynchronize(stream);
+    responder.join();
+  }
+  const bool ok = rc == hipSuccess && served;
+
+  if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec);
+  if (graph != nullptr) (void)hipGraphDestroy(graph);
+  if (stream != nullptr) (void)hipStreamDestroy(stream);
+  free_signal(ready);
+  free_signal(done);
+  return ok;
+}
+
+void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) {
+  auto s = reinterpret_cast<hipStream_t>(stream);
+  hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr;
+  // Reset done before raising ready, so the coordinator's completion for this step can never be wiped.
+  hip_check(write_value(s, impl_->done, 0, ops), "write done");
+  hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready");
+}
+
+void GpuSync::gpu_wait(uintptr_t stream) {
+  auto s = reinterpret_cast<hipStream_t>(stream);
+  hip_check(wait_done(s, impl_->done, needs_explicit_nodes(s) ? &impl_->wait_op : nullptr), "wait done");
+}
+
+bool GpuSync::host_func_capture_safe() { return false; }
+
+uint64_t GpuSync::cpu_poll() {
+  const int64_t task = flag_load_acquire(as_flag(impl_->ready));
+  if (task == 0) return 0;
+  flag_store_release(as_flag(impl_->ready), 0);
+  return static_cast<uint64_t>(task);
+}
+
+void GpuSync::cpu_done() { flag_store_release(as_flag(impl_->done), 1); }
+
+bool GpuSync::pending() const {
+  return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0;
+}
+
+uintptr_t GpuSync::ready_word() const { return reinterpret_cast<uintptr_t>(impl_->ready); }
+
+uintptr_t GpuSync::done_word() const { return reinterpret_cast<uintptr_t>(impl_->done); }
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu
new file mode 100644
index 0000000000..b41961f0bf
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu
@@ -0,0 +1,199 @@
+#include <algorithm>
+#include <cstring>
+#include <stdexcept>
+
+#include <freetoken/hip_compat.h>
+
+#include "gpu_sync_io.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+__device__ inline float bf16_bits_to_float(uint16_t b) { return __uint_as_float(static_cast<uint32_t>(b) << 16); }
+
+// Round to nearest even, the same rounding as a torch bf16 add.
+__device__ inline uint16_t float_to_bf16_bits(float f) {
+  const uint32_t u = __float_as_uint(f);
+  if ((u & 0x7fffffffu) > 0x7f800000u) return 0x7fc0;
+  return static_cast<uint16_t>((u + 0x7fffu + ((u >> 16) & 1u)) >> 16);
+}
+
+// One block. Copies x/ids/w into the IO rows, then done = 0 and ready = task.
+// A step without a CPU route writes done = 1 and rings nothing: the coordinator stays asleep and the GPU wait passes at once.
+__global__ void publish_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ ids, const float* __restrict__ w,
+                               int64_t n_x, int64_t n_ids, uint16_t* io_x, int32_t* io_ids, float* io_w,
+                               volatile int64_t* ready, volatile int64_t* done, int64_t task, int32_t* has_cpu) {
+  __shared__ int any;
+  if (threadIdx.x == 0) any = 0;
+  __syncthreads();
+  for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x)
+    if (ids[i] >= 0) any = 1;
+  __syncthreads();
+  if (!any) {
+    if (threadIdx.x == 0) {
+      *has_cpu = 0;
+      *done = 1;
+      __threadfence_system();
+    }
+    return;
+  }
+  const bool vec = (reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(io_x)) % 16 == 0 && n_x % 8 == 0;
+  if (vec) {
+    for (int64_t i = threadIdx.x; i < n_x / 8; i += blockDim.x)
+      reinterpret_cast<uint4*>(io_x)[i] = reinterpret_cast<const uint4*>(x)[i];
+  } else {
+    for (int64_t i = threadIdx.x; i < n_x; i += blockDim.x) io_x[i] = x[i];
+  }
+  for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) {
+    io_ids[i] = ids[i];
+    io_w[i] = w[i];
+  }
+  __threadfence_system();
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    *has_cpu = 1;
+    *done = 0;
+    __threadfence_system();
+    *ready = task;
+    __threadfence_system();
+  }
+}
+
+// out = CPU rows, or out += CPU rows with accumulate; a step publish skipped contributes zeros.
+// io_y is volatile: the same host rows hold a new step on every call.
+__global__ void combine_kernel(uint16_t* out, const volatile uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) {
+  const bool on = *has_cpu != 0;
+  if (accumulate && !on) return;
+  for (int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x; i < n; i += static_cast<int64_t>(gridDim.x) * blockDim.x) {
+    float v = accumulate ? bf16_bits_to_float(out[i]) : 0.0f;
+    if (on) v += bf16_bits_to_float(io_y[i]);
+    out[i] = float_to_bf16_bits(v);
+  }
+}
+
+cudaError_t launch_publish(cudaStream_t stream, const uint16_t* x, const int32_t* ids, const float* w, int64_t n_x, int64_t n_ids,
+                           uint16_t* io_x, int32_t* io_ids, float* io_w, volatile int64_t* ready, volatile int64_t* done, int64_t task,
+                           int32_t* has_cpu) {
+  publish_kernel<<<1, 256, 0, stream>>>(x, ids, w, n_x, n_ids, io_x, io_ids, io_w, ready, done, task, has_cpu);
+  return cudaGetLastError();
+}
+
+cudaError_t launch_combine(cudaStream_t stream, uint16_t* out, const uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) {
+  if (n == 0) return cudaSuccess;
+  const auto blocks = static_cast<unsigned>(std::min<int64_t>((n + 255) / 256, 1024));
+  combine_kernel<<<blocks, 256, 0, stream>>>(out, io_y, has_cpu, n, accumulate);
+  return cudaGetLastError();
+}
+
+void check(cudaError_t rc, const char* what) {
+  if (rc != cudaSuccess) throw std::runtime_error(std::string("CPU MoE kernel IO: ") + what + ": " + cudaGetErrorString(rc));
+}
+
+template <typename T>
+T* device_address(const void* host) {
+  void* dev = nullptr;
+  check(cudaHostGetDevicePointer(&dev, const_cast<void*>(host), 0), "device address of the IO rows");
+  return static_cast<T*>(dev);
+}
+
+}  // namespace
+
+KernelIo::KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden,
+                   int top_k)
+    : ready_(reinterpret_cast<volatile int64_t*>(ready)),
+      done_(reinterpret_cast<volatile int64_t*>(done)),
+      x_(device_address<uint16_t>(x)),
+      ids_(device_address<int32_t>(ids)),
+      w_(device_address<float>(w)),
+      y_(device_address<const uint16_t>(y)),
+      hidden_(hidden),
+      top_k_(top_k) {
+  check(cudaMalloc(reinterpret_cast<void**>(&has_cpu_), sizeof(int32_t)), "has_cpu allocation");
+  const cudaError_t rc = cudaMemset(has_cpu_, 0, sizeof(int32_t));
+  if (rc != cudaSuccess) {
+    (void)cudaFree(has_cpu_);
+    check(rc, "has_cpu reset");
+  }
+}
+
+KernelIo::~KernelIo() { (void)cudaFree(has_cpu_); }
+
+void KernelIo::publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task) {
+  check(launch_publish(reinterpret_cast<cudaStream_t>(stream), static_cast<const uint16_t*>(x), ids, w,
+                       static_cast<int64_t>(num_tokens) * hidden_, static_cast<int64_t>(num_tokens) * top_k_, x_, ids_, w_, ready_, done_,
+                       static_cast<int64_t>(task), has_cpu_),
+        "publish launch");
+}
+
+void KernelIo::combine(uintptr_t stream, void* out, int num_tokens, bool accumulate) {
+  check(launch_combine(reinterpret_cast<cudaStream_t>(stream), static_cast<uint16_t*>(out), y_, has_cpu_,
+                       static_cast<int64_t>(num_tokens) * hidden_, accumulate),
+        "combine launch");
+}
+
+// A ringing step, its combine and a skipped step on scratch mapped memory, checking every word they write.
+std::string KernelIo::probe(uintptr_t stream_handle) {
+  struct alignas(16) Area {
+    int64_t ready, done;
+    int32_t has_cpu, pad[3];
+    uint16_t x[8], io_x[8], io_y[8], out[8];
+    int32_t ids[2], io_ids[2];
+    float w[2], io_w[2];
+  };
+  const auto stream = reinterpret_cast<cudaStream_t>(stream_handle);
+  void* host = nullptr;
+  void* dev = nullptr;
+  if (cudaHostAlloc(&host, sizeof(Area), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess)
+    return "pinned scratch allocation failed";
+  std::string reason;
+  auto ok = [&](cudaError_t rc, const char* what) {
+    if (rc != cudaSuccess && reason.empty()) reason = std::string(what) + ": " + cudaGetErrorString(rc);
+    return rc == cudaSuccess;
+  };
+  auto fail = [&](const char* why) {
+    if (reason.empty()) reason = why;
+  };
+  if (ok(cudaHostGetDevicePointer(&dev, host, 0), "device address of the scratch")) {
+    auto* h = static_cast<Area*>(host);
+    auto* d = static_cast<Area*>(dev);
+    std::memset(h, 0, sizeof(Area));
+    for (int i = 0; i < 8; ++i) {
+      h->x[i] = static_cast<uint16_t>(0x3f80 + i);
+      h->io_y[i] = 0x3f80;  // 1.0
+      h->out[i] = 0x4000;  // 2.0
+    }
+    h->ids[0] = 3;
+    h->ids[1] = -1;
+    h->w[0] = 0.25f;
+    h->w[1] = 0.75f;
+    auto publish = [&] {
+      return launch_publish(stream, d->x, d->ids, d->w, 8, 2, d->io_x, d->io_ids, d->io_w, &d->ready, &d->done, 5, &d->has_cpu);
+    };
+
+    if (ok(publish(), "publish launch") && ok(cudaStreamSynchronize(stream), "publish")) {
+      if (h->ready != 5 || h->done != 0 || h->has_cpu != 1) fail("publish wrote wrong handshake words");
+      if (std::memcmp(h->io_x, h->x, sizeof(h->x)) || std::memcmp(h->io_ids, h->ids, sizeof(h->ids)) ||
+          std::memcmp(h->io_w, h->w, sizeof(h->w)))
+        fail("publish wrote wrong IO rows");
+    }
+    if (reason.empty() && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, true), "combine launch") &&
+        ok(cudaStreamSynchronize(stream), "combine")) {
+      for (int i = 0; i < 8; ++i)
+        if (h->out[i] != 0x4040) fail("combine returned a wrong sum");  // 3.0
+    }
+    if (reason.empty()) {
+      h->ready = 0;
+      h->ids[0] = -1;
+      if (ok(publish(), "skip launch") && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, false), "skip combine launch") &&
+          ok(cudaStreamSynchronize(stream), "skipped step")) {
+        bool zeros = true;
+        for (int i = 0; i < 8; ++i) zeros = zeros && h->out[i] == 0;
+        if (h->ready != 0 || h->done != 1 || h->has_cpu != 0 || !zeros) fail("a skipped step wrote wrong words");
+      }
+    }
+  }
+  (void)cudaFreeHost(host);
+  return reason;
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h
new file mode 100644
index 0000000000..6eacb3936b
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h
@@ -0,0 +1,39 @@
+// Kernel IO for the GPU <-> CPU handshake (gpu_sync.h): kernels on the executor's mapped pinned IO rows in place of the copies around submit and wait.
+// One source for both runtimes: nvcc builds gpu_sync_io.cu on CUDA, hipcc on ROCm through freetoken/hip_compat.h.
+#pragma once
+
+#include <cstdint>
+#include <string>
+
+namespace freetoken::cpu_moe {
+
+class KernelIo {
+ public:
+  // ready, done: device addresses of GpuSync's words. x, ids, w, y: host addresses of the pinned IO rows. Throws if setup fails.
+  KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, int top_k);
+  ~KernelIo();
+  KernelIo(const KernelIo&) = delete;
+  KernelIo& operator=(const KernelIo&) = delete;
+
+  // Copies the step's device inputs into the IO rows and does gpu_submit's writes, or skips the step when no route goes to the CPU.
+  // x [n, hidden] bf16, ids [n, top_k] int32 (< 0 = not a CPU route), w [n, top_k] fp32.
+  void publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task);
+  // Follows gpu_wait(): writes the CPU rows to out [n, hidden] bf16, or adds them to it; a skipped step gives zeros.
+  void combine(uintptr_t stream, void* out, int num_tokens, bool accumulate);
+
+  // Whether both kernels run and write the right words on this device: "" or the reason.
+  static std::string probe(uintptr_t stream);
+
+ private:
+  volatile int64_t* ready_;
+  volatile int64_t* done_;
+  uint16_t* x_;
+  int32_t* ids_;
+  float* w_;
+  const uint16_t* y_;
+  int hidden_;
+  int top_k_;
+  int32_t* has_cpu_ = nullptr;  // device word: whether the last published step has CPU routes
+};
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp
new file mode 100644
index 0000000000..273c394609
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp
@@ -0,0 +1,159 @@
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// ------------------------------- dot products -------------------------------
+// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function
+// pointer chosen at runtime; the per-row call overhead is negligible vs n.
+
+using dot_fn = float (*)(const bf16_t*, const bf16_t*, int);
+
+float dot_scalar(const bf16_t* w, const bf16_t* x, int n) {
+  float acc = 0.0f;
+  for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
+  return acc;
+}
+
+#if CPU_MOE_X86
+CPU_MOE_TARGET("avx512f")
+float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) {
+  // 4 independent accumulators -> more in-flight loads (memory-level parallelism),
+  // which is what lifts a bandwidth-bound GEMV toward peak.
+  __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
+  __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
+  int i = 0;
+  for (; i + 64 <= n; i += 64) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
+    for (int j = 0; j < 64; j += 16) {
+      __m256i wi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + i + j));
+      __m256i xi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + i + j));
+      __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16));
+      __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16));
+      __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3;
+      acc = _mm512_fmadd_ps(wf, xf, acc);
+    }
+  }
+  for (; i + 16 <= n; i += 16) {
+    __m256i wi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + i));
+    __m256i xi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + i));
+    __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16));
+    __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16));
+    a0 = _mm512_fmadd_ps(wf, xf, a0);
+  }
+  float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)));
+  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
+  return s;
+}
+
+#ifdef CPU_MOE_HAS_AVX512BF16
+CPU_MOE_TARGET("avx512bf16,avx512f")
+static inline __m512bh load_bh(const bf16_t* p) {
+  __m512i raw = _mm512_loadu_si512(reinterpret_cast<const void*>(p));
+  __m512bh out;
+  std::memcpy(&out, &raw, sizeof(out));
+  return out;
+}
+
+CPU_MOE_TARGET("avx512bf16,avx512f")
+float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) {
+  // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge.
+  __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
+  __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
+  int i = 0;
+  for (; i + 128 <= n; i += 128) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
+    a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i));
+    a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32));
+    a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64));
+    a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96));
+  }
+  for (; i + 32 <= n; i += 32) {
+    a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i));
+  }
+  float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)));
+  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
+  return s;
+}
+#endif  // avx512bf16 available
+
+// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ----------
+// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The
+// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the
+// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width).
+CPU_MOE_TARGET("avx2,fma")
+float dot_avx2(const bf16_t* w, const bf16_t* x, int n) {
+  __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
+  __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps();
+  int i = 0;
+  for (; i + 32 <= n; i += 32) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + i) + PF_AHEAD, _MM_HINT_T0);
+    for (int j = 0; j < 32; j += 8) {
+      __m128i wi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(w + i + j));
+      __m128i xi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + i + j));
+      __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16));
+      __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16));
+      __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3;
+      acc = _mm256_fmadd_ps(wf, xf, acc);
+    }
+  }
+  for (; i + 8 <= n; i += 8) {
+    __m128i wi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(w + i));
+    __m128i xi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + i));
+    __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16));
+    __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16));
+    a0 = _mm256_fmadd_ps(wf, xf, a0);
+  }
+  float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)));
+  for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]);
+  return s;
+}
+#endif  // CPU_MOE_X86
+
+dot_fn select_dot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+#ifdef CPU_MOE_HAS_AVX512BF16
+  if (t >= ISA_AVX512BF16) return dot_avx512bf16;
+#endif
+  if (t >= ISA_AVX512) return dot_avx512f;
+  if (t >= ISA_AVX2) return dot_avx2;
+#endif
+  (void)t;
+  return dot_scalar;
+}
+
+class Bf16Format final : public MoeFormat {
+ public:
+  Bf16Format(const MoeParams& p, const FormatArgs& a)
+      : MoeFormat(p), dot_(select_dot()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {}
+
+  std::string isa_name() const override { return tier_name(pick_isa()); }
+
+  void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float w_in,
+               bf16_t* g) const override {
+    const bf16_t* w = static_cast<const bf16_t*>(tbl_at(gate_up_, layer)) + (size_t)e * (2 * p.I) * p.H;
+    gated_rows(p, i0, i1, w_in, g, [&](int row) { return dot_(w + (size_t)row * p.H, x, p.H); });
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const bf16_t* w = static_cast<const bf16_t*>(tbl_at(down_, layer));
+    for (int h = h0; h < h1; ++h)
+      for (int j = 0; j < n; ++j)
+        out[j * kRowBlock + h - h0] = dot_(w + ((size_t)routes[j].e * p.H + h) * p.I, routes[j].g, p.I);
+  }
+
+ private:
+  dot_fn dot_;
+  // bf16 rows: gate_up [E, 2I, H], down [E, H, I].
+  const uint64_t* gate_up_;
+  const uint64_t* down_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_bf16_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<Bf16Format>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp
new file mode 100644
index 0000000000..a7afe9804d
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp
@@ -0,0 +1,269 @@
+#include <vector>
+
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) ---------------------------
+// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global
+// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128
+// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe +
+// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the
+// down output.
+
+// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per
+// token/route and reused across all output rows. This drops the hot dot to one
+// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16,
+// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5
+// bound and approaches the bf16 memory-bandwidth ceiling.
+using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int,
+                           const float*, const float*);
+
+float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe,
+                       const float* xo, int K, const float* e2m1, const float* e8m0) {
+  float acc = 0.0f;
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const float sc = e8m0[scale[b]];
+    const uint8_t* pk = packed + (size_t)b * 16;  // 16 bytes = 32 codes
+    const float* xeb = xe + (size_t)b * 16;
+    const float* xob = xo + (size_t)b * 16;
+    float bsum = 0.0f;
+    for (int j = 0; j < 16; ++j) {
+      const uint8_t byte = pk[j];
+      bsum += e2m1[byte & 0xF] * xeb[j];   // low nibble  -> even-K activation
+      bsum += e2m1[byte >> 4] * xob[j];    // high nibble -> odd-K activation
+    }
+    acc += sc * bsum;
+  }
+  return acc;
+}
+
+#if CPU_MOE_X86
+// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times
+// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale.
+CPU_MOE_TARGET("avx512f")
+inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut,
+                        __m512i loma, float sc) {
+  __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(pk)));
+  __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
+  __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
+  __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob)));
+  return _mm512_mul_ps(prod, _mm512_set1_ps(sc));
+}
+
+CPU_MOE_TARGET("avx512f")
+float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe,
+                       const float* xo, int K, const float* e2m1, const float* e8m0) {
+  const __m512 lut = _mm512_loadu_ps(e2m1);
+  const __m512i loma = _mm512_set1_epi32(0xF);
+  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
+  const int nb = K / 32;  // 16 packed bytes + one e8m0 scale per block
+  int b = 0;
+  for (; b + 2 <= nb; b += 2) {  // two independent accumulators hide FMA latency
+    acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16,
+                                         xo + (size_t)b * 16, lut, loma, e8m0[scale[b]]));
+    acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16,
+                                         xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]]));
+  }
+  for (; b < nb; ++b)
+    acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16,
+                                         xo + (size_t)b * 16, lut, loma, e8m0[scale[b]]));
+  return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
+}
+
+// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each).
+CPU_MOE_TARGET("avx2,fma")
+inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) {
+  __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk)));
+  __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
+  __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
+  return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob)));
+}
+
+CPU_MOE_TARGET("avx2,fma")
+float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe,
+                     const float* xo, int K, const float* e2m1, const float* e8m0) {
+  const __m256 mag8 = _mm256_loadu_ps(e2m1);
+  __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps();
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const uint8_t* pk = packed + (size_t)b * 16;
+    const float* xeb = xe + (size_t)b * 16;
+    const float* xob = xo + (size_t)b * 16;
+    const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]);
+    acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0);
+    acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1);
+  }
+  return hsum256(_mm256_add_ps(acc0, acc1));
+}
+#endif
+
+dsdot_fn select_dsdot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  if (t >= ISA_AVX512) return dot_dsfp4_avx512;
+  if (t >= ISA_AVX2) return dot_dsfp4_avx2;
+#endif
+  (void)t;
+  return dot_dsfp4_scalar;
+}
+
+// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches
+// torch.float8_e4m3fn / triton .to(float8e4nv).
+inline float e4m3_round(float x) {
+  const float sign = x < 0.0f ? -1.0f : 1.0f;
+  const float a = std::fabs(x);
+  if (a == 0.0f) return 0.0f;
+  if (a >= 448.0f) return sign * 448.0f;
+  int e;
+  std::frexp(a, &e);  // a in [2^(e-1), 2^e)
+  float step = std::ldexp(1.0f, e - 4);
+  const float min_step = std::ldexp(1.0f, -9);  // e4m3 subnormal step (2^-9)
+  if (step < min_step) step = min_step;
+  float r = std::nearbyint(a / step) * step;
+  if (r > 448.0f) r = 448.0f;
+  return sign * r;
+}
+
+// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale).
+inline int ceil_log2_pos(float v) {
+  uint32_t bits;
+  std::memcpy(&bits, &v, sizeof(bits));
+  const int exp = (int)((bits >> 23) & 0xFF);
+  const int man = (int)(bits & 0x7FFFFF);
+  return exp - 127 + (man != 0 ? 1 : 0);
+}
+
+// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's
+// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448));
+// y = round_e4m3(clamp(x/s,+-448)) * s.
+void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) {
+  for (int b0 = 0; b0 < K; b0 += block) {
+    const int b1 = std::min(K, b0 + block);
+    float amax = 1e-4f;
+    for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i])));
+    const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f)));
+    const float inv_s = 1.0f / s;
+    for (int i = b0; i < b1; ++i) {
+      float q = bf16_to_f32(src[i]) * inv_s;
+      q = std::min(448.0f, std::max(-448.0f, q));
+      dst[i] = f32_to_bf16(e4m3_round(q) * s);
+    }
+  }
+}
+
+class DsFp4Format final : public MoeFormat {
+ public:
+  DsFp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_dsdot()) {
+    gate_up_ = a.bank("gate_up");
+    gu_scale_ = a.bank("gate_up_scale");
+    down_ = a.bank("down");
+    dn_scale_ = a.bank("down_scale");
+    act_block_ = static_cast<int>(a.option("act_block"));
+    input_prequant_ = a.option("input_prequant", 0) != 0;
+    const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k;
+    xq_.resize(tokens * p.H);
+    xe_.resize(tokens * (p.H / 2));
+    xo_.resize(tokens * (p.H / 2));
+    ge_.resize(routes * (p.I / 2));
+    go_.resize(routes * (p.I / 2));
+    for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i];
+    for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127);
+  }
+
+  std::string isa_name() const override { return tier_name(pick_isa()); }
+
+
+  // DSV4 FP8-round-trips the input before the gate_up GEMV, unless the GPU already did.
+  void prep_input(int tok, const bf16_t* x) override {
+    const bf16_t* src = x;
+    if (!input_prequant_) {
+      bf16_t* xq = xq_.data() + (size_t)tok * p.H;
+      fp8_roundtrip_bf16(x, xq, p.H, act_block_);
+      src = xq;
+    }
+    deinterleave_bf16_f32(src, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H);
+  }
+
+  bool preps_inter() const override { return true; }
+
+  void prep_inter(int64_t r, bf16_t* g) override {
+    fp8_roundtrip_bf16(g, g, p.I, act_block_);
+    deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I);
+  }
+
+  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float /*w_in*/,
+               bf16_t* g) const override {
+    const int H = p.H, I = p.I;
+    const uint8_t* gu_packed_l = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
+    const uint8_t* gu_scale_l = static_cast<const uint8_t*>(tbl_at(gu_scale_, layer));
+    const int N2 = 2 * I, Hh = H / 2, Hs = H / 32;
+    const float* xe = xe_.data() + (size_t)tok * (H / 2);
+    const float* xo = xo_.data() + (size_t)tok * (H / 2);
+    const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh;
+    const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs;
+    const float lim = p.swiglu_limit;
+    for (int i = i0; i < i1; ++i) {
+      // gate_up is stored bf16 by the reference GEMV before swiglu; round to match.
+      float gate = bf16_to_f32(f32_to_bf16(
+          dot_(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_)));
+      float up = bf16_to_f32(f32_to_bf16(dot_(
+          gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_)));
+      if (lim > 0.0f) {
+        if (gate > lim) gate = lim;
+        if (up > lim) up = lim;
+        else if (up < -lim) up = -lim;
+      }
+      const float glu = gate / (1.0f + std::exp(-gate));  // silu(gate)
+      g[i] = f32_to_bf16(glu * up);
+    }
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const int H = p.H, I = p.I, Ih = I / 2, Is = I / 32;
+    const uint8_t* dn_packed_l = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    const uint8_t* dn_scale_l = static_cast<const uint8_t*>(tbl_at(dn_scale_, layer));
+    for (int h = h0; h < h1; ++h) {
+      for (int j = 0; j < n; ++j) {
+        const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r;
+        const uint8_t* dp = dn_packed_l + e * (size_t)H * Ih + (size_t)h * Ih;
+        const uint8_t* ds = dn_scale_l + e * (size_t)H * Is + (size_t)h * Is;
+        out[j * kRowBlock + h - h0] = dot_(dp, ds, ge_.data() + r * Ih, go_.data() + r * Ih, I, e2m1_lut_, e8m0_lut_);
+      }
+    }
+  }
+
+  // The reference rounds each route's weighted down output to bf16 before the fp32 sum over routes.
+  bool rounds_routes() const override { return true; }
+
+ private:
+  dsdot_fn dot_;
+  // e2m1 rows [E,2I,H/2] with e8m0 scales per 32 along K [E,2I,H/32]; down likewise with K = I.
+  const uint64_t* gate_up_ = nullptr;
+  const uint64_t* gu_scale_ = nullptr;
+  const uint64_t* down_ = nullptr;
+  const uint64_t* dn_scale_ = nullptr;
+  // Activation FP8 round-trip block, the checkpoint's fp8 block.
+  int act_block_ = 128;
+  // The caller already FP8-round-tripped the input activations on the GPU
+  // (same reference grid), so prep_input must not repeat it on the host-callback
+  // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL
+  // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving
+  // it to a captured GPU elementwise kernel removes it while keeping the official
+  // W4A8 numerics bit-exact (see cpu_executor.py).
+  bool input_prequant_ = false;
+  float e2m1_lut_[16];
+  float e8m0_lut_[256];
+  AlignedVector<bf16_t> xq_;
+  AlignedVector<float> xe_, xo_, ge_, go_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_ds_fp4_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<DsFp4Format>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp
new file mode 100644
index 0000000000..d27fb1868e
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp
@@ -0,0 +1,203 @@
+#include <string>
+#include <vector>
+
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// ------------------------ block-FP8 (W8A16) dequant ------------------------
+// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py.
+// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x.
+// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does.
+using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*,
+                                   const bf16_t*, int, const float*);
+
+constexpr int FP8_BLK = 128;  // K elements covered by one weight scale
+
+inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) {
+  float blk = 0.0f;
+  for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]);
+  return blk;
+}
+
+float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                           const bf16_t* s, int K, const float* e4m3) {
+  float acc = 0.0f;
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    // 8 independent sums: a single running sum is bound by the add latency
+    float a[8] = {};
+    const uint8_t* wb = w + k0;
+    const float* pe = xe + k0 / 2;
+    const float* po = xo + k0 / 2;
+    for (int j = 0; j < FP8_BLK / 2; j += 4)
+      for (int l = 0; l < 4; ++l) {
+        a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l];
+        a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l];
+      }
+    acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) *
+           bf16_to_f32(s[b]);
+  }
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+
+#if CPU_MOE_X86
+// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256.
+// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C.
+CPU_MOE_TARGET("avx2,fma,f16c")
+float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                         const bf16_t* s, int K, const float* e4m3) {
+  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
+  __m256 total = _mm256_setzero_ps();
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
+    __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
+    __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps();
+    for (int j = 0; j < FP8_BLK; j += 32) {
+      const __m256i v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
+      const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep);
+      const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep);
+      const float* pe = xe + (k0 + j) / 2;
+      const float* po = xo + (k0 + j) / 2;
+      a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0);
+      a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1);
+      a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2);
+      a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3);
+    }
+    const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3));
+    total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
+  }
+  float acc = hsum256(total);
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+
+CPU_MOE_TARGET("avx512f,avx2,fma,f16c")
+float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x,
+                            const bf16_t* s, int K, const float* e4m3) {
+  const __m256i keep = _mm256_set1_epi16((short)0xBF80);
+  __m512 total = _mm512_setzero_ps();
+  int b = 0, k0 = 0;
+  for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) {
+    _mm_prefetch(reinterpret_cast<const char*>(w + k0) + PF_AHEAD, _MM_HINT_T0);
+    __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps();
+    __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps();
+    for (int j = 0; j < FP8_BLK; j += 64) {
+      const __m256i v0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j));
+      const __m256i v1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(w + k0 + j + 32));
+      const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep);
+      const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep);
+      const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep);
+      const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep);
+      const float* pe = xe + (k0 + j) / 2;
+      const float* po = xo + (k0 + j) / 2;
+      a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0);
+      a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1);
+      a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2);
+      a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3);
+    }
+    const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3));
+    total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total);
+  }
+  float acc = _mm512_reduce_add_ps(total);
+  if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]);
+  return acc;
+}
+#endif  // CPU_MOE_X86
+
+fp8_block_dot_fn select_fp8_block_dot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16
+  if (t >= ISA_AVX512) return dot_fp8_block_avx512f;
+  if (t >= ISA_AVX2) return dot_fp8_block_avx2;
+#endif
+  (void)t;
+  return dot_fp8_block_scalar;
+}
+
+class Fp8BlockFormat final : public MoeFormat {
+ public:
+  Fp8BlockFormat(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_fp8_block_dot()) {
+    gate_up_ = a.bank("gate_up");
+    gu_scale_ = a.bank("gate_up_scale");
+    down_ = a.bank("down");
+    dn_scale_ = a.bank("down_scale");
+    gu_scale_stride_ = a.option("gu_scale_stride");
+    dn_scale_stride_ = a.option("dn_scale_stride");
+    const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k;
+    xe_.resize(tokens * (p.H / 2));
+    xo_.resize(tokens * (p.H / 2));
+    ge_.resize(routes * (p.I / 2));
+    go_.resize(routes * (p.I / 2));
+    for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i);
+  }
+
+  std::string isa_name() const override { return tier_name(pick_isa()); }
+
+
+  void prep_input(int tok, const bf16_t* x) override {
+    deinterleave_bf16_f32(x, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H);
+  }
+
+  bool preps_inter() const override { return true; }
+
+  void prep_inter(int64_t r, bf16_t* g) override {
+    deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I);
+  }
+
+  // The halves cover whole 128-blocks; the bf16 row covers a ragged tail.
+  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in,
+               bf16_t* g) const override {
+    const int H = p.H, I = p.I;
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
+    const bf16_t* scale = static_cast<const bf16_t*>(tbl_at(gu_scale_, layer));
+    const float* xe = xe_.data() + (size_t)tok * (H / 2);
+    const float* xo = xo_.data() + (size_t)tok * (H / 2);
+    gated_rows(p, i0, i1, w_in, g, [&](int row) {
+      const uint8_t* w = packed + ((size_t)e * (2 * I) + row) * (size_t)H;
+      const bf16_t* sc = scale + ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) *
+                                     (size_t)gu_scale_stride_;
+      return dot_(w, xe, xo, x, sc, H, e4m3_lut_);
+    });
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const int H = p.H, I = p.I;
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    const bf16_t* scale = static_cast<const bf16_t*>(tbl_at(dn_scale_, layer));
+    for (int h = h0; h < h1; ++h) {
+      for (int j = 0; j < n; ++j) {
+        const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r;
+        const uint8_t* w = packed + (e * H + h) * (size_t)I;
+        const bf16_t* sc = scale + (e * ((H + FP8_BLK - 1) / FP8_BLK) + h / FP8_BLK) *
+                                       (size_t)dn_scale_stride_;
+        out[j * kRowBlock + h - h0] =
+            dot_(w, ge_.data() + r * (I / 2), go_.data() + r * (I / 2), routes[j].g, sc, I, e4m3_lut_);
+      }
+    }
+  }
+
+ private:
+  fp8_block_dot_fn dot_;
+  // e4m3 rows [E,2I,H] and [E,H,I] with one bf16 scale per 128x128 block.
+  const uint64_t* gate_up_ = nullptr;
+  const uint64_t* gu_scale_ = nullptr;
+  const uint64_t* down_ = nullptr;
+  const uint64_t* dn_scale_ = nullptr;
+  // Row strides of the scale tensors, including any padding of the host bank layout.
+  int64_t gu_scale_stride_ = 0, dn_scale_stride_ = 0;
+  float e4m3_lut_[256];
+  AlignedVector<float> xe_, xo_, ge_, go_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_fp8_block_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<Fp8BlockFormat>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp
new file mode 100644
index 0000000000..484d6dc001
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp
@@ -0,0 +1,247 @@
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// ------------------------- mxfp4 (gpt-oss) GEMV -----------------------------
+// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2]
+// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1])
+// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of
+// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache-
+// sequential) loop. Used by both gate_up (K=H) and down (K=I).
+using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int,
+                           int, const float*, const float*);
+
+void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
+                       int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
+  for (int c = 0; c < ncol; ++c) out[c] = 0.0f;
+  for (int kb = 0; kb < Kpairs; ++kb) {
+    const uint8_t* w = blk + (size_t)kb * N2;
+    const uint8_t* s = scl + (size_t)(kb >> 4) * N2;
+    const float xl = bf16_to_f32(x[2 * kb]);
+    const float xh = bf16_to_f32(x[2 * kb + 1]);
+    for (int c = 0; c < ncol; ++c) {
+      const uint8_t byte = w[c];
+      out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]];
+    }
+  }
+}
+
+#if CPU_MOE_X86
+CPU_MOE_TARGET("avx512f")
+void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
+                       int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
+  (void)e8m0;  // e8m0[c]=2^(c-127) computed via bit construction (no gather)
+  const __m512 lut = _mm512_loadu_ps(e2m1);
+  const __m512i loma = _mm512_set1_epi32(0xF);
+  // K-outer / N-inner: each kb cache line is read once and all live column chunks
+  // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once.
+  int c0 = 0;
+  for (; c0 + 16 <= ncol; c0 += 64) {
+    const int nchunk = std::min(4, (ncol - c0) / 16);
+    __m512 acc[4];
+    for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps();
+    for (int kblk = 0; kblk < Kpairs; kblk += 16) {  // 16 K-pairs = 32 K = one scale row
+      __m512 sc[4];
+      for (int ci = 0; ci < nchunk; ++ci) {
+        __m128i sraw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(
+            scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16));
+        sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23));
+      }
+      __m512 blk_acc[4];
+      for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps();
+      for (int kk = 0; kk < 16; ++kk) {
+        const int kb = kblk + kk;
+        const uint8_t* wbase = blk + (size_t)kb * N2 + c0;
+        // The transposed layout strides K by N2 bytes; prefetch ahead so the strided
+        // reads are not exposed to DRAM latency (the HW streamer misses big strides).
+        constexpr int PFD = 8;
+        if (kb + PFD < Kpairs)
+          _mm_prefetch(reinterpret_cast<const char*>(blk + (size_t)(kb + PFD) * N2 + c0),
+                       _MM_HINT_T0);
+        const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb]));
+        const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1]));
+        for (int ci = 0; ci < nchunk; ++ci) {
+          __m512i wi = _mm512_cvtepu8_epi32(
+              _mm_loadu_si128(reinterpret_cast<const __m128i*>(wbase + ci * 16)));
+          __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
+          __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
+          blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]);
+          blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]);
+        }
+      }
+      for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]);
+    }
+    for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]);
+  }
+  for (int c = c0; c < ncol; ++c) {  // tail columns (< 16)
+    float o = 0.0f;
+    for (int kb = 0; kb < Kpairs; ++kb) {
+      const uint8_t byte = blk[(size_t)kb * N2 + c];
+      uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23;
+      float sc;
+      std::memcpy(&sc, &bits, 4);
+      o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) +
+            e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc;
+    }
+    out[c] = o;
+  }
+}
+
+CPU_MOE_TARGET("avx2,fma")
+void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x,
+                     int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) {
+  (void)e8m0;  // e8m0[s]=2^(s-127) built via s<<23 (no gather)
+  const __m256 mag8 = _mm256_loadu_ps(e2m1);
+  int c0 = 0;
+  for (; c0 + 8 <= ncol; c0 += 32) {  // up to 4 chunks of 8 = 32 cols
+    const int nchunk = std::min(4, (ncol - c0) / 8);
+    __m256 acc[4];
+    for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps();
+    for (int kblk = 0; kblk < Kpairs; kblk += 16) {  // 16 K-pairs = one scale row
+      __m256 sc[4];
+      for (int ci = 0; ci < nchunk; ++ci) {
+        __m128i sraw = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(
+            scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8));
+        sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23));
+      }
+      __m256 blk_acc[4];
+      for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps();
+      for (int kk = 0; kk < 16; ++kk) {
+        const int kb = kblk + kk;
+        const uint8_t* wbase = blk + (size_t)kb * N2 + c0;
+        constexpr int PFD = 8;
+        if (kb + PFD < Kpairs)
+          _mm_prefetch(reinterpret_cast<const char*>(blk + (size_t)(kb + PFD) * N2 + c0),
+                       _MM_HINT_T0);
+        const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb]));
+        const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1]));
+        for (int ci = 0; ci < nchunk; ++ci) {
+          __m256i wi = _mm256_cvtepu8_epi32(
+              _mm_loadl_epi64(reinterpret_cast<const __m128i*>(wbase + ci * 8)));
+          __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
+          __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
+          blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]);
+          blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]);
+        }
+      }
+      for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]);
+    }
+    for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]);
+  }
+  for (int c = c0; c < ncol; ++c) {  // tail columns (< 8); none when ncol%8==0
+    float o = 0.0f;
+    for (int kb = 0; kb < Kpairs; ++kb) {
+      const uint8_t byte = blk[(size_t)kb * N2 + c];
+      uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23;
+      float sc;
+      std::memcpy(&sc, &bits, 4);
+      o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) +
+            e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc;
+    }
+    out[c] = o;
+  }
+}
+#endif
+
+mxgemv_fn select_mxgemv() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  if (t >= ISA_AVX512) return mxfp4_gemv_avx512;
+  if (t >= ISA_AVX2) return mxfp4_gemv_avx2;
+#endif
+  (void)t;
+  return mxfp4_gemv_scalar;
+}
+
+// ----------------------------- mxfp4 (gpt-oss) -----------------------------
+// Transposed split-K layout (N innermost), so the GEMV streams K and accumulates
+// a contiguous block of N output columns -> cache-line-efficient, no repack, no
+// extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 adds the
+// down bias, and the executor applies the router weight and sums the routes.
+//
+// Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble
+// first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py.
+class Mxfp4Format final : public MoeFormat {
+ public:
+  Mxfp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), gemv_(select_mxgemv()) {
+    gate_up_ = a.bank("gate_up");
+    gu_scale_ = a.bank("gate_up_scale");
+    gu_bias_ = a.bank("gate_up_bias");
+    down_ = a.bank("down");
+    dn_scale_ = a.bank("down_scale");
+    dn_bias_ = a.bank("down_bias");
+    for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i];
+    // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254].
+    for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127);
+  }
+
+  std::string isa_name() const override { return tier_name(pick_isa()); }
+
+  void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float /*w_in*/,
+               bf16_t* g) const override {
+    const int H = p.H, I = p.I;
+    const uint8_t* gu_packed_l = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
+    const uint8_t* gu_scale_l = static_cast<const uint8_t*>(tbl_at(gu_scale_, layer));
+    const bf16_t* gu_bias_l = static_cast<const bf16_t*>(tbl_at(gu_bias_, layer));
+    const int N2 = 2 * I;            // gate_up output width (gate/up interleaved)
+    const int Hh = H / 2;            // packed-K rows
+    const int nunit = i1 - i0;       // intermediate units owned by this tile
+    const int col0 = 2 * i0;         // first gate_up column
+    const int ncol = 2 * nunit;      // gate_up columns owned by this tile
+    const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2;
+    const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2;
+    float gu[2 * kRowBlock];
+    gemv_(gu, blk_e + col0, scl_e + col0, x, Hh, N2, ncol, e2m1_lut_, e8m0_lut_);
+    const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0;
+    const float lim = p.swiglu_limit, alpha = p.swiglu_alpha;
+    for (int j = 0; j < nunit; ++j) {
+      float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]);
+      float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]);
+      if (gate > lim) gate = lim;
+      if (up > lim) up = lim;
+      else if (up < -lim) up = -lim;
+      const float glu = gate / (1.0f + std::exp(-gate * alpha));  // gate * sigmoid(alpha*gate)
+      g[i0 + j] = f32_to_bf16(glu * (up + 1.0f));
+    }
+  }
+
+  // The transposed layout already streams a whole column block per call, so the routes go one after another.
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const int H = p.H, I = p.I;
+    const int Ih = I / 2;            // packed-K rows
+    const int nh = h1 - h0;
+    const uint8_t* dn_packed_l = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    const uint8_t* dn_scale_l = static_cast<const uint8_t*>(tbl_at(dn_scale_, layer));
+    const bf16_t* dn_bias_l = static_cast<const bf16_t*>(tbl_at(dn_bias_, layer));
+    for (int j = 0; j < n; ++j) {
+      const size_t e = (size_t)routes[j].e;
+      const uint8_t* blk_e = dn_packed_l + e * Ih * H;
+      const uint8_t* scl_e = dn_scale_l + e * (size_t)(I / 32) * H;
+      float* part = out + j * kRowBlock;
+      gemv_(part, blk_e + h0, scl_e + h0, routes[j].g, Ih, H, nh, e2m1_lut_, e8m0_lut_);
+      const bf16_t* bias_e = dn_bias_l + e * H + h0;
+      for (int c = 0; c < nh; ++c) part[c] += bf16_to_f32(bias_e[c]);
+    }
+  }
+
+ private:
+  mxgemv_fn gemv_;
+  // gpt-oss split-K layout, N innermost: blocks [E,H/2,2I], e8m0 scales [E,H/32,2I], bf16 bias [E,2I]; down likewise with K = I, N = H.
+  const uint64_t* gate_up_ = nullptr;
+  const uint64_t* gu_scale_ = nullptr;
+  const uint64_t* gu_bias_ = nullptr;
+  const uint64_t* down_ = nullptr;
+  const uint64_t* dn_scale_ = nullptr;
+  const uint64_t* dn_bias_ = nullptr;
+  float e2m1_lut_[16];
+  float e8m0_lut_[256];
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_mxfp4_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<Mxfp4Format>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
new file mode 100644
index 0000000000..429955faf3
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
@@ -0,0 +1,466 @@
+#include <vector>
+
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// --------------------------- NVFP4 (W4A16) dequant ---------------------------
+// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in
+// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels
+// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global.
+// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop.
+
+
+// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI
+// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations;
+// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4.
+alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12};
+
+// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the
+// ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights).
+using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*,
+                           int, const float*, const float*);
+
+float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global,
+                       const float* xe, const float* xo, int K, const float* e2m1,
+                       const float* e4m3) {
+  float acc = 0.0f;
+  const int nb = K / 16;
+  for (int b = 0; b < nb; ++b) {
+    const float bs = e4m3[scale[b]];
+    const uint8_t* pk = packed + (size_t)b * 8;
+    const float* xeb = xe + (size_t)b * 8;  // 16 K -> 8 even + 8 odd
+    const float* xob = xo + (size_t)b * 8;
+    float bsum = 0.0f;
+    for (int j = 0; j < 8; ++j) {
+      const uint8_t byte = pk[j];
+      bsum += e2m1[byte & 0xF] * xeb[j];
+      bsum += e2m1[byte >> 4] * xob[j];
+    }
+    acc += bs * bsum;
+  }
+  return acc * global;
+}
+
+// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ----
+// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block
+// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale.
+using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int,
+                             const float*, const float*);
+
+[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale,
+                          float global, const int8_t* asi8, int K, const float* e4m3,
+                          const float* asb) {
+  float acc = 0.0f;
+  const int nb = K / 16;
+  for (int b = 0; b < nb; ++b) {
+    const uint8_t* pk = packed + (size_t)b * 8;
+    const int8_t* ae = asi8 + (size_t)b * 16;       // even(8)
+    const int8_t* ao = ae + 8;                      // odd(8)
+    int isum = 0;
+    for (int j = 0; j < 8; ++j) {
+      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j];
+      isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
+    }
+    acc += (e4m3[scale[b]] * asb[b]) * (float)isum;
+  }
+  return acc * (0.5f * global);
+}
+
+#if CPU_MOE_X86
+// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K,
+// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so
+// it is applied per lane (low 8 lanes = block b, high 8 = block b+1).
+CPU_MOE_TARGET("avx512f")
+inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut,
+                         __m512i loma, float s0, float s1) {
+  __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(pk)));
+  __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut);
+  __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut);
+  __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob)));
+  // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1  (pure AVX512F mask move)
+  __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1));
+  return _mm512_mul_ps(prod, scv);
+}
+
+CPU_MOE_TARGET("avx512f")
+float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global,
+                       const float* xe, const float* xo, int K, const float* e2m1,
+                       const float* e4m3) {
+  const __m512 lut = _mm512_loadu_ps(e2m1);  // 16 e2m1 values
+  const __m512i loma = _mm512_set1_epi32(0xF);
+  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
+  const int nb = K / 16;  // 8 packed bytes + one e4m3 scale per block
+  int b = 0;
+  for (; b + 4 <= nb; b += 4) {  // two blk2 calls -> 4 blocks, 2 accumulators
+    acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8,
+                                          xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]]));
+    acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8,
+                                          xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]]));
+  }
+  for (; b + 2 <= nb; b += 2)
+    acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8,
+                                          xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]]));
+  float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
+  for (; b < nb; ++b) {  // odd final 16-K block
+    const uint8_t* pk = packed + (size_t)b * 8;
+    const float* xeb = xe + (size_t)b * 8;
+    const float* xob = xo + (size_t)b * 8;
+    float bsum = 0.0f;
+    for (int j = 0; j < 8; ++j) {
+      bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j];
+    }
+    s += e4m3[scale[b]] * bsum;
+  }
+  return s * global;
+}
+
+// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes.
+CPU_MOE_TARGET("avx2,fma")
+inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob,
+                             __m256 mag8, float sc) {
+  __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk)));
+  __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8);
+  __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8);
+  __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob)));
+  return _mm256_mul_ps(prod, _mm256_set1_ps(sc));
+}
+
+CPU_MOE_TARGET("avx2,fma")
+float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global,
+                     const float* xe, const float* xo, int K, const float* e2m1,
+                     const float* e4m3) {
+  const __m256 mag8 = _mm256_loadu_ps(e2m1);  // e2m1[0..7] magnitudes
+  __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps();
+  const int nb = K / 16;
+  int b = 0;
+  for (; b + 2 <= nb; b += 2) {
+    acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8,
+                                              xo + (size_t)b * 8, mag8, e4m3[scale[b]]));
+    acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8,
+                                              xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]]));
+  }
+  for (; b < nb; ++b)
+    acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8,
+                                              xo + (size_t)b * 8, mag8, e4m3[scale[b]]));
+  return hsum256(_mm256_add_ps(acc0, acc1)) * global;
+}
+
+// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)]
+// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less
+// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout.
+CPU_MOE_TARGET("avx2,avxvnni,fma")
+inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) {
+  __m128i b = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(pk));   // 8 bytes
+  __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F));
+  __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F));
+  return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi));            // [lo(8),hi(8)] -> int8
+}
+
+// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a,
+// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32
+// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant
+// (W4A8) differs from the bf16 reference.
+CPU_MOE_TARGET("avx2,avxvnni,fma")
+float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global,
+                        const int8_t* asi8, int K, const float* e4m3, const float* asb) {
+  const __m128i lut = _mm_loadu_si128(reinterpret_cast<const __m128i*>(kE2M1x2));
+  __m256 accF = _mm256_setzero_ps();
+  const int nb = K / 16;
+  int b = 0;
+  for (; b + 2 <= nb; b += 2) {
+    __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut);
+    __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut);
+    __m256i w = _mm256_set_m128i(wb1, wb);                              // [blk b | blk b+1]
+    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(asi8 + (size_t)b * 16));
+    __m256i aw = _mm256_sign_epi8(w, w);                               // |w| (u8 operand)
+    __m256i sa = _mm256_sign_epi8(a, w);                               // sign(w)*a (s8 operand)
+    __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa);
+    __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]),
+                                 _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0);
+    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF);
+  }
+  float s = hsum256(accF);
+  for (; b < nb; ++b) {  // tail (odd block count)
+    const uint8_t* pk = packed + (size_t)b * 8;
+    const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8;
+    int isum = 0;
+    for (int j = 0; j < 8; ++j)
+      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
+    s += (e4m3[scale[b]] * asb[b]) * (float)isum;
+  }
+  return s * (0.5f * global);
+}
+
+#ifdef CPU_MOE_HAS_AVX512VNNI
+
+// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed
+// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then
+// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical
+// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row
+// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte
+// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that
+// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim
+// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20%
+// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it.
+// Prefetch never faults, so overshooting a row/bank tail is safe.
+static int nvfp4_pf_blocks() {
+  static const int v = [] {
+    const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS");
+    return (s && s[0]) ? atoi(s) : -1;
+  }();
+  return v;
+}
+
+// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI
+// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8
+// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so
+// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate.
+// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16.
+// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so
+// the caller can run several independent chains into separate accumulators (the decode ->
+// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound).
+CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2")
+static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale,
+                                   const int8_t* asi8, const float* e4m3, const float* asb,
+                                   int b, __m512i lut, __m512i idx, __m512i mask0F,
+                                   __m512i idxsc) {
+  const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL;  // bytes 8-15 of each 128b lane -> hi nibbles
+  __m256i raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(packed + (size_t)b * 8));
+  __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw));
+  __m512i lo = _mm512_and_si512(src, mask0F);
+  __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F);
+  __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi);
+  __m512i w = _mm512_shuffle_epi8(lut, comb);  // 64 int8 weights (e2m1*2)
+  // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0)
+  __m512i a = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(asi8 + (size_t)b * 16));
+  __m512i aw = _mm512_abs_epi8(w);
+  __mmask64 neg = _mm512_movepi8_mask(w);
+  __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a);
+  __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa);  // 16 int32 (groups of 4)
+  // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales
+  // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32
+  // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls +
+  // a set_ps assembly, which otherwise dominates the per-group op count.
+  int sc_raw;
+  memcpy(&sc_raw, scale + b, 4);
+  __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw));
+  __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b));
+  __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4));
+  return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv);
+}
+
+CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2")
+float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global,
+                              const int8_t* asi8, int K, const float* e4m3, const float* asb) {
+  const __m512i lut = _mm512_broadcast_i32x4(
+      _mm_loadu_si128(reinterpret_cast<const __m128i*>(kE2M1x2)));
+  const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0);  // block i -> 128b lane i
+  const __m512i mask0F = _mm512_set1_epi8(0x0F);
+  const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
+  __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps();
+  __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps();
+  const int nb = K / 16;
+  // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set
+  // env value is used verbatim so a per-machine sweep can reach operating points past
+  // the conservative default.
+  const int pfb = nvfp4_pf_blocks();
+  const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb;
+  int b = 0;
+  // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op
+  // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core
+  // rate sat at ~60% of the core's achievable DRAM stream rate).
+  for (; b + 16 <= nb; b += 16) {
+    // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss
+    // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids).
+    if (pf > 0) {
+      _mm_prefetch(reinterpret_cast<const char*>(packed + ((size_t)b + (size_t)pf) * 8),
+                   _MM_HINT_T0);
+      _mm_prefetch(reinterpret_cast<const char*>(packed + ((size_t)b + (size_t)pf) * 8 + 64),
+                   _MM_HINT_T0);
+    }
+    acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx,
+                                             mask0F, idxsc));
+    acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx,
+                                             mask0F, idxsc));
+    acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx,
+                                             mask0F, idxsc));
+    acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx,
+                                             mask0F, idxsc));
+  }
+  for (; b + 4 <= nb; b += 4)
+    acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx,
+                                             mask0F, idxsc));
+  float s = _mm512_reduce_add_ps(
+      _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3)));
+  for (; b < nb; ++b) {  // tail (<4 remaining 16-K blocks)
+    const uint8_t* pk = packed + (size_t)b * 8;
+    const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8;
+    int isum = 0;
+    for (int j = 0; j < 8; ++j)
+      isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j];
+    s += (e4m3[scale[b]] * asb[b]) * (float)isum;
+  }
+  return s * (0.5f * global);
+}
+#endif  // avx512vnni available
+#endif
+
+nvdot_fn select_nvdot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  if (t >= ISA_AVX512) return dot_nvfp4_avx512;
+  if (t >= ISA_AVX2) return dot_nvfp4_avx2;
+#endif
+  (void)t;
+  return dot_nvfp4_scalar;
+}
+
+// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the
+// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference.
+nvi8dot_fn select_nvi8dot() {
+#if CPU_MOE_X86
+#if defined(CPU_MOE_HAS_AVX512VNNI)
+  if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni;
+#endif
+  if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni;
+#endif
+  return nullptr;
+}
+
+// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the
+// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over
+// every output row), so a scalar pass is fine relative to the GEMV.
+void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) {
+  const int nb = K / 16;
+  for (int b = 0; b < nb; ++b) {
+    const float* xeb = xe + (size_t)b * 8;
+    const float* xob = xo + (size_t)b * 8;
+    float amax = 0.0f;
+    for (int j = 0; j < 8; ++j)
+      amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j])));
+    const float s = amax > 0.0f ? amax / 127.0f : 1.0f;
+    asb[b] = s;
+    const float inv = 1.0f / s;
+    int8_t* ae = asi8 + (size_t)b * 16;
+    for (int j = 0; j < 8; ++j) {
+      int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv);
+      ae[j] = (int8_t)std::max(-127, std::min(127, qe));
+      ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo));
+    }
+  }
+}
+
+class Nvfp4Format final : public MoeFormat {
+ public:
+  Nvfp4Format(const MoeParams& p, const FormatArgs& a)
+      : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) {
+    gate_up_ = a.bank("gate_up");
+    gu_scale_ = a.bank("gate_up_scale");
+    gu_global_ = a.bank("gate_up_global");
+    down_ = a.bank("down");
+    dn_scale_ = a.bank("down_scale");
+    dn_global_ = a.bank("down_global");
+    const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k;
+    xe_.resize(tokens * (p.H / 2));
+    xo_.resize(tokens * (p.H / 2));
+    ge_.resize(routes * (p.I / 2));
+    go_.resize(routes * (p.I / 2));
+    if (nvi8dot_ != nullptr) {
+      xi8_.resize(tokens * p.H);
+      xas_.resize(tokens * (p.H / 16));
+      gi8_.resize(routes * p.I);
+      gas_.resize(routes * (p.I / 16));
+    }
+    isa_ = tier_name(pick_isa());
+    if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)";
+    for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i];
+    for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i);
+  }
+
+  std::string isa_name() const override { return isa_; }
+
+
+  void prep_input(int tok, const bf16_t* x) override {
+    float* xe = xe_.data() + (size_t)tok * (p.H / 2);
+    float* xo = xo_.data() + (size_t)tok * (p.H / 2);
+    deinterleave_bf16_f32(x, xe, xo, p.H);
+    if (nvi8dot_ != nullptr)
+      quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16));
+  }
+
+  bool preps_inter() const override { return true; }
+
+  void prep_inter(int64_t r, bf16_t* g) override {
+    float* ge = ge_.data() + (size_t)r * (p.I / 2);
+    float* go = go_.data() + (size_t)r * (p.I / 2);
+    deinterleave_bf16_f32(g, ge, go, p.I);
+    if (nvi8dot_ != nullptr)
+      quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16));
+  }
+
+  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in,
+               bf16_t* g) const override {
+    const int H = p.H;
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
+    const uint8_t* scale = static_cast<const uint8_t*>(tbl_at(gu_scale_, layer));
+    const uint16_t* global = static_cast<const uint16_t*>(tbl_at(gu_global_, layer));
+    const float* xe = xe_.data() + (size_t)tok * (H / 2);
+    const float* xo = xo_.data() + (size_t)tok * (H / 2);
+    const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr;
+    const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr;
+    gated_rows(p, i0, i1, w_in, g, [&](int row) {
+      const size_t r = (size_t)e * (2 * p.I) + row;
+      if (nvi8dot_ != nullptr)
+        return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16),
+                        fp16_to_f32(global[r]), xi8, H, e4m3_lut_, xas);
+      return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16),
+                    fp16_to_f32(global[r]), xe, xo, H, e2m1_lut_, e4m3_lut_);
+    });
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const int I = p.I;
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    const uint8_t* scale = static_cast<const uint8_t*>(tbl_at(dn_scale_, layer));
+    const uint16_t* global = static_cast<const uint16_t*>(tbl_at(dn_global_, layer));
+    for (int h = h0; h < h1; ++h) {
+      for (int j = 0; j < n; ++j) {
+        const size_t r = (size_t)routes[j].r;
+        const size_t row = (size_t)routes[j].e * p.H + h;
+        out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr
+            ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]),
+                       gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16))
+            : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]),
+                     ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_);
+      }
+    }
+  }
+
+ private:
+  nvdot_fn nvdot_;
+  nvi8dot_fn nvi8dot_;  // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot)
+  // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp16 row globals [E,2I]; down likewise with K = I.
+  const uint64_t* gate_up_ = nullptr;
+  const uint64_t* gu_scale_ = nullptr;
+  const uint64_t* gu_global_ = nullptr;
+  const uint64_t* down_ = nullptr;
+  const uint64_t* dn_scale_ = nullptr;
+  const uint64_t* dn_global_ = nullptr;
+  std::string isa_;
+  float e2m1_lut_[16];
+  float e4m3_lut_[256];
+  // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale.
+  AlignedVector<float> xe_, xo_, ge_, go_;
+  AlignedVector<int8_t> xi8_, gi8_;
+  AlignedVector<float> xas_, gas_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_nvfp4_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<Nvfp4Format>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
new file mode 100644
index 0000000000..600a194ebc
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
@@ -0,0 +1,216 @@
+#include <stdexcept>
+#include <vector>
+
+#include "../format.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+// -------------------------------- Q4_0 (W4A8) --------------------------------
+// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed
+// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a
+// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches
+// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload
+// path streams.
+//
+// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0
+// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to
+// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high
+// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD
+// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then
+// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU
+// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close.
+using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int);
+
+float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
+  float acc = 0.0f;
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const uint8_t* blk = w + (size_t)b * 18;
+    uint16_t dh;
+    std::memcpy(&dh, blk, sizeof(dh));
+    const uint8_t* q = blk + 2;  // 16 nibble bytes
+    const int8_t* a = aq + (size_t)b * 32;
+    int isum = 0;
+    for (int j = 0; j < 16; ++j) {
+      isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j];       // elem j
+      isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j];    // elem 16+j
+    }
+    acc += fp16_to_f32(dh) * asb[b] * (float)isum;
+  }
+  return acc;
+}
+
+#if CPU_MOE_X86
+// fp16 block scale -> fp32 via HW F16C (single value in lane 0).
+CPU_MOE_TARGET("f16c")
+static inline float q4_scale(uint16_t h) {
+  return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h)));
+}
+
+// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low
+// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8).
+CPU_MOE_TARGET("avx2")
+static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) {
+  const __m128i qb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(blk + 2));
+  const __m128i lo = _mm_and_si128(qb, mask);
+  const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask);
+  return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight);
+}
+
+// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW +
+// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no
+// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the
+// avx512-tier fallback, since the block dot is 256-bit either way).
+CPU_MOE_TARGET("avx2,fma,f16c")
+float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
+  const __m128i mask = _mm_set1_epi8(0x0F);
+  const __m256i eight = _mm256_set1_epi8(8);
+  const __m256i ones16 = _mm256_set1_epi16(1);
+  __m256 accF = _mm256_setzero_ps();
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const uint8_t* blk = w + (size_t)b * 18;
+    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
+    uint16_t dh;
+    std::memcpy(&dh, blk, sizeof(dh));
+    __m256i wq = q4_unpack32(blk, mask, eight);
+    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
+    __m256i aw = _mm256_sign_epi8(wq, wq);              // |wq|          (unsigned operand)
+    __m256i sa = _mm256_sign_epi8(a, wq);               // sign(wq) * a  (signed operand)
+    __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16);  // 8 int32
+    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
+  }
+  return hsum256(accF);
+}
+
+// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs).
+CPU_MOE_TARGET("avx2,avxvnni,fma,f16c")
+float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
+  const __m128i mask = _mm_set1_epi8(0x0F);
+  const __m256i eight = _mm256_set1_epi8(8);
+  __m256 accF = _mm256_setzero_ps();
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const uint8_t* blk = w + (size_t)b * 18;
+    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
+    uint16_t dh;
+    std::memcpy(&dh, blk, sizeof(dh));
+    __m256i wq = q4_unpack32(blk, mask, eight);
+    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
+    __m256i aw = _mm256_sign_epi8(wq, wq);   // |wq|            (unsigned operand)
+    __m256i sa = _mm256_sign_epi8(a, wq);    // sign(wq) * a    (signed operand)
+    __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa);
+    // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and
+    // reduce at the end (equivalent to scale * block_total).
+    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
+  }
+  return hsum256(accF);
+}
+#endif  // CPU_MOE_X86
+
+// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to
+// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when
+// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers.
+q4dot_fn select_q4dot() {
+  const IsaTier t = pick_isa();
+#if CPU_MOE_X86
+  if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni;
+  if (t >= ISA_AVX2) return q4_0_dot_i8_avx2;
+#endif
+  (void)t;
+  return q4_0_dot_i8_scalar;
+}
+
+// Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in
+// natural order + one fp32 scale (amax/127) per block. Done once per token/route,
+// amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0.
+void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) {
+  const int nb = K / 32;
+  for (int b = 0; b < nb; ++b) {
+    const bf16_t* xb = x + (size_t)b * 32;
+    float xf[32], amax = 0.0f;
+    for (int j = 0; j < 32; ++j) {
+      xf[j] = bf16_to_f32(xb[j]);
+      amax = std::max(amax, std::fabs(xf[j]));
+    }
+    const float d = amax > 0.0f ? amax / 127.0f : 1.0f;
+    asb[b] = d;
+    const float inv = amax > 0.0f ? 1.0f / d : 0.0f;
+    int8_t* o = aq + (size_t)b * 32;
+    for (int j = 0; j < 32; ++j)
+      o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv)));
+  }
+}
+
+class Q4_0Format final : public MoeFormat {
+ public:
+  Q4_0Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_q4dot()) {
+    gate_up_ = a.bank("gate_up");
+    down_ = a.bank("down");
+    const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k;
+    xi8_.resize(tokens * p.H);
+    xas_.resize(tokens * (p.H / 32));
+    gi8_.resize(routes * p.I);
+    gas_.resize(routes * (p.I / 32));
+    if (p.H % 32 != 0 || p.I % 32 != 0)
+      throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32");
+    gu_row_bytes_ = (p.H / 32) * 18;  // K = H (gate_up rows)
+    dn_row_bytes_ = (p.I / 32) * 18;  // K = I (down rows)
+    // select_q4dot picks VPDPBUSD / VPMADDUBSW / scalar for the tier, so the tag reflects which of those it resolved to.
+    isa_ = std::string(tier_name(pick_isa())) + (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8");
+  }
+
+  std::string isa_name() const override { return isa_; }
+
+
+  void prep_input(int tok, const bf16_t* x) override {
+    quant_q8_0(x, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 32));
+  }
+
+  bool preps_inter() const override { return true; }
+
+  void prep_inter(int64_t r, bf16_t* g) override {
+    quant_q8_0(g, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 32));
+  }
+
+  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in,
+               bf16_t* g) const override {
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
+    const int8_t* xi8 = xi8_.data() + (size_t)tok * p.H;
+    const float* xas = xas_.data() + (size_t)tok * (p.H / 32);
+    gated_rows(p, i0, i1, w_in, g, [&](int row) {
+      return dot_(packed + ((size_t)e * (2 * p.I) + row) * (size_t)gu_row_bytes_, xi8, xas, p.H);
+    });
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    for (int h = h0; h < h1; ++h) {
+      for (int j = 0; j < n; ++j) {
+        const size_t r = (size_t)routes[j].r;
+        out[j * kRowBlock + h - h0] = dot_(packed + ((size_t)routes[j].e * p.H + h) * (size_t)dn_row_bytes_,
+                                          gi8_.data() + r * p.I, gas_.data() + r * (p.I / 32), p.I);
+      }
+    }
+  }
+
+ private:
+  q4dot_fn dot_;
+  // GGUF Q4_0 rows of 32-element blocks (fp16 scale + 16 nibble bytes): gate_up [E,2I,H/32*18], down [E,H,I/32*18].
+  const uint64_t* gate_up_ = nullptr;
+  const uint64_t* down_ = nullptr;
+  int gu_row_bytes_ = 0, dn_row_bytes_ = 0;
+  std::string isa_;
+  // Q8_0 activations: per-32-block int8 + fp32 scale, for the input and the intermediate.
+  AlignedVector<int8_t> xi8_, gi8_;
+  AlignedVector<float> xas_, gas_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_q4_0_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<Q4_0Format>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
index 1898b41189..8abcd1a098 100644
--- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
+++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
@@ -44,6 +44,18 @@
 #define cudaGetLastError hipGetLastError
 #endif
 
+#ifndef cudaMalloc
+#define cudaMalloc hipMalloc
+#endif
+
+#ifndef cudaFree
+#define cudaFree hipFree
+#endif
+
+#ifndef cudaMemset
+#define cudaMemset hipMemset
+#endif
+
 #ifndef cudaMallocHost
 #define cudaMallocHost hipHostMalloc
 #endif
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 4918a032fe..82f09a9bb8 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -346,8 +346,9 @@ def _decode_hybrid(
             alphas=cache.alphas_for_slots(self.layer_id),
             is_prefill=False,
         )
-        cpu_routed = cpu_routed_early if not _HYBRID_OVERLAP else executor.decode_sync(pending)
-        return gpu_routed + cpu_routed
+        if not _HYBRID_OVERLAP:
+            return gpu_routed + cpu_routed_early
+        return executor.decode_sync(pending, add_to=gpu_routed)
 
     def _prefill_routed(
         self,
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index e70cad31d3..9430ab2a84 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -490,17 +490,16 @@ def _cpu_moe_step_fns(ex, wl: Workload, E: int):
     disjoint experts (working set >> LLC) -- the DRAM-bound regime hybrid decode actually
     runs in; ``run()`` executes one step (releases the GIL inside the extension)."""
     bs, top_k = 1, wl.top_k
-    io = ex._io_for(bs)
+    io = ex._io
     io["x"].copy_(torch.randn(bs, wl.hidden, dtype=torch.bfloat16) * 0.1)
     io["w"].copy_(torch.rand(bs, top_k, dtype=torch.float32))
-    task = ex._task_for(0, bs)
     base = torch.arange(top_k, dtype=torch.int32)
 
     def set_ids(i: int) -> None:
         io["ids"].copy_((((i * top_k) + base) % E).view(bs, top_k))
 
     def run() -> None:
-        ex._ext.run_task(task)
+        ex._ext.run(0, bs)
 
     return set_ids, run
 
@@ -566,7 +565,7 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device,
     backend its bandwidth-matched fetch split: fetched : cpu-computed = pcie_ov : cpu_ov.
 
     Both sides hammer their own pinned banks flat out for ``seconds`` from a shared
-    barrier: a worker thread loops bs=1 CPU decode steps (``run_task`` releases the GIL)
+    barrier: a worker thread loops bs=1 CPU decode steps (``run`` releases the GIL)
     while the main thread loops full-layer ``copy_missing`` gathers (synchronized per copy
     so the DMA is really in flight, not just enqueued). Each side reports bytes / its own
     elapsed; the two windows differ by at most one CPU step + one gather.
@@ -663,7 +662,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
                 entry["isa_sweep"] = c["isa_sweep"]
                 # The tier label comes from the bf16 dot chooser, so it can over-report tiers
                 # the actual per-format kernel doesn't distinguish. Disclose the real story.
-                if fmt == "nvfp4":  # WF_NVFP4 is the only use_vnni format
+                if fmt == "nvfp4":  # nvfp4 is the only use_vnni format
                     _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI "
                                  "regardless of the forced tier")
                 elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"):
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index f519687ed6..bb9ee56385 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -1,18 +1,15 @@
 """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu).
 
-Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and
-the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the
-whole CUDA-graph-capturable sequence on the current stream:
+Owns the persistent CPU worker pool and one pinned IO block of ``max_tokens`` rows that
+every layer and batch size shares (one task is in flight at a time). ``decode`` issues
+the whole CUDA-graph-capturable sequence on the current stream:
 
-    D2H (hidden, topk_ids, topk_weights -> pinned)
-      -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool)
-      -> sync host node   (cudaLaunchHostFunc: spin until the pool drains)
+    D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs])
+      -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node)
+      -> sync   (wait for the CPU pool to finish)
       -> H2D (pinned expert output -> GPU)
 
-Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager
-``model.forward()`` at each batch size immediately before capturing it, so the
-first (eager) call materializes the stable pinned buffers + task pointers that
-the subsequent capture embeds in its host/memcpy nodes.
+With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h).
 """
 
 from __future__ import annotations
@@ -35,28 +32,27 @@
 # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a
 # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done"
 # flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are
-# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the
-# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so
+# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64 on CUDA, see
+# gpu_sync.h): they execute on the GPU front end with no SM-resident kernel, so
 # GPU utilization stays truthful during the CPU compute window. (The first cut used a
 # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic
 # power schedulers responded by clamping the CPU frequency -- a net decode regression on
-# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag
-# slot, so every captured decode graph rides the handshake. Where memops are unavailable
-# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot
-# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python
-# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] +
-# raise_if_unhealthy) instead of an indefinite stream stall.
+# power-coupled edge devices.) The doorbell word names the (layer, decode batch size)
+# task, so one ready/done pair serves every layer and every captured graph (gpu_sync.h).
+# Where memops are unavailable (some driver models, vGPU, old drivers -- functionally probed at
+# startup), decode keeps the host-func path (functional, slower). ROCm 7.14 cannot replay
+# that path from a graph, so capture fails loudly there instead. A Python watchdog thread
+# turns a wedged coordinator into a loud RuntimeError (via raise_if_unhealthy) instead of
+# an indefinite stream stall.
 # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff
 # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely.
 _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0"
-# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph
-# sizes plus any eager padded sizes); more than that is unheard of, and the overflow
-# just keeps the host-func path for the extra combos.
-_FLAG_SLOTS_PER_LAYER = 16
 # Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out.
 _SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0"
+# ds_fp4 runs its input FP8 round-trip on the GPU when it can; tests turn this off to compare against the CPU round-trip.
+_DSFP4_GPU_PREQUANT = True
 
-# Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the
+# Activation ids must match ActKind in csrc/cpu_moe/common.h. Id 3 is the
 # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue,
 # "gpt_oss_swiglu" is the same math fused inside the mxfp4 kernel.
 _ACT_IDS = {
@@ -70,10 +66,6 @@
     "swiglu_clamp": 4,
 }
 
-# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp.
-_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4,
-             "fp8_block": 5}
-
 
 def compiled_extension_supports(activation: str) -> bool:
     """Whether the compiled ``_cpu_moe`` extension can serve ``activation``
@@ -157,7 +149,7 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]:
 
 class CpuMoeExecutor:
     """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks
-    (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0 or fp8_block — see ``_WFMT_IDS`` / ``_resolve_banks``)."""
+    (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``)."""
 
     def __init__(
         self,
@@ -178,26 +170,20 @@ def __init__(
         from freetoken.moe.legacy_format import canonical_role
 
         fmt = fmt or cache.quant_format
-        if fmt not in _WFMT_IDS:
+        if fmt not in _cpu_moe.weight_formats():
             raise NotImplementedError(
                 f"--moe-strategy cpu/hybrid computes experts on the CPU and supports "
-                f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are "
+                f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are "
                 f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead."
             )
         if activation not in _ACT_IDS:
             raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}")
+        if apply_router_weight_on_input and fmt in ("mxfp4_triton", "ds_fp4"):
+            raise NotImplementedError(f"CPU MoE backend: {fmt} experts apply the router weight on the output only")
         # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without
         # error and silently computes the wrong activation in the generic
         # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4
         # handles its act inside the kernel and predates the marker.)
-        # A prebuilt .so that predates this format would accept the id and then take
-        # the wrong dequant branch -- silently wrong output rather than a crash.
-        if _WFMT_IDS[fmt] > getattr(_cpu_moe, "max_weight_format_id", lambda: 4)():
-            raise RuntimeError(
-                f"the compiled _cpu_moe extension predates the {fmt!r} expert format; "
-                "rebuild it with `python setup.py build_ext --inplace` (or reinstall "
-                "the wheel) before serving this checkpoint on the cpu/hybrid backend."
-            )
         if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton":
             supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)()
             if _ACT_IDS[activation] > supported:
@@ -218,28 +204,42 @@ def __init__(
         # The per-layer tensors and their pointer tables must outlive the executor
         # (C++ holds raw addresses into both).
         self._banks: list[torch.Tensor] = []
-        ptrs, (self.H, self.I) = self._resolve_banks(
+        banks, options, (self.H, self.I) = self._resolve_banks(
             {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt
         )
 
+        # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain
+        # that the C++ side runs single-threaded on the CUDA host-callback thread --
+        # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and
+        # the GPU waiting on it). When a GPU is present we run the numerically identical
+        # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell
+        # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode:
+        # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py).
+        self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" and _DSFP4_GPU_PREQUANT
+        # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1)
+        self._act_block = act_block
+        if fmt == "ds_fp4":
+            assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block"
+            options.update(act_block=int(act_block), input_prequant=int(self._gpu_prequant))
+        if self._gpu_prequant:
+            logger.info_rank0(
+                "ds_fp4: input FP8 round-trip moved to the GPU "
+                "(bit-identical grid; the CPU-side scalar round-trip is skipped)"
+            )
+
         # Decide the flag handshake up front (env + device + a functional stream-memop
         # probe): its coordinator needs a core of its own, which the auto thread sizing
         # below reserves (a coordinator time-slicing against the GEMV workers measurably
         # destabilizes throughput on fully-subscribed boxes).
         self._flag_sync = _FLAG_SYNC and device.type == "cuda"
-        self._cpu_moe = _cpu_moe  # module ref for the decode-path memop calls
-        if self._flag_sync:
-            probe_scratch = alloc_pinned_tensor(1, dtype=torch.int64)
-            probe_scratch.zero_()
-            if not _cpu_moe.memops_probe(
-                torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr()
-            ):
-                logger.info_rank0(
-                    "cpu-moe flag handshake unavailable: CUDA stream memory operations "
-                    "are not supported here (Windows WDDM / vGPU / old driver); using "
-                    "the cudaLaunchHostFunc sync"
-                )
-                self._flag_sync = False
+        if self._flag_sync and not _cpu_moe.memops_probe(torch.cuda.current_stream().cuda_stream):
+            logger.info_rank0(
+                "cpu-moe flag handshake unavailable: stream memory operations failed the "
+                "startup probe; using the cudaLaunchHostFunc sync"
+            )
+            self._flag_sync = False
+        # ROCm 7.14 returns wrong results when it replays host-function nodes, so there only the doorbell is graph-safe.
+        self._graph_capture_safe = self._flag_sync or _cpu_moe.host_func_capture_safe()
 
         nthreads, core_ids = resolve_threads_and_affinity(num_threads)
         coord_core = -1
@@ -250,6 +250,17 @@ def __init__(
             nthreads -= 1
             core_ids = core_ids[:-1]
         self._coord_core = coord_core
+        alloc_io = (
+            (lambda *shape, dtype: torch.empty(*shape, dtype=dtype))
+            if device.type == "cpu"
+            else alloc_pinned_tensor
+        )
+        self._io = {
+            "x": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16),
+            "ids": alloc_io(self.max_tokens, self.top_k, dtype=torch.int32),
+            "w": alloc_io(self.max_tokens, self.top_k, dtype=torch.float32),
+            "y": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16),
+        }
         self._ext = _cpu_moe.CpuMoeExecutor(
             num_threads=nthreads,
             num_layers=self.num_layers,
@@ -260,11 +271,16 @@ def __init__(
             max_tokens=self.max_tokens,
             activation_id=_ACT_IDS[activation],
             apply_router_weight_on_input=1 if apply_router_weight_on_input else 0,
-            weight_format=_WFMT_IDS[fmt],
+            weight_format=fmt,
+            banks=banks,
+            options=options,
             swiglu_alpha=float(swiglu_alpha),
             swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"),
             core_ids=core_ids,
-            **ptrs,
+            x_ptr=self._io["x"].data_ptr(),
+            ids_ptr=self._io["ids"].data_ptr(),
+            w_ptr=self._io["w"].data_ptr(),
+            y_ptr=self._io["y"].data_ptr(),
         )
         self.num_threads = nthreads
         self.core_ids = core_ids
@@ -281,34 +297,24 @@ def __init__(
             )
             torch.set_num_threads(clamp)
 
-        self._io: dict[int, dict[str, torch.Tensor]] = {}
-        self._tasks: dict[tuple[int, int], int] = {}
-
-        # Flag-based handshake: mapped-pinned ready/done/err int64 arrays (one slot per
-        # (MoE layer, decode batch size) pair, allocated as tasks are created) + a
-        # persistent CPU coordinator that polls ready[], runs the slot's task on the
-        # pool, and sets done[]. Binary per-step protocol (GPU memops: done=0, ready=1;
-        # coordinator: consume ready, run, done=1; GPU waits done>=1 -- the WAIT
-        # immediate is constant, so CUDA-graph replays are safe). err[] is raised by the
-        # WATCHDOG thread when a ready flag stays unanswered (dead coordinator): it
-        # poisons done to unblock the stream and raise_if_unhealthy() turns the step
-        # into a loud error instead of silent stale output. Buffers are kept alive on
-        # self so the coordinator's pinned pointers stay valid for the executor's
-        # lifetime (flag_sync itself was decided above, before thread sizing).
-        self._ready = self._done = self._err = None
-        self._flag_slots: dict[tuple[int, int], int] = {}  # (layer_id, bs) -> slot
-        self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER
+        # Flag-based handshake: the extension owns the ready/done words (gpu_sync.h) and a
+        # persistent CPU coordinator that polls the doorbell, runs the task it names on the
+        # pool, and sets done for the GPU. The WATCHDOG thread marks the executor unhealthy
+        # when a doorbell stays unanswered (dead coordinator): it poisons done to unblock the
+        # stream and raise_if_unhealthy() turns the step into a loud error instead of silent
+        # stale output (flag_sync itself was decided above, before thread sizing).
+        self._unhealthy = False
+        self._kernel_io = False
         if self._flag_sync:
-            self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64)
-            self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64)
-            self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64)
-            self._ready.zero_()
-            self._done.zero_()
-            self._err.zero_()
-            self._ext.start_flag_coordinator(
-                self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity,
-                self._coord_core,
-            )
+            self._ext.start_flag_coordinator(self._coord_core)
+            # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs.
+            if torch.version.hip is None:
+                reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream)
+                if reason:
+                    logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell")
+                else:
+                    self._ext.enable_kernel_io()
+                    self._kernel_io = True
             self._watchdog_stop = False
             # The thread target holds a WEAKREF and re-derefs it each tick: a bound
             # method would strong-reference the executor forever (the loop never ends
@@ -325,32 +331,14 @@ def __init__(
             )
             self._watchdog.start()
 
-        # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain
-        # that the C++ side runs single-threaded on the CUDA host-callback thread --
-        # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and
-        # the GPU waiting on it). When a GPU is present we run the numerically identical
-        # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell
-        # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode:
-        # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py).
-        self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda"
-        # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1)
-        self._act_block = act_block
-        if fmt == "ds_fp4":
-            assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block"
-            self._ext.set_act_block(act_block)
-        if self._gpu_prequant:
-            self._ext.set_input_prequant(True)
-            logger.info_rank0(
-                "ds_fp4: input FP8 round-trip moved to the GPU "
-                "(bit-identical grid; the CPU-side scalar round-trip is skipped)"
-            )
 
         logger.info_rank0(
             f"CPU MoE executor ready: threads={nthreads} (pinned to cores "
             f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} "
             f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} "
             f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} "
-            f"spin_wait={self.spin_wait}"
+            f"spin_wait={self.spin_wait} "
+            f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}"
         )
 
     def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor:
@@ -368,7 +356,11 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor:
         self._banks.extend(layers)
         return table
 
-    def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
+    def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]:
+        """Per-layer table addresses keyed by canonical bank role."""
+        return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()}
+
+    def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
         """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride."""
         gu, gus = banks["gate_up"], banks["gate_up_scale"]
         dn, dns = banks["down"], banks["down_scale"]
@@ -380,28 +372,19 @@ def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I)
         assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I)
         assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I)
-        ptrs = dict(
-            gate_up_ptr=self._make_table(gu).data_ptr(),
-            down_ptr=self._make_table(dn).data_ptr(),
-            gate_up_scale_ptr=self._make_table(gus).data_ptr(),
-            gate_up_global_ptr=0,
-            down_scale_ptr=self._make_table(dns).data_ptr(),
-            down_global_ptr=0,
-            gate_up_bias_ptr=0,
-            down_bias_ptr=0,
-            fp8_block_gu_scale_stride=int(gus[0].stride(1)),
-            fp8_block_dn_scale_stride=int(dns[0].stride(1)),
-        )
-        return ptrs, (H, I)
+        banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns)
+        options = dict(gu_scale_stride=int(gus[0].stride(1)), dn_scale_stride=int(dns[0].stride(1)))
+        return banks, options, (H, I)
 
-    def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]:
-        """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format.
+    def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int, int]]:
+        """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role
+        the format reads to its per-layer table's address, options holds its integer settings.
 
         ``banks[name]`` is a list of ``num_layers`` ``[num_experts, ...]`` tensors
         (the per-layer host bank contract); shapes are read from the first layer so
-        per-partition (TP) sizes are exact. Unused pointers are 0. Every pointer kwarg
-        is actually a per-layer table's address (see ``_make_table``), not a single
-        bank's -- the C++ ctor resolves ``tbl[layer_id]`` per task.
+        per-partition (TP) sizes are exact. Every bank address is a per-layer table
+        (see ``_make_table``), not a single bank's -- the C++ side resolves
+        ``tbl[layer_id]`` per task.
         """
         if fmt == "bf16":
             gate_up = banks["gate_up"]
@@ -414,17 +397,8 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]:
             I = int(gate_up[0].shape[1] // 2)
             assert gate_up[0].shape[1] == 2 * I
             assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I)
-            ptrs = dict(
-                gate_up_ptr=self._make_table(gate_up).data_ptr(),
-                down_ptr=self._make_table(down).data_ptr(),
-                gate_up_scale_ptr=0,
-                gate_up_global_ptr=0,
-                down_scale_ptr=0,
-                down_global_ptr=0,
-                gate_up_bias_ptr=0,
-                down_bias_ptr=0,
-            )
-            return ptrs, (H, I)
+            banks = self._tables(gate_up=gate_up, down=down)
+            return banks, {}, (H, I)
 
         if fmt == "q4_0":
             return self._resolve_q4_0_banks(banks)
@@ -452,19 +426,10 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]:
         assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H)
         assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I)
         assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,)
-        ptrs = dict(
-            gate_up_ptr=self._make_table(gup).data_ptr(),
-            down_ptr=self._make_table(dnp).data_ptr(),
-            gate_up_scale_ptr=self._make_table(gus).data_ptr(),
-            gate_up_global_ptr=self._make_table(gug).data_ptr(),
-            down_scale_ptr=self._make_table(dns).data_ptr(),
-            down_global_ptr=self._make_table(dng).data_ptr(),
-            gate_up_bias_ptr=0,
-            down_bias_ptr=0,
-        )
-        return ptrs, (H, I)
+        banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng)
+        return banks, {}, (H, I)
 
-    def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
+    def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
         """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble
         bytes), row-major over K -- the *same* packed banks the GPU offload path streams.
         gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a
@@ -479,19 +444,10 @@ def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         assert H % 32 == 0 and I % 32 == 0, (H, I)
         assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H)
         assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I)
-        ptrs = dict(
-            gate_up_ptr=self._make_table(gate_up).data_ptr(),
-            down_ptr=self._make_table(down).data_ptr(),
-            gate_up_scale_ptr=0,
-            gate_up_global_ptr=0,
-            down_scale_ptr=0,
-            down_global_ptr=0,
-            gate_up_bias_ptr=0,
-            down_bias_ptr=0,
-        )
-        return ptrs, (H, I)
+        banks = self._tables(gate_up=gate_up, down=down)
+        return banks, {}, (H, I)
 
-    def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
+    def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
         """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales
         (N innermost) + per-output-row biases. The C++ kernel streams K and
         accumulates a contiguous N-block, so the GPU-tiled layout is read in place
@@ -510,19 +466,10 @@ def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I)
         assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I)
         assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,)
-        ptrs = dict(
-            gate_up_ptr=self._make_table(gub).data_ptr(),
-            down_ptr=self._make_table(dnb).data_ptr(),
-            gate_up_scale_ptr=self._make_table(gus).data_ptr(),
-            gate_up_global_ptr=0,
-            down_scale_ptr=self._make_table(dns).data_ptr(),
-            down_global_ptr=0,
-            gate_up_bias_ptr=self._make_table(gob).data_ptr(),
-            down_bias_ptr=self._make_table(dob).data_ptr(),
-        )
-        return ptrs, (H, I)
+        banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob)
+        return banks, {}, (H, I)
 
-    def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
+    def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
         """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block
         scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row),
         so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the
@@ -538,57 +485,8 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]:
         assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I)
         assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H)
         assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I)
-        ptrs = dict(
-            gate_up_ptr=self._make_table(gup).data_ptr(),
-            down_ptr=self._make_table(dnp).data_ptr(),
-            gate_up_scale_ptr=self._make_table(gus).data_ptr(),
-            gate_up_global_ptr=0,
-            down_scale_ptr=self._make_table(dns).data_ptr(),
-            down_global_ptr=0,
-            gate_up_bias_ptr=0,
-            down_bias_ptr=0,
-        )
-        return ptrs, (H, I)
-
-    def _io_for(self, bs: int) -> dict[str, torch.Tensor]:
-        io = self._io.get(bs)
-        if io is None:
-            alloc_io = (
-                (lambda *shape, dtype: torch.empty(*shape, dtype=dtype))
-                if self.device.type == "cpu"
-                else alloc_pinned_tensor
-            )
-            io = {
-                "x": alloc_io(bs, self.H, dtype=torch.bfloat16),
-                "ids": alloc_io(bs, self.top_k, dtype=torch.int32),
-                "w": alloc_io(bs, self.top_k, dtype=torch.float32),
-                "y": alloc_io(bs, self.H, dtype=torch.bfloat16),
-            }
-            self._io[bs] = io
-        return io
-
-    def _task_for(self, layer_id: int, bs: int) -> int:
-        key = (layer_id, bs)
-        task = self._tasks.get(key)
-        if task is None:
-            io = self._io_for(bs)
-            task = self._ext.create_task(
-                layer_id,
-                bs,
-                io["x"].data_ptr(),
-                io["ids"].data_ptr(),
-                io["w"].data_ptr(),
-                io["y"].data_ptr(),
-            )
-            self._tasks[key] = task
-            # Allocate this (layer, bs) combo a flag slot and register its task with the
-            # coordinator. Combos past the slot capacity keep the host-func path.
-            if self._flag_sync and key not in self._flag_slots:
-                slot = len(self._flag_slots)
-                if slot < self._flag_capacity:
-                    self._flag_slots[key] = slot
-                    self._ext.register_flag_task(slot, task)
-        return task
+        banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns)
+        return banks, {}, (H, I)
 
     def decode(
         self,
@@ -612,7 +510,7 @@ def decode_submit(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> tuple:
-        """Issue the D2H copies + the CPU-pool submit host node, then return without
+        """Issue the D2H copies + the CPU-pool submit host node (or the publish kernel), then return without
         waiting. Lets a caller (the hybrid backend) enqueue GPU work between this and
         :meth:`decode_sync` so the CPU compute overlaps the GPU GEMM / PCIe fetch.
 
@@ -621,7 +519,14 @@ def decode_submit(
         :meth:`decode_sync`. The output tensor is allocated here so it stays live (and
         distinct from the interleaved GPU work) across the overlap window."""
         bs = hidden_states.shape[0]
-        io = self._io_for(bs)
+        if bs > self.max_tokens:
+            raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}")
+        if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop "
+                "doorbell is off or failed its startup probe, and the host-function sync is not "
+                "replay-safe on this runtime; run with --cuda-graph-max-bs 0"
+            )
 
         if self._gpu_prequant:
             # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the
@@ -631,91 +536,86 @@ def decode_submit(
 
             hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block)
 
+        stream = torch.cuda.current_stream().cuda_stream
+        if self._kernel_io and hidden_states.dtype == torch.bfloat16:
+            x = hidden_states.contiguous()
+            ids = topk_ids.to(torch.int32).contiguous()
+            w = topk_weights.to(torch.float32).contiguous()
+            assert x.shape == (bs, self.H) and ids.shape == w.shape == (bs, self.top_k), (x.shape, ids.shape, w.shape)
+            self._ext.flag_publish(stream, x.data_ptr(), ids.data_ptr(), w.data_ptr(), layer_id, bs)
+            return (bs, torch.empty_like(x), True)
+
         # D2H: ship this step's activations + routing to pinned host memory.
-        io["x"].copy_(hidden_states, non_blocking=True)
-        io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True)
-        io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True)
+        io = self._io
+        io["x"][:bs].copy_(hidden_states, non_blocking=True)
+        io["ids"][:bs].copy_(topk_ids.to(torch.int32), non_blocking=True)
+        io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True)
 
-        task = self._task_for(layer_id, bs)
         out = torch.empty_like(hidden_states)
-        slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None
-        if slot is not None:
-            # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's
-            # doorbell). No kernel launched; no host-func round trip.
-            self._cpu_moe.memop_submit(
-                torch.cuda.current_stream().cuda_stream,
-                self._done.data_ptr(), self._ready.data_ptr(), slot,
-            )
+        if self._flag_sync:
+            # Front-end memops: reset done, then ring the doorbell with (layer, bs).
+            # No kernel launched; no host-func round trip.
+            self._ext.flag_signal(stream, layer_id, bs)
         else:
-            stream = torch.cuda.current_stream().cuda_stream
-            self._ext.submit_with_cuda_stream(stream, task)
-        return (bs, task, out, slot)
+            self._ext.submit_with_cuda_stream(stream, layer_id, bs)
+        return (bs, out, False)
 
-    def decode_sync(self, pending: tuple) -> torch.Tensor:
+    def decode_sync(self, pending: tuple, add_to: torch.Tensor | None = None) -> torch.Tensor:
         """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`,
         and return the GPU output tensor. With flag-sync the wait is a front-end stream
-        memop on done[slot] (set by the CPU coordinator); otherwise a cudaLaunchHostFunc."""
-        bs, task, out, slot = pending
-        if slot is not None:
-            # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without
+        memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.
+        With ``add_to``, return ``add_to`` + the CPU output; the kernel path adds in place into ``add_to``."""
+        bs, out, kernel_io = pending
+        stream = torch.cuda.current_stream().cuda_stream
+        if self._flag_sync:
+            # Front-end WAIT(done >= 1): blocks this stream's later nodes without
             # occupying an SM, so GPU utilization stays truthful during the CPU window.
-            self._cpu_moe.memop_sync(
-                torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot,
-            )
+            self._ext.flag_wait(stream)
         else:
-            stream = torch.cuda.current_stream().cuda_stream
-            self._ext.sync_with_cuda_stream(stream, task)
-        io = self._io[bs]
-        out.copy_(io["y"], non_blocking=True)
-        return out
-
-    def _watchdog_tick(self, suspects: dict) -> None:
+            self._ext.sync_with_cuda_stream(stream)
+        if kernel_io:
+            if add_to is not None and add_to.dtype == torch.bfloat16 and add_to.is_contiguous() and add_to.shape == out.shape:
+                self._ext.flag_combine(stream, add_to.data_ptr(), bs, True)
+                return add_to
+            self._ext.flag_combine(stream, out.data_ptr(), bs, False)
+            return out if add_to is None else add_to + out
+        out.copy_(self._io["y"][:bs], non_blocking=True)
+        return out if add_to is None else add_to + out
+
+    def _watchdog_tick(self, suspect: dict) -> None:
         """One watchdog sampling round (called every 2 s by ``_watchdog_main``).
 
-        A slot is only declared dead when THREE things hold across >=10 s: its doorbell
-        is still pending (ready==1 && done==0), it was already pending when first
-        suspected, and the coordinator has served NOTHING on it since (flag_served_count
-        unchanged). The served-count criterion kills the false-positive window: two
-        point samples can land on the same slot's (independent, us-scale) pending
-        windows under heavy external load, but a coordinator that made progress in
-        between is alive by definition. ``suspects`` maps slot -> (first_seen,
-        served_at_first_sight) and persists across ticks."""
-        stuck = (self._ready == 1) & (self._done == 0)
-        if not bool(stuck.any()):
-            suspects.clear()
+        The doorbell is only declared dead when THREE things hold across >=10 s: it is
+        still pending (ready set, done==0), it was already pending when first suspected, and
+        the coordinator has served NOTHING since (flag_served unchanged). The served-count
+        criterion kills the false-positive window: two point samples can land on
+        (independent, us-scale) pending windows under heavy external load, but a
+        coordinator that made progress in between is alive by definition. ``suspect``
+        holds (first_seen, served_at_first_sight) and persists across ticks."""
+        if not self._ext.flag_pending():
+            suspect.clear()
             return
         now = time.monotonic()
-        pending = set(stuck.nonzero().flatten().tolist())
-        for slot in list(suspects):
-            if slot not in pending:
-                del suspects[slot]
-        dead = []
-        for slot in pending:
-            served = self._ext.flag_served_count(slot)
-            first_seen, served_then = suspects.get(slot, (None, None))
-            if first_seen is None or served != served_then:
-                suspects[slot] = (now, served)  # new suspect, or alive-but-loaded: rearm
-                continue
-            if now - first_seen >= 10.0:
-                dead.append(slot)
-        if not dead:
+        served = self._ext.flag_served()
+        if not suspect or suspect["served"] != served:
+            suspect.update(first_seen=now, served=served)  # new suspect, or alive-but-loaded: rearm
+            return
+        if now - suspect["first_seen"] < 10.0:
             return
         logger.error(
-            f"cpu-moe flag watchdog: slots {dead} unanswered for >10s with no coordinator "
-            "progress (wedged/dead); poisoning done[] and failing the next step"
+            "cpu-moe flag watchdog: doorbell unanswered for >10s with no coordinator "
+            "progress (wedged/dead); poisoning done and failing the next step"
         )
-        for i in dead:
-            self._err[i] = 1
-        for i in dead:
-            self._done[i] = 1  # after err: unblock the stream into a checked failure
-            suspects.pop(i, None)
+        self._unhealthy = True
+        self._ext.poison_flag()  # after the mark: unblock the stream into a checked failure
+        suspect.clear()
 
     def raise_if_unhealthy(self) -> None:
         """Raise if the flag watchdog fired (a doorbell stayed unanswered because the
         coordinator never responded). Called by the engine once per forward -- a single
-        pinned read -- so a dead coordinator surfaces as a loud error on the next step
+        flag read -- so a dead coordinator surfaces as a loud error on the next step
         instead of silently shipping stale expert outputs."""
-        if self._err is not None and bool((self._err != 0).any()):
+        if self._unhealthy:
             raise RuntimeError(
                 "CPU MoE flag-handshake watchdog fired: a decode step's doorbell was "
                 "never answered by the coordinator thread (its outputs cannot be "
@@ -728,13 +628,13 @@ def raise_if_unhealthy(self) -> None:
 def _watchdog_main(executor_ref) -> None:
     """Watchdog daemon body: weakref-deref per tick so the thread never keeps a dead
     executor alive (see the start site in ``CpuMoeExecutor.__init__``)."""
-    suspects: dict = {}
+    suspect: dict = {}
     while True:
         time.sleep(2.0)
         executor = executor_ref()
-        if executor is None or executor._watchdog_stop or executor._ready is None:
+        if executor is None or executor._watchdog_stop:
             return
         try:
-            executor._watchdog_tick(suspects)
+            executor._watchdog_tick(suspect)
         finally:
             del executor  # drop the strong ref before the next sleep
diff --git a/setup.py b/setup.py
index b8e4693c89..b968373151 100644
--- a/setup.py
+++ b/setup.py
@@ -77,15 +77,22 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
 if IS_ROCM:
     runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths()
     runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"]
-    # These extensions contain host code only. BuildExtension supplies the ROCm
+    # Host sources. BuildExtension supplies the ROCm
     # platform defines to the C++ compiler; offload architecture flags belong on
     # HIP device sources and would be rejected by the host compiler here.
     extra_compile = ["-O3", "-std=c++20"]
+    # Device sources (.cu, built by hipcc); BuildExtension adds the offload architectures.
+    device_compile = ["-O3", "-std=c++20"]
 else:
     runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths()
     runtime_lib = "cudart"
     runtime_link_args = []
     extra_compile = ["-O3", "-std=c++20"]
+    # Device sources (.cu, built by nvcc): SASS for each major arch, which also runs on the newer minors of that major,
+    # plus PTX for the oldest arch nvcc 13 supports, which the driver builds for any other GPU at load.
+    # Naming archs also stops torch from picking them, which fails on a build machine without a GPU.
+    device_compile = ["-O3", "-std=c++20"] + [f"-gencode=arch=compute_{a},code=sm_{a}" for a in (75, 80, 90, 100, 120)]
+    device_compile.append("-gencode=arch=compute_75,code=compute_75")
 
 _check_toolchain()
 
@@ -112,11 +119,17 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
             name="freetoken.kernel._cpu_moe",
             sources=[
                 "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp",
+                f"python/freetoken/kernel/csrc/cpu_moe/gpu_sync_{'hip' if IS_ROCM else 'cuda'}.cpp",
+                "python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu",
+                *sorted(
+                    str(p.relative_to(ROOT))
+                    for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp")
+                ),
             ],
             include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs],
             library_dirs=runtime_library_dirs,
             libraries=[runtime_lib],
-            extra_compile_args=extra_compile + ["-pthread"],
+            extra_compile_args={"cxx": extra_compile + ["-pthread"], "nvcc": device_compile},
             extra_link_args=runtime_link_args,
         ),
         # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies
diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp
new file mode 100644
index 0000000000..4426bf0d14
--- /dev/null
+++ b/tests/moe/hip_probe_faults.cpp
@@ -0,0 +1,75 @@
+#include <hip/hip_runtime.h>
+#include <dlfcn.h>
+
+#include <atomic>
+#include <cstdlib>
+#include <cstring>
+
+namespace {
+std::atomic<int> begins{0}, ends{0}, injected{0}, removed_nodes{0}, last_end_rc{0};
+
+bool mode(const char* value) {
+  const char* selected = std::getenv("FREETOKEN_TEST_HIP_PROBE_FAULT");
+  return selected != nullptr && std::strcmp(selected, value) == 0;
+}
+
+template <typename Function>
+Function real(const char* name) {
+  return reinterpret_cast<Function>(dlsym(RTLD_NEXT, name));
+}
+}  // namespace
+
+extern "C" hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode capture_mode) {
+  const hipError_t rc = real<decltype(&hipStreamBeginCapture)>("hipStreamBeginCapture")(stream, capture_mode);
+  if (rc == hipSuccess) ++begins;
+  return rc;
+}
+
+extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, uint64_t value, unsigned flags) {
+  // Force the compatibility path on runtimes that normally record native writes.
+  if (mode("dependency") && begins == 1) {
+    ++injected;
+    return hipSuccess;
+  }
+  return real<decltype(&hipStreamWriteValue64)>("hipStreamWriteValue64")(stream, address, value, flags);
+}
+
+extern "C" hipError_t hipStreamUpdateCaptureDependencies(
+    hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) {
+  if (mode("dependency")) {
+    ++injected;
+    return hipErrorNotSupported;
+  }
+  return real<decltype(&hipStreamUpdateCaptureDependencies)>("hipStreamUpdateCaptureDependencies")(
+      stream, dependencies, count, flags);
+}
+
+extern "C" hipError_t hipGraphDestroyNode(hipGraphNode_t node) {
+  const hipError_t rc = real<decltype(&hipGraphDestroyNode)>("hipGraphDestroyNode")(node);
+  if (rc == hipSuccess) ++removed_nodes;
+  return rc;
+}
+
+extern "C" hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* graph) {
+  ++ends;
+  hipError_t rc = real<decltype(&hipStreamEndCapture)>("hipStreamEndCapture")(stream, graph);
+  if (mode("end-capture") && ends == 2 && rc == hipSuccess) {
+    // A real invalid API call also sets HIP's thread-local last-error state.
+    rc = real<decltype(&hipGraphGetNodes)>("hipGraphGetNodes")(*graph, nullptr, nullptr);
+    ++injected;
+  }
+  last_end_rc = static_cast<int>(rc);
+  return rc;
+}
+
+extern "C" int probe_test_injected() { return injected; }
+extern "C" int probe_test_begins() { return begins; }
+extern "C" int probe_test_ends() { return ends; }
+extern "C" int probe_test_removed_nodes() { return removed_nodes; }
+extern "C" int probe_test_last_end_rc() { return last_end_rc; }
+extern "C" int probe_test_peek_error() {
+  return static_cast<int>(real<decltype(&hipPeekAtLastError)>("hipPeekAtLastError")());
+}
+extern "C" int probe_test_make_error() {
+  return static_cast<int>(real<decltype(&hipGraphGetNodes)>("hipGraphGetNodes")(nullptr, nullptr, nullptr));
+}
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index 547c94c60a..38e79b2743 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -440,6 +440,111 @@ def _reference(cache, layer, hidden, ids, w, act="silu", apply_in=False):
     return out
 
 
+@pytest.fixture(scope="module")
+def hip_probe_fault_library(tmp_path_factory):
+    import importlib.util
+    import os
+    from pathlib import Path
+    import shutil
+    import subprocess
+
+    if not torch.version.hip:
+        pytest.skip("requires the HIP sync backend")
+    compiler = shutil.which("c++")
+    if compiler is None:
+        pytest.skip("requires a C++ compiler for HIP fault injection")
+    candidates = [Path(os.getenv("ROCM_HOME", "/opt/rocm")), Path("/opt/rocm")]
+    spec = importlib.util.find_spec("_rocm_sdk_core")
+    if spec and spec.submodule_search_locations:
+        candidates.append(Path(next(iter(spec.submodule_search_locations))))
+    include = next((root / "include" for root in candidates if (root / "include/hip/hip_runtime.h").exists()), None)
+    if include is None:
+        pytest.skip("requires HIP SDK headers for fault injection")
+    library = tmp_path_factory.mktemp("hip_probe") / "faults.so"
+    command = [
+        compiler, "-std=c++20", "-shared", "-fPIC", "-D__HIP_PLATFORM_AMD__=1",
+        f"-I{include}", str(Path(__file__).with_name("hip_probe_faults.cpp")),
+        "-ldl", "-o", str(library),
+    ]
+    result = subprocess.run(command, capture_output=True, text=True, timeout=60)
+    assert result.returncode == 0, result.stdout + result.stderr
+    return library
+
+
+@pytest.mark.parametrize("fault", ["dependency", "end-capture", "pending"])
+def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault):
+    """Rejected probes leave eager usable and preserve errors predating the probe."""
+    import os
+    from pathlib import Path
+    import subprocess
+    import sys
+    import textwrap
+
+    script = textwrap.dedent("""
+        import ctypes
+        import os
+        import runpy
+        import sys
+
+        import pytest
+        import torch
+        from freetoken.kernel import _cpu_moe
+        from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+        shim = ctypes.CDLL(os.environ["LD_PRELOAD"])
+        stream_ptr = torch.cuda.current_stream().cuda_stream
+        fault = os.environ["FREETOKEN_TEST_HIP_PROBE_FAULT"]
+        if fault == "pending":
+            assert shim.probe_test_make_error() != 0
+            pending = shim.probe_test_peek_error()
+            with pytest.raises(RuntimeError, match="pending error before probe"):
+                _cpu_moe.memops_probe(stream_ptr)
+            assert shim.probe_test_peek_error() == pending
+            assert shim.probe_test_begins() == 0
+            sys.exit(0)
+
+        helpers = runpy.run_path(sys.argv[1])
+        torch.manual_seed(17)
+        cache = helpers["_make_cache"](2, 8, 256, 128, scale=0.03)
+        ex = CpuMoeExecutor(
+            cache, top_k=4, activation="silu", apply_router_weight_on_input=False,
+            num_threads=4, max_tokens=3, device=torch.device("cuda"),
+        )
+        assert not ex._flag_sync and not ex._graph_capture_safe
+        assert shim.probe_test_injected() > 0
+        assert shim.probe_test_begins() == shim.probe_test_ends()
+        assert shim.probe_test_peek_error() == 0
+        if fault == "dependency":
+            assert shim.probe_test_removed_nodes() == 1
+            assert shim.probe_test_last_end_rc() == 0
+        else:
+            assert shim.probe_test_last_end_rc() != 0
+
+        for index in range(2):
+            torch.manual_seed(23 + index)
+            hidden = (torch.randn(3, 256) * 0.3).bfloat16().cuda()
+            ids = torch.randint(0, 8, (3, 4), dtype=torch.int32).cuda()
+            weights = torch.rand(3, 4).cuda()
+            actual = ex.decode(1, hidden, weights, ids).float().cpu()
+            expected = helpers["_reference"](cache, 1, hidden, ids.cpu(), weights.cpu())
+            relative = (actual - expected).abs().max() / (expected.abs().max() + 1e-6)
+            assert relative < 2e-2, relative
+
+        stream = torch.cuda.Stream()
+        with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"):
+            with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream):
+                ex.decode(1, hidden, weights, ids)
+        assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all()
+    """)
+    env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library),
+               FREETOKEN_TEST_HIP_PROBE_FAULT=fault, FREETOKEN_CPU_MOE_FLAG_SYNC="1")
+    result = subprocess.run(
+        [sys.executable, "-c", script, str(Path(__file__).resolve())],
+        env=env, capture_output=True, text=True, timeout=30,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
 def test_cpu_moe_decode_cuda_graph_replay():
     from freetoken.moe.cpu_executor import CpuMoeExecutor
 
@@ -467,7 +572,7 @@ def test_cpu_moe_decode_cuda_graph_replay():
     ids = torch.randint(0, E, (bs, top_k), device=dev, dtype=torch.int32)
     w = torch.rand(bs, top_k, device=dev, dtype=torch.float32)
 
-    # Eager forward: also materializes the per-bs pinned buffers + per-(layer,bs) task.
+    # Eager forward.
     out_eager = ex.decode(layer, hidden, w, ids)
     torch.cuda.synchronize()
     ref_eager = _reference(cache, layer, hidden, ids.cpu(), w.cpu())
@@ -500,11 +605,9 @@ def test_cpu_moe_decode_cuda_graph_replay():
     # must actually have exercised it -- guard against a silent fallback to the
     # host-func path making this test vacuous for the capture-embedded memops.
     if ex._flag_sync:
-        assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task"
-        slot = ex._flag_slots[(layer, bs)]
-        assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected"
-        assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest"
-        assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation"
+        assert ex._ext.flag_served() >= 4, "1 eager + 3 replay dispatches expected"
+        assert not ex._ext.flag_pending(), "handshake at rest"
+        assert not ex._unhealthy, "watchdog must not fire in normal operation"
         ex.raise_if_unhealthy()
 
     print("cpu moe cuda graph replay OK")
diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py
index 06f954d8d9..baecb86f9a 100644
--- a/tests/moe/test_cpu_moe_fp8_block.py
+++ b/tests/moe/test_cpu_moe_fp8_block.py
@@ -74,11 +74,11 @@ def test_cpu_fp8_block_accepts_qwen38_padded_scale_bank():
         device=torch.device("cpu"),
     )
     hidden = torch.randn(1, H, dtype=torch.bfloat16)
-    io = ex._io_for(1)
+    io = ex._io
     io["x"].copy_(hidden)
     io["ids"].zero_()
     io["w"].fill_(1.0)
-    ex._ext.run_task(ex._task_for(0, 1))
+    ex._ext.run(0, 1)
     cpu_out = io["y"][0].float().to(torch.float64)
     banks = cache.bank_sources
     gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0])
@@ -140,11 +140,8 @@ def _reference_fp8_block_gemv(weights, scales, x):
 
 
 @pytest.mark.parametrize("codes", ["finite", "subnormal"])
-@pytest.mark.parametrize(
-    "isa, expected_isa",
-    [("scalar", "scalar"), ("avx2", "avx2"), ("avx512", "avx512f"), ("avx512bf16", "avx512f")],
-)
-def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, codes):
+@pytest.mark.parametrize("isa", ["scalar", "avx2", "avx512"])
+def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, codes):
     from freetoken.moe.cpu_executor import CpuMoeExecutor
 
     H, I = 259, 257  # ragged K for both GEMVs; 257 down rows exercise many reductions
@@ -159,17 +156,16 @@ def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa,
         max_tokens=1,
         device=torch.device("cpu"),
     )
-    tiers = {"scalar": 0, "avx2": 1, "avx512": 2, "avx512bf16": 3}
+    tiers = {"scalar": 0, "avx2": 1, "avx512": 2}
     if tiers[isa] > ("scalar", "avx2", "avx512f", "avx512bf16").index(ex.isa):
         pytest.skip(f"requested {isa}, host/build selected {ex.isa}")
-    assert ex._ext.fp8_block_isa_name() == expected_isa
 
     hidden = torch.randn(1, H, dtype=torch.bfloat16)
-    io = ex._io_for(1)
+    io = ex._io
     io["x"].copy_(hidden)
     io["ids"].zero_()
     io["w"].fill_(1.0)
-    ex._ext.run_task(ex._task_for(0, 1))
+    ex._ext.run(0, 1)
     cpu_out = io["y"][0].float().to(torch.float64)
 
     banks = cache.bank_sources
@@ -233,9 +229,8 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch):
         max_tokens=bs,
         device=dev,
     )
-    got = ex._ext.fp8_block_isa_name()
-    if got != _ISA_TIERS[isa]:
-        pytest.skip(f"asked for {isa}, kernel selected {got}: not supported here")
+    if ex.isa != _ISA_TIERS[isa]:
+        pytest.skip(f"asked for {isa}, host/build selected {ex.isa}: not supported here")
 
     hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16)
     ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32)
@@ -254,7 +249,7 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch):
     ).float()
 
     rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)
-    assert rel < 2e-2, f"bs={bs} isa={got} rel err {rel.item()}"
+    assert rel < 2e-2, f"bs={bs} isa={ex.isa} rel err {rel.item()}"
 
 
 def test_fp16_bits_carry_every_e4m3_byte():
diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py
index b5ec2f09b4..7babb5d4ce 100644
--- a/tests/moe/test_cpu_moe_q4_0.py
+++ b/tests/moe/test_cpu_moe_q4_0.py
@@ -1,7 +1,7 @@
 """CPU MoE executor -- native GGUF Q4_0 experts (``--moe-backend cpu``/``hybrid`` for
 the gemma4 GGUF checkpoint).
 
-The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/cpu_moe_ext.cpp) reads the *same*
+The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same*
 packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the
 K-loop. We check it against the reference dequant (models/gguf/dequant.py) + the
 production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the
diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py
index 11519502e2..8187402d3f 100644
--- a/tests/moe/test_dsfp4_prequant.py
+++ b/tests/moe/test_dsfp4_prequant.py
@@ -17,8 +17,9 @@
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-def test_dsfp4_gpu_prequant_bit_parity():
+def test_dsfp4_gpu_prequant_bit_parity(monkeypatch):
     from freetoken.kernel.pinned import alloc_pinned_tensor
+    from freetoken.moe import cpu_executor
     from freetoken.moe.cpu_executor import CpuMoeExecutor
 
     torch.manual_seed(0)
@@ -44,14 +45,13 @@ def test_dsfp4_gpu_prequant_bit_parity():
     wts = torch.rand(2, TOPK, device=dev)
 
     def run(cpu_roundtrip: bool) -> torch.Tensor:
+        # control: the original CPU-side scalar round-trip
+        monkeypatch.setattr(cpu_executor, "_DSFP4_GPU_PREQUANT", not cpu_roundtrip)
         ex = CpuMoeExecutor(
             cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False,
             num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128,
         )
-        assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant"
-        if cpu_roundtrip:  # control: the original CPU-side scalar round-trip
-            ex._gpu_prequant = False
-            ex._ext.set_input_prequant(False)
+        assert ex._gpu_prequant != cpu_roundtrip, "ds_fp4 on CUDA should enable GPU prequant"
         out = ex.decode(0, x, wts, ids.clone()).clone()
         torch.cuda.synchronize()
         del ex

From f02fbb16a213f5c81094e4017c988d41f0164141 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 13 Oct 2026 20:30:34 -0700
Subject: [PATCH 095/249] perf(moe): split hybrid decode routes inside the
 ensure kernel (#685)

* perf(moe): split hybrid decode routes inside the ensure kernel

* refactor(moe): always split routes in the hybrid ensure kernel

* test(moe): drop the fetch split test left without its kernel check
---
 python/freetoken/layers/moe.py          |  18 +---
 python/freetoken/moe/offload_cache.py   |  56 ++++------
 python/freetoken/moe/offload_kernels.py | 137 ++++++++++++------------
 tests/moe/test_hybrid_fetch.py          |  75 +------------
 4 files changed, 95 insertions(+), 191 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 82f09a9bb8..0595b1d451 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -307,7 +307,7 @@ def _decode_hybrid(
 
         The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so
         the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the
-        routing split is device-side elementwise and the CPU submit/sync are host nodes.
+        routing split runs inside the hybrid ensure kernel and the CPU submit/sync are host nodes.
         Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with
         weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads
         slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for
@@ -315,14 +315,10 @@ def _decode_hybrid(
         """
         executor = cache.cpu_executor
         assert executor is not None, "CPU MoE executor was not initialized"
-        raw_ids = topk_ids.clone()
-        cache.ensure_experts_hybrid(self.layer_id, topk_ids)
-        on_gpu = topk_ids >= 0
-        cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous()
-        gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous()
-        gpu_slots = topk_ids
-        if cache.collect_stats:
-            cache.record_decode_stats_hybrid(self.layer_id)
+        # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
+        # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
+        clamp = not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots)
+        gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp)
         pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids)
 
         # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the
@@ -332,10 +328,6 @@ def _decode_hybrid(
         )
 
         cache.copy_missing()
-        if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots):
-            # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
-            # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
-            gpu_slots = gpu_slots.clamp_min(0)
         gpu_routed = self._expert_gemm(
             cache,
             hidden_states,
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 0969f73829..984b7124ad 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -3,7 +3,7 @@
 import math
 import os
 from dataclasses import dataclass
-from typing import Iterator
+from typing import TYPE_CHECKING, Iterator
 
 import torch
 from flashlib.kernels.slot_cache import N_STATS, Stat
@@ -28,6 +28,9 @@
 from freetoken.kernel.fast_index_copy import default_blocks_per_bank
 from freetoken.utils import init_logger
 
+if TYPE_CHECKING:
+    from freetoken.moe.offload_kernels import HybridRoutes
+
 logger = init_logger(__name__)
 
 # quant_format -> bank names, in registration order: the single place a format's bank
@@ -224,24 +227,23 @@ def __post_init__(self) -> None:
         # is a one-off warm-up increment at capture time (<0.1% over a session).
         self.collect_stats = False
         # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to
-        # the kernel, which accumulates in the same launch. The stat_* tensors below stay
-        # for the hybrid path, whose kernel is still ours.
+        # the kernel, which accumulates in the same launch. The stat_* tensors below are
+        # the hybrid path's: views of one buffer its kernel accumulates into the same way.
         self.lru_stats = torch.zeros(
             (self.num_layers, N_STATS), dtype=torch.int64, device=self.device
         )
-        self.stat_missing = torch.zeros((), dtype=torch.int64, device=self.device)
-        self.stat_active = torch.zeros((), dtype=torch.int64, device=self.device)
-        self.stat_calls = torch.zeros((), dtype=torch.int64, device=self.device)
-        # hybrid only: experts actually fetched over PCIe (<= stat_missing). The CPU
-        # computes stat_missing - stat_fetched of them.
-        self.stat_fetched = torch.zeros((), dtype=torch.int64, device=self.device)
+        # (missing, fetched, active, calls), then the same per layer
+        stat_buf = torch.zeros(4 * (self.num_layers + 1), dtype=torch.int64, device=self.device)
+        self.hybrid_stats = (stat_buf[:4], stat_buf[4:])
+        # hybrid only: stat_fetched counts the experts actually fetched over PCIe (<= stat_missing).
+        # The CPU computes stat_missing - stat_fetched of them.
+        self.stat_missing, self.stat_fetched, self.stat_active, self.stat_calls = stat_buf[:4].unbind()
         # Per-layer counterparts of the scalars above (indexed by MoE-layer id). Same
         # device-side accumulation (graph-safe: layer_id is a static index per graph node),
         # so one req's per-layer miss rate is readable via decode_miss_stats_per_layer().
-        self.stat_missing_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_active_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_fetched_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_steps_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
+        per_layer = stat_buf[4:].view(self.num_layers, 4)
+        self.stat_missing_layer, self.stat_fetched_layer = per_layer[:, 0], per_layer[:, 1]
+        self.stat_active_layer, self.stat_steps_layer = per_layer[:, 2], per_layer[:, 3]
         # Opt-in decode routing histogram (per layer, per expert) for cache-skew
         # analysis. Accumulated in ``ensure_experts`` from the raw expert ids before the
         # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the
@@ -855,7 +857,9 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         self._pending_whole_layer = False
         ensure_experts(self, layer_id, expert_ids)
 
-    def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
+    def ensure_experts_hybrid(
+        self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False
+    ) -> HybridRoutes:
         """Capped-fetch LRU for the hybrid backend.
 
         Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at
@@ -864,7 +868,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None
         left non-resident and ``expert_ids`` is rewritten to their cache slot (hit or
         freshly fetched) or ``-1`` (overflow -> compute on the CPU). ``num_indices`` holds
         the capped fetch count (for ``copy_missing``); ``num_missing_full`` the pre-cap
-        miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe."""
+        miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe.
+        Returns the routes split between GPU and CPU (:class:`freetoken.moe.offload_kernels.HybridRoutes`)."""
         from freetoken.moe.offload_kernels import ensure_experts_hybrid
 
         if self.collect_decode_freq:
@@ -872,8 +877,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        ensure_experts_hybrid(
-            self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction
+        return ensure_experts_hybrid(
+            self, layer_id, expert_ids, weights, self.hybrid_max_fetch, self.hybrid_fetch_fraction, clamp
         )
 
     def materialize_layer(self, layer_id: int) -> None:
@@ -911,23 +916,6 @@ def record_decode_stats(self, layer_id: int) -> None:
         was eight torch ops per layer per step, all captured into the decode graph.
         """
 
-    def record_decode_stats_hybrid(self, layer_id: int) -> None:
-        """Hybrid stats: full miss count (pre-cap), the PCIe-fetched count (capped), and
-        the active count. The CPU computes (missing - fetched) experts. Device-side;
-        accumulates both the scalar totals and the per-layer breakdown."""
-        assert 0 <= layer_id < self.num_layers, f"layer_id {layer_id} out of range [0, {self.num_layers})"
-        missing = self.num_missing_full.sum()
-        fetched = self.num_indices.sum()
-        active = self.active_mask.sum()
-        self.stat_missing += missing
-        self.stat_fetched += fetched
-        self.stat_active += active
-        self.stat_calls += 1
-        self.stat_missing_layer[layer_id] += missing
-        self.stat_fetched_layer[layer_id] += fetched
-        self.stat_active_layer[layer_id] += active
-        self.stat_steps_layer[layer_id] += 1
-
     def decode_miss_stats(self) -> dict:
         if self.decode_target == "hybrid":
             active = int(self.stat_active.item())
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index cf513f52d1..ea12b6fb17 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from typing import NamedTuple
 
 import torch
 import triton
@@ -40,9 +41,23 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
     )
 
 
+class HybridRoutes(NamedTuple):
+    """One layer's routes split between the GPU and the CPU, each [bs, top_k]."""
+
+    gpu_slots: torch.Tensor  # cache slot of each route; a CPU route is -1, or 0 with clamp
+    gpu_w: torch.Tensor  # router weights, 0 for a CPU route
+    cpu_ids: torch.Tensor  # expert ids, -1 for a GPU route
+
+
 def ensure_experts_hybrid(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0
-) -> None:
+    cache,
+    layer_id: int,
+    expert_ids: torch.Tensor,
+    weights: torch.Tensor,
+    max_fetch: int,
+    fetch_fraction: float = 0.0,
+    clamp: bool = False,
+) -> HybridRoutes:
     """Capped-fetch variant of ``ensure_experts`` (hybrid backend).
 
     Identical LRU bookkeeping, but only the first ``max_fetch`` of this step's missing
@@ -52,12 +67,17 @@ def ensure_experts_hybrid(
     split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded to
     the integer that makes the PCIe fetch and the CPU overflow compute finish closest to
     together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` =
-    pre-cap miss count (stats)."""
-    # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float).
+    pre-cap miss count (stats).
+
+    The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and
+    ``clamp`` gives CPU routes slot 0 instead of -1 for GEMM kernels that read every slot."""
+    assert expert_ids.is_cuda, "the hybrid ensure kernel runs on the GPU"
+    # Q16 fixed point: the kernel caps in integers.
     frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16))))
-    if not expert_ids.is_cuda:
-        return _ensure_experts_hybrid_cpu(cache, layer_id, expert_ids, max_fetch, frac_q16)
-    _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16)
+    weights = weights.contiguous()
+    routes = HybridRoutes(expert_ids, torch.empty_like(weights), torch.empty_like(expert_ids))
+    _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, weights, routes, clamp, max_fetch, frac_q16)
+    return routes
 
 
 def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None:
@@ -95,7 +115,14 @@ def reset_cache(cache) -> None:
 
 
 def _ensure_experts_hybrid_gpu(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int
+    cache,
+    layer_id: int,
+    expert_ids: torch.Tensor,
+    weights: torch.Tensor,
+    routes: HybridRoutes,
+    clamp: bool,
+    max_fetch: int,
+    frac_q16: int,
 ) -> None:
     block_e = triton.next_power_of_2(cache.num_experts)
     block_c = triton.next_power_of_2(cache.cache_size)
@@ -112,6 +139,10 @@ def _ensure_experts_hybrid_gpu(
         cache.num_indices,
         cache.num_missing_full,
         cache.expert_recency,
+        weights,
+        routes.cpu_ids,
+        routes.gpu_w,
+        *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2),
         layer_id,
         expert_ids.numel(),
         int(max_fetch),
@@ -121,73 +152,12 @@ def _ensure_experts_hybrid_gpu(
         BLOCK_E=block_e,
         BLOCK_C=block_c,
         BY_RECENCY=_HYBRID_FETCH_BY_RECENCY,
+        CLAMP=clamp,
+        STATS=cache.collect_stats,
         num_warps=num_warps,
     )
 
 
-def _ensure_experts_hybrid_cpu(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int
-) -> None:
-    """CPU reference mirror of the hybrid kernel (eviction/fetch decisions bit-identical to
-    the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or
-    the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing
-    experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the
-    most-recently-active misses (ties -> lower id); else the lowest ids."""
-    seen = []
-    for expert in expert_ids.view(-1).tolist():
-        if expert not in seen:
-            seen.append(expert)
-
-    cache.active_mask.zero_()
-    step = int(cache.step.item()) + 1
-    cache.step.fill_(step)
-    for expert in seen:
-        cache.active_mask[expert] = 1
-
-    for expert in seen:
-        slot = int(cache.slot_for_id[layer_id, expert].item())
-        if slot != -1:
-            cache.usage[slot] = step
-
-    missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1]
-    if _HYBRID_FETCH_BY_RECENCY:
-        rec = cache.expert_recency[layer_id].tolist()
-        missing.sort(key=lambda e: (-rec[e], e))
-    else:
-        missing.sort()
-    if frac_q16 > 0:
-        m, q = len(missing), 1 << 16
-        lo = (m * frac_q16) >> 16
-        cost = lambda f: max(f * (q - frac_q16), (m - f) * frac_q16)  # noqa: E731
-        max_fetch = lo if cost(lo) <= cost(lo + 1) else lo + 1
-    num_fetch = min(len(missing), int(max_fetch))
-    cache.num_missing_full.fill_(len(missing))
-    cache.num_indices.fill_(num_fetch)
-
-    usage = cache.usage.tolist()
-    for idx in range(num_fetch):
-        expert = missing[idx]
-        victim = min(range(cache.cache_size), key=lambda s: (usage[s], s))
-        old_id = int(cache.id_of_slot[victim].item())
-        if old_id >= 0:
-            cache.slot_for_id.view(-1)[old_id] = -1
-        cache.id_of_slot[victim] = layer_id * cache.num_experts + expert
-        cache.slot_for_id[layer_id, expert] = victim
-        cache.usage[victim] = step
-        usage[victim] = step
-        cache.evict_slots[idx] = victim
-        cache.src_indices[idx] = expert  # layer-local row
-
-    if _HYBRID_FETCH_BY_RECENCY:
-        for expert in seen:
-            cache.expert_recency[layer_id, expert] = step
-
-    # Overflow misses keep slot_for_id == -1, so the rewrite below yields -1 for them.
-    flat = expert_ids.view(-1)
-    for i in range(flat.numel()):
-        flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item())
-
-
 def _materialize_layer_gpu(cache, layer_id: int) -> None:
     block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size))
     _materialize_layer_kernel[(1,)](
@@ -300,6 +270,11 @@ def _ensure_experts_hybrid_kernel(
     num_indices_ptr,
     num_missing_full_ptr,
     expert_recency_ptr,
+    weights_ptr,
+    cpu_ids_ptr,
+    gpu_w_ptr,
+    stats_ptr,
+    stats_layer_ptr,
     layer_id,
     num_active,
     max_fetch,
@@ -309,6 +284,8 @@ def _ensure_experts_hybrid_kernel(
     BLOCK_E: tl.constexpr,
     BLOCK_C: tl.constexpr,
     BY_RECENCY: tl.constexpr,
+    CLAMP: tl.constexpr,
+    STATS: tl.constexpr,
 ):
     """Capped-fetch timestamp-LRU (hybrid backend).
 
@@ -401,8 +378,26 @@ def _ensure_experts_hybrid_kernel(
     for i in tl.range(num_active):
         e = tl.load(expert_ids_ptr + i)
         s = tl.load(slot_for_id_ptr + base + e)
+        on_gpu = s >= 0
+        tl.store(cpu_ids_ptr + i, tl.where(on_gpu, -1, e))
+        tl.store(gpu_w_ptr + i, tl.where(on_gpu, tl.load(weights_ptr + i), 0.0))
+        if CLAMP:
+            s = tl.maximum(s, 0)
         tl.store(expert_ids_ptr + i, s)
 
+    # stats_ptr: (missing, fetched, active, calls); stats_layer_ptr: the same per layer, [num_layers, 4]
+    if STATS:
+        n_active = tl.sum((is_active & e_mask).to(tl.int64))
+        tl.store(stats_ptr + 0, tl.load(stats_ptr + 0) + num_missing)
+        tl.store(stats_ptr + 1, tl.load(stats_ptr + 1) + num_fetch)
+        tl.store(stats_ptr + 2, tl.load(stats_ptr + 2) + n_active)
+        tl.store(stats_ptr + 3, tl.load(stats_ptr + 3) + 1)
+        row = stats_layer_ptr + layer_id * 4
+        tl.store(row + 0, tl.load(row + 0) + num_missing)
+        tl.store(row + 1, tl.load(row + 1) + num_fetch)
+        tl.store(row + 2, tl.load(row + 2) + n_active)
+        tl.store(row + 3, tl.load(row + 3) + 1)
+
     # Bump every active expert's recency to this step (LRU on the expert): an overflow miss
     # computed on the CPU now ranks high if it recurs, so it gets fetched next time.
     if BY_RECENCY:
diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py
index a26630b894..fad4fa7843 100644
--- a/tests/moe/test_hybrid_fetch.py
+++ b/tests/moe/test_hybrid_fetch.py
@@ -1,43 +1,15 @@
 """Hybrid decode's bandwidth-matched fetch split.
 
-Covers the two halves of --moe-hybrid-max-fetch auto: the profile reader that turns
-`ft bench bw` kernel bandwidths into a fetch fraction, and the ensure kernel's
-per-step integer split (GPU kernel vs CPU reference mirror, and the balance rule).
+Covers the profile reader of --moe-hybrid-max-fetch auto, which turns `ft bench bw`
+kernel bandwidths into a fetch fraction.
 """
 
 import json
 import os
 
 import pytest
-import torch
 
 from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction
-from freetoken.moe.offload_cache import OffloadMoeCache
-
-Q = 1 << 16
-
-
-def _balanced_fetch(num_missing: int, frac_q16: int) -> int:
-    """Reference split: F ~ frac * misses, rounded to whichever integer neighbor
-    minimizes the slower overlapped side (fetch ~ F*(1-frac), CPU ~ (M-F)*frac)."""
-    lo = (num_missing * frac_q16) >> 16
-    cost = lambda f: max(f * (Q - frac_q16), (num_missing - f) * frac_q16)  # noqa: E731
-    return min(num_missing, lo if cost(lo) <= cost(lo + 1) else lo + 1)
-
-
-def test_balanced_fetch_tracks_fraction():
-    # The split follows fetched : cpu = pcie : (cpu - pcie) up to integer rounding, and
-    # never over/under-shoots by more than one expert.
-    for frac in (0.1, 0.415, 0.454, 0.7, 1.0):
-        q = round(frac * Q)
-        for m in range(0, 65):
-            f = _balanced_fetch(m, q)
-            assert 0 <= f <= m
-            assert abs(f - frac * m) <= 1.0
-    # ceil would over-fetch here (the regression this rule fixed): 41.5% of 3 misses is
-    # 1.24 -> fetching 2 makes the PCIe side ~1.6x slower than balance; keep it at 1.
-    assert _balanced_fetch(3, round(0.415 * Q)) == 1
-    assert _balanced_fetch(4, round(0.415 * Q)) == 2
 
 
 def test_load_hybrid_fetch_fraction(tmp_path):
@@ -87,46 +59,3 @@ def write(path, name, verdict):
     # this card's own file wins over the legacy one
     write(default_profile_path(uuid), "FAKE GPU", "offload")
     assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload"
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-def test_hybrid_fraction_gpu_matches_cpu_reference():
-    torch.manual_seed(0)
-    num_experts, cache_size, top_k, frac = 32, 40, 8, 0.415
-
-    def make():
-        return OffloadMoeCache(
-            num_layers=2, num_experts=num_experts, cache_size=cache_size,
-            device=torch.device("cuda"), quant_format="bf16", decode_target="hybrid",
-            hybrid_max_fetch=num_experts, hybrid_fetch_fraction=frac,
-        )
-
-    gpu, ref = make(), make()
-    frac_q16 = round(frac * Q)
-    for step in range(64):
-        ids = torch.randperm(num_experts)[:top_k].to(torch.int32)
-        g, c = ids.clone().cuda(), ids.clone()  # a CPU ids tensor drives the reference path
-        gpu.ensure_experts_hybrid(0, g)
-        ref.ensure_experts_hybrid(0, c)
-        missing = int(gpu.num_missing_full.item())
-        fetched = int(gpu.num_indices.item())
-        assert missing == int(ref.num_missing_full.item())
-        assert fetched == int(ref.num_indices.item()) == _balanced_fetch(missing, frac_q16)
-        # slot rewrites (hit/fetched -> slot, overflow -> -1) and LRU state stay identical
-        assert torch.equal(g.cpu(), c)
-        assert torch.equal(gpu.slot_for_id.cpu(), ref.slot_for_id.cpu())
-        assert torch.equal(gpu.id_of_slot.cpu(), ref.id_of_slot.cpu())
-        assert (g >= 0).sum().item() == len(set(ids.tolist())) - (missing - fetched)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-def test_hybrid_fixed_cap_unchanged():
-    # fraction 0 (no profile / explicit --moe-hybrid-max-fetch) keeps the fixed cap.
-    cache = OffloadMoeCache(
-        num_layers=1, num_experts=32, cache_size=40, device=torch.device("cuda"),
-        quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1,
-    )
-    ids = torch.arange(8, dtype=torch.int32).cuda()
-    cache.ensure_experts_hybrid(0, ids)
-    assert int(cache.num_missing_full.item()) == 8
-    assert int(cache.num_indices.item()) == 1

From b416f85c40970c033462b47773020872209978a3 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 14 Oct 2026 19:44:55 -0700
Subject: [PATCH 096/249] feat(moe): split the offload expert cache into slot
 pools (#689)

* feat(moe): split the offload expert cache into slot pools

* perf(moe): stop recompiling the hybrid ensure kernel for each slot pool size

* fix(engine): turn off prefill overlap when --moe-cache-mem leaves under two layers of slots

* fix(engine): load the --moe-cache-split profile before the expert banks
---
 docs/cli.md                                |   4 +-
 python/freetoken/engine/cache_budget.py    |  13 +-
 python/freetoken/engine/config.py          |   6 +
 python/freetoken/engine/engine.py          | 134 +++++--
 python/freetoken/kvcache/cache_status.py   |   3 +
 python/freetoken/layers/moe.py             |  17 +-
 python/freetoken/models/deepseek_v4/moe.py |   5 +-
 python/freetoken/moe/benchbw.py            |   9 +-
 python/freetoken/moe/expert_banks.py       |   2 +
 python/freetoken/moe/offload_cache.py      | 412 +++++++++++++++------
 python/freetoken/moe/offload_kernels.py    |  55 +--
 python/freetoken/moe/pool_policy.py        | 149 ++++++++
 python/freetoken/server/args.py            |  50 +++
 tests/engine/test_cache_budget.py          |  52 +++
 tests/moe/test_nvfp4_backends.py           |  12 +-
 tests/moe/test_offload.py                  | 186 +++++++++-
 16 files changed, 921 insertions(+), 188 deletions(-)
 create mode 100644 python/freetoken/moe/pool_policy.py

diff --git a/docs/cli.md b/docs/cli.md
index 2166e9e86e..44b0e30ad9 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -84,7 +84,9 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling |
 | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name |
 | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=<marlin\|b12x\|triton>` (`flashinfer` means b12x); cannot be combined with `--quant-backend` |
-| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) |
+| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-mem` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / memory such as `12GiB` / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) |
+| `--moe-cache-scope` | unified | `unified`: one LRU pool of slots per expert quant (quant type and bank layout), so all MoE layers share one pool when they share a quant format; `per-layer`: each layer has its own pool. `per-layer` does not combine with `--moe-strategy cpu` or `--moe-cpu-layers` |
+| `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:<path>`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k |
 | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts |
 | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor |
 | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) |
diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py
index b5f07ee8ab..891ac2ac11 100644
--- a/python/freetoken/engine/cache_budget.py
+++ b/python/freetoken/engine/cache_budget.py
@@ -14,8 +14,8 @@
     import torch
 
 
-def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int:
-    """Bytes one expert slot occupies on GPU: summed row bytes over all banks.
+def layer_slot_bytes(sources: dict[str, "list[torch.Tensor]"]) -> list[int]:
+    """Bytes one expert slot of each layer occupies on GPU: its row bytes summed over all banks.
 
     Each bank source is per-layer ``[num_experts, *row_shape]`` tensors and is
     already TP-sharded upstream, so the per-row byte count is the per-rank slot
@@ -25,7 +25,14 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int:
     # with cache_size), so they are intentionally excluded from the per-slot growth term.
     # tensor[0].numel() is the per-row element count (one expert slot); see the matching
     # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py.
-    return sum(t[0][0].numel() * t[0].element_size() for t in sources.values())
+    num_layers = len(next(iter(sources.values())))
+    return [sum(t[l][0].numel() * t[l].element_size() for t in sources.values()) for l in range(num_layers)]
+
+
+def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int:
+    """Bytes one expert slot occupies on GPU, rounded up to the mean over layers when their bank layouts differ."""
+    per_layer = layer_slot_bytes(sources)
+    return -(-sum(per_layer) // len(per_layer))
 
 
 def net_cache_budget_bytes(
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index b6fa8b8949..ae2f1fd629 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -45,6 +45,12 @@ class EngineConfig:
     moe_cache_size: int = 0
     moe_cache_rate: float | None = None
     moe_cache_auto: bool = False
+    # --moe-cache-mem: slot-cache budget in bytes, turned into a slot count once the expert banks are loaded
+    moe_cache_mem: int | None = None
+    # --moe-cache-scope: unified (one LRU pool per expert quant) or per-layer (one pool per layer)
+    moe_cache_scope: str = "unified"
+    # --moe-cache-split: how pools share the slots; equal or profile:<path> (by expert use from a profile file)
+    moe_cache_split: str = "equal"
     kv_reserve_tokens: int = 8192  # KV floor for --moe-cache-auto; small by design (MoE-priority)
     moe_cache_policy: str = "lru"
     moe_prefill_overlap: bool = True
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index feb33fbb94..2757b10dea 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -23,7 +23,8 @@
 from freetoken.moe import is_offload_moe_strategy
 from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks
 from freetoken.moe.host_banks import PinFailed
-from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache
+from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache, bank_layouts
+from freetoken.moe.pool_policy import group_layers, make_pools, pool_of_layer, resolve_profile, split_slots
 from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype
 
 from .config import EngineConfig
@@ -38,6 +39,25 @@
 
 logger = init_logger(__name__)
 
+def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None:
+    """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool."""
+    num_experts = config.model_config.num_experts
+    small = [s for s in pool_sizes if s < num_experts]
+    if not small:
+        return
+    limit = min(small) // config.model_config.num_experts_per_tok
+    if config.max_running_req > limit:
+        logger.warning_rank0(
+            f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req "
+            f"{config.max_running_req} -> {limit} so a decode step fits in it"
+        )
+        object.__setattr__(config, "max_running_req", limit)
+    if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > limit:
+        object.__setattr__(config, "cuda_graph_max_bs", limit)
+    if config.cuda_graph_bs is not None:
+        object.__setattr__(config, "cuda_graph_bs", [bs for bs in config.cuda_graph_bs if bs <= limit])
+
+
 def _require_offload_cache_size(cache_size: int, num_experts: int) -> None:
     """The offload MoE cache needs at least one slot per expert per layer. A too-small size
     (e.g. a bare offload run with moe_cache_size unset and auto disabled) must fail loudly."""
@@ -623,7 +643,7 @@ def _run_mm_encoder(self, batch: Batch) -> None:
         if parts:
             batch.mm_embeds = torch.cat([p.to(self.dtype) for p in parts], dim=0)
 
-    def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) -> tuple[int, int, bool]:
+    def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None, max_slots=None) -> tuple[int, int, bool]:
         """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap).
 
         Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable
@@ -647,7 +667,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None)
             prefill_overlap=config.moe_prefill_overlap,
             kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve),
             page_size=page_tokens,
-            max_slots=method.slot_limit() if method is not None else None,
+            max_slots=max_slots if max_slots is not None or method is None else method.slot_limit(),
         )
 
     def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
@@ -703,6 +723,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 else HostResidency.PINNED.value
                 for i in range(config.model_config.num_moe_layers)
             ]
+        # read before the expert load, so a bad profile fails in seconds
+        expert_counts = resolve_profile(
+            config.moe_cache_split, config.model_path,
+            config.model_config.num_moe_layers, config.model_config.num_experts,
+        )
         try:
             with _weight_load_context():
                 banks = load_expert_banks(
@@ -718,8 +743,27 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 )
         except PinFailed as exc:
             raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc
+        from freetoken.engine.cache_budget import expert_bytes_per_slot, layer_slot_bytes
+
+        self._moe_layer_slot_bytes = layer_slot_bytes(banks.sources)
+        self._moe_pool_groups = group_layers(config.moe_cache_scope, bank_layouts(banks.sources, banks.layer_tags), self._moe_layer_slot_bytes)
+        if len(self._moe_pool_groups) > 1 and cpu_layer_ids:
+            raise ValueError(
+                "several MoE slot pools (--moe-cache-scope per-layer, or experts in several quant formats) need every MoE "
+                "layer to decode from the GPU slot cache; they do not combine with --moe-strategy cpu or --moe-cpu-layers"
+            )
+        self._moe_expert_counts = expert_counts if len(self._moe_pool_groups) > 1 else None
+        if len(self._moe_pool_groups) == 1 and config.moe_cache_split != "equal":
+            logger.warning_rank0(
+                f"--moe-cache-split {config.moe_cache_split} is ignored: all MoE layers share one slot pool "
+                f"(--moe-cache-scope per-layer gives each layer its own)"
+            )
+        self._moe_slot_limit = method.slot_limit() if method is not None else None
+        # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step
+        pools = self._moe_pools(1)
+        size_cap = sum(p.max_num_slots for p in pools)
         if config.moe_cache_auto:
-            size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method)
+            size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap)
             object.__setattr__(config, "moe_cache_size", size)
             object.__setattr__(config, "moe_prefill_overlap", overlap)
             if config.num_page_override is None:
@@ -735,7 +779,28 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 f"--moe-cache-auto resolved moe_cache_size={size} "
                 f"num_pages={pages} (prefill_overlap={overlap})"
             )
+        elif config.moe_cache_mem is not None:
+            per_slot = expert_bytes_per_slot(banks.sources)
+            size = min(config.moe_cache_mem // per_slot, size_cap)
+            # per_slot is a mean over layers, so with pools of different slot sizes the split can exceed the budget: shrink until it fits
+            while size > 0 and self._moe_pool_bytes(split_slots(size, pools, self._moe_expert_counts)) > config.moe_cache_mem:
+                size -= 1
+            overlap = config.moe_prefill_overlap and size >= 2 * config.model_config.num_experts
+            object.__setattr__(config, "moe_cache_size", size)
+            object.__setattr__(config, "moe_prefill_overlap", overlap)
+            logger.info_rank0(
+                f"--moe-cache-mem {mem_GB(config.moe_cache_mem)} resolved moe_cache_size={size} "
+                f"({per_slot / 2**20:.2f} MiB per slot, prefill_overlap={overlap})"
+            )
         _require_offload_cache_size(config.moe_cache_size, config.model_config.num_experts)
+        pool_sizes = split_slots(config.moe_cache_size, pools, self._moe_expert_counts)
+        if len(pools) > 1:
+            logger.info_rank0(
+                f"MoE slot pools: {config.moe_cache_scope}, {len(pools)} pools of {min(pool_sizes)}-{max(pool_sizes)} "
+                f"slots ({config.moe_cache_split})"
+            )
+            if decode_target == "gpu":
+                _clamp_decode_batch_to_pools(config, pool_sizes)
         layout = max_slots = None
         if method is not None:
             if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name):
@@ -760,10 +825,12 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             hybrid_max_fetch=config.moe_hybrid_max_fetch,
             layout=layout,
             max_slots=max_slots,
+            pool_of_layer=pool_of_layer(pools, config.model_config.num_moe_layers),
+            pool_sizes=pool_sizes,
         )
         # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set
         cache.cpu_layer_ids = cpu_layer_ids
-        cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency)
+        cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags)
         cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
         if decode_target == "hybrid":
             self._resolve_hybrid_fetch(config, cache)
@@ -772,12 +839,23 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         cache.collect_stats = config.moe_collect_stats
         layers = attach_offload_moe_cache(self.model, cache)
         assert len(layers) == config.model_config.num_moe_layers
+        cache.zero_inactive_route_slots = cache.decode_target == "hybrid" and not all(layer.skips_inactive_slots for layer in layers)
         if cache.decode_target in ("cpu", "hybrid"):
             self._init_cpu_moe_executor(config, cache, layers)
         self.ctx.moe_offload_cache = cache
         self.moe_offload_cache = cache
         return cache
 
+    def _moe_pool_bytes(self, pool_sizes: list[int]) -> int:
+        return sum(s * self._moe_layer_slot_bytes[g[0]] for g, s in zip(self._moe_pool_groups, pool_sizes))
+
+    def _moe_pools(self, batch: int) -> list:
+        """The slot pools, each with room for the routed experts of ``batch`` decode tokens."""
+        model_config = self.config.model_config
+        num_experts = model_config.num_experts
+        min_num_slots = min(num_experts, model_config.num_experts_per_tok * batch)
+        return make_pools(self._moe_pool_groups, num_experts, min_num_slots, self._moe_slot_limit)
+
     def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None:
         """Resolve --moe-hybrid-max-fetch -1 (auto) into a bandwidth-matched fetch fraction.
 
@@ -867,19 +945,18 @@ def _sync_get_memory(self) -> Tuple[int, int]:
 
         return min_free_memory, max_free_memory
 
-    def _target_moe_and_expert_bytes(self, moe_cache_size: int | None) -> tuple[int, int]:
-        from freetoken.engine.cache_budget import expert_bytes_per_slot
-
-        target_moe = (
-            moe_cache_size
-            if moe_cache_size is not None
-            else (self.moe_offload_cache.cache_size if self.moe_offload_cache else 0)
-        )
-        per_expert_bytes = (
-            expert_bytes_per_slot(self.moe_offload_cache.bank_sources)
-            if self.moe_offload_cache is not None else 0
-        )
-        return target_moe, per_expert_bytes
+    def _target_moe_and_expert_bytes(self, moe_cache_size: int | None, pool_sizes: list[int] | None) -> tuple[int, int]:
+        cache = self.moe_offload_cache
+        if cache is None:
+            return moe_cache_size or 0, 0
+        if moe_cache_size is None:
+            target_moe = cache.cache_size
+            nbytes = sum(c.numel() * c.element_size() for _, c in cache.banks)
+        else:
+            target_moe = moe_cache_size
+            nbytes = self._moe_pool_bytes(pool_sizes)
+        # pools of several bank layouts have different slot sizes; round up so target * bytes covers them
+        return target_moe, -(-nbytes // target_moe)
 
     def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> None:
         # IN-PLACE, identity-preserving: the CacheManager's swa_pool reference, ctx.kv_cache and
@@ -938,7 +1015,10 @@ def rebuild_runtime_cache(
                     "moe_cache_size requested but this model has no MoE offload cache"
                 )
             try:
-                self.moe_offload_cache.validate_rebuild(moe_cache_size)
+                # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead
+                batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1
+                moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts)
+                self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes)
             except ValueError as e:
                 raise CacheRebuildRejected(str(e)) from e
         if num_pages is not None and num_pages <= 0:
@@ -975,7 +1055,9 @@ def rebuild_runtime_cache(
         #     must reject (recoverable) so the old caches stay intact and serving continues,
         #     rather than freeing and then OOMing into permanent failure. The engine supplies
         #     the memory account; the pool answers whether its target geometry fits.
-        target_moe, per_expert_bytes = self._target_moe_and_expert_bytes(moe_cache_size)
+        target_moe, per_expert_bytes = self._target_moe_and_expert_bytes(
+            moe_cache_size, moe_pool_sizes if moe_cache_size is not None else None
+        )
         # Price the sibling GDN state pool at ITS target (physical slots = usable + padding
         # sink) and hand the bytes in -- the KV pool only budgets its own tiers.
         target_mamba = (
@@ -1019,7 +1101,7 @@ def rebuild_runtime_cache(
             object.__setattr__(config, "swa_num_pages_override", num_swa_pages)
         if moe_cache_size is not None:
             assert self.moe_offload_cache is not None, "no MoE offload cache to resize"
-            self.moe_offload_cache.rebuild(moe_cache_size)
+            self.moe_offload_cache.rebuild(moe_cache_size, moe_pool_sizes)
         if num_pages is not None:
             # sets self.num_pages (rebuilds KV + window)
             self._resize_kv_pool(config, num_pages, num_swa_pages)
@@ -1455,6 +1537,9 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int
     "moe_cache_size": 0,
     "moe_cache_rate": None,
     "moe_cache_auto": False,
+    "moe_cache_mem": None,
+    "moe_cache_scope": "unified",
+    "moe_cache_split": "equal",
     "moe_cpu_layers": None,
     "moe_cpu_threads": 0,
     "moe_hybrid_max_fetch": -1,
@@ -1750,6 +1835,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             is_offload_moe_strategy(config.moe_strategy)
             and config.moe_cache_size <= 0
             and config.moe_cache_rate is None
+            and config.moe_cache_mem is None
             and not getattr(config, "moe_cache_auto", False)
         ):
             # args.py's "no sizing flag -> default --moe-cache-auto" only fires when the
@@ -1783,11 +1869,13 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         # An explicit 'fused' keeps the experts resident, so there is no slot cache to size. The
         # sizing flags no longer redirect the backend, so ignore them here and say so -- the
         # geometry the user asked for is what runs. Report the flag actually passed: --moe-cache-
-        # rate was already folded into moe_cache_size above, and the three are mutually exclusive.
+        # rate was already folded into moe_cache_size above, and the sizing flags are mutually exclusive.
         if config.moe_cache_rate:
             inert = f"--moe-cache-rate={config.moe_cache_rate}"
         elif config.moe_cache_size:
             inert = f"--moe-cache-size={config.moe_cache_size}"
+        elif config.moe_cache_mem is not None:
+            inert = f"--moe-cache-mem={config.moe_cache_mem}"
         elif getattr(config, "moe_cache_auto", False):
             inert = "--moe-cache-auto"
         else:
@@ -1799,6 +1887,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             )
             override("moe_cache_size", 0)
             override("moe_cache_rate", None)
+            override("moe_cache_mem", None)
             override("moe_cache_auto", False)
 
     if is_moe and config.moe_strategy == "cpu":
@@ -1809,6 +1898,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         num_experts = config.model_config.num_experts
         if getattr(config, "moe_cache_auto", False):
             override("moe_cache_auto", False)
+        override("moe_cache_mem", None)
         override("moe_cache_size", 2 * num_experts)
         override("moe_prefill_overlap", True)
         logger.info_rank0(
diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py
index 5b531ce265..8c29dccab5 100644
--- a/python/freetoken/kvcache/cache_status.py
+++ b/python/freetoken/kvcache/cache_status.py
@@ -40,6 +40,9 @@ def _moe() -> int:
         banks = getattr(engine.moe_offload_cache, "bank_caches", None)
         if not banks:
             return 0
+        if getattr(engine.moe_offload_cache, "_raw_banks", False):
+            # pools of several bank layouts share byte segments: report the mean slot size
+            return int(sum(t.numel() for t in banks.values()) // engine.moe_offload_cache.cache_size)
         # Each bank cache is (cache_size, *row_shape); one slot's bytes = row bytes summed over
         # the format's banks (== cache_budget.expert_bytes_per_slot on the source rows).
         return int(sum(t[0].numel() * t.element_size() for t in banks.values()))
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 0595b1d451..79ad9e3a5f 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -189,6 +189,10 @@ def __init__(
         )
         self.offload_cache: OffloadMoeCache | None = None
 
+    @property
+    def skips_inactive_slots(self) -> bool:
+        return self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -289,7 +293,7 @@ def _decode_routed(
             hidden_states,
             topk_weights,
             topk_ids,
-            views=cache.bank_views(),
+            views=cache.pool_bank_views(self.layer_id),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),
             is_prefill=False,
@@ -317,7 +321,8 @@ def _decode_hybrid(
         assert executor is not None, "CPU MoE executor was not initialized"
         # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
         # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
-        clamp = not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots)
+        # with several quants this holds only because the cache zeroes staged slot 0 rows (zero_inactive_route_slots)
+        clamp = not self.skips_inactive_slots
         gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp)
         pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids)
 
@@ -333,7 +338,7 @@ def _decode_hybrid(
             hidden_states,
             gpu_w,
             gpu_slots,
-            views=cache.bank_views(),
+            views=cache.pool_bank_views(self.layer_id),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),
             is_prefill=False,
@@ -370,16 +375,18 @@ def _prefill_routed(
             return out
         cache.materialize_layer(self.layer_id)
         cache.copy_missing()
-        return self._expert_gemm(
+        out = self._expert_gemm(
             cache,
             hidden_states,
             topk_weights,
             topk_ids,
-            views=cache.bank_views(self.num_experts),
+            views=cache.stage_views(self.layer_id, 0),
             n=self.num_experts,
             alphas=cache.alphas_for_layer(self.layer_id),
             is_prefill=True,
         )
+        cache.release_prefill_layer(self.layer_id)
+        return out
 
     def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]:
         """Double-buffer choreography for this layer's overlap prefill: kick off the
diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py
index 35a316f103..aadcdc0481 100644
--- a/python/freetoken/models/deepseek_v4/moe.py
+++ b/python/freetoken/models/deepseek_v4/moe.py
@@ -120,8 +120,9 @@ def _prefill_routed(
         cache = self.offload_cache
         assert cache is not None
         # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights)
+        # one threshold for all layers (the smallest pool): with overlap prefill, the layers of a chunk must all stream or all take the slot path
         if (
-            hidden_states.shape[0] * self.top_k >= self.num_experts
+            hidden_states.shape[0] * self.top_k >= min(self.num_experts, *cache.pool_sizes)
             or cache.is_unpinned_layer(self.layer_id)
         ):
             return super()._prefill_routed(hidden_states, topk_weights, topk_ids)
@@ -134,7 +135,7 @@ def _prefill_routed(
             hidden_states,
             topk_weights,
             topk_ids,
-            views=cache.bank_views(),
+            views=cache.pool_bank_views(self.layer_id),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),
             is_prefill=True,
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index 9430ab2a84..cd260cb2f6 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -413,14 +413,11 @@ def _build_gather_rig(fmt: str, wl: Workload, device: torch.device):
     specs = _offload_bank_specs(fmt, H, I)
     cache = OffloadMoeCache(num_layers=1, num_experts=E, cache_size=E, device=device, quant_format=fmt)
     total_bytes = 0
+    sources = {}
     for name, (elems, dtype) in specs.items():
-        src = alloc_pinned_tensor(E, elems, dtype=dtype)  # cudaHostAlloc (mapped) like the loaders
-        dst = torch.empty(E, elems, dtype=dtype, device=device)
-        cache.bank_sources[name] = [src]
-        cache.bank_caches[name] = dst
-        cache.banks.append(([src], dst))
+        sources[name] = [alloc_pinned_tensor(E, elems, dtype=dtype)]  # cudaHostAlloc (mapped) like the loaders
         total_bytes += E * elems * dtype.itemsize
-    cache._build_copy_plan()  # enable the fused path (production default) when aligned
+    cache.set_bank_sources(sources)  # builds the fused copy plan (production default) when aligned
 
     cache.evict_slots[:E].copy_(torch.arange(E, dtype=torch.int32, device=device))
     cache.src_indices[:E].copy_(torch.randperm(E, device=device).to(torch.int32))
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 7500174ae4..5521c34e1f 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -48,6 +48,8 @@ class ExpertBanks:
     kind: QuantKind | None = None
     kernel: str | None = None
     layout: dict | None = None
+    # per-layer quant tag (e.g. the gate_up and down GGUF types), for layers whose banks match in dtype and row shape but not in quant; None = one quant
+    layer_tags: list | None = None
 
 
 def _dummy_fill(role: str, tensor: torch.Tensor) -> None:
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 984b7124ad..701524e0a1 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -77,7 +77,7 @@
         "down_bias",
     ),
     # DeepSeek-V4 FP4: packed e2m1 codes + e8m0 per-32 block scales, no global scale
-    # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via bank_views().
+    # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via pool_bank_views().
     "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"),
 }
 
@@ -104,6 +104,14 @@
 MARLIN_MAX_CACHE_SIZE = 992
 
 
+def bank_layouts(sources: dict[str, list[torch.Tensor]], layer_tags: list | None = None) -> list[tuple]:
+    """Each layer's layout key: its quant tag and every bank's (dtype, row shape). A slot pool holds layers of one key."""
+    num_layers = len(next(iter(sources.values())))
+    tags = layer_tags if layer_tags is not None else [None] * num_layers
+    assert len(tags) == num_layers, (len(tags), num_layers)
+    return [(tags[l], tuple((t[l].dtype, tuple(t[l].shape[1:])) for t in sources.values())) for l in range(num_layers)]
+
+
 @dataclass
 class OffloadMoeCache:
     num_layers: int
@@ -148,6 +156,9 @@ class OffloadMoeCache:
     # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots
     layout: dict | None = None
     max_slots: int | None = None
+    # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers.
+    pool_of_layer: list[int] | None = None
+    pool_sizes: list[int] | None = None
 
     def __post_init__(self) -> None:
         policy_ids = {"lru": 0}
@@ -155,6 +166,14 @@ def __post_init__(self) -> None:
         assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target
         if self.layout is None:
             assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}"
+        if self.pool_sizes is None:
+            self.pool_of_layer = [0] * self.num_layers
+            self.pool_sizes = [self.cache_size]
+        self.pool_of_layer = list(self.pool_of_layer)
+        self.pool_sizes = list(self.pool_sizes)
+        assert len(self.pool_of_layer) == self.num_layers, (len(self.pool_of_layer), self.num_layers)
+        assert sorted(set(self.pool_of_layer)) == list(range(len(self.pool_sizes))), self.pool_of_layer
+        assert sum(self.pool_sizes) == self.cache_size and min(self.pool_sizes) > 0, (self.pool_sizes, self.cache_size)
         # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None
         # for the GPU decode path.
         self.cpu_executor = None
@@ -162,8 +181,16 @@ def __post_init__(self) -> None:
         # offload/PCIe path. Set by the engine after construction (empty = all-GPU,
         # all layers = the plain --moe-strategy cpu case).
         self.cpu_layer_ids: frozenset = frozenset()
+        # Bank row layouts, set by set_bank_sources: layouts[g][k] is (dtype, row shape) of bank k (layouts of two quant tags may be equal), layer_layout[l] is layer l's g.
+        # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several it is a byte segment each pool views with its own layout.
+        self.layouts: list[tuple[tuple[torch.dtype, tuple[int, ...]], ...]] = []
+        self.layer_layout: list[int] = [0] * self.num_layers
+        self._raw_banks = False
+        # Set by the engine when hybrid decode hands CPU routes to a kernel that reads slot 0 with weight 0 (no supports_inactive_slots).
+        # With several quants a staging buffer can leave another quant's bytes in a pool's slot 0, which may decode to NaN, so the last prefill layer zeroes those rows.
+        self.zero_inactive_route_slots = False
         # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path.
-        self.validate_rebuild(self.cache_size)
+        self.validate_rebuild(self.cache_size, self.pool_sizes)
         assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, (
             "Prefill overlap borrows two full expert-layer buffers from the unified MoE "
             "cache, so cache_size must be at least 2 * num_experts "
@@ -254,25 +281,28 @@ def __post_init__(self) -> None:
         )
         # (per-layer sources, cache) per bank, in schema order. Every piece of cache
         # machinery that moves bank bytes (copy_missing, the prefill double buffers,
-        # bank_views) iterates this list, so the slot cache is bank-count agnostic.
+        # the pool and staging views) iterates this list, so the slot cache is bank-count agnostic.
         self.banks: list[tuple[list[torch.Tensor], torch.Tensor]] = []
         # Fused multi-bank copy descriptor (built by set_bank_sources/_build_copy_plan).
         # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device
-        # tensor); dst/feat are layer-invariant.
+        # tensor); dst pointers are per pool and row bytes per layout.
         self._copy_fused_ok = False
         self._copy_blocks_per_bank = 8
         self._copy_dst_ptrs: torch.Tensor | None = None
+        self._copy_dst_ptrs_pool: list[torch.Tensor] | None = None
         self._copy_src_ptrs: list[torch.Tensor] | None = None
-        self._copy_feat_bytes: torch.Tensor | None = None
+        self._copy_feat_bytes: list[torch.Tensor] | None = None
         # The layer whose misses ensure_experts/materialize_layer staged last; consumed
         # by copy_missing to pick the per-layer source (part of the same pending-copy
         # state as evict_slots/src_indices/num_indices).
         # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer
         self._pending_src_layer: int | None = None
         self._pending_whole_layer = False
-        # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's
-        # first 2 * num_experts slots (set up when prefill_overlap is enabled).
-        self.prefill_bank_buffers: list[torch.Tensor] = []
+        # The two prefill staging buffers are the first 2 * _stage_bytes[k] bytes of every bank, on top of the first pools.
+        # _stage_views[g][b] are buffer b's per-bank [num_experts, ...] views for layout g; _inval[b] are the global slots it overwrites.
+        self._stage_bytes: list[int] = []
+        self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = []
+        self._inval: list[slice | torch.Tensor] = []
         self.prefill_copy_stream: torch.cuda.Stream | None = None
         self.prefill_begin_event: torch.cuda.Event | None = None
         self.prefill_ready_events: list[torch.cuda.Event] = []
@@ -282,7 +312,7 @@ def __post_init__(self) -> None:
         self._prefill_buffer_has_release_event: list[bool] = [False, False]
         # hit-D2D split state: pinned begin-of-chunk snapshot of slot_for_id (the
         # classification input; frozen for the chunk -- no decode runs inside one,
-        # and buffer invalidation only clears slot < 2E entries, which classify as
+        # and buffer invalidation only clears slots under the staging buffers, which classify as
         # miss regardless), the lazily resolved batch-memcpy entry point (False =
         # unavailable), and row counters for cache reports.
         self._prefill_slot_snapshot: torch.Tensor | None = None
@@ -292,11 +322,116 @@ def __post_init__(self) -> None:
         self._batch_memcpy = None
         self.prefill_hit_rows = 0
         self.prefill_total_rows = 0
+        self._slice_pools()
+
+    def _pool_layout(self, pool: int) -> int:
+        return self.layer_layout[self.pool_of_layer.index(pool)]
+
+    def _bank_rows(self, k: int, layout: int, offset: int, nrows: int) -> torch.Tensor:
+        """``nrows`` rows of bank ``k`` in ``layout``, starting ``offset`` bytes into the bank cache."""
+        cache = self.banks[k][1]
+        row = self._row_bytes[layout][k]
+        if not self._raw_banks:
+            first = offset // row
+            return cache if first == 0 and nrows == cache.shape[0] else cache[first : first + nrows]
+        dtype, shape = self.layouts[layout][k]
+        return cache[offset : offset + nrows * row].view(dtype).view(nrows, *shape)
+
+    def _stage_fits(self, buffers: int) -> bool:
+        return all(c.numel() * c.element_size() >= buffers * s for (_, c), s in zip(self.banks, self._stage_bytes))
+
+    def _slice_pools(self) -> None:
+        E = self.num_experts
+        bases = [0]
+        for size in self.pool_sizes[:-1]:
+            bases.append(bases[-1] + size)
+        self.pool_bases = bases
+        spans = list(zip(bases, self.pool_sizes))
+
+        def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
+            return t if s == t.shape[0] else t[b : b + s]
+
+        self._pool_id_of_slot = [cut(self.id_of_slot, b, s) for b, s in spans]
+        self._pool_usage = [cut(self.usage, b, s) for b, s in spans]
+        # rows[k][p]: bytes of one slot of pool p in bank k; before the banks exist, one unit per slot with a stage of E
+        if self.banks:
+            rows = [[self._row_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))]
+            stage = self._stage_bytes
+        else:
+            rows, stage = [[1] * len(spans)], [E]
+        offsets = []
+        for row in rows:
+            offsets.append([0])
+            for s, r in zip(self.pool_sizes[:-1], row):
+                offsets[-1].append(offsets[-1][-1] + s * r)
+        self._pool_offsets = offsets
+        # local slots below this index overlap the prefill staging area in some bank
+        self._pool_stable_from = [
+            max(min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows)))
+            for p, s in enumerate(self.pool_sizes)
+        ]
+        self._inval = []
+        masks = []
+        for b in range(2):
+            mask = torch.zeros(self.cache_size, dtype=torch.bool)
+            masks.append(mask)
+            for k in range(len(rows)):
+                lo, hi = b * stage[k], (b + 1) * stage[k]
+                for p, (base, s) in enumerate(spans):
+                    first = max((lo - offsets[k][p]) // rows[k][p], 0)
+                    end = min(-(-(hi - offsets[k][p]) // rows[k][p]), s)
+                    mask[base + first : base + max(first, end)] = True
+            idx = mask.nonzero().flatten()
+            if idx.numel() == 0 or int(idx[-1] - idx[0]) + 1 == idx.numel():
+                self._inval.append(slice(int(idx[0]), int(idx[-1]) + 1) if idx.numel() else slice(0, 0))
+            else:
+                self._inval.append(idx.to(self.device))
+        self._stage_dirty = [False, False]
+        self._stage_slot0_rows = [[], []]
+        if not self.banks:
+            self._pool_bank_views, self._stage_views = [], []
+            return
+        self._pool_bank_views = [
+            tuple(self._bank_rows(k, self._pool_layout(p), offsets[k][p], s) for k in range(len(self.banks)))
+            for p, s in enumerate(self.pool_sizes)
+        ]
+        self._stage_views = [
+            [
+                tuple(self._bank_rows(k, g, b * stage[k], E) for k in range(len(self.banks))) if self._stage_fits(b + 1) else None
+                for b in range(2)
+            ]
+            for g in range(len(self.layouts))
+        ]
+        if len(self.layouts) > 1:
+            self._stage_slot0_rows = [
+                [self._bank_rows(k, self._pool_layout(p), offsets[k][p], 1) for p, base in enumerate(bases) if masks[b][base] for k in range(len(self.banks))]
+                for b in range(2)
+            ]
+
+    def stage_views(self, layer_id: int, buffer_id: int) -> tuple[torch.Tensor, ...]:
+        """Per-bank ``[num_experts, ...]`` views of prefill staging buffer ``buffer_id`` in ``layer_id``'s layout; row == expert id."""
+        views = self._stage_views[self.layer_layout[layer_id]][buffer_id]
+        assert views is not None, f"the slot cache is too small for staging buffer {buffer_id}"
+        return views
+
+    def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(id_of_slot, usage)`` views of the pool that serves ``layer_id``."""
+        p = self.pool_of_layer[layer_id]
+        return self._pool_id_of_slot[p], self._pool_usage[p]
+
+    def pool_slots(self, layer_id: int) -> int:
+        return self.pool_sizes[self.pool_of_layer[layer_id]]
+
+    def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]:
+        """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids."""
+        assert self.banks, "set_bank_sources must register the banks first"
+        return self._pool_bank_views[self.pool_of_layer[layer_id]]
 
     def set_bank_sources(
         self,
         sources: dict[str, list[torch.Tensor]],
         layer_residency: list[str] | None = None,
+        layer_tags: list | None = None,
     ) -> None:
         """Attach the host (CPU pinned) expert source banks and allocate a GPU slot
         cache per bank, following the format's bank schema.
@@ -310,6 +445,8 @@ def set_bank_sources(
 
         ``layer_residency`` labels each layer with a ``HostResidency`` value (default: all pinned).
         Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them.
+
+        ``layer_tags`` names each layer's quant (see :func:`bank_layouts`); layers of different tags never share a pool.
         """
         from freetoken.moe.legacy_format import canonical_role
         from freetoken.moe.host_banks import HostResidency
@@ -343,43 +480,62 @@ def set_bank_sources(
         for name in self.bank_schema:
             per_layer = sources[name]
             assert len(per_layer) == self.num_layers, (name, len(per_layer))
-            head = per_layer[0]
-            if self.layout is not None:
-                spec = self.layout[name]
-                if tuple(head.shape[1:]) != tuple(spec.shape) or head.dtype != spec.dtype:
-                    raise ValueError(
-                        f"bank {name!r} rows are {tuple(head.shape[1:])} {head.dtype} but the expert kernel's layout "
-                        f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel"
-                    )
             for layer_id, source in enumerate(per_layer):
                 assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous"
                 assert source.size(0) == self.num_experts, (name, layer_id, source.shape)
-                assert source.shape == head.shape and source.dtype == head.dtype, (
-                    name, layer_id, source.shape, source.dtype,
-                )
+                if self.layout is not None:
+                    spec = self.layout[name]
+                    if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype:
+                        raise ValueError(
+                            f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout "
+                            f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel"
+                        )
             self.bank_sources[name] = list(per_layer)
-            self.bank_caches[name] = torch.zeros(
-                (self.cache_size, *head.shape[1:]),
-                dtype=head.dtype,
-                device=self.device,
-            )
-        self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema]
+        keys = bank_layouts(sources, layer_tags)
+        unique = list(dict.fromkeys(keys))
+        self.layouts = [banks for _, banks in unique]
+        self.layer_layout = [unique.index(key) for key in keys]
+        for p in range(len(self.pool_sizes)):
+            mixed = {self.layer_layout[l] for l, q in enumerate(self.pool_of_layer) if q == p}
+            assert len(mixed) == 1, f"slot pool {p} holds layers of {len(mixed)} layouts; a pool needs one"
+        self._raw_banks = len(set(self.layouts)) > 1
+        self._row_bytes = [[math.prod(shape) * dtype.itemsize for dtype, shape in layout] for layout in self.layouts]
+        if self._raw_banks:
+            # every pool view and staging buffer starts on a row boundary of its own layout
+            assert all(r % 16 == 0 for rows in self._row_bytes for r in rows), self._row_bytes
+        self._stage_bytes = [self.num_experts * max(rows[k] for rows in self._row_bytes) for k in range(len(self.bank_schema))]
+        self._alloc_banks()
+        if not self._stage_fits(1):
+            raise ValueError(f"MoE cache of {self.cache_size} slots cannot stage one whole layer of experts")
+        if self.prefill_overlap and not self._stage_fits(2):
+            logger.warning("Disabling MoE prefill overlap: the slot cache cannot hold two whole layers of experts")
+            self.prefill_overlap = False
+        self._slice_pools()
         self._build_copy_plan()
         if self.prefill_overlap:
             self._init_prefill_overlap_buffers()
 
+    def _alloc_banks(self) -> None:
+        for k, name in enumerate(self.bank_schema):
+            if self._raw_banks:
+                nbytes = sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(self.pool_sizes))
+                self.bank_caches[name] = torch.zeros(nbytes, dtype=torch.uint8, device=self.device)
+            else:
+                dtype, shape = self.layouts[0][k]
+                self.bank_caches[name] = torch.zeros((self.cache_size, *shape), dtype=dtype, device=self.device)
+        self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema]
+
     def _build_copy_plan(self) -> None:
         self._build_fused_copy_plan()
         if self._copy_fused_ok or self.device.type != "cuda" or not self.banks:
             return
-        for name in self.bank_schema:
-            cache = self.bank_caches[name]
-            feat = math.prod(cache.shape[1:]) * cache.element_size()
-            if feat % 128:
-                raise RuntimeError(
-                    f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): "
-                    f"only the fused multi-bank copy can move them, but it is disabled"
-                )
+        for k, name in enumerate(self.bank_schema):
+            for rows in self._row_bytes:
+                if rows[k] % 128:
+                    raise RuntimeError(
+                        f"MoE bank {name!r} rows are {rows[k]} bytes (not a multiple of 128): "
+                        f"only the fused multi-bank copy can move them, but it is disabled"
+                    )
 
     def _build_fused_copy_plan(self) -> None:
         """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes).
@@ -391,24 +547,25 @@ def _build_fused_copy_plan(self) -> None:
         """
         self._copy_fused_ok = False
         self._copy_dst_ptrs = None
+        self._copy_dst_ptrs_pool = None
         self._copy_src_ptrs = None
         self._copy_feat_bytes = None
         self._copy_dst_ptrs_host: list[int] = []
         self._copy_src_ptrs_host: list[list[int]] = []
-        self._copy_feat_bytes_host: list[int] = []
-        self._gather_bank_ids: list[int] = []
-        self._gather_dst_ptrs: torch.Tensor | None = None
-        self._gather_feat_bytes: torch.Tensor | None = None
+        self._gather_bank_ids: list[list[int]] = []
+        self._gather_feat_bytes: list[torch.Tensor | None] = []
+        self._gather_stage_ptrs: list[list[torch.Tensor] | None] = []
+        self._gather_pool_ptrs: list[torch.Tensor | None] = []
         if not _FUSED_COPY or self.device.type != "cuda" or not self.banks:
             return
         from freetoken.kernel.pinned import device_ptr
 
-        dst_ptrs, feats = [], []
+        aligned = [self._stage_bytes] + self._row_bytes + self._pool_offsets
+        if any(n % 16 for ns in aligned for n in ns) or any(c.data_ptr() % 16 for _, c in self.banks):
+            return  # leave fused disabled; copy_missing uses the per-bank path
+        dst_ptrs = [cache.data_ptr() for _, cache in self.banks]
         layer_src_ptrs = [[] for _ in range(self.num_layers)]
-        for per_layer, cache in self.banks:
-            feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size()
-            if feat % 16 != 0 or cache.data_ptr() % 16 != 0:
-                return  # leave fused disabled; copy_missing uses the per-bank path
+        for per_layer, _ in self.banks:
             for layer_id, source in enumerate(per_layer):
                 if layer_id in self._unpinned_layers:
                     # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill)
@@ -422,52 +579,65 @@ def _build_fused_copy_plan(self) -> None:
                 if src_dev % 16 != 0:
                     return
                 layer_src_ptrs[layer_id].append(src_dev)
-            dst_ptrs.append(cache.data_ptr())
-            feats.append(feat)
-        self._copy_dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=self.device)
-        self._copy_src_ptrs = [
-            torch.tensor(ptrs, dtype=torch.int64, device=self.device)
-            for ptrs in layer_src_ptrs
+
+        def ptrs(values: list[int]) -> torch.Tensor:
+            return torch.tensor(values, dtype=torch.int64, device=self.device)
+
+        self._copy_dst_ptrs = ptrs(dst_ptrs)
+        self._copy_dst_ptrs_pool = [
+            self._copy_dst_ptrs if all(o[p] == 0 for o in self._pool_offsets)
+            else ptrs([d + o[p] for d, o in zip(dst_ptrs, self._pool_offsets)])
+            for p in range(len(self.pool_sizes))
         ]
-        self._copy_feat_bytes = torch.tensor(feats, dtype=torch.int64, device=self.device)
+        self._copy_src_ptrs = [ptrs(values) for values in layer_src_ptrs]
+        self._copy_feat_bytes = [ptrs(rows) for rows in self._row_bytes]
         self._copy_dst_ptrs_host = dst_ptrs
         self._copy_src_ptrs_host = layer_src_ptrs
-        self._copy_feat_bytes_host = feats
         # hit-D2D gather serves only the big banks; small banks are whole-layer
         # H2D entries (see _SMALL_BANK_FEAT_BYTES), so their rows never need D2D.
-        self._gather_bank_ids = [i for i, f in enumerate(feats) if f >= _SMALL_BANK_FEAT_BYTES]
-        if len(self._gather_bank_ids) == len(feats):
-            self._gather_dst_ptrs = self._copy_dst_ptrs
-            self._gather_feat_bytes = self._copy_feat_bytes
-        elif self._gather_bank_ids:
-            self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous()
-            self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous()
+        # The gather copies cache rows of the layer's pool into its staging buffer, row == expert id.
+        self._gather_bank_ids = [[k for k, r in enumerate(rows) if r >= _SMALL_BANK_FEAT_BYTES] for rows in self._row_bytes]
+        for g, ids in enumerate(self._gather_bank_ids):
+            self._gather_feat_bytes.append(ptrs([self._row_bytes[g][k] for k in ids]) if ids else None)
+            self._gather_stage_ptrs.append(
+                [ptrs([dst_ptrs[k] + b * self._stage_bytes[k] for k in ids]) for b in range(2)] if ids else None
+            )
+        for p in range(len(self.pool_sizes)):
+            ids = self._gather_bank_ids[self._pool_layout(p)]
+            self._gather_pool_ptrs.append(ptrs([dst_ptrs[k] + self._pool_offsets[k][p] for k in ids]) if ids else None)
         self._copy_blocks_per_bank = default_blocks_per_bank(self.device)
         self._copy_fused_ok = True
 
-    def validate_rebuild(self, cache_size: int) -> None:
+    def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         """Pure geometry validation of a rebuild target (no GPU side effects).
 
         Raises ``ValueError`` if ``cache_size`` is below the ``num_experts`` floor or
-        above the marlin slot cap. Called by :meth:`rebuild` and by the engine's
+        a pool is above the expert kernel's slot cap, or the bank byte segments cannot stage one layer. Called by :meth:`rebuild` and by the engine's
         pre-teardown check, so an invalid target rejects with the old cache intact
         (no destructive free first).
         """
         if cache_size < self.num_experts:
             raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}")
-        if self.max_slots is not None and cache_size > self.max_slots:
+        # decode kernels see one pool's slots, so the cap applies per pool
+        largest = cache_size if pool_sizes is None else max(pool_sizes)
+        what = f"moe_cache_size={cache_size}" if largest == cache_size else f"a MoE slot pool of {largest} slots"
+        if self.max_slots is not None and largest > self.max_slots:
             raise ValueError(
-                f"moe_cache_size={cache_size} exceeds the expert kernel's slot limit of {self.max_slots}; "
+                f"{what} exceeds the expert kernel's slot limit of {self.max_slots}; "
                 f"pass --moe-cache-size {self.max_slots} or less, or let the default kernel serve the experts"
             )
-        if self.layout is None and self.quant_format == "nvfp4_marlin" and cache_size > MARLIN_MAX_CACHE_SIZE:
+        if self.layout is None and self.quant_format == "nvfp4_marlin" and largest > MARLIN_MAX_CACHE_SIZE:
             raise ValueError(
-                f"moe_cache_size={cache_size} exceeds the marlin backend's slot limit of "
+                f"{what} exceeds the marlin backend's slot limit of "
                 f"{MARLIN_MAX_CACHE_SIZE} (vLLM moe_align_block_size caps padded experts at "
                 "1024); reduce moe_cache_size or force --quant-backend moe.nvfp4=triton"
             )
+        if self._raw_banks and pool_sizes is not None:
+            for k, stage in enumerate(self._stage_bytes):
+                if sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(pool_sizes)) < stage:
+                    raise ValueError(f"moe_cache_size={cache_size} cannot stage one whole layer of experts")
 
-    def rebuild(self, cache_size: int) -> None:
+    def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         """Resize the GPU slot cache + bookkeeping to ``cache_size`` IN PLACE.
 
         Keeps the CPU/pinned ``bank_sources`` and the GPU-resident alphas; never
@@ -477,9 +647,10 @@ def rebuild(self, cache_size: int) -> None:
         ``ctx.moe_offload_cache`` stay valid.
         """
         assert self.bank_sources, "set_bank_sources must run before rebuild"
-        self.validate_rebuild(cache_size)
+        sizes = [cache_size] if pool_sizes is None else list(pool_sizes)
+        assert len(sizes) == len(self.pool_sizes) and sum(sizes) == cache_size and min(sizes) > 0, (sizes, cache_size)
+        self.validate_rebuild(cache_size, sizes)
         # 1. Tear down prefill-overlap (its buffer views alias the old bank_caches).
-        self.prefill_bank_buffers = []
         self.prefill_copy_stream = None
         self.prefill_begin_event = None
         self.prefill_ready_events = []
@@ -490,22 +661,22 @@ def rebuild(self, cache_size: int) -> None:
         # 2. Drop old GPU tensors (free-before-alloc).
         self.banks = []
         self.bank_caches = {}
+        self._pool_bank_views = []
+        self._stage_views = []
+        self._stage_slot0_rows = [[], []]
         self.cache_size = cache_size
+        self.pool_sizes = sizes
         if self.device.type == "cuda":
             torch.cuda.synchronize(self.device)
             torch.cuda.empty_cache()
         # 3. Reallocate the slot cache from the retained host sources.
-        for name in self.bank_schema:
-            head = self.bank_sources[name][0]
-            self.bank_caches[name] = torch.zeros(
-                (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device
-            )
-        self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema]
-        self._build_copy_plan()  # slot caches were reallocated -> refresh fused-copy addrs
+        self._alloc_banks()
         # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start).
         self.slot_for_id.fill_(-1)
         self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device)
         self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device)
+        self._slice_pools()
+        self._build_copy_plan()  # slot caches were reallocated -> refresh fused-copy addrs
         plan_slots = max(self.num_experts, cache_size)
         self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
         self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
@@ -529,10 +700,10 @@ def rebuild(self, cache_size: int) -> None:
         self.prefill_total_rows = 0
         self._hit_d2d_fallback_logged = False  # geometry changed; re-log if still unusable
         # 5. Re-evaluate prefill overlap against the new size.
-        if self.prefill_overlap and cache_size < 2 * self.num_experts:
+        if self.prefill_overlap and not self._stage_fits(2):
             logger.warning(
                 f"Disabling MoE prefill overlap on rebuild: cache_size {cache_size} "
-                f"< 2*num_experts {2 * self.num_experts}."
+                f"cannot hold two whole layers of experts (2*num_experts {2 * self.num_experts} slots of one layout)."
             )
             self.prefill_overlap = False
         if self.prefill_overlap:
@@ -585,8 +756,9 @@ def alphas_for_slots(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] |
         ``layer_id`` -- are ever read by the grouped GEMM."""
         if self.gate_up_alpha is None:
             return None
+        id_of_slot, _ = self.pool_state(layer_id)
         idx = layer_id * self.num_experts + (
-            self.id_of_slot.clamp(min=0).long() % self.num_experts
+            id_of_slot.clamp(min=0).long() % self.num_experts
         )
         return self.gate_up_alpha[idx], self.down_alpha[idx]
 
@@ -600,25 +772,12 @@ def alphas_for_layer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] |
         hi = lo + self.num_experts
         return self.gate_up_alpha[lo:hi], self.down_alpha[lo:hi]
 
-    def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]:
-        """Per-bank cache views in registration order: the full ``[S]`` slot cache
-        (decode), or its first ``n`` slots (materialized layer)."""
-        assert self.banks, "set_bank_sources must register the banks first"
-        if n is None:
-            return tuple(cache for _, cache in self.banks)
-        return tuple(cache[:n] for _, cache in self.banks)
-
     def _init_prefill_overlap_buffers(self) -> None:
         assert self.banks, "set_bank_sources must register the banks first"
         self._prefill_buffer_layer = [None, None]
         self._prefill_buffer_released = [True, True]
         self._prefill_buffer_has_release_event = [False, False]
-        # The double buffers borrow the slot cache's first 2 * num_experts slots
-        # (one full expert layer per buffer), one view per registered bank.
-        self.prefill_bank_buffers = [
-            cache[: 2 * self.num_experts].view(2, self.num_experts, *cache.shape[1:])
-            for _, cache in self.banks
-        ]
+        # the double buffers are the staging views (_slice_pools), on top of the first pools' slots
         if self.device.type == "cuda":
             self.prefill_copy_stream = torch.cuda.Stream(device=self.device)
             self.prefill_ready_events = [torch.cuda.Event() for _ in range(2)]
@@ -638,14 +797,14 @@ def _init_prefill_overlap_buffers(self) -> None:
             self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device)
 
     def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
-        slot_start = buffer_id * self.num_experts
-        slot_end = slot_start + self.num_experts
-        old_ids = self.id_of_slot[slot_start:slot_end]
+        slots = self._inval[buffer_id]
+        old_ids = self.id_of_slot[slots]
         self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
-        old_ids.fill_(-1)
+        self.id_of_slot[slots] = -1
         # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in
         # ensure_experts evicts them first.
-        self.usage[slot_start:slot_end].zero_()
+        self.usage[slots] = 0
+        self._stage_dirty[buffer_id] = True
 
     def begin_prefill(self) -> None:
         if not self.prefill_overlap:
@@ -676,7 +835,7 @@ def prefetch_prefill_layer(self, layer_id: int) -> None:
         if layer_id < 0:
             raise ValueError(f"Invalid prefill layer id: {layer_id}")
 
-        assert self.banks and self.prefill_bank_buffers
+        assert self.banks and self._stage_views
 
         buffer_id = layer_id % 2
         if self._prefill_buffer_layer[buffer_id] == layer_id:
@@ -688,8 +847,8 @@ def prefetch_prefill_layer(self, layer_id: int) -> None:
 
         def copy() -> None:
             self._invalidate_prefill_buffer(buffer_id)
-            for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers):
-                buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True)
+            for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)):
+                view.copy_(per_layer[layer_id], non_blocking=True)
 
         if self._prefill_hit_d2d_active:
             self._prefetch_split(layer_id, buffer_id)
@@ -720,10 +879,10 @@ def _hit_d2d_usable(self) -> bool:
             reason = "FREETOKEN_SKIP_FAST_INDEX_COPY is set (the hit gather would be a no-op)"
         elif not self._copy_fused_ok:
             reason = "the fused copy plan is unavailable (bank alignment or FREETOKEN_FUSED_COPY=0)"
-        elif self.cache_size <= 2 * self.num_experts:
+        elif all(stable >= s for stable, s in zip(self._pool_stable_from, self.pool_sizes)):
             reason = (
                 f"cache_size {self.cache_size} leaves no hit region "
-                f"(needs > {2 * self.num_experts} slots)"
+                f"(every slot lies under the prefill staging buffers of {2 * self.num_experts} experts)"
             )
         elif not self._resolve_batch_memcpy():
             reason = "cudaMemcpyBatchAsync is unavailable"  # resolve logged the specifics
@@ -761,7 +920,7 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
         run list comes from the begin-of-chunk snapshot because the batch API
         takes HOST pointer arrays. Live-vs-snapshot cannot disagree: the only
         chunk-internal writer (buffer invalidation) rewrites slots already below
-        the 2E threshold, and slots < 2E (including -1) are misses on both sides
+        the pool's staging threshold, and those slots (and -1) are misses on both sides
         -- the buffers own those slots, so their bytes are volatile within the
         chunk. Hit and miss row sets are disjoint, so the streams need no
         ordering against each other.
@@ -772,18 +931,20 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
         from freetoken.moe.offload_kernels import prefill_hit_compact
 
         E = self.num_experts
+        layout = self.layer_layout[layer_id]
+        pool = self.pool_of_layer[layer_id]
         snap = self._prefill_snapshot_np[layer_id]
-        hit_mask = snap >= 2 * E
+        hit_mask = snap >= self._pool_stable_from[pool]
         self.prefill_hit_rows += int(hit_mask.sum())
         self.prefill_total_rows += E
-        if self._gather_dst_ptrs is not None:
-            prefill_hit_compact(self, layer_id, buffer_id)
+        if self._gather_feat_bytes[layout] is not None:
+            prefill_hit_compact(self, layer_id)
             # blocks_per_bank=64 vs the PCIe-tuned default of 8: HBM D2D needs the
             # wider grid (~22 GB/s per 1024-thread block on H100).
             fast_index_copy_multi_jit(
-                self._gather_dst_ptrs,
-                self._gather_dst_ptrs,
-                self._gather_feat_bytes,
+                self._gather_stage_ptrs[layout][buffer_id],
+                self._gather_pool_ptrs[pool],
+                self._gather_feat_bytes[layout],
                 self._prefill_hit_dst,
                 self._prefill_hit_src,
                 self._prefill_hit_num,
@@ -799,16 +960,17 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
                 starts = miss[run_starts]
                 lengths = np.diff(np.concatenate((run_starts, [miss.size])))
             dst, src, nbytes = [], [], []
-            for b, feat in enumerate(self._copy_feat_bytes_host):
+            for b, feat in enumerate(self._row_bytes[layout]):
+                stage = self._copy_dst_ptrs_host[b] + buffer_id * self._stage_bytes[b]
                 if feat < _SMALL_BANK_FEAT_BYTES:
                     # Whole layer as one entry, EVEN with zero misses: it keeps every
                     # batch entry above the driver's async floor and covers the hit
                     # rows the gather skips for these banks.
-                    dst.append(self._copy_dst_ptrs_host[b] + buffer_id * E * feat)
+                    dst.append(stage)
                     src.append(self._copy_src_ptrs_host[layer_id][b])
                     nbytes.append(E * feat)
                 elif miss.size:
-                    dst.extend(self._copy_dst_ptrs_host[b] + (buffer_id * E + starts) * feat)
+                    dst.extend(stage + starts * feat)
                     src.extend(self._copy_src_ptrs_host[layer_id][b] + starts * feat)
                     nbytes.extend(lengths * feat)
             if dst:
@@ -826,15 +988,22 @@ def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]:
         marlin/b12x ``(gate_up_packed, gate_up_scale, down_packed, down_scale)``;
         nvfp4 native adds the two global banks after each scale bank."""
         assert self.prefill_overlap
-        assert self.prefill_bank_buffers
+        assert self._stage_views
         self.prefetch_prefill_layer(layer_id)
         buffer_id = layer_id % 2
         assert self._prefill_buffer_layer[buffer_id] == layer_id
         if self.prefill_ready_events:
             torch.cuda.current_stream(self.device).wait_event(self.prefill_ready_events[buffer_id])
-        return tuple(buffer[buffer_id] for buffer in self.prefill_bank_buffers)
+        return self.stage_views(layer_id, buffer_id)
 
     def release_prefill_layer(self, layer_id: int) -> None:
+        if self.zero_inactive_route_slots and layer_id == self.num_layers - 1:
+            # the staged slots were invalidated and nothing registers them within the forward, so their bytes are free to clear
+            for b in range(2):
+                if self._stage_dirty[b]:
+                    for row in self._stage_slot0_rows[b]:
+                        row.zero_()
+                    self._stage_dirty[b] = False
         if not self.prefill_overlap:
             return
         buffer_id = layer_id % 2
@@ -959,6 +1128,8 @@ def decode_miss_stats_per_layer(self) -> dict:
             s, m, a, f = steps[L], missing[L], active[L], fetched[L]
             per_layer.append({
                 "layer": L,
+                "pool": self.pool_of_layer[L],
+                "pool_slots": self.pool_slots(L),
                 "steps": s,
                 "active_per_step": (a / s) if s else 0.0,
                 "missing_per_step": (m / s) if s else 0.0,
@@ -1010,11 +1181,13 @@ def copy_missing(self) -> None:
                     f"pageable materialize (position == expert id); ensure_experts's "
                     f"LRU slot remap cannot be honored without a device alias"
                 )
-            # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy
+            # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into staging buffer 0 with position == expert id -- a plain synchronous pageable H2D copy
             # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor)
-            for per_layer, cache in self.banks:
-                cache[: self.num_experts].copy_(per_layer[layer_id])
+            for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)):
+                view.copy_(per_layer[layer_id])
             return
+        # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool
+        whole = self._pending_whole_layer
         if self._copy_fused_ok:
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
 
@@ -1023,9 +1196,9 @@ def copy_missing(self) -> None:
             # src_indices holds layer-local expert rows, resolved against this layer's
             # source pointers (layer_id is a static int per captured graph node).
             fast_index_copy_multi_jit(
-                self._copy_dst_ptrs,
+                self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]],
                 self._copy_src_ptrs[layer_id],
-                self._copy_feat_bytes,
+                self._copy_feat_bytes[self.layer_layout[layer_id]],
                 self.evict_slots,
                 self.src_indices,
                 self.num_indices,
@@ -1035,7 +1208,8 @@ def copy_missing(self) -> None:
 
         from freetoken.kernel import fast_index_copy_jit
 
-        for per_layer, cache in self.banks:
+        views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id)
+        for (per_layer, _), cache in zip(self.banks, views):
             fast_index_copy_jit(
                 cache,
                 self.evict_slots,
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index ea12b6fb17..1d863a2f54 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -26,11 +26,12 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
     tensor. ``out_indices`` aliases the input, preserving the in-place rewrite every
     downstream GEMM depends on.
     """
+    id_of_slot, usage = cache.pool_state(layer_id)
     lru_ensure(
         expert_ids,
         cache.slot_for_id.view(-1),
-        cache.id_of_slot,
-        cache.usage,
+        id_of_slot,
+        usage,
         cache.step,
         expert_ids,
         cache.src_indices,
@@ -80,13 +81,13 @@ def ensure_experts_hybrid(
     return routes
 
 
-def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None:
+def prefill_hit_compact(cache, layer_id: int) -> None:
     """Compact this layer's cache-resident experts into gather indices, device-side.
 
-    hit = slot_for_id[layer_id][e] >= 2 * num_experts (the double buffer owns the
+    hit = slot_for_id[layer_id][e] >= the pool's staging threshold (the double buffer owns the
     slots below, so those bytes are volatile within a prefill chunk and classify
     as miss). Writes fixed-shape ``_prefill_hit_dst``/``_prefill_hit_src`` (buffer
-    row / cache slot) and the count into ``_prefill_hit_num``; one launch on the
+    row / pool-local slot) and the count into ``_prefill_hit_num``; one launch on the
     current stream, no host sync. Safe against the concurrent buffer invalidation
     on the copy stream: that only rewrites entries already below the threshold."""
     num_experts = cache.num_experts
@@ -95,15 +96,23 @@ def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None:
         cache._prefill_hit_dst,
         cache._prefill_hit_src,
         cache._prefill_hit_num,
-        buffer_id * num_experts,
-        2 * num_experts,
+        cache._pool_stable_from[cache.pool_of_layer[layer_id]],
         num_experts,
         BLOCK=triton.next_power_of_2(num_experts),
     )
 
 
 def materialize_layer(cache, layer_id: int) -> None:
-    _materialize_layer_gpu(cache, layer_id)
+    if len(cache.pool_sizes) == 1:
+        _materialize_layer_gpu(cache, layer_id)
+        return
+    # staging buffer 0 may overlap several pools, so it only holds the layer for this forward and registers nothing
+    E = cache.num_experts
+    cache._invalidate_prefill_buffer(0)
+    rows = torch.arange(E, dtype=torch.int32, device=cache.device)
+    cache.evict_slots[:E].copy_(rows)
+    cache.src_indices[:E].copy_(rows)
+    cache.num_indices.fill_(E)
 
 
 def reset_cache(cache) -> None:
@@ -124,14 +133,16 @@ def _ensure_experts_hybrid_gpu(
     max_fetch: int,
     frac_q16: int,
 ) -> None:
+    id_of_slot, usage = cache.pool_state(layer_id)
+    num_slots = id_of_slot.numel()
     block_e = triton.next_power_of_2(cache.num_experts)
-    block_c = triton.next_power_of_2(cache.cache_size)
+    block_c = triton.next_power_of_2(num_slots)
     num_warps = 8 if block_c >= 2048 else 4
     _ensure_experts_hybrid_kernel[(1,)](
         expert_ids,
         cache.slot_for_id,
-        cache.id_of_slot,
-        cache.usage,
+        id_of_slot,
+        usage,
         cache.step,
         cache.active_mask,
         cache.evict_slots,
@@ -148,7 +159,7 @@ def _ensure_experts_hybrid_gpu(
         int(max_fetch),
         int(frac_q16),
         cache.num_experts,
-        cache.cache_size,
+        num_slots,
         BLOCK_E=block_e,
         BLOCK_C=block_c,
         BY_RECENCY=_HYBRID_FETCH_BY_RECENCY,
@@ -257,7 +268,7 @@ def _materialize_layer_kernel(
 
 
 
-@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16"])
+@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16", "cache_size", "id_of_slot_ptr", "usage_ptr"])
 def _ensure_experts_hybrid_kernel(
     expert_ids_ptr,
     slot_for_id_ptr,
@@ -280,7 +291,7 @@ def _ensure_experts_hybrid_kernel(
     max_fetch,
     fetch_frac_q16,
     num_experts: tl.constexpr,
-    cache_size: tl.constexpr,
+    cache_size,
     BLOCK_E: tl.constexpr,
     BLOCK_C: tl.constexpr,
     BY_RECENCY: tl.constexpr,
@@ -329,10 +340,11 @@ def _ensure_experts_hybrid_kernel(
             (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16
         )
         max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1)
-    num_fetch = tl.minimum(num_missing, max_fetch)
+    is_hit = is_active & (slot >= 0)
+    # a pool smaller than num_experts may hold fewer free slots than misses
+    num_fetch = tl.minimum(tl.minimum(num_missing, max_fetch), cache_size - tl.sum(is_hit.to(tl.int32)))
     tl.store(num_missing_full_ptr, num_missing.to(tl.int64))
     tl.store(num_indices_ptr, num_fetch.to(tl.int64))
-    is_hit = is_active & (slot >= 0)
     tl.store(usage_ptr + slot, step, mask=is_hit)
 
     # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered
@@ -405,14 +417,13 @@ def _ensure_experts_hybrid_kernel(
         tl.store(expert_recency_ptr + base + off_e, step_vec, mask=is_active & e_mask)
 
 
-@triton.jit(do_not_specialize=["buffer_base"])
+@triton.jit(do_not_specialize=["threshold"])
 def _prefill_hit_compact_kernel(
     slot_ptr,     # [num_experts] int32: this layer's slot_for_id row
-    dst_ptr,      # [num_experts] int32 out: buffer rows, compacted
-    src_ptr,      # [num_experts] int32 out: cache slots, compacted
+    dst_ptr,      # [num_experts] int32 out: buffer rows (expert ids), compacted
+    src_ptr,      # [num_experts] int32 out: pool-local slots, compacted
     num_ptr,      # [1] int64 out: hit count
-    buffer_base,  # buffer_id * num_experts
-    threshold,    # 2 * num_experts
+    threshold,    # first pool-local slot outside the staging buffers
     num_experts,
     BLOCK: tl.constexpr,
 ):
@@ -421,6 +432,6 @@ def _prefill_hit_compact_kernel(
     slots = tl.load(slot_ptr + offs, mask=lane, other=-1)
     is_hit = lane & (slots >= threshold)
     pos = tl.cumsum(is_hit.to(tl.int32)) - 1
-    tl.store(dst_ptr + pos, (buffer_base + offs).to(tl.int32), mask=is_hit)
+    tl.store(dst_ptr + pos, offs.to(tl.int32), mask=is_hit)
     tl.store(src_ptr + pos, slots, mask=is_hit)
     tl.store(num_ptr, tl.sum(is_hit.to(tl.int64)))
diff --git a/python/freetoken/moe/pool_policy.py b/python/freetoken/moe/pool_policy.py
new file mode 100644
index 0000000000..990c18ade2
--- /dev/null
+++ b/python/freetoken/moe/pool_policy.py
@@ -0,0 +1,149 @@
+"""Slot pool policy for the offload MoE cache: which layers share a pool, and how many slots each pool gets.
+
+Pure functions. The engine supplies the constraints and hands the result to ``OffloadMoeCache``:
+
+    groups = group_layers(scope, layouts, slot_bytes)                    # --moe-cache-scope, once per load
+    pools = make_pools(groups, num_experts, min_num_slots, slot_limit)   # the minimum follows the decode batch
+    sizes = split_slots(total, pools, counts)                            # --moe-cache-split, counts from resolve_profile
+    OffloadMoeCache(pool_of_layer=pool_of_layer(pools, num_layers), pool_sizes=sizes, ...)
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+SCOPES = ("unified", "per-layer")
+
+
+@dataclass(frozen=True)
+class PoolSpec:
+    layers: tuple[int, ...]
+    min_num_slots: int
+    max_num_slots: int
+
+
+def group_layers(scope: str, layouts: list, slot_bytes: list[int]) -> list[list[int]]:
+    """The layers of each pool, in placement order.
+
+    ``unified``: one pool per layout key (see ``bank_layouts``); ``per-layer``: one pool per layer.
+    Pools with wider slots come first: the prefill staging buffers sit at the start of every bank, so they then cover fewer slots.
+    """
+    if scope == "unified":
+        groups: dict = {}
+        for layer, layout in enumerate(layouts):
+            groups.setdefault(layout, []).append(layer)
+        pools = list(groups.values())
+    elif scope == "per-layer":
+        pools = [[layer] for layer in range(len(layouts))]
+    else:
+        raise ValueError(f"unknown --moe-cache-scope {scope!r}; want one of {SCOPES}")
+    return sorted(pools, key=lambda layers: (-slot_bytes[layers[0]], layers[0]))
+
+
+def make_pools(groups: list[list[int]], num_experts: int, min_num_slots: int, slot_limit: int | None) -> list[PoolSpec]:
+    """Each group's slot range: at least ``min_num_slots``, at most one slot per expert of its layers and the kernel's ``slot_limit``."""
+    pools = []
+    for layers in groups:
+        max_num_slots = len(layers) * num_experts
+        if slot_limit is not None:
+            max_num_slots = min(max_num_slots, slot_limit)
+        pools.append(PoolSpec(tuple(layers), min(min_num_slots, max_num_slots), max_num_slots))
+    return pools
+
+
+def pool_of_layer(pools: list[PoolSpec], num_layers: int) -> list[int]:
+    """Each layer's pool index, the form ``OffloadMoeCache`` takes."""
+    owner = [-1] * num_layers
+    for p, pool in enumerate(pools):
+        for layer in pool.layers:
+            owner[layer] = p
+    assert min(owner) >= 0, owner
+    return owner
+
+
+def split_slots(total: int, pools: list[PoolSpec], counts: torch.Tensor | None = None) -> list[int]:
+    """Slots per pool, each in [min_num_slots, max_num_slots], summing to ``total``.
+
+    With ``counts`` (per-layer expert use from ``resolve_profile``) the slots follow expert use, else the pools' layer counts.
+    """
+    if len(pools) == 1:
+        return [total]
+    _check_budget(total, pools)
+    return _split_equal(total, pools) if counts is None else _split_hot(total, pools, counts)
+
+
+def _check_budget(total: int, pools: list[PoolSpec]) -> None:
+    need = sum(p.min_num_slots for p in pools)
+    room = sum(p.max_num_slots for p in pools)
+    if total < need:
+        raise ValueError(f"MoE cache of {total} slots is below what the pools need ({need} slots)")
+    if total > room:
+        raise ValueError(f"MoE cache of {total} slots exceeds what the pools can hold ({room} slots)")
+
+
+def _split_equal(total: int, pools: list[PoolSpec]) -> list[int]:
+    sizes = [p.min_num_slots for p in pools]
+    weights = [len(p.layers) for p in pools]
+    left = total - sum(sizes)
+    while left > 0:
+        open_ = [i for i, p in enumerate(pools) if sizes[i] < p.max_num_slots]
+        # give the next slot to the pool with the fewest slots per layer
+        i = min(open_, key=lambda j: (sizes[j] / weights[j], j))
+        sizes[i] += 1
+        left -= 1
+    return sizes
+
+
+def _split_hot(total: int, pools: list[PoolSpec], counts: torch.Tensor) -> list[int]:
+    sizes = [p.min_num_slots for p in pools]
+    # a pool's minimum covers its hottest experts; the rest, up to its maximum, competes across pools
+    cand = []
+    for p, pool in enumerate(pools):
+        hot = counts[list(pool.layers)].flatten().sort(descending=True).values
+        cand.extend((float(c), p) for c in hot[pool.min_num_slots : pool.max_num_slots])
+    cand.sort(key=lambda c: -c[0])
+    for _, p in cand[: total - sum(sizes)]:
+        sizes[p] += 1
+    return sizes
+
+
+def resolve_profile(spec: str, model_path: str, num_layers: int, num_experts: int) -> torch.Tensor | None:
+    """Expert use for ``--moe-cache-split``: None for ``equal``, else the profile at ``profile:<path>``."""
+    if spec == "equal":
+        return None
+    path = Path(spec.removeprefix("profile:"))
+    logger.info_rank0(f"MoE pool profile: {path}")
+    return load_profile(path, num_layers, num_experts, model_path)
+
+
+def load_profile(path: str | Path, num_layers: int, num_experts: int, model_path: str = "") -> torch.Tensor:
+    """Per-layer expert use ``[num_layers, num_experts]`` from a profile, each row scaled to sum to 1."""
+    data = json.loads(Path(path).read_text())
+    model = str(data.get("model", ""))
+    if model and model_path and not any(n.lower().startswith(model.lower()) for n in _model_names(model_path)):
+        logger.warning_rank0(f"MoE pool profile {path} was taken on {model!r}, not on {model_path}")
+    counts = torch.tensor(data["expert_counts"], dtype=torch.float64)
+    if tuple(counts.shape) != (num_layers, num_experts):
+        raise ValueError(
+            f"profile {path} has {tuple(counts.shape)} expert counts; this model needs ({num_layers}, {num_experts})"
+        )
+    # every layer decodes the same tokens, so equal row sums keep the layers comparable
+    return counts / counts.sum(dim=1, keepdim=True).clamp(min=1e-12)
+
+
+def _model_names(model_path: str) -> list[str]:
+    # every path component, so an HF cache snapshot (models--Org--Name/snapshots/<hash>) still finds its name
+    names = []
+    for part in reversed(Path(model_path).parts):
+        name = part.split("--")[-1].removesuffix(".gguf")
+        if name:
+            names.append(name)
+    return names
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 7c0ed9b047..be244a117c 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -134,6 +134,7 @@ def parse_args(
     from freetoken.attention import validate_attn_backend
     from freetoken.kvcache import SUPPORTED_CACHE_MANAGER
     from freetoken.moe import MOE_STRATEGIES
+    from freetoken.moe.pool_policy import SCOPES
 
     def _parse_quant_backend(value: str) -> str:
         from freetoken.layers.quantization import QuantBackend
@@ -153,6 +154,27 @@ def _parse_moe_cache_rate(value: str) -> float:
             raise argparse.ArgumentTypeError("must be in [0, 1]")
         return rate
 
+    def _parse_mem_size(value: str) -> int:
+        units = {
+            "b": 1, "k": 1 << 10, "kb": 1 << 10, "kib": 1 << 10, "m": 1 << 20, "mb": 1 << 20, "mib": 1 << 20,
+            "g": 1 << 30, "gb": 1 << 30, "gib": 1 << 30,
+        }
+        text = value.strip().lower()
+        number = text.rstrip("abgikm")
+        unit = text[len(number):]
+        try:
+            size = int(float(number) * units[unit])
+        except (KeyError, ValueError, OverflowError) as exc:
+            raise argparse.ArgumentTypeError("must be a size with a unit, e.g. 12GiB or 800MiB") from exc
+        if size <= 0:
+            raise argparse.ArgumentTypeError("must be positive")
+        return size
+
+    def _parse_cache_split(value: str) -> str:
+        if value == "equal" or (value.startswith("profile:") and len(value) > len("profile:")):
+            return value
+        raise argparse.ArgumentTypeError("must be equal or profile:<path>")
+
     def _positive_int(value: str) -> int:
         try:
             n = int(value)
@@ -705,6 +727,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         default=ServerArgs.moe_cache_rate,
         help="The fraction of all MoE experts to keep in GPU cache.",
     )
+    moe_cache_group.add_argument(
+        "--moe-cache-mem",
+        type=_parse_mem_size,
+        default=ServerArgs.moe_cache_mem,
+        help="GPU memory for the MoE expert slot cache, e.g. 12GiB or 800MiB (G, GB and GiB all mean GiB); the slot count is this over the bytes per slot.",
+    )
     moe_cache_group.add_argument(
         "--moe-cache-auto",
         action="store_true",
@@ -722,6 +750,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="KV-cache token floor reserved before --moe-cache-auto fills experts.",
     )
 
+    parser.add_argument(
+        "--moe-cache-scope",
+        choices=SCOPES,
+        default=ServerArgs.moe_cache_scope,
+        help=(
+            "How MoE layers share the GPU expert slot cache: unified = one LRU pool per expert quant "
+            "(all layers when they share one quant format); per-layer = each layer has its own pool. "
+            "Pools are sized by --moe-cache-split."
+        ),
+    )
+    parser.add_argument(
+        "--moe-cache-split",
+        type=_parse_cache_split,
+        default=ServerArgs.moe_cache_split,
+        help=(
+            "How pools share the slots: equal (by layer count), or profile:<path> = by expert use from a "
+            "profile file. With --moe-strategy offload, a pool smaller than the expert count limits the "
+            "decode batch to its slots / top_k."
+        ),
+    )
+
     parser.add_argument(
         "--moe-cache-policy",
         default=ServerArgs.moe_cache_policy,
@@ -892,6 +941,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
     _no_cache_flag = (
         kwargs["moe_cache_size"] == 0
         and not kwargs["moe_cache_auto"]
+        and kwargs["moe_cache_mem"] is None
         and (kwargs["moe_cache_rate"] is None or kwargs["moe_cache_rate"] == 0)
     )
     if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag:
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index a77409f013..7b37af9e0e 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -523,3 +523,55 @@ def test_uncapped_platform_stays_uncapped(monkeypatch):
     if hasattr(os, "uname") and "microsoft" in os.uname().release.lower():
         pytest.skip("WSL caps pinning")
     assert _pin_budget_bytes(reserved=2**30) is None
+
+
+def test_moe_slot_pool_split(tmp_path):
+    import json
+
+    from freetoken.moe.pool_policy import group_layers, load_profile, make_pools, split_slots
+
+    pools = make_pools(group_layers("per-layer", [0] * 4, [1] * 4), num_experts=16, min_num_slots=8, slot_limit=None)
+    assert split_slots(41, pools) == [11, 10, 10, 10]
+    with pytest.raises(ValueError, match="need"):
+        split_slots(31, pools)
+    with pytest.raises(ValueError, match="hold"):
+        split_slots(65, pools)
+    # one pool takes the whole cache, as before pools existed
+    assert split_slots(100, make_pools(group_layers("unified", [0] * 4, [1] * 4), 16, 8, None)) == [100]
+    # unified pools one bank layout each, widest slots first
+    assert group_layers("unified", ["a", "b", "a", "b"], [1, 3, 1, 3]) == [[1, 3], [0, 2]]
+
+    # layer 0 uses two experts, layer 1 has a hot half and a cold half, layers 2 and 3 use all 16 evenly
+    counts = [[5] * 2 + [0] * 14, [3] * 8 + [1] * 8, [1] * 16, [1] * 16]
+    path = tmp_path / "profile.json"
+    path.write_text(json.dumps({"model": "toy", "expert_counts": counts}))
+    assert split_slots(48, pools, load_profile(path, 4, 16)) == [8, 8, 16, 16]
+    with pytest.raises(ValueError, match="needs"):
+        load_profile(path, 4, 32)
+
+
+
+def test_moe_pool_flags_and_decode_batch_clamp():
+    from unittest.mock import patch
+
+    from freetoken.engine.engine import _clamp_decode_batch_to_pools
+    from freetoken.server.args import parse_args
+
+    hf = SimpleNamespace(to_dict=lambda: {"architectures": ["Qwen3MoeForCausalLM"], "torch_dtype": "bfloat16"})
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: hf):
+        args, _ = parse_args([
+            "--model", "/models/anon", "--moe-cache-mem", "1.5GiB",
+            "--moe-cache-scope", "per-layer", "--moe-cache-split", "profile:/p.json",
+        ])
+        assert (args.moe_cache_mem, args.moe_cache_scope, args.moe_cache_split) == (3 << 29, "per-layer", "profile:/p.json")
+        for bad in (["--moe-cache-mem", "12"], ["--moe-cache-mem", "1GiB", "--moe-cache-size", "512"]):
+            with pytest.raises(SystemExit):
+                parse_args(["--model", "/models/anon", *bad])
+
+    # GPU decode needs batch * top_k distinct experts to fit the smallest pool below num_experts
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4),
+        max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4],
+    )
+    _clamp_decode_batch_to_pools(config, [12, 16, 9])
+    assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2])
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 0546fbcd6f..dbf06bcc6c 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -187,7 +187,7 @@ def test_marlin_prefill_matches_dequant_reference():
     cache.materialize_layer(0)
     cache.copy_missing()
     g1, g2 = cache.alphas_for_layer(0)
-    gu_p, gu_s, dn_p, dn_s = cache.bank_views(E)
+    gu_p, gu_s, dn_p, dn_s = cache.stage_views(0, 0)
     out = marlin_fused_experts(
         hidden, gu_p, gu_s, g1, dn_p, dn_s, g2,
         topk_weights, topk_ids, "silu", False,
@@ -214,7 +214,7 @@ def decode(layer_id, experts):
         cache.ensure_experts(layer_id, ids)  # rewrites ids -> slots in place
         cache.copy_missing()
         g1, g2 = cache.alphas_for_slots(layer_id)
-        gu_p, gu_s, dn_p, dn_s = cache.bank_views()
+        gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(layer_id)
         out = marlin_fused_experts(
             hidden, gu_p, gu_s, g1, dn_p, dn_s, g2,
             topk_weights, ids, "silu", False,
@@ -278,7 +278,7 @@ def test_marlin_overlap_prefill_matches_dequant_reference():
     assert int(cache.num_indices.item()) == 2, "stale slot mappings survived the prefetch"
     cache.copy_missing()
     g1, g2 = cache.alphas_for_slots(0)
-    gu_p, gu_s, dn_p, dn_s = cache.bank_views()
+    gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0)
     out = marlin_fused_experts(
         dec_hidden, gu_p, gu_s, g1, dn_p, dn_s, g2,
         dec_weights, ids, "silu", False,
@@ -403,7 +403,7 @@ def decode(layer_id, experts):
         ref = _ref_moe(ref_sources, layer_id, hidden, topk_weights, ids)
         cache.ensure_experts(layer_id, ids)  # rewrites ids -> slots in place
         cache.copy_missing()
-        gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.bank_views()
+        gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.pool_bank_views(layer_id)
         out = fused_experts_decode_nvfp4_marlin(
             hidden, gu_p, gu_s, gu_g, dn_p, dn_s, dn_g, topk_weights, ids, "silu", False
         )
@@ -433,7 +433,7 @@ def test_triton_decode_marlin_matches_baseline_kernel():
     ids = torch.tensor([[1, 6]], dtype=torch.int32, device=device)
     cache.ensure_experts(0, ids)
     cache.copy_missing()
-    banks = cache.bank_views()
+    banks = cache.pool_bank_views(0)
     marlin = fused_experts_decode_nvfp4_marlin(hidden, *banks, topk_weights, ids, "silu", False)
     base = fused_experts_decode_nvfp4_serial(hidden, *banks, topk_weights, ids, "silu", False)
     torch.testing.assert_close(marlin.float(), base.float(), rtol=2e-3, atol=2e-3)
@@ -461,7 +461,7 @@ def test_b12x_decode_matches_dequant_reference():
     cache.ensure_experts(0, ids)
     cache.copy_missing()
     g1, g2 = cache.alphas_for_slots(0)
-    gu_p, gu_s, dn_p, dn_s = cache.bank_views()
+    gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0)
     out = b12x_fused_experts(
         hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, ids, "silu", False
     )
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index cd5f5dc773..d3c8d9d843 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -209,7 +209,7 @@ def fake_fused(
     assert fused_calls[0]["w2_ptr"] == fused_calls[2]["w2_ptr"]
     assert fused_calls[0]["w1_ptr"] != fused_calls[1]["w1_ptr"]
     assert fused_calls[0]["w2_ptr"] != fused_calls[1]["w2_ptr"]
-    prefill_gate_up_buffer, prefill_down_buffer = cache.prefill_bank_buffers
+    prefill_gate_up_buffer, prefill_down_buffer = cache.stage_views(0, 0)
     assert prefill_gate_up_buffer.data_ptr() == cache.bank_caches["gate_up"].data_ptr()
     assert prefill_down_buffer.data_ptr() == cache.bank_caches["down"].data_ptr()
 
@@ -639,6 +639,136 @@ def test_offload_cache_rebuild_resizes_and_preserves_sources():
     assert torch.all(cache.id_of_slot == -1)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_slot_pools_of_two_bank_layouts_keep_their_rows_through_prefill():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    torch.manual_seed(0)
+    L, E, top_k = 4, 8, 4
+    # layers 1 and 3 keep wider down rows (uint8, like Q6_K next to Q4_K); their pools come first
+    wide = (1, 3)
+    sources = {
+        "gate_up": [torch.randn(E, 1024, 256, dtype=torch.bfloat16).pin_memory() for _ in range(L)],
+        "down": [
+            torch.randint(0, 255, (E, 512, 768), dtype=torch.uint8).pin_memory() if layer in wide
+            else torch.randn(E, 512, 256, dtype=torch.bfloat16).pin_memory()
+            for layer in range(L)
+        ],
+    }
+    cache = OffloadMoeCache(
+        num_layers=L, num_experts=E, cache_size=33, device=torch.device("cuda"),
+        prefill_overlap=True, prefill_hit_d2d=True, pool_of_layer=[2, 0, 3, 1], pool_sizes=[6, 5, 12, 10],
+    )
+    cache.set_bank_sources(sources)
+    assert [v[1].dtype for v in map(cache.pool_bank_views, range(L))] == [torch.bfloat16, torch.uint8] * 2
+    # staging buffers take 2 x 8 rows of every bank: all of the wide pools, the first 8 slots of layer 0's pool
+    assert cache._pool_stable_from == [6, 5, 8, 0]
+
+    def check_resident() -> None:
+        torch.cuda.synchronize()
+        for layer in range(L):
+            for e, s in enumerate(cache.slot_for_id[layer].tolist()):
+                if s >= 0:
+                    for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)):
+                        assert torch.equal(view[s].cpu(), src[layer][e]), (layer, e, s)
+
+    def decode(steps: int) -> None:
+        for _ in range(steps):
+            for layer in range(L):
+                cache.ensure_experts(layer, torch.randperm(E)[:top_k].to(torch.int32).view(1, top_k).cuda())
+                cache.copy_missing()
+        check_resident()
+
+    decode(20)
+    cache.zero_inactive_route_slots = True
+    cache.begin_prefill()
+    for layer in range(L):
+        cache.prefetch_prefill_layer(layer)
+        cache.prefetch_prefill_layer(layer + 1)
+        views = cache.wait_prefill_layer(layer)
+        torch.cuda.synchronize()
+        for (src, _), view in zip(cache.banks, views):
+            assert torch.equal(view.cpu(), src[layer])
+        cache.release_prefill_layer(layer)
+    assert cache.prefill_hit_rows > 0
+    for p, stable in enumerate(cache._pool_stable_from):
+        if stable:
+            assert not any(view[0].any() for view in cache.pool_bank_views(cache.pool_of_layer.index(p)))
+    check_resident()
+    decode(20)
+
+    cache.materialize_layer(3)
+    cache.copy_missing()
+    for (src, _), view in zip(cache.banks, cache.stage_views(3, 0)):
+        assert torch.equal(view.cpu(), src[3])
+    check_resident()
+
+    cache.rebuild(32, [5, 5, 10, 12])
+    decode(20)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_slot_pools_split_layers_of_one_bank_layout_by_quant_tag():
+    from freetoken.moe.offload_cache import OffloadMoeCache, bank_layouts
+    from freetoken.moe.pool_policy import group_layers
+
+    _init_tp()
+    L, E = 4, 8
+    # uint8 rows of one width, as GGUF Q4_K and IQ4_NL rows are
+    sources = {
+        "gate_up": [torch.randint(1, 255, (E, 64, 144), dtype=torch.uint8).pin_memory() for _ in range(L)],
+        "down": [torch.randint(1, 255, (E, 32, 144), dtype=torch.uint8).pin_memory() for _ in range(L)],
+    }
+    tags = [("ggml", "Q4_K", "Q4_K"), ("ggml", "IQ4_NL", "IQ4_NL")] * 2
+    assert group_layers("unified", bank_layouts(sources, tags), [1] * L) == [[0, 2], [1, 3]]
+
+    def make(pool_of_layer, pool_sizes):
+        return OffloadMoeCache(
+            num_layers=L, num_experts=E, cache_size=sum(pool_sizes), device=torch.device("cuda"),
+            prefill_overlap=True, pool_of_layer=pool_of_layer, pool_sizes=pool_sizes,
+        )
+
+    with pytest.raises(AssertionError, match="layouts"):
+        make([0] * L, [20]).set_bank_sources(sources, layer_tags=tags)
+    cache = make([0, 1, 0, 1], [10, 10])
+    cache.set_bank_sources(sources, layer_tags=tags)
+    assert cache.bank_caches["down"].shape == (20, 32, 144)
+    # the staging buffers cover pool 0 and slots [0, 6) of pool 1, so both slot 0 rows end up with another quant's bytes
+    cache.zero_inactive_route_slots = True
+    cache.begin_prefill()
+    for layer in range(L):
+        cache.prefetch_prefill_layer(layer)
+        cache.prefetch_prefill_layer(layer + 1)
+        cache.wait_prefill_layer(layer)
+        cache.release_prefill_layer(layer)
+    torch.cuda.synchronize()
+    for layer in (0, 1):
+        for view in cache.pool_bank_views(layer):
+            assert not view[0].any()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_offload_cache_rebuild_frees_the_old_slots_first():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=32, device=torch.device("cuda"),
+        pool_of_layer=[0, 1], pool_sizes=[16, 16],
+    )
+    cache.set_bank_sources({
+        "gate_up": [torch.zeros(4, 1024, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+        "down": [torch.zeros(4, 512, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+    })
+    torch.cuda.synchronize()
+    before = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    cache.rebuild(32, [8, 24])
+    # 96 MiB of slots: holding the old ones while allocating the new ones would double it
+    assert torch.cuda.max_memory_allocated() - before < 8 << 20
+
+
 def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small():
     from freetoken.moe.offload_cache import OffloadMoeCache
 
@@ -654,7 +784,7 @@ def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small():
 
     assert cache.cache_size == 5
     assert cache.prefill_overlap is False
-    assert cache.prefill_bank_buffers == []
+    assert cache._stage_views[0][1] is None  # no room for the second staging buffer
 
 
 def test_offload_cache_rebuild_keeps_overlap_at_boundary():
@@ -690,6 +820,58 @@ def test_offload_cache_validate_rebuild_enforces_marlin_cap_and_floor():
         bf16.validate_rebuild(3)  # below the num_experts floor
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("target", ["gpu", "hybrid"])
+def test_slot_pools_keep_each_layer_in_its_own_slots(target):
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    torch.manual_seed(0)
+    L, E, sizes = 3, 16, [5, 7, 20]
+    # hybrid routes more experts than the smallest pool holds; GPU decode must fit them
+    top_k = 2 if target == "gpu" else 8
+    sources = {
+        "gate_up": [torch.randn(E, 64, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)],
+        "down": [torch.randn(E, 32, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)],
+    }
+
+    cache = OffloadMoeCache(
+        num_layers=L, num_experts=E, cache_size=sum(sizes), device=torch.device("cuda"),
+        decode_target=target, hybrid_max_fetch=E, pool_of_layer=[0, 1, 2], pool_sizes=sizes,
+    )
+    cache.set_bank_sources(sources)
+    for _ in range(50):
+        for layer in range(L):
+            experts = torch.randperm(E)[:top_k].to(torch.int32)
+            slots = experts.cuda()
+            if target == "gpu":
+                cache.ensure_experts(layer, slots)
+            else:
+                cache.ensure_experts_hybrid(layer, slots, torch.rand(top_k, device="cuda"))
+            cache.copy_missing()
+            resident = [(e, s) for e, s in zip(experts.tolist(), slots.tolist()) if s >= 0]
+            assert len({s for _, s in resident}) == len(resident)
+            for e, s in resident:
+                assert s < sizes[layer]
+                for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)):
+                    assert torch.equal(view[s].cpu(), src[layer][e])
+
+    # with several pools a whole-layer prefill stages slots [0, E) without registering them in any pool
+    cache.materialize_layer(1)
+    cache.copy_missing()
+    for (src, _), view in zip(cache.banks, cache.stage_views(1, 0)):
+        assert torch.equal(view.cpu(), src[1])
+    slot_for_id, id_of_slot = cache.slot_for_id.cpu(), cache.id_of_slot.cpu()
+    for layer in range(L):
+        for e, s in enumerate(slot_for_id[layer].tolist()):
+            if s >= 0:
+                assert cache.pool_bases[layer] + s >= E
+                assert id_of_slot[cache.pool_bases[layer] + s] == layer * E + e
+
+    cache.rebuild(40, [8, 8, 24])
+    assert [v[0].shape[0] for v in map(cache.pool_bank_views, range(L))] == [8, 8, 24]
+
+
 def _make_split_cache(num_layers=2, locked=(1,), prefill_overlap=False, device="cpu"):
     """A [gate_up, down] bf16 cache with the given layers LOCKED (rest pinned)."""
     from freetoken.moe.host_banks import HostResidency

From 0bca88f70c4909624f7f3c3e3a4fe6b32d183d56 Mon Sep 17 00:00:00 2001
From: LZ_QAQ <32613092+LZ-QWQ@users.noreply.github.com>
Date: Thu, 15 Oct 2026 17:21:32 +0800
Subject: [PATCH 097/249] fix(rocm): reuse staged moe alignment on Triton 3.8
 (#691)

* fix(rocm): reuse staged moe alignment below triton 3.9

* fix(rocm): restrict moe alignment workaround to Triton 3.8
---
 python/freetoken/moe/fused.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py
index a90aa4c792..5c34fb33ca 100644
--- a/python/freetoken/moe/fused.py
+++ b/python/freetoken/moe/fused.py
@@ -84,7 +84,18 @@ def moe_align_block_size(
     - The padding ensures that the total number of tokens is now divisible
         by block_size for proper block matrix operations.
     """
-    from freetoken.kernel.backend import is_sgl_kernel_installed
+    from freetoken.kernel.backend import is_rocm, is_sgl_kernel_installed
+
+    if is_rocm():
+        import triton
+        from packaging.version import Version
+
+        # FIXME: remove once Triton 3.8 is unsupported. Its AMD backend folds compares on tl.histogram results to false.
+        # https://github.com/triton-lang/triton/pull/11246
+        if Version(triton.__version__).release[:2] == (3, 8):
+            from freetoken.kernel import moe_align_block_size_triton
+
+            return moe_align_block_size_triton(topk_ids, block_size, num_experts)
 
     if not is_sgl_kernel_installed():
         from freetoken.kernel.triton.moe_align import (

From f3c9982a88b860eff28a3bfb53cebece4f7fef68 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= <Avicennasis@gmail.com>
Date: Tue, 1 Sep 2026 15:36:05 -0400
Subject: [PATCH 098/249] fix(memory): honor cgroup memory limits when sizing
 expert banks

The expert-bank loader, the CPU bandwidth bench and the loader benchmark
all sized host-RAM work from /proc/meminfo MemAvailable alone.  Inside a
container started with --memory=N that figure is the host's, not the
cgroup's: the parallel expert reader was admitted against tens of GiB of
"available" RAM while the cgroup could hold a fraction of it, and the
process was OOM-killed mid-load.  benchbw.py carried its own cgroup-v2-only
parser; the loader had none.

Rule: effective available memory = min(host MemAvailable, remaining finite
cgroup budget).  The cgroup budget is the tightest finite limit-minus-usage
along the process's cgroup ancestry (v2 memory.max / memory.current with
the `max` sentinel; v1 memory.limit_in_bytes / memory.usage_in_bytes with
the unlimited sentinel as a fallback), resolved from /proc/self/cgroup
with bounded reads and a bounded ancestor walk.  Absent procfs or no
visible memory controller yields None (unknown), so non-Linux hosts keep
their previous best-effort behaviour and are never told they have 0 bytes;
a present but malformed or unreadable control file fails closed to 0.
The probe runs before any bank or buffer allocation: admission is the only
decision point, because once the parallel reader's whole-shard anonymous
buffers exist there is nothing reclaimable left to give back.

freetoken/memory.py is stdlib-only (no torch import) so it can be loaded
by path inside a slim container for diagnosis.  Every production
MemAvailable reader now goes through it: moe/expert_banks
._host_ram_fits_parallel, moe/benchbw._available_ram_bytes (the private
_cgroup_mem_headroom parser is removed so there is one cgroup parser; the
bench's clamping and synthetic-bank sizing are otherwise unchanged), and
benchmarks/bench_load_weight_generic.MemSampler.

Tests: tests/test_memory.py drives the resolver against fake /proc and
cgroup trees under tmp_path -- unlimited, finite, nested ancestor-min,
namespaced root, usage above limit, malformed, oversized, permission
denied, v1 fallback and sentinel, absent files, and the hybrid
`1:net_cls:/` + `0::/...` membership layout systemd and Docker both
produce.  tests/moe/test_expert_bank_memory_admission.py checks that the
auto-admission path drops to the serial build when the fake cgroup budget
(299 B) cannot hold the banks plus one shard transient (300 B) although
MemAvailable is large, that an explicit --expert-load override bypasses
the probe, and that the bench receives a finite budget unchanged while
keeping its legacy 8 GiB fallback only for an unknown (None) probe.

Validated on CPU-only Linux (Ubuntu 24.04, Docker 29.1, cgroup v2; no GPU
is involved in this path): 47 tests pass under python and python -O; the
full suite has the same two pre-existing failures as origin/main
(tests/models/test_muse_glimmer.py) and no new ones; python:3.12-slim with
--memory=2g / --memory=6g / no limit reports effective =
memory.max - memory.current (2136084480 / 6432010240) and host
MemAvailable respectively.

Mount-namespace hardening of the cgroup lookup is deliberately out of
scope for this change.
---
 benchmarks/bench_load_weight_generic.py       |  16 +-
 python/freetoken/memory.py                    | 294 +++++++++++++++
 python/freetoken/moe/benchbw.py               |  26 +-
 python/freetoken/moe/expert_banks.py          |  15 +-
 .../moe/test_expert_bank_memory_admission.py  | 108 ++++++
 tests/test_memory.py                          | 346 ++++++++++++++++++
 6 files changed, 763 insertions(+), 42 deletions(-)
 create mode 100644 python/freetoken/memory.py
 create mode 100644 tests/moe/test_expert_bank_memory_admission.py
 create mode 100644 tests/test_memory.py

diff --git a/benchmarks/bench_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py
index 4e4c5fe012..dfc19a9acd 100644
--- a/benchmarks/bench_load_weight_generic.py
+++ b/benchmarks/bench_load_weight_generic.py
@@ -41,6 +41,8 @@
 
 import torch
 
+from freetoken.memory import effective_memory_available
+
 GiB = float(1 << 30)
 ALL_MODES = ("parallel", "ftw")
 
@@ -54,14 +56,6 @@ def _default_ftw_dir(model_path: str) -> str:
 
 
 # ---------------- memory sampling + checksum ----------------
-def _meminfo_available_bytes() -> int:
-    with open("/proc/meminfo") as f:
-        for line in f:
-            if line.startswith("MemAvailable:"):
-                return int(line.split()[1]) * 1024
-    return 0
-
-
 def _status_kb(key: str) -> int:
     with open("/proc/self/status") as f:
         for line in f:
@@ -71,7 +65,7 @@ def _status_kb(key: str) -> int:
 
 
 class MemSampler(threading.Thread):
-    """Background sampler: peak process RSS + system MemAvailable low-water-mark."""
+    """Background sampler: peak RSS + effective available-memory low-water-mark."""
 
     def __init__(self, interval: float = 0.2):
         super().__init__(daemon=True)
@@ -82,7 +76,9 @@ def __init__(self, interval: float = 0.2):
 
     def run(self) -> None:
         while not self._stop_evt.wait(self.interval):
-            self.min_avail = min(self.min_avail, _meminfo_available_bytes())
+            available = effective_memory_available()
+            if available is not None:
+                self.min_avail = min(self.min_avail, available)
             self.max_rss = max(self.max_rss, _status_kb("VmRSS:"))
 
     def stop(self) -> None:
diff --git a/python/freetoken/memory.py b/python/freetoken/memory.py
new file mode 100644
index 0000000000..32832aa9a6
--- /dev/null
+++ b/python/freetoken/memory.py
@@ -0,0 +1,294 @@
+"""Conservative host-memory admission for bare-metal and cgrouped processes.
+
+The kernel's host-wide ``MemAvailable`` figure is not an allocation budget for a
+container.  This module combines it with the process's tightest finite cgroup
+budget and deliberately keeps the probing code independent of torch and the
+rest of the runtime so startup checks and diagnostic tools can share it.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+
+# Kernel pseudo-files are small, but bound every read and hierarchy walk.  A
+# container may expose a synthetic or partially mounted cgroup tree; admission
+# must not turn an unexpected file into an unbounded read or traversal.
+_CGROUP_VALUE_MAX_BYTES = 128
+_CGROUP_MEMBERSHIP_MAX_BYTES = 64 * 1024
+_CGROUP_MAX_ANCESTORS = 64
+_MEMINFO_MAX_BYTES = 256 * 1024
+_UINT64_MAX = (1 << 64) - 1
+
+# Cgroup v1 represents "unlimited" with a page-aligned value just below the
+# signed 64-bit maximum (commonly 9223372036854771712).  No supported host can
+# offer an exbibyte of usable RAM, so this threshold safely covers the kernel's
+# sentinel variants without mistaking a practical finite limit for unlimited.
+_CGROUP_V1_UNLIMITED_MIN = 1 << 60
+
+
+def _read_bounded_ascii(path: Path, max_bytes: int) -> tuple[bool, str | None]:
+    """Return ``(present, exact_text)`` for a small kernel pseudo-file.
+
+    A missing path means that source/controller is unavailable.  Other I/O
+    errors are authoritative probe failures and raise so the public resolver can
+    fail closed instead of silently treating the process as unconstrained.
+    ``None`` text means the present file was oversized or not ASCII.
+    """
+    path = Path(path)
+    try:
+        with path.open("rb") as stream:
+            raw = stream.read(max_bytes + 1)
+    except (FileNotFoundError, NotADirectoryError):
+        return False, None
+    except OSError as exc:
+        raise ValueError(f"cannot read memory control file {path}: {exc}") from exc
+    if len(raw) > max_bytes:
+        return True, None
+    try:
+        return True, raw.decode("ascii")
+    except UnicodeDecodeError:
+        return True, None
+
+
+def _parse_u64_counter(text: str | None) -> int | None:
+    """Parse one non-negative decimal counter with at most one final newline."""
+    match = re.fullmatch(r"(\d+)\n?", text or "")
+    if match is None:
+        return None
+    value = int(match.group(1))
+    return value if value <= _UINT64_MAX else None
+
+
+def _host_memory_available(meminfo_path: Path) -> int | None:
+    """Read host-wide ``MemAvailable`` bytes, or ``None`` when /proc is absent."""
+    present, text = _read_bounded_ascii(meminfo_path, _MEMINFO_MAX_BYTES)
+    if not present:
+        return None
+    if text is None:
+        raise ValueError(f"malformed or oversized memory information: {meminfo_path}")
+    match = re.search(
+        r"^MemAvailable:[ \t]+(\d+)[ \t]+kB[ \t]*(?:\n|$)",
+        text,
+        flags=re.MULTILINE,
+    )
+    if match is None:
+        raise ValueError(f"MemAvailable is unavailable or malformed: {meminfo_path}")
+    kib = int(match.group(1))
+    if kib > _UINT64_MAX // 1024:
+        raise ValueError(f"MemAvailable overflows a byte counter: {meminfo_path}")
+    return kib * 1024
+
+
+def _cgroup_pair_remaining(
+    directory: Path,
+    limit_name: str,
+    current_name: str,
+    *,
+    v1: bool = False,
+) -> tuple[bool, int | None]:
+    """Return ``(limit_present, finite_remaining_or_none)`` for one cgroup."""
+    limit_path = Path(directory) / limit_name
+    current_path = Path(directory) / current_name
+    limit_present, limit_text = _read_bounded_ascii(limit_path, _CGROUP_VALUE_MAX_BYTES)
+    if not limit_present:
+        return False, None
+    if not v1 and limit_text in ("max", "max\n"):
+        return True, None
+
+    limit = _parse_u64_counter(limit_text)
+    if limit is None:
+        raise ValueError(f"malformed cgroup memory limit: {limit_path}")
+    if v1 and limit >= _CGROUP_V1_UNLIMITED_MIN:
+        return True, None
+
+    current_present, current_text = _read_bounded_ascii(
+        current_path, _CGROUP_VALUE_MAX_BYTES
+    )
+    if not current_present:
+        raise ValueError(f"cgroup memory usage is unavailable: {current_path}")
+    current = _parse_u64_counter(current_text)
+    if current is None:
+        raise ValueError(f"malformed cgroup memory usage: {current_path}")
+    return True, max(0, limit - current)
+
+
+def _read_cgroup_memberships(proc_cgroup_path: Path) -> tuple[str | None, str | None]:
+    """Return the v2 and v1-memory paths from bounded ``/proc/self/cgroup``.
+
+    Hybrid hosts list leftover v1 hierarchies (``1:net_cls:/``) above the unified
+    ``0::/...`` line; non-memory v1 controllers are simply skipped.
+    """
+    present, text = _read_bounded_ascii(proc_cgroup_path, _CGROUP_MEMBERSHIP_MAX_BYTES)
+    if not present:
+        return None, None
+    if text is None:
+        raise ValueError(
+            f"malformed or oversized cgroup membership: {proc_cgroup_path}"
+        )
+
+    v2_path = None
+    v1_path = None
+    for line in text.splitlines():
+        fields = line.split(":", 2)
+        if len(fields) != 3:
+            continue
+        hierarchy, controllers, member_path = fields
+        if not member_path.startswith("/"):
+            continue
+        # Real cgroup membership paths are absolute and contain no dot
+        # components.  Reject lexical escapes in synthetic/malformed proc data.
+        parts = Path(member_path).parts[1:]
+        if any(part in ("", ".", "..") for part in parts):
+            continue
+        if hierarchy == "0" and not controllers:
+            v2_path = member_path
+        elif "memory" in controllers.split(","):
+            v1_path = member_path
+    return v2_path, v1_path
+
+
+def _control_file_exists(path: Path) -> bool:
+    """Existence probe that fails closed on errors other than a missing path."""
+    path = Path(path)
+    try:
+        path.stat()
+        return True
+    except (FileNotFoundError, NotADirectoryError):
+        return False
+    except OSError as exc:
+        raise ValueError(f"cannot inspect cgroup control file {path}: {exc}") from exc
+
+
+def _cgroup_control_directory(
+    root: Path, member_path: str | None, limit_name: str
+) -> Path | None:
+    """Map a proc membership path into a mounted, possibly namespaced tree."""
+    root = Path(root)
+    if member_path:
+        parts = Path(member_path).parts[1:]
+        if len(parts) >= _CGROUP_MAX_ANCESTORS:
+            raise ValueError(
+                f"cgroup membership nesting exceeds {_CGROUP_MAX_ANCESTORS - 1} levels"
+            )
+        candidate = root.joinpath(*parts)
+        # A leaf can omit controller files when only an ancestor has the memory
+        # controller enabled.  Find the nearest visible ancestor without a
+        # filesystem scan or a walk outside the mounted hierarchy.
+        for _ in range(_CGROUP_MAX_ANCESTORS):
+            if _control_file_exists(candidate / limit_name):
+                return candidate
+            if candidate == root:
+                break
+            parent = candidate.parent
+            if parent == candidate or (parent != root and root not in parent.parents):
+                break
+            candidate = parent
+
+    # Container runtimes may mount the process's own subgroup as the hierarchy
+    # root while /proc still reports a host-side membership path.
+    if _control_file_exists(root / limit_name):
+        return root
+    return None
+
+
+def _cgroup_hierarchy_remaining(
+    root: Path,
+    member_path: str | None,
+    limit_name: str,
+    current_name: str,
+    *,
+    v1: bool = False,
+) -> tuple[bool, int | None]:
+    """Return ``(hierarchy_seen, tightest finite ancestor headroom)``."""
+    root = Path(root)
+    current = _cgroup_control_directory(root, member_path, limit_name)
+    if current is None:
+        return False, None
+
+    seen = False
+    remaining = None
+    for _ in range(_CGROUP_MAX_ANCESTORS):
+        present, candidate = _cgroup_pair_remaining(
+            current, limit_name, current_name, v1=v1
+        )
+        seen = seen or present
+        if candidate is not None:
+            remaining = candidate if remaining is None else min(remaining, candidate)
+        if current == root:
+            break
+        parent = current.parent
+        if parent == current or (parent != root and root not in parent.parents):
+            break
+        current = parent
+    return seen, remaining
+
+
+def _cgroup_memory_remaining(
+    cgroup_root: Path = Path("/sys/fs/cgroup"),
+    proc_cgroup_path: Path = Path("/proc/self/cgroup"),
+) -> int | None:
+    """Return this process's tightest finite cgroup memory headroom.
+
+    Cgroup v2 is authoritative whenever its memory controller is visible.
+    ``max`` is a known-unlimited value.  Only when v2 is absent do we try the
+    conventional, bounded v1 controller roots; no filesystem scan is used.
+    Malformed, overflowing, or unreadable visible controls raise ``ValueError``.
+    """
+    root = Path(cgroup_root)
+    v2_path, v1_path = _read_cgroup_memberships(proc_cgroup_path)
+    v2_seen, remaining = _cgroup_hierarchy_remaining(
+        root, v2_path, "memory.max", "memory.current"
+    )
+    if v2_seen:
+        return remaining
+
+    for v1_root in (root / "memory", root):
+        v1_seen, remaining = _cgroup_hierarchy_remaining(
+            v1_root,
+            v1_path,
+            "memory.limit_in_bytes",
+            "memory.usage_in_bytes",
+            v1=True,
+        )
+        if v1_seen:
+            return remaining
+    return None
+
+
+def effective_memory_available(
+    *,
+    meminfo_path: Path = Path("/proc/meminfo"),
+    cgroup_root: Path = Path("/sys/fs/cgroup"),
+    proc_cgroup_path: Path = Path("/proc/self/cgroup"),
+) -> int | None:
+    """Return the effective available host-memory budget in bytes.
+
+    The result is ``min(MemAvailable, tightest finite cgroup headroom)`` when
+    both measurements exist, either known bound when only one exists, and
+    ``None`` only when neither source is available -- no procfs, a non-Linux
+    host, or a cgroup tree without a visible memory controller.  ``None`` means
+    *unknown* and is never conflated with ``0``; callers keep their historical
+    best-effort behaviour on it.  A known zero remains zero.
+
+    A *present* but malformed, overflowing, or permission-denied control file
+    is an authoritative signal that the budget cannot be trusted, so it fails
+    closed to ``0`` rather than being mistaken for an unlimited controller.
+    """
+    try:
+        host_available = _host_memory_available(Path(meminfo_path))
+        cgroup_remaining = _cgroup_memory_remaining(
+            Path(cgroup_root), Path(proc_cgroup_path)
+        )
+    except ValueError:
+        return 0
+
+    if host_available is None:
+        return cgroup_remaining
+    if cgroup_remaining is None:
+        return host_available
+    return min(host_available, cgroup_remaining)
+
+
+__all__ = ["effective_memory_available"]
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index cd260cb2f6..2f14e50525 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -56,6 +56,7 @@
     single_gpu_arg,
 )
 from freetoken.kernel.pinned import alloc_pinned_tensor
+from freetoken.memory import effective_memory_available
 from freetoken.moe.bench_profile import PROFILE_VERSION
 from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity
 from freetoken.utils import init_logger
@@ -156,29 +157,10 @@ def default_out_path(gpu_uuid: str | None = None) -> str:
     return default_profile_path(gpu_uuid)
 
 
-def _cgroup_mem_headroom() -> int | None:
-    """Free bytes under this process's cgroup v2 memory limit, or None if unlimited."""
-    try:
-        with open("/sys/fs/cgroup/memory.max") as f:
-            raw = f.read().strip()
-        if raw == "max":
-            return None
-        limit = int(raw)
-        with open("/sys/fs/cgroup/memory.current") as f:
-            used = int(f.read().strip())
-        return max(0, limit - used)
-    except (OSError, ValueError):
-        return None
-
-
 def _available_ram_bytes() -> int:
-    """Free host RAM, clamped to the cgroup memory limit so a container isn't over-estimated."""
-    try:
-        host = os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")
-    except (ValueError, OSError, AttributeError):
-        host = 8 << 30
-    cg = _cgroup_mem_headroom()
-    return min(host, cg) if cg is not None else host
+    """Effective free RAM, retaining the legacy fallback only when probing is absent."""
+    available = effective_memory_available()
+    return (8 << 30) if available is None else available
 
 
 # ============================== ceilings (hardware) ==============================
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 5521c34e1f..2c76354724 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -16,6 +16,7 @@
 import torch
 
 from freetoken.layers.quantization import QuantKind
+from freetoken.memory import effective_memory_available
 from freetoken.utils import init_logger
 
 from .host_banks import alloc_layer_banks
@@ -217,16 +218,10 @@ def _host_ram_fits_parallel(model_path: str) -> bool:
     """Best-effort: can free host RAM hold the expert banks plus the parallel reader's one
     extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True,
     i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the
-    largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure."""
-    avail = None
-    try:
-        with open("/proc/meminfo") as f:
-            for line in f:
-                if line.startswith("MemAvailable:"):
-                    avail = int(line.split()[1]) * 1024
-                    break
-    except OSError:
-        pass
+    largest shard. Uses the effective available memory (host ``MemAvailable``, which counts
+    reclaimable cache, clamped by the process's tightest finite cgroup budget) -- the
+    OOM-relevant figure on bare metal and inside a ``--memory``-limited container alike."""
+    avail = effective_memory_available()
     if avail is None:
         return True
     try:  # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards
diff --git a/tests/moe/test_expert_bank_memory_admission.py b/tests/moe/test_expert_bank_memory_admission.py
new file mode 100644
index 0000000000..221814212d
--- /dev/null
+++ b/tests/moe/test_expert_bank_memory_admission.py
@@ -0,0 +1,108 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.memory import effective_memory_available
+
+
+def test_loader_auto_admission_consumes_effective_memory(
+    tmp_path, monkeypatch: pytest.MonkeyPatch
+):
+    """Regression: production auto-load must use the shared effective budget."""
+    import freetoken.checkpoint.ftw as ftw
+    import freetoken.models.weight as model_weight
+    import freetoken.moe.expert_banks as expert_banks
+    import freetoken.utils.hf as hf
+
+    # The bank estimate is 200 bytes plus a 100-byte parallel-reader transient.
+    # Host MemAvailable is large, but the fake cgroup has only 299 bytes left.
+    (tmp_path / "a.safetensors").write_bytes(b"a" * 100)
+    (tmp_path / "b.safetensors").write_bytes(b"b" * 100)
+    meminfo = tmp_path / "meminfo"
+    meminfo.write_text("MemAvailable: 1024 kB\n", encoding="ascii")
+    proc_cgroup = tmp_path / "proc-self-cgroup"
+    proc_cgroup.write_text("0::/\n", encoding="ascii")
+    cgroup_root = tmp_path / "cgroup"
+    cgroup_root.mkdir()
+    (cgroup_root / "memory.max").write_text("1000\n", encoding="ascii")
+    (cgroup_root / "memory.current").write_text("701\n", encoding="ascii")
+
+    def effective_fixture():
+        return effective_memory_available(
+            meminfo_path=meminfo,
+            cgroup_root=cgroup_root,
+            proc_cgroup_path=proc_cgroup,
+        )
+
+    monkeypatch.setattr(expert_banks, "_PARALLEL_READER_SUPPORTED", True)
+    monkeypatch.setattr(expert_banks, "effective_memory_available", effective_fixture)
+    monkeypatch.setattr(model_weight, "experts_scattered", lambda _path: True)
+    monkeypatch.setattr(hf, "download_hf_weight", lambda path: path)
+    monkeypatch.setattr(ftw, "is_ftw_checkpoint", lambda _path: False)
+
+    chosen = []
+
+    def build(*args, **kwargs):
+        chosen.append(args[5])  # _legacy_expert_banks(..., dummy, parallel, ...)
+        return expert_banks.ExpertBanks("bf16", {})
+
+    monkeypatch.setattr(expert_banks, "_legacy_expert_banks", build)
+    config = SimpleNamespace(num_moe_layers=1, expert_quant="none")
+
+    expert_banks.load_expert_banks(
+        str(tmp_path), config, device=torch.device("cpu"), dtype=torch.bfloat16
+    )
+
+    assert chosen == [False]
+
+
+def test_explicit_parallel_loader_override_bypasses_auto_admission(
+    tmp_path, monkeypatch: pytest.MonkeyPatch
+):
+    import freetoken.checkpoint.ftw as ftw
+    import freetoken.moe.expert_banks as expert_banks
+
+    monkeypatch.setattr(expert_banks, "_PARALLEL_READER_SUPPORTED", True)
+    monkeypatch.setattr(ftw, "is_ftw_checkpoint", lambda _path: False)
+    monkeypatch.setattr(
+        expert_banks,
+        "effective_memory_available",
+        lambda: (_ for _ in ()).throw(AssertionError("auto admission should not run")),
+    )
+    chosen = []
+
+    def build(*args, **kwargs):
+        chosen.append(args[5])
+        return expert_banks.ExpertBanks("bf16", {})
+
+    monkeypatch.setattr(expert_banks, "_legacy_expert_banks", build)
+    config = SimpleNamespace(num_moe_layers=1, expert_quant="none")
+
+    expert_banks.load_expert_banks(
+        str(tmp_path),
+        config,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        parallel=True,
+    )
+
+    assert chosen == [True]
+
+
+def test_benchmark_ram_estimate_uses_effective_headroom(monkeypatch):
+    import freetoken.moe.benchbw as benchbw
+
+    # A finite cgroup budget (here 3 GiB) must reach the benchmark unchanged
+    # instead of the host-wide figure the old sysconf probe reported.
+    monkeypatch.setattr(benchbw, "effective_memory_available", lambda: 3 << 30)
+
+    assert benchbw._available_ram_bytes() == 3 << 30
+
+
+def test_benchmark_legacy_fallback_is_only_for_unknown_probe(monkeypatch):
+    import freetoken.moe.benchbw as benchbw
+
+    monkeypatch.setattr(benchbw, "effective_memory_available", lambda: None)
+
+    assert benchbw._available_ram_bytes() == 8 << 30
diff --git a/tests/test_memory.py b/tests/test_memory.py
new file mode 100644
index 0000000000..c136b88e88
--- /dev/null
+++ b/tests/test_memory.py
@@ -0,0 +1,346 @@
+from pathlib import Path
+
+import pytest
+
+from freetoken.memory import (
+    _CGROUP_MAX_ANCESTORS,
+    _CGROUP_V1_UNLIMITED_MIN,
+    _UINT64_MAX,
+    _cgroup_memory_remaining,
+    effective_memory_available,
+)
+
+
+class FakeMemoryFiles:
+    """Pure fake /proc and cgroup trees; never touches the runner's hierarchy."""
+
+    def __init__(self, base: Path):
+        self.root = base / "cgroup"
+        self.root.mkdir()
+        self.proc_cgroup = base / "proc-self-cgroup"
+        self.meminfo = base / "meminfo"
+        self.meminfo.write_text(
+            "MemTotal: 32768 kB\nMemAvailable: 16384 kB\n", encoding="ascii"
+        )
+
+    def membership(self, *lines: str) -> None:
+        self.proc_cgroup.write_text("\n".join(lines) + "\n", encoding="ascii")
+
+    @staticmethod
+    def pair(
+        directory: Path,
+        limit_name: str,
+        current_name: str,
+        limit: object,
+        current: object,
+    ) -> None:
+        directory.mkdir(parents=True, exist_ok=True)
+        (directory / limit_name).write_text(f"{limit}\n", encoding="ascii")
+        (directory / current_name).write_text(f"{current}\n", encoding="ascii")
+
+    def v2(self, directory: Path, limit: object, current: object) -> None:
+        self.pair(directory, "memory.max", "memory.current", limit, current)
+
+    def v1(self, directory: Path, limit: object, current: object) -> None:
+        self.pair(
+            directory,
+            "memory.limit_in_bytes",
+            "memory.usage_in_bytes",
+            limit,
+            current,
+        )
+
+    def remaining(self) -> int | None:
+        return _cgroup_memory_remaining(self.root, self.proc_cgroup)
+
+    def effective(self) -> int | None:
+        return effective_memory_available(
+            meminfo_path=self.meminfo,
+            cgroup_root=self.root,
+            proc_cgroup_path=self.proc_cgroup,
+        )
+
+
+@pytest.fixture
+def memory_files(tmp_path: Path) -> FakeMemoryFiles:
+    return FakeMemoryFiles(tmp_path)
+
+
+def test_v2_finite_budget_clamps_host_memavailable(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/tenant/job")
+    memory_files.v2(memory_files.root / "tenant" / "job", 12_000_000, 3_000_000)
+
+    assert memory_files.remaining() == 9_000_000
+    assert memory_files.effective() == 9_000_000
+
+
+def test_host_memavailable_wins_when_lower(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/tenant/job")
+    memory_files.v2(memory_files.root / "tenant" / "job", 100_000_000, 1)
+
+    assert memory_files.effective() == 16_384 * 1024
+
+
+def test_nested_v2_uses_tightest_finite_ancestor(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/tenant/job")
+    memory_files.v2(memory_files.root, "max", "not-read-for-unlimited")
+    memory_files.v2(memory_files.root / "tenant", 20_000_000, 14_000_000)
+    memory_files.v2(memory_files.root / "tenant" / "job", 12_000_000, 3_000_000)
+
+    assert memory_files.remaining() == 6_000_000
+
+
+def test_unlimited_v2_leaf_still_honors_finite_parent(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/tenant/job")
+    memory_files.v2(memory_files.root / "tenant", 9_000_000, 4_000_000)
+    memory_files.v2(
+        memory_files.root / "tenant" / "job", "max", "not-read-for-unlimited"
+    )
+
+    assert memory_files.remaining() == 5_000_000
+
+
+def test_namespaced_v2_mount_uses_root_controls(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/host/tenant/job")
+    memory_files.v2(memory_files.root, 10_000_000, 4_000_000)
+
+    assert memory_files.remaining() == 6_000_000
+
+
+def test_v2_max_is_known_unlimited(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, "max", "malformed-but-irrelevant")
+
+    assert memory_files.remaining() is None
+    assert memory_files.effective() == 16_384 * 1024
+
+
+def test_counter_accepts_no_trailing_newline(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/")
+    (memory_files.root / "memory.max").write_text("10000", encoding="ascii")
+    (memory_files.root / "memory.current").write_text("4000", encoding="ascii")
+
+    assert memory_files.remaining() == 6_000
+
+
+@pytest.mark.parametrize("current", [10_000, 10_001])
+def test_v2_current_at_or_above_limit_is_zero(
+    memory_files: FakeMemoryFiles, current: int
+):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, 10_000, current)
+
+    assert memory_files.remaining() == 0
+    assert memory_files.effective() == 0
+
+
+@pytest.mark.parametrize(
+    "value",
+    ["", "-1", "+1", "1.5", "garbage", " 1", "1 ", "\t1", "1\n", "1\r", "max "],
+)
+def test_malformed_v2_limit_fails_closed(memory_files: FakeMemoryFiles, value: str):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, value, 1)
+
+    with pytest.raises(ValueError, match="malformed cgroup memory limit"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+@pytest.mark.parametrize("value", ["-1", "+1", "1.5", "unknown", "1 "])
+def test_malformed_v2_usage_fails_closed(memory_files: FakeMemoryFiles, value: str):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, 10_000, value)
+
+    with pytest.raises(ValueError, match="malformed cgroup memory usage"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_overflow_and_oversized_controls_fail_closed(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, _UINT64_MAX + 1, 1)
+    assert memory_files.effective() == 0
+
+    memory_files.v2(memory_files.root, "9" * 256, 1)
+    with pytest.raises(ValueError, match="malformed cgroup memory limit"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_finite_limit_without_usage_fails_closed(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/")
+    (memory_files.root / "memory.max").write_text("10000\n", encoding="ascii")
+
+    with pytest.raises(ValueError, match="usage is unavailable"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_permission_error_on_visible_counter_fails_closed(
+    memory_files: FakeMemoryFiles, monkeypatch: pytest.MonkeyPatch
+):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, 10_000, 1_000)
+    denied = memory_files.root / "memory.current"
+    original_open = Path.open
+
+    def deny_current(path: Path, *args, **kwargs):
+        if path == denied:
+            raise PermissionError("fixture denied")
+        return original_open(path, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "open", deny_current)
+    with pytest.raises(ValueError, match="cannot read memory control file"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_excessive_membership_nesting_is_bounded(memory_files: FakeMemoryFiles):
+    deep = "/" + "/".join(f"level-{i}" for i in range(_CGROUP_MAX_ANCESTORS))
+    memory_files.membership(f"0::{deep}")
+    memory_files.v2(memory_files.root, "max", 1)
+
+    with pytest.raises(ValueError, match="membership nesting exceeds"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_oversized_membership_read_fails_closed(memory_files: FakeMemoryFiles):
+    memory_files.proc_cgroup.write_text("0::/" + "x" * (65 * 1024), encoding="ascii")
+
+    with pytest.raises(ValueError, match="oversized cgroup membership"):
+        memory_files.remaining()
+    assert memory_files.effective() == 0
+
+
+def test_malformed_leaf_is_not_masked_by_valid_parent(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/tenant/job")
+    memory_files.v2(memory_files.root / "tenant", 8_000_000, 3_000_000)
+    memory_files.v2(memory_files.root / "tenant" / "job", "bad", 1)
+
+    assert memory_files.effective() == 0
+
+
+def test_v2_is_authoritative_over_stale_v1(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/", "7:memory:/legacy")
+    memory_files.v2(memory_files.root, "max", 1)
+    memory_files.v1(memory_files.root / "memory" / "legacy", 4_000_000, 3_000_000)
+
+    assert memory_files.remaining() is None
+    assert memory_files.effective() == 16_384 * 1024
+
+
+def test_missing_v2_controller_uses_bounded_v1(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/unified", "7:memory:/legacy")
+    memory_files.v1(memory_files.root / "memory" / "legacy", 8_000_000, 3_000_000)
+
+    assert memory_files.remaining() == 5_000_000
+    assert memory_files.effective() == 5_000_000
+
+
+def test_v1_honors_tighter_parent(memory_files: FakeMemoryFiles):
+    memory_files.membership("7:memory:/tenant/job")
+    v1_root = memory_files.root / "memory"
+    memory_files.v1(v1_root / "tenant", 10_000_000, 7_000_000)
+    memory_files.v1(v1_root / "tenant" / "job", 8_000_000, 2_000_000)
+
+    assert memory_files.remaining() == 3_000_000
+
+
+def test_v1_unlimited_sentinel_is_known_unlimited(memory_files: FakeMemoryFiles):
+    memory_files.membership("7:memory:/")
+    memory_files.v1(
+        memory_files.root / "memory",
+        _CGROUP_V1_UNLIMITED_MIN,
+        "malformed-but-irrelevant",
+    )
+
+    assert memory_files.remaining() is None
+    assert memory_files.effective() == 16_384 * 1024
+
+
+def test_v1_current_above_limit_is_zero(memory_files: FakeMemoryFiles):
+    memory_files.membership("7:memory:/")
+    memory_files.v1(memory_files.root / "memory", 10_000, 20_000)
+
+    assert memory_files.remaining() == 0
+    assert memory_files.effective() == 0
+
+
+def test_malformed_membership_can_use_namespaced_root(memory_files: FakeMemoryFiles):
+    memory_files.v2(memory_files.root, 9_000_000, 2_000_000)
+    for proc_text in ("", "not:a:valid:line\n"):
+        memory_files.proc_cgroup.write_text(proc_text, encoding="ascii")
+        assert memory_files.remaining() == 7_000_000
+
+
+def test_absent_cgroup_files_leave_host_measurement(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/missing")
+
+    assert memory_files.remaining() is None
+    assert memory_files.effective() == 16_384 * 1024
+
+
+def test_finite_cgroup_is_bound_when_meminfo_is_absent(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/")
+    memory_files.v2(memory_files.root, 7_000_000, 2_000_000)
+
+    assert (
+        effective_memory_available(
+            meminfo_path=memory_files.meminfo.with_name("missing"),
+            cgroup_root=memory_files.root,
+            proc_cgroup_path=memory_files.proc_cgroup,
+        )
+        == 5_000_000
+    )
+
+
+def test_malformed_or_overflowing_meminfo_fails_closed(memory_files: FakeMemoryFiles):
+    memory_files.membership("0::/missing")
+    for available in ("not-a-number", str(_UINT64_MAX // 1024 + 1)):
+        memory_files.meminfo.write_text(
+            f"MemTotal: 32768 kB\nMemAvailable: {available} kB\n", encoding="ascii"
+        )
+        assert memory_files.effective() == 0
+
+
+def test_permission_error_on_meminfo_fails_closed(
+    memory_files: FakeMemoryFiles, monkeypatch: pytest.MonkeyPatch
+):
+    memory_files.membership("0::/missing")
+    denied = memory_files.meminfo
+    original_open = Path.open
+
+    def deny_meminfo(path: Path, *args, **kwargs):
+        if path == denied:
+            raise PermissionError("fixture denied")
+        return original_open(path, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "open", deny_meminfo)
+    assert memory_files.effective() == 0
+
+
+def test_completely_absent_probe_is_unknown(tmp_path: Path):
+    assert (
+        effective_memory_available(
+            meminfo_path=tmp_path / "missing-meminfo",
+            cgroup_root=tmp_path / "missing-cgroup",
+            proc_cgroup_path=tmp_path / "missing-proc-cgroup",
+        )
+        is None
+    )
+
+
+def test_hybrid_v1_and_v2_membership_is_not_a_failure(memory_files: FakeMemoryFiles):
+    """systemd hosts keep leftover v1 hierarchies (``1:net_cls:/``) above the unified
+    ``0::/...`` line; those lines must be skipped, not treated as malformed or as a
+    memory controller, and the real unified root carries no ``memory.max`` of its own."""
+    memory_files.membership("2:cpu,cpuacct:/", "1:net_cls:/", "0::/tenant/job")
+    memory_files.v2(memory_files.root / "tenant", 20_000_000, 14_000_000)
+    memory_files.v2(
+        memory_files.root / "tenant" / "job", "max", "not-read-for-unlimited"
+    )
+
+    assert memory_files.remaining() == 6_000_000
+    assert memory_files.effective() == 6_000_000

From e5d3bd3fafb7c9b5a026bb59daa9ab324f18f246 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 16:33:16 -0700
Subject: [PATCH 099/249] feat(gguf): serve GGUF checkpoints on vendored
 ggml-cuda kernels (#690)

* build(kernel): vendor ggml-cuda from llama.cpp 8c2134a8 behind a tvm-ffi binding

* build(deps): vendor the gguf-py reader

* feat(quant): add the GGUF quant kind and an embedding quant method

* feat(gguf): serve Flash-Next and Gemma4 GGUF on the GGUF quant kind and drop the old Q4_0 path

* fix(kernel): build ggml-cuda for plain sm_120 by turning off the native fp4 MMA

* perf(gguf): write ggml results in the activation dtype and keep MoE intermediates in bf16

* build(kernel): compress the ggml fatbin to keep the multi-arch kernel cache small

* fix(gguf): read F32, F16 and BF16 PLE tables without ggml

* fix(moe): pad the offload slot cache for ggml MMQ tail reads

* fix(gguf): support hip builds and runtime devices

* fix(gguf): copy same-type dequantize and check for missing ggml converters

* fix(gguf): load quantized qwen4exp ple_key and ple_value

* feat(gguf): run float parts of a fused GGUF linear as bf16 segments

* fix(gguf): load quantized qwen4exp hyper-connection and gdn ba weights

* fix(moe): zero the prefill staging area after the last prefill layer

---------

Co-authored-by: zihaomu <zihaomu@outlook.com>
---
 benchmarks/bench_load_weight_generic.py       |    2 +-
 docs/cli.md                                   |    2 +-
 docs/models.md                                |    4 +-
 pyproject.toml                                |    3 +-
 python/freetoken/checkpoint/convert.py        |    2 +-
 python/freetoken/engine/config.py             |    2 +-
 python/freetoken/engine/engine.py             |   29 +-
 python/freetoken/kernel/aot.py                |   10 +
 python/freetoken/kernel/aot_models.py         |   12 +-
 python/freetoken/kernel/csrc/ggml/LICENSE     |   21 +
 python/freetoken/kernel/csrc/ggml/VERSION     |   13 +
 .../kernel/csrc/ggml/freetoken.patch          |  994 +++
 .../kernel/csrc/ggml/freetoken/ggml_ops.cu    |  426 ++
 .../kernel/csrc/ggml/freetoken/ggml_shim.cu   |  255 +
 .../kernel/csrc/ggml/ggml-cuda/common.cuh     | 1722 +++++
 .../kernel/csrc/ggml/ggml-cuda/convert.cu     |  743 ++
 .../kernel/csrc/ggml/ggml-cuda/convert.cuh    |   66 +
 .../kernel/csrc/ggml/ggml-cuda/dequantize.cuh |  452 ++
 .../kernel/csrc/ggml/ggml-cuda/getrows.cu     |  488 ++
 .../kernel/csrc/ggml/ggml-cuda/getrows.cuh    |   15 +
 .../kernel/csrc/ggml/ggml-cuda/mma.cuh        | 1594 +++++
 .../kernel/csrc/ggml/ggml-cuda/mmid.cu        |  221 +
 .../kernel/csrc/ggml/ggml-cuda/mmid.cuh       |    8 +
 .../csrc/ggml/ggml-cuda/mmq-config-ampere.cuh |  383 +
 .../ggml/ggml-cuda/mmq-config-blackwell.cuh   |   37 +
 .../csrc/ggml/ggml-cuda/mmq-config-cdna.cuh   |  185 +
 .../csrc/ggml/ggml-cuda/mmq-config-gcn.cuh    |  281 +
 .../ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh |  273 +
 .../ggml-cuda/mmq-config-pascal-older.cuh     |  273 +
 .../csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh  |  273 +
 .../ggml/ggml-cuda/mmq-config-rdna3-5.cuh     |  290 +
 .../csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh  |  274 +
 .../csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh  |  290 +
 .../csrc/ggml/ggml-cuda/mmq-load-tiles.cuh    | 1768 +++++
 .../csrc/ggml/ggml-cuda/mmq-vec-dot.cuh       | 1237 ++++
 .../kernel/csrc/ggml/ggml-cuda/mmq.cu         |  522 ++
 .../kernel/csrc/ggml/ggml-cuda/mmq.cuh        | 1661 +++++
 .../kernel/csrc/ggml/ggml-cuda/mmvq.cu        | 1626 +++++
 .../kernel/csrc/ggml/ggml-cuda/mmvq.cuh       |   20 +
 .../kernel/csrc/ggml/ggml-cuda/quantize.cu    |  729 ++
 .../kernel/csrc/ggml/ggml-cuda/quantize.cuh   |   74 +
 .../template-instances/mmq-instance-iq1_s.cu  |    5 +
 .../template-instances/mmq-instance-iq2_s.cu  |    5 +
 .../template-instances/mmq-instance-iq2_xs.cu |    5 +
 .../mmq-instance-iq2_xxs.cu                   |    5 +
 .../template-instances/mmq-instance-iq3_s.cu  |    5 +
 .../mmq-instance-iq3_xxs.cu                   |    5 +
 .../template-instances/mmq-instance-iq4_nl.cu |    5 +
 .../template-instances/mmq-instance-iq4_xs.cu |    5 +
 .../template-instances/mmq-instance-mxfp4.cu  |    7 +
 .../template-instances/mmq-instance-nvfp4.cu  |    7 +
 .../template-instances/mmq-instance-q1_0.cu   |    5 +
 .../template-instances/mmq-instance-q2_0.cu   |    5 +
 .../template-instances/mmq-instance-q2_k.cu   |    5 +
 .../template-instances/mmq-instance-q3_k.cu   |    5 +
 .../template-instances/mmq-instance-q4_0.cu   |    5 +
 .../template-instances/mmq-instance-q4_1.cu   |    5 +
 .../template-instances/mmq-instance-q4_k.cu   |    5 +
 .../template-instances/mmq-instance-q5_0.cu   |    5 +
 .../template-instances/mmq-instance-q5_1.cu   |    5 +
 .../template-instances/mmq-instance-q5_k.cu   |    5 +
 .../template-instances/mmq-instance-q6_k.cu   |    5 +
 .../template-instances/mmq-instance-q8_0.cu   |    5 +
 .../kernel/csrc/ggml/ggml-cuda/unary.cuh      |  123 +
 .../kernel/csrc/ggml/ggml-cuda/vecdotq.cuh    | 1380 ++++
 .../kernel/csrc/ggml/ggml-cuda/vendors/cuda.h |   28 +
 .../kernel/csrc/ggml/ggml-cuda/vendors/hip.h  |  313 +
 .../kernel/csrc/ggml/ggml-cuda/vendors/musa.h |  161 +
 .../kernel/csrc/ggml/include/ggml-alloc.h     |   87 +
 .../kernel/csrc/ggml/include/ggml-backend.h   |  463 ++
 .../kernel/csrc/ggml/include/ggml-cuda.h      |   47 +
 .../freetoken/kernel/csrc/ggml/include/ggml.h | 3027 ++++++++
 .../freetoken/kernel/csrc/ggml/include/gguf.h |  211 +
 .../kernel/csrc/ggml/src/ggml-common.h        | 1911 +++++
 .../kernel/csrc/ggml/src/ggml-impl.h          |  795 +++
 .../freetoken/kernel/csrc/gguf/dequantize.cuh |  583 --
 python/freetoken/kernel/csrc/gguf/dispatch.h  |   28 -
 .../freetoken/kernel/csrc/gguf/ggml-common.h  | 1029 ---
 .../freetoken/kernel/csrc/gguf/gguf_kernel.cu |  848 ---
 python/freetoken/kernel/csrc/gguf/mmq.cuh     |  881 ---
 python/freetoken/kernel/csrc/gguf/mmvq.cuh    |  352 -
 python/freetoken/kernel/csrc/gguf/moe.cuh     | 1379 ----
 python/freetoken/kernel/csrc/gguf/moe_vec.cuh |  413 --
 python/freetoken/kernel/csrc/gguf/vecdotq.cuh | 2037 ------
 python/freetoken/kernel/ggml.py               |  155 +
 python/freetoken/kernel/gguf.py               |  142 -
 python/freetoken/layers/activation.py         |   19 +-
 python/freetoken/layers/embedding.py          |   27 +-
 python/freetoken/layers/gguf.py               |  128 -
 python/freetoken/layers/linear.py             |    2 +
 python/freetoken/layers/moe.py                |   46 +-
 .../freetoken/layers/quantization/__init__.py |    5 +-
 .../layers/quantization/configs/__init__.py   |    7 +-
 .../layers/quantization/configs/base.py       |    8 +-
 .../layers/quantization/configs/gguf.py       |   82 +
 .../layers/quantization/embedding/__init__.py |   10 +
 .../layers/quantization/embedding/base.py     |   64 +
 .../layers/quantization/embedding/gguf.py     |   45 +
 .../quantization/embedding/unquantized.py     |   33 +
 .../layers/quantization/ggml_types.py         |   88 +
 .../layers/quantization/linear/__init__.py    |    7 +-
 .../layers/quantization/linear/gguf.py        |   56 +
 .../layers/quantization/moe/__init__.py       |    7 +-
 .../freetoken/layers/quantization/moe/gguf.py |  118 +
 .../freetoken/layers/quantization/registry.py |    1 +
 .../freetoken/layers/quantization/scheme.py   |   11 +
 python/freetoken/models/gemma4/__init__.py    |   10 +-
 python/freetoken/models/gemma4/gguf.py        |  363 +-
 python/freetoken/models/gemma4/model.py       |    9 +-
 python/freetoken/models/gemma4/weight.py      |    2 +-
 python/freetoken/models/gguf/__init__.py      |    7 -
 python/freetoken/models/gguf/config.py        |    1 +
 python/freetoken/models/gguf/dequant.py       |  153 -
 python/freetoken/models/gguf/experts.py       |  104 +
 python/freetoken/models/gguf/reader.py        |  115 +-
 python/freetoken/models/gguf/tensors.py       |   56 +
 python/freetoken/models/gguf/tokenizer.py     |    2 +-
 python/freetoken/models/qwen4_exp/__init__.py |    4 +
 python/freetoken/models/qwen4_exp/config.py   |    2 +
 python/freetoken/models/qwen4_exp/gdn.py      |    7 +-
 python/freetoken/models/qwen4_exp/gguf.py     |  264 +
 python/freetoken/models/qwen4_exp/hc.py       |    1 +
 python/freetoken/models/qwen4_exp/model.py    |    8 +-
 python/freetoken/models/qwen4_exp/ple_disk.py |   42 +-
 python/freetoken/models/register.py           |   32 +-
 python/freetoken/models/weight.py             |   19 -
 python/freetoken/moe/expert_banks.py          |   67 +-
 python/freetoken/moe/fused_q4_0.py            |   53 -
 python/freetoken/moe/offload_cache.py         |   85 +-
 python/freetoken/moe/offload_kernels.py       |    2 +-
 python/freetoken/third_party/__init__.py      |    1 +
 python/freetoken/third_party/gguf/LICENSE     |   21 +
 python/freetoken/third_party/gguf/VERSION     |   10 +
 python/freetoken/third_party/gguf/__init__.py |    4 +
 .../freetoken/third_party/gguf/constants.py   | 6332 +++++++++++++++++
 .../third_party/gguf/freetoken.patch          |   23 +
 .../freetoken/third_party/gguf/gguf_reader.py |  390 +
 python/freetoken/third_party/gguf/lazy.py     |  293 +
 python/freetoken/third_party/gguf/quants.py   | 1378 ++++
 tests/kernels/test_ggml.py                    |   84 +
 tests/models/deepseek_v41/harness.py          |    4 +-
 tests/moe/test_cpu_moe_q4_0.py                |   23 +-
 tests/moe/test_offload.py                     |   52 +-
 143 files changed, 35946 insertions(+), 8628 deletions(-)
 create mode 100644 python/freetoken/kernel/csrc/ggml/LICENSE
 create mode 100644 python/freetoken/kernel/csrc/ggml/VERSION
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken.patch
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/ggml-backend.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/ggml.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/gguf.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/src/ggml-common.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/src/ggml-impl.h
 delete mode 100644 python/freetoken/kernel/csrc/gguf/dequantize.cuh
 delete mode 100644 python/freetoken/kernel/csrc/gguf/dispatch.h
 delete mode 100644 python/freetoken/kernel/csrc/gguf/ggml-common.h
 delete mode 100644 python/freetoken/kernel/csrc/gguf/gguf_kernel.cu
 delete mode 100644 python/freetoken/kernel/csrc/gguf/mmq.cuh
 delete mode 100644 python/freetoken/kernel/csrc/gguf/mmvq.cuh
 delete mode 100644 python/freetoken/kernel/csrc/gguf/moe.cuh
 delete mode 100644 python/freetoken/kernel/csrc/gguf/moe_vec.cuh
 delete mode 100644 python/freetoken/kernel/csrc/gguf/vecdotq.cuh
 create mode 100644 python/freetoken/kernel/ggml.py
 delete mode 100644 python/freetoken/kernel/gguf.py
 delete mode 100644 python/freetoken/layers/gguf.py
 create mode 100644 python/freetoken/layers/quantization/configs/gguf.py
 create mode 100644 python/freetoken/layers/quantization/embedding/__init__.py
 create mode 100644 python/freetoken/layers/quantization/embedding/base.py
 create mode 100644 python/freetoken/layers/quantization/embedding/gguf.py
 create mode 100644 python/freetoken/layers/quantization/embedding/unquantized.py
 create mode 100644 python/freetoken/layers/quantization/ggml_types.py
 create mode 100644 python/freetoken/layers/quantization/linear/gguf.py
 create mode 100644 python/freetoken/layers/quantization/moe/gguf.py
 delete mode 100644 python/freetoken/models/gguf/dequant.py
 create mode 100644 python/freetoken/models/gguf/experts.py
 create mode 100644 python/freetoken/models/gguf/tensors.py
 create mode 100644 python/freetoken/models/qwen4_exp/gguf.py
 delete mode 100644 python/freetoken/moe/fused_q4_0.py
 create mode 100644 python/freetoken/third_party/__init__.py
 create mode 100644 python/freetoken/third_party/gguf/LICENSE
 create mode 100644 python/freetoken/third_party/gguf/VERSION
 create mode 100644 python/freetoken/third_party/gguf/__init__.py
 create mode 100644 python/freetoken/third_party/gguf/constants.py
 create mode 100644 python/freetoken/third_party/gguf/freetoken.patch
 create mode 100644 python/freetoken/third_party/gguf/gguf_reader.py
 create mode 100644 python/freetoken/third_party/gguf/lazy.py
 create mode 100644 python/freetoken/third_party/gguf/quants.py
 create mode 100644 tests/kernels/test_ggml.py

diff --git a/benchmarks/bench_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py
index 4e4c5fe012..89ddb30708 100644
--- a/benchmarks/bench_load_weight_generic.py
+++ b/benchmarks/bench_load_weight_generic.py
@@ -179,7 +179,7 @@ def _bench_load(mode: str, model_path: str, *, parallel: bool, workers: int, chu
     t = time.perf_counter()
     try:
         banks = load_expert_banks(model_path, mc, device=torch.device("cuda", torch.cuda.current_device()),
-                                  dtype=torch.bfloat16, parallel=parallel,
+                                  parallel=parallel,
                                   workers=workers, chunk=chunk)
     except NotImplementedError as e:
         s.stop()
diff --git a/docs/cli.md b/docs/cli.md
index 44b0e30ad9..4918a3177a 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -189,7 +189,7 @@ ft checkpoint --model <hf_dir> --out <ftw_dir> [--dtype bfloat16] [--moe-backend
 Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained
 fast-load format; point `ft serve --model` at the output dir. MoE experts are
 always packed into expert banks, which serve the fused and the offload
-strategies alike (GGUF and DeepSeek-V4 experts serve offload only);
+strategies alike (DeepSeek-V4 experts serve offload only);
 `--moe-backend` is kept for compatibility. See the FTW caveats in
 [models.md](models.md#notes); FTW files from older builds can be repaired with
 [scripts/ftw_hotfix.py](ftw-hotfix.md) instead of reconverting.
diff --git a/docs/models.md b/docs/models.md
index cdadc2dcb1..d3f26a086b 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -42,7 +42,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 `ft serve --moe-strategy {auto,fused,offload,cpu,hybrid}` (`--moe-backend` is the deprecated old spelling):
 
 - **fused** — experts resident on GPU (needs the VRAM). Serves every expert
-  format except GGUF and DeepSeek-V4's.
+  format except DeepSeek-V4's.
 - **offload** — experts live in host RAM, an LRU cache of expert slots on GPU;
   misses stream over PCIe.
 - **cpu** — misses are computed on the CPU instead of fetched.
@@ -56,7 +56,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 
 The GPU and the CPU share one DRAM, so offload only copies experts between two
 names for the same memory. On these GPUs `auto` resolves every MoE model to
-`fused`: the default is the resident path only. GGUF and DeepSeek-V4 experts
+`fused`: the default is the resident path only. DeepSeek-V4 experts
 have no resident path and stay on `offload`. Offload-only flags
 (`--moe-cpu-layers`, `--moe-cache-size` / `-rate` / `-auto`,
 `--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a
diff --git a/pyproject.toml b/pyproject.toml
index b308045f6c..ab1dfd0731 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,7 +40,6 @@ dependencies = [
     "fastapi>=0.115,<1",
     # slot_cache: the device-side LRU admission kernel behind the MoE expert cache.
     "flashlib==0.4.0", # pin it
-    "gguf>=0.19,<1",
     "huggingface_hub>=1.5,<2",
     "msgpack>=1.1,<2",
     "modelscope>=1.37,<2",
@@ -113,7 +112,7 @@ version = {attr = "freetoken.version.__version__"}
 where = ["python"]
 
 [tool.setuptools.package-data]
-"*" = ["csrc/**/*", "moe/configs/**/*.json"]
+"*" = ["csrc/**/*", "moe/configs/**/*.json", "third_party/*/LICENSE", "third_party/*/VERSION", "third_party/*/freetoken.patch"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py
index 152f3871c0..054fade4eb 100644
--- a/python/freetoken/checkpoint/convert.py
+++ b/python/freetoken/checkpoint/convert.py
@@ -247,7 +247,7 @@ def convert_checkpoint(
         names = legacy_bank_names(legacy_format_for(method.kind, method.kernel.name)) if method is not None else {}
         sink = _ConvertSink(writer, names=names)
         banks = load_expert_banks(
-            model_path, mc, method=method, device=dev, dtype=dtype, layer_sink=sink
+            model_path, mc, method=method, device=dev, layer_sink=sink
         )
         quant_format = banks.quant_format
         if banks.streamed:
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index ae2f1fd629..a6c47dd0fb 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -141,7 +141,7 @@ def model_config(self) -> ModelConfig:
         hf_config = copy.copy(self.hf_config)
         built = {e.config_key for e in self.active_encoders}
         for key in set(ENCODER_SECTIONS) | {e.config_key for e in self.model_spec.encoders}:
-            if key not in built:
+            if key not in built and getattr(hf_config, key, None) is not None:
                 setattr(hf_config, key, None)
         spec = self.model_spec
         quant = checkpoint_quant_config(self.model_path, hf_config, spec)
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 2757b10dea..81f42dbadd 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -15,7 +15,7 @@
 from freetoken.gpu_select import gpu_identity
 from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint
 from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device
-from freetoken.layers.quantization import LayerKind, QuantBackend, finalize_quant, set_quant_backend
+from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind, finalize_quant, set_quant_backend
 from freetoken.moe.offload_cache import iter_offload_moe_layers
 from freetoken.mm.config import ENCODER_SECTIONS
 from freetoken.models import create_model, load_weight
@@ -304,7 +304,7 @@ def _make_dummy_weight_state_dict(
             # MXFP8 e8m0 exponent codes: 127 encodes scale 1.0; zeros would collapse
             # every scale to 2^-127 and zero the model. Scoped BY NAME: other uint8
             # buffers are packed payloads whose bytes mean something else entirely
-            # (GGUF qweight blocks embed fp16 scales -- 0x7F7F is fp16 NaN), so they
+            # (GGUF blocks embed fp16 scales -- 0x7F7F is fp16 NaN), so they
             # keep the benign all-zeros fill below.
             state_dict[key] = torch.full(param.shape, 127, dtype=param.dtype, device=device)
         else:
@@ -603,7 +603,6 @@ def _load_resident_experts(self, config: EngineConfig, layers) -> None:
                 config.model_config,
                 method=layers[0].quant_method,
                 device=self.device,
-                dtype=self.dtype,
                 dummy=config.use_dummy_weight,
                 parallel={"serial": False, "parallel": True}.get(config.expert_load, None),
                 resident=True,
@@ -735,10 +734,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                     config.model_config,
                     method=method,
                     device=self.device,
-                    dtype=self.dtype,
                     dummy=config.use_dummy_weight,
                     parallel=expert_parallel,
-                    decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"),
                     layer_residency=requested_residency,
                 )
         except PinFailed as exc:
@@ -808,7 +805,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                     f"expert banks were packed for {banks.kind} / {banks.kernel} but the model "
                     f"binds {method.kind} / {method.kernel.name}; reconvert or select that kernel"
                 )
-            layout = method.layout()
+            # GGUF layers keep their own ggml types, so the cache reads each layer's layout off its banks
+            layout = None if method.kind is QuantKind.GGUF else method.layout()
             max_slots = method.slot_limit()
         cache = OffloadMoeCache(
             # Models with leading dense layers (GLM-4) only have experts on the MoE
@@ -830,7 +828,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         )
         # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set
         cache.cpu_layer_ids = cpu_layer_ids
-        cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags)
+        cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags, row_tail_bytes=banks.row_tail_bytes)
         cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
         if decode_target == "hybrid":
             self._resolve_hybrid_fetch(config, cache)
@@ -839,7 +837,6 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         cache.collect_stats = config.moe_collect_stats
         layers = attach_offload_moe_cache(self.model, cache)
         assert len(layers) == config.model_config.num_moe_layers
-        cache.zero_inactive_route_slots = cache.decode_target == "hybrid" and not all(layer.skips_inactive_slots for layer in layers)
         if cache.decode_target in ("cpu", "hybrid"):
             self._init_cpu_moe_executor(config, cache, layers)
         self.ctx.moe_offload_cache = cache
@@ -907,6 +904,9 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         # round a batch up to the largest captured size; cover both.
         max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1)
         method = sample.quant_method
+        formats = {layer.quant_method.cpu_format for layer in layers}
+        if len(formats) > 1:
+            raise NotImplementedError(f"CPU MoE backend needs one expert format across the MoE layers, got {sorted(map(str, formats))}")
         executor = CpuMoeExecutor(
             cache,
             top_k=sample.top_k,
@@ -917,9 +917,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
             device=self.device,
             swiglu_alpha=float(sample.alpha),
             swiglu_limit=sample.limit,
-            # FIXME: the None branch serves GGUF q4_0 banks, which have no quant method yet; drop it once GGUF joins the quant path
-            fmt=method.cpu_format if method is not None else None,
-            act_block=method.cfg.act_block if method is not None else None,
+            fmt=method.cpu_format,
+            act_block=method.cfg.act_block,
         )
         cache.set_cpu_executor(executor)
         self.cpu_moe_executor = executor
@@ -1255,11 +1254,11 @@ def _is_unified_memory_gpu(index: "int | None" = None) -> bool:
         return False
 
 
-_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4")
+_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4", "gguf")
 
 
 def _fused_resident_ok(model_config) -> bool:
-    """Whether the resident ('fused') MoE path can hold this model's experts; GGUF q4_0 and DeepSeek-V4 experts stay offload."""
+    """Whether the resident ('fused') MoE path can hold this model's experts; DeepSeek-V4 experts stay offload."""
     if getattr(model_config, "expert_quant", "none") not in _RESIDENT_EXPERT_QUANTS:
         return False
     return getattr(model_config, "moe_weight_format", None) in (None, "bf16", "mxfp4")
@@ -1752,7 +1751,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     if is_moe and config.moe_strategy == "auto":
         # A MoE model always defaults to the offload family: experts stream from pinned host
         # banks into an auto-sized GPU slot cache, which is the only default that serves a model
-        # bigger than the GPU. The resident 'fused' path (every expert format but GGUF q4_0 and
+        # bigger than the GPU. The resident 'fused' path (every expert format but
         # DeepSeek-V4's) is still reachable, but only when asked for explicitly
         # -- auto never picks it, because nothing here knows whether the experts would fit in
         # HBM and a wrong guess is a weight-load OOM rather than a slower-but-working run.
@@ -1790,7 +1789,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         # offload -- same auto-sized GPU slot cache (_resolve_auto_moe_cache_size), plus a
         # host-RAM CPU executor -- so this never raises the OOM risk; with no profile (or one
         # from different hardware) it stays offload. offload remains the always-safe fallback.
-        # Key the lookup on the real expert format: mxfp4/q4_0 live in moe_weight_format when
+        # Key the lookup on the real expert format: mxfp4 lives in moe_weight_format when
         # expert_quant is "none", and "none" with no weight format means plain bf16 experts.
         moe_wfmt = getattr(model_config, "moe_weight_format", None)
         bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py
index bbc2d2c669..76250af7b6 100644
--- a/python/freetoken/kernel/aot.py
+++ b/python/freetoken/kernel/aot.py
@@ -143,6 +143,15 @@ def build(build_directory: pathlib.Path) -> object:
     return KernelSpec(name=_make_name("radix"), build=build)
 
 
+def _ggml_spec() -> KernelSpec:
+    def build(build_directory: pathlib.Path) -> object:
+        from .ggml import load_module
+
+        return load_module(str(build_directory))
+
+    return KernelSpec(name=_make_name("ggml"), build=build)
+
+
 def default_kernel_specs() -> tuple[KernelSpec, ...]:
     specs: list[KernelSpec] = []
     specs.extend(_store_spec(element_size) for element_size in DEFAULT_STORE_ELEMENT_SIZES)
@@ -158,6 +167,7 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]:
     specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64))
     specs.append(_batch_memcpy_spec())
     specs.append(_radix_spec())
+    specs.append(_ggml_spec())
 
     names = [spec.name for spec in specs]
     duplicates = {name for name in names if names.count(name) > 1}
diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index ecd1f8502b..308c2d6810 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -16,7 +16,7 @@
   DSV4 writes its MLA latent via torch scatter and contributes nothing.
 - index: ``element_size = hidden_size * 2`` (bf16 embedding row) paired with
   the runtime ``num_splits_for`` rule (layers/embedding.py -> kernel/index.py).
-  DSV4 (plain nn.Embedding) and GGUF embeddings (GGUFEmbedding) bypass it.
+  DSV4 (plain nn.Embedding) and GGUF embeddings (ggml get_rows) bypass it.
 
 The whole table targets the shipped serving configuration: TP=1 (TP>1 shards
 kv heads, shrinking the store row) and a 2-byte compute dtype (``--dtype
@@ -91,9 +91,6 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
             "down": H * I,
             "down_scale": (H // B) * fp8_block_scale_pad(H // B, I // B) * 2,
         }
-    if fmt == "q4_0":
-        # gemma4/gguf.py _q4_0_expert_specs: GGML Q4_0 rows, 32 elems -> 18 bytes
-        return {"gate_up": 2 * I * (H // 32 * 18), "down": H * (I // 32 * 18)}
     if fmt in ("nvfp4", "nvfp4_marlin", "nvfp4_b12x"):
         # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp16
         # per-row globals; marlin/b12x repacks are byte-identical with the globals
@@ -205,6 +202,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         top_k=10,
         moe_intermediate_size=640,
         expert_formats=(*_NVFP4_FORMATS, "fp8_block"),
+        # GGUF expert rows differ per layer and reach the fused copy as runtime sizes
+        arch_aliases=("Qwen4ExpGGUFForCausalLM",),
     ),
     AotModel(
         name="google/gemma-4-26B-A4B-it",
@@ -215,9 +214,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         kv_groups=((2, 512), (8, 256)),
         top_k=8,
         moe_intermediate_size=704,
-        # q4_0 = the Q4_0 GGUF release (Gemma4GGUFForCausalLM); its GGUFEmbedding
-        # path skips indexing(), but the bf16 safetensors path needs the variant.
-        expert_formats=("bf16", "q4_0"),
+        expert_formats=("bf16",),
+        # GGUF expert rows differ per layer and reach the fused copy as runtime sizes
         arch_aliases=("Gemma4ForCausalLM", "Gemma4GGUFForCausalLM"),
     ),
     AotModel(
diff --git a/python/freetoken/kernel/csrc/ggml/LICENSE b/python/freetoken/kernel/csrc/ggml/LICENSE
new file mode 100644
index 0000000000..e7dca554bc
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2026 The ggml authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION
new file mode 100644
index 0000000000..6eabb0c19f
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/VERSION
@@ -0,0 +1,13 @@
+llama.cpp 8c2134a8755d8a9792af513dab20282affe903e4 (2026-10-14)
+
+Only the ggml files that the GGUF ops (freetoken/ggml_ops.cu, built by kernel/ggml.py) compile or include:
+ggml/include -> include, ggml/src -> src, ggml/src/ggml-cuda -> ggml-cuda.
+
+Local changes, kept in freetoken.patch (already applied):
+- quantize.cu/.cuh: the q8_1 quantize kernels read fp32, bf16 or fp16 activations.
+- mmq.cu/.cuh: MMQ takes bf16 or fp16 activations, except on the native fp4 path. Its kernels are built for fp32 and bf16 results (dst_t); a bf16 result that needs the stream-k fixup is computed in fp32 first, and so is an fp16 result unless GGML_MMQ_F16_DST builds fp16 kernels too.
+- mmvq.cu/.cuh: MMVQ takes bf16 or fp16 activations, can write bf16 or fp16 results, and can reuse activations already quantized to q8_1 (src1_q8_1_pre).
+- mmq.cu/.cuh, mmid.cu/.cuh: MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold tokens, listed by mm_ids_tile_map; ncols_max > 0 keeps the regular grid.
+- common.cuh: the native fp4 MMA (BLACKWELL_MMA_AVAILABLE, blackwell_mma_available) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build.
+
+To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file.
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken.patch b/python/freetoken/kernel/csrc/ggml/freetoken.patch
new file mode 100644
index 0000000000..3dc715fe0f
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken.patch
@@ -0,0 +1,994 @@
+--- a/ggml/src/ggml-cuda/common.cuh
++++ b/ggml/src/ggml-cuda/common.cuh
+@@ -293,7 +293,8 @@
+ #define AMPERE_MMA_AVAILABLE
+ #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+ 
+-#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN
++// FreeToken: the native fp4 MMA needs an sm_120a/sm_120f build, so it is off unless GGML_CUDA_FP4_MMA is defined; FP4 weights then take Q8_1 activations
++#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN && defined(GGML_CUDA_FP4_MMA)
+ #    define BLACKWELL_MMA_AVAILABLE
+ #endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL
+ 
+@@ -374,8 +375,13 @@
+ }
+ 
+ static bool blackwell_mma_available(const int cc) {
++#ifdef GGML_CUDA_FP4_MMA
+     return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_BLACKWELL &&
+            ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
++#else
++    GGML_UNUSED(cc);
++    return false;
++#endif // GGML_CUDA_FP4_MMA
+ }
+ 
+ // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+--- a/ggml/src/ggml-cuda/mmid.cu
++++ b/ggml/src/ggml-cuda/mmid.cu
+@@ -177,3 +177,45 @@
+             break;
+     }
+ }
++
++// Lists the (expert, column tile) pairs that hold columns for a flat MoE grid, in expert order.
++// The other ntiles grid columns get tile jt_none of expert 0, which starts past its columns, so mul_mat_q skips them.
++template <int nthreads>
++__launch_bounds__(nthreads, 1)
++static __global__ void mm_ids_tile_map(
++        const int32_t * __restrict__ expert_bounds, int2 * __restrict__ tile_map, const int n_experts, const int J, const int ntiles, const int jt_none) {
++    __shared__ int scan[nthreads];
++
++    int start = 0; // First tile of the current chunk of experts.
++    for (int e0 = 0; e0 < n_experts; e0 += nthreads) {
++        const int e = e0 + threadIdx.x;
++        const int n = e < n_experts ? (expert_bounds[e + 1] - expert_bounds[e] + J - 1) / J : 0;
++
++        scan[threadIdx.x] = n;
++        __syncthreads();
++#pragma unroll
++        for (int offset = 1; offset < nthreads; offset *= 2) {
++            const int add = threadIdx.x >= offset ? scan[threadIdx.x - offset] : 0;
++            __syncthreads();
++            scan[threadIdx.x] += add;
++            __syncthreads();
++        }
++
++        const int first = start + scan[threadIdx.x] - n;
++        for (int jt = 0; jt < n; ++jt) {
++            tile_map[first + jt] = make_int2(e, jt);
++        }
++        start += scan[nthreads - 1];
++        __syncthreads();
++    }
++
++    for (int t = start + threadIdx.x; t < ntiles; t += nthreads) {
++        tile_map[t] = make_int2(0, jt_none);
++    }
++}
++
++void ggml_cuda_launch_mm_ids_tile_map(
++        const int32_t * expert_bounds, int2 * tile_map, const int n_experts, const int J, const int ntiles, const int jt_none, cudaStream_t stream) {
++    constexpr int nthreads = 512;
++    mm_ids_tile_map<nthreads><<<1, nthreads, 0, stream>>>(expert_bounds, tile_map, n_experts, J, ntiles, jt_none);
++}
+--- a/ggml/src/ggml-cuda/mmid.cuh
++++ b/ggml/src/ggml-cuda/mmid.cuh
+@@ -3,3 +3,6 @@
+ void ggml_cuda_launch_mm_ids_helper(
+         const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
+         int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
++
++void ggml_cuda_launch_mm_ids_tile_map(
++        const int32_t * expert_bounds, int2 * tile_map, int n_experts, int J, int ntiles, int jt_none, cudaStream_t stream);
+--- a/ggml/src/ggml-cuda/mmq.cu
++++ b/ggml/src/ggml-cuda/mmq.cu
+@@ -132,9 +132,11 @@
+ }
+ 
+ void ggml_cuda_mul_mat_q(
+-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
+-    GGML_ASSERT(        src1->type == GGML_TYPE_F32);
+-    GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
++        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
++        int64_t ncols_max) {
++    // FreeToken: bf16 and fp16 activations are quantized directly, except for the native fp4 paths; dst is fp32, bf16 or fp16
++    GGML_ASSERT(        src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16);
++    GGML_ASSERT(        dst->type  == GGML_TYPE_F32 || dst->type  == GGML_TYPE_BF16 || dst->type  == GGML_TYPE_F16);
+     GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
+ 
+     GGML_TENSOR_BINARY_OP_LOCALS;
+@@ -156,7 +158,9 @@
+ 
+     const char  * src0_d = (const char  *) src0->data;
+     const float * src1_d = (const float *) src1->data;
+-    float       *  dst_d = (float       *)  dst->data;
++    const mmq_dst  dst_d = {dst->data, dst->type};
++    const nv_bfloat16 * src1_bf16 = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1->data : nullptr;
++    const half        * src1_f16  = src1->type == GGML_TYPE_F16  ? (const half        *) src1->data : nullptr;
+ 
+     // If src0 is a temporary compute buffer, clear any potential padding.
+     if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+@@ -183,6 +187,7 @@
+     const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+ 
+     const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
++    GGML_ASSERT(!use_native_fp4 || src1->type == GGML_TYPE_F32);
+     const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
+     const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ            : QK8_1_MMQ;
+ 
+@@ -192,7 +197,8 @@
+         int64_t ncols_opt = ne11;
+         if (ids) {
+             const int64_t n_expert_used = ids->ne[0];
+-            ncols_opt = ne12;
++            // a flat grid runs only the tiles that hold tokens, so it sizes them for the average expert
++            ncols_opt = ncols_max > 0 ? ncols_max : (ne12*n_expert_used + ne02 - 1) / ne02;
+ 
+             // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
+             // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
+@@ -249,6 +255,12 @@
+                 quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
+                                         ne11, ne12, ne13, stream);
+ 
++            } else if (src1_bf16) {
++                quantize_mmq_q8_1_cuda(src1_bf16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
++                                       ne11, ne12, ne13, stream);
++            } else if (src1_f16) {
++                quantize_mmq_q8_1_cuda(src1_f16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
++                                       ne11, ne12, ne13, stream);
+             } else {
+                 quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                                        ne11, ne12, ne13, stream);
+@@ -299,6 +311,17 @@
+         CUDA_CHECK(cudaGetLastError());
+     }
+ 
++    // Every expert's column tiles fit in ntx_flat grid columns, so a flat grid needs no host sync to size it.
++    ggml_cuda_pool_alloc<int2> tile_map(ctx.pool());
++    int ntx_flat = 0;
++    if (ncols_max <= 0) {
++        const int jt_none = (ne_get_rows + J_best - 1) / J_best;
++        ntx_flat = jt_none + std::min<int64_t>(ne02, ne_get_rows);
++        tile_map.alloc(ntx_flat);
++        ggml_cuda_launch_mm_ids_tile_map(expert_bounds.get(), tile_map.get(), ne02, J_best, ntx_flat, jt_none, stream);
++        CUDA_CHECK(cudaGetLastError());
++    }
++
+     const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
+     ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+     ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+@@ -325,9 +348,21 @@
+                 quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
+                                         ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+             }
++        } else if (dedup_bcast && src1_bf16) {
++            quantize_scatter_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
++                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
++        } else if (dedup_bcast && src1_f16) {
++            quantize_scatter_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
++                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+         } else if (dedup_bcast) {
+             quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                     /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
++        } else if (src1_bf16) {
++            quantize_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
++                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
++        } else if (src1_f16) {
++            quantize_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
++                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+         } else {
+             quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
+                                    ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+@@ -347,7 +382,7 @@
+         ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
+         ne02, ne02, s02, s12, s2,
+         ne03, ne13, s03, s13, s3,
+-        ne12, J_best};
++        ncols_max > 0 ? ncols_max : ne12, J_best, tile_map.ptr, ntx_flat};
+ 
+     ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
+ }
+--- a/ggml/src/ggml-cuda/mmq.cuh
++++ b/ggml/src/ggml-cuda/mmq.cuh
+@@ -10,6 +10,17 @@
+ #define MMQ_ITER_K_FP4         512
+ #define MMQ_NWARPS               8
+ 
++// FreeToken: MMQ writes fp32, bf16 or fp16 results (dst_t); the sums stay fp32 until the store
++template <typename T> static __device__ __forceinline__ T mmq_cast(float v);
++template <> __device__ __forceinline__ float mmq_cast<float>(float v) { return v; }
++template <> __device__ __forceinline__ nv_bfloat16 mmq_cast<nv_bfloat16>(float v) { return __float2bfloat16(v); }
++template <> __device__ __forceinline__ half mmq_cast<half>(float v) { return __float2half(v); }
++
++struct mmq_dst {
++    void * p;
++    ggml_type type;
++};
++
+ typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+ typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
+ typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
+@@ -429,8 +440,8 @@
+ #include "mmq-load-tiles.cuh"
+ #include "mmq-vec-dot.cuh"
+ 
+-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
+-        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
++template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t = float> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
++        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, dst_t * __restrict__ dst,
+         const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+     constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+@@ -456,21 +467,21 @@
+ 
+             if constexpr (type == GGML_TYPE_NVFP4) {
+                 if (y_scale_used) {
+-                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
++                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                 } else {
+-                    dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
++                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                 }
+             } else {
+-                dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
++                dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                 GGML_UNUSED(y_scale_used);
+             }
+         }
+     }
+ }
+ 
+-template<ggml_type type, int J, bool fallback, ggml_prec prec_src1>
++template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t = float>
+ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
+-            const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
++            const float * __restrict__ sum, const int * __restrict__ ids_dst, dst_t * __restrict__ dst,
+             const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+ 
+ #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+@@ -506,12 +517,12 @@
+ 
+                 if constexpr (type == GGML_TYPE_NVFP4) {
+                     if (y_scale_used) {
+-                        dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
++                        dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                     } else {
+-                        dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
++                        dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                     }
+                 } else {
+-                    dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
++                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                     GGML_UNUSED(y_scale_used);
+                 }
+             }
+@@ -869,10 +880,10 @@
+ 
+ // ---------------------------------------------------------------------------------------------
+ 
+-template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1>
++template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1, typename dst_t>
+ static __device__ __forceinline__ void mul_mat_q_process_tile(
+         const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+-        const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
++        const int * __restrict__ ids_dst, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup,
+         const float * __restrict__ y_scale,
+         const int stride_row_x, const int ncols_y, const int stride_col_dst,
+         const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+@@ -942,19 +953,32 @@
+ 
+     if (fixup) {
+         write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
+-    } else {
++    } else if constexpr (std::is_same_v<dst_t, float>) {
+         write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
++    } else if constexpr (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
++        ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, dst_t>(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
++    } else {
++        ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1, dst_t>(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+     }
+ }
+ 
+ 
++// A flat MoE grid runs only the (expert, column tile) pairs listed in tile_map, see ggml_cuda_launch_mm_ids_tile_map.
++static __device__ __forceinline__ void mmq_flat_tile(const int2 * __restrict__ tile_map, int & zt, int & jt) {
++    if (tile_map) {
++        const int2 t = tile_map[jt];
++        zt = t.x;
++        jt = t.y;
++    }
++}
++
+ // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
+ 
+-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
++template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t>
+ __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
+ static __global__ void mul_mat_q(
+         const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
+-        const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
++        const int32_t * __restrict__ expert_bounds, const int2 * __restrict__ tile_map, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup,
+         const float * __restrict__ y_scale,
+         const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
+         const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
+@@ -993,9 +1017,10 @@
+     if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+         const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
+         const int wt = tmp2.x;
+-        const int zt = tmp2.y;
+-        const int jt = blockIdx.y;
++        int zt = tmp2.y;
++        int jt = blockIdx.y;
+         const int it = blockIdx.x;
++        mmq_flat_tile(tile_map, zt, jt);
+ 
+         // Defaults for regular matrix multiplication:
+         int col_low    = 0;
+@@ -1076,14 +1101,15 @@
+     while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) {
+         int tmp = fastdiv(kbc, blocks_per_ne00);
+         uint2 tmp2 = fast_div_modulo(tmp, ntx);
+-        const int jt = tmp2.y;
++        int jt = tmp2.y;
+         tmp = tmp2.x;
+         tmp2 = fast_div_modulo(tmp, nchannels_y);
+-        const int zt = tmp2.y;
++        int zt = tmp2.y;
+         tmp = tmp2.x;
+         tmp2 = fast_div_modulo(tmp, nsamples_y);
+         const int wt = tmp2.y;
+         const int it = tmp2.x;
++        mmq_flat_tile(tile_map, zt, jt);
+ 
+         // Defaults for regular matrix multiplication:
+         int col_low    = 0;
+@@ -1165,14 +1191,15 @@
+ 
+     int tmp = fastdiv(kbc, blocks_per_ne00);
+     uint2 tmp2 = fast_div_modulo(tmp, ntx);
+-    const int jt = tmp2.y;
++    int jt = tmp2.y;
+     tmp = tmp2.x;
+     tmp2 = fast_div_modulo(tmp, nchannels_y);
+-    const int zt = tmp2.y;
++    int zt = tmp2.y;
+     tmp = tmp2.x;
+     tmp2 = fast_div_modulo(tmp, nsamples_y);
+     const int wt = tmp2.y;
+     const int it = tmp2.x;
++    mmq_flat_tile(tile_map, zt, jt);
+ 
+     // Defaults for regular matrix multiplication:
+     int col_low    = 0;
+@@ -1375,14 +1402,26 @@
+     }
+ }
+ 
++// FreeToken: copies an fp32 result with dst's layout into dst
++template <typename dst_t>
++static __global__ void mmq_store_dst(const float * __restrict__ src, dst_t * __restrict__ dst, const int nrows, const int stride) {
++    const int64_t col = blockIdx.x;
++    const int     row = blockIdx.y*blockDim.x + threadIdx.x;
++    if (row < nrows) {
++        dst[col*stride + row] = mmq_cast<dst_t>(src[col*stride + row]);
++    }
++}
++
+ struct mmq_args {
+-    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
++    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; mmq_dst dst;
+     const float * y_scale;
+     int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
+     int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
+     int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
+     int64_t ncols_max;
+     int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
++    const int2 * tile_map = nullptr; // Flat MoE grid: the (expert, column tile) of each of the ntx_flat grid columns.
++    int ntx_flat = 0;
+ };
+ 
+ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
+@@ -1406,12 +1445,11 @@
+ 
+     const dim3 block_dims(warp_size, nwarps, 1);
+ 
+-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1>), nbytes_shared);
+-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J,  true, prec_src1>), nbytes_shared);
+-
++    // a flat MoE grid puts all experts into one channel of ntx_flat column tiles
++    const int64_t nchannels_grid = args.tile_map ? 1 : args.nchannels_y;
+     const int nty  = (args.nrows_x   + config.I - 1) / config.I;
+-    const int ntx  = (args.ncols_max + config.J - 1) / config.J;
+-    const int ntzw = args.nchannels_y * args.nsamples_y;
++    const int ntx  = args.tile_map ? args.ntx_flat : (args.ncols_max + config.J - 1) / config.J;
++    const int ntzw = nchannels_grid * args.nsamples_y;
+     const dim3 block_nums_xy_tiling(nty, ntx, ntzw);
+ 
+     GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
+@@ -1421,18 +1459,71 @@
+ 
+     const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits<type>::qk);
+     const uint3 ntx_fd             = init_fastdiv_values(ntx);
+-    const uint3 nchannels_y_fd     = init_fastdiv_values(args.nchannels_y);
++    const uint3 nchannels_y_fd     = init_fastdiv_values(nchannels_grid);
+     const uint3 nsamples_y_fd      = init_fastdiv_values(args.nsamples_y);
+     const uint3 channel_ratio_fd   = init_fastdiv_values(channel_ratio);
+     const uint3 sample_ratio_fd    = init_fastdiv_values(sample_ratio);
+ 
+-    if (!config.stream_k) {
+-        mul_mat_q<type, J, fallback, prec_src1><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+-            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
++    const auto run = [&](auto * dst, const dim3 & block_nums, float * tmp_fixup) {
++        using dst_t = std::remove_pointer_t<decltype(dst)>;
++        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1, dst_t>), nbytes_shared);
++        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J,  true, prec_src1, dst_t>), nbytes_shared);
++        mul_mat_q<type, J, fallback, prec_src1, dst_t><<<block_nums, block_dims, nbytes_shared, stream>>>
++            (args.x, args.y, args.ids_dst, args.expert_bounds, args.tile_map, dst, tmp_fixup, args.y_scale,
+              blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
+              channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
+              sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
+              ntx_fd);
++    };
++
++    // FreeToken: the fixup kernel adds partial sums into an fp32 dst, so a bf16 dst that needs it is computed in fp32 first.
++    // An fp16 dst always is, unless GGML_MMQ_F16_DST builds the fp16 kernels.
++    const auto launch = [&](const dim3 & block_nums, const bool fixup, float * tmp_fixup) {
++        bool via_f32 = fixup && args.dst.type != GGML_TYPE_F32;
++#ifndef GGML_MMQ_F16_DST
++        via_f32 = via_f32 || args.dst.type == GGML_TYPE_F16;
++#endif // GGML_MMQ_F16_DST
++        ggml_cuda_pool_alloc<float> dst_f32(ctx.pool(id));
++        float * dst_float = args.dst.type == GGML_TYPE_F32 ? (float *) args.dst.p : nullptr;
++        if (via_f32) {
++            GGML_ASSERT(args.ids_dst || (args.nchannels_y == 1 && args.nsamples_y == 1));
++            dst_float = dst_f32.alloc(args.ncols_dst * args.nrows_dst);
++        }
++
++        if (dst_float) {
++            run(dst_float, block_nums, tmp_fixup);
++        } else if (args.dst.type == GGML_TYPE_BF16) {
++            run((nv_bfloat16 *) args.dst.p, block_nums, tmp_fixup);
++        }
++#ifdef GGML_MMQ_F16_DST
++        else {
++            run((half *) args.dst.p, block_nums, tmp_fixup);
++        }
++#endif // GGML_MMQ_F16_DST
++
++        if (fixup) {
++            CUDA_CHECK(cudaGetLastError());
++            const dim3 block_nums_fixup(block_nums.x, config.I/warp_size, 1);
++            const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
++            mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
++                (args.ids_dst, args.expert_bounds, dst_float, tmp_fixup, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
++                 args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
++                 ntx_fd);
++        }
++
++        if (via_f32) {
++            CUDA_CHECK(cudaGetLastError());
++            const dim3 block_nums_store(args.ncols_dst, (args.nrows_x + 255) / 256, 1);
++            if (args.dst.type == GGML_TYPE_BF16) {
++                mmq_store_dst<<<block_nums_store, 256, 0, stream>>>(dst_float, (nv_bfloat16 *) args.dst.p, args.nrows_x, args.nrows_dst);
++            } else {
++                mmq_store_dst<<<block_nums_store, 256, 0, stream>>>(dst_float, (half *) args.dst.p, args.nrows_x, args.nrows_dst);
++            }
++        }
++    };
++
++    if (!config.stream_k) {
++        launch(block_nums_xy_tiling, false, nullptr);
+         return;
+     }
+ 
+@@ -1441,7 +1532,8 @@
+     const int ntiles_dst = ntx * nty * ntzw;
+     const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
+     const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);
+-    const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : nsm, 1, 1);
++    // The fixup kernel does not read tile_map, so a flat MoE grid always runs one tile per block.
++    const dim3 block_nums_stream_k(args.tile_map || (GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90) ? ntiles_dst : nsm, 1, 1);
+ 
+     GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow.
+ 
+@@ -1453,25 +1545,7 @@
+         tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I);
+     }
+ 
+-    const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
+-    const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
+-
+-    mul_mat_q<type, J, fallback, prec_src1><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+-        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
+-         blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
+-         channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
+-         sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
+-         ntx_fd);
+-
+-    if (!fixup_needed) {
+-        return;
+-    }
+-
+-    CUDA_CHECK(cudaGetLastError());
+-    mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
+-        (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
+-         args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
+-         ntx_fd);
++    launch(block_nums_stream_k, fixup_needed, tmp_fixup.ptr);
+ }
+ 
+ template <ggml_type type, bool fallback, ggml_prec prec_src1>
+@@ -1579,7 +1653,9 @@
+ 
+ // -------------------------------------------------------------------------------------------------------------------------
+ 
++// FreeToken: under ids, ncols_max > 0 bounds the tokens any one expert gets and sizes the tile grid; otherwise the grid is flat, see ggml_cuda_launch_mm_ids_tile_map
+ void ggml_cuda_mul_mat_q(
+-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);
++        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
++        int64_t ncols_max = -1);
+ 
+ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);
+--- a/ggml/src/ggml-cuda/mmvq.cu
++++ b/ggml/src/ggml-cuda/mmvq.cu
+@@ -1,4 +1,5 @@
+ #include "mmvq.cuh"
++#include "convert.cuh"
+ #include "quantize.cuh"
+ #include "unary.cuh"
+ #include "vecdotq.cuh"
+@@ -596,10 +597,11 @@
+     return 1;
+ }
+ 
+-template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool halve_iters = false>
++// FreeToken: dst_t is float, nv_bfloat16 or half
++template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool halve_iters = false, typename dst_t = float>
+ __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1)
+ static __global__ void mul_mat_vec_q(
+-        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr,
++        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst_ptr,
+         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
+         uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x,
+         const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio,
+@@ -608,7 +610,7 @@
+     const void    * GGML_CUDA_RESTRICT vx  = vx_ptr;
+     const void    * GGML_CUDA_RESTRICT vy  = vy_ptr;
+     const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr;
+-    float         * GGML_CUDA_RESTRICT dst = dst_ptr;
++    dst_t         * GGML_CUDA_RESTRICT dst = dst_ptr;
+ 
+     constexpr int qk  = ggml_cuda_type_traits<type>::qk;
+     constexpr int qi  = ggml_cuda_type_traits<type>::qi;
+@@ -629,7 +631,9 @@
+     const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y;
+     if (shared_expert) {
+         vx = fusion.shared_up;
+-        dst = fusion.shared_dst;
++        if constexpr (std::is_same_v<dst_t, float>) {
++            dst = fusion.shared_dst;
++        }
+         stride_col_dst = fusion.shared_stride_col_dst;
+     }
+ 
+@@ -837,7 +841,7 @@
+                         }
+                     }
+                 }
+-                dst[j*stride_col_dst + i] = result;
++                dst[j*stride_col_dst + i] = ggml_cuda_cast<dst_t>(result);
+             }
+         }
+     }
+@@ -854,11 +858,11 @@
+ // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst)
+ // Block: (warp_size, ncols_dst) - each warp handles one token independently.
+ // No shared memory reduction needed since each warp works alone.
+-template <ggml_type type, int c_rows_per_block, bool has_fusion = false>
++template <ggml_type type, int c_rows_per_block, bool has_fusion = false, typename dst_t = float>
+ __launch_bounds__(get_mmvq_mmid_max_batch_for_device<type>()*ggml_cuda_get_physical_warp_size(), 1)
+ static __global__ void mul_mat_vec_q_moe(
+         const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion,
+-        float * dst_ptr,
++        dst_t * dst_ptr,
+         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
+         const uint32_t stride_row_x, const uint32_t stride_col_y, uint32_t stride_col_dst,
+         const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+@@ -866,7 +870,7 @@
+     const void    * GGML_CUDA_RESTRICT vx  = vx_ptr;
+     const void    * GGML_CUDA_RESTRICT vy  = vy_ptr;
+     const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr;
+-    float         * GGML_CUDA_RESTRICT dst = dst_ptr;
++    dst_t         * GGML_CUDA_RESTRICT dst = dst_ptr;
+ 
+     constexpr int qk  = ggml_cuda_type_traits<type>::qk;
+     constexpr int qi  = ggml_cuda_type_traits<type>::qi;
+@@ -878,7 +882,9 @@
+     const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1;
+     if (shared_expert) {
+         vx = fusion.shared_up;
+-        dst = fusion.shared_dst;
++        if constexpr (std::is_same_v<dst_t, float>) {
++            dst = fusion.shared_dst;
++        }
+         stride_col_dst = fusion.shared_stride_col_dst;
+     }
+ 
+@@ -998,7 +1004,7 @@
+                 }
+             }
+         }
+-        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result;
++        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = ggml_cuda_cast<dst_t>(result);
+     }
+ 
+     if constexpr (!has_fusion) {
+@@ -1020,9 +1026,9 @@
+     return {block_nums, block_dims};
+ }
+ 
+-template<ggml_type type, int c_ncols_dst, bool small_k = false, bool halve_iters = false>
++template<ggml_type type, int c_ncols_dst, bool small_k = false, bool halve_iters = false, typename dst_t>
+ static void mul_mat_vec_q_switch_fusion(
+-        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
++        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
+         const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x,
+         const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio,
+@@ -1035,7 +1041,7 @@
+     if constexpr (c_ncols_dst == 1) {
+         if (has_fusion) {
+             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
+-            ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, true, small_k, halve_iters>, launch_params,
++            ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, true, small_k, halve_iters, dst_t>, launch_params,
+                  vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
+                  channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
+                  sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
+@@ -1046,15 +1052,15 @@
+     GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1");
+ 
+     const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
+-    ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, false, small_k, halve_iters>, launch_params,
++    ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, false, small_k, halve_iters, dst_t>, launch_params,
+         vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
+         channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
+         sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
+ }
+ 
+-template <ggml_type type>
++template <ggml_type type, typename dst_t>
+ static void mul_mat_vec_q_moe_launch(
+-        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
++        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
+         const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
+         const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+@@ -1071,13 +1077,13 @@
+                             fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
+ 
+     if (has_fusion) {
+-        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true>, launch_params,
++        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true, dst_t>, launch_params,
+             vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+             stride_row_x, stride_col_y, stride_col_dst,
+             stride_channel_x, stride_channel_y, stride_channel_dst,
+             ncols_dst, ids_stride);
+     } else {
+-        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false>, launch_params,
++        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false, dst_t>, launch_params,
+             vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+             stride_row_x, stride_col_y, stride_col_dst,
+             stride_channel_x, stride_channel_y, stride_channel_dst,
+@@ -1085,9 +1091,9 @@
+     }
+ }
+ 
+-template <ggml_type type>
++template <ggml_type type, typename dst_t>
+ static void mul_mat_vec_q_switch_ncols_dst(
+-        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
++        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+         const int ncols_x, const int nrows_x, const int ncols_dst,
+         const int stride_row_x, const int stride_col_y, const int stride_col_dst,
+         const int nchannels_x, const int nchannels_y, const int nchannels_dst,
+@@ -1277,8 +1283,10 @@
+             GGML_ABORT("fatal error");
+     }
+ }
++
++template <typename dst_t>
+ static void mul_mat_vec_q_switch_type(
+-        const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
++        const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+         const int ncols_x, const int nrows_x, const int ncols_dst,
+         const int stride_row_x, const int stride_col_y, const int stride_col_dst,
+         const int nchannels_x, const int nchannels_y, const int nchannels_dst,
+@@ -1431,9 +1439,10 @@
+ 
+ void ggml_cuda_mul_mat_vec_q(
+         ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+-        const ggml_cuda_mm_fusion_args_host * fusion) {
+-    GGML_ASSERT(        src1->type == GGML_TYPE_F32);
+-    GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
++        const ggml_cuda_mm_fusion_args_host * fusion, const void * src1_q8_1_pre) {
++    // FreeToken: activations and results may also be bf16 or fp16
++    GGML_ASSERT(        src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16);
++    GGML_ASSERT(        dst->type  == GGML_TYPE_F32 || dst->type  == GGML_TYPE_BF16 || dst->type  == GGML_TYPE_F16);
+     GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
+ 
+     GGML_TENSOR_BINARY_OP_LOCALS;
+@@ -1526,12 +1535,22 @@
+     }
+ 
+     const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
+-    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1);
+-    {
++    // FreeToken: src1_q8_1_pre is src1 already quantized by quantize_row_q8_1_cuda with this padding and layout
++    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool());
++    const void * src1_q8 = src1_q8_1_pre;
++    if (src1_q8 == nullptr) {
++        src1_q8_1.alloc(ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1);
++        src1_q8 = src1_q8_1.get();
+         const int64_t s11 = src1->nb[1] / ts_src1;
+         const int64_t s12 = src1->nb[2] / ts_src1;
+         const int64_t s13 = src1->nb[3] / ts_src1;
+-        quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
++        if (src1->type == GGML_TYPE_BF16) {
++            quantize_row_q8_1_cuda((const nv_bfloat16 *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
++        } else if (src1->type == GGML_TYPE_F16) {
++            quantize_row_q8_1_cuda((const half *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
++        } else {
++            quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
++        }
+     }
+ 
+     const int64_t s01 = src0->nb[1] / ts_src0;
+@@ -1556,11 +1575,23 @@
+ 
+     const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0;
+ 
+-    mul_mat_vec_q_switch_type(
+-        src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00,
+-        ne01,              ncols_dst,     s01, stride_col_y,     stride_col_dst,
+-        ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst,
+-        ne03,              ne3,           s03, s13,              s3,               ids_stride, stream);
++    const auto launch = [&](auto * dst_ptr) {
++        mul_mat_vec_q_switch_type(
++            src0->data, src0->type, src1_q8, ids_d, fusion_local, dst_ptr, ne00,
++            ne01,              ncols_dst,     s01, stride_col_y,     stride_col_dst,
++            ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst,
++            ne03,              ne3,           s03, s13,              s3,               ids_stride, stream);
++    };
++    if (dst->type == GGML_TYPE_BF16) {
++        // FreeToken: the fused shared expert writes an fp32 shared_dst
++        GGML_ASSERT(!fusion || !fusion->shared_up);
++        launch((nv_bfloat16 *) dst->data);
++    } else if (dst->type == GGML_TYPE_F16) {
++        GGML_ASSERT(!fusion || !fusion->shared_up);
++        launch((half *) dst->data);
++    } else {
++        launch(dst_d);
++    }
+ }
+ 
+ void ggml_cuda_op_mul_mat_vec_q(
+--- a/ggml/src/ggml-cuda/mmvq.cuh
++++ b/ggml/src/ggml-cuda/mmvq.cuh
+@@ -8,8 +8,10 @@
+ // based on the quantization type and GPU architecture (compute capability).
+ int get_mmvq_mmid_max_batch(ggml_type type, int cc);
+ 
++// FreeToken: src1_q8_1_pre skips the internal quantization of src1
+ void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx,
+-    const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr);
++    const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr,
++    const void * src1_q8_1_pre = nullptr);
+ 
+ void ggml_cuda_op_mul_mat_vec_q(
+     ggml_backend_cuda_context & ctx,
+--- a/ggml/src/ggml-cuda/quantize.cu
++++ b/ggml/src/ggml-cuda/quantize.cu
+@@ -1,4 +1,5 @@
+ #include "quantize.cuh"
++#include "convert.cuh"
+ #include <cstdint>
+ 
+ #if defined(BLACKWELL_MMA_AVAILABLE)
+@@ -50,13 +51,14 @@
+ #endif // CUDART_VERSION >= 12080
+ #endif // defined(BLACKWELL_MMA_AVAILABLE)
+ 
++template <typename src_t>
+ __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
+ static __global__ void quantize_q8_1(
+-        const float * x_ptr, void * vy_ptr,
++        const src_t * x_ptr, void * vy_ptr,
+         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+         const int64_t ne0, const uint32_t ne1, const uint3 ne2) {
+     ggml_cuda_pdl_lc();
+-    const float * GGML_CUDA_RESTRICT x  = x_ptr;
++    const src_t * GGML_CUDA_RESTRICT x  = x_ptr;
+     void        * GGML_CUDA_RESTRICT vy = vy_ptr;
+     const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
+ 
+@@ -81,7 +83,7 @@
+     const int64_t iqs = i_cont % QK8_1; // quant index
+ 
+     ggml_cuda_pdl_sync();
+-    const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
++    const float xi = i0 < ne00 ? ggml_cuda_cast<float>(x[i03*s03 + i02*s02 + i01*s01 + i00]) : 0.0f;
+     float amax = fabsf(xi);
+     float sum = xi;
+ 
+@@ -452,10 +454,29 @@
+     GGML_UNUSED(n_expert_used);
+ }
+ 
++// FreeToken: 4 consecutive activations as float, from fp32, bf16 or fp16
++static __device__ __forceinline__ float4 load_float4(const float * x) {
++    return *(const float4 *) x;
++}
++
++static __device__ __forceinline__ float4 load_float4(const nv_bfloat16 * x) {
++    const uint2 v = *(const uint2 *) x;
++    const float2 lo = __bfloat1622float2(*(const nv_bfloat162 *) &v.x);
++    const float2 hi = __bfloat1622float2(*(const nv_bfloat162 *) &v.y);
++    return make_float4(lo.x, lo.y, hi.x, hi.y);
++}
++
++static __device__ __forceinline__ float4 load_float4(const half * x) {
++    const uint2 v = *(const uint2 *) x;
++    const float2 lo = __half22float2(*(const half2 *) &v.x);
++    const float2 hi = __half22float2(*(const half2 *) &v.y);
++    return make_float4(lo.x, lo.y, hi.x, hi.y);
++}
++
+ // scatter: grid over tokens, quantize once, write to all the token's compact rows
+-template <mmq_q8_1_ds_layout ds_layout, bool scatter>
++template <mmq_q8_1_ds_layout ds_layout, bool scatter, typename src_t>
+ static __global__ void quantize_mmq_q8_1(
+-        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
++        const src_t * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+         const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
+ 
+@@ -481,14 +502,13 @@
+         base_idx = i3*s03 + i2*s02 + i01*s01;
+     }
+ 
+-    const float4 * x4 = (const float4 *) x;
+     block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
+ 
+     const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel
+     const int64_t iqs     = i0 % QK8_1_MMQ; // quant index in block
+ 
+     // Load 4 floats per thread and calculate max. abs. value between them:
+-    const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
++    const float4 xi = i0 < ne00 ? load_float4(x + base_idx + i00) : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+     float amax = fabsf(xi.x);
+     amax = fmaxf(amax, fabsf(xi.y));
+     amax = fmaxf(amax, fabsf(xi.z));
+@@ -554,8 +574,9 @@
+     GGML_UNUSED(n_expert_used);
+ }
+ 
++template <typename src_t>
+ void quantize_row_q8_1_cuda(
+-        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
++        const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+     GGML_ASSERT(!ids);
+@@ -567,12 +588,13 @@
+     const dim3 num_blocks(block_num_x, ne1, ne2*ne3);
+     const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);
+     const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream);
+-    ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
++    ggml_cuda_kernel_launch(quantize_q8_1<src_t>, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
+     GGML_UNUSED(type_src0);
+ }
+ 
++template <typename src_t>
+ void quantize_mmq_q8_1_cuda(
+-        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
++        const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+     GGML_ASSERT(ne00 % 4 == 0);
+@@ -584,15 +606,15 @@
+     const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+     switch (mmq_get_q8_1_ds_layout(type_src0)) {
+         case MMQ_Q8_1_DS_LAYOUT_D4:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false>
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false, src_t>
+                 <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+             break;
+         case MMQ_Q8_1_DS_LAYOUT_DS4:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false>
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false, src_t>
+                 <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+             break;
+         case MMQ_Q8_1_DS_LAYOUT_D2S6:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false>
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false, src_t>
+                 <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+             break;
+         default:
+@@ -601,8 +623,9 @@
+ }
+ 
+ // scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
++template <typename src_t>
+ void quantize_scatter_mmq_q8_1_cuda(
+-        const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
++        const src_t * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+         const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+         const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+     GGML_ASSERT(ne00 % 4 == 0);
+@@ -613,15 +636,15 @@
+     const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+     switch (mmq_get_q8_1_ds_layout(type_src0)) {
+         case MMQ_Q8_1_DS_LAYOUT_D4:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true><<<num_blocks, block_size, 0, stream>>>(
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                 x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+             break;
+         case MMQ_Q8_1_DS_LAYOUT_DS4:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true><<<num_blocks, block_size, 0, stream>>>(
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                 x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+             break;
+         case MMQ_Q8_1_DS_LAYOUT_D2S6:
+-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true><<<num_blocks, block_size, 0, stream>>>(
++            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                 x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+             break;
+         default:
+@@ -692,3 +715,15 @@
+         quantize_mmq_mxfp4<false><<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+     }
+ }
++
++#define FREETOKEN_QUANTIZE_INSTANCES(src_t)                                                                              \
++    template void quantize_row_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t,    \
++        int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t);                                             \
++    template void quantize_mmq_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t,    \
++        int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t);                                             \
++    template void quantize_scatter_mmq_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t,     \
++        int64_t, int64_t, int64_t, int64_t, int, cudaStream_t);
++
++FREETOKEN_QUANTIZE_INSTANCES(float)
++FREETOKEN_QUANTIZE_INSTANCES(nv_bfloat16)
++FREETOKEN_QUANTIZE_INSTANCES(half)
+--- a/ggml/src/ggml-cuda/quantize.cuh
++++ b/ggml/src/ggml-cuda/quantize.cuh
+@@ -16,13 +16,16 @@
+         ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+         int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+ 
++// FreeToken: src_t is float, nv_bfloat16 or half
++template <typename src_t>
+ void quantize_row_q8_1_cuda(
+-        const float * x, const int32_t * ids, void * vy,
++        const src_t * x, const int32_t * ids, void * vy,
+         ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+         int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+ 
++template <typename src_t>
+ void quantize_mmq_q8_1_cuda(
+-        const float * x, const int32_t * ids, void * vy,
++        const src_t * x, const int32_t * ids, void * vy,
+         ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+         int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+ 
+@@ -57,7 +60,8 @@
+                                    int             n_expert_used,
+                                    cudaStream_t    stream);
+ 
+-void quantize_scatter_mmq_q8_1_cuda(const float *   x,
++template <typename src_t>
++void quantize_scatter_mmq_q8_1_cuda(const src_t *   x,
+                                     const int32_t * ids_src1_inv,
+                                     void *          vy,
+                                     ggml_type       type_src0,
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
new file mode 100644
index 0000000000..cc735c3b0e
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
@@ -0,0 +1,426 @@
+// tvm-ffi ops over the vendored ggml-cuda kernels. Weights stay in their packed ggml blocks.
+// Activations go through the same q8_1 quantization and kernel choice as llama.cpp's mul_mat / mul_mat_id.
+// Activations are fp32, bf16 or fp16 with contiguous rows. Results have the activation dtype and are accumulated in fp32; the caller allocates every output.
+
+#include <freetoken/utils.h>
+
+#include "common.cuh"
+#include "convert.cuh"
+#include "getrows.cuh"
+#include "mmq.cuh"
+#include "mmvq.cuh"
+#include "quantize.cuh"
+
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/optional.h>
+#include <tvm/ffi/string.h>
+
+#include <algorithm>
+#include <array>
+#include <memory>
+#include <string>
+
+using tvm::ffi::Optional;
+using tvm::ffi::String;
+using tvm::ffi::TensorView;
+
+namespace {
+
+using host::RuntimeCheck;
+
+struct DeviceGuard {
+    int prev, device;
+
+    explicit DeviceGuard(DLDevice d) : prev(ggml_cuda_get_device()), device(d.device_id) {
+#ifdef GGML_USE_HIP
+        RuntimeCheck(d.device_type == kDLROCM, "ggml ops take ROCm tensors");
+#else
+        RuntimeCheck(d.device_type == kDLCUDA, "ggml ops take CUDA tensors");
+#endif
+        ggml_cuda_set_device(device);
+    }
+
+    ~DeviceGuard() {
+        if (prev != device) {
+            cudaSetDevice(prev);
+        }
+    }
+};
+
+// ggml treats a null stream as unset and would create its own; name the default stream explicitly
+cudaStream_t current_stream(DLDevice d) {
+    auto stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(d.device_type, d.device_id));
+#ifdef GGML_USE_HIP
+    return stream == nullptr ? hipStreamLegacy : stream;
+#else
+    return stream == nullptr ? cudaStreamLegacy : stream;
+#endif
+}
+
+ggml_backend_cuda_context & context(DLDevice d) {
+    static std::unique_ptr<ggml_backend_cuda_context> ctxs[GGML_CUDA_MAX_DEVICES];
+    const int dev = d.device_id;
+    if (!ctxs[dev]) {
+        ctxs[dev] = std::make_unique<ggml_backend_cuda_context>(dev);
+    }
+    ctxs[dev]->curr_stream_no = 0;
+    ctxs[dev]->streams[dev][0] = current_stream(d);
+    return *ctxs[dev];
+}
+
+ggml_tensor describe(ggml_type type, const void * data, std::array<int64_t, 4> ne, std::array<size_t, 4> nb) {
+    ggml_tensor t{};
+    t.type = type;
+    t.data = const_cast<void *>(data);
+    for (int i = 0; i < 4; ++i) {
+        t.ne[i] = ne[i];
+        t.nb[i] = nb[i];
+    }
+    return t;
+}
+
+int device_cc() { return ggml_cuda_info().devices[ggml_cuda_get_device()].cc; }
+
+bool is_dtype(DLDataType dt, DLDataTypeCode code, int bits) { return dt.code == code && dt.bits == bits && dt.lanes == 1; }
+
+ggml_type float_type(TensorView t) {
+    const DLDataType dt = t.dtype();
+    if (is_dtype(dt, kDLFloat, 32)) {
+        return GGML_TYPE_F32;
+    }
+    if (is_dtype(dt, kDLBfloat, 16)) {
+        return GGML_TYPE_BF16;
+    }
+    RuntimeCheck(is_dtype(dt, kDLFloat, 16), "ggml tensors must be fp32, bf16 or fp16");
+    return GGML_TYPE_F16;
+}
+
+// the activation dtype, which the result must share
+ggml_type io_type(TensorView x, TensorView out) {
+    const ggml_type type = float_type(x);
+    RuntimeCheck(float_type(out) == type, "the result must have the activation dtype");
+    return type;
+}
+
+int64_t check_weight(TensorView w, ggml_type type, int64_t k) {
+    RuntimeCheck(is_dtype(w.dtype(), kDLUInt, 8), "ggml weight must be a uint8 tensor");
+    RuntimeCheck(w.stride(w.ndim() - 1) == 1, "ggml weight rows must be contiguous");
+    const int64_t bs = ggml_blck_size(type);
+    RuntimeCheck(k % bs == 0, "K=", k, " is not a multiple of the ", ggml_type_name(type), " block size ", bs);
+    const int64_t row_bytes = k / bs * (int64_t) ggml_type_size(type);
+    RuntimeCheck(w.size(w.ndim() - 1) == row_bytes, "ggml weight row is ", w.size(w.ndim() - 1), " bytes, ", ggml_type_name(type),
+                 " K=", k, " needs ", row_bytes);
+    return row_bytes;
+}
+
+int64_t q8_1_row_bytes(int64_t k) { return GGML_PAD(k, MATRIX_ROW_PADDING) / QK8_1 * (int64_t) sizeof(block_q8_1); }
+
+// MMQ tiles read up to MATRIX_ROW_PADDING values past the last row; the caller keeps this many bytes after it
+int64_t mmq_tail_bytes(int64_t type_id, int64_t k) {
+    const int64_t rem = k % MATRIX_ROW_PADDING;
+    return rem == 0 ? 0 : (int64_t) ggml_row_size((ggml_type) type_id, MATRIX_ROW_PADDING - rem);
+}
+
+struct Activation {
+    const void * data;
+    ggml_type type;
+    size_t esz;
+};
+
+// The native fp4 MMQ path quantizes fp32 activations only; it is built only with GGML_CUDA_FP4_MMA.
+bool mmq_reads_16bit(ggml_type type) {
+#ifdef GGML_CUDA_FP4_MMA
+    return type != GGML_TYPE_MXFP4 && type != GGML_TYPE_NVFP4;
+#else
+    GGML_UNUSED(type);
+    return true;
+#endif
+}
+
+// x as the kernel reads it: x itself, or an fp32 copy in ``tmp`` for a kernel that reads fp32 only
+Activation activation(TensorView x, bool f16_ok, ggml_cuda_pool_alloc<float> & tmp, cudaStream_t stream) {
+    RuntimeCheck(x.IsContiguous(), "ggml activations must be contiguous");
+    const ggml_type type = float_type(x);
+    if (type == GGML_TYPE_F32) {
+        return {x.data_ptr(), type, 4};
+    }
+    if (f16_ok) {
+        return {x.data_ptr(), type, 2};
+    }
+    float * f32 = tmp.alloc(x.numel());
+    ggml_get_to_fp32_cuda(type)(x.data_ptr(), f32, x.numel(), stream);
+    return {f32, GGML_TYPE_F32, 4};
+}
+
+}  // namespace
+
+// out[M, q8_1_row_bytes(K)] = x[M, K] as MMVQ's q8_1 blocks, for GEMVs that share one input
+void op_quantize_q8_1(TensorView x, TensorView out) {
+    RuntimeCheck(x.ndim() == 2 && x.stride(1) == 1, "quantize_q8_1 takes x[M, K] with contiguous rows");
+    const DeviceGuard guard(x.device());
+    const int64_t m = x.size(0), k = x.size(1), k_padded = GGML_PAD(k, MATRIX_ROW_PADDING);
+    RuntimeCheck(is_dtype(out.dtype(), kDLUInt, 8) && out.IsContiguous() && out.ndim() == 2 && out.size(0) == m &&
+                     out.size(1) == q8_1_row_bytes(k),
+                 "out must be a contiguous uint8 [M, q8_1_row_bytes(K)] tensor");
+    if (m == 0) {
+        return;
+    }
+    const int64_t s1 = x.stride(0);
+    const cudaStream_t stream = current_stream(x.device());
+    const ggml_type type = float_type(x);
+    if (type == GGML_TYPE_BF16) {
+        quantize_row_q8_1_cuda((const nv_bfloat16 *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream);
+    } else if (type == GGML_TYPE_F16) {
+        quantize_row_q8_1_cuda((const half *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream);
+    } else {
+        quantize_row_q8_1_cuda((const float *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream);
+    }
+}
+
+// out[M, N] = x[M, K] @ W[N, K]^T; out may be a view with strided rows; MMVQ reads ``xq`` (op_quantize_q8_1 of x) instead of quantizing x again
+void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Optional<TensorView> xq) {
+    RuntimeCheck(x.ndim() == 2 && w.ndim() == 2, "mul_mat takes w[N, row_bytes] and x[M, K]");
+    const DeviceGuard guard(x.device());
+    const ggml_type type = (ggml_type) type_id;
+    const int64_t m = x.size(0), k = x.size(1), n = w.size(0);
+    check_weight(w, type, k);
+    RuntimeCheck(out.ndim() == 2 && out.size(0) == m && out.size(1) == n && out.stride(1) == 1, "mul_mat out must be a [M, N] view with contiguous rows");
+    const ggml_type otype = io_type(x, out);
+    if (m == 0) {
+        return;
+    }
+    const size_t s1 = w.stride(0);
+    ggml_tensor src0 = describe(type, w.data_ptr(), {k, n, 1, 1}, {ggml_type_size(type), s1, s1 * n, s1 * n});
+    const size_t oe = ggml_type_size(otype), os = out.stride(0) * oe;
+    ggml_tensor dst = describe(otype, out.data_ptr(), {n, m, 1, 1}, {oe, os, os * m, os * m});
+    auto & ctx = context(x.device());
+    const cudaStream_t stream = ctx.stream();
+    const int cc = device_cc();
+    ggml_cuda_pool_alloc<float> x32(ctx.pool());
+    if (ggml_cuda_should_use_mmvq(type, cc, m)) {
+        const Activation a = activation(x, true, x32, stream);
+        ggml_tensor src1 = describe(a.type, a.data, {k, m, 1, 1}, {a.esz, a.esz * k, a.esz * k * m, a.esz * k * m});
+        const void * q8 = nullptr;
+        if (xq.has_value()) {
+            const TensorView q = xq.value();
+            RuntimeCheck(is_dtype(q.dtype(), kDLUInt, 8) && q.IsContiguous() && q.ndim() == 2 && q.size(0) == m && q.size(1) == q8_1_row_bytes(k),
+                         "xq must be quantize_q8_1 of x");
+            q8 = q.data_ptr();
+        }
+        ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, nullptr, &dst, nullptr, q8);
+    } else if (ggml_cuda_should_use_mmq(type, cc, m, 0)) {
+        const Activation a = activation(x, mmq_reads_16bit(type), x32, stream);
+        ggml_tensor src1 = describe(a.type, a.data, {k, m, 1, 1}, {a.esz, a.esz * k, a.esz * k * m, a.esz * k * m});
+        ggml_cuda_mul_mat_q(ctx, &src0, &src1, nullptr, &dst);
+    } else {
+        RuntimeCheck(false, "no ggml kernel for ", ggml_type_name(type), " with M=", m);
+    }
+}
+
+// out[T, k, N]: row (t, j) multiplies x[t, j or 0] by expert ids[t, j] of W[E, N, K]; the matrix kernel runs a flat grid over the tiles that hold routes
+void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out) {
+    RuntimeCheck(w.ndim() == 3, "mul_mat_id takes w[E, N, row_bytes]");
+    RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
+    const DeviceGuard guard(x.device());
+    const ggml_type type = (ggml_type) type_id;
+    const int64_t t = ids.size(0), top_k = ids.size(1);
+    RuntimeCheck(x.ndim() == 3 && x.size(0) == t && (x.size(1) == 1 || x.size(1) == top_k), "x must be [T, 1, K] or [T, k, K]");
+    const int64_t ky = x.size(1), k = x.size(2), e = w.size(0), n = w.size(1);
+    check_weight(w, type, k);
+    RuntimeCheck(out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n && out.IsContiguous(),
+                 "out must be a contiguous [T, k, N] tensor");
+    const ggml_type otype = io_type(x, out);
+    if (t == 0) {
+        return;
+    }
+    const size_t s1 = w.stride(1), s2 = w.stride(0);
+    ggml_tensor src0 = describe(type, w.data_ptr(), {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e});
+    const size_t ids_stride = sizeof(int32_t) * top_k, ids_bytes = ids_stride * t;
+    ggml_tensor tids = describe(GGML_TYPE_I32, ids.data_ptr(), {top_k, t, 1, 1}, {sizeof(int32_t), ids_stride, ids_bytes, ids_bytes});
+    auto & ctx = context(x.device());
+    const cudaStream_t stream = ctx.stream();
+    const int cc = device_cc();
+    const bool mmvq = t <= MMVQ_MAX_BATCH_SIZE && t <= get_mmvq_mmid_max_batch(type, cc);
+    RuntimeCheck(mmvq || ggml_cuda_should_use_mmq(type, cc, t, e), "no ggml MoE kernel for ", ggml_type_name(type), " with T=", t);
+    ggml_cuda_pool_alloc<float> x32(ctx.pool());
+    const Activation a = activation(x, mmvq || mmq_reads_16bit(type), x32, stream);
+    ggml_tensor src1 = describe(a.type, a.data, {k, ky, t, 1}, {a.esz, a.esz * k, a.esz * k * ky, a.esz * k * ky * t});
+    const size_t oe = ggml_type_size(otype);
+    ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t});
+    if (mmvq) {
+        ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst);
+    } else {
+        ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0);
+    }
+}
+
+namespace {
+
+// The GLUs the MMVQ epilogue computes; any other ggml_glu_op falls through to a plain gate * up there.
+ggml_glu_op fused_glu_op(const String & name) {
+    if (name == "swiglu") {
+        return GGML_GLU_OP_SWIGLU;
+    }
+    if (name == "geglu") {
+        return GGML_GLU_OP_GEGLU;
+    }
+    if (name == "swiglu_oai") {
+        return GGML_GLU_OP_SWIGLU_OAI;
+    }
+    RuntimeCheck(name == "swiglu_clamp", "MMVQ fuses swiglu, geglu, swiglu_oai or swiglu_clamp, not ", std::string(name));
+    return GGML_GLU_OP_SWIGLU_CLAMP;
+}
+
+}  // namespace
+
+// out[T, k, I] = glu(x @ gate^T, x @ up^T) for small T in one MMVQ pass; w[E, 2I, row_bytes] holds each expert's gate rows then its up rows; limit is swiglu_clamp's
+void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit) {
+    const ggml_glu_op glu_op = fused_glu_op(glu);
+    RuntimeCheck(w.ndim() == 3 && w.size(1) % 2 == 0, "mul_mat_id_glu takes w[E, 2I, row_bytes]");
+    RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
+    RuntimeCheck(x.ndim() == 2 && x.size(0) == ids.size(0), "x must be [T, K]");
+    const DeviceGuard guard(x.device());
+    const ggml_type type = (ggml_type) type_id;
+    const int64_t t = ids.size(0), top_k = ids.size(1), k = x.size(1), e = w.size(0), n = w.size(1) / 2;
+    check_weight(w, type, k);
+    RuntimeCheck(out.IsContiguous() && out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n,
+                 "out must be a contiguous [T, k, I] tensor");
+    const ggml_type otype = io_type(x, out);
+    const int cc = device_cc();
+    RuntimeCheck(t <= MMVQ_MAX_BATCH_SIZE && t <= get_mmvq_mmid_max_batch(type, cc), "mul_mat_id_glu is the small-batch path, T=", t);
+    if (t == 0) {
+        return;
+    }
+    auto & ctx = context(x.device());
+    ggml_cuda_pool_alloc<float> x32(ctx.pool());
+    const Activation a = activation(x, true, x32, ctx.stream());
+    const size_t s1 = w.stride(1), s2 = w.stride(0);
+    const char * base = (const char *) w.data_ptr();
+    ggml_tensor up = describe(type, base + n * s1, {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e});
+    ggml_tensor gate = describe(type, base, {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e});
+    ggml_tensor src1 = describe(a.type, a.data, {k, 1, t, 1}, {a.esz, a.esz * k, a.esz * k, a.esz * k * t});
+    const size_t ids_stride = sizeof(int32_t) * top_k, ids_bytes = ids_stride * t;
+    ggml_tensor tids = describe(GGML_TYPE_I32, ids.data_ptr(), {top_k, t, 1, 1}, {sizeof(int32_t), ids_stride, ids_bytes, ids_bytes});
+    const size_t oe = ggml_type_size(otype);
+    ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t});
+    ggml_cuda_mm_fusion_args_host fusion{};
+    fusion.gate = &gate;
+    fusion.glu_op = glu_op;
+    fusion.glu_limit = (float) limit;
+    ggml_cuda_mul_mat_vec_q(ctx, &up, &src1, &tids, &dst, &fusion);
+}
+
+namespace {
+
+template <typename T>
+__global__ void moe_weighted_sum_kernel(const T * __restrict__ y, const float * __restrict__ w, T * __restrict__ out,
+                                        const int64_t n, const int top_k) {
+    const int64_t token = blockIdx.x;
+    const int64_t col = (int64_t) blockIdx.y * blockDim.x + threadIdx.x;
+    if (col >= n) {
+        return;
+    }
+    const int64_t row0 = token * top_k;
+    float sum = 0.0f;
+    for (int j = 0; j < top_k; ++j) {
+        sum += ggml_cuda_cast<float>(y[(row0 + j) * n + col]) * w[row0 + j];
+    }
+    out[token * n + col] = ggml_cuda_cast<T>(sum);
+}
+
+}  // namespace
+
+// out[T, N] = sum_j weights[t, j] * y[t, j, :], summed in fp32
+void op_moe_weighted_sum(TensorView y, TensorView weights, TensorView out) {
+    RuntimeCheck(y.ndim() == 3 && y.IsContiguous(), "y must be a contiguous [T, k, N] tensor");
+    const int64_t t = y.size(0), top_k = y.size(1), n = y.size(2);
+    RuntimeCheck(float_type(weights) == GGML_TYPE_F32 && weights.IsContiguous() && weights.numel() == t * top_k, "weights must be a contiguous fp32 [T, k] tensor");
+    RuntimeCheck(out.ndim() == 2 && out.size(0) == t && out.size(1) == n && out.IsContiguous(), "out must be a contiguous [T, N] tensor");
+    const ggml_type otype = io_type(y, out);
+    const DeviceGuard guard(y.device());
+    if (t == 0) {
+        return;
+    }
+    constexpr int threads = 256;
+    const dim3 blocks(t, (n + threads - 1) / threads);
+    const cudaStream_t stream = current_stream(y.device());
+    const float * wp = (const float *) weights.data_ptr();
+    if (otype == GGML_TYPE_BF16) {
+        moe_weighted_sum_kernel<<<blocks, threads, 0, stream>>>((const nv_bfloat16 *) y.data_ptr(), wp, (nv_bfloat16 *) out.data_ptr(), n, top_k);
+    } else if (otype == GGML_TYPE_F16) {
+        moe_weighted_sum_kernel<<<blocks, threads, 0, stream>>>((const half *) y.data_ptr(), wp, (half *) out.data_ptr(), n, top_k);
+    } else {
+        moe_weighted_sum_kernel<<<blocks, threads, 0, stream>>>((const float *) y.data_ptr(), wp, (float *) out.data_ptr(), n, top_k);
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
+// out[rows, K] = W[rows, row_bytes] dequantized to out's dtype
+void op_dequantize(TensorView w, int64_t type_id, int64_t k, TensorView out) {
+    const DeviceGuard guard(w.device());
+    const ggml_type type = (ggml_type) type_id;
+    check_weight(w, type, k);
+    RuntimeCheck(w.IsContiguous(), "dequantize needs a contiguous weight");
+    const int64_t rows = w.numel() / w.size(w.ndim() - 1);
+    RuntimeCheck(out.IsContiguous() && out.numel() == rows * k, "out must be a contiguous [rows, K] tensor");
+    const ggml_type otype = float_type(out);
+    if (rows == 0) {
+        return;
+    }
+    const cudaStream_t stream = current_stream(w.device());
+    if (type == otype) {
+        CUDA_CHECK(cudaMemcpyAsync(out.data_ptr(), w.data_ptr(), rows * k * ggml_type_size(type), cudaMemcpyDeviceToDevice, stream));
+        return;
+    }
+    if (otype == GGML_TYPE_BF16) {
+        const auto convert = ggml_get_to_bf16_cuda(type);
+        RuntimeCheck(convert != nullptr, "no ggml bf16 conversion for ", ggml_type_name(type));
+        convert(w.data_ptr(), (nv_bfloat16 *) out.data_ptr(), rows * k, stream);
+    } else if (otype == GGML_TYPE_F16) {
+        const auto convert = ggml_get_to_fp16_cuda(type);
+        RuntimeCheck(convert != nullptr, "no ggml fp16 conversion for ", ggml_type_name(type));
+        convert(w.data_ptr(), (half *) out.data_ptr(), rows * k, stream);
+    } else {
+        const auto convert = ggml_get_to_fp32_cuda(type);
+        RuntimeCheck(convert != nullptr, "no ggml fp32 conversion for ", ggml_type_name(type));
+        convert(w.data_ptr(), (float *) out.data_ptr(), rows * k, stream);
+    }
+}
+
+// out[i] = dequant(W[ids[i]]) for a table W[rows, row_bytes]
+void op_get_rows(TensorView w, int64_t type_id, int64_t k, TensorView ids, TensorView out) {
+    RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 1 && ids.IsContiguous(), "ids must be a contiguous int32 vector");
+    RuntimeCheck(w.ndim() == 2, "get_rows takes w[rows, row_bytes]");
+    const DeviceGuard guard(w.device());
+    const ggml_type type = (ggml_type) type_id;
+    check_weight(w, type, k);
+    const int64_t n = ids.size(0);
+    RuntimeCheck(out.IsContiguous() && out.ndim() == 2 && out.size(0) == n && out.size(1) == k, "out must be a contiguous [n, K] tensor");
+    const ggml_type otype = float_type(out);
+    if (n == 0) {
+        return;
+    }
+    const size_t nb01 = w.stride(0), esz = otype == GGML_TYPE_F32 ? 4 : 2;
+    get_rows_cuda(w.data_ptr(), type, (const int32_t *) ids.data_ptr(), out.data_ptr(), otype, k, nb01, nb01 * w.size(0), nb01 * w.size(0),
+                  n, 1, 1, 4, 4 * n, 4 * n, k * esz, k * esz * n, k * esz * n, current_stream(w.device()));
+}
+
+bool op_uses_mmvq(int64_t type_id, int64_t m) { return ggml_cuda_should_use_mmvq((ggml_type) type_id, device_cc(), m); }
+
+int64_t op_mmvq_max_tokens(int64_t type_id) {
+    return std::min<int64_t>(MMVQ_MAX_BATCH_SIZE, get_mmvq_mmid_max_batch((ggml_type) type_id, device_cc()));
+}
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(quantize_q8_1, op_quantize_q8_1);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(q8_1_row_bytes, q8_1_row_bytes);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(mmq_tail_bytes, mmq_tail_bytes);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat, op_mul_mat);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat_id, op_mul_mat_id);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat_id_glu, op_mul_mat_id_glu);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_weighted_sum, op_moe_weighted_sum);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(dequantize, op_dequantize);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_rows, op_get_rows);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(uses_mmvq, op_uses_mmvq);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(mmvq_max_tokens, op_mmvq_max_tokens);
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu
new file mode 100644
index 0000000000..80ef72855c
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu
@@ -0,0 +1,255 @@
+// The ggml-base and ggml-cuda.cu host symbols that the vendored ggml-cuda kernels reference.
+// Tensor descriptors built by ggml_ops.cu never have a backend buffer, and temporary buffers come from tvm-ffi's env allocator (torch's caching allocator under torch).
+
+#include "common.cuh"
+
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <cstdarg>
+#include <cstdio>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+
+namespace {
+
+struct TypeInfo {
+    const char * name = nullptr;
+    int64_t blck_size = 0;
+    size_t type_size = 0;
+    bool quantized = false;
+};
+
+const TypeInfo * type_infos() {
+    static TypeInfo t[GGML_TYPE_COUNT];
+    static const bool init = [] {
+        auto set = [](ggml_type type, const char * name, int64_t blck, size_t size, bool quantized) {
+            t[type] = {name, blck, size, quantized};
+        };
+        set(GGML_TYPE_F32, "f32", 1, sizeof(float), false);
+        set(GGML_TYPE_F16, "f16", 1, sizeof(ggml_fp16_t), false);
+        set(GGML_TYPE_BF16, "bf16", 1, sizeof(ggml_bf16_t), false);
+        set(GGML_TYPE_F64, "f64", 1, sizeof(double), false);
+        set(GGML_TYPE_I8, "i8", 1, sizeof(int8_t), false);
+        set(GGML_TYPE_I16, "i16", 1, sizeof(int16_t), false);
+        set(GGML_TYPE_I32, "i32", 1, sizeof(int32_t), false);
+        set(GGML_TYPE_I64, "i64", 1, sizeof(int64_t), false);
+        set(GGML_TYPE_Q1_0, "q1_0", QK1_0, sizeof(block_q1_0), true);
+        set(GGML_TYPE_Q2_0, "q2_0", QK2_0, sizeof(block_q2_0), true);
+        set(GGML_TYPE_Q4_0, "q4_0", QK4_0, sizeof(block_q4_0), true);
+        set(GGML_TYPE_Q4_1, "q4_1", QK4_1, sizeof(block_q4_1), true);
+        set(GGML_TYPE_Q5_0, "q5_0", QK5_0, sizeof(block_q5_0), true);
+        set(GGML_TYPE_Q5_1, "q5_1", QK5_1, sizeof(block_q5_1), true);
+        set(GGML_TYPE_Q8_0, "q8_0", QK8_0, sizeof(block_q8_0), true);
+        set(GGML_TYPE_Q8_1, "q8_1", QK8_1, sizeof(block_q8_1), true);
+        set(GGML_TYPE_MXFP4, "mxfp4", QK_MXFP4, sizeof(block_mxfp4), true);
+        set(GGML_TYPE_NVFP4, "nvfp4", QK_NVFP4, sizeof(block_nvfp4), true);
+        set(GGML_TYPE_TQ1_0, "tq1_0", QK_K, sizeof(block_tq1_0), true);
+        set(GGML_TYPE_TQ2_0, "tq2_0", QK_K, sizeof(block_tq2_0), true);
+        set(GGML_TYPE_Q2_K, "q2_K", QK_K, sizeof(block_q2_K), true);
+        set(GGML_TYPE_Q3_K, "q3_K", QK_K, sizeof(block_q3_K), true);
+        set(GGML_TYPE_Q4_K, "q4_K", QK_K, sizeof(block_q4_K), true);
+        set(GGML_TYPE_Q5_K, "q5_K", QK_K, sizeof(block_q5_K), true);
+        set(GGML_TYPE_Q6_K, "q6_K", QK_K, sizeof(block_q6_K), true);
+        set(GGML_TYPE_Q8_K, "q8_K", QK_K, sizeof(block_q8_K), true);
+        set(GGML_TYPE_IQ2_XXS, "iq2_xxs", QK_K, sizeof(block_iq2_xxs), true);
+        set(GGML_TYPE_IQ2_XS, "iq2_xs", QK_K, sizeof(block_iq2_xs), true);
+        set(GGML_TYPE_IQ2_S, "iq2_s", QK_K, sizeof(block_iq2_s), true);
+        set(GGML_TYPE_IQ3_XXS, "iq3_xxs", QK_K, sizeof(block_iq3_xxs), true);
+        set(GGML_TYPE_IQ3_S, "iq3_s", QK_K, sizeof(block_iq3_s), true);
+        set(GGML_TYPE_IQ1_S, "iq1_s", QK_K, sizeof(block_iq1_s), true);
+        set(GGML_TYPE_IQ1_M, "iq1_m", QK_K, sizeof(block_iq1_m), true);
+        set(GGML_TYPE_IQ4_NL, "iq4_nl", QK4_NL, sizeof(block_iq4_nl), true);
+        set(GGML_TYPE_IQ4_XS, "iq4_xs", QK_K, sizeof(block_iq4_xs), true);
+        return true;
+    }();
+    (void) init;
+    return t;
+}
+
+const TypeInfo & type_info(ggml_type type) {
+    if (type < 0 || type >= GGML_TYPE_COUNT || type_infos()[type].name == nullptr) {
+        throw std::invalid_argument("unknown ggml type " + std::to_string((int) type));
+    }
+    return type_infos()[type];
+}
+
+struct EnvPool : ggml_cuda_pool {
+    explicit EnvPool(int device) : device(device) {}
+
+    void * alloc(size_t size, size_t * actual_size) override {
+        *actual_size = size;
+        if (size == 0) {
+            return nullptr;
+        }
+#ifdef GGML_USE_HIP
+        constexpr DLDeviceType device_type = kDLROCM;
+#else
+        constexpr DLDeviceType device_type = kDLCUDA;
+#endif
+        auto buf = tvm::ffi::Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, tvm::ffi::Shape({(int64_t) size}), DLDataType{kDLUInt, 8, 1},
+                                                DLDevice{device_type, device});
+        void * ptr = buf.data_ptr();
+        live.emplace(ptr, std::move(buf));
+        return ptr;
+    }
+
+    void free(void * ptr, size_t) override { live.erase(ptr); }
+
+    int device;
+    std::unordered_map<void *, tvm::ffi::Tensor> live;
+};
+
+}  // namespace
+
+void ggml_abort(const char * file, int line, const char * fmt, ...) {
+    char msg[1024];
+    va_list args;
+    va_start(args, fmt);
+    vsnprintf(msg, sizeof(msg), fmt, args);
+    va_end(args);
+    throw std::runtime_error(std::string("ggml: ") + file + ":" + std::to_string(line) + ": " + msg);
+}
+
+void ggml_log_internal(enum ggml_log_level level, const char * format, ...) {
+    if (level < GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_CONT) {
+        return;
+    }
+    va_list args;
+    va_start(args, format);
+    vfprintf(stderr, format, args);
+    va_end(args);
+}
+
+int64_t ggml_blck_size(enum ggml_type type) { return type_info(type).blck_size; }
+
+size_t ggml_type_size(enum ggml_type type) { return type_info(type).type_size; }
+
+const char * ggml_type_name(enum ggml_type type) { return type_info(type).name; }
+
+bool ggml_is_quantized(enum ggml_type type) { return type_info(type).quantized; }
+
+size_t ggml_row_size(enum ggml_type type, int64_t ne) { return ggml_type_size(type) * ne / ggml_blck_size(type); }
+
+size_t ggml_element_size(const struct ggml_tensor * tensor) { return ggml_type_size(tensor->type); }
+
+int64_t ggml_nelements(const struct ggml_tensor * tensor) {
+    return tensor->ne[0] * tensor->ne[1] * tensor->ne[2] * tensor->ne[3];
+}
+
+size_t ggml_nbytes(const struct ggml_tensor * tensor) {
+    for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+        if (tensor->ne[i] <= 0) {
+            return 0;
+        }
+    }
+    size_t nbytes;
+    const size_t blck_size = ggml_blck_size(tensor->type);
+    if (blck_size == 1) {
+        nbytes = ggml_type_size(tensor->type);
+        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+            nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    } else {
+        nbytes = tensor->ne[0] * tensor->nb[0] / blck_size;
+        for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+            nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    }
+    return nbytes;
+}
+
+bool ggml_is_matrix(const struct ggml_tensor * tensor) { return tensor->ne[2] == 1 && tensor->ne[3] == 1; }
+
+bool ggml_is_contiguous(const struct ggml_tensor * tensor) {
+    size_t next_nb = ggml_type_size(tensor->type);
+    if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) {
+        return false;
+    }
+    next_nb *= tensor->ne[0] / ggml_blck_size(tensor->type);
+    for (int i = 1; i < GGML_MAX_DIMS; i++) {
+        if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) {
+            return false;
+        }
+        next_nb *= tensor->ne[i];
+    }
+    return true;
+}
+
+bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) {
+    return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type) / ggml_blck_size(tensor->type);
+}
+
+bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) {
+    return t0->ne[0] == t1->ne[0] && t0->ne[1] == t1->ne[1] && t0->ne[2] == t1->ne[2] && t0->ne[3] == t1->ne[3];
+}
+
+bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) {
+    return t0->nb[0] == t1->nb[0] && t0->nb[1] == t1->nb[1] && t0->nb[2] == t1->nb[2] && t0->nb[3] == t1->nb[3];
+}
+
+enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage(ggml_backend_buffer_t) {
+    return GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+}
+
+size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t, const struct ggml_tensor * tensor) {
+    return ggml_nbytes(tensor);
+}
+
+[[noreturn]] void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) {
+    throw std::runtime_error(std::string("ggml-cuda: ") + msg + " in " + func + " at " + file + ":" + std::to_string(line) + ": " + stmt);
+}
+
+int ggml_cuda_get_device() {
+    int id;
+    CUDA_CHECK(cudaGetDevice(&id));
+    return id;
+}
+
+void ggml_cuda_set_device(int device) {
+    int current;
+    CUDA_CHECK(cudaGetDevice(&current));
+    if (current != device) {
+        CUDA_CHECK(cudaSetDevice(device));
+    }
+}
+
+const ggml_cuda_device_info & ggml_cuda_info() {
+    static const ggml_cuda_device_info info = [] {
+        ggml_cuda_device_info in = {};
+        CUDA_CHECK(cudaGetDeviceCount(&in.physical_device_count));
+        GGML_ASSERT(in.physical_device_count <= GGML_CUDA_MAX_DEVICES);
+        in.device_count = in.physical_device_count;
+        for (int id = 0; id < in.device_count; ++id) {
+            cudaDeviceProp prop;
+            CUDA_CHECK(cudaGetDeviceProperties(&prop, id));
+            auto & d = in.devices[id];
+#ifdef GGML_USE_HIP
+            GGML_ASSERT(std::string(prop.gcnArchName).starts_with("gfx"));
+            d.cc = GGML_CUDA_CC_OFFSET_AMD + static_cast<int>(std::strtol(prop.gcnArchName + 3, nullptr, 16));
+#else
+            d.cc = 100 * prop.major + 10 * prop.minor;
+#endif
+            d.nsm = prop.multiProcessorCount;
+            d.smpb = prop.sharedMemPerBlock;
+            d.smpbo = prop.sharedMemPerBlockOptin;
+            d.warp_size = prop.warpSize;
+            d.total_vram = prop.totalGlobalMem;
+            d.physical_device = id;
+            d.physical_share_count = 1;
+            int coop = 0;
+            CUDA_CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, id));
+            d.supports_cooperative_launch = coop != 0;
+        }
+        return in;
+    }();
+    return info;
+}
+
+// The streams belong to the caller, so there is nothing to release.
+ggml_backend_cuda_context::~ggml_backend_cuda_context() {}
+
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, int) {
+    return std::make_unique<EnvPool>(device);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
new file mode 100644
index 0000000000..36912ce4b6
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
@@ -0,0 +1,1722 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-impl.h"
+#include "ggml-cuda.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <mutex>
+
+#if defined(GGML_USE_HIP)
+#define GGML_COMMON_DECL_HIP
+#define GGML_COMMON_IMPL_HIP
+#else
+#define GGML_COMMON_DECL_CUDA
+#define GGML_COMMON_IMPL_CUDA
+#if defined(GGML_USE_MUSA)
+#define GGML_COMMON_DECL_MUSA
+#define GGML_COMMON_IMPL_MUSA
+#endif
+#endif
+#include "ggml-common.h"
+
+#include <array>
+#include <algorithm>
+#include <cassert>
+#include <cfloat>
+#include <cstdio>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#if defined(GGML_USE_HIP)
+#include "vendors/hip.h"
+#elif defined(GGML_USE_MUSA)
+#include "vendors/musa.h"
+#else
+#include "vendors/cuda.h"
+#endif // defined(GGML_USE_HIP)
+
+#define STRINGIZE_IMPL(...) #__VA_ARGS__
+#define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__)
+
+#define WARP_SIZE 32
+#define CUDART_HMAX   11070 // CUDA 11.7, min. ver. for which __hmax and __hmax2 are known to work (may be higher than needed)
+#define CUDART_HMASK  12000 // CUDA 12.0, min. ver. for half2 -> uint mask comparisons
+
+#define GGML_CUDA_CC_PASCAL          600
+#define GGML_CUDA_CC_DP4A            610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products
+#define GGML_CUDA_CC_VOLTA           700
+#define GGML_CUDA_CC_TURING          750
+#define GGML_CUDA_CC_AMPERE          800
+#define GGML_CUDA_CC_ORIN            870
+#define GGML_CUDA_CC_ADA_LOVELACE    890
+#define GGML_CUDA_CC_HOPPER          900
+// While BW spans CC 1000, 1100 & 1200, we are integrating Tensor Core instructions available to 1200 family, see
+// https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#blackwell-sm120-gemms
+#define GGML_CUDA_CC_BLACKWELL       1200
+#define GGML_CUDA_CC_DGX_SPARK       1210
+#define GGML_CUDA_CC_RUBIN           1300
+#define GGML_CUDA_CC_OFFSET_AMD      0x1000000
+#define GGML_CUDA_CC_OFFSET_MTHREADS 0x0100000
+#define GGML_CUDA_CC_IS_NVIDIA(cc)   (cc < GGML_CUDA_CC_OFFSET_MTHREADS)
+
+// AMD
+// GCN/CDNA, wave size is 64
+#define GGML_CUDA_CC_GCN4       (GGML_CUDA_CC_OFFSET_AMD + 0x803)  // Tonga, Fiji, Polaris, minimum for fast fp16
+#define GGML_CUDA_CC_VEGA       (GGML_CUDA_CC_OFFSET_AMD + 0x900)  // Vega56/64, minimum for fp16 dual issue
+#define GGML_CUDA_CC_VEGA20     (GGML_CUDA_CC_OFFSET_AMD + 0x906)  // MI50/Radeon VII, minimum for dp4a
+#define GGML_CUDA_CC_GFX909     (GGML_CUDA_CC_OFFSET_AMD + 0x909)  // GCN APU
+#define GGML_CUDA_CC_GFX90C     (GGML_CUDA_CC_OFFSET_AMD + 0x90c)  // GCN APU
+#define GGML_CUDA_CC_CDNA1      (GGML_CUDA_CC_OFFSET_AMD + 0x908)  // MI100, minimum for MFMA, acc registers
+#define GGML_CUDA_CC_CDNA2      (GGML_CUDA_CC_OFFSET_AMD + 0x90a)  // MI210 (gfx90a), minimum acc register renaming
+#define GGML_CUDA_CC_CDNA3      (GGML_CUDA_CC_OFFSET_AMD + 0x942)  // MI300
+#define GGML_CUDA_CC_CDNA4      (GGML_CUDA_CC_OFFSET_AMD + 0x950)  // MI350X/MI355X
+
+// RDNA removes MFMA, dp4a, xnack, acc registers, wave size is 32
+#define GGML_CUDA_CC_RDNA1      (GGML_CUDA_CC_OFFSET_AMD + 0x1010) // RX 5000
+#define GGML_CUDA_CC_RDNA2      (GGML_CUDA_CC_OFFSET_AMD + 0x1030) // RX 6000, minimum for dp4a
+#define GGML_CUDA_CC_RDNA3      (GGML_CUDA_CC_OFFSET_AMD + 0x1100) // RX 7000, minimum for WMMA
+#define GGML_CUDA_CC_RDNA3_5    (GGML_CUDA_CC_OFFSET_AMD + 0x1150) // AI 370, AI Max 395 laptops.
+#define GGML_CUDA_CC_RDNA4      (GGML_CUDA_CC_OFFSET_AMD + 0x1200) // RX 9000
+
+#define GGML_CUDA_CC_IS_AMD(cc)     (cc >= GGML_CUDA_CC_OFFSET_AMD)
+#define GGML_CUDA_CC_IS_RDNA(cc)    (cc >= GGML_CUDA_CC_RDNA1)
+#define GGML_CUDA_CC_IS_RDNA1(cc)   (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2)
+#define GGML_CUDA_CC_IS_RDNA2(cc)   (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3)
+#define GGML_CUDA_CC_IS_RDNA3_0(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA3_5)
+#define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4)
+#define GGML_CUDA_CC_IS_RDNA3(cc)   (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc))
+#define GGML_CUDA_CC_IS_RDNA4(cc)   (cc >= GGML_CUDA_CC_RDNA4)
+#define GGML_CUDA_CC_IS_GCN_APU(cc) ((cc) == GGML_CUDA_CC_GFX909 || (cc) == GGML_CUDA_CC_GFX90C)
+#define GGML_CUDA_CC_IS_GCN(cc)     ((cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA1) || GGML_CUDA_CC_IS_GCN_APU(cc))
+#define GGML_CUDA_CC_IS_CDNA(cc)    (!GGML_CUDA_CC_IS_GCN_APU(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_RDNA1)
+#define GGML_CUDA_CC_IS_CDNA1(cc)   (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_CDNA2)
+#define GGML_CUDA_CC_IS_CDNA2(cc)   (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA2 && cc < GGML_CUDA_CC_CDNA3)
+#define GGML_CUDA_CC_IS_CDNA3(cc)   (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA3 && cc < GGML_CUDA_CC_CDNA4)
+#define GGML_CUDA_CC_IS_CDNA4(cc)   (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA4 && cc < GGML_CUDA_CC_RDNA1)
+
+// Moore Threads
+#define MUSART_HMASK 40300 // MUSA rc4.3, min. ver. for half2 -> uint mask comparisons
+
+#define GGML_CUDA_CC_QY1 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x210) // MTT S80, MTT S3000
+#define GGML_CUDA_CC_QY2 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x220) // MTT S4000
+#define GGML_CUDA_CC_PH1 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x310) // MTT S5000
+
+#define GGML_CUDA_CC_IS_MTHREADS(cc) (cc >= GGML_CUDA_CC_OFFSET_MTHREADS && cc < GGML_CUDA_CC_OFFSET_AMD)
+#define GGML_CUDA_CC_IS_QY1(cc)      (cc >= GGML_CUDA_CC_QY1 && cc < GGML_CUDA_CC_QY2)
+#define GGML_CUDA_CC_IS_QY2(cc)      (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1)
+#define GGML_CUDA_CC_IS_PH1(cc)      (cc >= GGML_CUDA_CC_PH1)
+
+#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)
+#    define GGML_CUDA_USE_CUB
+#endif  // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)
+
+// PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8.
+// However, this has been bugged in CTK < 12.3 for MSVC builds, see
+// https://github.com/ggml-org/llama.cpp/pull/22522#discussion_r3302393293
+// __CUDA_ARCH__  is undefined in host passes; GPU arch check happens in device-side code.
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && \
+    (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080))
+#    define GGML_CUDA_USE_PDL
+#endif  // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080))
+
+static __device__ __forceinline__ void ggml_cuda_syncwarp() {
+#ifndef GGML_USE_HIP
+    __syncwarp();
+#endif // GGML_USE_HIP
+}
+
+static __device__ __forceinline__ void ggml_cuda_pdl_sync() {
+#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+    cudaGridDependencySynchronize();
+#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+}
+
+static __device__ __forceinline__ void ggml_cuda_pdl_lc() {
+#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+}
+
+#ifdef __CUDA_ARCH_LIST__
+constexpr bool ggml_cuda_has_arch_impl(int) {
+    return false;
+}
+
+template<class ... Archs>
+constexpr bool ggml_cuda_has_arch_impl(const int arch, const int first, Archs... rest) {
+    return arch == first || ggml_cuda_has_arch_impl(arch, rest...);
+}
+
+constexpr bool ggml_cuda_has_arch(const int arch) {
+    return ggml_cuda_has_arch_impl(arch, __CUDA_ARCH_LIST__);
+}
+
+constexpr int ggml_cuda_highest_compiled_arch_impl(const int /*arch*/, const int cur) {
+    if (cur == 0) {
+        return -1;
+    }
+    return cur;
+}
+
+template<class ... Archs>
+constexpr int ggml_cuda_highest_compiled_arch_impl(const int arch, const int cur, const int first, Archs... rest) {
+    if (first <= arch && first > cur) {
+        return ggml_cuda_highest_compiled_arch_impl(arch, first, rest...);
+    } else {
+        return ggml_cuda_highest_compiled_arch_impl(arch, cur, rest...);
+    }
+}
+
+constexpr int ggml_cuda_highest_compiled_arch(const int arch) {
+    return ggml_cuda_highest_compiled_arch_impl(arch, 0, __CUDA_ARCH_LIST__);
+}
+#else
+static int ggml_cuda_highest_compiled_arch(const int arch) {
+    return arch;
+}
+#endif // __CUDA_ARCH_LIST__
+
+// ---------------------------------------------------------------------------------------------------------
+
+#define MATRIX_ROW_PADDING 512 // last row of quant. matrices is a multiple of this to avoid out-of-bounds memory accesses
+
+#define GGML_CUDA_MAX_STREAMS 8
+
+[[noreturn]]
+void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg);
+
+#define CUDA_CHECK_GEN(err, success, error_fn)                                      \
+     do {                                                                           \
+        auto err_ = (err);                                                          \
+        if (err_ != (success)) {                                                    \
+            ggml_cuda_error(#err, __func__, __FILE__, __LINE__, error_fn(err_));    \
+        }                                                                           \
+    } while (0)
+
+#define CUDA_CHECK(err) CUDA_CHECK_GEN(err, cudaSuccess, cudaGetErrorString)
+
+
+#if CUDART_VERSION >= 12000 || defined(GGML_USE_MUSA)
+    static const char * cublas_get_error_str(const cublasStatus_t err) {
+        return cublasGetStatusString(err);
+    }
+#else
+    static const char * cublas_get_error_str(const cublasStatus_t err) {
+        switch (err) {
+            case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS";
+            case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED";
+            case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
+            case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
+            case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
+            case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
+            case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
+            case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
+            case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
+            default: return "unknown error";
+        }
+    }
+#endif // CUDART_VERSION >= 12000
+
+#define CUBLAS_CHECK(err) CUDA_CHECK_GEN(err, CUBLAS_STATUS_SUCCESS, cublas_get_error_str)
+
+#ifdef GGML_USE_NCCL
+#define NCCL_CHECK(err) CUDA_CHECK_GEN(err, ncclSuccess, ncclGetErrorString)
+#endif // GGML_USE_NCCL
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)
+static const char * cu_get_error_str(CUresult err) {
+    const char * err_str;
+    cuGetErrorString(err, &err_str);
+    return err_str;
+}
+#define CU_CHECK(err) CUDA_CHECK_GEN(err, CUDA_SUCCESS, cu_get_error_str)
+#endif
+
+#if !defined(GGML_USE_HIP)
+#    define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes)                                                       \
+        do {                                                                                                   \
+            static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = { false };                         \
+            const int   id                                                = ggml_cuda_get_device();            \
+            if (!shared_memory_limit_raised[id]) {                                                             \
+                CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes)); \
+                shared_memory_limit_raised[id] = true;                                                         \
+            }                                                                                                  \
+        } while (0)
+#else
+#    define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes) \
+        do {                                             \
+            GGML_UNUSED(nbytes);                         \
+        } while (0)
+#endif // !defined(GGML_USE_HIP)
+
+#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA)
+#define GGML_CUDA_ASSUME(x) __builtin_assume(x)
+#else
+#define GGML_CUDA_ASSUME(x)
+#endif // CUDART_VERSION >= 11010
+
+#if (!defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)) || (defined(GGML_USE_HIP) && !defined(GGML_HIP_NO_VMM))
+#define GGML_USE_VMM
+#endif // (!defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)) || (defined(GGML_USE_HIP) && !defined(GGML_HIP_NO_VMM))
+
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL
+#define FP16_AVAILABLE
+#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL
+
+#if defined(FP16_AVAILABLE) && __CUDA_ARCH__ != 610
+#define FAST_FP16_AVAILABLE
+#endif // defined(FP16_AVAILABLE) && __CUDA_ARCH__ != 610
+
+#if defined(GGML_USE_HIP) && defined(CDNA) && !defined(GGML_HIP_NO_MMQ_MFMA)
+#define AMD_MFMA_AVAILABLE
+#endif // defined(GGML_USE_HIP) && defined(CDNA) && !defined(GGML_HIP_NO_MMQ_MFMA)
+
+#if defined(GGML_USE_HIP) && (defined(RDNA4) || defined(RDNA3))
+#define AMD_WMMA_AVAILABLE
+#endif // defined(GGML_USE_HIP) && defined(RDNA4)
+
+// The Volta instructions are in principle available on Turing or newer but they are effectively unusable:
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
+#define VOLTA_MMA_AVAILABLE
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+#define TURING_MMA_AVAILABLE
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#define AMPERE_MMA_AVAILABLE
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+
+// FreeToken: the native fp4 MMA needs an sm_120a/sm_120f build, so it is off unless GGML_CUDA_FP4_MMA is defined; FP4 weights then take Q8_1 activations
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN && defined(GGML_CUDA_FP4_MMA)
+#    define BLACKWELL_MMA_AVAILABLE
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#define CP_ASYNC_AVAILABLE
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+
+#if !defined(GGML_CUDA_NO_FA) && !(defined(GGML_USE_MUSA) && __MUSA_ARCH__ < 220)
+#define FLASH_ATTN_AVAILABLE
+#endif // !defined(GGML_CUDA_NO_FA) && !(defined(GGML_USE_MUSA) && __MUSA_ARCH__ < 220)
+
+static bool fp16_available(const int cc) {
+    return ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_PASCAL ||
+        (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_PH1);
+}
+
+static bool fast_fp16_available(const int cc) {
+    return GGML_CUDA_CC_IS_AMD(cc) ||
+        (GGML_CUDA_CC_IS_NVIDIA(cc) && fp16_available(cc) && ggml_cuda_highest_compiled_arch(cc) != 610) ||
+        (GGML_CUDA_CC_IS_MTHREADS(cc) && fp16_available(cc));
+}
+
+// To be used for feature selection of external libraries, e.g. cuBLAS.
+static bool fast_fp16_hardware_available(const int cc) {
+    return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_PASCAL && cc != 610) || GGML_CUDA_CC_IS_AMD(cc) ||
+        (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2);
+}
+
+// To be used for feature selection of external libraries, e.g. cuBLAS.
+static bool fp16_mma_hardware_available(const int cc) {
+    return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_VOLTA) ||
+        GGML_CUDA_CC_IS_CDNA(cc) || GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc) ||
+        (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2);
+}
+
+// To be used for feature selection of external libraries, e.g. cuBLAS.
+static bool fast_bf16_hardware_available(const int cc) {
+        return (GGML_CUDA_CC_IS_AMD(cc) && (cc >= GGML_CUDA_CC_RDNA3 || GGML_CUDA_CC_IS_CDNA(cc)))
+            || (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE);
+}
+
+static bool bf16_mma_hardware_available(const int cc) {
+    return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) ||
+        GGML_CUDA_CC_IS_CDNA(cc) || cc >= GGML_CUDA_CC_RDNA3 ||
+        (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_PH1);
+}
+
+static bool fp32_mma_hardware_available(const int cc) {
+    return GGML_CUDA_CC_IS_CDNA(cc);
+}
+
+static bool amd_mfma_available(const int cc) {
+#if !defined(GGML_HIP_NO_MMQ_MFMA)
+    return GGML_CUDA_CC_IS_CDNA(cc);
+#else
+    return false;
+#endif //!defined(GGML_HIP_NO_MMQ_MFMA)
+}
+
+static bool amd_wmma_available(const int cc) {
+    return (GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_RDNA3(cc));
+}
+
+static bool volta_mma_available(const int cc) {
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA;
+}
+
+static bool turing_mma_available(const int cc) {
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING;
+}
+
+static bool ampere_mma_available(const int cc) {
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE;
+}
+
+static bool cp_async_available(const int cc) {
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE;
+}
+
+static bool blackwell_mma_available(const int cc) {
+#ifdef GGML_CUDA_FP4_MMA
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_BLACKWELL &&
+           ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
+#else
+    GGML_UNUSED(cc);
+    return false;
+#endif // GGML_CUDA_FP4_MMA
+}
+
+// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) {
+    GGML_ASSERT(tensor != nullptr);
+    return (reinterpret_cast<uintptr_t>(tensor->data) % alignment) == 0 &&
+           tensor->nb[1] % alignment == 0 &&
+           tensor->nb[2] % alignment == 0 &&
+           tensor->nb[3] % alignment == 0;
+}
+
+static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
+#if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
+    return 64;
+#else
+    return 32;
+#endif // defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
+}
+
+// Maximum number of bytes that can be copied in a single instruction.
+static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    return 16;
+#else
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
+    return 16;
+#else
+    return 8;
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
+#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+}
+
+
+[[noreturn]]
+static __device__ void no_device_code(
+    const char * file_name, const int line, const char * function_name, const int arch, const char * arch_list) {
+
+#if defined(GGML_USE_HIP)
+    printf("%s:%d: ERROR: HIP kernel %s has no device code compatible with HIP arch %d.\n",
+           file_name, line, function_name, arch);
+    GGML_UNUSED(arch_list);
+#else
+    printf("%s:%d: ERROR: CUDA kernel %s has no device code compatible with CUDA arch %d. ggml-cuda.cu was compiled for: %s\n",
+           file_name, line, function_name, arch, arch_list);
+#endif // defined(GGML_USE_HIP)
+    __trap();
+
+    GGML_UNUSED(no_device_code); // suppress unused function warning
+}
+
+#ifdef __CUDA_ARCH__
+#define NO_DEVICE_CODE no_device_code(__FILE__, __LINE__, __FUNCTION__, __CUDA_ARCH__, STRINGIZE(__CUDA_ARCH_LIST__))
+#else
+#define NO_DEVICE_CODE //GGML_ABORT("NO_DEVICE_CODE not valid in host code.")
+#endif // __CUDA_ARCH__
+
+// The compiler is always able to unroll loops if they contain continue expressions.
+// In such cases loop unrolling can still be achieved via recursion:
+template <int n>
+struct ggml_cuda_unroll {
+    template <typename Func, typename... Args>
+    __device__ void operator()(const Func & f, Args... args) const {
+        f(n - 1, args...);
+        ggml_cuda_unroll<n - 1>{}(f, args...);
+    }
+};
+
+template <>
+struct ggml_cuda_unroll<1> {
+    template <typename Func, typename... Args>
+    __device__ void operator()(const Func & f, Args... args) const {
+        f(0, args...);
+    }
+};
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ int warp_reduce_sum(int x) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+    return __reduce_add_sync(0xffffffff, x);
+#else
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        x += __shfl_xor_sync(0xffffffff, x, offset, width);
+    }
+    return x;
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        x += __shfl_xor_sync(0xffffffff, x, offset, width);
+    }
+    return x;
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ float2 warp_reduce_sum(float2 a) {
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        a.x += __shfl_xor_sync(0xffffffff, a.x, offset, width);
+        a.y += __shfl_xor_sync(0xffffffff, a.y, offset, width);
+    }
+    return a;
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ half2 warp_reduce_sum(half2 a) {
+#ifdef FP16_AVAILABLE
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        a = __hadd2(a, __shfl_xor_sync(0xffffffff, a, offset, width));
+    }
+    return a;
+
+#else
+    NO_DEVICE_CODE;
+    return a;
+#endif // FP16_AVAILABLE
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ int warp_reduce_all(int x) {
+    if (width == ggml_cuda_get_physical_warp_size()) {
+        return __all_sync(0xffffffff, x);
+    } else {
+#pragma unroll
+        for (int offset = width/2; offset > 0; offset >>= 1) {
+            x = __shfl_xor_sync(0xffffffff, x, offset, width) && x;
+        }
+        return x;
+    }
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ int warp_reduce_any(int x) {
+    if (width == ggml_cuda_get_physical_warp_size()) {
+        return __any_sync(0xffffffff, x);
+    } else {
+#pragma unroll
+        for (int offset = width/2; offset > 0; offset >>= 1) {
+            x = __shfl_xor_sync(0xffffffff, x, offset, width) || x;
+        }
+        return x;
+    }
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, width));
+    }
+    return x;
+}
+
+template<typename T, int width = WARP_SIZE>
+static __device__ __forceinline__ T warp_prefix_inclusive_sum(T x) {
+    const int lane_id = threadIdx.x % width;
+#pragma unroll
+    for (int offset = 1; offset < width; offset <<= 1) {
+        const T t = __shfl_up_sync(0xffffffff, x, offset, width);
+        if (lane_id >= offset) {
+            x += t;
+        }
+    }
+    return x;
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ float2 warp_prefix_inclusive_sum(float2 a) {
+    const int lane_id = threadIdx.x % width;
+#pragma unroll
+    for (int offset = 1; offset < width; offset <<= 1) {
+        const float t_x = __shfl_up_sync(0xffffffff, a.x, offset, width);
+        const float t_y = __shfl_up_sync(0xffffffff, a.y, offset, width);
+        if (lane_id >= offset) {
+            a.x += t_x;
+            a.y += t_y;
+        }
+    }
+    return a;
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ half2 warp_prefix_inclusive_sum(half2 a) {
+#ifdef FP16_AVAILABLE
+    const int lane_id = threadIdx.x % width;
+#pragma unroll
+    for (int offset = 1; offset < width; offset <<= 1) {
+        const half2 t = __shfl_up_sync(0xffffffff, a, offset, width);
+        if (lane_id >= offset) {
+            a = __hadd2(a, t);
+        }
+    }
+    return a;
+
+#else
+    NO_DEVICE_CODE;
+    return a;
+#endif // FP16_AVAILABLE
+}
+
+enum class block_reduce_method {
+    MAX,
+    SUM,
+};
+
+template<block_reduce_method method_t, typename T>
+struct block_reduce_policy;
+
+template <typename T, typename... Ts>
+inline constexpr bool is_any = (std::is_same_v<T, Ts> || ...);
+
+template<typename...>
+inline constexpr bool ggml_cuda_dependent_false_v = false;
+
+template <typename T> struct block_reduce_policy<block_reduce_method::SUM, T> {
+    static __device__ T reduce(T val) {
+        if constexpr(is_any<T, float, float2, half2, int>) {
+            return warp_reduce_sum(val);
+        } else {
+            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce sum");
+        }
+    }
+
+    static __device__ T sentinel() {
+        if constexpr (std::is_same_v<T, float>) {
+            return 0.0f;
+        } else if constexpr (std::is_same_v<T, float2>) {
+            return make_float2(0.0f, 0.0f);
+        } else if constexpr (std::is_same_v<T, half2>) {
+            return make_half2(0.0f, 0.0f);
+        } else if constexpr (std::is_same_v<T, int>) {
+            return 0;
+        } else {
+            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce sum");
+        }
+    }
+};
+
+template <typename T> struct block_reduce_policy<block_reduce_method::MAX, T> {
+    static __device__ T reduce(T val) {
+        if constexpr (is_any<T, float, half2>) {
+            return warp_reduce_max(val);
+        } else {
+            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce max");
+        }
+    }
+
+    static __device__ T sentinel() {
+        if constexpr (std::is_same_v<T, float>) {
+            return -INFINITY;
+        } else if constexpr (std::is_same_v<T, half2>) {
+            return make_half2(-INFINITY, -INFINITY);
+        } else {
+            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce max");
+        }
+    }
+};
+
+template <block_reduce_method reduce_method_t, const unsigned int block_size_template = 0, typename T>
+static __device__ T block_reduce(T val, [[maybe_unused]] T * shared_vals) {
+    // for multi-warp reductions, callers must not reuse shared_vals until all reads from this invocation have completed
+    val                           = block_reduce_policy<reduce_method_t, T>::reduce(val);
+    const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
+    if (block_size > WARP_SIZE) {
+        assert((block_size <= 1024) && (block_size % WARP_SIZE) == 0);
+        const int warp_id = threadIdx.x / WARP_SIZE;
+        const int lane_id = threadIdx.x % WARP_SIZE;
+        if (lane_id == 0) {
+            shared_vals[warp_id] = val;
+        }
+        __syncthreads();
+        val = block_reduce_policy<reduce_method_t, T>::sentinel();
+        if (lane_id < (static_cast<int>(block_size) / WARP_SIZE)) {
+            val = shared_vals[lane_id];
+        }
+        return block_reduce_policy<reduce_method_t, T>::reduce(val);
+    }
+
+    return val;
+}
+
+static __device__ __forceinline__ half ggml_cuda_hmax(const half a, const half b) {
+#ifdef FP16_AVAILABLE
+
+#if !defined(GGML_USE_HIP) && CUDART_VERSION < CUDART_HMAX
+    return __float2half(fmaxf(__half2float(a), __half2float(b)));
+#else
+    return __hmax(a, b);
+#endif // !defined(GGML_USE_HIP) && CUDART_VERSION < CUDART_HMAX
+
+#else
+   NO_DEVICE_CODE;
+   GGML_UNUSED(b);
+   return a;
+#endif // FP16_AVAILABLE
+}
+
+static __device__ __forceinline__ half2 ggml_cuda_hmax2(const half2 a, const half2 b) {
+#if defined(GGML_USE_HIP)
+    return half2(__hmax(a.x, b.x), __hmax(a.y, b.y));
+#elif CUDART_VERSION >= CUDART_HMAX
+    return __hmax2(a, b);
+#else
+    half2 ret;
+    reinterpret_cast<half&>(ret.x) = __float2half(fmaxf( __low2float(a),  __low2float(b)));
+    reinterpret_cast<half&>(ret.y) = __float2half(fmaxf(__high2float(a), __high2float(b)));
+    return ret;
+#endif
+}
+
+template<int width = WARP_SIZE>
+static __device__ __forceinline__ half2 warp_reduce_max(half2 x) {
+#pragma unroll
+   for (int offset = width/2; offset > 0; offset >>= 1) {
+       x = ggml_cuda_hmax2(x, __shfl_xor_sync(0xffffffff, x, offset, width));
+   }
+   return x;
+}
+
+#if (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || \
+    (defined(MUSART_VERSION) && MUSART_VERSION < MUSART_HMASK)
+static __device__ __forceinline__ uint32_t __hgt2_mask(const half2 a, const half2 b) {
+    const uint32_t mask_low  = 0x0000FFFF * (float( __low2half(a)) > float( __low2half(b)));
+    const uint32_t mask_high = 0xFFFF0000 * (float(__high2half(a)) > float(__high2half(b)));
+    return mask_low | mask_high;
+}
+#endif // (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || (defined(MUSART_VERSION) && MUSART_VERSION < MUSART_HMASK)
+
+static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) {
+#if defined(GGML_USE_HIP)
+#if defined(CDNA) || defined(RDNA2) || defined(__gfx906__)
+    c = __builtin_amdgcn_sdot4(a, b, c, false);
+#elif defined(RDNA3) || defined(RDNA4)
+    c = __builtin_amdgcn_sudot4( true, a, true, b, c, false);
+#elif defined(RDNA1) || defined(__gfx900__)
+    int tmp1;
+    int tmp2;
+    asm("\n \
+        v_mul_i32_i24 %1, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:BYTE_0 \n \
+        v_mul_i32_i24 %2, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_1 src1_sel:BYTE_1 \n \
+        v_add3_u32 %0, %1, %2, %0 \n \
+        v_mul_i32_i24 %1, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_2 src1_sel:BYTE_2 \n \
+        v_mul_i32_i24 %2, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_3 src1_sel:BYTE_3 \n \
+        v_add3_u32 %0, %1, %2, %0 \n \
+        "
+        : "+v"(c), "=&v"(tmp1), "=&v"(tmp2)
+        : "v"(a), "v"(b)
+    );
+#else
+    const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
+    const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
+    c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3];
+#endif
+    return c;
+
+#else // defined(GGML_USE_HIP)
+
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA)
+    return __dp4a(a, b, c);
+#else // __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA)
+    const int8_t * a8 = (const int8_t *) &a;
+    const int8_t * b8 = (const int8_t *) &b;
+    return c + a8[0]*b8[0] + a8[1]*b8[1] + a8[2]*b8[2] + a8[3]*b8[3];
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA)
+
+#endif // defined(GGML_USE_HIP)
+}
+
+static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const float v, const float u) {
+    acc += v*u;
+}
+
+static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const float2 v, const float2 u) {
+    acc += v.x*u.x;
+    acc += v.y*u.y;
+}
+
+#if defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
+#define V_DOT2_F32_F16_AVAILABLE
+#endif // defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
+
+static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const half2 v, const half2 u) {
+#ifdef V_DOT2_F32_F16_AVAILABLE
+    asm volatile("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(acc) : "v"(v), "v"(u));
+#else
+#ifdef FAST_FP16_AVAILABLE
+    const float2 tmp = __half22float2(v*u);
+    acc += tmp.x + tmp.y;
+#else
+    const float2 tmpv = __half22float2(v);
+    const float2 tmpu = __half22float2(u);
+    acc += tmpv.x * tmpu.x;
+    acc += tmpv.y * tmpu.y;
+#endif // FAST_FP16_AVAILABLE
+#endif // V_DOT2_F32_F16_AVAILABLE
+}
+
+static __device__ __forceinline__ void ggml_cuda_mad(half2 & acc, const half2 v, const half2 u) {
+#ifdef FAST_FP16_AVAILABLE
+    acc += v*u;
+#else
+    const float2 tmpv = __half22float2(v);
+    const float2 tmpu = __half22float2(u);
+    float2 tmpacc = __half22float2(acc);
+    tmpacc.x += tmpv.x * tmpu.x;
+    tmpacc.y += tmpv.y * tmpu.y;
+    acc = make_half2(tmpacc.x, tmpacc.y);
+#endif // FAST_FP16_AVAILABLE
+}
+
+// Aligned memory transfers of 8/16 bytes can be faster than 2 transfers with 4 bytes, especially on AMD.
+// Important: do not use this function if dst and src both point at registers.
+//     Due to the strict aliasing rule the compiler can do incorrect optimizations if src and dst have different types.
+//     The function is intended for copies between registers and SRAM/VRAM to make the compiler emit the right instructions.
+//     If dst and src point at different address spaces then they are guaranteed to not be aliased.
+template <int nbytes, int alignment = 0>
+static __device__ __forceinline__ void ggml_cuda_memcpy_1(void * __restrict__ dst, const void * __restrict__ src) {
+    static_assert(
+        nbytes <= ggml_cuda_get_max_cpy_bytes() || alignment == 0,
+        "You are misusing the alignment parameter for ggml_cuda_memcpy_1. "
+        "The intent is for the parameter is only as a workaround if either one of the pointers is not properly aligned. "
+        "If you use it to do more bytes per copy than ggml_cuda_max_cpy_bytes() the reads and writes may not be coalesced. "
+        "Call ggml_cuda_memcpy_1 in a loop instead.");
+    if constexpr (alignment != 0) {
+        static_assert(nbytes % alignment == 0, "bad alignment");
+    }
+    constexpr int nb_per_cpy = alignment == 0 ? nbytes : alignment;
+
+#pragma unroll
+    for (int i = 0; i < nbytes/nb_per_cpy; ++i) {
+        if constexpr (nb_per_cpy == 1) {
+            ((char *) dst)[i] = ((const char *) src)[i];
+        } else if constexpr (nb_per_cpy == 2) {
+            ((short *) dst)[i] = ((const short *) src)[i];
+        } else if constexpr (nb_per_cpy == 4) {
+            ((int *) dst)[i] = ((const int *) src)[i];
+        } else if constexpr (nb_per_cpy == 8) {
+            ((int2 *) dst)[i] = ((const int2 *) src)[i];
+        } else if constexpr (nb_per_cpy == 16) {
+            ((int4 *) dst)[i] = ((const int4 *) src)[i];
+        } else {
+            static_assert(nbytes == 0 && nbytes == -1, "bad nbytes");
+        }
+    }
+}
+
+static __device__ __forceinline__ float ggml_cuda_e8m0_to_fp32(uint8_t x) {
+#if CUDART_VERSION >= 12080
+    const nv_bfloat16 e = __nv_cvt_e8m0_to_bf16raw(x);
+    return (float) e;
+#else
+    uint32_t bits;
+    if (x == 0) {
+        bits = 0x00400000;
+    } else {
+        bits = (uint32_t) x << 23;
+    }
+
+    float result;
+    memcpy(&result, &bits, sizeof(float));
+    return result;
+#endif // CUDART_VERSION >= 12050
+}
+
+static __device__ __forceinline__ float ggml_cuda_ue4m3_to_fp32(uint8_t x) {
+#if defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000
+    // ROCm does not support fp8 in software on devices with fp8 hardware,
+    // but CDNA3 supports only e4m3_fnuz (no inf).
+    const uint32_t bits = x * (x != 0x7F && x != 0xFF); // Convert NaN to 0.0f to match CPU implementation.
+    const __hip_fp8_e4m3_fnuz xf = *reinterpret_cast<const __hip_fp8_e4m3_fnuz *>(&bits);
+    return static_cast<float>(xf) / 2;
+#else
+#if defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP)
+    const uint32_t bits = x * (x != 0x7F && x != 0xFF); // Convert NaN to 0.0f to match CPU implementation.
+    const __nv_fp8_e4m3 xf = *reinterpret_cast<const __nv_fp8_e4m3 *>(&bits);
+    return static_cast<float>(xf) / 2;
+#else
+    if (x == 0 || (x == 0x7F && x != 0xFF)) { // Convert NaN to 0.0f
+        return 0.0f;
+    }
+    const int exp = (x >> 3) & 0xF;
+    const int man = x & 0x7;
+    float raw;
+    if (exp == 0) {
+        raw = ldexpf((float) man, -9);
+    } else {
+        raw = ldexpf(1.0f + (float) man / 8.0f, exp - 7);
+    }
+    return static_cast<float>(raw / 2);
+#endif // defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP)
+#endif // defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000
+}
+
+static __device__ __forceinline__ uint8_t ggml_cuda_fp32_to_ue4m3(float x) {
+#if defined(BLACKWELL_MMA_AVAILABLE) // This is used for NVFP4 subblock scale quantizations only
+    if (!(x > 0.0f)) {
+        return 0;
+    }
+    const __nv_fp8_e4m3 xf(x);
+    return xf.__x;
+#else
+     GGML_UNUSED(x);
+     NO_DEVICE_CODE; // Used only for NVFP4 Scales for Activations, only for Blackwell
+#endif // defined(BLACKWELL_MMA_AVAILABLE)
+}
+
+__device__ __forceinline__ uint8_t ggml_cuda_float_to_fp4_e2m1(float x, float e) {
+    const uint8_t sign_bit = (x < 0.0f) << 3;
+    float         ax       = fabsf(x) * e;
+
+    // Positive LUT
+    static constexpr float pos_lut[8] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f };
+
+    int   best_i   = 0;
+    float best_err = fabsf(ax - pos_lut[0]);
+
+#pragma unroll
+    for (int i = 1; i < 8; ++i) {
+        const float err = fabsf(ax - pos_lut[i]);
+        if (err < best_err) {
+            best_err = err;
+            best_i   = i;
+        }
+    }
+
+    return static_cast<uint8_t>(best_i | sign_bit);
+}
+
+// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1.
+// Precompute mp (m' in the paper) and L such that division
+// can be computed using a multiply (high 32b of 64b result)
+// and a shift:
+//
+// n/d = (mulhi(n, mp) + n) >> L;
+static const uint3 init_fastdiv_values(uint64_t d_64) {
+    GGML_ASSERT(d_64 != 0);
+    GGML_ASSERT(d_64 <= std::numeric_limits<uint32_t>::max());
+
+    uint32_t d = (uint32_t)d_64;
+
+    // compute L = ceil(log2(d));
+    uint32_t L = 0;
+    while (L < 32 && (uint32_t{ 1 } << L) < d) {
+        L++;
+    }
+
+    uint32_t mp = (uint32_t) ((uint64_t{ 1 } << 32) * ((uint64_t{ 1 } << L) - d) / d + 1);
+    // pack divisor as well to reduce error surface
+    return make_uint3(mp, L, d);
+}
+
+static __device__ __forceinline__ uint32_t fastdiv(uint32_t n, const uint3 fastdiv_values) {
+    // expects fastdiv_values to contain <mp, L, divisor> in <x, y, z>
+    // fastdiv_values.z is unused and optimized away by the compiler.
+    // Compute high 32 bits of n * mp
+    const uint32_t hi = __umulhi(n, fastdiv_values.x);
+    // add n, apply bit shift
+    return (hi + n) >> fastdiv_values.y;
+}
+
+static __device__ __forceinline__ uint32_t fastmodulo(uint32_t n, const uint3 fastdiv_values) {
+    // expects  fastdiv_values to contain <mp, L, divisor> in <x, y, z> (see init_fastdiv_values)
+    return n - fastdiv(n, fastdiv_values) * fastdiv_values.z;
+}
+
+// Calculate both division and modulo at once, returns <n/divisor, n%divisor>
+static __device__ __forceinline__ uint2 fast_div_modulo(uint32_t n, const uint3 fastdiv_values) {
+    // expects  fastdiv_values to contain <mp, L, divisor> in <x, y, z> (see init_fastdiv_values)
+    const uint32_t div_val = fastdiv(n, fastdiv_values);
+    const uint32_t mod_val = n - div_val * fastdiv_values.z;
+    return make_uint2(div_val, mod_val);
+}
+
+typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);
+
+template<typename dst_t>
+using dequantize_kq_t = void (*)(const void * vx, const int64_t ib, dst_t * y, const int tid);
+
+static __device__ __forceinline__ float get_alibi_slope(
+    const float max_bias, const uint32_t h, const uint32_t n_head_log2, const float m0, const float m1
+) {
+    if (max_bias <= 0.0f) {
+        return 1.0f;
+    }
+    const float base = h < n_head_log2 ? m0 : m1;
+    const int   exph = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
+
+    return powf(base, exph);
+}
+
+template <ggml_type type>
+struct ggml_cuda_type_traits;
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_F16> {
+    static constexpr int qk = 1;
+    static constexpr int qr = 1;
+    static constexpr int bs = sizeof(ggml_half);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q1_0> {
+    static constexpr int qk = QK1_0;
+    static constexpr int qr = QR1_0;
+    static constexpr int qi = QI1_0;
+    static constexpr int bs = sizeof(block_q1_0);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q2_0> {
+    static constexpr int qk = QK2_0;
+    static constexpr int qr = QR2_0;
+    static constexpr int qi = QI2_0;
+    static constexpr int bs = sizeof(block_q2_0);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q4_0> {
+    static constexpr int qk = QK4_0;
+    static constexpr int qr = QR4_0;
+    static constexpr int qi = QI4_0;
+    static constexpr int bs = sizeof(block_q4_0);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q4_1> {
+    static constexpr int qk = QK4_1;
+    static constexpr int qr = QR4_1;
+    static constexpr int qi = QI4_1;
+    static constexpr int bs = sizeof(block_q4_1);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q5_0> {
+    static constexpr int qk = QK5_0;
+    static constexpr int qr = QR5_0;
+    static constexpr int qi = QI5_0;
+    static constexpr int bs = sizeof(block_q5_0);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q5_1> {
+    static constexpr int qk = QK5_1;
+    static constexpr int qr = QR5_1;
+    static constexpr int qi = QI5_1;
+    static constexpr int bs = sizeof(block_q5_1);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q8_0> {
+    static constexpr int qk = QK8_0;
+    static constexpr int qr = QR8_0;
+    static constexpr int qi = QI8_0;
+    static constexpr int bs = sizeof(block_q8_0);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_MXFP4> {
+    static constexpr int qk = QK_MXFP4;
+    static constexpr int qr = QR_MXFP4;
+    static constexpr int qi = QI_MXFP4;
+    static constexpr int bs = sizeof(block_mxfp4);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_NVFP4> {
+    static constexpr int qk = QK_NVFP4;
+    static constexpr int qr = QR_NVFP4;
+    static constexpr int qi = QI_NVFP4;
+    static constexpr int bs = sizeof(block_nvfp4);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q2_K> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR2_K;
+    static constexpr int qi = QI2_K;
+    static constexpr int bs = sizeof(block_q2_K);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q3_K> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR3_K;
+    static constexpr int qi = QI3_K;
+    static constexpr int bs = sizeof(block_q3_K);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q4_K> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR4_K;
+    static constexpr int qi = QI4_K;
+    static constexpr int bs = sizeof(block_q4_K);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q5_K> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR5_K;
+    static constexpr int qi = QI5_K;
+    static constexpr int bs = sizeof(block_q5_K);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q6_K> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR6_K;
+    static constexpr int qi = QI6_K;
+    static constexpr int bs = sizeof(block_q6_K);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ2_XXS> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR2_XXS;
+    static constexpr int qi = QI2_XXS;
+    static constexpr int bs = sizeof(block_iq2_xxs);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ2_XS> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR2_XS;
+    static constexpr int qi = QI2_XS;
+    static constexpr int bs = sizeof(block_iq2_xs);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ2_S> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR2_S;
+    static constexpr int qi = QI2_S;
+    static constexpr int bs = sizeof(block_iq2_s);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ3_XXS> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR3_XXS;
+    static constexpr int qi = QI3_XXS;
+    static constexpr int bs = sizeof(block_iq3_xxs);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ1_S> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR1_S;
+    static constexpr int qi = QI1_S;
+    static constexpr int bs = sizeof(block_iq1_s);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ1_M> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR1_M;
+    static constexpr int qi = QI1_M;
+    static constexpr int bs = sizeof(block_iq1_m);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ4_NL> {
+    static constexpr int qk = QK4_NL;
+    static constexpr int qr = QR4_NL;
+    static constexpr int qi = QI4_NL;
+    static constexpr int bs = sizeof(block_iq4_nl);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ4_XS> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR4_XS;
+    static constexpr int qi = QI4_XS;
+    static constexpr int bs = sizeof(block_iq4_xs);
+};
+
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ3_S> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR3_S;
+    static constexpr int qi = QI3_S;
+    static constexpr int bs = sizeof(block_iq3_s);
+};
+
+//////////////////////
+
+struct ggml_cuda_device_info {
+    int device_count;           // number of (possibly virtual) devices exposed to the rest of ggml
+    int physical_device_count;  // number of physical CUDA devices actually present
+
+    struct cuda_device_info {
+        int     cc;                             // compute capability
+        int     nsm;                            // number of streaming multiprocessors
+        size_t  smpb;                           // max. shared memory per block
+        size_t  smpbo;                          // max. shared memory per block (with opt-in)
+        bool    integrated;                     // Device is integrated as opposed to discrete
+        bool    vmm;                            // virtual memory support
+        size_t  vmm_granularity;                // granularity of virtual memory
+        size_t  total_vram;
+        int     warp_size;                      // Number of threads in a dispatch
+        bool    supports_cooperative_launch;    // whether cooperative launch is supported
+        int     physical_device;                // backing physical CUDA device for this (virtual) device
+        int     physical_share_count;           // number of (virtual) devices sharing this device's physical GPU
+        int     virtual_index;                  // index of this (virtual) device among those sharing its physical GPU
+    };
+
+    cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};
+
+    std::array<float, GGML_CUDA_MAX_DEVICES> default_tensor_split = {};
+};
+
+const ggml_cuda_device_info & ggml_cuda_info();
+
+void ggml_cuda_set_device(int device);
+int ggml_cuda_get_device();
+
+struct ggml_cuda_pool {
+    virtual ~ggml_cuda_pool() = default;
+
+    virtual void * alloc(size_t size, size_t * actual_size) = 0;
+    virtual void free(void * ptr, size_t size) = 0;
+};
+
+template<typename T>
+struct ggml_cuda_pool_alloc {
+    ggml_cuda_pool * pool = nullptr;
+    T * ptr = nullptr;
+    size_t actual_size = 0;
+
+    ggml_cuda_pool_alloc() = default;
+
+    explicit ggml_cuda_pool_alloc(ggml_cuda_pool & pool) : pool(&pool) {
+    }
+
+    ggml_cuda_pool_alloc(ggml_cuda_pool & pool, size_t size) : pool(&pool) {
+        alloc(size);
+    }
+
+    ~ggml_cuda_pool_alloc() {
+        if (ptr != nullptr) {
+            pool->free(ptr, actual_size);
+        }
+    }
+
+    // size is in number of elements
+    T * alloc(size_t size) {
+        GGML_ASSERT(pool != nullptr);
+        GGML_ASSERT(ptr == nullptr);
+        ptr = (T *) pool->alloc(size * sizeof(T), &this->actual_size);
+        return ptr;
+    }
+
+    T * alloc(ggml_cuda_pool & pool, size_t size) {
+        this->pool = &pool;
+        return alloc(size);
+    }
+
+    T * get() {
+        return ptr;
+    }
+
+    ggml_cuda_pool_alloc(const ggml_cuda_pool_alloc &) = delete;
+    ggml_cuda_pool_alloc(ggml_cuda_pool_alloc &&) = delete;
+    ggml_cuda_pool_alloc& operator=(const ggml_cuda_pool_alloc &) = delete;
+    ggml_cuda_pool_alloc& operator=(ggml_cuda_pool_alloc &&) = delete;
+};
+
+
+// backend interface
+
+struct ggml_tensor_extra_gpu {
+    void * data_device[GGML_CUDA_MAX_DEVICES]; // 1 pointer for each device for split tensors
+    cudaEvent_t events[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS]; // events for synchronizing multiple GPUs
+};
+
+
+#if (defined(GGML_CUDA_USE_GRAPHS) || defined(GGML_HIP_GRAPHS)) || defined(GGML_MUSA_GRAPHS)
+#define USE_CUDA_GRAPH
+#endif
+
+struct ggml_cuda_graph {
+#ifdef USE_CUDA_GRAPH
+    ~ggml_cuda_graph() {
+        if (instance != nullptr) {
+            CUDA_CHECK(cudaGraphExecDestroy(instance));
+        }
+        if (graph != nullptr) {
+            CUDA_CHECK(cudaGraphDestroy(graph));
+        }
+    }
+    cudaGraph_t graph = nullptr;
+    cudaGraphExec_t instance = nullptr;
+    size_t num_nodes = 0;
+    std::vector<cudaGraphNode_t> nodes;
+    bool disable_due_to_gpu_arch = false;
+    bool warmup_complete = false;
+    uint64_t uid = 0;
+    int64_t last_used_time = 0;
+    struct node_properties {
+        ggml_tensor node;
+        void *   node_src_data_ptrs[GGML_MAX_SRC];
+        int64_t  node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS];
+        size_t   node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS];
+    };
+    std::vector<node_properties> node_props;
+
+    bool is_enabled() const {
+        static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
+        return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env);
+    }
+#endif
+};
+
+struct ggml_cuda_concurrent_event {
+    std::vector<cudaEvent_t> join_events;
+    cudaEvent_t              fork_event = nullptr;
+
+    int                                          n_streams = 0;
+    std::unordered_map<const ggml_tensor *, int> stream_mapping;
+
+    // Original order of nodes in this concurrent region (before interleaving)
+    // Used to restore grouping for fusion within streams
+    std::vector<const ggml_tensor *> original_order;
+
+    const ggml_tensor * join_node;
+
+    ggml_cuda_concurrent_event() = default;
+
+    ggml_cuda_concurrent_event(const ggml_cuda_concurrent_event &) = delete;
+    ggml_cuda_concurrent_event & operator=(const ggml_cuda_concurrent_event &) = delete;
+
+    explicit ggml_cuda_concurrent_event(int n_streams) : n_streams(n_streams) {
+        join_events.resize(n_streams);
+
+        for (size_t i = 0; i < join_events.size(); ++i) {
+            CUDA_CHECK(cudaEventCreateWithFlags(&join_events[i], cudaEventDisableTiming));
+        }
+
+        CUDA_CHECK(cudaEventCreateWithFlags(&fork_event, cudaEventDisableTiming));
+    }
+
+    ggml_cuda_concurrent_event(ggml_cuda_concurrent_event && other) noexcept
+    : join_events(std::move(other.join_events))
+    , fork_event(other.fork_event)
+    , n_streams(other.n_streams)
+    , stream_mapping(std::move(other.stream_mapping))
+    , original_order(std::move(other.original_order))
+    , join_node(other.join_node) {
+        other.fork_event = nullptr;
+    }
+
+    // 1. check if any branches write to overlapping memory ranges (except the join node)
+    // 2. check whether all srcs are either within the branch or outside the nodes covered by ggml_cuda_concurrent_event
+    // we assume all nodes have the same buffer
+    bool is_valid() const {
+        std::vector<std::vector<std::pair<int64_t, int64_t>>> write_ranges;
+        write_ranges.resize(n_streams);
+
+        // get join_node's memory range to exclude from overlap checking.
+        // multiple nodes can use join_node's buffer; we synchronize on the join node.
+        const ggml_tensor * join_t     = join_node->view_src ? join_node->view_src : join_node;
+        const int64_t       join_start = (int64_t) join_t->data;
+        const int64_t       join_end   = join_start + ggml_nbytes(join_t);
+
+        for (const auto & [tensor, stream] : stream_mapping) {
+            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
+            const int64_t       t_start = (int64_t) t->data;
+            const int64_t       t_end   = t_start + ggml_nbytes(t);
+
+            // skip tensors that overlap with join_node's buffer.
+            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
+                continue;
+            }
+
+            // concurrent streams begin from 1
+            write_ranges[stream - 1].emplace_back(t_start, t_end);
+        }
+
+        for (int i = 0; i < n_streams; ++i) {
+            // sorts first by start then by end of write range
+            std::sort(write_ranges[i].begin(), write_ranges[i].end());
+        }
+
+        bool writes_overlap = false;
+        bool dependent_srcs = false;
+        for (const auto & [tensor, stream] : stream_mapping) {
+            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
+            const int64_t       t_start = (int64_t) t->data;
+            const int64_t       t_end   = t_start + ggml_nbytes(t);
+
+            // skip tensors that overlap with join_node's buffer
+            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
+                continue;
+            }
+
+            // check if this buffer's write data overlaps with another stream's
+            std::pair<int64_t, int64_t> data_range = std::make_pair(t_start, t_end);
+            for (int i = 0; i < n_streams; ++i) {
+                if (i == stream - 1) {
+                    continue;
+                }
+                auto it = std::lower_bound(write_ranges[i].begin(), write_ranges[i].end(), data_range);
+
+                if (it != write_ranges[i].end()) {
+                    const std::pair<int64_t, int64_t> & other = *it;
+
+                    // std::lower_bound returns the first element where other >= data_range (lexicographically).
+                    // This guarantees other.first >= data_range.first.
+                    // Therefore, overlap occurs iff other.first < data_range.second
+                    // (i.e., the other range starts before this range ends).
+                    if (other.first < data_range.second) {
+                        GGML_LOG_DEBUG("Writes overlap for %s", tensor->name);
+                        writes_overlap = true;
+                        break;
+                    }
+                }
+            }
+
+            //check if all srcs are either in branch or don't have a branch
+            for (int i = 0; i < GGML_MAX_SRC; ++i) {
+                if (!tensor->src[i]) {
+                    continue;
+                }
+
+                auto it = stream_mapping.find(tensor->src[i]);
+
+                if (it == stream_mapping.end()) {
+                    continue;
+                }
+
+                if (it->second != stream) {
+                    dependent_srcs = true;
+                    break;
+                }
+            }
+
+            if (dependent_srcs || writes_overlap) {
+                break;
+            }
+        }
+
+        return !writes_overlap && !dependent_srcs;
+    }
+
+    ~ggml_cuda_concurrent_event() {
+        if (fork_event != nullptr) {
+            CUDA_CHECK(cudaEventDestroy(fork_event));
+        }
+        for (cudaEvent_t e : join_events) {
+            if (e != nullptr) {
+                CUDA_CHECK(cudaEventDestroy(e));
+            }
+        }
+    }
+};
+
+struct ggml_cuda_stream_context {
+    std::unordered_map<const ggml_tensor *, ggml_cuda_concurrent_event> concurrent_events;
+
+    void reset() {
+        concurrent_events.clear();
+    }
+};
+
+struct ggml_backend_cuda_context {
+    int device;
+    std::string name;
+    cudaEvent_t copy_event = nullptr;
+
+    cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {};
+
+    int curr_stream_no = 0;
+
+#ifdef USE_CUDA_GRAPH
+    // Map from first_node_ptr to cuda_graph - allows multiple graphs per context
+    // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe)
+    std::unordered_map<const void *, std::unique_ptr<ggml_cuda_graph>> cuda_graphs;
+
+    int64_t last_graph_eviction_sweep = 0;
+
+    ggml_cuda_graph * cuda_graph(const void * first_node_ptr) {
+        const int64_t time_now = ggml_time_us();
+
+        // sweep every 5s, evicting cuda graphs unused for >=10s
+        if (time_now - last_graph_eviction_sweep >= 5'000'000) {
+            last_graph_eviction_sweep = time_now;
+            for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
+                if (time_now - it->second->last_used_time >= 10'000'000) {
+                    it = cuda_graphs.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+        }
+
+        auto it = cuda_graphs.find(first_node_ptr);
+        if (it == cuda_graphs.end()) {
+            it = cuda_graphs.emplace(first_node_ptr, std::make_unique<ggml_cuda_graph>()).first;
+        }
+        it->second->last_used_time = time_now;
+        return it->second.get();
+    }
+
+    // Check if any CUDA graph is enabled for this context (used by kernels that need to know
+    // if graphs are in use without having access to the specific graph key)
+    bool any_cuda_graph_enabled() const {
+        for (const auto & [key, graph] : cuda_graphs) {
+            if (graph && graph->is_enabled()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    // Check if any CUDA graph has an instance for this context
+    bool any_cuda_graph_has_instance() const {
+        for (const auto & [key, graph] : cuda_graphs) {
+            if (graph && graph->instance != nullptr) {
+                return true;
+            }
+        }
+        return false;
+    }
+#endif // USE_CUDA_GRAPH
+
+    explicit ggml_backend_cuda_context(int device) :
+        device(device),
+        name(GGML_CUDA_NAME + std::to_string(device)) {
+    }
+
+    ggml_cuda_stream_context concurrent_stream_context;
+
+    ~ggml_backend_cuda_context();
+
+    cudaStream_t stream(int device, int stream) {
+        if (streams[device][stream] == nullptr) {
+            ggml_cuda_set_device(device);
+            CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking));
+        }
+        return streams[device][stream];
+    }
+
+    cudaStream_t stream() { return stream(device, curr_stream_no); }
+
+    ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; }
+
+    cublasHandle_t cublas_handle() {
+        if (cublas_handles[device][curr_stream_no] == nullptr) {
+            ggml_cuda_set_device(device);
+            CUBLAS_CHECK(cublasCreate(&cublas_handles[device][curr_stream_no]));
+            CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device][curr_stream_no], CUBLAS_TF32_TENSOR_OP_MATH));
+            CUBLAS_CHECK(cublasSetStream(cublas_handles[device][curr_stream_no], stream()));
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUBLAS_VER_MAJOR > 11 || (CUBLAS_VER_MAJOR == 11 && CUBLAS_VER_MINOR >= 2))
+            if (cublas_workspace_sizes[device] == 0) {
+                const int cc = ggml_cuda_info().devices[device].cc;
+                cublas_workspace_sizes[device] = (cc >= GGML_CUDA_CC_HOPPER) ? 32 * 1024 * 1024 : 4 * 1024 * 1024;
+            }
+            CUDA_CHECK(cudaMalloc(&cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device]));
+            CUBLAS_CHECK(cublasSetWorkspace(cublas_handles[device][curr_stream_no], cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device]));
+#endif
+        }
+        return cublas_handles[device][curr_stream_no];
+    }
+
+    // pool
+    std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS];
+
+    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, int stream_no);
+
+    ggml_cuda_pool & pool(int device) {
+        if (pools[device][curr_stream_no] == nullptr) {
+            pools[device][curr_stream_no] = new_pool_for_device(device, curr_stream_no);
+        }
+        return *pools[device][curr_stream_no];
+    }
+
+    ggml_cuda_pool & pool() {
+        return pool(device);
+    }
+};
+
+struct ggml_cuda_mm_fusion_args_host {
+    const ggml_tensor * x_bias = nullptr;
+    const ggml_tensor * gate = nullptr;
+    const ggml_tensor * gate_bias = nullptr;
+    const ggml_tensor * x_scale = nullptr;
+    const ggml_tensor * gate_scale = nullptr;
+    ggml_glu_op glu_op;
+    float glu_limit = 0.0f;
+    const ggml_tensor * shared_up = nullptr;
+    const ggml_tensor * shared_gate = nullptr;
+    ggml_tensor * shared_dst = nullptr;
+};
+struct ggml_cuda_mm_fusion_args_device {
+    const void * x_bias = nullptr;
+    const void * gate = nullptr;
+    const void * gate_bias = nullptr;
+    const void * x_scale = nullptr;
+    const void * gate_scale = nullptr;
+    ggml_glu_op glu_op;
+    float glu_limit = 0.0f;
+    const void * shared_up = nullptr;
+    const void * shared_gate = nullptr;
+    float * shared_dst = nullptr;
+    uint32_t shared_stride_col_dst = 0;
+};
+
+struct ggml_cuda_kernel_launch_params {
+    dim3 block_nums;
+    dim3 block_dims;
+    size_t shmem;
+    cudaStream_t stream;
+
+    // size_t shmem
+    ggml_cuda_kernel_launch_params(const dim3& block_nums_, const dim3& block_dims_, const size_t shmem_, const cudaStream_t stream_)
+        : block_nums(block_nums_), block_dims(block_dims_), shmem(shmem_), stream(stream_) {}
+
+    // Some call sites pass ints instead of the required size_t. This 2nd constructor casts int->size_t to avoid these -Wnarrowing warnings.
+    ggml_cuda_kernel_launch_params(const dim3& block_nums_, const dim3& block_dims_, const int shmem_, const cudaStream_t stream_)
+        : block_nums(block_nums_), block_dims(block_dims_), shmem((size_t)shmem_), stream(stream_) {}
+};
+
+#if defined(GGML_CUDA_USE_PDL)
+struct ggml_cuda_pdl_config {
+    cudaLaunchAttribute attr;
+    cudaLaunchConfig_t  cfg;
+
+    ggml_cuda_pdl_config(const ggml_cuda_kernel_launch_params & params) {
+        attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attr.val.programmaticStreamSerializationAllowed = 1;
+
+        cfg = {};
+        cfg.gridDim          = params.block_nums;
+        cfg.blockDim         = params.block_dims;
+        cfg.dynamicSmemBytes = params.shmem;
+        cfg.stream           = params.stream;
+        cfg.attrs            = &attr;
+        cfg.numAttrs         = 1;
+    }
+
+    // Delete due to &attr
+    ggml_cuda_pdl_config(const ggml_cuda_pdl_config&) = delete;
+    ggml_cuda_pdl_config& operator=(const ggml_cuda_pdl_config&) = delete;
+    ggml_cuda_pdl_config& operator=(ggml_cuda_pdl_config&&) = delete;
+
+};
+
+static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
+    const int device = ggml_cuda_get_device();
+
+    struct cache_key {
+        int          device;
+        const void * kernel;
+
+        bool operator==(const cache_key & other) const { return device == other.device && kernel == other.kernel; }
+    };
+
+    struct cache_key_hash {
+        // MurmurHash3 mixing function for better hash distribution (vs. just std::hash which in some implementations simply returns the identity)
+        static size_t hash_mix(size_t x) {
+            std::uint64_t       y = x;
+            const std::uint64_t m = 0xe9846af9b1a615d;
+
+            y ^= y >> 32;
+            y *= m;
+            y ^= y >> 32;
+            y *= m;
+            y ^= y >> 28;
+
+            return static_cast<size_t>(y);
+        }
+
+        size_t operator()(const cache_key & key) const {
+            // Use a nonzero seed to avoid mapping all-zero keys to zero
+            size_t h = 42;
+            h        = hash_mix(h + key.device);
+            h        = hash_mix(h + reinterpret_cast<size_t>(key.kernel));
+            return h;
+        }
+    };
+
+    static std::mutex                                          cache_mutex;
+    static std::unordered_map<cache_key, bool, cache_key_hash> cache;
+
+    const cache_key             key = { device, kernel };
+    std::lock_guard<std::mutex> lock(cache_mutex);
+    const auto                  it = cache.find(key);
+    if (it != cache.end()) {
+        return it->second;
+    }
+
+    cudaFuncAttributes attr = {};
+    CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel));
+
+    // PDL device-side primitives are emitted only for PTX versions >= 90.
+    // We have to guard on a loaded kernel's PTX version so a kernel forward-JIT'ed
+    // from pre-Hopper PTX to a Hopper-or-newer GPU does not opt into PDL.
+    const bool can_use_pdl = attr.ptxVersion >= 90;
+    cache.emplace(key, can_use_pdl);
+    return can_use_pdl;
+}
+
+#endif //defined(GGML_CUDA_USE_PDL)
+
+// PDL and __restrict__ need to be mutually exclusive, see https://github.com/ggml-org/llama.cpp/pull/24030
+# if (defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER)
+# define GGML_CUDA_RESTRICT
+# else
+# define GGML_CUDA_RESTRICT __restrict__
+# endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+
+template<typename Kernel, typename... Args>
+static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_kernel_launch_params & launch_params, Args&&... args) {
+#if defined(GGML_CUDA_USE_PDL)
+
+    static const bool env_pdl_enabled = []() {
+        const char * env = getenv("GGML_CUDA_PDL");
+        return env == nullptr || std::atoi(env) != 0;
+    }();
+
+    if (env_pdl_enabled && ggml_cuda_kernel_can_use_pdl(reinterpret_cast<const void *>(kernel))) {
+        auto pdl_cfg = ggml_cuda_pdl_config(launch_params);
+
+        CUDA_CHECK(cudaLaunchKernelEx(&pdl_cfg.cfg, kernel, std::forward<Args>(args)... ));
+        return;
+    }
+#endif //defined(GGML_CUDA_USE_PDL)
+
+    kernel<<<launch_params.block_nums, launch_params.block_dims, launch_params.shmem, launch_params.stream>>>(std::forward<Args>(args)... );
+    CUDA_CHECK(cudaGetLastError());
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu
new file mode 100644
index 0000000000..77fd288945
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu
@@ -0,0 +1,743 @@
+#include "convert.cuh"
+#include "dequantize.cuh"
+
+#include <cstdint>
+
+#define CUDA_Q8_0_NE_ALIGN 2048
+
+template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
+static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y,
+        const int64_t ne00, const int64_t ne01,
+        const int64_t ne0203, const uint3 ne02,
+        const int64_t s01, const int64_t s02, const int64_t s03) {
+    const int64_t i00 = 2 * (int64_t(blockDim.x)*blockIdx.x + threadIdx.x);
+
+    if (i00 >= ne00) {
+        return;
+    }
+
+    for (int64_t i01 = blockIdx.y; i01 < ne01; i01 += gridDim.y) {
+        for (int64_t i0203 = blockIdx.z; i0203 < ne0203; i0203 += gridDim.z) {
+            const uint2 dm = fast_div_modulo((uint32_t)i0203, ne02);
+            const int64_t i02 = dm.y;
+            const int64_t i03 = dm.x;
+
+            const int64_t ibx0 = i03*s03 + i02*s02 + i01*s01;
+
+            const int64_t ib = ibx0 + i00/qk; // block index
+            const int64_t iqs = (i00%qk)/qr; // quant index
+            const int64_t iybs = i00 - i00%qk; // y block start index
+            const int64_t y_offset = qr == 1 ? 1 : qk/2;
+
+            // dequantize
+            float2 v;
+            dequantize_kernel(vx, ib, iqs, v);
+
+            const int64_t iy0 = (i0203*ne01 + i01)*ne00 + iybs + iqs;
+            y[iy0 + 0]        = ggml_cuda_cast<dst_t>(v.x);
+            y[iy0 + y_offset] = ggml_cuda_cast<dst_t>(v.y);
+        }
+    }
+}
+
+template <bool need_check>
+static __global__ void dequantize_block_q8_0_f16(const void * __restrict__ vx, half * __restrict__ y, const int64_t k) {
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL
+    constexpr int nint = CUDA_Q8_0_NE_ALIGN/sizeof(int) + WARP_SIZE;
+
+    const int64_t   i0 = CUDA_Q8_0_NE_ALIGN*blockIdx.x;
+    const int * x0 = ((int *) vx) + blockIdx.x * nint;
+    half2 * y2 = (half2 *) (y + i0);
+
+    __shared__ int vals[nint];
+
+#pragma unroll
+    for (int ix0 = 0; ix0 < nint; ix0 += WARP_SIZE) {
+        if (need_check && i0*sizeof(block_q8_0)/QK8_0 + sizeof(int)*(ix0 + threadIdx.x) >= k*sizeof(block_q8_0)/QK8_0) {
+            break;
+        }
+
+        const int ix = ix0 + threadIdx.x;
+        vals[ix] = x0[ix];
+    }
+
+    __syncthreads();
+
+#pragma unroll
+    for (int iy = 0; iy < CUDA_Q8_0_NE_ALIGN; iy += 2*WARP_SIZE) {
+        if (need_check && i0 + iy + 2*threadIdx.x >= k) {
+            return;
+        }
+
+        const half * b0 = ((const half  *) vals) + (sizeof(block_q8_0)/sizeof(half)) * ((iy + 2*threadIdx.x)/QK8_0);
+        const half    d = *b0;
+        const char2  qs = ((const char2 *) (b0 + 1))[threadIdx.x % (QK8_0/2)];
+
+        y2[iy/2 + threadIdx.x] = __hmul2(make_half2(qs.x, qs.y), __half2half2(d));
+    }
+#else
+    GGML_UNUSED_VARS(vx, y, k);
+    NO_DEVICE_CODE;
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q4_0(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) {
+
+    const int64_t i = blockIdx.x;
+
+    // assume 32 threads
+    const int64_t tid = threadIdx.x;
+    const int64_t il  = tid/8;
+    const int64_t ir  = tid%8;
+    const int64_t ib = 8*i + ir;
+    if (ib >= nb32) {
+        return;
+    }
+
+    dst_t * y = yy + 256*i + 32*ir + 4*il;
+
+    const block_q4_0 * x = (const block_q4_0 *)vx + ib;
+    const float d = __half2float(x->d);
+    const float dm = -8*d;
+
+    const uint8_t * q = x->qs + 4*il;
+
+    for (int l = 0; l < 4; ++l) {
+        y[l+ 0] = ggml_cuda_cast<dst_t>(d * (q[l] & 0xF) + dm);
+        y[l+16] = ggml_cuda_cast<dst_t>(d * (q[l] >>  4) + dm);
+    }
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) {
+
+    const int64_t i = blockIdx.x;
+
+    // assume 32 threads
+    const int64_t tid = threadIdx.x;
+    const int64_t il  = tid/8;
+    const int64_t ir  = tid%8;
+    const int64_t ib = 8*i + ir;
+    if (ib >= nb32) {
+        return;
+    }
+
+    dst_t * y = yy + 256*i + 32*ir + 4*il;
+
+    const block_q4_1 * x = (const block_q4_1 *)vx + ib;
+    const float2 d = __half22float2(x->dm);
+
+    const uint8_t * q = x->qs + 4*il;
+
+    for (int l = 0; l < 4; ++l) {
+        y[l+ 0] = ggml_cuda_cast<dst_t>(d.x * (q[l] & 0xF) + d.y);
+        y[l+16] = ggml_cuda_cast<dst_t>(d.x * (q[l] >>  4) + d.y);
+    }
+}
+
+//================================== k-quants
+
+template<typename dst_t>
+static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q4_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q5_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q6_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq2_xxs(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq2_xs(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq2_s(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq3_xxs(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq3_s(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq1_s(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq1_m(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) {
+    const int64_t i = blockIdx.x;
+
+    const int64_t ib = 8*i + threadIdx.x%8;
+    if (ib >= nb32) {
+        return;
+    }
+
+    dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_iq4_xs(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
+static void dequantize_block_cuda(const void * vx, dst_t * y,
+        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
+        const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) {
+    const int64_t ne0203 = ne02*ne03;
+    const uint3 ne02_fdv = init_fastdiv_values(ne02);
+    const dim3 num_blocks((ne00 + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE), (int)std::min(ne01, (int64_t)65535), (int)std::min(ne0203, (int64_t)65535));
+    dequantize_block<qk, qr, dequantize_kernel><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>
+        (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03);
+}
+
+template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
+static void dequantize_block_cont_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k, cudaStream_t stream) {
+    dequantize_block_cuda<qk, qr, dequantize_kernel, dst_t>(vx, y, k, 1, 1, 1, k/qk, k/qk, k/qk, stream);
+}
+
+static void dequantize_block_q8_0_f16_cuda(const void * __restrict__ vx, half * __restrict__ y, const int64_t k, cudaStream_t stream) {
+    const int num_blocks = (k + CUDA_Q8_0_NE_ALIGN - 1) / CUDA_Q8_0_NE_ALIGN;
+    if (k % CUDA_Q8_0_NE_ALIGN == 0) {
+        const bool need_check = false;
+        dequantize_block_q8_0_f16<need_check><<<num_blocks, WARP_SIZE, 0, stream>>>(vx, y, k);
+    } else {
+        const bool need_check = true;
+        dequantize_block_q8_0_f16<need_check><<<num_blocks, WARP_SIZE, 0, stream>>>(vx, y, k);
+    }
+}
+
+template<typename dst_t>
+static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_q4_0_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb32 = k / 32;
+    const int nb = (k + 255) / 256;
+    dequantize_block_q4_0<<<nb, 32, 0, stream>>>(vx, y, nb32);
+}
+
+template<typename dst_t>
+static void dequantize_row_q4_1_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb32 = k / 32;
+    const int nb = (k + 255) / 256;
+    dequantize_block_q4_1<<<nb, 32, 0, stream>>>(vx, y, nb32);
+}
+
+template<typename dst_t>
+static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q4_K<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q5_K<<<nb, 64, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q6_K<<<nb, 64, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq2_xxs<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq2_xs<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq2_s<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq3_xxs<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq3_s<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq1_s<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb32 = k / QK4_NL;
+    const int nb = (k + QK_K - 1) / QK_K;
+    dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y, nb32);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_iq1_m<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = (k + QK_K - 1) / QK_K;
+    dequantize_block_iq4_xs<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = (k + QK_K - 1) / QK_K;
+    dequantize_block_mxfp4<<<nb, 32, 0, stream>>>(vx, y);
+}
+
+template <typename dst_t>
+static __global__ void dequantize_block_nvfp4(
+        const void * __restrict__ vx,
+        dst_t * __restrict__ yy,
+        const int64_t ne) {
+    const int64_t i = blockIdx.x;
+    const int     tid = threadIdx.x;
+
+    const int64_t base = i * QK_NVFP4;
+    if (base >= ne) {
+        return;
+    }
+
+    const block_nvfp4 * x = (const block_nvfp4 *) vx;
+    const block_nvfp4 & xb = x[i];
+
+    const int sub = tid / (QK_NVFP4_SUB / 2);
+    const int j = tid % (QK_NVFP4_SUB / 2);
+
+    const float d = ggml_cuda_ue4m3_to_fp32(xb.d[sub]);
+    const uint8_t q = xb.qs[sub * (QK_NVFP4_SUB / 2) + j];
+
+    const int64_t y0 = base + sub * QK_NVFP4_SUB + j;
+    const int64_t y1 = y0 + QK_NVFP4_SUB / 2;
+
+    yy[y0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q & 0x0F]);
+    yy[y1] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q >> 4]);
+}
+
+template <typename dst_t>
+static void dequantize_row_nvfp4_cuda(
+        const void * vx,
+        dst_t * y,
+        const int64_t k,
+        cudaStream_t stream) {
+    GGML_ASSERT(k % QK_NVFP4 == 0);
+    const int nb = k / QK_NVFP4;
+    dequantize_block_nvfp4<<<nb, 32, 0, stream>>>(vx, y, k);
+}
+template <typename src_t, typename dst_t>
+static __global__ void convert_unary(
+        const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t ne00, const int64_t ne01,
+        const int64_t ne0203, const uint3 ne02,
+        const int64_t s01, const int64_t s02, const int64_t s03) {
+    const int64_t i00 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
+
+    if (i00 >= ne00) {
+        return;
+    }
+
+    const src_t * x = (const src_t *) vx;
+
+    for (int64_t i01 = blockIdx.y; i01 < ne01; i01 += gridDim.y) {
+        for (int64_t i0203 = blockIdx.z; i0203 < ne0203; i0203 += gridDim.z) {
+            const uint2 dm = fast_div_modulo((uint32_t)i0203, ne02);
+            const int64_t i02 = dm.y;
+            const int64_t i03 = dm.x;
+
+            const int64_t ix = i03*s03 + i02*s02 + i01*s01 + i00;
+            const int64_t iy = (i0203*ne01 + i01)*ne00 + i00;
+            y[iy] = ggml_cuda_cast<dst_t>(x[ix]);
+        }
+    }
+}
+
+template <typename T> struct alignas(sizeof(T)*4) cvt_vec4 { T v[4]; };
+
+// four elements per thread, so a warp moves 512B (RDNA) / 1k (CDNA) per load
+template <typename src_t, typename dst_t>
+static __global__ void convert_unary_cont_vec4(
+        const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k4) {
+    const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
+
+    if (i >= k4) {
+        return;
+    }
+
+    const cvt_vec4<src_t> xv = ((const cvt_vec4<src_t> *) vx)[i];
+
+    cvt_vec4<dst_t> yv;
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        yv.v[j] = ggml_cuda_cast<dst_t>(xv.v[j]);
+    }
+
+    ((cvt_vec4<dst_t> *) y)[i] = yv;
+}
+
+template <typename src_t, typename dst_t>
+static void convert_unary_cuda(const void * vx, dst_t * y,
+        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
+        const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) {
+    const int64_t ne0203 = ne02*ne03;
+    const uint3 ne02_fdv = init_fastdiv_values(ne02);
+    const dim3 num_blocks((ne00 + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE, (int)std::min(ne01, (int64_t)65535), (int)std::min(ne0203, (int64_t)65535));
+    convert_unary<src_t><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>
+        (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03);
+}
+
+template <typename src_t, typename dst_t>
+static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    if (k % 4 == 0 &&
+        (uintptr_t) vx % alignof(cvt_vec4<src_t>) == 0 &&
+        (uintptr_t) y  % alignof(cvt_vec4<dst_t>) == 0) {
+        const int64_t k4 = k/4;
+        const int64_t num_blocks = (k4 + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE;
+        convert_unary_cont_vec4<src_t, dst_t><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(vx, y, k4);
+        return;
+    }
+
+    convert_unary_cuda<src_t>(vx, y, k, 1, 1, 1, k, k, k, stream);
+}
+
+to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_row_q4_0_cuda;
+        case GGML_TYPE_Q4_1:
+            return dequantize_row_q4_1_cuda;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cont_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cont_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cont_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_Q2_K:
+            return dequantize_row_q2_K_cuda;
+        case GGML_TYPE_Q3_K:
+            return dequantize_row_q3_K_cuda;
+        case GGML_TYPE_Q4_K:
+            return dequantize_row_q4_K_cuda;
+        case GGML_TYPE_Q5_K:
+            return dequantize_row_q5_K_cuda;
+        case GGML_TYPE_Q6_K:
+            return dequantize_row_q6_K_cuda;
+        case GGML_TYPE_IQ2_XXS:
+            return dequantize_row_iq2_xxs_cuda;
+        case GGML_TYPE_IQ2_XS:
+            return dequantize_row_iq2_xs_cuda;
+        case GGML_TYPE_IQ2_S:
+            return dequantize_row_iq2_s_cuda;
+        case GGML_TYPE_IQ3_XXS:
+            return dequantize_row_iq3_xxs_cuda;
+        case GGML_TYPE_IQ1_S:
+            return dequantize_row_iq1_s_cuda;
+        case GGML_TYPE_IQ1_M:
+            return dequantize_row_iq1_m_cuda;
+        case GGML_TYPE_IQ4_NL:
+            return dequantize_row_iq4_nl_cuda;
+        case GGML_TYPE_IQ4_XS:
+            return dequantize_row_iq4_xs_cuda;
+        case GGML_TYPE_IQ3_S:
+            return dequantize_row_iq3_s_cuda;
+        case GGML_TYPE_MXFP4:
+            return dequantize_row_mxfp4_cuda;
+        case GGML_TYPE_NVFP4:
+            return dequantize_row_nvfp4_cuda;
+        case GGML_TYPE_F32:
+            return convert_unary_cont_cuda<float>;
+        case GGML_TYPE_F16:
+            return convert_unary_cont_cuda<half>;
+        default:
+            return nullptr;
+    }
+}
+
+to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_row_q4_0_cuda;
+        case GGML_TYPE_Q4_1:
+            return dequantize_row_q4_1_cuda;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cont_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cont_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            if (fp16_available(ggml_cuda_info().devices[ggml_cuda_get_device()].cc)) {
+                return dequantize_block_q8_0_f16_cuda;
+            }
+            return dequantize_block_cont_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_Q2_K:
+            return dequantize_row_q2_K_cuda;
+        case GGML_TYPE_Q3_K:
+            return dequantize_row_q3_K_cuda;
+        case GGML_TYPE_Q4_K:
+            return dequantize_row_q4_K_cuda;
+        case GGML_TYPE_Q5_K:
+            return dequantize_row_q5_K_cuda;
+        case GGML_TYPE_Q6_K:
+            return dequantize_row_q6_K_cuda;
+        case GGML_TYPE_IQ2_XXS:
+            return dequantize_row_iq2_xxs_cuda;
+        case GGML_TYPE_IQ2_XS:
+            return dequantize_row_iq2_xs_cuda;
+        case GGML_TYPE_IQ2_S:
+            return dequantize_row_iq2_s_cuda;
+        case GGML_TYPE_IQ3_XXS:
+            return dequantize_row_iq3_xxs_cuda;
+        case GGML_TYPE_IQ1_S:
+            return dequantize_row_iq1_s_cuda;
+        case GGML_TYPE_IQ1_M:
+            return dequantize_row_iq1_m_cuda;
+        case GGML_TYPE_IQ4_NL:
+            return dequantize_row_iq4_nl_cuda;
+        case GGML_TYPE_IQ4_XS:
+            return dequantize_row_iq4_xs_cuda;
+        case GGML_TYPE_IQ3_S:
+            return dequantize_row_iq3_s_cuda;
+        case GGML_TYPE_MXFP4:
+            return dequantize_row_mxfp4_cuda;
+        case GGML_TYPE_NVFP4:
+            return dequantize_row_nvfp4_cuda;
+        case GGML_TYPE_F32:
+            return convert_unary_cont_cuda<float>;
+        case GGML_TYPE_BF16:
+            return convert_unary_cont_cuda<nv_bfloat16>;
+        default:
+            return nullptr;
+    }
+}
+
+to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_row_q4_0_cuda;
+        case GGML_TYPE_Q4_1:
+            return dequantize_row_q4_1_cuda;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cont_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cont_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cont_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_Q2_K:
+            return dequantize_row_q2_K_cuda;
+        case GGML_TYPE_Q3_K:
+            return dequantize_row_q3_K_cuda;
+        case GGML_TYPE_Q4_K:
+            return dequantize_row_q4_K_cuda;
+        case GGML_TYPE_Q5_K:
+            return dequantize_row_q5_K_cuda;
+        case GGML_TYPE_Q6_K:
+            return dequantize_row_q6_K_cuda;
+        case GGML_TYPE_IQ2_XXS:
+            return dequantize_row_iq2_xxs_cuda;
+        case GGML_TYPE_IQ2_XS:
+            return dequantize_row_iq2_xs_cuda;
+        case GGML_TYPE_IQ2_S:
+            return dequantize_row_iq2_s_cuda;
+        case GGML_TYPE_IQ3_XXS:
+            return dequantize_row_iq3_xxs_cuda;
+        case GGML_TYPE_IQ1_S:
+            return dequantize_row_iq1_s_cuda;
+        case GGML_TYPE_IQ1_M:
+            return dequantize_row_iq1_m_cuda;
+        case GGML_TYPE_IQ4_NL:
+            return dequantize_row_iq4_nl_cuda;
+        case GGML_TYPE_IQ4_XS:
+            return dequantize_row_iq4_xs_cuda;
+        case GGML_TYPE_IQ3_S:
+            return dequantize_row_iq3_s_cuda;
+        case GGML_TYPE_MXFP4:
+            return dequantize_row_mxfp4_cuda;
+        case GGML_TYPE_NVFP4:
+            return dequantize_row_nvfp4_cuda;
+        case GGML_TYPE_F16:
+            return convert_unary_cont_cuda<half>;
+        case GGML_TYPE_BF16:
+            return convert_unary_cont_cuda<nv_bfloat16>;
+        default:
+            return nullptr;
+    }
+}
+
+to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+            return convert_unary_cuda<float>;
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
+        case GGML_TYPE_Q4_1:
+            return dequantize_block_cuda<QK4_1, QR4_1, dequantize_q4_1>;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_BF16:
+            return convert_unary_cuda<nv_bfloat16>;
+        default:
+            return nullptr;
+    }
+}
+
+to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+            return convert_unary_cuda<float, nv_bfloat16>;
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
+        case GGML_TYPE_Q4_1:
+            return dequantize_block_cuda<QK4_1, QR4_1, dequantize_q4_1>;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_F16:
+            return convert_unary_cuda<half, nv_bfloat16>;
+        default:
+            return nullptr;
+    }
+}
+
+to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F16:
+            return convert_unary_cuda<half, float>;
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
+        case GGML_TYPE_Q4_1:
+            return dequantize_block_cuda<QK4_1, QR4_1, dequantize_q4_1>;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_BF16:
+            return convert_unary_cuda<nv_bfloat16, float>;
+        default:
+            return nullptr;
+    }
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh
new file mode 100644
index 0000000000..91255ccf31
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh
@@ -0,0 +1,66 @@
+#pragma once
+#include "common.cuh"
+
+#define CUDA_DEQUANTIZE_BLOCK_SIZE 256
+
+template<typename T>
+using to_t_cuda_t = void (*)(const void * x, T * y, int64_t k, cudaStream_t stream);
+
+typedef to_t_cuda_t<float> to_fp32_cuda_t;
+typedef to_t_cuda_t<half> to_fp16_cuda_t;
+typedef to_t_cuda_t<nv_bfloat16> to_bf16_cuda_t;
+
+to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type);
+
+to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type);
+
+to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type);
+
+// TODO more general support for non-contiguous inputs
+
+template<typename T>
+using to_t_nc_cuda_t = void (*)(const void * x, T * y,
+    int64_t ne00, int64_t ne01, int64_t ne02, int64_t ne03,
+    int64_t s01, int64_t s02, int64_t s03, cudaStream_t stream);
+
+typedef to_t_nc_cuda_t<float> to_fp32_nc_cuda_t;
+typedef to_t_nc_cuda_t<half> to_fp16_nc_cuda_t;
+typedef to_t_nc_cuda_t<nv_bfloat16> to_bf16_nc_cuda_t;
+
+to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type);
+to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type);
+to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type);
+
+template<typename dst_t, typename src_t>
+ __host__ __device__ inline dst_t ggml_cuda_cast(src_t x) {
+    if constexpr (std::is_same_v<dst_t, src_t>) {
+        return x;
+    } else if constexpr(std::is_same_v<dst_t, nv_bfloat16>) {
+        return __float2bfloat16(float(x));
+    } else if constexpr(std::is_same_v<src_t, nv_bfloat16>) {
+        return __bfloat162float(x);
+    } else if constexpr(std::is_same_v<src_t, float2> && std::is_same_v<dst_t, half2>) {
+        return __float22half2_rn(x);
+    } else if constexpr(std::is_same_v<src_t, nv_bfloat162> && std::is_same_v<dst_t, float2>) {
+#ifdef GGML_USE_HIP
+        return make_float2(__bfloat162float(__low2bfloat16(x)), __bfloat162float(__high2bfloat16(x)));
+#else
+#if !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800
+        return __bfloat1622float2(x);
+#else
+        return make_float2(__bfloat162float(x.x), __bfloat162float(x.y));
+#endif // !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800
+#endif // GGML_USE_HIP
+    } else if constexpr(std::is_same_v<src_t, float2> && std::is_same_v<dst_t, nv_bfloat162>) {
+        // bypass compile error on cuda 12.0.1
+#ifdef GGML_USE_HIP
+        return __float22bfloat162_rn(x);
+#else
+        return {x.x, x.y};
+#endif // GGML_USE_HIP
+    } else if constexpr(std::is_same_v<dst_t, int32_t>) {
+        return int32_t(x);
+    } else {
+        return float(x);
+    }
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh
new file mode 100644
index 0000000000..651524be7b
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh
@@ -0,0 +1,452 @@
+#include "common.cuh"
+#include "convert.cuh"
+
+static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q1_0 * x = (const block_q1_0 *) vx;
+
+    const float d = x[ib].d;
+
+    const int bit_index_0 = iqs;
+    const int bit_index_1 = iqs + 1;
+
+    const int byte_index_0 = bit_index_0 / 8;
+    const int bit_offset_0 = bit_index_0 % 8;
+
+    const int byte_index_1 = bit_index_1 / 8;
+    const int bit_offset_1 = bit_index_1 % 8;
+
+    // Extract bits: 1 = +d, 0 = -d (branchless)
+    const int bit_0 = (x[ib].qs[byte_index_0] >> bit_offset_0) & 1;
+    const int bit_1 = (x[ib].qs[byte_index_1] >> bit_offset_1) & 1;
+
+    v.x = (2*bit_0 - 1) * d;
+    v.y = (2*bit_1 - 1) * d;
+}
+
+static __device__ __forceinline__ void dequantize_q2_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q2_0 * x = (const block_q2_0 *) vx;
+
+    const float d = x[ib].d;
+
+    // Q2_0: 2 bits per element, 4 elements per byte.
+    // Stored code c in {0,1,2,3} maps to symbol s = c - 1 in {-1, 0, +1, +2}.
+    const int byte_index_0 = iqs / 4;
+    const int bit_offset_0 = (iqs % 4) * 2;
+
+    const int byte_index_1 = (iqs + 1) / 4;
+    const int bit_offset_1 = ((iqs + 1) % 4) * 2;
+
+    const int c0 = (x[ib].qs[byte_index_0] >> bit_offset_0) & 0x3;
+    const int c1 = (x[ib].qs[byte_index_1] >> bit_offset_1) & 0x3;
+
+    v.x = (c0 - 1) * d;
+    v.y = (c1 - 1) * d;
+}
+
+static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q4_0 * x = (const block_q4_0 *) vx;
+
+    const float d = x[ib].d;
+
+    const int vui = x[ib].qs[iqs];
+
+    v.x = vui & 0xF;
+    v.y = vui >> 4;
+
+    v.x = (v.x - 8.0f) * d;
+    v.y = (v.y - 8.0f) * d;
+}
+
+static __device__ __forceinline__ void dequantize_q4_1(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q4_1 * x = (const block_q4_1 *) vx;
+
+    const float2 dm = __half22float2(x[ib].dm);
+
+    const int vui = x[ib].qs[iqs];
+
+    v.x = vui & 0xF;
+    v.y = vui >> 4;
+
+    v.x = (v.x * dm.x) + dm.y;
+    v.y = (v.y * dm.x) + dm.y;
+}
+
+static __device__ __forceinline__ void dequantize_q5_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q5_0 * x = (const block_q5_0 *) vx;
+
+    const float d = x[ib].d;
+
+    uint32_t qh;
+    memcpy(&qh, x[ib].qh, sizeof(qh));
+
+    const int xh_0 = ((qh >> (iqs +  0)) << 4) & 0x10;
+    const int xh_1 = ((qh >> (iqs + 12))     ) & 0x10;
+
+    v.x = ((x[ib].qs[iqs] & 0xf) | xh_0);
+    v.y = ((x[ib].qs[iqs] >>  4) | xh_1);
+
+    v.x = (v.x - 16.0f) * d;
+    v.y = (v.y - 16.0f) * d;
+}
+
+static __device__ __forceinline__ void dequantize_q5_1(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q5_1 * x = (const block_q5_1 *) vx;
+
+    const float2 dm = __half22float2(x[ib].dm);
+
+    uint32_t qh;
+    memcpy(&qh, x[ib].qh, sizeof(qh));
+
+    const int xh_0 = ((qh >> (iqs +  0)) << 4) & 0x10;
+    const int xh_1 = ((qh >> (iqs + 12))     ) & 0x10;
+
+    v.x = ((x[ib].qs[iqs] & 0xf) | xh_0);
+    v.y = ((x[ib].qs[iqs] >>  4) | xh_1);
+
+    v.x = (v.x * dm.x) + dm.y;
+    v.y = (v.y * dm.x) + dm.y;
+}
+
+static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+    const block_q8_0 * x = (const block_q8_0 *) vx;
+
+    const float d = x[ib].d;
+
+    v.x = x[ib].qs[iqs + 0];
+    v.y = x[ib].qs[iqs + 1];
+
+    v.x *= d;
+    v.y *= d;
+}
+
+//================================== k-quants
+
+// Each call dequantizes one super-block of QK_K values into y using the
+// thread layout of the caller: 32 threads for q4_K, 64 threads otherwise.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q2_K * x = (const block_q2_K *) vx;
+
+    const int64_t n   = tid/32;
+    const int64_t l   = tid - 32*n;
+    const int64_t is  = 8*n + l/16;
+
+    const uint8_t q = x[ib].qs[32*n + l];
+    dst_t * y = yy + 128*n;
+
+    float dall = __low2half(x[ib].dm);
+    float dmin = __high2half(x[ib].dm);
+    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[ib].scales[is+0] >> 4));
+    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[ib].scales[is+2] >> 4));
+    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[ib].scales[is+4] >> 4));
+    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4));
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q3_K * x = (const block_q3_K *) vx;
+
+    const int64_t r = tid/4;
+    const int64_t t = r/2;
+    const int64_t is0 = r%2;
+    const int64_t l0 = 16*is0 + 4*(tid%4);
+    const int64_t n = t / 4;
+    const int64_t j = t - 4*n;
+
+    uint8_t m = 1 << (4*n + j);
+    int64_t is = 8*n + 2*j + is0;
+    int shift = 2*j;
+
+    int8_t us = is <  4 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+8] >> 0) & 3) << 4) :
+                is <  8 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+4] >> 2) & 3) << 4) :
+                is < 12 ? (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is+0] >> 4) & 3) << 4) :
+                          (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is-4] >> 6) & 3) << 4);
+    float d_all = x[ib].d;
+    float dl = d_all * (us - 32);
+
+    dst_t * y = yy + 128*n + 32*j;
+    const uint8_t * q = x[ib].qs + 32*n;
+    const uint8_t * hm = x[ib].hmask;
+
+    for (int l = l0; l < l0+4; ++l) {
+        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
+    }
+}
+
+static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
+    if (j < 4) {
+        d = q[j] & 63; m = q[j + 4] & 63;
+    } else {
+        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
+        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q4_K * x = (const block_q4_K *) vx;
+
+    // assume 32 threads
+    const int64_t il  = tid/8;
+    const int64_t ir  = tid%8;
+    const int64_t is  = 2*il;
+    const int64_t n   = 4;
+
+    dst_t * y = yy + 64*il + n*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * q = x[ib].qs + 32*il + n*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+    for (int l = 0; l < n; ++l) {
+        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
+        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q5_K * x = (const block_q5_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t il  = tid/16;   // il is in 0...3
+    const int64_t ir  = tid%16;   // ir is in 0...15
+    const int64_t is  = 2*il;     // is is in 0...6
+
+    dst_t * y = yy + 64*il + 2*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * ql = x[ib].qs + 32*il + 2*ir;
+    const uint8_t * qh = x[ib].qh + 2*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+
+    uint8_t   hm  = 1 << (2*il);
+    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
+    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
+    hm <<= 1;
+    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
+    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q6_K * x = (const block_q6_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t ip  = tid/32;   // ip is 0 or 1
+    const int64_t il  = tid - 32*ip; // 0...32
+    const int64_t is  = 8*ip + il/16;
+
+    dst_t * y = yy + 128*ip + il;
+
+    const float d = x[ib].d;
+
+    const uint8_t * ql = x[ib].ql + 64*ip + il;
+    const uint8_t   qh = x[ib].qh[32*ip + il];
+    const int8_t  * sc = x[ib].scales + is;
+
+    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
+    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
+    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
+    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+}
+
+//================================== i-quants
+
+// Each call dequantizes one super-block of QK_K values into y with 32
+// threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * aux8 = (const uint8_t *)q2;
+    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
+    const uint32_t aux32 = q2[2] | (q2[3] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xs * x = (const block_iq2_xs *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_s * x = (const block_iq2_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300)));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t  * q3 = x[ibs].qs + 8*ib;
+    const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib;
+    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
+    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
+    const uint32_t aux32 = gas[0] | (gas[1] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.5f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_s * x = (const block_iq3_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * qs = x[ibs].qs + 8*ib;
+    const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256)));
+    const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256)));
+    const float d = (float)x[ibs].d * (1 + 2*((x[ibs].scales[ib/2] >> 4*(ib%2)) & 0xf));
+    const uint8_t signs = x[ibs].signs[4*ib + il];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_s * x = (const block_iq1_s  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
+    const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1);
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[ib] >> 3*il) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_m * x = (const block_iq1_m  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * sc = (const uint16_t *)x[ibs].scales;
+    iq1m_scale_t scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
+    const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
+    const float delta = x[ibs].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = (float)x[ib].d;
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+    const block_iq4_xs * x = (const block_iq4_xs *)vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ibs].qs + 16*ib + 4*il;
+    const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
+    }
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu
new file mode 100644
index 0000000000..d97e8fdb0f
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu
@@ -0,0 +1,488 @@
+#include "getrows.cuh"
+#include "dequantize.cuh"
+#include "convert.cuh"
+
+template<int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
+static __global__ void k_get_rows(
+        const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
+        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+        /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+    ggml_cuda_pdl_sync();
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        for (int64_t i00 = 2*(blockIdx.y*blockDim.x + threadIdx.x); i00 < ne00; i00 += gridDim.y*blockDim.x) {
+            // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+            const int i10 =  blockIdx.x;
+            const uint2 dm  = fast_div_modulo((uint32_t)z, ne12_fdv);
+            const int i11 =  dm.x;
+            const int i12 =  dm.y;
+
+            const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+            dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+            const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+            const int ib   =  i00/qk;      // block index
+            const int iqs  = (i00%qk)/qr;  // quant index
+            const int iybs = i00 - i00%qk; // dst block start index
+            const int y_offset = qr == 1 ? 1 : qk/2;
+
+            // dequantize
+            float2 v;
+            dequantize_kernel(src0_row, ib, iqs, v);
+
+            dst_row[iybs + iqs + 0]        = ggml_cuda_cast<dst_t>(v.x);
+            dst_row[iybs + iqs + y_offset] = ggml_cuda_cast<dst_t>(v.y);
+        }
+    }
+}
+
+template<typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static __global__ void k_get_rows_kq(
+        const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
+        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+        /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+    ggml_cuda_pdl_sync();
+    const int64_t nsb = ne00/QK_K; // super-blocks per row
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm  = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+        for (int64_t ib = blockIdx.y; ib < nsb; ib += gridDim.y) {
+            dequantize_kq(src0_row, ib, dst_row + ib*QK_K, threadIdx.x);
+        }
+    }
+}
+
+template<typename src0_t, typename dst_t>
+static __global__ void k_get_rows_float(
+        const src0_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
+        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+        /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+    ggml_cuda_pdl_lc();
+    const src0_t  * GGML_CUDA_RESTRICT src0 = src0_ptr;
+    const int32_t * GGML_CUDA_RESTRICT src1 = src1_ptr;
+    dst_t         * GGML_CUDA_RESTRICT dst  = dst_ptr;
+    ggml_cuda_pdl_sync();
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
+
+        for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
+            dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
+        }
+    }
+}
+
+template<typename dst_t>
+static __global__ void k_get_rows_float_vec(
+        const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
+        const int64_t ne00v,
+        const int64_t ne11, const uint3 ne12_fdv,
+        const size_t s1, const size_t s2, const size_t s3,
+        const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12) {
+
+    ggml_cuda_pdl_lc();
+    ggml_cuda_pdl_sync();
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];
+
+        int4       * GGML_CUDA_RESTRICT dst_row  = (int4 *)      (dst_ptr + i10*s1 + i11*s2 + i12*s3);
+        const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03);
+
+        for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) {
+            dst_row[i] = src0_row[i];
+        }
+    }
+}
+
+template<typename grad_t, typename dst_t>
+static __global__ void k_get_rows_back_float(
+        const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst,
+        const int64_t ncols, const int64_t nrows_grad, const int64_t nrows_dst) {
+    const int col = blockIdx.x*blockDim.x + threadIdx.x;
+
+    if (col >= ncols) {
+        return;
+    }
+
+    ggml_cuda_pdl_sync();
+
+    // grid.y is clamped to the CUDA grid limit, so stride over the destination rows
+    for (int64_t dst_row = blockIdx.y; dst_row < nrows_dst; dst_row += gridDim.y) {
+        float sum = 0.0f;
+
+        for (int64_t i = 0; i < nrows_grad; ++i) {
+            if (rows[i] != dst_row) {
+                continue;
+            }
+            sum += grad[i*ncols + col];
+        }
+
+        dst[dst_row*ncols + col] = sum;
+    }
+}
+
+template<int qk, int qr, dequantize_kernel_t dq, typename dst_t>
+static void get_rows_cuda_q(
+        const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
+    const int block_num_y = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE);
+    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
+    // strides in elements
+    // const size_t s0 = nb0 / sizeof(dst_t);
+    const size_t s1 = nb1 / sizeof(dst_t);
+    const size_t s2 = nb2 / sizeof(dst_t);
+    const size_t s3 = nb3 / sizeof(dst_t);
+
+    const size_t s10 = nb10 / sizeof(int32_t);
+    const size_t s11 = nb11 / sizeof(int32_t);
+    const size_t s12 = nb12 / sizeof(int32_t);
+    // const size_t s13 = nb13 / sizeof(int32_t);
+
+    GGML_ASSERT(ne00 % 2 == 0);
+
+    GGML_ASSERT(ne12 > 0);
+    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+    const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+    k_get_rows<qk, qr, dq><<<block_nums, block_dims, 0, stream>>>(
+        src0_d, src1_d, dst_d,
+        ne00, /*ne01, ne02, ne03,*/
+        /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+        /* s0,*/ s1, s2, s3,
+        /* nb00,*/ nb01, nb02, nb03,
+        s10, s11, s12/*, s13*/);
+}
+
+template<int block_dim, typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static void get_rows_cuda_kq(
+        const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    GGML_ASSERT(ne00 % QK_K == 0);
+    const int64_t nsb = ne00/QK_K;
+
+    const dim3 block_dims(block_dim, 1, 1);
+    const dim3 block_nums(ne10, MIN(nsb, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
+    // strides in elements
+    // const size_t s0 = nb0 / sizeof(dst_t);
+    const size_t s1 = nb1 / sizeof(dst_t);
+    const size_t s2 = nb2 / sizeof(dst_t);
+    const size_t s3 = nb3 / sizeof(dst_t);
+
+    const size_t s10 = nb10 / sizeof(int32_t);
+    const size_t s11 = nb11 / sizeof(int32_t);
+    const size_t s12 = nb12 / sizeof(int32_t);
+    // const size_t s13 = nb13 / sizeof(int32_t);
+
+    GGML_ASSERT(ne12 > 0);
+    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+    const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+    k_get_rows_kq<dst_t, dequantize_kq><<<block_nums, block_dims, 0, stream>>>(
+        src0_d, src1_d, dst_d,
+        ne00, /*ne01, ne02, ne03,*/
+        /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+        /* s0,*/ s1, s2, s3,
+        /* nb00,*/ nb01, nb02, nb03,
+        s10, s11, s12/*, s13*/);
+}
+
+template<typename src0_t, typename dst_t>
+static void get_rows_cuda_float(
+        const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
+
+    // strides in elements
+    // const size_t s0 = nb0 / sizeof(dst_t);
+    const size_t s1 = nb1 / sizeof(dst_t);
+    const size_t s2 = nb2 / sizeof(dst_t);
+    const size_t s3 = nb3 / sizeof(dst_t);
+
+    const size_t s10 = nb10 / sizeof(int32_t);
+    const size_t s11 = nb11 / sizeof(int32_t);
+    const size_t s12 = nb12 / sizeof(int32_t);
+    // const size_t s13 = nb13 / sizeof(int32_t);
+
+    GGML_ASSERT(ne12 > 0);
+    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+    const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+    if constexpr (std::is_same<src0_t, dst_t>::value) {
+        constexpr int VEC = 16 / sizeof(dst_t);
+        const int64_t ne00v = ne00 / VEC;
+        const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+        const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128;
+        const bool can_vec = VEC > 1 && enough_blocks &&
+            (ne00 % VEC == 0) &&
+            (nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) &&
+            (nb1  % 16 == 0) && (nb2  % 16 == 0) && (nb3  % 16 == 0) &&
+            (((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0);
+
+        if (can_vec) {
+            const int block_num_y = vec_block_num_y;
+            const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
+            ggml_cuda_kernel_launch(k_get_rows_float_vec<dst_t>, launch_params,
+                (const dst_t *) src0_d, src1_d, dst_d,
+                ne00v, ne11, ne12_fdv,
+                s1, s2, s3,
+                nb01, nb02, nb03,
+                s10, s11, s12);
+            return;
+        }
+    }
+
+    const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
+    ggml_cuda_kernel_launch(k_get_rows_float<src0_t, dst_t>, launch_params,
+        src0_d, src1_d, dst_d,
+        ne00, /*ne01, ne02, ne03,*/
+        /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+        /* s0,*/ s1, s2, s3,
+        /* nb00,*/ nb01, nb02, nb03,
+        s10, s11, s12/*, s13*/);
+}
+
+template <typename dst_t>
+static void ggml_cuda_get_rows_switch_src0_type(
+        const void * src0_d, const ggml_type src0_type, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    switch (src0_type) {
+        case GGML_TYPE_F16:
+            get_rows_cuda_float((const half *) src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_F32:
+            get_rows_cuda_float((const float *) src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_I32:
+            get_rows_cuda_float((const int32_t *) src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_BF16:
+            get_rows_cuda_float((const nv_bfloat16 *) src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q1_0:
+            get_rows_cuda_q<QK1_0, QR1_0, dequantize_q1_0>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q2_0:
+            get_rows_cuda_q<QK2_0, QR2_0, dequantize_q2_0>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q4_0:
+            get_rows_cuda_q<QK4_0, QR4_0, dequantize_q4_0>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q4_1:
+            get_rows_cuda_q<QK4_1, QR4_1, dequantize_q4_1>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q5_0:
+            get_rows_cuda_q<QK5_0, QR5_0, dequantize_q5_0>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q5_1:
+            get_rows_cuda_q<QK5_1, QR5_1, dequantize_q5_1>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q8_0:
+            get_rows_cuda_q<QK8_0, QR8_0, dequantize_q8_0>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q2_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q2_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q3_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            get_rows_cuda_kq<32, dst_t, dequantize_q4_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q5_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q6_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_M:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_m<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_nl<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_MXFP4:
+            get_rows_cuda_kq<32, dst_t, dequantize_mxfp4<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        default:
+            GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type));
+    }
+}
+
+void get_rows_cuda(
+        const void * src0_d, ggml_type src0_type, const int32_t * src1_d, void * dst_d, ggml_type dst_type,
+        int64_t ne00, size_t nb01, size_t nb02, size_t nb03,
+        int64_t ne10, int64_t ne11, int64_t ne12, size_t nb10, size_t nb11, size_t nb12,
+        size_t nb1, size_t nb2, size_t nb3,
+        cudaStream_t stream) {
+    switch (dst_type) {
+        case GGML_TYPE_F32:
+            ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (float *) dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_I32:
+            ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (int32_t *) dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_F16:
+            ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (half *) dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_BF16:
+            ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (nv_bfloat16 *) dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        default:
+            GGML_ABORT("%s: unsupported dst type: %s\n", __func__, ggml_type_name(dst_type));
+    }
+}
+
+void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    cudaStream_t stream = ctx.stream();
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    GGML_ASSERT(src1->type == GGML_TYPE_I32);
+    GGML_ASSERT(ne13 == 1);
+
+    GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type));
+    GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type));
+    GGML_ASSERT(dst->nb[0]  == ggml_type_size(dst->type));
+
+    get_rows_cuda(src0->data, src0->type, (const int32_t *) src1->data, dst->data, dst->type,
+        ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+}
+
+void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0]; // gradients of forward pass output
+    const ggml_tensor * src1 = dst->src[1]; // src1 in forward pass
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const float   * src0_d = (const float   *) src0->data;
+    const int32_t * src1_d = (const int32_t *) src1->data;
+    float         * dst_d  = (float         *) dst->data;
+
+    cudaStream_t stream = ctx.stream();
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_I32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    GGML_ASSERT(ne02*ne03 == 1);
+    GGML_ASSERT(ne12*ne13 == 1);
+    GGML_ASSERT(ne2*ne3 == 1);
+
+    const dim3 block_dims(CUDA_GET_ROWS_BACK_BLOCK_SIZE, 1, 1);
+    const int block_num_x = (ne00 + CUDA_GET_ROWS_BACK_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BACK_BLOCK_SIZE;
+    const dim3 block_nums(block_num_x, MIN(ne1, (int64_t)UINT16_MAX), 1);
+
+    k_get_rows_back_float<<<block_nums, block_dims, 0, stream>>>(src0_d, src1_d, dst_d, ne00, ne10, ne1);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh
new file mode 100644
index 0000000000..3c5bea5f48
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh
@@ -0,0 +1,15 @@
+#include "common.cuh"
+
+#define CUDA_GET_ROWS_BLOCK_SIZE 256
+#define CUDA_GET_ROWS_BACK_BLOCK_SIZE 256
+
+void get_rows_cuda(
+        const void * src0_d, ggml_type src0_type, const int32_t * src1_d, void * dst_d, ggml_type dst_type,
+        int64_t ne00, size_t nb01, size_t nb02, size_t nb03,
+        int64_t ne10, int64_t ne11, int64_t ne12, size_t nb10, size_t nb11, size_t nb12,
+        size_t nb1, size_t nb2, size_t nb3,
+        cudaStream_t stream);
+
+void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh
new file mode 100644
index 0000000000..93a51aa0be
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh
@@ -0,0 +1,1594 @@
+#pragma once
+// This file contains primitives that expose the tensor core PTX instructions for CUDA code.
+// The primitives can be used in a similar way as the nvcuda::wmma interface but with a well-defined memory layout.
+// The documentation for the PTX instructions can be found under:
+//   https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-multiply-accumulate-operation-using-mma-instruction
+//
+// Like with nvcuda::wmma there are three types of matrix tiles: A, B, and C with A @ B = C.
+// A is a row-major matrix with shape M x K.
+// B is a column-major matrix with shape K x N.
+// C is a column-major matrix with shape M x N.
+// A, B, and C are represented using the same fundamental data type: a row-major matrix with I rows and J columns.
+// Note that J is measured in physical 32 bit elements instead of logical elements.
+// The methods get_i and get_j can be used to get the physical 32 bit index of the lth element of a thread within a tile.
+// All matrix tiles have ne physical 32 bit elements per warp.
+//
+// As described in the PTX documentation, all pointers for load_ldmatrix must be to shared memory and aligned to 16 bytes.
+// The API in this file also assumes that the pointers for load_generic are aligned to 16 bytes, unaligned pointers are considered undefined behavior.
+
+#include "common.cuh"
+
+// On Volta each warp is doing 4 8x8 mma operations in parallel.
+// The basic memory layout for a 32x8 output tile is to stack 4 input tiles in I direction and to mirror the B tile.
+// However, the i indices in this file are by default permuted to simplify the index calculations.
+// #define GGML_CUDA_MMA_NO_VOLTA_PERM
+
+#if CUDART_VERSION >= 11080
+
+static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) {
+    int ret = 0;
+
+#ifdef TURING_MMA_AVAILABLE
+    asm("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;"
+        : "=r"(ret) : "r"(x));
+#else
+    GGML_UNUSED(x);
+    NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+    return ret;
+}
+
+#else
+
+static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) {
+    // Imagine transposing row-major matrix to column-major matrix.
+    const int src_i_low  = 2 * (threadIdx.x % 4);
+    const int src_i_high = src_i_low + 1;
+    const int src_j      = threadIdx.x / 4;
+
+    const int src_laneid_low  = src_i_low  * 4 + src_j / 2;
+    const int src_laneid_high = src_i_high * 4 + src_j / 2;
+
+    const int shift_low  = ((src_j + 0) % 2) * 16;
+    const int shift_high = ((src_j + 1) % 2) * 16;
+
+    const int ret_low  = (__shfl_sync(0xFFFFFFFF, x, src_laneid_low,  WARP_SIZE) >> shift_low)  & 0x0000FFFF;
+    const int ret_high = (__shfl_sync(0xFFFFFFFF, x, src_laneid_high, WARP_SIZE) << shift_high) & 0xFFFF0000;
+
+    return ret_low | ret_high;
+}
+
+#endif // CUDART_VERSION >= 11080
+
+static __device__ __forceinline__ half2 ggml_cuda_movmatrix(const half2 x) {
+    half2 ret;
+    *((int *) &ret) = ggml_cuda_movmatrix(*((const int *) &x));
+    return ret;
+}
+
+namespace ggml_cuda_mma {
+
+    // Some architectures like Volta or CDNA3 perform multiple matrix multiplications per warp in parallel,
+    //     effectively the warp is being split into subgroups of threads that each perform a single mma instruction.
+    // In those cases the data can be split in different ways across the warp.
+    enum data_layout {
+        // By default the data uses the I direction as its major dimension and the J direction as its minor dimension.
+        // For the A/C matrices this means I major == row major, J major == column major.
+        // For the B matrix this means I major == column major, J major == row major.
+        // MIRRORED == Each data value is held exactly once per thread subgroup.
+        DATA_LAYOUT_I_MAJOR           =  0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA.
+        DATA_LAYOUT_J_MAJOR           = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3.
+        DATA_LAYOUT_I_MAJOR_MIRRORED  = 20, // Volta, matrix A&B for RDNA3.
+        DATA_LAYOUT_J_MAJOR_MIRRORED  = 30,
+        DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (RDNA4/CDNA), convert to float to unscramble.
+    };
+    // Implemented mma combinations are:
+    //   - (I_MAJOR, I_MAJOR)          -> I_MAJOR
+    //   - (I_MAJOR, I_MAJOR_MIRRORED) -> I_MAJOR
+    //   - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR
+
+    static constexpr __device__ data_layout get_input_data_layout() {
+#if defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE)
+        return DATA_LAYOUT_I_MAJOR_MIRRORED;
+#else
+        return DATA_LAYOUT_I_MAJOR;
+#endif // defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE)
+    }
+
+    template <int I_, int J_, typename T, data_layout ds_=DATA_LAYOUT_I_MAJOR>
+    struct tile {};
+
+    template <int I_, int J_, typename T>
+    struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR;
+
+#if defined(AMD_MFMA_AVAILABLE)
+        static constexpr int ne = I * J / 64;
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J ==  8) return true;
+            if (I == 32 && J ==  4) return true;
+            if (I == 16 && J == 16) return true;
+            if (I == 32 && J == 32) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 16 && J == 4) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 16 && J == 8) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 32 && J == 4) {
+                return threadIdx.x % 32;
+            } else if constexpr (I == 16 && J == 16) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 32 && J == 32) {
+                return threadIdx.x % 32;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 16 && J == 4) {
+                return threadIdx.x / 16;
+            } else if constexpr (I == 16 && J == 8) {
+                return 2 * (threadIdx.x / 16) + l;
+            } else if constexpr (I == 32 && J == 4) {
+                return 2 * (threadIdx.x / 32) + l;
+            } else if constexpr (I == 16 && J == 16) {
+                return 4 * (threadIdx.x / 16) + l;
+            } else if constexpr (I == 32 && J == 32) {
+                return 4 * (threadIdx.x / 32) + 8 * (l / 4) + (l % 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#elif defined(VOLTA_MMA_AVAILABLE)
+        static constexpr int ne = I * J / 32;
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I == 32 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 32 && J == 8) {
+#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM
+                return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (l & 2) + (threadIdx.x % 2);
+#else
+                return (l & 2) + (threadIdx.x & ~2);
+#endif // GGML_CUDA_MMA_NO_VOLTA_PERM
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 32 && J == 8) {
+                return (threadIdx.x & 2) + (l & (4 + 1));
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#elif defined(AMD_WMMA_AVAILABLE)
+        static constexpr int ne = I * J / 32;
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J == 16) return true;
+            if (I == 16 && J == 8) return true;
+            if (I == 16 && J == 4) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (supported()) {
+                return threadIdx.x % 16;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 16 && J == 16) {
+#if defined(RDNA3)
+                if constexpr (std::is_same_v<T, float> || std::is_same_v<T, int>) {
+                    // matrix C
+                    return 2 * l + (threadIdx.x / 16);
+                } else {
+                    // matrix A&B
+                    return l;
+                }
+#else
+                // matrix C is the transposed matrix A&B on RDNA4
+                return ne * (threadIdx.x / 16) + l;
+#endif // defined(RDNA3)
+            } else if constexpr (I == 16 && J == 8) {
+                // mmq input for RDNA4
+                return ne * (threadIdx.x / 16) + l;
+            } else if constexpr (I == 16 && J == 4) {
+                return ne * (threadIdx.x / 16) + l;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#else
+        static constexpr int ne = I * J / 32;
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I ==  8 && J ==  4) return true;
+            if (I ==  8 && J ==  8) return true;
+            if (I == 16 && J ==  8) return true;
+            if (I == 16 && J == 16) return true;
+            if (I == 32 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 8 && J == 4) {
+                return threadIdx.x / 4;
+            } else if constexpr (I == 8 && J == 8) {
+                return threadIdx.x / 4;
+            } else if constexpr (I == 16 && J == 8) {
+                return ((l / 2) * 8) + (threadIdx.x / 4);
+            } else if constexpr (I == 16 && J == 16) {
+                return (((l / 2) % 2) * 8) + (threadIdx.x / 4);
+            } else if constexpr (I == 32 && J == 8) {
+                return tile<16, 8, T>::get_i(l); // Memory layout simply repeated with same pattern in i direction.
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 8 && J == 4) {
+                return threadIdx.x % 4;
+            } else if constexpr (I == 8 && J == 8) {
+                return (l * 4) + (threadIdx.x % 4);
+            } else if constexpr (I == 16 && J == 8) {
+                return ((threadIdx.x % 4) * 2) + (l % 2);
+            } else if constexpr (I == 16 && J == 16) {
+                return ((l / 4) * 8) + ((threadIdx.x % 4) * 2) + (l % 2);
+            } else if constexpr (I == 32 && J == 8) {
+                return tile<16, 8, T>::get_j(l); // Memory layout simply repeated with same pattern in i direction.
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#endif // defined(GGML_USE_HIP)
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR;
+
+#if defined(VOLTA_MMA_AVAILABLE)
+        static constexpr int ne = I * J / WARP_SIZE;
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I == 32 && J ==  4) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 32 && J == 4) {
+#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM
+                return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (threadIdx.x % 4);
+#else
+                return threadIdx.x;
+#endif // GGML_CUDA_MMA_NO_VOLTA_PERM
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 32 && J == 4) {
+                return l;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#elif defined(AMD_WMMA_AVAILABLE)
+        static constexpr int ne = I * J / 32;
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J ==  8) return true;
+            if (I == 16 && J == 16) return true;
+            if (I == 32 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 16 && J == 8) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 16 && J == 16) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 32 && J == 8) {
+                return (threadIdx.x % 16) * 2 + l / (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 16 && J == 8) {
+                return (threadIdx.x / 16) * ne + l;
+            } else if constexpr (I == 16 && J == 16) {
+#ifdef RDNA3
+                return l*2 + (threadIdx.x / 16);
+#else
+                return (threadIdx.x / 16) * ne + l;
+#endif // RDNA3
+            } else if constexpr (I == 32 && J == 8) {
+                return (threadIdx.x / 16) * (ne/2) + l % (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#elif defined(AMD_MFMA_AVAILABLE)
+        static constexpr int ne = I * J / 64;
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J ==  8) return true;
+            if (I == 16 && J == 16) return true;
+            if (I == 32 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 16 && J == 8) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 16 && J == 16) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 32 && J == 8) {
+                return (threadIdx.x % 16) * 2 + l / (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 16 && J == 8) {
+                return (threadIdx.x / 16) * ne + l;
+            } else if constexpr (I == 16 && J == 16) {
+                return (threadIdx.x / 16) * ne + l;
+            } else if constexpr (I == 32 && J == 8) {
+                return (threadIdx.x / 16) * (ne/2) + l % (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#else
+        static constexpr int ne = I * J / WARP_SIZE;
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I ==  8 && J ==  4) return true;
+            if (I ==  8 && J ==  8) return true;
+            if (I == 16 && J ==  8) return true;
+            if (I == 16 && J == 16) return true;
+            if (I == 32 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 8 && J == 8) {
+                return threadIdx.x / 4;
+            } else if constexpr (I == 16 && J == 4) {
+                return (l * 8) + (threadIdx.x / 4);
+            } else if constexpr (I == 16 && J == 8) {
+                return ((l % 2) * 8) + (threadIdx.x / 4);
+            } else if constexpr (I == 32 && J == 8) {
+                return ((l / 4) * 16) + ((l % 2) * 8) + (threadIdx.x / 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 8 && J == 8) {
+                return (l * 4) + (threadIdx.x % 4);
+            } else if constexpr (I == 16 && J == 4) {
+                return threadIdx.x % 4;
+            } else if constexpr (I == 16 && J == 8) {
+                return ((l / 2) * 4) + (threadIdx.x % 4);
+            } else if constexpr (I == 32 && J == 8) {
+                return ((l & 2) * 2) + (threadIdx.x % 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#endif // defined(VOLTA_MMA_AVAILABLE)
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, nv_bfloat162, DATA_LAYOUT_I_MAJOR> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR;
+
+#if defined(AMD_WMMA_AVAILABLE)
+        static constexpr int ne = tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::ne;
+        nv_bfloat162 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::supported();
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l);
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_j(l);
+        }
+#elif defined(AMD_MFMA_AVAILABLE)
+        static constexpr int ne = tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::ne;
+        nv_bfloat162 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::supported();
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l);
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_j(l);
+        }
+#else
+        static constexpr int ne = I * J / WARP_SIZE;
+        nv_bfloat162 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I ==  8 && J ==  8) return true;
+            if (I == 16 && J ==  4) return true;
+            if (I == 16 && J ==  8) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 8 && J == 8) {
+                return threadIdx.x / 4;
+            } else if constexpr (I == 16 && J == 4) {
+                return (l * 8) + (threadIdx.x / 4);
+            } else if constexpr (I == 16 && J == 8) {
+                return ((l % 2) * 8) + (threadIdx.x / 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 8 && J == 8) {
+                return (l * 4) + (threadIdx.x % 4);
+            } else if constexpr (I == 16 && J == 4) {
+                return threadIdx.x % 4;
+            } else if constexpr (I == 16 && J == 8) {
+                return ((l / 2) * 4) + (threadIdx.x % 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#endif  // defined(AMD_WMMA_AVAILABLE)
+    };
+
+    template <int I_, int J_, typename T>
+    struct tile<I_, J_, T, DATA_LAYOUT_J_MAJOR> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR;
+
+        static constexpr int ne = tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::ne;
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::supported();
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::get_j(l);
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::get_i(l);
+        }
+    };
+
+    template <int I_, int J_, typename T>
+    struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR_MIRRORED> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
+
+        // RDNA3
+        static constexpr int         ne = I * J / 32 * 2;
+
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J == 16) return true;
+            if (I == 16 && J == 8)  return true;
+            if (I == 16 && J == 4)  return true;
+            if (I == 32 && J == 8)  return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 16) {
+                return threadIdx.x % 16;
+            } else if constexpr (I == 32) {
+                return (threadIdx.x % 16) * 2 + l / (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 16) {
+                return l;
+            } else if constexpr (I == 32) {
+                return l % (ne/2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
+#if defined(RDNA3)
+        static constexpr int         ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne;
+
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported();
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l);
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l);
+        }
+#else // Volta
+        static constexpr int         ne = I * J / (WARP_SIZE/4);
+
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I ==  8 && J ==  4) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int /*l*/) {
+            if constexpr (I == 8 && J == 4) {
+                return ((threadIdx.x / 16) * 4) + (threadIdx.x % 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 8 && J == 4) {
+                return l;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+#endif // defined(RDNA3)
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, nv_bfloat162, DATA_LAYOUT_I_MAJOR_MIRRORED> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
+        static constexpr int         ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne;
+
+        nv_bfloat162 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported();
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l);
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l);
+        }
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR_MIRRORED;
+        static constexpr int         ne = I * J / (WARP_SIZE/4);
+
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I ==  8 && J ==  4) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (I == 8 && J == 4) {
+                return ((l / 2) * 4) + (threadIdx.x % 4);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (I == 8 && J == 4) {
+                return ((threadIdx.x / 16) * 2) + (l % 2);
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+    };
+
+    template <int I_, int J_>
+    struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_SCRAMBLED;
+
+        static constexpr int ne = I * J / ggml_cuda_get_physical_warp_size();
+        half2 x[ne] = {{0.0f, 0.0f}};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J == 16) return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l);
+        }
+    };
+
+    static __device__ __forceinline__ tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> unscramble(const tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & t) {
+#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+        tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> ret;
+#pragma unroll
+        for (int l0 = 0; l0 < t.ne/2; ++l0) {
+            ret.x[2*l0 + 0] =  __lows2half2(t.x[l0], t.x[l0 + t.ne/2]);
+            ret.x[2*l0 + 1] = __highs2half2(t.x[l0], t.x[l0 + t.ne/2]);
+        }
+        return ret;
+#else
+        NO_DEVICE_CODE;
+        GGML_UNUSED(t);
+#endif // defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+    }
+
+#if defined(TURING_MMA_AVAILABLE)
+    template <int I, int J>
+    static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) {
+        tile<I, J/2, half2> ret;
+#pragma unroll
+        for (int l0 = 0; l0 < tile_float.ne; l0 += 2) {
+            ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]);
+        }
+        return ret;
+    }
+
+    static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) {
+        tile<8, 8, half2> ret;
+        ret.x[0] = ggml_cuda_movmatrix(t.x[0]);
+        ret.x[1] = ggml_cuda_movmatrix(t.x[1]);
+
+        return ret;
+    }
+#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
+    static __device__ __forceinline__ tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> get_half2(
+            const tile<16, 16, float, DATA_LAYOUT_I_MAJOR> & tile_float) {
+        tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> ret;
+#pragma unroll
+        for (int l = 0; l < tile_float.ne; ++l) {
+            float tmp[2];
+            int i = threadIdx.x / 16;
+            tmp[i] = tile_float.x[l];
+            i ^= 1;
+            tmp[i] = __shfl_xor_sync(0xFFFFFFFF, tile_float.x[l], 16, WARP_SIZE);
+            ret.x[l] = make_half2(tmp[0], tmp[1]);
+        }
+        return ret;
+    }
+#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
+    template <int I, int J>
+    static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) {
+        tile<I, J/2, half2> ret;
+#pragma unroll
+        for (int l0 = 0; l0 < tile_float.ne; l0 += 2) {
+            ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]);
+        }
+        return ret;
+    }
+
+    static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) {
+        NO_DEVICE_CODE;
+        return tile<8, 8, half2>{};
+    }
+#else // Volta
+    template <int I, int J>
+    static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) {
+        tile<I, J/2, half2> ret;
+#pragma unroll
+        for (int l0 = 0; l0 < tile_float.ne; l0 += 4) {
+            ret.x[l0/2 + 0] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]);
+            ret.x[l0/2 + 1] = make_half2(tile_float.x[l0 + 2], tile_float.x[l0 + 3]);
+
+            // On Volta FP16 and FP32 tiles have a different memory layout,
+            //     for the conversion threads with an offset of 2 need to exchange half their values:
+            ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = __shfl_xor_sync(
+                0xFFFFFFFF, ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2, WARP_SIZE);
+        }
+        return ret;
+    }
+#endif // defined(TURING_MMA_AVAILABLE)
+
+    template <int I, int J, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_generic(tile<I, J, T, dl> & t, const T * __restrict__ xs0, const int stride) {
+#pragma unroll
+        for (int l = 0; l < t.ne; ++l) {
+            t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)];
+        }
+    }
+
+    template <int stride, typename T>
+    static __device__ __forceinline__ uint32_t swizzle(const uint32_t offset, const uint32_t i) {
+        static_assert(sizeof(T) <= 4, "unsupported type size");
+        constexpr int stride_bytes = stride*sizeof(T);
+        static_assert(stride_bytes % 16 == 0, "bad stride");
+        constexpr uint32_t shift = sizeof(T) == 1 ? 4 : (sizeof(T) == 2 ? 3 : 2);
+        if (stride_bytes % 32 != 0) {
+            return offset; // Equivalent to padding with 16 bytes.
+        }
+        if (stride_bytes % 64 != 0) {
+            return offset ^ (((i / 4) % 2) << shift);
+        }
+        if (stride_bytes % 128 != 0) {
+            return offset ^ (((i / 2) % 4) << shift);
+        }
+        return offset ^ ((i % 8) << shift);
+    }
+
+    template <int stride, typename T>
+    static __device__ __forceinline__ T * swizzle(T * ptr, const uint32_t offset, const uint32_t i) {
+        return ptr + swizzle<stride, T>(offset, i);
+    }
+
+    template <typename T>
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<8, 8, T> & t, const T * __restrict__ xs0, const int stride) {
+#ifdef TURING_MMA_AVAILABLE
+        int * xi = (int *) t.x;
+        const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + ((threadIdx.x / t.I) * (t.J / 2)) % t.J;
+        asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];"
+            : "=r"(xi[0]), "=r"(xi[1])
+            : "l"(xs));
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    template <typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<16, 4, T, dl> & t, const T * __restrict__ xs0, const int stride) {
+#ifdef TURING_MMA_AVAILABLE
+        int * xi = (int *) t.x;
+        const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride;
+        asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];"
+            : "=r"(xi[0]), "=r"(xi[1])
+            : "l"(xs));
+#elif defined(AMD_WMMA_AVAILABLE)
+#ifdef RDNA3
+        static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        static_assert(sizeof(t.x) == 16, "bad ne");
+        ggml_cuda_memcpy_1<8>(t.x + 0, xs0 + t.get_i(0)*stride + 0);
+        ggml_cuda_memcpy_1<8>(t.x + 2, xs0 + t.get_i(0)*stride + 2);
+#else
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        static_assert(sizeof(t.x) == 8, "bad ne");
+        ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0));
+#endif // RDNA3
+#elif defined(AMD_MFMA_AVAILABLE)
+        static_assert(sizeof(t.x) == 4, "bad ne");
+        ggml_cuda_memcpy_1<4>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0));
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    template <typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<16, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) {
+#if defined(TURING_MMA_AVAILABLE)
+        int * xi = (int * ) t.x;
+        const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2);
+        asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
+            : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
+            : "l"(xs));
+#elif defined(VOLTA_MMA_AVAILABLE)
+        ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 0, xs0 + t.get_i(0)*stride + 0);
+        ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 4, xs0 + t.get_i(4)*stride + 4);
+#elif defined(AMD_WMMA_AVAILABLE)
+#ifdef RDNA3
+        static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        static_assert(sizeof(t.x) == 32, "bad ne");
+        ggml_cuda_memcpy_1<16>(t.x + 0, xs0 + t.get_i(0)*stride + 0);
+        ggml_cuda_memcpy_1<16>(t.x + 4, xs0 + t.get_i(0)*stride + 4);
+#else
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        static_assert(sizeof(t.x) == 16, "bad ne");
+        ggml_cuda_memcpy_1<16>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0));
+#endif // RDNA3
+#elif defined(AMD_MFMA_AVAILABLE)
+        static_assert(sizeof(t.x) == 8, "bad ne");
+        ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0));
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) {
+        ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride);
+    }
+
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) {
+#ifdef VOLTA_MMA_AVAILABLE
+#pragma unroll
+        for (int l0 = 0; l0 < t.ne; l0 += 2) {
+            ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, xs0 + t.get_i(l0)*stride + t.get_j(l0));
+        }
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // VOLTA_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void load_ldmatrix(
+            tile<32, 4, half2> & t, const half2 * __restrict__ xs0, const int stride) {
+#if defined(VOLTA_MMA_AVAILABLE)
+        ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride);
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // defined(VOLTA_MMA_AVAILABLE)
+    }
+
+    template <int I, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix_trans(
+            tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) {
+#ifdef TURING_MMA_AVAILABLE
+        static_assert(I == 16, "bad tile width");
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        int * xi = (int *) t.x;
+        const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2);
+        asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
+            : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
+            : "l"(xs));
+#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        if constexpr (I == 32) {
+#pragma unroll
+            for (int l0 = 0; l0 < t.ne/2; ++l0) {
+                const half2 tmp0 = xs0[(2*t.get_j(l0) + 0)*stride + t.get_i(l0)/2];
+                const half2 tmp1 = xs0[(2*t.get_j(l0) + 1)*stride + t.get_i(l0)/2];
+
+                t.x[l0]          =  __lows2half2(tmp0, tmp1);
+                t.x[l0 + t.ne/2] = __highs2half2(tmp0, tmp1);
+            }
+        } else {
+            half * xh = (half *) t.x;
+#pragma unroll
+            for (int l = 0; l < t.ne; ++l) {
+                xh[2*l + 0] = ((const half *) xs0)[(2*t.get_j(l) + 0)*(2*stride) + t.get_i(l)];
+                xh[2*l + 1] = ((const half *) xs0)[(2*t.get_j(l) + 1)*(2*stride) + t.get_i(l)];
+            }
+        }
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    template <int stride, int I, int J, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix_swizzled(
+            tile<I, J, T, dl> & t, const T * __restrict__ xs0, const int offset) {
+#if defined(TURING_MMA_AVAILABLE)
+        static_assert(I == 16, "bad tile width");
+        static_assert(J ==  8, "bad tile height");
+        const int i = threadIdx.x % t.I;
+        const int j = (threadIdx.x / t.I) * (t.J / 2);
+        int offset_ij = offset + i * stride + j;
+        offset_ij = swizzle<stride, T>(offset_ij, i);
+        int * xi = (int *) t.x;
+        asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
+            : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
+            : "l"(xs0 + offset_ij));
+#elif defined(VOLTA_MMA_AVAILABLE)
+#pragma unroll
+        for (int o = 0; o < t.ne; o += 4) {
+            const int offset_ij = offset + t.get_i(o) * stride + o;
+            ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + o, swizzle<stride>(xs0, offset_ij, t.get_i(o)));
+        }
+#elif defined(AMD_WMMA_AVAILABLE)
+#ifdef RDNA3
+        static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        static_assert(sizeof(t.x) == 32, "bad ne");
+        static_assert(I == 16, "bad tile width");
+        static_assert(J ==  8, "bad tile height");
+#pragma unroll
+        for (int o = 0; o < 8; o += 4) {
+            const int offset_ij = offset + t.get_i(0) * stride + o;
+            ggml_cuda_memcpy_1<16>(t.x + o, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+        }
+#else
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        static_assert(sizeof(t.x) == 16, "bad ne");
+        const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0);
+        ggml_cuda_memcpy_1<16>(t.x, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+#endif // RDNA3
+#elif defined(AMD_MFMA_AVAILABLE)
+        static_assert(sizeof(t.x) == 8, "bad ne");
+        const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0);
+        ggml_cuda_memcpy_1<8>(t.x, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+#else
+        GGML_UNUSED_VARS(t, xs0, offset);
+        NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+    }
+
+    template <int stride>
+    static __device__ __forceinline__ void load_ldmatrix_swizzled(
+            tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int offset) {
+#ifdef VOLTA_MMA_AVAILABLE
+#pragma unroll
+        for (int l0 = 0; l0 < t.ne; l0 += 2) {
+            const int offset_ij = offset + t.get_i(l0)*stride + t.get_j(l0);
+            ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, swizzle<stride>(xs0, offset_ij, t.get_i(l0)));
+        }
+#else
+        GGML_UNUSED_VARS(t, xs0, offset);
+        NO_DEVICE_CODE;
+#endif // VOLTA_MMA_AVAILABLE
+    }
+
+    template <int stride, int I, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix_trans_swizzled(
+            tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int offset) {
+#if defined(TURING_MMA_AVAILABLE)
+        static_assert(I == 16, "bad tile width");
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        const int i = threadIdx.x % t.I;
+        const int j = (threadIdx.x / t.I) * (t.J / 2);
+        int offset_ij = offset + i * stride + j;
+        offset_ij = swizzle<stride, T>(offset_ij, i);
+        int * xi = (int *) t.x;
+        asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
+            : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
+            : "l"(xs0 + offset_ij));
+#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        if constexpr (I == 32) {
+#pragma unroll
+            for (int l0 = 0; l0 < t.ne/2; ++l0) {
+                half2 tmp[2];
+#pragma unroll
+                for (int o = 0; o < 2; ++o) {
+                    const int j = 2*t.get_j(l0) + o;
+                    int offset_ij = offset + j*stride + t.get_i(l0)/2;
+                    offset_ij = swizzle<stride, T>(offset_ij, j);
+                    tmp[o] = xs0[offset_ij];
+                }
+
+                t.x[l0]          =  __lows2half2(tmp[0], tmp[1]);
+                t.x[l0 + t.ne/2] = __highs2half2(tmp[0], tmp[1]);
+            }
+        } else {
+            half * xh = (half *) t.x;
+#pragma unroll
+            for (int l = 0; l < t.ne; ++l) {
+#pragma unroll
+                for (int o = 0; o < 2; ++o) {
+                    const int j = 2*t.get_j(l) + o;
+                    xh[2*l + o] = ((const half *) xs0)[swizzle<2*stride, half>(2*offset + j*(2*stride) + t.get_i(l), j)];
+                }
+            }
+        }
+#else
+        GGML_UNUSED_VARS(t, xs0, offset);
+        NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, int> & D, const tile<16, 4, int> & A, const tile<8, 4, int> & B) {
+#ifdef TURING_MMA_AVAILABLE
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3])
+            : "r"(A.x[0]), "r"(A.x[1]), "r"(B.x[0]));
+#else
+        // On Turing m16n8k16 mma is not available, use 2x m8n8k16 mma instead:
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[0]), "+r"(D.x[1])
+            : "r"(A.x[0]), "r"(B.x[0]));
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[2]), "+r"(D.x[3])
+            : "r"(A.x[1]), "r"(B.x[0]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, int> & D, const tile<16, 8, int> & A, const tile<8, 8, int> & B) {
+#ifdef TURING_MMA_AVAILABLE
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3])
+            : "r"(A.x[0]), "r"(A.x[1]), "r"(A.x[2]), "r"(A.x[3]), "r"(B.x[0]), "r"(B.x[1]));
+#else
+        // On Turing m16n8k32 mma is not available, use 4x m8n8k16 mma instead:
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[0]), "+r"(D.x[1])
+            : "r"(A.x[0]), "r"(B.x[0]));
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[2]), "+r"(D.x[3])
+            : "r"(A.x[1]), "r"(B.x[0]));
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[0]), "+r"(D.x[1])
+            : "r"(A.x[2]), "r"(B.x[1]));
+        asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
+            : "+r"(D.x[2]), "+r"(D.x[3])
+            : "r"(A.x[3]), "r"(B.x[1]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 4, half2> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) {
+#ifdef TURING_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]));
+#else
+        // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead:
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, half2> & D, const tile<16, 8, half2> & A, const tile<16, 8, half2> & B) {
+#ifdef TURING_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2]));
+        asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};"
+            : "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3]));
+#else
+        // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead:
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};"
+            : "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#elif defined(AMD_WMMA_AVAILABLE)
+#if defined(RDNA4)
+        using halfx8_t = __attribute__((ext_vector_type(8))) _Float16;
+        halfx8_t& acc_frag = reinterpret_cast<halfx8_t&>(D.x[0]);
+        const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]);
+        const halfx8_t& b_frag = reinterpret_cast<const halfx8_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(RDNA4)
+#elif defined(AMD_MFMA_AVAILABLE)
+        // MFMA: FP16 input, FP32 accumulate, convert back to half2.
+        using halfx4_t = __attribute__((ext_vector_type(4))) _Float16;
+        using floatx4_t = __attribute__((ext_vector_type(4))) float;
+
+        // Convert existing half2 accumulator to float for MFMA:
+        floatx4_t acc_f32;
+        {
+            const halfx4_t acc_h = reinterpret_cast<const halfx4_t&>(D.x[0]);
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                acc_f32[i] = (float)acc_h[i];
+            }
+        }
+
+        const halfx4_t& a_frag = reinterpret_cast<const halfx4_t&>(A.x[0]);
+        const halfx4_t& b_frag = reinterpret_cast<const halfx4_t&>(B.x[0]);
+        acc_f32 = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_f32, 0, 0, 0);
+
+        // Convert back to half2:
+        {
+            halfx4_t result_h;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                result_h[i] = (_Float16)acc_f32[i];
+            }
+            reinterpret_cast<halfx4_t&>(D.x[0]) = result_h;
+        }
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR> & A,
+            const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> & B) {
+#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+        tile<16, 8, half2>       * D16 = (tile<16, 8, half2>       *) &D;
+        const tile<16, 8, half2> * A16 = (const tile<16, 8, half2> *) &A;
+        mma(D16[0], A16[0], B);
+        mma(D16[1], A16[1], B);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)
+    }
+
+    template <data_layout dl_ab, data_layout dl_d>
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<8, 8, float, dl_ab> & B) {
+#ifdef AMPERE_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]));
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // AMPERE_MMA_AVAILABLE
+    }
+
+    template <data_layout dl_ab, data_layout dl_d>
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<16, 8, float, dl_ab> & B) {
+#ifdef AMD_MFMA_AVAILABLE
+        using floatx4_t = __attribute__((ext_vector_type(4))) float;
+        floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]);
+#if defined(CDNA3)
+        using floatx2_t = __attribute__((ext_vector_type(2))) float;
+        const floatx2_t& a_frag = reinterpret_cast<const floatx2_t&>(A.x[0]);
+        const floatx2_t& b_frag = reinterpret_cast<const floatx2_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_mfma_f32_16x16x8_xf32(a_frag, b_frag, acc_frag, 0, 0, 0);
+#elif defined(CDNA4) || defined(CDNA2) || defined(CDNA1)
+        // CDNA4 (gfx950) does not support xf32 MFMA, use f32 path like CDNA2/CDNA1
+#pragma unroll
+        for (int i = 0; i < 2; ++i) {
+            acc_frag = __builtin_amdgcn_mfma_f32_16x16x4f32(A.x[i], B.x[i], acc_frag, 0, 0, 0);
+        }
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(CDNA3)
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE
+    }
+
+    template <ggml_type type>
+    static __device__ __forceinline__ void mma_block_scaled_fp4(tile<16, 8, float> &     D,
+                                                                const tile<16, 8, int> & A,
+                                                                const tile<8, 8, int> &  B,
+                                                                uint32_t                 a_scale,
+                                                                uint32_t                 b_scale) {
+#ifdef BLACKWELL_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        float *     Dxi = (float *) D.x;
+
+        if constexpr (type == GGML_TYPE_MXFP4) {
+            asm volatile(
+                "mma.sync.aligned.kind::mxf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 "
+                "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, "
+                "%10, {0, 0}, %11, {0, 0};"
+                : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3])
+                : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale));
+        } else {
+            asm volatile(
+                "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue4m3 "
+                "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, "
+                "%10, {0, 0}, %11, {0, 0};"
+                : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3])
+                : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale));
+        }
+#else
+        GGML_UNUSED_VARS(D, A, B, a_scale, b_scale);
+#endif // BLACKWELL_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, float> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) {
+#ifdef TURING_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]));
+#else
+        // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead:
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 8, float> & D, const tile<16, 8, nv_bfloat162> & A, const tile<8, 8, nv_bfloat162> & B) {
+#ifdef AMPERE_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]));
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // AMPERE_MMA_AVAILABLE
+    }
+
+    template <data_layout dl_ab, data_layout dl_d>
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, float, dl_d> & D, const tile<16, 8, half2, dl_ab> & A, const tile<16, 8, half2, dl_ab> & B) {
+#ifdef TURING_MMA_AVAILABLE
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2]));
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3]));
+#else
+        // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead:
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1]));
+        asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3]));
+#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#elif defined(AMD_WMMA_AVAILABLE)
+#if defined(RDNA4)
+        using halfx8_t = __attribute__((ext_vector_type(8))) _Float16;
+        using floatx8_t = __attribute__((ext_vector_type(8))) float;
+        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
+        const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]);
+        const halfx8_t& b_frag = reinterpret_cast<const halfx8_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag);
+#elif defined(RDNA3)
+        using halfx16_t = __attribute__((ext_vector_type(16))) _Float16;
+        using floatx8_t = __attribute__((ext_vector_type(8))) float;
+        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
+        const halfx16_t& a_frag = reinterpret_cast<const halfx16_t&>(A.x[0]);
+        const halfx16_t& b_frag = reinterpret_cast<const halfx16_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_frag, b_frag, acc_frag);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // RDNA4
+#elif defined(AMD_MFMA_AVAILABLE)
+        using halfx4_t = __attribute__((ext_vector_type(4))) _Float16;
+        using floatx4_t = __attribute__((ext_vector_type(4))) float;
+        floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]);
+        const halfx4_t& a_frag = reinterpret_cast<const halfx4_t&>(A.x[0]);
+        const halfx4_t& b_frag = reinterpret_cast<const halfx4_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_frag, 0, 0, 0);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    template <data_layout dl_ab, data_layout dl_d>
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) {
+#if defined(AMD_WMMA_AVAILABLE)
+#if defined(RDNA4)
+        using bf16x8_t = __attribute__((ext_vector_type(8))) __bf16;
+        using floatx8_t = __attribute__((ext_vector_type(8))) float;
+        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
+        const bf16x8_t& a_frag = reinterpret_cast<const bf16x8_t&>(A.x[0]);
+        const bf16x8_t& b_frag = reinterpret_cast<const bf16x8_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag);
+#elif defined(RDNA3)
+        using bf16x16_t = __attribute__((ext_vector_type(16))) __bf16;
+        using floatx8_t = __attribute__((ext_vector_type(8))) float;
+        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
+        const bf16x16_t& a_frag = reinterpret_cast<const bf16x16_t&>(A.x[0]);
+        const bf16x16_t& b_frag = reinterpret_cast<const bf16x16_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a_frag, b_frag, acc_frag);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(RDNA4)
+#elif defined(AMD_MFMA_AVAILABLE)
+        using floatx4_t = __attribute__((ext_vector_type(4))) float;
+        floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]);
+#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2)
+        using bf16x4_t = __attribute__((ext_vector_type(4))) __bf16;
+        const bf16x4_t& a_frag = reinterpret_cast<const bf16x4_t&>(A.x[0]);
+        const bf16x4_t& b_frag = reinterpret_cast<const bf16x4_t&>(B.x[0]);
+        acc_frag = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a_frag, b_frag, acc_frag, 0, 0, 0);
+#elif defined(CDNA1)
+#pragma unroll
+        for (int i = 0; i < 2; ++i) {
+            using bf16x2_t = __attribute__((ext_vector_type(2))) __bf16;
+            const bf16x2_t& a_frag = reinterpret_cast<const bf16x2_t&>(A.x[i]);
+            const bf16x2_t& b_frag = reinterpret_cast<const bf16x2_t&>(B.x[i]);
+            acc_frag = __builtin_amdgcn_mfma_f32_16x16x8bf16(a_frag, b_frag, acc_frag, 0, 0, 0);
+        }
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(CDNA3) || defined(CDNA2)
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(AMD_WMMA_AVAILABLE)
+    }
+
+    template <data_layout dl_d, data_layout dl_ab>
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, int, dl_d> & D, const tile<16, 8, int, dl_ab> & A, const tile<16, 8, int, dl_ab> & B) {
+#if defined(AMD_MFMA_AVAILABLE)
+        using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
+        int32x4_t * acc = (int32x4_t *) D.x;
+#if defined(CDNA4) || defined(CDNA3)
+        acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0);
+#elif defined(CDNA2) || defined(CDNA1)
+        acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0);
+        acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[1], B.x[1], acc[0], 0, 0, 0);
+#endif // defined(CDNA4) || defined(CDNA3)
+#elif defined(AMD_WMMA_AVAILABLE)
+        using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int;
+        int32x8_t * acc = (int32x8_t *) D.x;
+#if defined(RDNA4)
+        using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
+        int32x2_t * a_vec = (int32x2_t *) A.x;
+        int32x2_t * b_vec = (int32x2_t *) B.x;
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], true);
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[1], true, b_vec[1], acc[0], true);
+#elif defined(RDNA3)
+        using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
+        int32x4_t * a_vec = (int32x4_t *) A.x;
+        int32x4_t * b_vec = (int32x4_t *) B.x;
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], true);
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[1], true, b_vec[1], acc[0], true);
+#endif // RDNA4
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<32, 32, int> & D, const tile<32, 4, int> & A, const tile<32, 4, int> & B) {
+#if defined(AMD_MFMA_AVAILABLE)
+        using int32x16_t = __attribute__((__vector_size__(16 * sizeof(int)))) int;
+        int32x16_t * acc = (int32x16_t *) D.x;
+#if defined(CDNA4) || defined(CDNA3)
+        acc[0] = __builtin_amdgcn_mfma_i32_32x32x16_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0);
+#elif defined(CDNA2) || defined(CDNA1)
+        acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[0], B.x[0], acc[0], 0, 0, 0);
+        acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[1], B.x[1], acc[0], 0, 0, 0);
+#endif // defined(CDNA4) || defined(CDNA3)
+
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE
+    }
+
+    template <typename T1, typename T2, int J, int K>
+    static __device__ __forceinline__ void mma(
+            tile<32, J, T1> & D, const tile<32, K, T2> & A, const tile<J, K, T2> & B) {
+        tile      <16, J, T1> * D16 = reinterpret_cast<      tile<16, J, T1> *>(&D);
+        const tile<16, K, T2> * A16 = reinterpret_cast<const tile<16, K, T2> *>(&A);
+        mma(D16[0], A16[0], B);
+        mma(D16[1], A16[1], B);
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<32, 8, float> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) {
+#if defined(VOLTA_MMA_AVAILABLE)
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+        asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 "
+            "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1]));
+        asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 "
+            "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3]));
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(VOLTA_MMA_AVAILABLE)
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<32, 4, half2> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & B) {
+#if defined(VOLTA_MMA_AVAILABLE)
+        const int * Axi = (const int *) A.x;
+        const int * Bxi = (const int *) B.x;
+        int       * Dxi = (int       *) D.x;
+        asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 "
+            "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1]));
+        asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 "
+            "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};"
+            : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3])
+            : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3]));
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // defined(VOLTA_MMA_AVAILABLE)
+    }
+
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> & D, const tile<32,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & A,
+            const tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) {
+#if defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
+        using halfx16_t = __attribute__((ext_vector_type(16))) _Float16;
+        halfx16_t       * xD = (halfx16_t       *) D.x;
+        const halfx16_t * xA = (const halfx16_t *) A.x;
+        const halfx16_t * xB = (const halfx16_t *) B.x;
+        xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[0], xB[0], xD[0], /*opsel =*/ 0);
+        xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[1], xB[0], xD[0], /*opsel =*/ 1);
+#else
+        GGML_UNUSED_VARS(D, A, B);
+        NO_DEVICE_CODE;
+#endif // TURING_MMA_AVAILABLE
+    }
+
+    template <data_layout dl_d, data_layout dl_ab>
+    static __device__ __forceinline__ void mma(
+            tile<16, 16, int, dl_d> & D, const tile<16, 4, int, dl_ab> & A, const tile<16, 4, int, dl_ab> & B) {
+#if defined(AMD_MFMA_AVAILABLE)
+        using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
+        int32x4_t * acc = (int32x4_t *) D.x;
+#if defined(CDNA4) || defined(CDNA3)
+        const int64_t xA = uint32_t(A.x[0]);
+        const int64_t xB = uint32_t(B.x[0]);
+        acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(xA, xB, acc[0], 0, 0, 0);
+#elif defined(CDNA2) || defined(CDNA1)
+        acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0);
+#endif // defined(CDNA4) || defined(CDNA3)
+#elif defined(AMD_WMMA_AVAILABLE)
+        using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int;
+        int32x8_t * acc = (int32x8_t *) D.x;
+#if defined(RDNA4)
+        using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
+        int32x2_t * a_vec = (int32x2_t *) A.x;
+        int32x2_t * b_vec = (int32x2_t *) B.x;
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], false);
+#elif defined(RDNA3)
+        using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
+        int32x4_t * a_vec = (int32x4_t *) A.x;
+        int32x4_t * b_vec = (int32x4_t *) B.x;
+        acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], false);
+#endif // RDNA4
+#else
+        GGML_UNUSED(D);
+        GGML_UNUSED(A);
+        GGML_UNUSED(B);
+        NO_DEVICE_CODE;
+#endif // AMD_WMMA_AVAILABLE
+    }
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu
new file mode 100644
index 0000000000..345b39eb5f
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu
@@ -0,0 +1,221 @@
+#include "common.cuh"
+#include "mmid.cuh"
+
+// To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each.
+struct mm_ids_helper_store {
+    uint32_t data;
+
+    __device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
+        data = (it & 0x003FFFFF) | (iex_used << 22);
+    }
+
+    __device__ uint32_t it() const {
+        return data & 0x003FFFFF;
+    }
+
+    __device__ uint32_t iex_used() const {
+        return data >> 22;
+    }
+};
+static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store");
+
+// the generic path passes 0, which needs no padding since it never groups lanes by token
+template <int n> struct mm_ids_pow2 { static constexpr int value = 2*mm_ids_pow2<(n + 1)/2>::value; };
+template <>      struct mm_ids_pow2<1> { static constexpr int value = 1; };
+template <>      struct mm_ids_pow2<0> { static constexpr int value = 1; };
+
+// Helper function for mul_mat_id, converts ids to a more convenient format.
+// ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert.
+// ids_dst describes the same mapping but for the dst tensor.
+// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
+template <int n_expert_used_template>
+__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
+    const int expert = blockIdx.x;
+
+    extern __shared__ char data_mm_ids_helper[];
+    mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
+
+    int nex_prev   = 0; // Number of columns for experts with a lower index.
+    int it_compact = 0; // Running index for the compact slice of this expert.
+
+    if constexpr (n_expert_used_template == 0) {
+        // Generic implementation:
+        for (int it = 0; it < n_tokens; ++it) {
+            int iex_used = -1; // The index at which the expert is used, if any.
+            for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
+                const int expert_used = ids[it*si1 + iex];
+                nex_prev += expert_used < expert;
+                if (expert_used == expert) {
+                    iex_used = iex;
+                }
+            }
+
+            if (iex_used != -1) {
+                store[it_compact] = mm_ids_helper_store(it, iex_used);
+            }
+
+            if (warp_reduce_any<warp_size>(iex_used != -1)) {
+                it_compact++;
+            }
+        }
+    } else {
+        // Implementation optimized for specific numbers of experts used:
+        // a warp holds a whole number of token slots, so the slot count is padded to a power of 2
+        constexpr int neu_padded = mm_ids_pow2<n_expert_used_template>::value;
+        static_assert(neu_padded <= warp_size && warp_size % neu_padded == 0, "bad n_expert_used");
+        for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+            const int it = it0 + threadIdx.x / neu_padded;
+
+            const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
+            const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
+                ids[it*si1 + iex] : INT_MAX;
+            const int iex_used = expert_used == expert ? iex : -1;
+            nex_prev += expert_used < expert;
+
+            // Whether the threads at this token position have used the expert:
+            const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
+
+            // Do a scan over threads at lower token positions in warp to get the correct index for writing data:
+            int it_compact_add_lower = 0;
+#pragma unroll
+            for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
+                const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
+                if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                    it_compact_add_lower += tmp;
+                }
+            }
+
+            if (iex_used != -1) {
+                store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
+            }
+
+            // The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
+            it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
+        }
+    }
+    nex_prev = warp_reduce_sum<warp_size>(nex_prev);
+    ggml_cuda_syncwarp();
+
+    for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
+        const mm_ids_helper_store store_it = store[itc];
+        const int it       = store_it.it();
+        const int iex_used = store_it.iex_used();
+        ids_dst[nex_prev + itc] = it*n_expert_used + iex_used;
+        // ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup
+        if (write_inverse) {
+            ids_src1[it*n_expert_used + iex_used] = nex_prev + itc;
+        } else {
+            ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
+        }
+    }
+
+    if (threadIdx.x != 0) {
+        return;
+    }
+
+    expert_bounds[expert] = nex_prev;
+
+    if (expert < static_cast<int>(gridDim.x) - 1) {
+        return;
+    }
+
+    expert_bounds[gridDim.x] = nex_prev + it_compact;
+}
+
+template <int n_expert_used_template>
+static void launch_mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+    GGML_ASSERT(n_tokens          < (1 << 22) && "too few bits in mm_ids_helper_store");
+    GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
+
+    const int id = ggml_cuda_get_device();
+    const int warp_size = ggml_cuda_info().devices[id].warp_size;
+    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
+    CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
+
+    const dim3 num_blocks(n_experts, 1, 1);
+    const dim3 block_size(warp_size, 1, 1);
+    const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
+    GGML_ASSERT(nbytes_shared <= smpbo);
+    mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
+        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
+}
+
+void ggml_cuda_launch_mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+    switch (n_expert_used) {
+        case  2:
+            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case  4:
+            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case  6:
+            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case  8:
+            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case 10:
+            launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case 16:
+            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        case 32:
+            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+        default:
+            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            break;
+    }
+}
+
+// Lists the (expert, column tile) pairs that hold columns for a flat MoE grid, in expert order.
+// The other ntiles grid columns get tile jt_none of expert 0, which starts past its columns, so mul_mat_q skips them.
+template <int nthreads>
+__launch_bounds__(nthreads, 1)
+static __global__ void mm_ids_tile_map(
+        const int32_t * __restrict__ expert_bounds, int2 * __restrict__ tile_map, const int n_experts, const int J, const int ntiles, const int jt_none) {
+    __shared__ int scan[nthreads];
+
+    int start = 0; // First tile of the current chunk of experts.
+    for (int e0 = 0; e0 < n_experts; e0 += nthreads) {
+        const int e = e0 + threadIdx.x;
+        const int n = e < n_experts ? (expert_bounds[e + 1] - expert_bounds[e] + J - 1) / J : 0;
+
+        scan[threadIdx.x] = n;
+        __syncthreads();
+#pragma unroll
+        for (int offset = 1; offset < nthreads; offset *= 2) {
+            const int add = threadIdx.x >= offset ? scan[threadIdx.x - offset] : 0;
+            __syncthreads();
+            scan[threadIdx.x] += add;
+            __syncthreads();
+        }
+
+        const int first = start + scan[threadIdx.x] - n;
+        for (int jt = 0; jt < n; ++jt) {
+            tile_map[first + jt] = make_int2(e, jt);
+        }
+        start += scan[nthreads - 1];
+        __syncthreads();
+    }
+
+    for (int t = start + threadIdx.x; t < ntiles; t += nthreads) {
+        tile_map[t] = make_int2(0, jt_none);
+    }
+}
+
+void ggml_cuda_launch_mm_ids_tile_map(
+        const int32_t * expert_bounds, int2 * tile_map, const int n_experts, const int J, const int ntiles, const int jt_none, cudaStream_t stream) {
+    constexpr int nthreads = 512;
+    mm_ids_tile_map<nthreads><<<1, nthreads, 0, stream>>>(expert_bounds, tile_map, n_experts, J, ntiles, jt_none);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh
new file mode 100644
index 0000000000..b0e8edb88e
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh
@@ -0,0 +1,8 @@
+#pragma once
+
+void ggml_cuda_launch_mm_ids_helper(
+        const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
+        int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
+
+void ggml_cuda_launch_mm_ids_tile_map(
+        const int32_t * expert_bounds, int2 * tile_map, int n_experts, int J, int ntiles, int jt_none, cudaStream_t stream);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh
new file mode 100644
index 0000000000..9f9fd19738
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh
@@ -0,0 +1,383 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_ampere(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh
new file mode 100644
index 0000000000..9fbe32b697
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh
@@ -0,0 +1,37 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_blackwell(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+    CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+
+    return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh
new file mode 100644
index 0000000000..4a8d89f720
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh
@@ -0,0 +1,185 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q1_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_1, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_1, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q8_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q2_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q3_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q4_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q5_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_Q6_K, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_MXFP4, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+    CASE(GGML_TYPE_NVFP4, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh
new file mode 100644
index 0000000000..24af2ef2b0
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh
@@ -0,0 +1,281 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gcn(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 3,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh
new file mode 100644
index 0000000000..83eb7c146e
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh
@@ -0,0 +1,273 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh
new file mode 100644
index 0000000000..2a8dc9e1a9
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh
@@ -0,0 +1,273 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh
new file mode 100644
index 0000000000..8324d9e1a8
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh
@@ -0,0 +1,273 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna2(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh
new file mode 100644
index 0000000000..ead3825d21
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh
@@ -0,0 +1,290 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh
new file mode 100644
index 0000000000..3a3ef7bd9c
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh
@@ -0,0 +1,274 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 4,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 4,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 4,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 4,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 4,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 4,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 4,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 4,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 1,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 4,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 4,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 4,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 4,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 4,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 4,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 4,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 4,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 1,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 4,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 4,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 4,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 1,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 1,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 1,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 1,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh
new file mode 100644
index 0000000000..9293d9d558
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh
@@ -0,0 +1,290 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh
new file mode 100644
index 0000000000..017b93b546
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh
@@ -0,0 +1,1768 @@
+#pragma once
+
+#include "vecdotq.cuh"
+
+#include "mmq.cuh"
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int blocks_per_iter = MMQ_ITER_K / QK1_0;
+    constexpr int threads_per_row = blocks_per_iter * QI1_0;
+    constexpr int nrows = warp_size / threads_per_row;
+    constexpr int scale_entries_per_block = QK1_0 / QK8_1;
+    constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
+
+    const int txi  = threadIdx.x % threads_per_row;
+    const int kbx  = txi / QI1_0;
+    const int kqsx = txi % QI1_0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + kbx;
+        const int16_t    * qxi = (const int16_t *) bxi->qs + kqsx * 2;
+
+        const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
+#pragma unroll
+        for (int j = 0; j < 2; ++j) {
+            const int q  = qxi[j];
+
+            // unpack crumbs into nibble indices
+            const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
+            const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+            // unpack nibbles into byte values
+            const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >>  0);
+            const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >>  0);
+            const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+            const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+            // unshuffle values
+            const int v0 = __byte_perm(s0, s1, 0x5410);
+            const int v1 = __byte_perm(s0, s1, 0x7632);
+            const int v2 = __byte_perm(s2, s3, 0x5410);
+            const int v3 = __byte_perm(s2, s3, 0x7632);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride           + dst_offset + j*4+0] = v0;
+            x_qs[i*sram_stride           + dst_offset + j*4+1] = v1;
+            x_qs[i*sram_stride           + dst_offset + j*4+2] = v2;
+            x_qs[i*sram_stride           + dst_offset + j*4+3] = v3;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+0] = v0;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+1] = v1;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+2] = v2;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+3] = v3;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+    }
+
+    const int ksx = threadIdx.x % scale_entries_per_row;
+    const int scale_block = ksx / scale_entries_per_block;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps) {
+        int i = i0 + threadIdx.y;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + scale_block;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                           + ksx] = bxi->d;
+#else
+        x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int blocks_per_iter = MMQ_ITER_K / QK2_0;
+    constexpr int threads_per_row = blocks_per_iter * QI2_0;
+    constexpr int nrows = warp_size / threads_per_row;
+    constexpr int scale_entries_per_block = QK2_0 / QK8_1;
+    constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
+
+    const int txi  = threadIdx.x % threads_per_row;
+    const int kbx  = txi / QI2_0;
+    const int kqsx = txi % QI2_0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + kbx;
+        const int16_t    * qxi = (const int16_t *) bxi->qs + kqsx * 4;
+
+        const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
+
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const int q  = qxi[j];
+
+#if defined(GGML_USE_HIP)
+            const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18);
+            const uint32_t qy_bits    = q >> 8;
+            const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18);
+            const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices);
+            const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices);
+#else
+            // unpack even and odd crumbs into byte values
+            const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
+            const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2);
+            // unshuffle values
+            const int qx = __byte_perm(qe, qo, 0x5140);
+            const int qy = __byte_perm(qe, qo, 0x7362);
+#endif // defined(GGML_USE_HIP)
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride           + dst_offset + j*2+0] = qx;
+            x_qs[i*sram_stride           + dst_offset + j*2+1] = qy;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+0] = qx;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+1] = qy;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+    }
+
+    const int ksx = threadIdx.x % scale_entries_per_row;
+    const int scale_block = ksx / scale_entries_per_block;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps) {
+        int i = i0 + threadIdx.y;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + scale_block;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                           + ksx] = bxi->d;
+#else
+        x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_0);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI4_0;
+    const int kqsx = txi % QI4_0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx;
+        const int qs0 = get_int_b2(bxi->qs, kqsx);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0]     = __vsub4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808);
+        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsub4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808);
+#else
+        x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                     + kbxd] = bxi->d;
+#else
+        x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)  || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_1);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI4_1;
+    const int kqsx = txi % QI4_1;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbx;
+        const int qs0 = get_int_b4(bxi->qs, kqsx);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + 0]     = (qs0 >> 0) & 0x0F0F0F0F;
+        x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + QI4_1] = (qs0 >> 4) & 0x0F0F0F0F;
+#else
+        x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_1;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_dm[i*sram_stride                     + kbxd] = bxi->dm;
+#else
+        x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + kbxd] = bxi->dm;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_0, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_0);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI5_0;
+    const int kqsx = txi % QI5_0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbx;
+
+        const int ql = get_int_b2(bxi->qs, kqsx);
+        const int qh = get_int_b2(bxi->qh, 0) >> (4 * kqsx);
+
+        int qs0 = (ql >>  0)   & 0x0F0F0F0F;
+        qs0    |= (qh <<  4)   & 0x00000010;  // 0 ->  4
+        qs0    |= (qh << 11)   & 0x00001000;  // 1 -> 12
+        qs0    |= (qh << 18)   & 0x00100000;  // 2 -> 20
+        qs0    |= (qh << 25)   & 0x10000000;  // 3 -> 28
+        qs0     = __vsub4(qs0, 0x10101010); // subtract 16
+
+        int qs1 = (ql >>  4)   & 0x0F0F0F0F;
+        qs1    |= (qh >> 12)   & 0x00000010;  // 16 ->  4
+        qs1    |= (qh >>  5)   & 0x00001000;  // 17 -> 12
+        qs1    |= (qh <<  2)   & 0x00100000;  // 18 -> 20
+        qs1    |= (qh <<  9)   & 0x10000000;  // 19 -> 28
+        qs1     = __vsub4(qs1, 0x10101010); // subtract 16
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + 0]     = qs0;
+        x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + 0]     = qs0;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_0;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                     + kbxd] = bxi->d;
+#else
+        x_df[i*(MMQ_TILE_NE_K/QI5_0) + i/QI5_0 + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)  || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_1);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI5_1;
+    const int kqsx = txi % QI5_1;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbx;
+
+        const int ql = get_int_b4(bxi->qs, kqsx);
+        const int qh = get_int_b4(bxi->qh, 0) >> (4 * kqsx);
+
+        int qs0 = (ql >>  0) & 0x0F0F0F0F;
+        qs0    |= (qh <<  4) & 0x00000010; // 0 ->  4
+        qs0    |= (qh << 11) & 0x00001000; // 1 -> 12
+        qs0    |= (qh << 18) & 0x00100000; // 2 -> 20
+        qs0    |= (qh << 25) & 0x10000000; // 3 -> 28
+
+        int qs1 = (ql >>  4) & 0x0F0F0F0F;
+        qs1    |= (qh >> 12) & 0x00000010; // 16 ->  4
+        qs1    |= (qh >>  5) & 0x00001000; // 17 -> 12
+        qs1    |= (qh <<  2) & 0x00100000; // 18 -> 20
+        qs1    |= (qh <<  9) & 0x10000000; // 19 -> 28
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + 0]     = qs0;
+        x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + 0]     = qs0;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_1;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_dm[i*sram_stride                     + kbxd] = bxi->dm;
+#else
+        x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + kbxd] = bxi->dm;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    // MMQ_ITER_K / (4 * QR8_0) == 64 required. but NV has only 32 threads per warp
+    constexpr int threads_per_row = 32;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI8_0;
+    const int kqsx = txi % QI8_0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + 0             + txi] = get_int_b2(bxi[0].qs,                   kqsx);
+        x_qs[i*sram_stride + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + 0             + txi] = get_int_b2(bxi[0].qs,                   kqsx);
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                           + kbxd] = bxi->d;
+#else
+        x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR2_K);
+    constexpr int nrows = ggml_cuda_get_physical_warp_size() / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q2_K * bxi = (const block_q2_K *) x + kbx0 + i*stride;
+
+        const int x_ql_0 = get_int_b2(bxi->qs, kqsx);
+
+#pragma unroll
+        for (int l = 0; l < QR2_K; ++l) {
+            const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
+
+            const int x_qs_k = (x_ql_0 >> (2*l)) & 0x03030303;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride           + k] = x_qs_k;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int sc_m = bxi->scales[kqsx];
+#ifdef FAST_FP16_AVAILABLE
+        const half2 x_dm_ik = __hmul2(bxi->dm, make_half2(sc_m & 0x0F, sc_m >> 4));
+#else
+        const float2 bxi_dmf = __half22float2(bxi->dm);
+        const half2 x_dm_ik = make_half2(bxi_dmf.x*(sc_m & 0x0F), bxi_dmf.y*(sc_m >> 4));
+#endif // FAST_FP16_AVAILABLE
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_dm[i*sram_stride         + kqsx] = x_dm_ik;
+#else
+        x_dm[i*(MMQ_TILE_NE_K + 1) + kqsx] = x_dm_ik;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+    int   * x_sc = (int   *) (x_df + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR3_K);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+        const int x_ql_0 = get_int_b2(bxi->qs,    kqsx);
+        const int x_qh_0 = get_int_b2(bxi->hmask, kqsx % (QI3_K/2)) >> (4 * (kqsx / (QI3_K/2)));
+
+#pragma unroll
+        for (int l = 0; l < QR3_K; ++l) {
+            const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
+
+            const int x_ql_k =  (x_ql_0 >> (2*l))       & 0x03030303;
+            const int x_qh_k = ((x_qh_0 >>    l)  << 2) & 0x04040404;
+
+            const int x_qs_k = __vsub4(x_ql_k | x_qh_k, 0x04040404);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride           + k] = x_qs_k;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+    }
+
+    constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+        int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/4;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+        const int ksc = threadIdx.x % 4;
+
+        const int ksc_low = ksc % (QI3_K/8);
+        const int shift_low = 4 * (ksc / (QI3_K/8));
+        const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F;
+
+        const int ksc_high = QI3_K/8;
+        const int shift_high = 2 * ksc;
+        const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;
+
+        const int sc = __vsub4(sc_low | sc_high, 0x20202020);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        const int8_t * sc8 = (const int8_t *) &sc;
+        const float d = bxi->d;
+
+#pragma unroll
+        for (int l = 0; l < int(sizeof(int)); ++l) {
+            x_df[i*sram_stride + sizeof(int)*ksc + l] = d*sc8[l];
+        }
+#else
+        x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = sc;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+#if !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE))
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+        int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+        x_df[i] = bxi->d;
+    }
+#endif // !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)) || defined(AMD_WMMA_AVAILABLE)
+}
+
+static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, const int ksc) {
+    // scale arrangement after the following two lines:
+    //   - ksc == 0: sc0, sc1, sc2, sc3
+    //   - ksc == 1: sc4, sc5, sc6, sc7
+    //   - ksc == 2:  m0,  m1,  m2,  m3
+    //   - ksc == 3:  m4,  m5,  m6,  m7
+    return ((scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F) | // lower 4 bits
+           ((scales[ksc/2]              >> (2 * (ksc % 2)))       & 0x30303030);  // upper 2 bits
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + txs.qs);
+    int   * x_sc = (int   *) (x_dm + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+        const int qs0 = get_int_b4(bxi->qs, txi);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F;
+        x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F;
+#else
+        x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int rows_per_warp = warp_size / 2;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        // Need if on AMD instead of % because warp_size == 64
+        // This causes double work and throughput loss (MI300X)
+        // H100 loses about 100 t/s with 'if' condition over '%'
+        int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
+        if (i < I) {
+#else
+        int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I;
+        {
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            if (fallback) {
+                i = min(i, i_max);
+            }
+
+            const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+
+            const int * scales = (const int *) bxi->scales;
+            const int ksc = threadIdx.x % 2;
+
+            const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
+            const int  m32 = unpack_scales_q45_K(scales, ksc + 2);
+
+            const uint8_t * sc8 = (const uint8_t *) &sc32;
+            const uint8_t *  m8 = (const uint8_t *)  &m32;
+
+            const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
+
+    #pragma unroll
+            for (int l = 0; l < int(sizeof(int)); ++l) {
+                x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
+            }
+        }
+    }
+#else
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+        int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+
+        x_dm[i] = bxi->dm;
+    }
+    constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+        int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
+
+        const int * scales = (const int *) bxi->scales;
+
+        const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
+        const int scales8 = unpack_scales_q45_K(scales, ksc);
+
+        x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
+    }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_dm = (half2 *) (x_qs + txs.qs);
+    int   * x_sc = (int   *) (x_dm + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_K);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+        const int ky = QR5_K*txi;
+
+        const int ql = get_int_b4(bxi->qs, txi);
+        const int ql0 = (ql >> 0) & 0x0F0F0F0F;
+        const int ql1 = (ql >> 4) & 0x0F0F0F0F;
+
+        const int qh = get_int_b4(bxi->qh, txi % (QI5_K/4));
+        const int qh0 = ((qh >> (2 * (txi / (QI5_K/4)) + 0)) << 4) & 0x10101010;
+        const int qh1 = ((qh >> (2 * (txi / (QI5_K/4)) + 1)) << 4) & 0x10101010;
+
+        const int kq0 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + 0;
+        const int kq1 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + QI5_K/4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kq0] = ql0 | qh0;
+        x_qs[i*sram_stride + kq1] = ql1 | qh1;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = ql0 | qh0;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = ql1 | qh1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int rows_per_warp = warp_size / 2;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+#if defined(AMD_MFMA_AVAILABLE)
+        // Need if on AMD instead of % because warp_size == 64
+        // This causes double work and throughput loss (MI300X)
+        // H100 loses about 100 t/s with 'if' condition over '%'
+        int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
+        if (i < I) {
+#else
+        int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I;
+        {
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            if (fallback) {
+                i = min(i, i_max);
+            }
+
+            const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+            const int * scales = (const int *) bxi->scales;
+            const int ksc = threadIdx.x % 2;
+
+            const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
+            const int  m32 = unpack_scales_q45_K(scales, ksc + 2);
+
+            const uint8_t * sc8 = (const uint8_t *) &sc32;
+            const uint8_t *  m8 = (const uint8_t *)  &m32;
+
+            const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
+
+#pragma unroll
+            for (int l = 0; l < int(sizeof(int)); ++l) {
+                x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
+            }
+        }
+    }
+#else
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+        int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+        x_dm[i] = bxi->dm;
+    }
+
+    constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+        int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+        const int * scales = (const int *) bxi->scales;
+
+        const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
+        const int scales8 = unpack_scales_q45_K(scales, ksc);
+
+        x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
+    }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+    int   * x_sc = (int   *) (x_df + MMQ_TILE_NE_K/QI6_K);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+    int   * x_sc = (int   *) (x_df + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR6_K);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+
+        const int ql = get_int_b2(bxi->ql, txi);
+        const int ql0 = (ql >> 0) & 0x0F0F0F0F;
+        const int ql1 = (ql >> 4) & 0x0F0F0F0F;
+
+        const int qh = get_int_b2(bxi->qh, (QI6_K/4) * (txi / (QI6_K/2)) + txi % (QI6_K/4));
+        const int qh0 = ((qh >> ((txi & 0x08) >> 2)) << 4) & 0x30303030;
+        const int qh1 =  (qh >> ((txi & 0x08) >> 2))       & 0x30303030;
+
+        const int kq0 = 2*txi - txi % (QI6_K/2) + 0;
+        const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + kq0] = __vsub4(ql0 | qh0, 0x20202020);
+        x_qs[i*sram_stride + kq1] = __vsub4(ql1 | qh1, 0x20202020);
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsub4(ql0 | qh0, 0x20202020);
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsub4(ql1 | qh1, 0x20202020);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+        int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride]                     = bxi->d;
+#else
+        x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+        int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
+#else
+        x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + threadIdx.x%(MMQ_TILE_NE_K/8)] = get_int_b2(bxi->scales, threadIdx.x%(QI6_K/8));
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I);
+    int   * x_qs = (int   *)  x_tile;
+    half2 * x_ds = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR1_S);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq1_s * bxi = (const block_iq1_s *) x + kbx0 + i*stride;
+
+        const int       qs_packed = get_int_b2(bxi->qs, kqsx);
+        const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+        const int qh = bxi->qh[kqsx];
+
+    #pragma unroll
+        for (int l = 0; l < QR1_S/2; ++l) {
+            const int grid = iq1s_grid_gpu[qs[l] | (((qh >> (3*l)) & 0x07) << 8)];
+
+            const int grid0 = (grid >> 0) & 0x0F0F0F0F;
+            const int grid1 = (grid >> 4) & 0x0F0F0F0F;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid0;
+            x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid1;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid0;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const float  d1q   = __half2float(bxi->d) * (((qh >> 11) & 0x0E) + 1);
+        const float  delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_ds[i*sram_stride             + kqsx] = make_half2(d1q, d1q*delta);
+#else
+        x_ds[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = make_half2(d1q, d1q*delta);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XXS, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XXS)) / 2;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq2_xxs * bxi = (const block_iq2_xxs *) x + kbx0 + i*stride;
+
+        const int q2 = get_int_b2(bxi->qs, 2*kqsx+0);
+        const uint8_t * aux8 = (const uint8_t *) &q2;
+        const uint32_t aux32 = get_int_b2(bxi->qs, 2*kqsx+1);
+
+#pragma unroll
+        for (int l = 0; l < QR2_XXS; ++l) {
+            const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[l]];
+            const uint32_t signs = unpack_ksigns(aux32 >> (7 * l));
+
+            const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+            const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0);
+
+            const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+            const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid0;
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid1;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid0;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int ls = aux32 >> 27 | 1; // (scale * 2 + 1)
+        const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride             + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
+#else
+        x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)  || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XS, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XS)) / 2;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq2_xs * bxi = (const block_iq2_xs *) x + kbx0 + i*stride;
+
+        const int2 q2_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+        const uint16_t * q2 = (const uint16_t *) &q2_packed;
+
+    #pragma unroll
+        for (int l = 0; l < QR2_XS; ++l) {
+            const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l] & 0x1FF];
+            const uint32_t signs = unpack_ksigns(q2[l] >> 9);
+
+            const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+            const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+
+            const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+            const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int ls = bxi->scales[kqsx];
+        const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                             + 2*kqsx+0] = ((ls &  0x0F)*d + d/2)/4;
+        x_df[i*sram_stride                             + 2*kqsx+1] = ((ls >>    4)*d + d/2)/4;
+#else
+        x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls &  0x0F)*d + d/2)/4;
+        x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >>    4)*d + d/2)/4;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_S, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_S)) / 2;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq2_s * bxi = (const block_iq2_s *) x + kbx0 + i*stride;
+
+        const int       qs_packed = get_int_b2(bxi->qs, kqsx);
+        const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+        const int qh = bxi->qh[kqsx];
+
+        const int       signs_packed_32 = get_int_b2(bxi->qs, QK_K/32 + kqsx);
+        const uint8_t * signs_packed_8  = (const uint8_t *) &signs_packed_32;
+
+#pragma unroll
+        for (int l = 0; l < QR2_S; ++l) {
+            const int * grid_pos = (const int *)(iq2s_grid + (qs[l] | ((qh << (8-2*l)) & 0x300)));
+
+            const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
+            const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
+
+            const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0);
+            const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int ls = bxi->scales[kqsx];
+        const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                             + 2*kqsx+0] = ((ls &  0x0F)*d + d/2)/4;
+        x_df[i*sram_stride                             + 2*kqsx+1] = ((ls >>    4)*d + d/2)/4;
+#else
+        x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls &  0x0F)*d + d/2)/4;
+        x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >>    4)*d + d/2)/4;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_XXS, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_XXS)) / 2;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq3_xxs * bxi = (const block_iq3_xxs *) x + kbx0 + i*stride;
+
+        const int2 q3_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+        const uint8_t * q3 = (const uint8_t *) &q3_packed;
+        const uint32_t aux32 = get_int_b2(bxi->qs, QK_K/16 + kqsx);
+
+#pragma unroll
+        for (int l = 0; l < QR3_XXS; ++l) {
+            const int2 grid_pos = make_int2(iq3xxs_grid[q3[2*l+0]], iq3xxs_grid[q3[2*l+1]]);
+            const uint32_t signs = unpack_ksigns(aux32 >> (7*l));
+
+            const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+            const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+
+            const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+            const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int ls = aux32 >> 28;
+        const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride             + kqsx] = (ls*d + d/2)/2;
+#else
+        x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = (ls*d + d/2)/2;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_S)) / 2;
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+        int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq3_s * bxi = (const block_iq3_s *) x + kbx0 + i*stride;
+
+        const int2      qs_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+        const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+        const int qh = bxi->qh[kqsx];
+
+        const int       signs_packed_32 = get_int_b2(bxi->signs, kqsx);
+        const uint8_t * signs_packed_8  = (const uint8_t *) &signs_packed_32;
+
+#pragma unroll
+        for (int l = 0; l < QR3_S; ++l) {
+            const int2 grid_pos = make_int2(
+                iq3s_grid[qs[2*l+0] | ((qh << (8 - 2*l)) & 0x100)],
+                iq3s_grid[qs[2*l+1] | ((qh << (7 - 2*l)) & 0x100)]);
+
+            const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
+            const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
+
+            const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+            const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid_l;
+            x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid_h;
+#else
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid_l;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+
+        const int ls = 1 + 2*((bxi->scales[kqsx/2] >> (((2*kqsx) << 1) & 0x04)) & 0x0F);
+        const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride             + kqsx] = ls*d;
+#else
+        x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = ls*d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_XS);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+
+        const int aux_q4 = get_int_b4(bxi->qs, kqsx);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+        const int k0 = 8 * (kqsx / 4) + kqsx % 4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + k0 + 0] = v.x;
+        x_qs[i*sram_stride + k0 + 4] = v.y;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 4] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int rows_per_warp = warp_size / 8;
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / (MMQ_TILE_NE_K/4);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+
+        const float d = __half2float(bxi->d);
+
+        const int ls = ((bxi->scales_l[(threadIdx.x % 8)/2] >> (4*(threadIdx.x % 2))) & 0x0F)
+            | (((bxi->scales_h >> (2*(threadIdx.x % 8))) & 0x03) << 4);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride             + threadIdx.x % 8] = d * (ls - 32);
+#else
+        x_df[i*(MMQ_TILE_NE_K/4) + i/4 + threadIdx.x % 8] = d * (ls - 32);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_NL, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_NL);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI4_NL;
+    const int kqsx = txi % QI4_NL;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbx;
+
+        const int aux_q4 = get_int_b2(bxi->qs, kqsx);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+        const int k0 = kbx * (2 * QI4_NL) + kqsx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + k0 + 0]      = v.x;
+        x_qs[i*sram_stride + k0 + QI4_NL] = v.y;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0]      = v.x;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI4_NL] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_NL;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                       + kbxd] = __half2float(bxi->d);
+#else
+        x_df[i*(MMQ_TILE_NE_K/QI4_NL) + i/QI4_NL + kbxd] = __half2float(bxi->d);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_MXFP4, I);
+    int   * x_qs = (int   *)  x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_MXFP4);
+    constexpr int nrows = warp_size / threads_per_row;
+    const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+    const int kbx  = txi / QI_MXFP4;
+    const int kqsx = txi % QI_MXFP4;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+        int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbx;
+
+        const int aux_q4 = get_int_b1(bxi->qs, kqsx);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4);
+        const int k0 = kbx * (2 * QI_MXFP4) + kqsx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_qs[i*sram_stride + k0 + 0]        = v.x;
+        x_qs[i*sram_stride + k0 + QI_MXFP4] = v.y;
+#else
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0]        = v.x;
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_MXFP4] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)  || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_MXFP4;
+    constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+    const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+        int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+        if (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        x_df[i*sram_stride                           + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
+#else
+        x_df[i*(MMQ_TILE_NE_K/QI_MXFP4) + i/QI_MXFP4 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
+
+    int *      x_qs = (int *) x_tile;
+    uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
+
+    const int txi = threadIdx.x;
+
+    constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4);
+
+    constexpr int threads_per_row = iter_k / QK_MXFP4;  // each thread processes 1 block
+    constexpr int rows_per_warp   = warp_size / threads_per_row;
+    const int     kbx             = txi % threads_per_row;
+    const int     row_in_warp     = txi / threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+        int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+        if constexpr (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx;
+
+        // quantize_mxfp4_mmq permutes nibbles to match the quantized format
+        const int k0 = kbx * 4;
+        memcpy(x_qs + i*sram_stride + k0, bxi->qs, 16);
+
+        // Load E8M0 scales: pack 2 consecutive scales into one uint32
+        if (kbx % 2 == 0) {
+            uint32_t e = bxi->e;
+            e |= ((bxi + 1)->e << 8);
+            x_sc[i*sram_stride + kbx / 2] = e;
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
+    int   * x_qs = (int   *) x_tile;
+    float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_NVFP4, I);
+    int   * x_qs = (int   *) x_tile;
+    float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4;
+    constexpr int rows_per_warp = warp_size / threads_per_row;
+    const int kbx = threadIdx.x % threads_per_row;
+    const int row_in_warp = threadIdx.x / threads_per_row;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+        int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+        if constexpr (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx;
+        const uint32_t * __restrict__ src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
+        const int kqs = 16 * kbx;
+        const int ksc = 4 * kbx;
+
+#pragma unroll
+        for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
+            const int2 q0 = get_int_from_table_16(src_qs[2 * sub + 0], kvalues_mxfp4);
+            const int2 q1 = get_int_from_table_16(src_qs[2 * sub + 1], kvalues_mxfp4);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+            x_qs[i*sram_stride + kqs + 4 * sub + 0] = q0.x;
+            x_qs[i*sram_stride + kqs + 4 * sub + 1] = q1.x;
+            x_qs[i*sram_stride + kqs + 4 * sub + 2] = q0.y;
+            x_qs[i*sram_stride + kqs + 4 * sub + 3] = q1.y;
+            x_df[i*sram_stride + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
+#else
+            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x;
+            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x;
+            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y;
+            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y;
+            x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
+        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    constexpr int warp_size       = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps          = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
+    constexpr int I               = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
+    constexpr int iter_k          = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4);
+    constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block
+    constexpr int rows_per_warp   = warp_size / threads_per_row;
+    constexpr int sram_stride     = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
+
+    uint32_t * x_u32 = (uint32_t *) x_tile;
+
+    const int txi = threadIdx.x;
+    const int kbx = txi % threads_per_row;
+    const int row_in_warp = txi / threads_per_row;
+
+    const block_nvfp4 * bxi_base = (const block_nvfp4 *) x + kbx0 + kbx;
+    uint32_t * x_u32_scale = x_u32 + 64 + kbx;
+
+#pragma unroll
+    for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+        int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+        if constexpr (fallback) {
+            i = min(i, i_max);
+        }
+
+        const block_nvfp4 * bxi = bxi_base + i * stride;
+
+        const uint32_t * src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
+
+#pragma unroll
+        for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
+            x_u32[i*sram_stride + 8*kbx + 2 * sub + 0] = src_qs[2 * sub + 0];
+            x_u32[i*sram_stride + 8*kbx + 2 * sub + 1] = src_qs[2 * sub + 1];
+        }
+
+        x_u32_scale[i*sram_stride] = get_int_b4(bxi->d, 0);
+    }
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh
new file mode 100644
index 0000000000..a75dd046e5
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh
@@ -0,0 +1,1237 @@
+#pragma once
+
+#include "vecdotq.cuh"
+#include "mma.cuh"
+
+using namespace ggml_cuda_mma;
+
+#include "mmq.cuh"
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I);
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+                const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
+
+                int u[2*VDR_Q4_0_Q8_1_MMQ];
+
+                constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
+                constexpr int mcpy_int = max_cpy / sizeof(int);
+                static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
+
+                int tmp0[4], tmp1[4];
+
+                #pragma unroll
+                for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
+                    ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int]  );
+                    ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_0 + l0 * mcpy_int]);
+                }
+
+                u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
+                u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMQ>
+                    (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u,
+                     x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I);
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+                const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
+
+                int u[2*VDR_Q4_1_Q8_1_MMQ];
+
+                constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
+                constexpr int mcpy_int = max_cpy / sizeof(int);
+                static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
+
+                int tmp0[4], tmp1[4];
+
+                #pragma unroll
+                for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
+                    ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int]  );
+                    ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_1 + l0 * mcpy_int]);
+                }
+
+                u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
+                u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMQ>
+                    (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u,
+                     x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ>
+                    (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K],
+                     x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback, mmq_q8_1_ds_layout ds_layout>
+static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
+    const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr data_layout input_layout = get_input_data_layout();
+    typedef tile<16,  8, int, input_layout>        tile_A;
+    typedef tile<16,  8, int, input_layout>        tile_B;
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+    const half2 * y_ds = (const half2 *) y;
+
+    const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+        const int k0 = k00 + k01;
+
+        tile_A A[ntx];
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+            tile_B B;
+            load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+            float dB;
+            const int j = j0 + tile_C::get_j(0);
+            if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+                dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+            } else {
+                dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    const int i = i0 + n*tile_A::I + tile_C::get_i(l);
+                    const float dA = x_df[i*sram_stride + k0/QI8_0];
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB;
+                }
+            }
+        }
+    }
+#else
+    typedef tile<16, 8, int> tile_A;
+    typedef tile< 8, 8, int> tile_B;
+    typedef tile<16, 8, int> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+    const half2 * y_ds = (const half2 *) y;
+
+    tile_A A[ntx][MMQ_TILE_NE_K/QI8_0];
+    float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0];
+
+    const int i0 = (threadIdx.y/ntx)*rows_per_warp;
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+            const int k0 = k00 + k01;
+
+            load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
+
+#pragma unroll
+            for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+                const int k0 = k00 + k01;
+
+                dA[n][l][k01/QI8_0] = x_df[i*sram_stride + k0/QI8_0];
+            }
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+            tile_B B;
+            float dB[tile_C::ne/2];
+
+            load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int j = j0 + tile_C::get_j(l);
+
+                if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+                    dB[l] =             y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+                } else {
+                    dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+                }
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n][k01/QI8_0], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2];
+                }
+            }
+        }
+    }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I);
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl<QR5_1*VDR_Q5_1_Q8_1_MMQ>
+                    (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+                    x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr data_layout input_layout = get_input_data_layout();
+    typedef tile<16,  8, int, input_layout>        tile_A;
+    typedef tile<16,  8, int, input_layout>        tile_B;
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_dm = (const half2 *) y;
+
+    const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+        const int k0 = k00 + k01;
+
+        tile_A A[ntx];
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+            tile_B B;
+            load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+            const int j = j0 + tile_C::get_j(0);
+            const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    const int i = i0 + n*tile_A::I + tile_C::get_i(l);
+                    float2 dmA = __half22float2(x_dm[i*sram_stride + k0/QI8_1]);
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l];
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y;
+                }
+            }
+        }
+    }
+#else
+    typedef tile<16,  8, int> tile_A;
+    typedef tile< 8,  8, int> tile_B;
+    typedef tile<16,  8, int> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_dm = (const half2 *) y;
+
+    tile_A   A[ntx][MMQ_TILE_NE_K/QI8_1];
+    float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1];
+
+    const int i0 = (threadIdx.y/ntx)*rows_per_warp;
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+            const int k0 = k00 + k01;
+
+            load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
+
+#pragma unroll
+            for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+                const int k0 = k00 + k01;
+
+                dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*sram_stride + k0/QI8_1]);
+            }
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+            tile_B   B;
+            float2 dsB[tile_C::ne/2];
+
+            load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int j = j0 + tile_C::get_j(l);
+
+                dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n][k01/QI8_1], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l];
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y;
+                }
+            }
+        }
+    }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, I);
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl<QI8_0>(
+                    &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0],
+                    &y_qs[j*MMQ_TILE_Y_K + k01],
+                    &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)],
+                    y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+// Used for Q3_K, IQ2_S, and IQ2_XS:
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr data_layout input_layout = get_input_data_layout();
+    typedef tile<16,  4, int, input_layout>        tile_A;
+    typedef tile<16,  4, int, input_layout>        tile_B;
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+    const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+        const int k0 = k00 + k01;
+
+        tile_A A[ntx];
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+            tile_B B;
+            load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+            const int j = j0 + tile_C::get_j(0);
+            const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*sram_stride + k0/4] * dB;
+                }
+            }
+        }
+    }
+#elif defined(TURING_MMA_AVAILABLE)
+
+    typedef tile<16, 4, int> tile_A;
+    typedef tile<16, 8, int> tile_A_8;
+    typedef tile< 8, 4, int> tile_B;
+    typedef tile<16, 8, int> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+    const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+    tile_A  A[ntx][8];
+    float  dA[ntx][tile_C::ne/2][8];
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+            const int k0 = k00 + k01;
+
+            load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+#pragma unroll
+            for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+                const int k0 = k00 + k01;
+
+                dA[n][l][k01/4] = x_df[i*sram_stride + k0/4];
+            }
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
+            tile_B B[2];
+            float dB[tile_C::ne/2];
+
+            // Here load_generic is faster than load_ldmatrix.
+            load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0),         MMQ_TILE_Y_K);
+            load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int j = j0 + tile_C::get_j(l);
+
+                dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C[2];
+                mma(C[0], A[n][k01/4 + 0], B[0]);
+                mma(C[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]);
+                }
+            }
+        }
+    }
+#else
+    GGML_UNUSED_VARS(x, y, sum, k00);
+    NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I);
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+    float2 y_df[J/nwarps];
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += nwarps) {
+        const int j = j0 + threadIdx.y;
+
+        y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
+    }
+
+#pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                constexpr int ns = 2;
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
+                    &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+                    &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
+                    &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+            }
+        }
+    }
+
+    // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop.
+    // As a workaround 2 separate loops are used instead.
+#pragma unroll
+    for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                constexpr int ns = 1;
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
+                    &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+                    &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
+                    &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr data_layout input_layout = get_input_data_layout();
+    typedef tile<16,  4, int, input_layout>        tile_A;
+    typedef tile<16,  4, int, input_layout>        tile_B;
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+    const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+        const int k0 = k00 + k01;
+
+        tile_A A[ntx];
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+            tile_B B;
+            load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+            const int j = j0 + tile_C::get_j(0);
+            const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y;
+            const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0
+                                              : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y
+                                                             : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x);
+
+            tile_C Cm;
+            if (k01 >= MMQ_TILE_NE_K * 3/4) {
+                tile_A A1;
+#pragma unroll
+                for (int l = 0; l < tile_A::ne; ++l) {
+                    A1.x[l] = 0x01010101;
+                }
+                mma(Cm, A1, B);
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C Cd;
+                mma(Cd, A[n], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+                    const float2 dm = __half22float2(x_dm[i*sram_stride + k0/4]);
+                    float tmp = Cd.x[l]*dm.x;
+                    if (k01 >= MMQ_TILE_NE_K * 3/4) {
+                        tmp -= Cm.x[l]*dm.y;
+                    }
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB;
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB;
+                }
+            }
+        }
+    }
+#elif defined(TURING_MMA_AVAILABLE)
+
+    typedef tile<16, 4, int> tile_A;
+    typedef tile<16, 8, int> tile_A_8;
+    typedef tile< 8, 4, int> tile_B;
+    typedef tile<16, 8, int> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+    const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+    tile_A  A[ntx][8];
+    float  dA[ntx][tile_C::ne/2][8];
+    float  mA[ntx][tile_C::ne/2][8];
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+            const int k0 = k00 + k01;
+
+            load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+    }
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+#pragma unroll
+            for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) {
+                const int k0 = k00 + k01;
+
+                const float2 dm = __half22float2(x_dm[i*sram_stride + k0/(QI8_1/2)]);
+
+                dA[n][l][k01/(QI8_1/2)] = dm.x;
+                mA[n][l][k01/(QI8_1/2)] = dm.y;
+            }
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+        float2 dB[tile_C::ne/2];
+
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int j = j0 + tile_C::get_j(l);
+
+            dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
+        }
+
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+            tile_B B[2];
+
+            // Here load_generic is faster than load_ldmatrix.
+            load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0),         MMQ_TILE_Y_K);
+            load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
+
+            tile_C Cm[2];
+            if (k01 >= MMQ_TILE_NE_K * 3/4) {
+                tile_A A1;
+                A1.x[0] = 0x01010101;
+                A1.x[1] = 0x01010101;
+                mma(Cm[0], A1, B[0]);
+                mma(Cm[1], A1, B[1]);
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C Cd[2];
+
+                mma(Cd[0], A[n][k01/4 + 0], B[0]);
+                mma(Cd[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1];
+                    if (k01 >= MMQ_TILE_NE_K * 3/4) {
+                        tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1];
+                    }
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y);
+                }
+            }
+        }
+
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) {
+            float2 sB[tile_C::ne/2];
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int j = j0 + tile_C::get_j(l);
+
+                sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x;
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y;
+                }
+            }
+        }
+    }
+#else
+    GGML_UNUSED_VARS(x, y, sum, k00);
+    NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I);
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + txs.qs;
+    const int   * x_sc = (const int   *) x_df + txs.dm;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4;
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq(
+                    &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales,
+                    x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I);
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+    const int   * x_sc = (const int   *) x_dm + txs.dm;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16);
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq(
+                    &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
+                    x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I);
+    const int   * x_qs = (const int   *) x;
+    const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+    const int   * x_sc = (const int   *) x_dm + txs.dm;
+    const int   * y_qs = (const int   *) y + 4;
+    const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16);
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq(
+                    &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
+                    x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
+
+    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I);
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + txs.qs;
+    const int   * x_sc = (const int   *) x_df + txs.dm;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+// #pragma unroll
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) {
+        const int k0 = k00 + k01;
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+#pragma unroll
+            for (int i0 = 0; i0 < I; i0 += warp_size) {
+                const int i = i0 + threadIdx.x;
+
+                const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]);
+
+                sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq(
+                    &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc,
+                    x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+            }
+        }
+    }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr data_layout input_layout = get_input_data_layout();
+    typedef tile<16,  4, int, input_layout>        tile_A;
+    typedef tile<16,  4, int, input_layout>        tile_B;
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * x_sc = (const int   *) x_df + MMQ_TILE_NE_K/QI6_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+    const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+        const int k0 = k00 + k01;
+
+        tile_A A[ntx];
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+            tile_B B;
+            load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+            const int j = j0 + tile_C::get_j(0);
+            const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C;
+                mma(C, A[n], B);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+                    const int8_t * sc = (const int8_t *) (x_sc + i*sram_stride + k00/16);
+                    sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*sram_stride] * dB;
+                }
+            }
+        }
+    }
+#elif defined(TURING_MMA_AVAILABLE)
+
+    typedef tile<16, 4, int> tile_A;
+    typedef tile< 8, 4, int> tile_B;
+    typedef tile<16, 8, int> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+    const int   * x_qs = (const int   *) x;
+    const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+    const int   * x_sc = (const int   *) x_df + MMQ_TILE_NE_K/QI6_K;
+    const int   * y_qs = (const int   *) y + 4;
+    const float * y_df = (const float *) y;
+
+    const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+    tile_A   A[ntx][8];
+    int    scA[ntx][tile_C::ne/2][8];
+    float   dA[ntx][tile_C::ne/2];
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+            const int k0 = k00 + k01;
+
+            load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + 0),         sram_stride);
+            load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + tile_A::J), sram_stride);
+        }
+
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) {
+            const int k0 = k00 + k01;
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+                const int      sc_packed = x_sc[i*sram_stride + k0/16];
+                const int8_t * sc        = (const int8_t *) &sc_packed;
+
+#pragma unroll
+                for (int ksc = 0; ksc < sizeof(int); ++ksc) {
+                    scA[n][l][k01/4 + ksc] = sc[ksc];
+                }
+            }
+        }
+
+#pragma unroll
+        for (int l = 0; l < tile_C::ne/2; ++l) {
+            const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+            dA[n][l] = x_df[i*sram_stride];
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+        float tmp[ntx][tile_C::ne] = {{0.0f}};
+
+#pragma unroll
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+            tile_B B[2];
+            float dB[tile_C::ne/2];
+
+            // Here load_generic is faster than load_ldmatrix.
+            load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0         + k01, MMQ_TILE_Y_K);
+            load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K);
+
+#pragma unroll
+            for (int l = 0; l < tile_C::ne/2; ++l) {
+                const int j = j0 + tile_C::get_j(l);
+
+                dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+            }
+
+#pragma unroll
+            for (int n = 0; n < ntx; ++n) {
+                tile_C C[2];
+                mma(C[0], A[n][k01/4 + 0], B[0]);
+                mma(C[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+                for (int l = 0; l < tile_C::ne; ++l) {
+                    tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2];
+                }
+            }
+        }
+
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+            for (int l = 0; l < tile_C::ne; ++l) {
+                sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2];
+            }
+        }
+    }
+#else
+    GGML_UNUSED_VARS(x, y, sum, k00);
+    NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+// ---------------------------------------------------------------------------------------------
+
+// Shared MMA kernel for MXFP4 and NVFP4 on Blackwell.
+// Both quantizations encode values as e2m1 (FP4) and produce one uint32 scale per
+// m16n8k64 MMA call; only the PTX kind (scale_vec::2X ue8m0 vs scale_vec::4X ue4m3)
+// and the per-type stride constant differ.
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_fp4_fp4_mma(
+        const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+
+    typedef tile<16, 8, int>   tile_A;
+    typedef tile<8,  8, int>   tile_B;
+    typedef tile<16, 8, float> tile_C;
+
+    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q4);
+    constexpr int ntx           = rows_per_warp / tile_C::I;
+    constexpr int nfrags        = MMQ_TILE_NE_K / tile_A::J;
+
+    y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_K);
+
+    const int *      x_qs = (const int *) x;
+    const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
+    const int *      y_qs = (const int *) y + 4;
+    const uint32_t * y_sc = (const uint32_t *) y;
+
+    // 2 threads per quad supply the packed scale register to the block_scale MMA,
+    // see https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling
+    const int tidx_A = threadIdx.x / 4 + (threadIdx.x % 2) * 8;
+    const int tidx_B = threadIdx.x / 4;
+    const int i0     = (threadIdx.y / ntx) * rows_per_warp;
+
+    tile_A   A[ntx][nfrags];
+    uint32_t scaleA[ntx][nfrags];
+
+#pragma unroll
+    for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+        for (int frag = 0; frag < nfrags; ++frag) {
+            const int k0 = k00 + frag * tile_A::J;
+            load_ldmatrix(A[n][frag], x_qs + (i0 + n * tile_A::I) * sram_stride + k0, sram_stride);
+            scaleA[n][frag] = x_sc[(i0 + n * tile_A::I + tidx_A) * sram_stride + k0 / tile_A::J];
+        }
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx * tile_C::J) {
+        tile_B   B[nfrags];
+        uint32_t scaleB[nfrags];
+
+#pragma unroll
+        for (int frag = 0; frag < nfrags; ++frag) {
+            const int k0 = frag * tile_B::J;
+            load_generic(B[frag], y_qs + j0 * MMQ_TILE_Y_K + k0, MMQ_TILE_Y_K);
+            scaleB[frag] = y_sc[(j0 + tidx_B) * MMQ_TILE_Y_K + frag];
+        }
+
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+            // accumulate in place into the output sum array
+            tile_C & C = *reinterpret_cast<tile_C *>(sum + (j0 / tile_C::J + n) * tile_C::ne);
+#pragma unroll
+            for (int frag = 0; frag < nfrags; ++frag) {
+                mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]);
+            }
+        }
+    }
+}
+
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu
new file mode 100644
index 0000000000..f754837dd3
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu
@@ -0,0 +1,522 @@
+#include "common.cuh"
+#include "mmq.cuh"
+#include "quantize.cuh"
+#include "mmid.cuh"
+
+#include <cstdint>
+
+static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream, const ggml_prec prec_src1) {
+    switch (args.type_x) {
+        case GGML_TYPE_Q1_0:
+            mul_mat_q_case<GGML_TYPE_Q1_0, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q2_0:
+            mul_mat_q_case<GGML_TYPE_Q2_0, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q4_0:
+            mul_mat_q_case<GGML_TYPE_Q4_0, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q4_1:
+            mul_mat_q_case<GGML_TYPE_Q4_1, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q5_0:
+            mul_mat_q_case<GGML_TYPE_Q5_0, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q5_1:
+            mul_mat_q_case<GGML_TYPE_Q5_1, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q8_0:
+            mul_mat_q_case<GGML_TYPE_Q8_0, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+// -----------------------------------------------------------------------
+        case GGML_TYPE_Q2_K:
+            mul_mat_q_case<GGML_TYPE_Q2_K, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            mul_mat_q_case<GGML_TYPE_Q3_K, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            mul_mat_q_case<GGML_TYPE_Q4_K, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            mul_mat_q_case<GGML_TYPE_Q5_K, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            mul_mat_q_case<GGML_TYPE_Q6_K, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+// -----------------------------------------------------------------------
+        case GGML_TYPE_IQ1_S:
+            mul_mat_q_case<GGML_TYPE_IQ1_S, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ2_XXS:
+            mul_mat_q_case<GGML_TYPE_IQ2_XXS, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ2_XS:
+            mul_mat_q_case<GGML_TYPE_IQ2_XS, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ2_S:
+            mul_mat_q_case<GGML_TYPE_IQ2_S, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ3_XXS:
+            mul_mat_q_case<GGML_TYPE_IQ3_XXS, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            mul_mat_q_case<GGML_TYPE_IQ3_S, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            mul_mat_q_case<GGML_TYPE_IQ4_XS, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            mul_mat_q_case<GGML_TYPE_IQ4_NL, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+// -----------------------------------------------------------------------
+        case GGML_TYPE_MXFP4:
+            // src1 at Q4 uses the native FP4 instructions, which are Blackwell-only
+            if (prec_src1 == GGML_PREC_Q4) {
+                mul_mat_q_case<GGML_TYPE_MXFP4, GGML_PREC_Q4>(ctx, args, stream);
+                break;
+            }
+            mul_mat_q_case<GGML_TYPE_MXFP4, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        case GGML_TYPE_NVFP4:
+            if (prec_src1 == GGML_PREC_Q4) {
+                mul_mat_q_case<GGML_TYPE_NVFP4, GGML_PREC_Q4>(ctx, args, stream);
+                break;
+            }
+            mul_mat_q_case<GGML_TYPE_NVFP4, GGML_PREC_Q8>(ctx, args, stream);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+// overrides the src1 precision requested by the graph, "auto" keeps the requested one
+static ggml_prec ggml_cuda_mmq_get_prec_env() {
+    const char * env_c = getenv("GGML_CUDA_MMQ_PREC");
+    if (env_c == nullptr) {
+        return GGML_PREC_UNDEFINED;
+    }
+    std::string env_cpp = env_c;
+    for (char & c : env_cpp) {
+        c = std::tolower(c);
+    }
+    if (env_cpp == "q4") {
+        return GGML_PREC_Q4;
+    }
+    if (env_cpp == "q8") {
+        return GGML_PREC_Q8;
+    }
+    if (env_cpp != "auto") {
+        GGML_LOG_WARN("%s: Unknown value for GGML_CUDA_MMQ_PREC: '%s'. Available: 'q4', 'q8', 'auto'.\n", __func__, env_cpp.c_str());
+    }
+    return GGML_PREC_UNDEFINED;
+}
+
+// src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they
+// default to the native W4A4 instructions on Blackwell.
+static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) {
+    static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env();
+
+    ggml_prec prec = prec_env;
+    if (prec == GGML_PREC_UNDEFINED) {
+        prec = (ggml_prec) ggml_get_op_params_i32(dst, 3);
+    }
+
+    // Q4 only for the FP4 types on Blackwell
+    GGML_ASSERT(prec == GGML_PREC_UNDEFINED || prec == GGML_PREC_Q8 || prec == GGML_PREC_Q4);
+    const bool can_use_q4 = (src0->type == GGML_TYPE_NVFP4 || src0->type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc);
+    if (prec == GGML_PREC_Q8 || !can_use_q4) {
+        return GGML_PREC_Q8;
+    }
+    return GGML_PREC_Q4;
+}
+
+void ggml_cuda_mul_mat_q(
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        int64_t ncols_max) {
+    // FreeToken: bf16 and fp16 activations are quantized directly, except for the native fp4 paths; dst is fp32, bf16 or fp16
+    GGML_ASSERT(        src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16);
+    GGML_ASSERT(        dst->type  == GGML_TYPE_F32 || dst->type  == GGML_TYPE_BF16 || dst->type  == GGML_TYPE_F16);
+    GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
+
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    cudaStream_t stream = ctx.stream();
+
+    const int    id    = ggml_cuda_get_device();
+    const int    cc    = ggml_cuda_info().devices[id].cc;
+    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
+
+    const size_t ts_src0 = ggml_type_size(src0->type);
+    const size_t ts_src1 = ggml_type_size(src1->type);
+    const size_t ts_dst  = ggml_type_size(dst->type);
+
+    GGML_ASSERT(        nb00       == ts_src0);
+    GGML_ASSERT(        nb10       == ts_src1);
+    GGML_ASSERT(        nb0        == ts_dst);
+    GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type));
+
+    const char  * src0_d = (const char  *) src0->data;
+    const float * src1_d = (const float *) src1->data;
+    const mmq_dst  dst_d = {dst->data, dst->type};
+    const nv_bfloat16 * src1_bf16 = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1->data : nullptr;
+    const half        * src1_f16  = src1->type == GGML_TYPE_F16  ? (const half        *) src1->data : nullptr;
+
+    // If src0 is a temporary compute buffer, clear any potential padding.
+    if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+        const size_t size_data  = ggml_nbytes(src0);
+        const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
+        if (size_alloc > size_data) {
+            GGML_ASSERT(ggml_is_contiguously_allocated(src0));
+            GGML_ASSERT(!src0->view_src);
+            CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
+        }
+    }
+
+    const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
+
+    const int64_t s01 = src0->nb[1] / ts_src0;
+    const int64_t s1  =  dst->nb[1] / ts_dst;
+    const int64_t s02 = src0->nb[2] / ts_src0;
+    const int64_t s2  =  dst->nb[2] / ts_dst;
+    const int64_t s03 = src0->nb[3] / ts_src0;
+    const int64_t s3  =  dst->nb[3] / ts_dst;
+
+    const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);
+
+    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+
+    const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
+    GGML_ASSERT(!use_native_fp4 || src1->type == GGML_TYPE_F32);
+    const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
+    const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ            : QK8_1_MMQ;
+
+    int J_best        = 0;
+    int nthreads_best = 0;
+    {
+        int64_t ncols_opt = ne11;
+        if (ids) {
+            const int64_t n_expert_used = ids->ne[0];
+            // a flat grid runs only the tiles that hold tokens, so it sizes them for the average expert
+            ncols_opt = ncols_max > 0 ? ncols_max : (ne12*n_expert_used + ne02 - 1) / ne02;
+
+            // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
+            // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
+            if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
+                ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
+            }
+        }
+
+        int ntiles_J_best = INT_MAX;
+
+        for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+            const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1);
+            if (config.type == GGML_TYPE_COUNT) {
+                continue;
+            }
+
+            if (mmq_get_nbytes_shared(config, cc) > smpbo) {
+                continue;
+            }
+
+            const int ntiles_x = (ncols_opt + config.J - 1) / config.J;
+
+            if (ntiles_x < ntiles_J_best) {
+                J_best = J;
+                nthreads_best = config.nthreads;
+                ntiles_J_best = ntiles_x;
+            }
+        }
+    }
+    GGML_ASSERT(J_best > 0);
+
+    // A tile of size J can read in at most J - 1 extra columns.
+    // For simplicity, round up the padding of a full tile to a multiple of the number of bytes that nthreads can load in parallel.
+    const size_t src1_load_chunk_size = nthreads_best * sizeof(int);
+    const size_t src1_q8_1_padding = ((J_best * sizeof(block_q8_1_mmq) + src1_load_chunk_size - 1) / src1_load_chunk_size)
+        * src1_load_chunk_size;
+
+    if (!ids) {
+        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
+        ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+        ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+        if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+            src1_scale.alloc(ne13*ne12*ne11);
+        }
+
+        {
+            const int64_t s11 = src1->nb[1] / ts_src1;
+            const int64_t s12 = src1->nb[2] / ts_src1;
+            const int64_t s13 = src1->nb[3] / ts_src1;
+            if (use_native_fp4) {
+                static constexpr size_t align_float8 = 32;
+                const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
+                static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
+                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
+                                        ne11, ne12, ne13, stream);
+
+            } else if (src1_bf16) {
+                quantize_mmq_q8_1_cuda(src1_bf16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                                       ne11, ne12, ne13, stream);
+            } else if (src1_f16) {
+                quantize_mmq_q8_1_cuda(src1_f16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                                       ne11, ne12, ne13, stream);
+            } else {
+                quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                                       ne11, ne12, ne13, stream);
+            }
+            CUDA_CHECK(cudaGetLastError());
+        }
+
+        // Stride depends on quantization format
+        const int64_t s12 = use_native_fp4 ?
+                                ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
+                                ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
+        const int64_t s13 = ne12*s12;
+
+        const mmq_args args = {
+            src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
+            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
+            ne00, ne01, ne1, s01, ne11, s1,
+            ne02, ne12, s02, s12, s2,
+            ne03, ne13, s03, s13, s3,
+            ne1, J_best};
+        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
+        return;
+    }
+
+    GGML_ASSERT(ne13 == 1);
+    GGML_ASSERT(nb12 % nb11 == 0);
+    GGML_ASSERT(nb2  % nb1  == 0);
+
+    const int64_t n_expert_used = ids->ne[0];
+    const int64_t ne_get_rows = ne12 * n_expert_used;
+    GGML_ASSERT(ne1 == n_expert_used);
+
+    ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
+    ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access.
+    ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
+
+    // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
+    // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
+    const bool dedup_bcast = ne11 == 1 && n_expert_used > 1;
+
+    {
+        GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
+        const int si1  = ids->nb[1] / ggml_element_size(ids);
+        const int sis1 = nb12 / nb11;
+
+        ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
+            ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
+        CUDA_CHECK(cudaGetLastError());
+    }
+
+    // Every expert's column tiles fit in ntx_flat grid columns, so a flat grid needs no host sync to size it.
+    ggml_cuda_pool_alloc<int2> tile_map(ctx.pool());
+    int ntx_flat = 0;
+    if (ncols_max <= 0) {
+        const int jt_none = (ne_get_rows + J_best - 1) / J_best;
+        ntx_flat = jt_none + std::min<int64_t>(ne02, ne_get_rows);
+        tile_map.alloc(ntx_flat);
+        ggml_cuda_launch_mm_ids_tile_map(expert_bounds.get(), tile_map.get(), ne02, J_best, ntx_flat, jt_none, stream);
+        CUDA_CHECK(cudaGetLastError());
+    }
+
+    const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
+    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+    ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+    if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+        src1_scale.alloc(ne12*n_expert_used);
+    }
+
+    const int64_t ne11_flat = ne12*n_expert_used;
+    const int64_t ne12_flat = 1;
+    const int64_t ne13_flat = 1;
+
+    {
+        const int64_t s11 = src1->nb[1] / ts_src1;
+        const int64_t s12 = src1->nb[2] / ts_src1;
+        const int64_t s13 = src1->nb[3] / ts_src1;
+
+        if (use_native_fp4) {
+            static constexpr size_t align_float8 = 32;
+            const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
+            if (dedup_bcast) {
+                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
+                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+            } else {
+                quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
+                                        ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+            }
+        } else if (dedup_bcast && src1_bf16) {
+            quantize_scatter_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+        } else if (dedup_bcast && src1_f16) {
+            quantize_scatter_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+        } else if (dedup_bcast) {
+            quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+        } else if (src1_bf16) {
+            quantize_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
+                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+        } else if (src1_f16) {
+            quantize_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
+                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+        } else {
+            quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
+                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+        }
+        CUDA_CHECK(cudaGetLastError());
+    }
+
+    static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4");
+    const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
+                                         ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
+    const int64_t s13 = ne12*s12;
+
+    // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
+    const mmq_args args = {
+        src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
+        src1_scale.ptr,
+        ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
+        ne02, ne02, s02, s12, s2,
+        ne03, ne13, s03, s13, s3,
+        ncols_max > 0 ? ncols_max : ne12, J_best, tile_map.ptr, ntx_flat};
+
+    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
+}
+
+bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) {
+#ifdef GGML_CUDA_FORCE_CUBLAS
+    return false;
+#endif // GGML_CUDA_FORCE_CUBLAS
+
+    bool mmq_supported;
+
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+// -------------------------------------------------
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+// -------------------------------------------------
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_IQ4_NL:
+// -------------------------------------------------
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+            mmq_supported = true;
+            break;
+        default:
+            mmq_supported = false;
+            break;
+    }
+
+    if (!mmq_supported) {
+        return false;
+    }
+
+    // MMQ tiles require at least 48 KiB per-block shared memory; fall back to BLAS otherwise.
+    {
+        const int    id    = ggml_cuda_get_device();
+        const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
+        if (smpbo < 48 * 1024) {
+            return false;
+        }
+    }
+
+    if (turing_mma_available(cc)) {
+        return true;
+    }
+
+    if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) {
+        // for MoE, mmq is faster even without native dp4a
+        // TODO: check if cards older than pascal might benefit from this as well
+        return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0;
+    }
+
+#ifdef GGML_CUDA_FORCE_MMQ
+    return true;
+#endif //GGML_CUDA_FORCE_MMQ
+
+    if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
+        return !fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
+    }
+
+    if (amd_mfma_available(cc)) {
+        // As of ROCM 7.0 rocblas/tensile performs very poorly on CDNA3 and hipblaslt (via ROCBLAS_USE_HIPBLASLT)
+        // performs better but is currently suffering from a crash on this architecture.
+        // TODO: Revisit when hipblaslt is fixed on CDNA3
+        if (GGML_CUDA_CC_IS_CDNA3(cc)) {
+            return true;
+        }
+        if (n_experts > 64 || ne11 <= 128) {
+            return true;
+        }
+        if (type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1) {
+            return true;
+        }
+        if (ne11 <= 256 && (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K)) {
+            return true;
+        }
+        return false;
+    }
+
+    if (amd_wmma_available(cc)) {
+        if (GGML_CUDA_CC_IS_RDNA3(cc)) {
+            // High expert counts are almost always better on MMQ due to
+            //     the synchronization overhead in the cuBLAS/hipBLAS path:
+            // https://github.com/ggml-org/llama.cpp/pull/18202
+            if (n_experts >= 64) {
+                return true;
+            }
+
+            // For some quantization types MMQ can have lower peak TOPS than hipBLAS
+            //     so it's only faster for sufficiently small batch sizes:
+            switch (type) {
+                case GGML_TYPE_Q2_K:
+                    return ne11 <= 128;
+                case GGML_TYPE_Q6_K:
+                    return ne11 <= (GGML_CUDA_CC_IS_RDNA3_0(cc) ? 128 : 256);
+                case GGML_TYPE_IQ2_XS:
+                case GGML_TYPE_IQ2_S:
+                    return GGML_CUDA_CC_IS_RDNA3_5(cc) || ne11 <= 128;
+                default:
+                    return true;
+            }
+        }
+
+        // For RDNA4 MMQ is consistently faster than dequantization + hipBLAS:
+        // https://github.com/ggml-org/llama.cpp/pull/18537#issuecomment-3706422301
+        return true;
+    }
+
+    // gfx900 (Vega 10), gfx909, and gfx90c lack native dp4a, losing to dequant + hipBLAS
+    // for dense matrices; keep MMQ only for MoE, where the
+    // hipBLAS path is much slower.
+    if (cc == GGML_CUDA_CC_VEGA || GGML_CUDA_CC_IS_GCN_APU(cc)) {
+        return n_experts > 0;
+    }
+
+    // MUSA: the MMQ kernels compute wrong values on PH1 (MTT S5000).
+    if (cc == GGML_CUDA_CC_PH1) {
+        return false;
+    }
+
+    return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh
new file mode 100644
index 0000000000..596a54eabe
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh
@@ -0,0 +1,1661 @@
+#pragma once
+
+#include "common.cuh"
+
+#include <climits>
+#include <cstdint>
+
+#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available.
+#define MMQ_ITER_K             256
+#define MMQ_ITER_K_FP4         512
+#define MMQ_NWARPS               8
+
+// FreeToken: MMQ writes fp32, bf16 or fp16 results (dst_t); the sums stay fp32 until the store
+template <typename T> static __device__ __forceinline__ T mmq_cast(float v);
+template <> __device__ __forceinline__ float mmq_cast<float>(float v) { return v; }
+template <> __device__ __forceinline__ nv_bfloat16 mmq_cast<nv_bfloat16>(float v) { return __float2bfloat16(v); }
+template <> __device__ __forceinline__ half mmq_cast<half>(float v) { return __float2half(v); }
+
+struct mmq_dst {
+    void * p;
+    ggml_type type;
+};
+
+typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
+typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
+    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);
+
+enum mmq_q8_1_ds_layout {
+    MMQ_Q8_1_DS_LAYOUT_D4,
+    MMQ_Q8_1_DS_LAYOUT_DS4,
+    MMQ_Q8_1_DS_LAYOUT_D2S6,
+};
+
+static constexpr int QK8_1_MMQ  = 4*QK8_1;
+static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ;
+
+struct block_q8_1_mmq {
+    // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block.
+    // The y float data is first grouped as blocks of 128 values.
+    // These blocks are then treated as individual data values and transposed.
+    //
+    // To avoid shared memory bank conflicts each block is padded with 16 bytes.
+    // This padding is also used to store block scales/partial sums.
+    // The scales multiplied with the quantized data are equal to the unquantized values.
+    // The partial sums are obtained by summing up a subgroup of the contained values (prior to quantization)
+    //     and are only needed for performance reasons.
+    //
+    // The exact data stored depends on the x data type.
+    union {
+        float d4[4];    // 1 32 bit scale per 32 values, stored as d0,d1,d2,d3
+        half2 ds4[4];   // 1 16 bit scale + 1 16 bit partial sum per 32 values, stored as d0,s0,d1,s1,d2,s2,d3,s3
+        half  d2s6[8];  // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values,
+                        //     stored as d0,d1,s1,s2,s3,s4,s5
+    };
+    int8_t qs[QK8_1_MMQ];
+};
+
+// this struct is used for fp4 data types (currently only used for Blackwell)
+// mxfp4 has block size 32, each int32 of d4 contains 2 e8m0 scales in the lower 16 bits
+// nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales
+struct block_fp4_mmq {
+    uint32_t d4[4];
+    int8_t   qs[QK_FP4_MMQ / 2];
+};
+
+static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size");
+static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1),      "Unexpected block_q8_1_mmq size");
+static_assert(sizeof(block_fp4_mmq)  == sizeof(block_q8_1_mmq),    "Unexpected block_fp4_mmq size");
+
+static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
+    switch (type_x) {
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+            return MMQ_Q8_1_DS_LAYOUT_DS4;
+        case GGML_TYPE_Q5_0:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_Q5_1:
+            return MMQ_Q8_1_DS_LAYOUT_DS4;
+        case GGML_TYPE_Q8_0:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_MXFP4:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_NVFP4:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_Q2_K:
+            return MMQ_Q8_1_DS_LAYOUT_D2S6;
+        case GGML_TYPE_Q3_K:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+            return MMQ_Q8_1_DS_LAYOUT_DS4;
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        case GGML_TYPE_IQ1_S:
+            return MMQ_Q8_1_DS_LAYOUT_DS4;
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_IQ4_NL:
+            return MMQ_Q8_1_DS_LAYOUT_D4;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+struct tile_x_sizes {
+    int qs;
+    int dm;
+    int sc;
+};
+
+// Decouple shared memory tile sizes from WARP_SIZE to allow for different warp sizes.
+// The K dimension of the tiles has either,
+// 1*MMQ_TILE_NE_K==32 (always for TILE_Y_K) or 2*MMQ_TILE_NE_K==64 (typically for TILE_X_K),
+// 32 bit elements for the quantized data (does not include scales).
+// In other words, the size of the quantized data in the K dimension is a multiple of MMQ_TILE_NE_K.
+// The final tile size in K direction is padded to avoid shared memory bank conflicts,
+// in terms of 32 bit elements that means K % 2 == 1 for dp4a or K % 8 == 4 for mma.
+#define MMQ_TILE_NE_K 32
+
+// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales)
+#define MMQ_TILE_Y_K     (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1)
+#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K
+
+enum ggml_cuda_mmq_sram_layout {
+    GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0,
+    GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1,
+    GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K,
+    GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K,
+    GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K,
+    GGML_CUDA_MMQ_SRAM_LAYOUT_FP4,   // MXFP4 and NVFP4 on Blackwell.
+    GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, // Generic NVFP4
+};
+
+static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_sram_layout sram_layout) {
+    switch (sram_layout) {
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0:
+            return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1:
+            return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_1 + 4;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K:
+            return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K         + 4;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K:
+            return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2       + 4;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K:
+            return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K   + MMQ_TILE_NE_K/8 + 7;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_FP4:
+            return 2*MMQ_TILE_NE_K + 8                     + 4;
+        case GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4:
+            return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2       + 4;
+        default:
+            return -1;
+    }
+}
+
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0)  % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1)  % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K)  % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K)  % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K)  % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4)   % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4) % 8 == 4, "Wrong padding.");
+
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) == ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1), "Wrong tile size for MXFP4");
+
+// Config options for the MMQ kernel.
+// Should not affect results, only speed/register pressure/shared memory use.
+struct ggml_cuda_mmq_config {
+    ggml_type                 type;        // src0->type
+    int                       nthreads;    // Number of threads per CUDA block.
+    int                       occupancy;   // Targeted occupancy for the MMA kernel.
+    int                       I;           // SRAM tile width in src0->ne[1]/dst->ne[0] direction.
+    int                       J;           // SRAM tile width in src1->ne[1]/dst->ne[1] direction.
+    ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements).
+    int                       K_vram;      // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements).
+    bool                      stream_k;    // Whether or not to use stream-k decomposition.
+    bool                      fallback;    // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed.
+
+    constexpr __host__ __device__ ggml_cuda_mmq_config(
+            ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) :
+        type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {}
+
+    constexpr __device__ int rows_per_warp() const {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        return 16;
+#else
+        return J >= 48 && J % 16 == 0 ? 32 : 16;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    }
+
+    // TODO transition all combinations of GPUs and quantizations to the MMA data layout.
+    __host__ int use_mma_data_layout(const int cc) const {
+        if (amd_mfma_available(cc) || amd_wmma_available(cc) || turing_mma_available(cc)) {
+            return true;
+        }
+        return false;
+    }
+
+    constexpr __device__ bool use_mma_data_layout() const {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+        return true;
+#else
+        return false;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+    }
+
+};
+
+#define CASE(type_, nthreads_, occupancy_, I_, J_, sram_layout_, K_vram_, stream_k_, fallback_)                                           \
+    if (type == (type_) && J == (J_) && fallback == (fallback_)) {                                                                        \
+        static_assert((nthreads_) %  32 == 0 && (nthreads_)       <= 512, "bad nthreads");                                                \
+        static_assert(                          (occupancy_)      <=   8, "bad occupancy");                                               \
+        static_assert((I_)        %  32 == 0,                             "bad I");                                                       \
+        static_assert((J_)        %   8 == 0 && (J_)              <= 128, "bad J");                                                       \
+        static_assert((K_vram_)   % 256 == 0,                             "bad K_vram");                                                  \
+        return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
+    }                                                                                                                                     \
+
+#include "mmq-config-pascal-older.cuh"
+#include "mmq-config-pascal-dp4a.cuh"
+#include "mmq-config-ampere.cuh"
+#include "mmq-config-blackwell.cuh"
+
+#include "mmq-config-gcn.cuh"
+#include "mmq-config-cdna.cuh"
+#include "mmq-config-rdna2.cuh"
+#include "mmq-config-rdna3.cuh"
+#include "mmq-config-rdna3-5.cuh"
+#include "mmq-config-rdna4.cuh"
+
+#undef CASE
+
+static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    if (GGML_CUDA_CC_IS_AMD(cc)) {
+        if (GGML_CUDA_CC_IS_GCN(cc)) {
+            return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+        }
+        if (GGML_CUDA_CC_IS_CDNA(cc)) {
+            return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
+        }
+        if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+            return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
+        }
+        if (GGML_CUDA_CC_IS_RDNA3_5(cc)) {
+            return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
+        }
+        if (GGML_CUDA_CC_IS_RDNA3(cc)) {  // covers RDNA 3.0
+            return ggml_cuda_mmq_get_config_rdna3(type, J, fallback);
+        }
+        return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
+    }
+    if (blackwell_mma_available(cc)) {
+        // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1
+        if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) {
+            return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+        }
+        return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
+    }
+    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) {
+        return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+    }
+    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) {
+        return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
+    }
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
+}
+
+static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+#ifdef GGML_USE_HIP
+#ifdef GCN
+    return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+#elif defined(CDNA)
+    return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
+#elif defined(RDNA4)
+    return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
+#elif defined(RDNA3_5)
+    return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
+#elif defined(RDNA3)
+    return ggml_cuda_mmq_get_config_rdna3(type, J, fallback);
+#else
+    return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
+#endif // CDNA
+#else
+#ifdef BLACKWELL_MMA_AVAILABLE
+    // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1
+    if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) {
+        return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+    }
+    return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
+#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
+    return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A
+    return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
+#else
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
+#endif // BLACKWELL_MMA_AVAILABLE
+#endif // GGML_USE_HIP
+    GGML_UNUSED_VARS(type, J, fallback, prec_src1);
+}
+
+static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).type;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).nthreads;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).occupancy;
+}
+
+static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).I;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).I;
+}
+
+static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).J;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).J;
+}
+
+static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).sram_layout;
+}
+
+static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).sram_layout;
+}
+
+static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).K_vram;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).K_vram;
+}
+
+static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).stream_k;
+}
+
+static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).stream_k;
+}
+
+static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).fallback;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).fallback;
+}
+
+// ---------------------------------------------------------------------------------------------
+
+static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc, prec_src1));
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, prec_src1));
+}
+
+static __host__ bool ggml_cuda_mmq_needs_fallback(const int64_t nrows_x) {
+    return nrows_x % 128 != 0;
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type, int J, bool fallback, ggml_prec prec_src1) {
+    return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).rows_per_warp();
+}
+
+#define MMQ_DP4A_TXS_Q4_0    tile_x_sizes{I*MMQ_TILE_NE_K   + I, I*MMQ_TILE_NE_K/QI4_0   + I/QI4_0,     0}
+#define MMQ_DP4A_TXS_Q4_1    tile_x_sizes{I*MMQ_TILE_NE_K   + I, I*MMQ_TILE_NE_K/QI4_1   + I/QI4_1,     0}
+#define MMQ_DP4A_TXS_Q8_0    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_0 + I/(QI8_0/2), 0}
+#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I/(QI8_0/4), 0}
+#define MMQ_DP4A_TXS_Q8_1    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_1 + I/(QI8_1/2), 0}
+#define MMQ_DP4A_TXS_Q2_K    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K         + I,           0}
+#define MMQ_DP4A_TXS_Q3_K    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I,                                     I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q4_K    tile_x_sizes{I*MMQ_TILE_NE_K   + I, I*MMQ_TILE_NE_K/QI4_K,                 I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q5_K    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI5_K   + I/QI5_K,     I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q6_K    tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI6_K   + I/QI6_K,     I*MMQ_TILE_NE_K/8 + I/8}
+
+static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int I) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:    return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_Q2_0:    return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_Q4_0:    return MMQ_DP4A_TXS_Q4_0;
+        case GGML_TYPE_Q4_1:    return MMQ_DP4A_TXS_Q4_1;
+        case GGML_TYPE_Q5_0:    return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_Q5_1:    return MMQ_DP4A_TXS_Q8_1;
+        case GGML_TYPE_Q8_0:    return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_MXFP4:   return MMQ_DP4A_TXS_Q8_1;
+        case GGML_TYPE_NVFP4:   return MMQ_DP4A_TXS_Q8_0_16;
+        case GGML_TYPE_Q2_K:    return MMQ_DP4A_TXS_Q2_K;
+        case GGML_TYPE_Q3_K:    return MMQ_DP4A_TXS_Q3_K;
+        case GGML_TYPE_Q4_K:    return MMQ_DP4A_TXS_Q4_K;
+        case GGML_TYPE_Q5_K:    return MMQ_DP4A_TXS_Q5_K;
+        case GGML_TYPE_Q6_K:    return MMQ_DP4A_TXS_Q6_K;
+        case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_IQ2_XS:  return MMQ_DP4A_TXS_Q8_0_16;
+        case GGML_TYPE_IQ2_S:   return MMQ_DP4A_TXS_Q8_0_16;
+        case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_IQ3_S:   return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_IQ1_S:   return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_IQ4_XS:  return MMQ_DP4A_TXS_Q8_0;
+        case GGML_TYPE_IQ4_NL:  return MMQ_DP4A_TXS_Q8_0;
+        default:                return tile_x_sizes{0, 0, 0};
+    }
+}
+
+// FIXME temporary until all combinations of data types and GPUs can use the MMA data layout
+static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config & config, const int cc) {
+    if (config.use_mma_data_layout(cc)) {
+        return config.I * ggml_cuda_mmq_get_sram_stride(config.sram_layout) * 4;
+    }
+    const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(config.type, config.I);
+    return (txs.qs + txs.dm + txs.sc) * 4;
+}
+
+// ------------------------------------------------------------
+
+#include "mmq-load-tiles.cuh"
+#include "mmq-vec-dot.cuh"
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t = float> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
+        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, dst_t * __restrict__ dst,
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+
+    const bool y_scale_used = y_scale != nullptr;
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += nwarps) {
+        const int j = j0 + threadIdx.y;
+
+        if (j > j_max) {
+            return;
+        }
+
+#pragma unroll
+        for (int i0 = 0; i0 < I; i0 += warp_size) {
+            const int i = i0 + threadIdx.x;
+
+            if (fallback && i > i_max) {
+                continue;
+            }
+
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (y_scale_used) {
+                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                } else {
+                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                }
+            } else {
+                dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]);
+                GGML_UNUSED(y_scale_used);
+            }
+        }
+    }
+}
+
+template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t = float>
+static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
+            const float * __restrict__ sum, const int * __restrict__ ids_dst, dst_t * __restrict__ dst,
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+#else
+    typedef tile<16,  8, int> tile_C;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
+    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+    const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);
+
+    const bool y_scale_used = y_scale != nullptr;
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+        for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+            for (int l = 0; l < tile_C::ne; ++l) {
+                const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l);
+
+                if (j > j_max) {
+                    continue;
+                }
+
+                const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+
+                if (fallback && i > i_max) {
+                    continue;
+                }
+
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (y_scale_used) {
+                        dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                    } else {
+                        dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                    }
+                } else {
+                    dst[ids_dst[j]*stride + i] = mmq_cast<dst_t>(sum[(j0/tile_C::J + n)*tile_C::ne + l]);
+                    GGML_UNUSED(y_scale_used);
+                }
+            }
+        }
+    }
+}
+
+// -------------------------------------------------------------------------------------------------------------------------------------
+
+// TODO remove this struct and use ggml_cuda_mmq_sram_layout instead.
+struct ggml_cuda_mmq_util_funcs {
+    int              vdr;
+    ggml_cuda_mmq_load_tiles_t load_tiles;
+    ggml_cuda_mmq_vec_dot_t    vec_dot;
+    ggml_cuda_mmq_write_back_t write_back;
+
+    constexpr __host__ __device__ ggml_cuda_mmq_util_funcs(
+            int vdr, ggml_cuda_mmq_load_tiles_t load_tiles, ggml_cuda_mmq_vec_dot_t vec_dot, ggml_cuda_mmq_write_back_t write_back) :
+        vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
+};
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
+    if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
+        switch (type) {
+            case GGML_TYPE_Q1_0:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q1_0_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q2_0:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q2_0_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q4_0:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q4_0_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q4_1:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q4_1_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q5_0:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q5_0_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q5_1:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q5_1_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q8_0:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q8_0_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+            case GGML_TYPE_Q2_K:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q2_K_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q3_K:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q3_K_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q4_K:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q4_K_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q5_K:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q5_K_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_Q6_K:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_Q6_K_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+            case GGML_TYPE_IQ1_S:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ1_S_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ2_XXS:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ2_XXS_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ2_XS:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ2_XS_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ2_S:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ2_S_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ3_XXS:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ3_XXS_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ3_S:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ3_S_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ4_XS:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ4_XS_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_IQ4_NL:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_IQ4_NL_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+            case GGML_TYPE_MXFP4:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_MXFP4_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            case GGML_TYPE_NVFP4:
+                return ggml_cuda_mmq_util_funcs(
+                    VDR_NVFP4_Q8_1_MMQ,
+                    ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
+                    ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
+            default:
+                return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
+        }
+    }
+
+// ---------------------------------------------------------------------------------------------
+
+#ifdef BLACKWELL_MMA_AVAILABLE
+    switch (type) {
+        case GGML_TYPE_MXFP4:
+            if (prec_src1 == GGML_PREC_Q4) {
+                return ggml_cuda_mmq_util_funcs(
+                    -1,
+                    ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+            }
+            break;
+        case GGML_TYPE_NVFP4:
+            if (prec_src1 == GGML_PREC_Q4) {
+                return ggml_cuda_mmq_util_funcs(
+                    -1,
+                    ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
+                    ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+            }
+            break;
+        default:
+            break;
+    }
+#endif // BLACKWELL_MMA_AVAILABLE
+
+// ---------------------------------------------------------------------------------------------
+
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q2_0:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q4_0:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q4_1:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q5_0:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q5_1:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q8_0:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+        case GGML_TYPE_Q2_K:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q3_K:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q4_K:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q5_K:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_Q6_K:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+        case GGML_TYPE_IQ1_S:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ2_XXS:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ2_XS:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ2_S:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ3_XXS:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ3_S:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ4_XS:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_IQ4_NL:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+// ---------------------------------------------------------------------------------------------
+        case GGML_TYPE_MXFP4:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
+                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        case GGML_TYPE_NVFP4:
+            return ggml_cuda_mmq_util_funcs(
+                -1,
+                ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+        default:
+            return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
+    }
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vdr;
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().load_tiles;
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vec_dot;
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().write_back;
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1, typename dst_t>
+static __device__ __forceinline__ void mul_mat_q_process_tile(
+        const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+        const int * __restrict__ ids_dst, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
+        const int stride_row_x, const int ncols_y, const int stride_col_dst,
+        const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+
+    constexpr int              warp_size  = ggml_cuda_get_physical_warp_size();
+    constexpr int              nwarps     = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+    constexpr int              qk         = ggml_cuda_type_traits<type>::qk;
+    constexpr int              I          = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+    constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_vec_dot_t    vec_dot    = ggml_cuda_mmq_get_vec_dot<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1>();
+
+    extern __shared__ int data_mul_mat_q[];
+    int * tile_y = data_mul_mat_q + J;
+    int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size);
+
+#if defined(BLACKWELL_MMA_AVAILABLE)
+    // FP4 tile stores 8 blocks. src1 above Q4 uses the generic
+    // Q8_1 tile layout instead of the packed FP4 tile.
+    constexpr int ne_block = ((type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) && prec_src1 == GGML_PREC_Q4) ?
+        QK_FP4_MMQ : QK8_1_MMQ;
+#else
+    constexpr int ne_block = QK8_1_MMQ;
+#endif  // defined(BLACKWELL_MMA_AVAILABLE)
+
+    constexpr int ITER_K          = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
+    constexpr int blocks_per_iter = ITER_K / qk;
+
+    float sum[J*I / (nwarps*warp_size)] = {0.0f};
+
+    constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
+
+    for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
+        load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+        {
+            const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
+#pragma unroll
+            for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
+                int l = l0 + threadIdx.y*warp_size + threadIdx.x;
+
+                tile_y[l] = by0[l];
+            }
+        }
+
+        __syncthreads();
+
+        vec_dot(tile_x, tile_y, sum, 0);
+
+        __syncthreads();
+
+        {
+            const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz);
+#pragma unroll
+            for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
+                int l = l0 + threadIdx.y*warp_size + threadIdx.x;
+
+                tile_y[l] = by0[l];
+            }
+        }
+
+        __syncthreads();
+
+        vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K);
+
+        __syncthreads();
+    }
+
+    if (fixup) {
+        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
+    } else if constexpr (std::is_same_v<dst_t, float>) {
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+    } else if constexpr (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
+        ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, dst_t>(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+    } else {
+        ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1, dst_t>(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+    }
+}
+
+
+// A flat MoE grid runs only the (expert, column tile) pairs listed in tile_map, see ggml_cuda_launch_mm_ids_tile_map.
+static __device__ __forceinline__ void mmq_flat_tile(const int2 * __restrict__ tile_map, int & zt, int & jt) {
+    if (tile_map) {
+        const int2 t = tile_map[jt];
+        zt = t.x;
+        jt = t.y;
+    }
+}
+
+// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, typename dst_t>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
+static __global__ void mul_mat_q(
+        const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
+        const int32_t * __restrict__ expert_bounds, const int2 * __restrict__ tile_map, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
+        const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
+        const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
+        const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
+        const uint3 ntx) {
+
+    // Skip unused template specializations for faster compilation:
+    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) {
+        NO_DEVICE_CODE;
+        return;
+    }
+
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+    constexpr int qk        = ggml_cuda_type_traits<type>::qk;
+    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+
+    const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y
+
+    // Initialize the ids for writing back data with just the index.
+    // For regular matrix multiplications this is never changed.
+    // For MoE the correct indices are loaded from ids_dst.
+    extern __shared__ int ids_dst_shared[]; // Stored at beginning of shared memory.
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
+        const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
+
+        if (j0 + nwarps*warp_size > J && j >= J) {
+            break;
+        }
+
+        ids_dst_shared[j] = j;
+    }
+    __syncthreads();
+
+    if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+        const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
+        const int wt = tmp2.x;
+        int zt = tmp2.y;
+        int jt = blockIdx.y;
+        const int it = blockIdx.x;
+        mmq_flat_tile(tile_map, zt, jt);
+
+        // Defaults for regular matrix multiplication:
+        int col_low    = 0;
+        int col_high   = ncols_dst;
+        int col_diff   = ncols_dst;
+        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }
+
+        if (ids_dst) {
+            col_low  = expert_bounds[zt + 0];
+            col_high = expert_bounds[zt + 1];
+            col_diff = col_high - col_low;
+
+            offset_y   = 0;
+            offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }
+
+            if (jt*J >= col_diff) {
+                return;
+            }
+
+            // __syncthreads(); // There is no previous tile that could cause a race condition.
+#pragma unroll
+            for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
+                const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
+
+                if (j0 + nwarps*warp_size > J && j >= J) {
+                    break;
+                }
+
+                ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+            }
+            __syncthreads();
+        }
+
+        offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
+        offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt*J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }
+
+        const int tile_x_max_i = nrows_x  - it*I - 1;
+        const int tile_y_max_j = col_diff - jt*J - 1;
+
+        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
+
+        constexpr bool fixup = false;
+        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
+             tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
+        return;
+    }
+
+    constexpr int ITER_K          = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
+    constexpr int blocks_per_iter = ITER_K / qk;
+
+    // kbc == k block continuous, current index in continuous ijk space.
+    int kbc      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+    int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+
+    kbc      -= fastmodulo(kbc,      blocks_per_ne00) % blocks_per_iter;
+    kbc_stop -= fastmodulo(kbc_stop, blocks_per_ne00) % blocks_per_iter;
+
+    // kb0 == k index when doing the matrix multiplication for an output tile.
+    int kb0_start = fastmodulo(kbc, blocks_per_ne00);
+    int kb0_stop  = min(blocks_per_ne00.z, uint32_t(kb0_start + kbc_stop - kbc));
+    while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) {
+        int tmp = fastdiv(kbc, blocks_per_ne00);
+        uint2 tmp2 = fast_div_modulo(tmp, ntx);
+        int jt = tmp2.y;
+        tmp = tmp2.x;
+        tmp2 = fast_div_modulo(tmp, nchannels_y);
+        int zt = tmp2.y;
+        tmp = tmp2.x;
+        tmp2 = fast_div_modulo(tmp, nsamples_y);
+        const int wt = tmp2.y;
+        const int it = tmp2.x;
+        mmq_flat_tile(tile_map, zt, jt);
+
+        // Defaults for regular matrix multiplication:
+        int col_low    = 0;
+        int col_high   = ncols_dst;
+        int col_diff   = ncols_dst;
+        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }
+
+        if (ids_dst) {
+            col_low  = expert_bounds[zt + 0];
+            col_high = expert_bounds[zt + 1];
+            col_diff = col_high - col_low;
+
+            offset_y   = 0;
+            offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }
+
+            if (jt*J >= col_diff) {
+                kbc += blocks_per_ne00.z;
+                kbc -= fastmodulo(kbc, blocks_per_ne00);
+
+                kb0_start = 0;
+                kb0_stop  = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc));
+
+                continue;
+            }
+
+            __syncthreads();
+#pragma unroll
+            for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
+                const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
+
+                if (j0 + nwarps*warp_size > J && j >= J) {
+                    break;
+                }
+
+                ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+            }
+            __syncthreads();
+        }
+
+        offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
+        offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt * J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }
+
+        const int tile_x_max_i = nrows_x  - it*I - 1;
+        const int tile_y_max_j = col_diff - jt*J - 1;
+
+        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
+
+        constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
+        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
+             tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+
+        kbc += blocks_per_ne00.z;
+        kbc -= fastmodulo(kbc, blocks_per_ne00);
+
+        kb0_start = 0;
+        kb0_stop  = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc));
+    }
+
+    if (kbc >= kbc_stop) {
+        return;
+    }
+
+    int tmp = fastdiv(kbc, blocks_per_ne00);
+    uint2 tmp2 = fast_div_modulo(tmp, ntx);
+    int jt = tmp2.y;
+    tmp = tmp2.x;
+    tmp2 = fast_div_modulo(tmp, nchannels_y);
+    int zt = tmp2.y;
+    tmp = tmp2.x;
+    tmp2 = fast_div_modulo(tmp, nsamples_y);
+    const int wt = tmp2.y;
+    const int it = tmp2.x;
+    mmq_flat_tile(tile_map, zt, jt);
+
+    // Defaults for regular matrix multiplication:
+    int col_low    = 0;
+    int col_high   = ncols_dst;
+    int col_diff   = ncols_dst;
+    int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+    int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+    int offset_y_scale;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+    } else {
+        GGML_UNUSED(offset_y_scale);
+    }
+
+    if (ids_dst) {
+        col_low  = expert_bounds[zt + 0];
+        col_high = expert_bounds[zt + 1];
+        col_diff = col_high - col_low;
+
+        offset_y   = 0;
+        offset_dst = 0;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = 0;
+        }
+
+        if (jt*J >= col_diff) {
+            return;
+        }
+
+        // The memory layout for the fixup buffer is always contiguous, therefore reset ids:
+        __syncthreads();
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
+            const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
+
+            if (j0 + nwarps*warp_size > J && j >= J) {
+                break;
+            }
+
+            ids_dst_shared[j] = j;
+        }
+        __syncthreads();
+    }
+
+    offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
+    offset_dst += it*I;
+    const float * y_scale_tile = nullptr;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale += col_low + jt * J;
+        y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+    }
+
+    const int tile_x_max_i = nrows_x  - it*I - 1;
+    const int tile_y_max_j = col_diff - jt*J - 1;
+
+    const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
+
+    constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
+    mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
+        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+         stride_row_x, ncols_y, stride_col_dst,
+         tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1)
+static __global__ void mul_mat_q_stream_k_fixup(
+        const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
+        float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst,
+        const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y,
+        const int stride_sample_dst, const uint3 ntx) {
+    constexpr int warp_size       = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps          = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size;
+    constexpr int I               = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+    constexpr int qk              = ggml_cuda_type_traits<type>::qk;
+    constexpr int ITER_K          = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
+    constexpr int blocks_per_iter = ITER_K / qk;
+
+    float sum[J / nwarps] = {0.0f};
+    const int i = blockIdx.y*warp_size + threadIdx.x;
+
+    const int nty = (nrows_x + I - 1) / I;
+
+    const int bidx0 = blockIdx.x;
+
+    // kbc == k block continuous, current index in continuous ijk space.
+    int kbc0      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+    int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+
+    kbc0      -= fastmodulo(kbc0,      blocks_per_ne00) % blocks_per_iter;
+    kbc0_stop -= fastmodulo(kbc0_stop, blocks_per_ne00) % blocks_per_iter;
+
+    const bool did_not_have_any_data   = kbc0 == kbc0_stop;
+    const bool wrote_beginning_of_tile = fastmodulo(kbc0, blocks_per_ne00) == 0;
+    const bool did_not_write_last      = fastdiv(kbc0, blocks_per_ne00) == fastdiv(kbc0_stop, blocks_per_ne00) && fastmodulo(kbc0_stop, blocks_per_ne00) != 0;
+    if (did_not_have_any_data || wrote_beginning_of_tile || did_not_write_last) {
+        return;
+    }
+
+    bool any_fixup = false;
+
+    // Iterate over previous blocks and sum up partial sums written to fixup buffer.
+    // All CUDA blocks that get here must have a previous block that needs a fixup.
+    int bidx = bidx0 - 1;
+    int kbc_stop = kbc0;
+    while(true) {
+        int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+        kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter;
+
+        if (kbc == kbc_stop) { // Did not have any data.
+            bidx--;
+            kbc_stop = kbc;
+            continue;
+        }
+
+        any_fixup = true;
+
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+            sum[j0/nwarps] += tmp_last_tile[bidx*(J*I) + j*I + i];
+        }
+
+        // If this block started in a previous tile we are done and don't need to combine additional partial results.
+        if (fastmodulo(kbc, blocks_per_ne00) == 0 || fastdiv(kbc, blocks_per_ne00) < fastdiv(kbc0, blocks_per_ne00)) {
+            break;
+        }
+        bidx--;
+        kbc_stop = kbc;
+    }
+
+    if (!any_fixup) {
+        return;
+    }
+
+    int tmp = fastdiv(kbc0, blocks_per_ne00);
+    uint2 tmp2 = fast_div_modulo(tmp, ntx);
+    const int jt = tmp2.y;
+    tmp = tmp2.x;
+    tmp2 = fast_div_modulo(tmp, nchannels_y);
+    const int zt = tmp2.y;
+    tmp = tmp2.x;
+    tmp2 = fast_div_modulo(tmp, nsamples_y);
+    const int wt = tmp2.y;
+    const int it = tmp2.x;
+
+    if (!ids_dst) {
+        const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I;
+        dst += offset_dst;
+
+        const int i_max = nrows_x   - it*I - 1;
+        const int j_max = ncols_dst - jt*J - 1;
+        if (fallback && i > i_max) {
+            return;
+        }
+
+#pragma unroll
+        for (int j0 = 0; j0 < J; j0 += nwarps) {
+            const int j = j0 + threadIdx.y;
+
+            if (j > j_max) {
+                return;
+            }
+
+            dst[j*stride_col_dst + i] += sum[j0/nwarps];
+        }
+        return;
+    }
+
+    __shared__ int ids_dst_shared[J];
+    const int col_low  = expert_bounds[zt + 0];
+    const int col_high = expert_bounds[zt + 1];
+    const int col_diff = col_high - col_low;
+
+    for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) {
+        ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+    }
+    __syncthreads();
+
+    const int offset_dst = it*I;
+    dst += offset_dst;
+
+    const int i_max = nrows_x  - it*I - 1;
+    const int j_max = col_diff - jt*J - 1;
+    if (fallback && i > i_max) {
+        return;
+    }
+
+#pragma unroll
+    for (int j0 = 0; j0 < J; j0 += nwarps) {
+        const int j = j0 + threadIdx.y;
+
+        if (j > j_max) {
+            return;
+        }
+
+        dst[ids_dst_shared[j]*stride_col_dst + i] += sum[j0/nwarps];
+    }
+}
+
+// FreeToken: copies an fp32 result with dst's layout into dst
+template <typename dst_t>
+static __global__ void mmq_store_dst(const float * __restrict__ src, dst_t * __restrict__ dst, const int nrows, const int stride) {
+    const int64_t col = blockIdx.x;
+    const int     row = blockIdx.y*blockDim.x + threadIdx.x;
+    if (row < nrows) {
+        dst[col*stride + row] = mmq_cast<dst_t>(src[col*stride + row]);
+    }
+}
+
+struct mmq_args {
+    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; mmq_dst dst;
+    const float * y_scale;
+    int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
+    int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
+    int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
+    int64_t ncols_max;
+    int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
+    const int2 * tile_map = nullptr; // Flat MoE grid: the (expert, column tile) of each of the ntx_flat grid columns.
+    int ntx_flat = 0;
+};
+
+static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
+    const size_t nbs_ids = config.J*sizeof(int);
+    const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc);
+    const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq));
+    return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
+}
+
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+    const int id = ggml_cuda_get_device();
+    const int cc = ggml_cuda_info().devices[id].cc;
+    const int nsm = ggml_cuda_info().devices[id].nsm;
+    const int warp_size = ggml_cuda_info().devices[id].warp_size;
+
+    const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
+    GGML_ASSERT(config.nthreads % warp_size == 0);
+    const int nwarps = config.nthreads / warp_size;
+    const int nbytes_shared = mmq_get_nbytes_shared(config, cc);
+
+    const dim3 block_dims(warp_size, nwarps, 1);
+
+    // a flat MoE grid puts all experts into one channel of ntx_flat column tiles
+    const int64_t nchannels_grid = args.tile_map ? 1 : args.nchannels_y;
+    const int nty  = (args.nrows_x   + config.I - 1) / config.I;
+    const int ntx  = args.tile_map ? args.ntx_flat : (args.ncols_max + config.J - 1) / config.J;
+    const int ntzw = nchannels_grid * args.nsamples_y;
+    const dim3 block_nums_xy_tiling(nty, ntx, ntzw);
+
+    GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
+    GGML_ASSERT(args.nsamples_y  % args.nsamples_x  == 0);
+    const int channel_ratio = args.nchannels_y / args.nchannels_x;
+    const int sample_ratio  = args.nsamples_y  / args.nsamples_x;
+
+    const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits<type>::qk);
+    const uint3 ntx_fd             = init_fastdiv_values(ntx);
+    const uint3 nchannels_y_fd     = init_fastdiv_values(nchannels_grid);
+    const uint3 nsamples_y_fd      = init_fastdiv_values(args.nsamples_y);
+    const uint3 channel_ratio_fd   = init_fastdiv_values(channel_ratio);
+    const uint3 sample_ratio_fd    = init_fastdiv_values(sample_ratio);
+
+    const auto run = [&](auto * dst, const dim3 & block_nums, float * tmp_fixup) {
+        using dst_t = std::remove_pointer_t<decltype(dst)>;
+        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1, dst_t>), nbytes_shared);
+        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J,  true, prec_src1, dst_t>), nbytes_shared);
+        mul_mat_q<type, J, fallback, prec_src1, dst_t><<<block_nums, block_dims, nbytes_shared, stream>>>
+            (args.x, args.y, args.ids_dst, args.expert_bounds, args.tile_map, dst, tmp_fixup, args.y_scale,
+             blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
+             channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
+             sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
+             ntx_fd);
+    };
+
+    // FreeToken: the fixup kernel adds partial sums into an fp32 dst, so a bf16 dst that needs it is computed in fp32 first.
+    // An fp16 dst always is, unless GGML_MMQ_F16_DST builds the fp16 kernels.
+    const auto launch = [&](const dim3 & block_nums, const bool fixup, float * tmp_fixup) {
+        bool via_f32 = fixup && args.dst.type != GGML_TYPE_F32;
+#ifndef GGML_MMQ_F16_DST
+        via_f32 = via_f32 || args.dst.type == GGML_TYPE_F16;
+#endif // GGML_MMQ_F16_DST
+        ggml_cuda_pool_alloc<float> dst_f32(ctx.pool(id));
+        float * dst_float = args.dst.type == GGML_TYPE_F32 ? (float *) args.dst.p : nullptr;
+        if (via_f32) {
+            GGML_ASSERT(args.ids_dst || (args.nchannels_y == 1 && args.nsamples_y == 1));
+            dst_float = dst_f32.alloc(args.ncols_dst * args.nrows_dst);
+        }
+
+        if (dst_float) {
+            run(dst_float, block_nums, tmp_fixup);
+        } else if (args.dst.type == GGML_TYPE_BF16) {
+            run((nv_bfloat16 *) args.dst.p, block_nums, tmp_fixup);
+        }
+#ifdef GGML_MMQ_F16_DST
+        else {
+            run((half *) args.dst.p, block_nums, tmp_fixup);
+        }
+#endif // GGML_MMQ_F16_DST
+
+        if (fixup) {
+            CUDA_CHECK(cudaGetLastError());
+            const dim3 block_nums_fixup(block_nums.x, config.I/warp_size, 1);
+            const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
+            mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
+                (args.ids_dst, args.expert_bounds, dst_float, tmp_fixup, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
+                 args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
+                 ntx_fd);
+        }
+
+        if (via_f32) {
+            CUDA_CHECK(cudaGetLastError());
+            const dim3 block_nums_store(args.ncols_dst, (args.nrows_x + 255) / 256, 1);
+            if (args.dst.type == GGML_TYPE_BF16) {
+                mmq_store_dst<<<block_nums_store, 256, 0, stream>>>(dst_float, (nv_bfloat16 *) args.dst.p, args.nrows_x, args.nrows_dst);
+            } else {
+                mmq_store_dst<<<block_nums_store, 256, 0, stream>>>(dst_float, (half *) args.dst.p, args.nrows_x, args.nrows_dst);
+            }
+        }
+    };
+
+    if (!config.stream_k) {
+        launch(block_nums_xy_tiling, false, nullptr);
+        return;
+    }
+
+    // For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles.
+    // This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important.
+    const int ntiles_dst = ntx * nty * ntzw;
+    const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
+    const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);
+    // The fixup kernel does not read tile_map, so a flat MoE grid always runs one tile per block.
+    const dim3 block_nums_stream_k(args.tile_map || (GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90) ? ntiles_dst : nsm, 1, 1);
+
+    GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow.
+
+    const bool fixup_needed = ntiles_dst % block_nums_stream_k.x != 0;
+
+    ggml_cuda_pool & pool = ctx.pool(id);
+    ggml_cuda_pool_alloc<float> tmp_fixup(pool);
+    if (fixup_needed) {
+        tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I);
+    }
+
+    launch(block_nums_stream_k, fixup_needed, tmp_fixup.ptr);
+}
+
+template <ggml_type type, bool fallback, ggml_prec prec_src1>
+void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+    switch (args.J_best) {
+        case   8:
+            launch_mul_mat_q<type,   8, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  16:
+            launch_mul_mat_q<type,  16, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  24:
+            launch_mul_mat_q<type,  24, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  32:
+            launch_mul_mat_q<type,  32, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  40:
+            launch_mul_mat_q<type,  40, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  48:
+            launch_mul_mat_q<type,  48, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  56:
+            launch_mul_mat_q<type,  56, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  64:
+            launch_mul_mat_q<type,  64, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  72:
+            launch_mul_mat_q<type,  72, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  80:
+            launch_mul_mat_q<type,  80, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  88:
+            launch_mul_mat_q<type,  88, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case  96:
+            launch_mul_mat_q<type,  96, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case 104:
+            launch_mul_mat_q<type, 104, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case 112:
+            launch_mul_mat_q<type, 112, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case 120:
+            launch_mul_mat_q<type, 120, fallback, prec_src1>(ctx, args, stream);
+            break;
+        case 128:
+            launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
+            break;
+        default:
+            fprintf(stderr, "J_best=%d\n", args.J_best);
+            GGML_ABORT("fatal error");
+    }
+}
+
+template <ggml_type type, ggml_prec prec_src1>
+void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+    if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) {
+        constexpr bool fallback = true;
+        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+    } else {
+        constexpr bool fallback = false;
+        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+    }
+}
+
+#define DECL_MMQ_CASE(type)                                                        \
+    template void mul_mat_q_case<type, GGML_PREC_Q8>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
+
+// FP4 variant: uses native FP4 MMA instead of keeping src1 at Q8_1.
+#define DECL_MMQ_CASE_W4A4(type)                                                   \
+    template void mul_mat_q_case<type, GGML_PREC_Q4>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
+
+extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
+extern DECL_MMQ_CASE(GGML_TYPE_Q2_0);
+extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
+extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
+extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
+extern DECL_MMQ_CASE(GGML_TYPE_Q5_1);
+extern DECL_MMQ_CASE(GGML_TYPE_Q8_0);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_Q2_K);
+extern DECL_MMQ_CASE(GGML_TYPE_Q3_K);
+extern DECL_MMQ_CASE(GGML_TYPE_Q4_K);
+extern DECL_MMQ_CASE(GGML_TYPE_Q5_K);
+extern DECL_MMQ_CASE(GGML_TYPE_Q6_K);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
+extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
+extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+extern DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4);
+extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
+
+// -------------------------------------------------------------------------------------------------------------------------
+
+// FreeToken: under ids, ncols_max > 0 bounds the tokens any one expert gets and sizes the tile grid; otherwise the grid is flat, see ggml_cuda_launch_mm_ids_tile_map
+void ggml_cuda_mul_mat_q(
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        int64_t ncols_max = -1);
+
+bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
new file mode 100644
index 0000000000..f3a3ccc683
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
@@ -0,0 +1,1626 @@
+#include "mmvq.cuh"
+#include "convert.cuh"
+#include "quantize.cuh"
+#include "unary.cuh"
+#include "vecdotq.cuh"
+
+#include <cstdint>
+#include <type_traits>
+
+// only enabled on DGX Spark, where it is a gain on every type below. On the higher-bandwidth parts the kernel
+// has little exposed latency left to hide and the extra requests cost more than they save.
+// For perf data, see https://github.com/ggml-org/llama.cpp/pull/26705#issuecomment-5569335031
+#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+// returns true only for those quants that benefit from prefetch and false otherwise
+static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+            return true;
+        default:
+            return false;
+    }
+}
+
+static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) {
+    asm volatile("prefetch.global.L2 [%0];" :: "l"(p));
+}
+#endif
+
+typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs);
+
+static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:    return vec_dot_q1_0_q8_1;
+        case GGML_TYPE_Q2_0:    return vec_dot_q2_0_q8_1;
+        case GGML_TYPE_Q4_0:    return vec_dot_q4_0_q8_1;
+        case GGML_TYPE_Q4_1:    return vec_dot_q4_1_q8_1;
+        case GGML_TYPE_Q5_0:    return vec_dot_q5_0_q8_1;
+        case GGML_TYPE_Q5_1:    return vec_dot_q5_1_q8_1;
+        case GGML_TYPE_Q8_0:    return vec_dot_q8_0_q8_1;
+        case GGML_TYPE_MXFP4:   return vec_dot_mxfp4_q8_1;
+        case GGML_TYPE_NVFP4:   return vec_dot_nvfp4_q8_1;
+        case GGML_TYPE_Q2_K:    return vec_dot_q2_K_q8_1;
+        case GGML_TYPE_Q3_K:    return vec_dot_q3_K_q8_1;
+        case GGML_TYPE_Q4_K:    return vec_dot_q4_K_q8_1;
+        case GGML_TYPE_Q5_K:    return vec_dot_q5_K_q8_1;
+        case GGML_TYPE_Q6_K:    return vec_dot_q6_K_q8_1;
+        case GGML_TYPE_IQ2_XXS: return vec_dot_iq2_xxs_q8_1;
+        case GGML_TYPE_IQ2_XS:  return vec_dot_iq2_xs_q8_1;
+        case GGML_TYPE_IQ2_S:   return vec_dot_iq2_s_q8_1;
+        case GGML_TYPE_IQ3_XXS: return vec_dot_iq3_xxs_q8_1;
+        case GGML_TYPE_IQ1_S:   return vec_dot_iq1_s_q8_1;
+        case GGML_TYPE_IQ1_M:   return vec_dot_iq1_m_q8_1;
+        case GGML_TYPE_IQ4_NL:  return vec_dot_iq4_nl_q8_1;
+        case GGML_TYPE_IQ4_XS:  return vec_dot_iq4_xs_q8_1;
+        case GGML_TYPE_IQ3_S:   return vec_dot_iq3_s_q8_1;
+        default:                return nullptr;
+    }
+}
+
+static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:    return VDR_Q1_0_Q8_1_MMVQ;
+        case GGML_TYPE_Q2_0:    return VDR_Q2_0_Q8_1_MMVQ;
+        case GGML_TYPE_Q4_0:    return VDR_Q4_0_Q8_1_MMVQ;
+        case GGML_TYPE_Q4_1:    return VDR_Q4_1_Q8_1_MMVQ;
+        case GGML_TYPE_Q5_0:    return VDR_Q5_0_Q8_1_MMVQ;
+        case GGML_TYPE_Q5_1:    return VDR_Q5_1_Q8_1_MMVQ;
+        case GGML_TYPE_Q8_0:    return VDR_Q8_0_Q8_1_MMVQ;
+        case GGML_TYPE_MXFP4:   return VDR_MXFP4_Q8_1_MMVQ;
+        case GGML_TYPE_NVFP4:   return VDR_NVFP4_Q8_1_MMVQ;
+        case GGML_TYPE_Q2_K:    return VDR_Q2_K_Q8_1_MMVQ;
+        case GGML_TYPE_Q3_K:    return VDR_Q3_K_Q8_1_MMVQ;
+        case GGML_TYPE_Q4_K:    return VDR_Q4_K_Q8_1_MMVQ;
+        case GGML_TYPE_Q5_K:    return VDR_Q5_K_Q8_1_MMVQ;
+        case GGML_TYPE_Q6_K:    return VDR_Q6_K_Q8_1_MMVQ;
+        case GGML_TYPE_IQ2_XXS: return VDR_IQ2_XXS_Q8_1_MMVQ;
+        case GGML_TYPE_IQ2_XS:  return VDR_IQ2_XS_Q8_1_MMVQ;
+        case GGML_TYPE_IQ2_S:   return VDR_IQ2_S_Q8_1_MMVQ;
+        case GGML_TYPE_IQ3_XXS: return VDR_IQ3_XXS_Q8_1_MMVQ;
+        case GGML_TYPE_IQ3_S:   return VDR_IQ3_S_Q8_1_MMVQ;
+        case GGML_TYPE_IQ4_NL:  return VDR_IQ4_NL_Q8_1_MMVQ;
+        case GGML_TYPE_IQ4_XS:  return VDR_IQ4_XS_Q8_1_MMVQ;
+        default:                return 1;
+    }
+}
+
+enum mmvq_parameter_table_id {
+    MMVQ_PARAMETERS_GENERIC = 0,
+    MMVQ_PARAMETERS_TURING,
+    MMVQ_PARAMETERS_GCN,
+    MMVQ_PARAMETERS_RDNA2,
+    MMVQ_PARAMETERS_RDNA3_0,
+    MMVQ_PARAMETERS_RDNA4,
+    MMVQ_PARAMETERS_GB10
+};
+
+static constexpr __device__ mmvq_parameter_table_id get_device_table_id() {
+#if defined(RDNA4)
+    return MMVQ_PARAMETERS_RDNA4;
+#elif defined(RDNA3_0)
+    return MMVQ_PARAMETERS_RDNA3_0;
+#elif defined(RDNA2) || defined(RDNA3_5)
+    return MMVQ_PARAMETERS_RDNA2;
+#elif defined(GCN) || defined(CDNA)
+    return MMVQ_PARAMETERS_GCN;
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE
+    return MMVQ_PARAMETERS_TURING;
+#elif __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+    return MMVQ_PARAMETERS_GB10;
+#else
+    return MMVQ_PARAMETERS_GENERIC;
+#endif
+}
+
+static __host__ mmvq_parameter_table_id get_device_table_id(int cc) {
+    if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+        return MMVQ_PARAMETERS_RDNA4;
+    }
+    if (GGML_CUDA_CC_IS_RDNA3_0(cc)) {
+        return MMVQ_PARAMETERS_RDNA3_0;
+    }
+    if (GGML_CUDA_CC_IS_RDNA2(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) {
+        return MMVQ_PARAMETERS_RDNA2;
+    }
+    if (GGML_CUDA_CC_IS_GCN(cc) || GGML_CUDA_CC_IS_CDNA(cc)) {
+        return MMVQ_PARAMETERS_GCN;
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_AMPERE) {
+        return MMVQ_PARAMETERS_TURING;
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) {
+        return MMVQ_PARAMETERS_GB10;
+    }
+    return MMVQ_PARAMETERS_GENERIC;
+}
+
+// Per-architecture maximum batch size for which MMVQ should be used for MUL_MAT_ID.
+// Returns a value <= MMVQ_MAX_BATCH_SIZE. Default is MMVQ_MAX_BATCH_SIZE.
+// Check https://github.com/ggml-org/llama.cpp/pull/20905#issuecomment-4145835627 for details
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_pascal_older(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ1_S:   return 6;
+        case GGML_TYPE_IQ1_M:   return 6;
+        case GGML_TYPE_IQ2_S:   return 4;
+        case GGML_TYPE_IQ2_XS:  return 5;
+        case GGML_TYPE_IQ2_XXS: return 5;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 4;
+        case GGML_TYPE_IQ4_NL:  return 6;
+        case GGML_TYPE_IQ4_XS:  return 5;
+        case GGML_TYPE_MXFP4:   return 4;
+        case GGML_TYPE_NVFP4:   return 4;
+        case GGML_TYPE_Q2_K:    return 4;
+        case GGML_TYPE_Q3_K:    return 4;
+        case GGML_TYPE_Q4_0:    return 6;
+        case GGML_TYPE_Q4_1:    return 6;
+        case GGML_TYPE_Q4_K:    return 5;
+        case GGML_TYPE_Q5_0:    return 6;
+        case GGML_TYPE_Q5_1:    return 6;
+        case GGML_TYPE_Q5_K:    return 5;
+        case GGML_TYPE_Q6_K:    return 4;
+        case GGML_TYPE_Q8_0:    return 4;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_turing_plus(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ2_S:   return 7;
+        case GGML_TYPE_IQ3_S:   return 6;
+        case GGML_TYPE_IQ3_XXS: return 7;
+        case GGML_TYPE_MXFP4:   return 7;
+        case GGML_TYPE_NVFP4:   return 8;
+        case GGML_TYPE_Q2_K:    return 7;
+        case GGML_TYPE_Q3_K:    return 5;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_gcn(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ1_S:   return 5;
+        case GGML_TYPE_IQ1_M:   return 5;
+        case GGML_TYPE_IQ2_S:   return 4;
+        case GGML_TYPE_IQ2_XS:  return 4;
+        case GGML_TYPE_IQ2_XXS: return 4;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 4;
+        case GGML_TYPE_IQ4_NL:  return 6;
+        case GGML_TYPE_IQ4_XS:  return 4;
+        case GGML_TYPE_Q2_K:    return 4;
+        case GGML_TYPE_Q3_K:    return 4;
+        case GGML_TYPE_Q4_0:    return 5;
+        case GGML_TYPE_Q4_1:    return 5;
+        case GGML_TYPE_Q4_K:    return 4;
+        case GGML_TYPE_Q5_K:    return 4;
+        case GGML_TYPE_Q6_K:    return 4;
+        case GGML_TYPE_Q8_0:    return 4;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_cdna(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ2_S:   return 5;
+        case GGML_TYPE_IQ2_XS:  return 5;
+        case GGML_TYPE_IQ2_XXS: return 5;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 5;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna1_rdna2(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ2_S:   return 4;
+        case GGML_TYPE_IQ2_XS:  return 4;
+        case GGML_TYPE_IQ2_XXS: return 4;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 4;
+        case GGML_TYPE_Q2_K:    return 7;
+        case GGML_TYPE_Q3_K:    return 4;
+        case GGML_TYPE_Q4_K:    return 5;
+        case GGML_TYPE_Q5_K:    return 6;
+        case GGML_TYPE_Q6_K:    return 5;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna3(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ1_S:   return 6;
+        case GGML_TYPE_IQ1_M:   return 6;
+        case GGML_TYPE_IQ2_S:   return 4;
+        case GGML_TYPE_IQ2_XS:  return 4;
+        case GGML_TYPE_IQ2_XXS: return 4;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 4;
+        case GGML_TYPE_IQ4_NL:  return 6;
+        case GGML_TYPE_IQ4_XS:  return 6;
+        case GGML_TYPE_Q4_K:    return 4;
+        case GGML_TYPE_Q5_K:    return 4;
+        case GGML_TYPE_Q6_K:    return 4;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_IQ1_S:   return 7;
+        case GGML_TYPE_IQ1_M:   return 7;
+        case GGML_TYPE_IQ2_S:   return 4;
+        case GGML_TYPE_IQ2_XS:  return 4;
+        case GGML_TYPE_IQ2_XXS: return 4;
+        case GGML_TYPE_IQ3_S:   return 4;
+        case GGML_TYPE_IQ3_XXS: return 4;
+        case GGML_TYPE_IQ4_NL:  return 7;
+        case GGML_TYPE_IQ4_XS:  return 5;
+        case GGML_TYPE_MXFP4:   return 5;
+        case GGML_TYPE_NVFP4:   return 5;
+        case GGML_TYPE_Q3_K:    return 4;
+        case GGML_TYPE_Q4_0:    return 7;
+        case GGML_TYPE_Q4_1:    return 7;
+        case GGML_TYPE_Q4_K:    return 4;
+        case GGML_TYPE_Q5_0:    return 7;
+        case GGML_TYPE_Q5_1:    return 7;
+        case GGML_TYPE_Q5_K:    return 5;
+        case GGML_TYPE_Q6_K:    return 5;
+        case GGML_TYPE_Q8_0:    return 7;
+        default:                return MMVQ_MAX_BATCH_SIZE;
+    }
+}
+
+// Host function: returns the max batch size for the current arch+type at runtime.
+int get_mmvq_mmid_max_batch(ggml_type type, int cc) {
+    // NVIDIA: P100, Volta, Ada Lovelace, and Blackwell always use MMVQ for MUL_MAT_ID.
+    if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
+        if (cc == GGML_CUDA_CC_PASCAL || cc == GGML_CUDA_CC_VOLTA || cc >= GGML_CUDA_CC_ADA_LOVELACE) {
+            return MMVQ_MAX_BATCH_SIZE;
+        }
+        if (cc >= GGML_CUDA_CC_TURING) {
+            return get_mmvq_mmid_max_batch_turing_plus(type);
+        }
+        return get_mmvq_mmid_max_batch_pascal_older(type);
+    }
+
+    // AMD
+    if (GGML_CUDA_CC_IS_AMD(cc)) {
+        if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+            return get_mmvq_mmid_max_batch_rdna4(type);
+        }
+        if (GGML_CUDA_CC_IS_RDNA3(cc)) {
+            return get_mmvq_mmid_max_batch_rdna3(type);
+        }
+        if (GGML_CUDA_CC_IS_RDNA1(cc) || GGML_CUDA_CC_IS_RDNA2(cc)) {
+            return get_mmvq_mmid_max_batch_rdna1_rdna2(type);
+        }
+        if (GGML_CUDA_CC_IS_CDNA(cc)) {
+            return get_mmvq_mmid_max_batch_cdna(type);
+        }
+        if (GGML_CUDA_CC_IS_GCN(cc)) {
+            return get_mmvq_mmid_max_batch_gcn(type);
+        }
+    }
+    return MMVQ_MAX_BATCH_SIZE;
+}
+
+bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
+    if (!ggml_is_quantized(type)) {
+        return false;
+    }
+    // k-quants cost more to decode and mvq redoes that per column, so MMQ wins sooner.
+    // Only list quant-types MMQ supports, others would fall back to cuBLAS.
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ADA_LOVELACE) {
+        switch (type) { // tuned on RTX 4090
+            case GGML_TYPE_Q2_K:
+                return ne11 <= 4;
+            case GGML_TYPE_Q3_K:
+                return ne11 <= 6;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_BLACKWELL) {
+        switch (type) { // tuned on RTX 5090
+            case GGML_TYPE_Q2_K:
+            case GGML_TYPE_Q3_K:
+            case GGML_TYPE_Q4_K:
+                return ne11 <= 5;
+            case GGML_TYPE_Q5_K:
+                return ne11 <= 6;
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 7;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+        switch (type) { // tuned on DGX Spark GB10
+            case GGML_TYPE_Q2_K:
+                return ne11 <= 6;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ORIN) {
+        switch (type) { // tuned for Jetson Orin
+            case GGML_TYPE_Q2_K:
+            case GGML_TYPE_Q3_K:
+            case GGML_TYPE_Q4_K:
+            case GGML_TYPE_Q5_K:
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 1;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_VOLTA) {
+        switch (type) {
+            case GGML_TYPE_Q2_K:
+                return ne11 <= 4;
+            case GGML_TYPE_Q3_K:
+                return ne11 <= 6;
+            case GGML_TYPE_Q4_K:
+                return ne11 <= 5;
+            case GGML_TYPE_Q5_K:
+                return ne11 <= 6;
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 7;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    if (GGML_CUDA_CC_IS_CDNA(cc)) {
+        if (GGML_CUDA_CC_IS_CDNA1(cc)) {
+            switch (type) {
+                case GGML_TYPE_Q4_0:
+                case GGML_TYPE_Q4_1:
+                    return ne11 <= 7;
+                case GGML_TYPE_Q5_1:
+                    return ne11 <= 7;
+                case GGML_TYPE_Q8_0:
+                    return ne11 <= 6;
+                case GGML_TYPE_Q2_K:
+                    return ne11 <= 4;
+                case GGML_TYPE_Q3_K:
+                    return ne11 <= 3;
+                case GGML_TYPE_Q4_K:
+                    return ne11 <= 2;
+                case GGML_TYPE_Q5_K:
+                    return ne11 <= 3;
+                case GGML_TYPE_Q6_K:
+                    return ne11 <= 4;
+                case GGML_TYPE_IQ1_S:
+                    return ne11 <= 5;
+                case GGML_TYPE_IQ2_XXS:
+                case GGML_TYPE_IQ3_S:
+                case GGML_TYPE_IQ4_XS:
+                    return ne11 <= 6;
+                default:
+                    return ne11 <= MMVQ_MAX_BATCH_SIZE;
+            }
+        }
+        switch (type) { // tuned for CDNA2
+            case GGML_TYPE_Q2_K:
+                return ne11 <= 5;
+            case GGML_TYPE_Q3_K:
+            case GGML_TYPE_Q4_K:
+            case GGML_TYPE_Q5_K:
+                return ne11 <= 3;
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 5;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
+    return ne11 <= MMVQ_MAX_BATCH_SIZE;
+}
+
+// Device constexpr: returns the max batch size for the current arch+type at compile time.
+template <ggml_type type>
+static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() {
+#if defined(RDNA4)
+    return get_mmvq_mmid_max_batch_rdna4(type);
+#elif defined(RDNA3)
+    return get_mmvq_mmid_max_batch_rdna3(type);
+#elif defined(RDNA2) || defined(RDNA1)
+    return get_mmvq_mmid_max_batch_rdna1_rdna2(type);
+#elif defined(CDNA)
+    return get_mmvq_mmid_max_batch_cdna(type);
+#elif defined(GCN)
+    return get_mmvq_mmid_max_batch_gcn(type);
+#elif !defined(GGML_USE_MUSA) && (__CUDA_ARCH__ == GGML_CUDA_CC_PASCAL || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA || __CUDA_ARCH__ >= GGML_CUDA_CC_ADA_LOVELACE)
+    return MMVQ_MAX_BATCH_SIZE;
+#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+    return get_mmvq_mmid_max_batch_turing_plus(type);
+#else
+    return get_mmvq_mmid_max_batch_pascal_older(type);
+#endif
+}
+
+static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) {
+    if (table_id == MMVQ_PARAMETERS_GENERIC) {
+        switch (ncols_dst) {
+            case 1:
+            case 2:
+            case 3:
+            case 4:
+                return 4;
+            case 5:
+            case 6:
+            case 7:
+            case 8:
+                return 2;
+            default:
+                return 1;
+        }
+    } else if (table_id == MMVQ_PARAMETERS_GCN) {
+        switch (ncols_dst) {
+            case 1:
+            case 2:
+            case 3:
+            case 4:
+                return 2;
+            case 5:
+            case 6:
+            case 7:
+            case 8:
+            default:
+                return 1;
+        }
+    }
+    if (table_id == MMVQ_PARAMETERS_RDNA4) {
+        // nwarps=8 benefits types with simple vec_dot on RDNA4 (ncols_dst=1).
+        // Types with complex vec_dot (Q3_K, IQ2_*, IQ3_*) regress due to register
+        // pressure and lookup table contention at higher thread counts.
+        if (ncols_dst == 1) {
+            switch (type) {
+                case GGML_TYPE_Q4_0:
+                case GGML_TYPE_Q4_1:
+                case GGML_TYPE_Q5_0:
+                case GGML_TYPE_Q5_1:
+                case GGML_TYPE_Q8_0:
+                case GGML_TYPE_Q2_K:
+                case GGML_TYPE_Q4_K:
+                case GGML_TYPE_Q5_K:
+                case GGML_TYPE_Q6_K:
+                case GGML_TYPE_IQ4_NL:
+                case GGML_TYPE_IQ4_XS:
+                    return 8;
+                default:
+                    return 1;
+            }
+        }
+        return 1;
+    }
+    if (table_id == MMVQ_PARAMETERS_RDNA3_0) {
+        // RDNA3 (W7900): stricter whitelist than RDNA4.
+        // Q2_K / Q5_K / IQ4_XS regress in full quant sweeps.
+        if (ncols_dst == 1) {
+            switch (type) {
+                case GGML_TYPE_Q4_0:
+                case GGML_TYPE_Q4_1:
+                case GGML_TYPE_Q5_0:
+                case GGML_TYPE_Q5_1:
+                case GGML_TYPE_Q8_0:
+                    return 8;
+                case GGML_TYPE_Q6_K:
+                    return 2;
+                case GGML_TYPE_IQ4_NL:
+                    return 8;
+                default:
+                    return 1;
+            }
+        }
+        return 1;
+    }
+    if (table_id == MMVQ_PARAMETERS_TURING) {
+        if (ncols_dst == 1) {
+            switch (type) {
+                case GGML_TYPE_Q2_K:
+                case GGML_TYPE_Q3_K:
+                case GGML_TYPE_Q4_K:
+                case GGML_TYPE_Q5_K:
+                case GGML_TYPE_Q6_K:
+                    return 2;
+                default:
+                    return 4;
+            }
+        }
+        switch (ncols_dst) {
+            case 2:
+            case 3:
+            case 4:
+                return 4;
+            case 5:
+            case 6:
+            case 7:
+            case 8:
+                return 2;
+            default:
+                return 1;
+        }
+    }
+    if (table_id == MMVQ_PARAMETERS_GB10) {
+        const int generic = calc_nwarps(type, ncols_dst, MMVQ_PARAMETERS_GENERIC);
+        // Only worth the wider block when it actually retires the K loop in half the trips (Observation)
+        if (ncols_dst == 1 && !small_k && halve_iters) {
+            switch (type) {
+                case GGML_TYPE_Q4_0:
+                case GGML_TYPE_Q4_1:
+                case GGML_TYPE_Q5_0:
+                case GGML_TYPE_Q5_1:
+                case GGML_TYPE_Q8_0:
+                case GGML_TYPE_Q4_K:
+                case GGML_TYPE_Q5_K:
+                case GGML_TYPE_Q6_K:
+                case GGML_TYPE_IQ4_NL:
+                    return 2 * generic;
+                default:
+                    break;
+            }
+        }
+        return generic;
+    }
+    return 1;
+}
+
+static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) {
+    if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) {
+        switch (ncols_dst) {
+            case 1:
+                return small_k ? nwarps : 1;
+            case 2:
+            case 3:
+            case 4:
+            case 5:
+            case 6:
+            case 7:
+            case 8:
+                return 2;
+            default:
+                return 1;
+        }
+    }
+    return 1;
+}
+
+// FreeToken: dst_t is float, nv_bfloat16 or half
+template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool halve_iters = false, typename dst_t = float>
+__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mul_mat_vec_q(
+        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst_ptr,
+        const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
+        uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x,
+        const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio,
+        const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst,
+        const uint32_t ids_stride) {
+    const void    * GGML_CUDA_RESTRICT vx  = vx_ptr;
+    const void    * GGML_CUDA_RESTRICT vy  = vy_ptr;
+    const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr;
+    dst_t         * GGML_CUDA_RESTRICT dst = dst_ptr;
+
+    constexpr int qk  = ggml_cuda_type_traits<type>::qk;
+    constexpr int qi  = ggml_cuda_type_traits<type>::qi;
+    constexpr int vdr = get_vdr_mmvq(type);
+    constexpr mmvq_parameter_table_id table_id = get_device_table_id();
+    constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters);
+    constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+
+    constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
+
+    const     int tid = warp_size*threadIdx.y + threadIdx.x;
+    const     int row0 = rows_per_cuda_block*blockIdx.x;
+    const     int blocks_per_row_x = ncols_x / qk;
+    constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi;
+
+    const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1;
+    const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y;
+    if (shared_expert) {
+        vx = fusion.shared_up;
+        if constexpr (std::is_same_v<dst_t, float>) {
+            dst = fusion.shared_dst;
+        }
+        stride_col_dst = fusion.shared_stride_col_dst;
+    }
+
+    uint32_t channel_x;
+    uint32_t channel_y;
+    uint32_t sample_dst;
+
+    ggml_cuda_pdl_sync();
+    channel_x  = shared_expert ? 0 : ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio);
+    channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst;
+    sample_dst = blockIdx.z;
+
+    const uint32_t sample_x    = fastdiv(sample_dst, sample_ratio);
+    const uint32_t sample_y    = sample_dst;
+
+    bool use_gate = false;
+    bool use_bias = false;
+    bool use_gate_bias = false;
+    bool use_scale = false;
+    bool use_gate_scale = false;
+    [[maybe_unused]] const void * vgate = nullptr;
+    const float * x_bias = nullptr;
+    const float * gate_bias = nullptr;
+    const float * x_scale = nullptr;
+    const float * gate_scale = nullptr;
+    ggml_glu_op active_glu;
+    float glu_limit = 0.0f;
+
+    if constexpr (has_fusion) {
+        use_gate      = fusion.gate      != nullptr;
+        use_bias      = fusion.x_bias    != nullptr;
+        use_gate_bias = fusion.gate_bias != nullptr && use_gate;
+        vgate         = shared_expert ? fusion.shared_gate : fusion.gate;
+        x_bias        = (const float *) fusion.x_bias;
+        gate_bias     = (const float *) fusion.gate_bias;
+        active_glu    = fusion.glu_op;
+        glu_limit     = fusion.glu_limit;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            use_scale      = fusion.x_scale    != nullptr;
+            use_gate_scale = fusion.gate_scale != nullptr && use_gate;
+            x_scale        = (const float *) fusion.x_scale;
+            gate_scale     = (const float *) fusion.gate_scale;
+        }
+    }
+
+
+    [[maybe_unused]] float x_biases[ncols_dst]    = { 0.0f };
+    [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f };
+    [[maybe_unused]] float x_scales = 1.0f;
+    [[maybe_unused]] float gate_scales = 1.0f;
+    if constexpr (has_fusion) {
+        // 1. Hide latency by prefetching bias, gates and scales here
+        // 2. load only on threads that won't die after partial sum calculation
+        const uint32_t channel_bias = ids ? channel_x : channel_dst;
+        if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
+            (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+            if (use_bias) {
+                x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
+#pragma unroll
+                for (int j = 0; j < ncols_dst; ++j) {
+                    x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x];
+                }
+            }
+            if (use_gate_bias) {
+                gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
+#pragma unroll
+                for (int j = 0; j < ncols_dst; ++j) {
+                    gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x];
+                }
+            }
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (use_scale) {
+                    x_scales = x_scale[ids ? channel_x : 0];
+                }
+                if (use_gate_scale) {
+                    gate_scales = gate_scale[ids ? channel_x : 0];
+                }
+            }
+        }
+    }
+
+    // partial sum for each thread
+    float tmp[ncols_dst][rows_per_cuda_block] = {{0.0f}};
+    float tmp_gate[ncols_dst][rows_per_cuda_block] = {{0.0f}};
+
+    const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y;
+    const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x;
+
+    for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
+        const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx
+
+        // x block quant index when casting the quants to int
+        const int kqs = vdr * (tid % (qi/vdr));
+
+#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+        // start the next iterations' weight loads early
+        if constexpr (mmvq_should_prefetch(type)) {
+            constexpr int pf_dist = 2; // loop iterations, not blocks
+            const int kbx_pf = kbx + pf_dist*blocks_per_iter;
+            if (kbx_pf < blocks_per_row_x) {
+#pragma unroll
+                for (int i = 0; i < rows_per_cuda_block; ++i) {
+                    const size_t off = (size_t)(kbx_offset + i*stride_row_x + kbx_pf) * ggml_cuda_type_traits<type>::bs;
+                    mmvq_prefetch_l2((const char *) vx + off);
+                    if constexpr (has_fusion) {
+                        if (use_gate) {
+                            mmvq_prefetch_l2((const char *) vgate + off);
+                        }
+                    }
+                }
+            }
+        }
+#endif
+
+#pragma unroll
+        for (int j = 0; j < ncols_dst; ++j) {
+#pragma unroll
+            for (int i = 0; i < rows_per_cuda_block; ++i) {
+                tmp[j][i] += vec_dot_q_cuda(
+                    vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs);
+                if constexpr (has_fusion) {
+                    if (use_gate) {
+                        tmp_gate[j][i] += vec_dot_q_cuda(
+                            vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs);
+                    }
+                }
+            }
+        }
+    }
+
+    __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size];
+    [[maybe_unused]] __shared__ float tmp_shared_gate[(has_fusion && (nwarps-1 > 0)) ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size];
+
+    if (threadIdx.y > 0) {
+#pragma unroll
+        for (int j = 0; j < ncols_dst; ++j) {
+#pragma unroll
+            for (int i = 0; i < rows_per_cuda_block; ++i) {
+                tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i];
+                if constexpr (has_fusion) {
+                    if (use_gate) {
+                        tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i];
+                    }
+                }
+            }
+        }
+    }
+    __syncthreads();
+    if (threadIdx.y > 0) {
+        return;
+    }
+
+    dst += sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row0;
+
+    // sum up partial sums and write back result
+#pragma unroll
+    for (int j = 0; j < ncols_dst; ++j) {
+#pragma unroll
+        for (int i = 0; i < rows_per_cuda_block; ++i) {
+#pragma unroll
+            for (int l = 0; l < nwarps-1; ++l) {
+                tmp[j][i] += tmp_shared[l][j][i][threadIdx.x];
+                if constexpr (has_fusion) {
+                    if (use_gate) {
+                        tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x];
+                    }
+                }
+            }
+            tmp[j][i] = warp_reduce_sum<warp_size>(tmp[j][i]);
+            if constexpr (has_fusion) {
+                if (use_gate) {
+                    tmp_gate[j][i] = warp_reduce_sum<warp_size>(tmp_gate[j][i]);
+                }
+            }
+
+            if (threadIdx.x == i && (rows_per_cuda_block == 1 || uint32_t(row0 + i) < stride_col_dst)) {
+                float result = tmp[j][i];
+                if constexpr (has_fusion) {
+                    if constexpr (type == GGML_TYPE_NVFP4) {
+                        result *= x_scales;
+                    }
+                    result += x_biases[j];
+                    if (use_gate) {
+                        float gate_value = tmp_gate[j][i];
+                        if constexpr (type == GGML_TYPE_NVFP4) {
+                            gate_value *= gate_scales;
+                        }
+                        gate_value += gate_biases[j];
+                        switch (active_glu) {
+                            case GGML_GLU_OP_SWIGLU:
+                                result *= ggml_cuda_op_silu_single(gate_value);
+                                break;
+                            case GGML_GLU_OP_GEGLU:
+                                result *= ggml_cuda_op_gelu_single(gate_value);
+                                break;
+                            case GGML_GLU_OP_SWIGLU_OAI:
+                                result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
+                                break;
+                            case GGML_GLU_OP_SWIGLU_CLAMP:
+                                result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
+                                break;
+                            default:
+                                result = result * gate_value;
+                                break;
+                        }
+                    }
+                }
+                dst[j*stride_col_dst + i] = ggml_cuda_cast<dst_t>(result);
+            }
+        }
+    }
+
+    if constexpr (!has_fusion) {
+        GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, glu_limit, gate_bias, x_bias, x_scale, gate_scale, tmp_gate);
+    }
+    if constexpr (type != GGML_TYPE_NVFP4) {
+        GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales);
+    }
+}
+
+// Dedicated MoE multi-token kernel.
+// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst)
+// Block: (warp_size, ncols_dst) - each warp handles one token independently.
+// No shared memory reduction needed since each warp works alone.
+template <ggml_type type, int c_rows_per_block, bool has_fusion = false, typename dst_t = float>
+__launch_bounds__(get_mmvq_mmid_max_batch_for_device<type>()*ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mul_mat_vec_q_moe(
+        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion,
+        dst_t * dst_ptr,
+        const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
+        const uint32_t stride_row_x, const uint32_t stride_col_y, uint32_t stride_col_dst,
+        const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+        const uint32_t ncols_dst, const uint32_t ids_stride) {
+    const void    * GGML_CUDA_RESTRICT vx  = vx_ptr;
+    const void    * GGML_CUDA_RESTRICT vy  = vy_ptr;
+    const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr;
+    dst_t         * GGML_CUDA_RESTRICT dst = dst_ptr;
+
+    constexpr int qk  = ggml_cuda_type_traits<type>::qk;
+    constexpr int qi  = ggml_cuda_type_traits<type>::qi;
+    constexpr int vdr = get_vdr_mmvq(type);
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+
+    constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
+
+    const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1;
+    if (shared_expert) {
+        vx = fusion.shared_up;
+        if constexpr (std::is_same_v<dst_t, float>) {
+            dst = fusion.shared_dst;
+        }
+        stride_col_dst = fusion.shared_stride_col_dst;
+    }
+
+    // fuse gate, bias, scales, and glu_op into the up projection
+    bool use_gate = false;
+    const void  * vgate      = nullptr;
+    const float * x_bias     = nullptr;
+    const float * gate_bias  = nullptr;
+    const float * x_scale    = nullptr;
+    const float * gate_scale = nullptr;
+    ggml_glu_op   active_glu = GGML_GLU_OP_SWIGLU;
+    float         glu_limit  = 0.0f;
+
+    if constexpr (has_fusion) {
+        use_gate   = fusion.gate != nullptr;
+        vgate      = shared_expert ? fusion.shared_gate : fusion.gate;
+        x_bias     = (const float *) fusion.x_bias;
+        gate_bias  = (const float *) fusion.gate_bias;
+        active_glu = fusion.glu_op;
+        glu_limit  = fusion.glu_limit;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            x_scale    = (const float *) fusion.x_scale;
+            gate_scale = (const float *) fusion.gate_scale;
+        }
+    }
+
+    const uint32_t token_idx   = threadIdx.y;
+    const int      row0        = c_rows_per_block*blockIdx.x;
+    const int      blocks_per_row_x = ncols_x / qk;
+    constexpr int  blocks_per_iter  = vdr * warp_size / qi;
+
+    const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y;
+
+    if (token_idx >= ncols_dst) {
+        return;
+    }
+
+    ggml_cuda_pdl_sync();
+    const uint32_t channel_x = shared_expert ? 0 : ids[channel_dst + token_idx * ids_stride];
+    const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y);
+
+    const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
+    const int kbx_offset  = channel_x*stride_channel_x + row0*stride_row_x;
+
+    // partial sum for each thread
+    float tmp[c_rows_per_block] = {0.0f};
+    float tmp_gate[c_rows_per_block] = {0.0f};
+
+    for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
+        const int kby = kbx * (qk/QK8_1);
+        const int kqs = vdr * (threadIdx.x % (qi/vdr));
+
+#pragma unroll
+        for (int i = 0; i < c_rows_per_block; ++i) {
+            tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
+            if constexpr (has_fusion) {
+                if (use_gate) {
+                    tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
+                }
+            }
+        }
+    }
+
+    ggml_cuda_pdl_lc();
+
+    // Warp-level reduction only - no shared memory needed
+#pragma unroll
+    for (int i = 0; i < c_rows_per_block; ++i) {
+        tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
+        if constexpr (has_fusion) {
+            if (use_gate) {
+                tmp_gate[i] = warp_reduce_sum<warp_size>(tmp_gate[i]);
+            }
+        }
+    }
+
+    // Write results
+    if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
+        float result = tmp[threadIdx.x];
+        if constexpr (has_fusion) {
+            const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x;
+
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (x_scale) {
+                    result *= x_scale[channel_x];
+                }
+            }
+            if (x_bias) {
+                result += x_bias[bias_idx];
+            }
+            if (use_gate) {
+                float gate_value = tmp_gate[threadIdx.x];
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (gate_scale) {
+                        gate_value *= gate_scale[channel_x];
+                    }
+                }
+                if (gate_bias) {
+                    gate_value += gate_bias[bias_idx];
+                }
+                switch (active_glu) {
+                    case GGML_GLU_OP_SWIGLU:
+                        result *= ggml_cuda_op_silu_single(gate_value);
+                        break;
+                    case GGML_GLU_OP_GEGLU:
+                        result *= ggml_cuda_op_gelu_single(gate_value);
+                        break;
+                    case GGML_GLU_OP_SWIGLU_OAI:
+                        result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
+                        break;
+                    case GGML_GLU_OP_SWIGLU_CLAMP:
+                        result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
+                        break;
+                    default:
+                        result = result * gate_value;
+                        break;
+                }
+            }
+        }
+        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = ggml_cuda_cast<dst_t>(result);
+    }
+
+    if constexpr (!has_fusion) {
+        GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale);
+    } else if constexpr (type != GGML_TYPE_NVFP4) {
+        GGML_UNUSED_VARS(x_scale, gate_scale);
+    }
+}
+
+template<ggml_type type>
+static std::pair<dim3, dim3> calc_launch_params(
+        const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens,
+        const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) {
+    const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters);
+    const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
+    const int64_t nblocks = (nrows_x + rpb - 1) / rpb;
+    const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens);
+    const dim3 block_dims(warp_size, nwarps, 1);
+    return {block_nums, block_dims};
+}
+
+template<ggml_type type, int c_ncols_dst, bool small_k = false, bool halve_iters = false, typename dst_t>
+static void mul_mat_vec_q_switch_fusion(
+        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+        const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
+        const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x,
+        const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio,
+        const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst,
+        const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared,
+        const uint32_t ids_stride, cudaStream_t stream) {
+
+    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
+                            fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
+    if constexpr (c_ncols_dst == 1) {
+        if (has_fusion) {
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
+            ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, true, small_k, halve_iters, dst_t>, launch_params,
+                 vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
+            return;
+        }
+    }
+
+    GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1");
+
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
+    ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, false, small_k, halve_iters, dst_t>, launch_params,
+        vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
+        channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
+        sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
+}
+
+template <ggml_type type, typename dst_t>
+static void mul_mat_vec_q_moe_launch(
+        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+        const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
+        const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
+        const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+        const uint32_t ncols_dst, const uint32_t ids_stride,
+        const int warp_size, const int nchannels_dst, cudaStream_t stream) {
+
+    constexpr int rows_per_block = 2; // 2 gives best perf based on tuning
+    const int64_t nblocks_rows = (nrows_x + rows_per_block - 1) / rows_per_block;
+    const dim3 block_nums(nblocks_rows, nchannels_dst + (fusion.shared_up != nullptr));
+    const dim3 block_dims(warp_size, ncols_dst);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
+
+    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
+                            fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
+
+    if (has_fusion) {
+        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true, dst_t>, launch_params,
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+            stride_row_x, stride_col_y, stride_col_dst,
+            stride_channel_x, stride_channel_y, stride_channel_dst,
+            ncols_dst, ids_stride);
+    } else {
+        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false, dst_t>, launch_params,
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+            stride_row_x, stride_col_y, stride_col_dst,
+            stride_channel_x, stride_channel_y, stride_channel_dst,
+            ncols_dst, ids_stride);
+    }
+}
+
+template <ggml_type type, typename dst_t>
+static void mul_mat_vec_q_switch_ncols_dst(
+        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+        const int ncols_x, const int nrows_x, const int ncols_dst,
+        const int stride_row_x, const int stride_col_y, const int stride_col_dst,
+        const int nchannels_x, const int nchannels_y, const int nchannels_dst,
+        const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
+        const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
+        const int ids_stride, cudaStream_t stream) {
+
+    GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0);
+    GGML_ASSERT(ncols_dst <= MMVQ_MAX_BATCH_SIZE);
+
+    const uint3 nchannels_y_fd   = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0);
+    const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0)              : init_fastdiv_values(nchannels_dst / nchannels_x);
+    const uint3 sample_ratio_fd  = init_fastdiv_values(nsamples_dst  / nsamples_x);
+
+    const int device = ggml_cuda_get_device();
+    const int                     cc        = ggml_cuda_info().devices[device].cc;
+    const int warp_size = ggml_cuda_info().devices[device].warp_size;
+    const mmvq_parameter_table_id table_id  = get_device_table_id(cc);
+
+    const bool has_ids = ids != nullptr;
+
+    // How the K loop divides up at the baseline block width, both decisions below use these.
+    constexpr int qk                    = ggml_cuda_type_traits<type>::qk;
+    constexpr int qi                    = ggml_cuda_type_traits<type>::qi;
+    constexpr int vdr                   = get_vdr_mmvq(type);
+    const int     blocks_per_row_x      = ncols_x / qk;
+    const int     blocks_per_iter_1warp = vdr * warp_size / qi;
+
+    const auto should_use_small_k = [&](int c_ncols_dst) {
+        // When K is small, increase rows_per_block to match nwarps so each warp has more work to do
+        // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle.
+        const int  nwarps = calc_nwarps(type, c_ncols_dst, table_id);
+        bool       use    = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp;
+
+        constexpr std::array<ggml_type, 2> iq_slow_turing = {
+            GGML_TYPE_IQ3_XXS,
+            GGML_TYPE_IQ3_S,
+        };
+        constexpr std::array<ggml_type, 8> iq_slow_other = {
+            GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M,   GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS,
+            GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S,   GGML_TYPE_IQ4_XS,
+        };
+        constexpr std::array<ggml_type, 3> slow_pascal = {
+            GGML_TYPE_IQ3_S,
+            GGML_TYPE_Q2_K,
+            GGML_TYPE_Q3_K,
+        };
+
+        const bool is_nvidia_turing_plus  = GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING;
+        const bool is_nvidia_pascal_older = GGML_CUDA_CC_IS_NVIDIA(cc) && cc < GGML_CUDA_CC_VOLTA;
+
+        if (is_nvidia_turing_plus) {
+            if (ncols_dst == 1 &&
+                    std::find(iq_slow_turing.begin(), iq_slow_turing.end(), type) != iq_slow_turing.end()) {
+                use = false;
+            }
+        } else if ((ncols_dst == 1 && std::find(iq_slow_other.begin(), iq_slow_other.end(), type) != iq_slow_other.end()) ||
+                (is_nvidia_pascal_older && std::find(slow_pascal.begin(), slow_pascal.end(), type) != slow_pascal.end()) ||
+                GGML_CUDA_CC_IS_RDNA(cc)) {
+            use = false;
+        }
+
+        return use;
+    };
+
+    // Whether doubling nwarps pays off on the ncols_dst == 1 path, where K sets the K loop trip count.
+    const auto should_halve_iters = [&] {
+        if (table_id != MMVQ_PARAMETERS_GB10) {
+            return false;
+        }
+
+        // Expert rows are gathered per token, so a wider block adds reduction work without reuse.
+        if (has_ids) {
+            return false;
+        }
+
+        const int blocks_per_iter = calc_nwarps(type, 1, table_id) * blocks_per_iter_1warp;
+        const int iters           = (blocks_per_row_x + blocks_per_iter - 1) /  blocks_per_iter;
+        const int iters_wide      = (blocks_per_row_x + blocks_per_iter * 2 - 1) / (blocks_per_iter * 2);
+
+        // An odd trip count leaves half the wider block idle for its last iteration, that tail is
+        // only affordable once the loop is long enough to dilute it to an eighth of the work (observation).
+        const int idle = iters_wide * 2 - iters;
+
+        return idle * 8 <= iters_wide * 2;
+    };
+
+    if (has_ids && ncols_dst > 1) {
+        // Multi-token MUL_MAT_ID path - dedicated MoE kernel
+        mul_mat_vec_q_moe_launch<type>(
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x,
+            stride_row_x, stride_col_y, stride_col_dst,
+            stride_channel_x, stride_channel_y, stride_channel_dst,
+            ncols_dst, ids_stride, warp_size, nchannels_dst, stream);
+        return;
+    }
+
+    switch (ncols_dst) {
+        case 1: {
+            // static, else MSVC lambda capture breaks the constexpr uses below
+            static constexpr int c_ncols_dst = 1;
+
+            // Tag types keep the flags compile-time, so __launch_bounds__ matches what is launched.
+            const auto launch = [&](auto small_k_tag, auto halve_iters_tag) {
+                constexpr bool c_small_k = decltype(small_k_tag)::value;
+                // Types the table does not promote would compile a second, identical kernel.
+                constexpr bool c_promoted =
+                    calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, true) !=
+                    calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, false);
+
+                constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted;
+
+                const std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst + (fusion.shared_up != nullptr),
+                                                                              nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters);
+                mul_mat_vec_q_switch_fusion<type, c_ncols_dst, c_small_k, c_halve_iters>(
+                    vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                    channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd,
+                    stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride,
+                    stream);
+            };
+
+            if (should_use_small_k(c_ncols_dst)) {
+                launch(std::true_type{},  std::false_type{});
+            } else if (should_halve_iters()) {
+                launch(std::false_type{}, std::true_type{});
+            } else {
+                launch(std::false_type{}, std::false_type{});
+            }
+        } break;
+        case 2: {
+            constexpr int c_ncols_dst = 2;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 3: {
+            constexpr int c_ncols_dst = 3;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 4: {
+            constexpr int c_ncols_dst = 4;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 5: {
+            constexpr int c_ncols_dst = 5;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 6: {
+            constexpr int c_ncols_dst = 6;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 7: {
+            constexpr int c_ncols_dst = 7;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        case 8: {
+            constexpr int c_ncols_dst = 8;
+            std::pair<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
+            mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst,
+                 channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst,
+                 dims.first, dims.second, 0, ids_stride, stream);
+        } break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+template <typename dst_t>
+static void mul_mat_vec_q_switch_type(
+        const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst,
+        const int ncols_x, const int nrows_x, const int ncols_dst,
+        const int stride_row_x, const int stride_col_y, const int stride_col_dst,
+        const int nchannels_x, const int nchannels_y, const int nchannels_dst,
+        const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
+        const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
+        const int ids_stride, cudaStream_t stream) {
+    switch (type_x) {
+        case GGML_TYPE_Q1_0:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q1_0>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q2_0:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q2_0>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q4_0:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q4_0>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q4_1:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q4_1>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q5_0:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q5_0>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q5_1:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q5_1>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q8_0:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q8_0>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_MXFP4:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_MXFP4>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_NVFP4:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_NVFP4>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q2_K:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q2_K>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q3_K>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q4_K>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q5_K>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q6_K>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ2_XXS:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ2_XXS>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ2_XS:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ2_XS>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ2_S:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ2_S>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ3_XXS:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ3_XXS>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ1_S:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ1_S>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ1_M:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ1_M>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ4_NL>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ4_XS>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_IQ3_S>
+                (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+                 nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+void ggml_cuda_mul_mat_vec_q(
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion, const void * src1_q8_1_pre) {
+    // FreeToken: activations and results may also be bf16 or fp16
+    GGML_ASSERT(        src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16);
+    GGML_ASSERT(        dst->type  == GGML_TYPE_F32 || dst->type  == GGML_TYPE_BF16 || dst->type  == GGML_TYPE_F16);
+    GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
+
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    cudaStream_t stream = ctx.stream();
+
+    const size_t ts_src0 = ggml_type_size(src0->type);
+    const size_t ts_src1 = ggml_type_size(src1->type);
+    const size_t ts_dst  = ggml_type_size(dst->type);
+
+    GGML_ASSERT(        nb00       == ts_src0);
+    GGML_ASSERT(        nb10       == ts_src1);
+    GGML_ASSERT(        nb0        == ts_dst);
+    GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type));
+
+    GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE);
+
+    const float   * src1_d =       (const float   *) src1->data;
+    const int32_t *  ids_d = ids ? (const int32_t *)  ids->data : nullptr;
+    float         *  dst_d =       (float         *)  dst->data;
+
+    ggml_cuda_mm_fusion_args_device fusion_local{};
+
+    if (fusion) {
+        const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+        GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc));
+        GGML_ASSERT(  ids || dst->ne[1] == 1);
+        // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
+        // non-negligible for some models such as gpt-oss-20b
+        GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4);
+
+        if (fusion->shared_up) {
+            GGML_ASSERT(ids && fusion->gate && fusion->shared_gate && fusion->shared_dst);
+            GGML_ASSERT(!fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
+            GGML_ASSERT(ne11 == 1 && ne03 == 1 && ne13 == 1);
+            GGML_ASSERT(fusion->shared_up->type == src0->type && fusion->shared_gate->type == src0->type);
+            GGML_ASSERT(ggml_are_same_shape(fusion->shared_up, fusion->shared_gate));
+            GGML_ASSERT(ggml_is_contiguous(fusion->shared_up) && ggml_is_contiguous(fusion->shared_gate));
+            GGML_ASSERT(fusion->shared_up->ne[0] == ne00 && fusion->shared_up->ne[1] == ne01);
+            GGML_ASSERT(fusion->shared_up->nb[1] == nb01 && ggml_is_matrix(fusion->shared_up));
+            GGML_ASSERT(fusion->shared_dst->type == GGML_TYPE_F32 && ggml_is_contiguous(fusion->shared_dst));
+            GGML_ASSERT(fusion->shared_dst->ne[0] == ne0 && fusion->shared_dst->ne[1] == ne2);
+            fusion_local.shared_up   = fusion->shared_up->data;
+            fusion_local.shared_gate = fusion->shared_gate->data;
+            fusion_local.shared_dst  = (float *) fusion->shared_dst->data;
+            fusion_local.shared_stride_col_dst = fusion->shared_dst->nb[1] / ts_dst;
+        }
+
+        if (fusion->x_bias) {
+            GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32);
+            GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]);
+            GGML_ASSERT(!ids || fusion->x_bias->ne[1] == src0->ne[2]);
+            fusion_local.x_bias = fusion->x_bias->data;
+        }
+        if (fusion->gate) {
+            GGML_ASSERT(fusion->gate->type == src0->type && ggml_are_same_stride(fusion->gate, src0));
+            fusion_local.gate = fusion->gate->data;
+        }
+        if (fusion->gate_bias) {
+            GGML_ASSERT(fusion->gate_bias->type == GGML_TYPE_F32);
+            GGML_ASSERT(fusion->gate_bias->ne[0] == dst->ne[0]);
+            GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
+            fusion_local.gate_bias = fusion->gate_bias->data;
+        }
+        if (fusion->x_scale) {
+            GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
+            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
+            GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
+            fusion_local.x_scale = fusion->x_scale->data;
+        }
+        if (fusion->gate_scale) {
+            GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
+            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
+            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
+            fusion_local.gate_scale = fusion->gate_scale->data;
+        }
+        fusion_local.glu_op = fusion->glu_op;
+        fusion_local.glu_limit = fusion->glu_limit;
+    }
+
+    // If src0 is a temporary compute buffer, clear any potential padding.
+    if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+        const size_t size_data  = ggml_nbytes(src0);
+        const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
+        if (size_alloc > size_data) {
+            GGML_ASSERT(ggml_is_contiguously_allocated(src0));
+            GGML_ASSERT(!src0->view_src);
+            CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
+        }
+    }
+
+    const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
+    // FreeToken: src1_q8_1_pre is src1 already quantized by quantize_row_q8_1_cuda with this padding and layout
+    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool());
+    const void * src1_q8 = src1_q8_1_pre;
+    if (src1_q8 == nullptr) {
+        src1_q8_1.alloc(ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1);
+        src1_q8 = src1_q8_1.get();
+        const int64_t s11 = src1->nb[1] / ts_src1;
+        const int64_t s12 = src1->nb[2] / ts_src1;
+        const int64_t s13 = src1->nb[3] / ts_src1;
+        if (src1->type == GGML_TYPE_BF16) {
+            quantize_row_q8_1_cuda((const nv_bfloat16 *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
+        } else if (src1->type == GGML_TYPE_F16) {
+            quantize_row_q8_1_cuda((const half *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
+        } else {
+            quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
+        }
+    }
+
+    const int64_t s01 = src0->nb[1] / ts_src0;
+    const int64_t s11 = ne10_padded / QK8_1;
+    const int64_t s1  =  dst->nb[1] / ts_dst;
+    const int64_t s02 = src0->nb[2] / ts_src0;
+    const int64_t s2  =  dst->nb[2] / ts_dst;
+    const int64_t s03 = src0->nb[3] / ts_src0;
+    const int64_t s3  =  dst->nb[3] / ts_dst;
+
+    const int64_t s12 = ne11*s11;
+    const int64_t s13 = ne12*s12;
+
+    // For MUL_MAT_ID the memory layout is different than for MUL_MAT:
+    const int64_t ncols_dst          = ids ? ne2  : ne1;
+    const int64_t nchannels_y        = ids ? ne11 : ne12;
+    const int64_t nchannels_dst      = ids ? ne1  : ne2;
+    const int64_t stride_col_dst     = ids ? s2   : s1;
+    const int64_t stride_col_y       = ids ? s12  : s11;
+    const int64_t stride_channel_dst = ids ? s1   : s2;
+    const int64_t stride_channel_y   = ids ? s11  : s12;
+
+    const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0;
+
+    const auto launch = [&](auto * dst_ptr) {
+        mul_mat_vec_q_switch_type(
+            src0->data, src0->type, src1_q8, ids_d, fusion_local, dst_ptr, ne00,
+            ne01,              ncols_dst,     s01, stride_col_y,     stride_col_dst,
+            ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst,
+            ne03,              ne3,           s03, s13,              s3,               ids_stride, stream);
+    };
+    if (dst->type == GGML_TYPE_BF16) {
+        // FreeToken: the fused shared expert writes an fp32 shared_dst
+        GGML_ASSERT(!fusion || !fusion->shared_up);
+        launch((nv_bfloat16 *) dst->data);
+    } else if (dst->type == GGML_TYPE_F16) {
+        GGML_ASSERT(!fusion || !fusion->shared_up);
+        launch((half *) dst->data);
+    } else {
+        launch(dst_d);
+    }
+}
+
+void ggml_cuda_op_mul_mat_vec_q(
+    ggml_backend_cuda_context & ctx,
+    const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
+    const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
+    const int64_t src1_padded_row_size, cudaStream_t stream) {
+
+    const int64_t ne00 = src0->ne[0];
+    const int64_t row_diff = row_high - row_low;
+
+    const int64_t ne10 = src1->ne[0];
+    GGML_ASSERT(ne10 % QK8_1 == 0);
+
+    const int64_t ne0 = dst->ne[0];
+
+    int id = ggml_cuda_get_device();
+
+    // the main device has a larger memory buffer to hold the results from all GPUs
+    // nrows_dst == nrows of the matrix that the kernel writes into
+    const int64_t nrows_dst = id == ctx.device ? ne0 : row_diff;
+
+    const int stride_row_x = ne00 / ggml_blck_size(src0->type);
+    const int stride_col_y = src1_padded_row_size / QK8_1;
+
+    ggml_cuda_mm_fusion_args_device fusion_local{};
+    mul_mat_vec_q_switch_type(
+        src0_dd_i, src0->type, src1_ddq_i, nullptr, fusion_local, dst_dd_i, ne00, row_diff, src1_ncols, stride_row_x, stride_col_y, nrows_dst,
+        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, stream);
+
+    GGML_UNUSED_VARS(src1, dst, src1_ddf_i, src1_ncols, src1_padded_row_size);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh
new file mode 100644
index 0000000000..6973780c73
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh
@@ -0,0 +1,20 @@
+#include "common.cuh"
+
+#define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels.
+
+bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11);
+
+// Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID,
+// based on the quantization type and GPU architecture (compute capability).
+int get_mmvq_mmid_max_batch(ggml_type type, int cc);
+
+// FreeToken: src1_q8_1_pre skips the internal quantization of src1
+void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx,
+    const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr,
+    const void * src1_q8_1_pre = nullptr);
+
+void ggml_cuda_op_mul_mat_vec_q(
+    ggml_backend_cuda_context & ctx,
+    const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
+    const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
+    const int64_t src1_padded_row_size, cudaStream_t stream);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu
new file mode 100644
index 0000000000..46c9f861a1
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu
@@ -0,0 +1,729 @@
+#include "quantize.cuh"
+#include "convert.cuh"
+#include <cstdint>
+
+#if defined(BLACKWELL_MMA_AVAILABLE)
+// this maps to 256-bit loads in PTX on supported devices,
+// and otherwise falls back to 2 128-bit loads
+struct __builtin_align__(32) float8 {
+    float x; float y; float z; float w;
+    float p; float q; float r; float s;
+};
+
+#if CUDART_VERSION >= 12080
+static __device__ __forceinline__ float nvfp4_native_scale_error(
+        const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) {
+    const float scale_dequant = 2.0f * scale;
+    float err = 0.0f;
+
+#pragma unroll
+    for (int k = 0; k < QK_NVFP4_SUB; k += 4) {
+        const float v0 = vals[k + 0] * inv_col_scale;
+        const float v1 = vals[k + 1] * inv_col_scale;
+        const float v2 = vals[k + 2] * inv_col_scale;
+        const float v3 = vals[k + 3] * inv_col_scale;
+
+        const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale));
+        const __nv_fp4x4_storage_t q_storage = q.__x;
+        const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage);
+        const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U);
+
+        const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1);
+        const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1);
+        const __half2 h2_lo = static_cast<__half2>(hraw2_lo);
+        const __half2 h2_hi = static_cast<__half2>(hraw2_hi);
+        const float2 dq_lo = __half22float2(h2_lo);
+        const float2 dq_hi = __half22float2(h2_hi);
+
+        const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant;
+        const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant;
+        const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant;
+        const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant;
+
+        err = fmaf(err0, err0, err);
+        err = fmaf(err1, err1, err);
+        err = fmaf(err2, err2, err);
+        err = fmaf(err3, err3, err);
+    }
+
+    return err;
+}
+#endif // CUDART_VERSION >= 12080
+#endif // defined(BLACKWELL_MMA_AVAILABLE)
+
+template <typename src_t>
+__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
+static __global__ void quantize_q8_1(
+        const src_t * x_ptr, void * vy_ptr,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const uint32_t ne1, const uint3 ne2) {
+    ggml_cuda_pdl_lc();
+    const src_t * GGML_CUDA_RESTRICT x  = x_ptr;
+    void        * GGML_CUDA_RESTRICT vy = vy_ptr;
+    const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
+
+    if (i0 >= ne0) {
+        return;
+    }
+
+    const int64_t i3 = fastdiv(blockIdx.z, ne2);
+    const int64_t i2 = blockIdx.z - i3*ne2.z;
+    const int64_t i1 = blockIdx.y;
+
+    const int64_t & i00 = i0;
+    const int64_t & i01 = i1;
+    const int64_t & i02 = i2;
+    const int64_t & i03 = i3;
+
+    const int64_t i_cont = ((i3*ne2.z + i2) * ne1 + i1) * ne0 + i0;
+
+    block_q8_1 * y = (block_q8_1 *) vy;
+
+    const int64_t ib  = i_cont / QK8_1; // block index
+    const int64_t iqs = i_cont % QK8_1; // quant index
+
+    ggml_cuda_pdl_sync();
+    const float xi = i0 < ne00 ? ggml_cuda_cast<float>(x[i03*s03 + i02*s02 + i01*s01 + i00]) : 0.0f;
+    float amax = fabsf(xi);
+    float sum = xi;
+
+    amax = warp_reduce_max<QK8_1>(amax);
+    sum  = warp_reduce_sum<QK8_1>(sum);
+
+    const float  d = amax / 127.0f;
+    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
+
+    y[ib].qs[iqs] = q;
+
+    if (iqs > 0) {
+        return;
+    }
+
+    y[ib].ds = make_half2(d, sum);
+}
+
+__device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
+    if (!(amax > 0.0f)) {
+        return 0;
+    }
+
+    // FP4 E2M1: max exponent (unbiased) is 2.
+    constexpr int FP4_E2M1_EMAX = 2;
+
+    const float e = log2f(amax);
+
+    // "even" -> round-to-nearest integer, ties-to-even
+    const int e_int = __float2int_rn(e);
+
+    const int shared_exp = e_int - FP4_E2M1_EMAX;
+
+    int biased = shared_exp + 127;
+
+    biased = max(biased, 0);
+    biased = min(biased, 254);
+
+    return static_cast<uint8_t>(biased);
+}
+
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter, bool use_aligned_float8>
+static __global__ void quantize_mmq_nvfp4(
+        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
+#if defined(BLACKWELL_MMA_AVAILABLE)
+
+    int64_t base_idx;
+    if constexpr (scatter) {
+        base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.y % ne2;
+        const int64_t i3  = blockIdx.y / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_idx = i3 * s03 + i2 * s02 + i01 * s01;
+    }
+    const float * __restrict__  x_row = x + base_idx;
+
+    float amax = 0.0f;
+    if constexpr (use_aligned_float8) {
+        for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) {
+            const float * x_base = x_row + i0;
+            const float8 v = reinterpret_cast<const float8 *>(x_base)[0];
+            amax = fmaxf(amax, fabsf(v.x));
+            amax = fmaxf(amax, fabsf(v.y));
+            amax = fmaxf(amax, fabsf(v.z));
+            amax = fmaxf(amax, fabsf(v.w));
+            amax = fmaxf(amax, fabsf(v.p));
+            amax = fmaxf(amax, fabsf(v.q));
+            amax = fmaxf(amax, fabsf(v.r));
+            amax = fmaxf(amax, fabsf(v.s));
+        }
+    } else {
+        for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) {
+            amax = fmaxf(amax, fabsf(x_row[i0]));
+        }
+    }
+
+    amax = warp_reduce_max<WARP_SIZE>(amax);
+
+    __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE];
+    const int lane = threadIdx.x % WARP_SIZE;
+    const int warp = threadIdx.x / WARP_SIZE;
+
+    if (lane == 0) {
+        warp_amax[warp] = amax;
+    }
+    __syncthreads();
+
+    if (warp == 0) {
+        amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f;
+        amax = warp_reduce_max<WARP_SIZE>(amax);
+        if (lane == 0) {
+            warp_amax[0] = amax / (6.0f * 448.0f);
+            if constexpr (scatter) {
+#pragma unroll
+                for (int slot = 0; slot < n_expert_used; ++slot) {
+                    const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                    scale[i] = warp_amax[0];
+                }
+            } else {
+                scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0];
+            }
+        }
+    }
+    __syncthreads();
+
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB;
+
+    for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) {
+        const int64_t i0_base = isb * QK_NVFP4_SUB;
+        const int64_t k_block = i0_base / QK_FP4_MMQ;
+        const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
+
+        const float row_scale = warp_amax[0];
+        const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f;
+
+        float vals[QK_NVFP4_SUB];
+        if constexpr (use_aligned_float8) {
+            const float * x_base = x_row + i0_base;
+            const float8 v0 = i0_base +  7 < ne00 ? reinterpret_cast<const float8 *>(x_base)[0]     : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast<const float8 *>(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w;
+            vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s;
+            vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w;
+            vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s;
+        } else {
+#pragma unroll
+            for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+                const int64_t i00 = i0_base + k;
+                vals[k] = i00 < ne00 ? x_row[i00] : 0.0f;
+            }
+        }
+
+        uint32_t q0 = 0;
+        uint32_t q1 = 0;
+
+        float amax_sub = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale));
+        }
+
+        static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 };
+        const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f);
+
+        uint8_t fp8_code = (uint8_t) first_fp8_code;
+        float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code);
+        float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+        float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale);
+#else
+        float best_err = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            const float v = vals[k] * inv_col_scale;
+            const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err);
+            const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale;
+            best_err = fmaf(err_diff, err_diff, best_err);
+        }
+#endif // CUDART_VERSION >= 12080
+
+#pragma unroll
+        for (int i = 1; i < 5; ++i) {
+            const int test_code = first_fp8_code + test_offsets[i];
+            if (test_code < 0 || test_code > 0x7e) {
+                continue;
+            }
+
+            const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code);
+            const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+            const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale);
+#else
+            float cur_err = 0.0f;
+#pragma unroll
+            for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+                const float v = vals[k] * inv_col_scale;
+                const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
+                const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale;
+                cur_err = fmaf(err_diff, err_diff, cur_err);
+            }
+#endif // CUDART_VERSION >= 12080
+
+            if (cur_err < best_err) {
+                best_err = cur_err;
+                fp8_code = (uint8_t) test_code;
+                subblock_scale = test_scale;
+            }
+        }
+#if CUDART_VERSION >= 12080
+        const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+        const float s = inv_col_scale * inv_scale;
+
+        __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8]  * s, vals[1] * s, vals[9]  * s));
+        __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s));
+        __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s));
+        __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s));
+
+        const char2 q0_lo_c = *reinterpret_cast<char2 *>(&q0_lo);
+        const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
+        const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
+        const char2 q1_hi_c = *reinterpret_cast<char2 *>(&q1_hi);
+
+        q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24);
+        q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24);
+#else
+        const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k);
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4);
+        }
+#endif // CUDART_VERSION >= 12080
+
+        if constexpr (scatter) {
+#pragma unroll
+            for (int slot = 0; slot < n_expert_used; ++slot) {
+                const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                block_fp4_mmq * yb = y + (k_block * ne1 + i);
+                uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+                yqs[2 * sub + 0] = q0;
+                yqs[2 * sub + 1] = q1;
+                reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+            }
+        } else {
+            const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
+            block_fp4_mmq * yb = y + (blockIdx.y * (blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
+            uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+            yqs[2 * sub + 0] = q0;
+            yqs[2 * sub + 1] = q1;
+            reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+        }
+    }
+#else
+    GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used);
+    NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
+#endif // defined(BLACKWELL_MMA_AVAILABLE)
+
+}
+
+// quantize values in the format mxfp4 is stored which is interleaved nibbles
+// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter>
+static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
+                                          const int32_t * __restrict__ ids,
+                                          void * __restrict__ vy,
+                                          const int64_t ne00,
+                                          const int64_t s01,
+                                          const int64_t s02,
+                                          const int64_t s03,
+                                          const int64_t ne0,
+                                          const int     ne1,
+                                          const int     ne2,
+                                          const int     n_expert_used) {
+    constexpr int vals_per_scale = 32;
+    constexpr int vals_per_warp  = 2 * vals_per_scale;  // Each warp processes 2 blocks of 32 = 64 values
+
+    const int warp_id = threadIdx.y;
+    const int lane_id_32 = threadIdx.x;
+
+    const int nwarps = blockDim.y;
+
+    const int64_t warp_start_offset = (blockIdx.y * nwarps + warp_id) * vals_per_warp;
+
+    if (warp_start_offset >= ne0) {
+        return;
+    }
+
+    const int64_t block_fp4_mmq_size = QK_FP4_MMQ;
+    const int64_t k_block            = warp_start_offset / block_fp4_mmq_size;
+    const int64_t quad_idx_in_block  = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
+
+    const int group_id = lane_id_32 / 4;
+    const int lane_in_group = lane_id_32 % 4;
+    const int base = group_id * 2;
+
+    ggml_cuda_pdl_sync();
+    int64_t base_pos;
+    if constexpr (scatter) {
+        base_pos = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_pos = i3 * s03 + i2 * s02 + i01 * s01;
+    }
+
+    uint8_t scales[2];
+    char2   packed[2];
+
+#pragma unroll
+    for (int b = 0; b < 2; ++b) {
+        const int64_t i0 = warp_start_offset + b * vals_per_scale + lane_id_32;
+        const float xi = (i0 < ne00) ? x[base_pos + i0] : 0.0f;
+
+        float amax = fabsf(xi);
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1) {
+            amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, WARP_SIZE));
+        }
+
+        const uint8_t e = compute_e8m0_scale(amax);
+        scales[b] = e;
+        const float inv_s = (amax == 0.0f) ? 0.0f : __frcp_rn(ggml_cuda_e8m0_to_fp32(e));
+
+#if CUDART_VERSION >= 12080
+        const float scaled_val = xi * inv_s;
+
+        const float val0 = __shfl_sync(0xFFFFFFFF, scaled_val, base, WARP_SIZE);
+        const float val1 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 16, WARP_SIZE);
+        const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
+        const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
+
+        __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
+        packed[b] = *(char2 *) &fp4_packed;
+#else
+        // Fallback: manual FP4 conversion using LUT
+        const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
+
+        const uint8_t q_lo_0 = __shfl_sync(0xFFFFFFFF, q_val, base,      WARP_SIZE);
+        const uint8_t q_lo_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 1,  WARP_SIZE);
+        const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
+        const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
+
+        char2 q;
+        q.x = (q_hi_0 << 4) | q_lo_0;
+        q.y = (q_hi_1 << 4) | q_lo_1;
+        packed[b] = q;
+#endif // CUDART_VERSION >= 12080
+    }
+
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    if constexpr (scatter) {
+#pragma unroll
+        for (int slot = 0; slot < n_expert_used; ++slot) {
+            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            block_fp4_mmq * yb = y + (k_block * ne1 + i);
+            char2 * yqs2 = (char2 *) yb->qs;
+            if (lane_in_group == 0) {
+                yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+                yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+            }
+            if (lane_id_32 == 0) {
+                yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+            }
+        }
+    } else {
+        const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
+        block_fp4_mmq * yb = y + (ib0 + k_block * ne1 + blockIdx.x);
+        char2 * yqs2 = (char2 *) yb->qs;
+        if (lane_in_group == 0) {
+            yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+            yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+        }
+        if (lane_id_32 == 0) {
+            yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+        }
+    }
+    GGML_UNUSED(n_expert_used);
+}
+
+// FreeToken: 4 consecutive activations as float, from fp32, bf16 or fp16
+static __device__ __forceinline__ float4 load_float4(const float * x) {
+    return *(const float4 *) x;
+}
+
+static __device__ __forceinline__ float4 load_float4(const nv_bfloat16 * x) {
+    const uint2 v = *(const uint2 *) x;
+    const float2 lo = __bfloat1622float2(*(const nv_bfloat162 *) &v.x);
+    const float2 hi = __bfloat1622float2(*(const nv_bfloat162 *) &v.y);
+    return make_float4(lo.x, lo.y, hi.x, hi.y);
+}
+
+static __device__ __forceinline__ float4 load_float4(const half * x) {
+    const uint2 v = *(const uint2 *) x;
+    const float2 lo = __half22float2(*(const half2 *) &v.x);
+    const float2 hi = __half22float2(*(const half2 *) &v.y);
+    return make_float4(lo.x, lo.y, hi.x, hi.y);
+}
+
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <mmq_q8_1_ds_layout ds_layout, bool scatter, typename src_t>
+static __global__ void quantize_mmq_q8_1(
+        const src_t * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
+
+    constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
+    constexpr int vals_per_sum   = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
+
+    const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4;
+
+    if (i0 >= ne0) {
+        return;
+    }
+
+    const int64_t i00 = i0;
+    ggml_cuda_pdl_sync();
+
+    int64_t base_idx;
+    if constexpr (scatter) {
+        base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_idx = i3*s03 + i2*s02 + i01*s01;
+    }
+
+    block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
+
+    const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel
+    const int64_t iqs     = i0 % QK8_1_MMQ; // quant index in block
+
+    // Load 4 floats per thread and calculate max. abs. value between them:
+    const float4 xi = i0 < ne00 ? load_float4(x + base_idx + i00) : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float amax = fabsf(xi.x);
+    amax = fmaxf(amax, fabsf(xi.y));
+    amax = fmaxf(amax, fabsf(xi.z));
+    amax = fmaxf(amax, fabsf(xi.w));
+
+    // Exchange max. abs. value between vals_per_scale/4 threads.
+#pragma unroll
+    for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) {
+        amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
+    }
+
+    float sum;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        sum = xi.x + xi.y + xi.z + xi.w;
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
+        }
+    }
+
+    const float d_inv = 127.0f / amax;
+    char4 q;
+    q.x = roundf(xi.x*d_inv);
+    q.y = roundf(xi.y*d_inv);
+    q.z = roundf(xi.z*d_inv);
+    q.w = roundf(xi.w*d_inv);
+    const float d = 1.0f / d_inv;
+
+    // write the block once (normal) or to each of the token's compact rows (scatter)
+    const int nwrite = scatter ? n_expert_used : 1;
+#pragma unroll
+    for (int slot = 0; slot < nwrite; ++slot) {
+        int64_t ib;
+        if constexpr (scatter) {
+            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            ib = k_block*ne1 + i;
+        } else {
+            const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
+            ib = ib0 + k_block*ne1 + blockIdx.x;
+        }
+
+        // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
+        char4 * yqs4 = (char4 *) y[ib].qs;
+        yqs4[iqs/4] = q;
+
+        if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
+            if (iqs % 16 == 0 && iqs < 96) {
+                y[ib].d2s6[2 + iqs/16] = sum;
+                if (iqs % 64 == 0) {
+                    y[ib].d2s6[iqs/64] = d;
+                }
+            }
+        } else if (iqs % 32 == 0) {
+            if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
+                y[ib].ds4[iqs/32] = make_half2(d, sum);
+            } else {
+                y[ib].d4[iqs/32]  = d;
+            }
+        }
+    }
+    GGML_UNUSED(n_expert_used);
+}
+
+template <typename src_t>
+void quantize_row_q8_1_cuda(
+        const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(!ids);
+    GGML_ASSERT(ne0 % QK8_1 == 0);
+
+    const uint3 ne2_fastdiv = init_fastdiv_values(ne2);
+
+    const int64_t block_num_x = (ne0 + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
+    const dim3 num_blocks(block_num_x, ne1, ne2*ne3);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream);
+    ggml_cuda_kernel_launch(quantize_q8_1<src_t>, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
+    GGML_UNUSED(type_src0);
+}
+
+template <typename src_t>
+void quantize_mmq_q8_1_cuda(
+        const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % 4 == 0);
+    GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
+
+    // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid:
+    const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(ne1, block_num_y, ne2*ne3);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    switch (mmq_get_q8_1_ds_layout(type_src0)) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false, src_t>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false, src_t>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_D2S6:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false, src_t>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+// scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
+template <typename src_t>
+void quantize_scatter_mmq_q8_1_cuda(
+        const src_t * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+        const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % 4 == 0);
+    GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
+
+    const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(n_tokens, block_num_y, 1);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    switch (mmq_get_q8_1_ds_layout(type_src0)) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_D2S6:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true, src_t><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+// scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
+void quantize_scatter_mmq_fp4_cuda(
+        const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
+        const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+        const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+    GGML_ASSERT(ne0 > 0);
+    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
+        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+        const dim3 num_blocks(n_tokens, 1, 1);
+        if (use_aligned_float8) {
+            quantize_mmq_nvfp4<true, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        } else {
+            quantize_mmq_nvfp4<true, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        }
+    } else {
+        GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
+        constexpr int nwarps = 8;
+        constexpr int vals_per_block = nwarps * 2 * QK_MXFP4;
+        const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
+        const dim3 block_size(WARP_SIZE, nwarps, 1);
+        const dim3 num_blocks(n_tokens, block_num_y, 1);
+        quantize_mmq_mxfp4<true><<<num_blocks, block_size, 0, stream>>>(
+            x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+    }
+}
+
+void quantize_mmq_fp4_cuda(
+        const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
+    GGML_ASSERT(ne0 > 0);
+
+    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
+        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+        const dim3 num_blocks(ne1, ne2 * ne3, 1);
+        if (use_aligned_float8) {
+            quantize_mmq_nvfp4<false, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        } else {
+            quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        }
+    } else {
+        GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
+
+        constexpr int nwarps = 8;
+        constexpr int vals_per_warp  = 2 * QK_MXFP4;
+        constexpr int vals_per_block = nwarps * vals_per_warp;
+
+        const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
+        const dim3    num_blocks(ne1, block_num_y, ne2 * ne3);
+        const dim3    block_size(WARP_SIZE, nwarps, 1);
+
+        quantize_mmq_mxfp4<false><<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+    }
+}
+
+#define FREETOKEN_QUANTIZE_INSTANCES(src_t)                                                                              \
+    template void quantize_row_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t,    \
+        int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t);                                             \
+    template void quantize_mmq_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t,    \
+        int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t);                                             \
+    template void quantize_scatter_mmq_q8_1_cuda<src_t>(const src_t *, const int32_t *, void *, ggml_type, int64_t,     \
+        int64_t, int64_t, int64_t, int64_t, int, cudaStream_t);
+
+FREETOKEN_QUANTIZE_INSTANCES(float)
+FREETOKEN_QUANTIZE_INSTANCES(nv_bfloat16)
+FREETOKEN_QUANTIZE_INSTANCES(half)
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh
new file mode 100644
index 0000000000..99283ea6e1
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh
@@ -0,0 +1,74 @@
+#pragma once
+
+#include "common.cuh"
+#include "mmq.cuh"
+
+#include <cstdint>
+
+#define CUDA_QUANTIZE_BLOCK_SIZE     256
+#define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128
+
+static_assert(MATRIX_ROW_PADDING %    CUDA_QUANTIZE_BLOCK_SIZE      == 0, "Risk of out-of-bounds access.");
+static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access.");
+
+typedef void (*quantize_cuda_t)(
+        const float * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+// FreeToken: src_t is float, nv_bfloat16 or half
+template <typename src_t>
+void quantize_row_q8_1_cuda(
+        const src_t * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+template <typename src_t>
+void quantize_mmq_q8_1_cuda(
+        const src_t * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+void quantize_mmq_fp4_cuda(const float *   x,
+                             const int32_t * ids,
+                             void *          vy,
+                             float *         scale,
+                             ggml_type       type_src0,
+                             bool            use_aligned_float8,
+                             int64_t         ne00,
+                             int64_t         s01,
+                             int64_t         s02,
+                             int64_t         s03,
+                             int64_t         ne0,
+                             int64_t         ne1,
+                             int64_t         ne2,
+                             int64_t         ne3,
+                             cudaStream_t    stream);
+
+// quantize each token once and scatter the block to its compact rows (via the inverse map)
+void quantize_scatter_mmq_fp4_cuda(const float *   x,
+                                   const int32_t * ids_src1_inv,
+                                   void *          vy,
+                                   float *         scale,
+                                   ggml_type       type_src0,
+                                   bool            use_aligned_float8,
+                                   int64_t         ne00,
+                                   int64_t         stride_token,
+                                   int64_t         ne0,
+                                   int64_t         n_tokens,
+                                   int64_t         nrows_dst,
+                                   int             n_expert_used,
+                                   cudaStream_t    stream);
+
+template <typename src_t>
+void quantize_scatter_mmq_q8_1_cuda(const src_t *   x,
+                                    const int32_t * ids_src1_inv,
+                                    void *          vy,
+                                    ggml_type       type_src0,
+                                    int64_t         ne00,
+                                    int64_t         stride_token,
+                                    int64_t         ne0,
+                                    int64_t         n_tokens,
+                                    int64_t         nrows_dst,
+                                    int             n_expert_used,
+                                    cudaStream_t    stream);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu
new file mode 100644
index 0000000000..84ec850294
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu
new file mode 100644
index 0000000000..583c4e5a51
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ2_S);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu
new file mode 100644
index 0000000000..edaf1560de
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ2_XS);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu
new file mode 100644
index 0000000000..233d9342c9
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu
new file mode 100644
index 0000000000..6092dc7136
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ3_S);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu
new file mode 100644
index 0000000000..1d5bd201fb
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu
new file mode 100644
index 0000000000..eb02fab002
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu
new file mode 100644
index 0000000000..1eb3b74307
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
new file mode 100644
index 0000000000..0c1c7f291b
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
@@ -0,0 +1,7 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_MXFP4);
+
+DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
new file mode 100644
index 0000000000..bed8c8e2e4
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
@@ -0,0 +1,7 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+
+DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu
new file mode 100644
index 0000000000..f0686b0d0d
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q1_0);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu
new file mode 100644
index 0000000000..750180e330
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q2_0);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu
new file mode 100644
index 0000000000..6415369dc1
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q2_K);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu
new file mode 100644
index 0000000000..ffb6213af8
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q3_K);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu
new file mode 100644
index 0000000000..0c0b0c8a8e
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q4_0);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu
new file mode 100644
index 0000000000..ee67f6942a
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q4_1);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu
new file mode 100644
index 0000000000..9eeb3cd7f3
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q4_K);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu
new file mode 100644
index 0000000000..cc57fb9753
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q5_0);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu
new file mode 100644
index 0000000000..721ac790c4
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q5_1);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu
new file mode 100644
index 0000000000..a2e90ffd5d
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q5_K);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu
new file mode 100644
index 0000000000..470938fef8
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q6_K);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu
new file mode 100644
index 0000000000..974477bbb7
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu
@@ -0,0 +1,5 @@
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q8_0);
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh
new file mode 100644
index 0000000000..04f3af6443
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh
@@ -0,0 +1,123 @@
+#pragma once
+#include "common.cuh"
+
+#define CUDA_NEG_BLOCK_SIZE 256
+#define CUDA_STEP_BLOCK_SIZE 256
+#define CUDA_GELU_BLOCK_SIZE 256
+#define CUDA_SILU_BLOCK_SIZE 256
+#define CUDA_SILU_BACK_BLOCK_SIZE 256
+#define CUDA_TANH_BLOCK_SIZE 256
+#define CUDA_RELU_BLOCK_SIZE 256
+#define CUDA_SIGMOID_BLOCK_SIZE 256
+#define CUDA_HARDSIGMOID_BLOCK_SIZE 256
+#define CUDA_EXP_BLOCK_SIZE 256
+#define CUDA_HARDSWISH_BLOCK_SIZE 256
+#define CUDA_SQR_BLOCK_SIZE 256
+#define CUDA_SQRT_BLOCK_SIZE 256
+#define CUDA_SIN_BLOCK_SIZE 256
+#define CUDA_COS_BLOCK_SIZE 256
+#define CUDA_GLU_BLOCK_SIZE 256
+#define CUDA_XIELU_BLOCK_SIZE 256
+
+void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_neg(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_silu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_silu_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_tanh(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_exp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_hardswish(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_leaky_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sqrt(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sin(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_cos(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_log(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_elu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_floor(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_ceil(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_swiglu_clamp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node);
+
+void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node);
+
+__device__ __forceinline__ float ggml_cuda_op_silu_single(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__device__ __forceinline__ float ggml_cuda_op_gelu_single(float x) {
+    const float GELU_COEF_A    = 0.044715f;
+    const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
+
+    return 0.5f * x * (1.0f + tanhf(SQRT_2_OVER_PI * x * (1.0f + GELU_COEF_A * x * x)));
+}
+
+__device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) {
+    x = fminf(x, limit);
+    g = fmaxf(fminf(g, limit), -limit);
+
+    float out_glu = x / (1.0f + expf(-x * alpha));
+    out_glu = out_glu * (1.0f + g);
+    return out_glu;
+}
+
+__device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, float up, float limit) {
+    gate = fminf(gate, limit);
+    up = fmaxf(fminf(up, limit), -limit);
+
+    return ggml_cuda_op_silu_single(gate) * up;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh
new file mode 100644
index 0000000000..f65fe530eb
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh
@@ -0,0 +1,1380 @@
+#pragma once
+
+#include "common.cuh"
+
+#include <cstdint>
+
+static __device__ __forceinline__ int get_int_b1(const void * x, const int & i32) {
+    const uint8_t * x8 = (const uint8_t *) x;
+
+    int x32  = x8[4*i32 + 0] <<  0;
+    x32     |= x8[4*i32 + 1] <<  8;
+    x32     |= x8[4*i32 + 2] << 16;
+    x32     |= x8[4*i32 + 3] << 24;
+
+    return x32;
+}
+
+static __device__ __forceinline__ int get_int_b2(const void * x, const int & i32) {
+    const uint16_t * x16 = (const uint16_t *) x; // assume at least 2 byte alignment
+
+    int x32  = x16[2*i32 + 0] <<  0;
+    x32     |= x16[2*i32 + 1] << 16;
+
+    return x32;
+}
+
+static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32) {
+    return ((const int *) x)[i32]; // assume at least 4 byte alignment
+}
+
+// q4 contains 8 indices with 4 bit each.
+// This function selects those bytes from table that are at those indices and returns them as int2.
+// The first int contains the bytes with even indices in q4, the second int contains the bytes with odd indices in q4.
+static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, const int8_t * table) {
+#if defined(GGML_USE_HIP)
+    // Load the 16-byte table into four 32-bit unsigned integers.
+    const uint32_t *values = (const uint32_t *)table;
+
+    const uint32_t q_even = q4;
+    const uint32_t q_odd  = (q4 >> 4);
+
+    // Perform lookups in the lower half of the table (indices 0-7).
+    uint32_t v_even_low = __builtin_amdgcn_perm(values[1], values[0], q_even & 0x07070707);
+    uint32_t v_odd_low = __builtin_amdgcn_perm(values[1], values[0], q_odd & 0x07070707);
+
+    // Perform lookups in the upper half of the table (indices 8-15).
+    uint32_t v_even_high = __builtin_amdgcn_perm(values[3], values[2], q_even & 0x07070707);
+    uint32_t v_odd_high = __builtin_amdgcn_perm(values[3], values[2], q_odd & 0x07070707);
+
+    // Select between the low and high results based on the MSB of each index nibble.
+    uint32_t mask_even = 0x03020100 | ((q_even & 0x08080808) >> 1);
+    uint32_t res_x = __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even);
+    uint32_t mask_odd = 0x03020100 | ((q_odd & 0x08080808) >> 1);
+    uint32_t res_y = __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd);
+
+    return make_int2(res_x, res_y);
+#elif !defined(GGML_USE_MUSA)
+    // CUDA does not have an instruction for selecting bytes with 4 bit indices.
+    // However, __byte_perm is an instruction that selects bytes with 3 bit indices that can be used instead.
+    const uint32_t * table32 = (const uint32_t *) table;
+
+    // __byte_perm selects bytes based on the lower 16 bits in its third argument.
+    // Therefore, do 2 iterations over the 32 bits in q4 with 0 and 16 shift.
+    // To handle the fourth bit, first call _byte_perm both for the low and the high 64 bit of table, using the low 3 bits.
+    // Then, call __byte_perm again to select from the low and high bytes based on the fourth bit.
+    uint32_t tmp[2];
+    const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1));
+#pragma unroll
+    for (uint32_t i = 0; i < 2; ++i) {
+        const uint32_t shift = 16 * i;
+
+        const uint32_t low  = __byte_perm(table32[0], table32[1], q4 >> shift);
+        const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift);
+        tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift);
+    }
+
+    // tmp contains the bytes from tyble in the same order as the 4 bit indices in q4.
+    // However, for the result we need ints with all even/odd 4 bit indices in q4.
+    // Therefore, 2 more calls to __byte_perm to put the bytes in the correct order.
+    return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531));
+#else
+    // Generic implementation.
+    const int      q0_32  = (q4 >> 0) & 0x0F0F0F0F;
+    const int8_t * q0_8   = (const int8_t *) &q0_32;
+    const char4    val0_8 = make_char4(
+        table[q0_8[0]], table[q0_8[1]], table[q0_8[2]], table[q0_8[3]]);
+
+    const int      q1_32  = (q4 >> 4) & 0x0F0F0F0F;
+    const int8_t * q1_8   = (const int8_t *) &q1_32;
+    const char4    val1_8 = make_char4(
+        table[q1_8[0]], table[q1_8[1]], table[q1_8[2]], table[q1_8[3]]);
+
+    return make_int2(*((const int *) &val0_8), *((const int *) &val1_8));
+#endif
+}
+
+static __device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) {
+    // v is a 7 bit int, with the 8th sign being encodable as popcnt
+    // with xor we can "correct" the bit instead of having to mask
+    const uint32_t p = __popc(v) & 1;
+    const uint32_t s = v ^ p << 7;
+    // broadcast over uint to allow for 0x08040201 / 0x80402010 as selectors
+    return s * 0x01010101;
+}
+
+// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called
+// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q
+
+#define VDR_Q1_0_Q8_1_MMVQ 1  // Process one 32-element chunk at a time for parallelism
+#define VDR_Q1_0_Q8_1_MMQ  4  // Q1_0 has 128 bits (4 ints) per block
+
+#define VDR_Q2_0_Q8_1_MMVQ 1  // Process one 32-element chunk at a time for parallelism
+#define VDR_Q2_0_Q8_1_MMQ  2  // Q2_0 group 64: 128 bits (4 ints) per block, 2 32-element chunks
+
+#define VDR_Q4_0_Q8_1_MMVQ 2
+#define VDR_Q4_0_Q8_1_MMQ  4
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl(
+    const int * v, const int * u, const float & d4, const half2 & ds8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        const int vi0 = (v[i] >> 0) & 0x0F0F0F0F;
+        const int vi1 = (v[i] >> 4) & 0x0F0F0F0F;
+
+        // SIMD dot product of quantized values
+        sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi);
+        sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi);
+    }
+
+    const float2 ds8f = __half22float2(ds8);
+
+    // second part effectively subtracts 8 from each quant value
+    return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y);
+}
+
+#define VDR_Q4_1_Q8_1_MMVQ 2
+#define VDR_Q4_1_Q8_1_MMQ  4
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl(
+    const int * v, const int * u, const half2 & dm4, const half2 & ds8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        const int vi0 = (v[i] >> 0) & 0x0F0F0F0F;
+        const int vi1 = (v[i] >> 4) & 0x0F0F0F0F;
+
+        // SIMD dot product of quantized values
+        sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi);
+        sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi);
+    }
+
+#ifdef FAST_FP16_AVAILABLE
+    const float2 tmp = __half22float2(__hmul2(dm4, ds8));
+    const float d4d8 = tmp.x;
+    const float m4s8 = tmp.y;
+#else
+    const float2 dm4f = __half22float2(dm4);
+    const float2 ds8f = __half22float2(ds8);
+    const float d4d8 = dm4f.x * ds8f.x;
+    const float m4s8 = dm4f.y * ds8f.y;
+#endif // FAST_FP16_AVAILABLE
+
+    // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it
+    return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1));
+}
+
+#define VDR_Q5_0_Q8_1_MMVQ 2
+#define VDR_Q5_0_Q8_1_MMQ  4
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl(
+    const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        int vi0 = (vl[i] >>  0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits
+        vi0    |= (vh[i] <<  4) & 0x00000010; // 0 ->  4
+        vi0    |= (vh[i] << 11) & 0x00001000; // 1 -> 12
+        vi0    |= (vh[i] << 18) & 0x00100000; // 2 -> 20
+        vi0    |= (vh[i] << 25) & 0x10000000; // 3 -> 28
+        sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values
+
+        int vi1 = (vl[i] >>  4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits
+        vi1    |= (vh[i] >> 12) & 0x00000010; // 16 ->  4
+        vi1    |= (vh[i] >>  5) & 0x00001000; // 17 -> 12
+        vi1    |= (vh[i] <<  2) & 0x00100000; // 18 -> 20
+        vi1    |= (vh[i] <<  9) & 0x10000000; // 19 -> 28
+        sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values
+    }
+
+    const float2 ds8f = __half22float2(ds8);
+
+    // second part effectively subtracts 16 from each quant value
+    return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y);
+}
+
+#define VDR_Q5_1_Q8_1_MMVQ 2
+#define VDR_Q5_1_Q8_1_MMQ  4
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl(
+    const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        int vi0 = (vl[i] >>  0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits
+        vi0    |= (vh[i] <<  4) & 0x00000010; // 0 ->  4
+        vi0    |= (vh[i] << 11) & 0x00001000; // 1 -> 12
+        vi0    |= (vh[i] << 18) & 0x00100000; // 2 -> 20
+        vi0    |= (vh[i] << 25) & 0x10000000; // 3 -> 28
+        sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values
+
+        int vi1 = (vl[i] >>  4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits
+        vi1    |= (vh[i] >> 12) & 0x00000010; // 16 ->  4
+        vi1    |= (vh[i] >>  5) & 0x00001000; // 17 -> 12
+        vi1    |= (vh[i] <<  2) & 0x00100000; // 18 -> 20
+        vi1    |= (vh[i] <<  9) & 0x10000000; // 19 -> 28
+        sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values
+    }
+
+#ifdef FAST_FP16_AVAILABLE
+    const float2 tmp = __half22float2(__hmul2(dm5, ds8));
+    const float d5d8 = tmp.x;
+    const float m5s8 = tmp.y;
+#else
+    const float2 dm5f = __half22float2(dm5);
+    const float2 ds8f = __half22float2(ds8);
+    const float d5d8 = dm5f.x * ds8f.x;
+    const float m5s8 = dm5f.y * ds8f.y;
+#endif // FAST_FP16_AVAILABLE
+
+    // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it
+    return sumi*d5d8 + m5s8 / (QI5_1 / vdr);
+}
+
+#define VDR_Q8_0_Q8_1_MMVQ 2
+#define VDR_Q8_0_Q8_1_MMQ 8
+
+template <typename T, int vdr> static __device__ __forceinline__ T vec_dot_q8_0_q8_1_impl(
+    const int * v, const int * u, const T & d8_0, const T & d8_1) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        // SIMD dot product of quantized values
+        sumi = ggml_cuda_dp4a(v[i], u[i], sumi);
+    }
+
+    return d8_0*d8_1 * ((T) sumi);
+}
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl(
+    const int * v, const int * u, const half2 & dm8, const half2 & ds8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+        // SIMD dot product of quantized values
+        sumi = ggml_cuda_dp4a(v[i], u[i], sumi);
+    }
+
+#ifdef FAST_FP16_AVAILABLE
+    const float2 tmp = __half22float2(__hmul2(dm8, ds8));
+    const float d8d8 = tmp.x;
+    const float m8s8 = tmp.y;
+#else
+    const float2 dm8f = __half22float2(dm8);
+    const float2 ds8f = __half22float2(ds8);
+    const float d8d8 = dm8f.x * ds8f.x;
+    const float m8s8 = dm8f.y * ds8f.y;
+#endif // FAST_FP16_AVAILABLE
+
+    // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it
+    return sumi*d8d8 + m8s8 / (QI8_1 / vdr);
+}
+
+template <int vdr> static __device__ __forceinline__ float vec_dot_q8_0_16_q8_1_impl(
+    const int * v, const int * u, const float * d8_0, const float & d8_1) {
+
+    float sumf = 0.0f;
+
+#pragma unroll
+    for (int i0 = 0; i0 < vdr; i0 += QI8_0/2) {
+        int sumi = 0;
+
+#pragma unroll
+        for (int i = i0; i < i0 + QI8_0/2; ++i) {
+            // SIMD dot product of quantized values
+            sumi = ggml_cuda_dp4a(v[i], u[i], sumi);
+        }
+
+        sumf += d8_0[i0/(QI8_0/2)]*sumi;
+    }
+
+    return d8_1*sumf;
+}
+
+#define VDR_MXFP4_Q8_1_MMVQ 2
+#define VDR_MXFP4_Q8_1_MMQ  4
+
+static __device__ __forceinline__ float vec_dot_mxfp4_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_mxfp4 * bq4 = (const block_mxfp4 *) vbq + kbx;
+
+    const int * q8 = (const int *) bq8_1->qs + iqs;
+
+    int sumi = 0;
+#pragma unroll
+    for (int l = 0; l < VDR_MXFP4_Q8_1_MMVQ; ++l) {
+        const int aux_q4 = get_int_b1(bq4->qs, iqs + l);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4);
+
+        sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi);
+        sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi);
+    }
+
+    const float d = ggml_cuda_e8m0_to_fp32(bq4->e) * 0.5f * __low2float(bq8_1->ds);
+    return d * sumi;
+}
+
+#define VDR_NVFP4_Q8_1_MMVQ 4
+#define VDR_NVFP4_Q8_1_MMQ  8
+
+static __device__ __forceinline__ float vec_dot_nvfp4_q8_1(
+                                        const void * __restrict__ vbq,
+                                        const block_q8_1 * __restrict__ bq8_1,
+                                        const int32_t & kbx,
+                                        const int32_t & iqs) {
+
+    const block_nvfp4 * bq4 = (const block_nvfp4 *) vbq + kbx;
+    float sum = 0.0f;
+#pragma unroll
+    for (int i = 0; i < VDR_NVFP4_Q8_1_MMVQ/2; i++) {
+        const int32_t iqs0 = iqs + 2*i;
+        const int32_t iqs1 = iqs0 + 1;
+        const int32_t is = iqs0 >> 1;
+        const int2 v0 = get_int_from_table_16(get_int_b4(bq4->qs, iqs0), kvalues_mxfp4);
+        const int2 v1 = get_int_from_table_16(get_int_b4(bq4->qs, iqs1), kvalues_mxfp4);
+        const block_q8_1 * bq8 = bq8_1 + (is >> 1);
+        const int32_t i8 = ((is & 1) << 2);
+
+        int sumi = ggml_cuda_dp4a(v0.x, get_int_b4(bq8->qs, i8 + 0), 0);
+        sumi = ggml_cuda_dp4a(v0.y, get_int_b4(bq8->qs, i8 + 2), sumi);
+        sumi = ggml_cuda_dp4a(v1.x, get_int_b4(bq8->qs, i8 + 1), sumi);
+        sumi = ggml_cuda_dp4a(v1.y, get_int_b4(bq8->qs, i8 + 3), sumi);
+
+        const float d = ggml_cuda_ue4m3_to_fp32(bq4->d[is]) * __low2float(bq8->ds);
+        sum += d * float(sumi);
+    }
+
+    return sum;
+}
+#define VDR_Q2_K_Q8_1_MMVQ 1
+#define VDR_Q2_K_Q8_1_MMQ  4
+
+// contiguous v/x values
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
+    const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales,
+    const half2 & dm2, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR2_K; ++i) {
+        const int sc = scales[2*i];
+
+        const int vi = (v >> (2*i)) & 0x03030303;
+
+        sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product
+
+        // fill int with 4x m
+        int m = sc >> 4;
+        m |= m <<  8;
+        m |= m << 16;
+        sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values
+    }
+
+    const float2 dm2f = __half22float2(dm2);
+
+    return dm2f.x*sumf_d - dm2f.y*sumf_m;
+}
+
+// contiguous v/x + u/y values
+template <int ns8>
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq(
+    const int * __restrict__ v, const int * __restrict__ u, const half2 * dm2, const float & d8, const half2 * s8) {
+
+    float sumf    = 0.0f;
+    float sumf_d8 = 0.0f;
+
+#pragma unroll
+    for (int i0 = 0; i0 < QR2_K*VDR_Q2_K_Q8_1_MMQ; i0 += QI8_1) {
+        const float2 dm2f0 = __half22float2(dm2[i0/(QI8_1/2) + 0]);
+        int sumi_d0 = 0;
+
+        const float2 dm2f1 = __half22float2(dm2[i0/(QI8_1/2) + 1]);
+        int sumi_d1 = 0;
+
+#pragma unroll
+        for (int i = i0; i < i0 + QI8_1/2; ++i) {
+            sumi_d0 = ggml_cuda_dp4a(v[i], u[i], sumi_d0);
+        }
+        sumf_d8 += dm2f0.x * sumi_d0;
+
+#pragma unroll
+        for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) {
+            sumi_d1 = ggml_cuda_dp4a(v[i], u[i], sumi_d1);
+        }
+        sumf_d8 += dm2f1.x * sumi_d1;
+
+        if (i0/QI8_1 < ns8) {
+            const float2 s8f = __half22float2(s8[i0/QI8_1]);
+            sumf -= dm2f0.y*s8f.x;
+            sumf -= dm2f1.y*s8f.y;
+        } else {
+            int sumi_m0 = 0;
+#pragma unroll
+            for (int i = i0; i < i0 + QI8_1/2; ++i) {
+                sumi_m0 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m0);
+            }
+            sumf_d8 -= dm2f0.y * sumi_m0;
+
+            int sumi_m1 = 0;
+#pragma unroll
+            for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) {
+                sumi_m1 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m1);
+            }
+            sumf_d8 -= dm2f1.y * sumi_m1;
+        }
+    }
+
+    return sumf + d8*sumf_d8;
+}
+
+#define VDR_Q3_K_Q8_1_MMVQ 1
+#define VDR_Q3_K_Q8_1_MMQ  2
+
+// contiguous v/x values
+static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq(
+    const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales,
+    const int & scale_offset, const float & d3, const float * __restrict__ d8) {
+
+    float sumf = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR3_K; ++i) {
+        const int isc = scale_offset + 2*i;
+
+        const int isc_low = isc % (QK_K/32);
+        const int sc_shift_low = 4 * (isc / (QK_K/32));
+        const int sc_low  = (scales[isc_low] >> sc_shift_low) & 0xF;
+
+        const int isc_high = isc % (QK_K/64);
+        const int sc_shift_high = 2 * (isc / (QK_K/64));
+        const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4;
+
+        const int sc = (sc_low | sc_high) - 32;
+
+        const int vil = (vl >> (2*i)) & 0x03030303;
+
+        const int vih = ((vh >> i) << 2) & 0x04040404;
+
+        const int vi = __vsub4(vil, vih);
+
+        sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
+    }
+
+    return d3 * sumf;
+}
+
+// contiguous v/x + u/y values
+static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq(
+    const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales,
+    const float & d3, const float & d8) {
+
+    int sumi = 0;
+
+#pragma unroll
+    for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) {
+        int sumi_sc = 0;
+
+#pragma unroll
+        for (int i = i0; i < i0 + QI8_1/2; ++i) {
+            sumi_sc = ggml_cuda_dp4a(v[i], u[i], sumi_sc); // SIMD dot product
+        }
+
+        sumi += sumi_sc * scales[i0 / (QI8_1/2)];
+    }
+
+    return d3*d8 * sumi;
+}
+
+#define VDR_Q4_K_Q8_1_MMVQ 2
+#define VDR_Q4_K_Q8_1_MMQ  8
+
+// contiguous v/x values
+static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq(
+    const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc,
+    const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR4_K; ++i) {
+        const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F;
+        const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F;
+
+        const int dot1 = ggml_cuda_dp4a(v1i, u[2*i+1], ggml_cuda_dp4a(v0i, u[2*i+0], 0)); // SIMD dot product
+        const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+1], ggml_cuda_dp4a(0x01010101, u[2*i+0], 0)); // sum of u
+
+        sumf_d += d8[i] * (dot1 * sc[i]);
+        sumf_m += d8[i] * (dot2 * m[i]);  // multiply constant part of q4_K with sum of q8_1 values
+    }
+
+    const float2 dm4f = __half22float2(dm4);
+
+    return dm4f.x*sumf_d - dm4f.y*sumf_m;
+}
+
+// contiguous v/x + u/y values
+static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq(
+    const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc,
+    const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) {
+        int sumi_d = 0;
+
+#pragma unroll
+        for (int j = 0; j < QI8_1; ++j) {
+            sumi_d = ggml_cuda_dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product
+        }
+
+        const float2 ds8f = __half22float2(ds8[i]);
+
+        sumf_d += ds8f.x * (sc[i] * sumi_d);
+        sumf_m += ds8f.y *   m[i]; // sum of q8_1 block * q4_K min val
+    }
+
+    const float2 dm4f = __half22float2(dm4);
+
+    return dm4f.x*sumf_d - dm4f.y*sumf_m;
+}
+
+#define VDR_Q5_K_Q8_1_MMVQ 2
+#define VDR_Q5_K_Q8_1_MMQ  8
+
+// contiguous v/x values
+static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq(
+    const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc,
+    const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR5_K; ++i) {
+        const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F;
+        const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F;
+
+        const int vh0i = ((vh[0] >> i) << 4) & 0x10101010;
+        const int vh1i = ((vh[1] >> i) << 4) & 0x10101010;
+
+        const int v0i = vl0i | vh0i;
+        const int v1i = vl1i | vh1i;
+
+        const int dot1 = ggml_cuda_dp4a(v0i, u[2*i+0], ggml_cuda_dp4a(v1i, u[2*i+1], 0)); // SIMD dot product
+        const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+0], ggml_cuda_dp4a(0x01010101, u[2*i+1], 0)); // sum of u
+
+        sumf_d += d8[i] * (dot1 * sc[i]);
+        sumf_m += d8[i] * (dot2 * m[i]);
+
+    }
+
+    const float2 dm5f = __half22float2(dm5);
+
+    return dm5f.x*sumf_d - dm5f.y*sumf_m;
+}
+
+// contiguous v/x + u/y values
+static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq(
+    const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc,
+    const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) {
+        int sumi_d = 0;
+
+#pragma unroll
+        for (int j = 0; j < QI8_1; ++j) {
+            sumi_d = ggml_cuda_dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product
+        }
+
+        const float2 ds8f = __half22float2(ds8[i]);
+
+        sumf_d += ds8f.x * (sc[i] * sumi_d);
+        sumf_m += ds8f.y *   m[i]; // sum of q8_1 block * q4_K min val
+    }
+
+    const float2 dm4f = __half22float2(dm4);
+
+    return dm4f.x*sumf_d - dm4f.y*sumf_m;
+}
+
+#define VDR_Q6_K_Q8_1_MMVQ 1
+#define VDR_Q6_K_Q8_1_MMQ  8
+
+// contiguous v/x values
+static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq(
+    const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales,
+    const float & d, const float * __restrict__ d8) {
+
+    float sumf = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR6_K; ++i) {
+        const int sc = scales[4*i];
+
+        const int vil = (vl >> (4*i)) & 0x0F0F0F0F;
+
+        const int vih = ((vh >> (4*i)) << 4) & 0x30303030;
+
+        const int vi = __vsub4((vil | vih), 0x20202020); // vi = (vil | vih) - 32
+
+        sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
+    }
+
+    return d*sumf;
+}
+
+// contiguous v/x + u/y values
+static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq(
+    const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc,
+    const float & d6, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+
+    const int      sc_packed = get_int_b4(sc, 0);
+    const int8_t * sc_reg    = (const int8_t *) &sc_packed;
+
+#pragma unroll
+    for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) {
+        int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale
+
+#pragma unroll
+        for (int i = i0; i < i0 + 2; ++i) {
+            sumi_d.x = ggml_cuda_dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product
+            sumi_d.x = ggml_cuda_dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product
+
+            sumi_d.y = ggml_cuda_dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product
+            sumi_d.y = ggml_cuda_dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product
+        }
+
+        sumf_d += d8[i0/4] * (sc_reg[i0/2+0]*sumi_d.x + sc_reg[i0/2+1]*sumi_d.y);
+    }
+
+    return d6 * sumf_d;
+}
+
+static __device__ __forceinline__ float vec_dot_q1_0_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q1_0 * bq1_0 = (const block_q1_0 *) vbq + kbx;
+
+    // Q1_0: 128 elements with ONE scale
+    // Q8_1: 32 elements per block with individual scales
+    // iqs selects which of the 4 chunks of 32 elements to process (0-3)
+
+    const float     d1 = bq1_0->d;
+    const int16_t * qs = (const int16_t *) bq1_0->qs + iqs * 2;
+
+    // Process only the chunk specified by iqs
+    const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
+
+    int sumi = 0;
+#pragma unroll
+    for (int j = 0; j < 2; ++j) {
+        const int q  = qs[j];
+
+        const int u0 = get_int_b4(bq8_1_chunk->qs, j*4+0);
+        const int u1 = get_int_b4(bq8_1_chunk->qs, j*4+1);
+        const int u2 = get_int_b4(bq8_1_chunk->qs, j*4+2);
+        const int u3 = get_int_b4(bq8_1_chunk->qs, j*4+3);
+
+        // unpack crumbs into nibble indices
+        const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
+        const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+        // unpack nibbles into byte values
+        const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >>  0);
+        const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >>  0);
+        const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+        const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+        // unshuffle values
+        const int v0 = __byte_perm(s0, s1, 0x5410);
+        const int v1 = __byte_perm(s0, s1, 0x7632);
+        const int v2 = __byte_perm(s2, s3, 0x5410);
+        const int v3 = __byte_perm(s2, s3, 0x7632);
+
+        sumi = ggml_cuda_dp4a(v0, u0, sumi);
+        sumi = ggml_cuda_dp4a(v1, u1, sumi);
+        sumi = ggml_cuda_dp4a(v2, u2, sumi);
+        sumi = ggml_cuda_dp4a(v3, u3, sumi);
+    }
+
+    // Apply Q1_0's single scale and this chunk's Q8_1 scale
+    const float d8 = __low2float(bq8_1_chunk->ds);
+    return d1 * d8 * sumi;
+}
+
+static __device__ __forceinline__ float vec_dot_q2_0_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q2_0 * bq2_0 = (const block_q2_0 *) vbq + kbx;
+
+    // Q2_0 (group 64): 64 elements with ONE scale, 2 bits per element (4 elements per byte)
+    // Q8_1: 32 elements per block with individual scales
+    // iqs selects which of the 2 chunks of 32 elements to process (0-1)
+
+    const float     d2 = bq2_0->d;
+    const int16_t * qs = (const int16_t *) bq2_0->qs + iqs * 4;
+
+    // Process only the chunk specified by iqs
+    const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
+
+    int sumi = 0;
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        const int q  = qs[j];
+        const int u  = get_int_b4(bq8_1_chunk->qs, j*2+0);
+        const int v  = get_int_b4(bq8_1_chunk->qs, j*2+1);
+
+#if defined(GGML_USE_HIP)
+        const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18);
+        const uint32_t qy_bits    = q >> 8;
+        const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18);
+        const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices);
+        const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices);
+#else
+        // unpack even and odd crumbs into byte values
+        const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
+        const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2);
+        // unshuffle values
+        const int qx = __byte_perm(qe, qo, 0x5140);
+        const int qy = __byte_perm(qe, qo, 0x7362);
+#endif // defined(GGML_USE_HIP)
+
+        sumi = ggml_cuda_dp4a(u, qx, sumi);
+        sumi = ggml_cuda_dp4a(v, qy, sumi);
+    }
+
+    // Apply Q2_0's single scale and this chunk's Q8_1 scale
+    const float d8 = __low2float(bq8_1_chunk->ds);
+    return d2 * d8 * sumi;
+}
+
+static __device__ __forceinline__ float vec_dot_q4_0_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq + kbx;
+
+    int v[VDR_Q4_0_Q8_1_MMVQ];
+    int u[2*VDR_Q4_0_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) {
+        v[i]     = get_int_b2(bq4_0->qs, iqs + i);
+        u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i);
+        u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_0);
+    }
+
+    return vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMVQ>(v, u, bq4_0->d, bq8_1->ds);
+}
+
+
+static __device__ __forceinline__ float vec_dot_q4_1_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq + kbx;
+
+    int v[VDR_Q4_1_Q8_1_MMVQ];
+    int u[2*VDR_Q4_1_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) {
+        v[i]     = get_int_b4(bq4_1->qs, iqs + i);
+        u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i);
+        u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_1);
+    }
+
+    return vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMVQ>(v, u, bq4_1->dm, bq8_1->ds);
+}
+
+static __device__ __forceinline__ float vec_dot_q5_0_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq + kbx;
+
+    int vl[VDR_Q5_0_Q8_1_MMVQ];
+    int vh[VDR_Q5_0_Q8_1_MMVQ];
+    int  u[2*VDR_Q5_0_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) {
+        vl[i]    = get_int_b2(bq5_0->qs, iqs + i);
+        vh[i]    = get_int_b2(bq5_0->qh, 0) >> (4 * (iqs + i));
+        u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i);
+        u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_0);
+    }
+
+    return vec_dot_q5_0_q8_1_impl<VDR_Q5_0_Q8_1_MMVQ>(vl, vh, u, bq5_0->d, bq8_1->ds);
+}
+
+static __device__ __forceinline__ float vec_dot_q5_1_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq + kbx;
+
+    int vl[VDR_Q5_1_Q8_1_MMVQ];
+    int vh[VDR_Q5_1_Q8_1_MMVQ];
+    int  u[2*VDR_Q5_1_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) {
+        vl[i]    = get_int_b4(bq5_1->qs, iqs + i);
+        vh[i]    = get_int_b4(bq5_1->qh, 0) >> (4 * (iqs + i));
+        u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i);
+        u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_1);
+    }
+
+    return vec_dot_q5_1_q8_1_impl<VDR_Q5_1_Q8_1_MMVQ>(vl, vh, u, bq5_1->dm, bq8_1->ds);
+}
+
+static __device__ __forceinline__ float vec_dot_q8_0_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq + kbx;
+
+    int v[VDR_Q8_0_Q8_1_MMVQ];
+    int u[VDR_Q8_0_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) {
+        v[i] = get_int_b2(bq8_0->qs, iqs + i);
+        u[i] = get_int_b4(bq8_1->qs, iqs + i);
+    }
+
+    return vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMVQ>(v, u, bq8_0->d, __low2half(bq8_1->ds));
+}
+
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx;
+
+    const int bq8_offset = QR2_K * (iqs / QI8_1);
+    const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
+
+    const uint8_t * scales = bq2_K->scales + scale_offset;
+
+    const int v = get_int_b4(bq2_K->qs, iqs);
+    int    u[QR2_K];
+    float d8[QR2_K];
+
+#pragma unroll
+    for (int i = 0; i < QR2_K; ++ i) {
+        u[i]  = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
+        d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
+    }
+
+    return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8);
+}
+
+static __device__ __forceinline__ float vec_dot_q3_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q3_K * bq3_K = (const block_q3_K *) vbq + kbx;
+
+    const int bq8_offset = QR3_K * (iqs / (QI3_K/2));
+    const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
+
+    const float d = bq3_K->d;
+
+    const int vl = get_int_b2(bq3_K->qs, iqs);
+
+    // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted
+    const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset;
+
+    int    u[QR3_K];
+    float d8[QR3_K];
+
+#pragma unroll
+    for (int i = 0; i < QR3_K; ++i) {
+        u[i]  = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
+        d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
+    }
+
+    return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8);
+}
+
+static __device__ __forceinline__ float vec_dot_q4_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q4_K * bq4_K = (const block_q4_K *) vbq + kbx;
+
+    int    v[2];
+    int    u[2*QR4_K];
+    float d8[QR4_K];
+
+    // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6
+    const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2));
+
+    // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12
+    // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44
+    // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76
+    // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108
+
+    const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4));
+    v[0] = q4[0];
+    v[1] = q4[4];
+
+    // branchless so nvcc can hoist this out of the ncols_dst loop
+    const uint16_t * scales = (const uint16_t *)bq4_K->scales;
+    const int j  = bq8_offset/2;
+    const int jm = j & 1;
+
+    const uint32_t s0 = scales[jm + 0];
+    const uint32_t s2 = scales[jm + 2];
+    const uint32_t s4 = scales[jm + 4];
+
+    const uint32_t hi = (uint32_t) -(int32_t) (j >= 2);
+
+    uint16_t aux[2];
+    aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi));
+    aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi));
+    const uint8_t * sc = (const uint8_t *)aux;
+    const uint8_t * m  = sc + 2;
+
+    for (int i = 0; i < QR4_K; ++i) {
+        const block_q8_1 * bq8i = bq8_1 + bq8_offset + i;
+        d8[i] = __low2float(bq8i->ds);
+
+        const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4);
+        u[2*i+0] = q8[0];
+        u[2*i+1] = q8[4];
+    }
+
+    return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8);
+}
+
+static __device__ __forceinline__ float vec_dot_q5_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q5_K * bq5_K = (const block_q5_K *) vbq + kbx;
+
+    int   vl[2];
+    int   vh[2];
+    int    u[2*QR5_K];
+    float d8[QR5_K];
+
+    const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2));
+    const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4));
+    const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4));
+
+    vl[0] = ql[0];
+    vl[1] = ql[4];
+
+    vh[0] = qh[0] >> bq8_offset;
+    vh[1] = qh[4] >> bq8_offset;
+
+    // same as q4_K
+    const uint16_t * scales = (const uint16_t *)bq5_K->scales;
+    const int j  = bq8_offset/2;
+    const int jm = j & 1;
+
+    const uint32_t s0 = scales[jm + 0];
+    const uint32_t s2 = scales[jm + 2];
+    const uint32_t s4 = scales[jm + 4];
+
+    const uint32_t hi = (uint32_t) -(int32_t) (j >= 2);
+
+    uint16_t aux[2];
+    aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi));
+    aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi));
+
+    const uint8_t * sc = (const uint8_t *)aux;
+    const uint8_t * m  = sc + 2;
+
+#pragma unroll
+    for (int i = 0; i < QR5_K; ++i) {
+        const block_q8_1 * bq8i = bq8_1 + bq8_offset + i;
+        d8[i] = __low2float(bq8i->ds);
+
+        const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4);
+        u[2*i+0] = q8[0];
+        u[2*i+1] = q8[4];
+    }
+
+    return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8);
+}
+
+static __device__ __forceinline__ float vec_dot_q6_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q6_K * bq6_K = (const block_q6_K *) vbq + kbx;
+
+    const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4);
+    const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8);
+    const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4));
+
+    const int vl = get_int_b2(bq6_K->ql, iqs);
+    const int vh = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift;
+
+    const int8_t * scales = bq6_K->scales + scale_offset;
+
+    int    u[QR6_K];
+    float d8[QR6_K];
+
+#pragma unroll
+    for (int i = 0; i < QR6_K; ++i) {
+        u[i]  = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1);
+        d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds);
+    }
+
+    return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8);
+}
+
+#define VDR_IQ2_XXS_Q8_1_MMVQ 2
+#define VDR_IQ2_XXS_Q8_1_MMQ  2
+
+static __device__ __forceinline__ float vec_dot_iq2_xxs_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq2_xxs * bq2 = (const block_iq2_xxs *) vbq + kbx;
+
+    const int q2 = get_int_b2(bq2->qs, iqs);
+    const uint8_t * aux8 = (const uint8_t *) &q2;
+    const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1);
+
+    int sumi = 0;
+#pragma unroll
+    for (int k0 = 0; k0 < 8; k0 += 2) {
+        const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[k0/2]];
+        const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2));
+
+        const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+        const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0);
+        const int u0 = get_int_b4(bq8_1[iqs/2].qs, k0 + 0);
+        sumi = ggml_cuda_dp4a(grid0, u0, sumi);
+
+        const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+        const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1);
+        const int u1 = get_int_b4(bq8_1[iqs/2].qs, k0 + 1);
+        sumi = ggml_cuda_dp4a(grid1, u1, sumi);
+    }
+
+    const int ls = aux32 >> 27 | 1; // (scale * 2 + 1)
+    sumi = sumi * ls / 8;           // (sumi * scale + sumi / 2) / 4
+    const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds);
+    return d * sumi;
+}
+
+#define VDR_IQ2_XS_Q8_1_MMVQ 2
+#define VDR_IQ2_XS_Q8_1_MMQ  2
+
+static __device__ __forceinline__ float vec_dot_iq2_xs_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq2_xs * bq2 = (const block_iq2_xs *) vbq + kbx;
+
+    const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1));
+    const uint16_t * q2 = (const uint16_t *) &q2_packed;
+    const int ls0 = bq2->scales[iqs/2] & 0x0F;
+    const int ls1 = bq2->scales[iqs/2] >> 4;
+
+    int sumi0 = 0;
+    int sumi1 = 0;
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l0/2] & 0x1FF];
+        const uint32_t signs = unpack_ksigns(q2[l0/2] >> 9);
+
+        const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+        const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+        const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0);
+
+        const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+        const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+        const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1);
+
+        if (l0 < 4) {
+            sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0);
+            sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0);
+        } else {
+            sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1);
+            sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1);
+        }
+    }
+    const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4;
+    const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds);
+    return d * sumi;
+}
+
+#define VDR_IQ2_S_Q8_1_MMVQ 2
+#define VDR_IQ2_S_Q8_1_MMQ  2
+
+static __device__ __forceinline__ float vec_dot_iq2_s_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq2_s * bq2 = (const block_iq2_s *) vbq + kbx;
+
+    const int       qs_packed = get_int_b2(bq2->qs, iqs/2);
+    const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+    const int qh = bq2->qh[iqs/2];
+
+    const int       signs_packed_32 = get_int_b2(bq2->qs, QK_K/32 + iqs/2);
+    const uint8_t * signs_packed_8  = (const uint8_t *) &signs_packed_32;
+
+    const int ls0 = bq2->scales[iqs/2] & 0x0F;
+    const int ls1 = bq2->scales[iqs/2] >> 4;
+
+    int sumi0 = 0;
+    int sumi1 = 0;
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const int * grid_pos = (const int *)(iq2s_grid + (qs[l0/2] | ((qh << (8-l0)) & 0x300)));
+
+        const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000);
+        const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000);
+
+        const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0);
+        const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1);
+
+        const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0);
+        const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1);
+
+        if (l0 < 4) {
+            sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0);
+            sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0);
+        } else {
+            sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1);
+            sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1);
+        }
+    }
+    const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4;
+
+    const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds);
+    return d * sumi;
+}
+
+#define VDR_IQ3_XXS_Q8_1_MMVQ 2
+#define VDR_IQ3_XXS_Q8_1_MMQ  2
+
+static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq3_xxs * bq3 = (const block_iq3_xxs *) vbq + kbx;
+
+    const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs+1));
+    const uint8_t * q3 = (const uint8_t *) &q3_packed;
+    const uint32_t aux32 = get_int_b2(bq3->qs, QK_K/16 + iqs/2);
+
+    int sumi = 0;
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]);
+        const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2));
+
+        const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+        const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+
+        const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0);
+
+        const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+        const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+        const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1);
+
+        sumi = ggml_cuda_dp4a(grid_l, u0, sumi);
+        sumi = ggml_cuda_dp4a(grid_h, u1, sumi);
+    }
+
+    const int ls = aux32 >> 28;
+    sumi = (ls*sumi + sumi/2)/2;
+    const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds);
+    return d * sumi;
+}
+
+#define VDR_IQ3_S_Q8_1_MMVQ 2
+#define VDR_IQ3_S_Q8_1_MMQ  2
+
+// TODO: don't use lookup table for signs
+static __device__ __forceinline__ float vec_dot_iq3_s_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq3_s * bq3 = (const block_iq3_s *) vbq + kbx;
+
+    const int2      qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1));
+    const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+    const int qh = bq3->qh[iqs/2];
+
+    const int       signs_packed_32 = get_int_b2(bq3->signs, iqs/2);
+    const uint8_t * signs_packed_8  = (const uint8_t *) &signs_packed_32;
+
+    int sumi = 0;
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const int2 grid_pos = make_int2(
+            iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)],
+            iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]);
+
+        const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000);
+        const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000);
+
+        const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+        const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+        const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0);
+        const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1);
+
+        sumi = ggml_cuda_dp4a(grid_l, u0, sumi);
+        sumi = ggml_cuda_dp4a(grid_h, u1, sumi);
+    }
+
+    sumi *= 1 + 2*((bq3->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F);
+
+    const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds);
+    return d * sumi;
+}
+
+#define VDR_IQ1_S_Q8_1_MMVQ 1
+#define VDR_IQ1_S_Q8_1_MMQ  1
+
+static __device__ __forceinline__ float vec_dot_iq1_s_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+    const block_iq1_s * bq1 = (const block_iq1_s *) vbq + kbx;
+
+    const int       qs_packed = get_int_b2(bq1->qs, iqs);
+    const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+    const int qh = bq1->qh[iqs];
+
+    int sumi = 0;
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const int grid = iq1s_grid_gpu[qs[l0/2] | (((qh >> 3*(l0/2)) & 0x07) << 8)];
+
+        const int grid0 = (grid >> 0) & 0x0F0F0F0F;
+        const int grid1 = (grid >> 4) & 0x0F0F0F0F;
+
+        const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0);
+        const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1);
+
+        sumi = ggml_cuda_dp4a(grid0, u0, sumi);
+        sumi = ggml_cuda_dp4a(grid1, u1, sumi);
+    }
+
+    const float  d1q   = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1);
+    const float  delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000);
+    const float2 ds    = __half22float2(bq8_1[iqs].ds);
+    return d1q * (ds.x*sumi + ds.y*delta);
+}
+
+#define VDR_IQ1_M_Q8_1_MMVQ 1
+#define VDR_IQ1_M_Q8_1_MMQ  1
+
+static __device__ __forceinline__ float vec_dot_iq1_m_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq1_m * bq1 = (const block_iq1_m *) vbq + kbx;
+
+    const int       qs_packed = get_int_b4(bq1->qs, iqs);
+    const uint8_t * qs        = (const uint8_t *) &qs_packed;
+
+    int   sumi[2] = {0};
+    float sumf[2] = {0.0f};
+#pragma unroll
+    for (int l0 = 0; l0 < 8; l0 += 2) {
+        const int qhl = bq1->qh[2*iqs + l0/4] >> (4 * ((l0/2) % 2));
+
+        const int grid = iq1s_grid_gpu[qs[l0/2] | ((qhl & 0x07) << 8)];
+
+        const int grid0 = (grid >> 0) & 0x0F0F0F0F;
+        const int grid1 = (grid >> 4) & 0x0F0F0F0F;
+
+        const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0);
+        const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1);
+
+        sumi[l0/4] = ggml_cuda_dp4a(grid0, u0, sumi[l0/4]);
+        sumi[l0/4] = ggml_cuda_dp4a(grid1, u1, sumi[l0/4]);
+
+        const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f*IQ1M_DELTA/0x08);
+        int sumy = 0;
+        sumy = ggml_cuda_dp4a(u0, 0x01010101, sumy);
+        sumy = ggml_cuda_dp4a(u1, 0x01010101, sumy);
+        sumf[l0/4] += delta*sumy;
+    }
+
+    const uint16_t * sc = (const uint16_t *) bq1->scales;
+
+    iq1m_scale_t scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000);
+    const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds);
+
+    const int tmp = sc[iqs/2] >> (6*(iqs%2));
+    const int sc0 = 2*((tmp >> 0) & 0x07) + 1;
+    const int sc1 = 2*((tmp >> 3) & 0x07) + 1;
+    return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1);
+}
+
+#define VDR_IQ4_NL_Q8_1_MMVQ 2
+#define VDR_IQ4_NL_Q8_1_MMQ  4
+
+static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq4_nl * bq4 = (const block_iq4_nl *) vbq + kbx;
+
+    const int * q8 = (const int *) bq8_1->qs + iqs;
+
+    int sumi = 0;
+#pragma unroll
+    for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) {
+        const int aux_q4 = get_int_b2(bq4->qs, iqs + l);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+
+        sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi);
+        sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi);
+    }
+
+    const float d = __half2float(bq4->d) * __low2float(bq8_1->ds);
+    return d * sumi;
+}
+
+#define VDR_IQ4_XS_Q8_1_MMVQ 4
+#define VDR_IQ4_XS_Q8_1_MMQ  4
+
+static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx;
+
+    int sumi = 0;
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        const int aux_q4 = get_int_b4(bq4->qs, iqs + j);
+        const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+
+        const int u0 = get_int_b4(bq8_1[iqs/4].qs, j + 0);
+        const int u1 = get_int_b4(bq8_1[iqs/4].qs, j + 4);
+
+        sumi = ggml_cuda_dp4a(v.x, u0, sumi);
+        sumi = ggml_cuda_dp4a(v.y, u1, sumi);
+    }
+
+    const int ls = ((bq4->scales_l[iqs/8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs/2)) & 0x03) << 4);
+    sumi *= ls - 32;
+
+    const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs/4].ds);
+    return d * sumi;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h
new file mode 100644
index 0000000000..323c980193
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cuda.h>
+#include <cublas_v2.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#ifdef GGML_USE_NCCL
+#include <nccl.h>
+#endif // GGML_USE_NCCL
+
+#if CUDART_VERSION >= 11080
+#include <cuda_fp8.h>
+#define FP8_AVAILABLE
+#endif // CUDART_VERSION >= 11080
+
+#if CUDART_VERSION >= 12080
+#include <cuda_fp4.h>
+#endif // CUDART_VERSION >= 12080
+
+#if CUDART_VERSION < 11020
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
+#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH
+#define CUBLAS_COMPUTE_16F CUDA_R_16F
+#define CUBLAS_COMPUTE_32F CUDA_R_32F
+#define cublasComputeType_t cudaDataType_t
+#endif // CUDART_VERSION < 11020
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h
new file mode 100644
index 0000000000..da7e3da0b9
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h
@@ -0,0 +1,313 @@
+#pragma once
+
+#define HIP_DISABLE_WARP_SYNC_BUILTINS 1
+#include <hip/hip_runtime.h>
+#include <hipblas/hipblas.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_bf16.h>
+
+#ifdef GGML_USE_NCCL
+#include <rccl/rccl.h>
+#endif // GGML_USE_NCCL
+
+
+#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
+#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT
+#define CUBLAS_OP_N HIPBLAS_OP_N
+#define CUBLAS_OP_T HIPBLAS_OP_T
+#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
+#define CUBLAS_TF32_TENSOR_OP_MATH 0
+#define CUDA_R_16F  HIPBLAS_R_16F
+#define CUDA_R_16BF HIPBLAS_R_16B
+#define CUDA_R_32F  HIPBLAS_R_32F
+#define CUBLAS_SIDE_RIGHT HIPBLAS_SIDE_RIGHT
+#define CUBLAS_FILL_MODE_UPPER HIPBLAS_FILL_MODE_UPPER
+#define CUBLAS_DIAG_NON_UNIT HIPBLAS_DIAG_NON_UNIT
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported
+#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended
+#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned
+#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice
+#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite
+#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }}
+#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width)
+#define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width)
+#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)
+#define __all_sync(mask, var) __all(var)
+#define __any_sync(mask, var) __any(var)
+#define cublasStrsmBatched hipblasStrsmBatched
+#define cublasCreate hipblasCreate
+#define cublasDestroy hipblasDestroy
+#define cublasGemmEx hipblasGemmEx
+#define cublasGemmBatchedEx hipblasGemmBatchedEx
+#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx
+#define cublasHandle_t hipblasHandle_t
+#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS
+#define cublasSetStream hipblasSetStream
+#define cublasSgemm hipblasSgemm
+#define cublasSgemmBatched hipblasSgemmBatched
+#define cublasSgemmStridedBatched hipblasSgemmStridedBatched
+#define cublasStatus_t hipblasStatus_t
+#define cublasOperation_t hipblasOperation_t
+#define cudaDevAttrCooperativeLaunch hipDeviceAttributeCooperativeLaunch
+#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer
+#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess
+#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#define cudaDeviceGetPCIBusId hipDeviceGetPCIBusId
+#define cudaDeviceProp hipDeviceProp_t
+#define cudaDeviceSynchronize hipDeviceSynchronize
+#define cudaError_t hipError_t
+#define cudaErrorMemoryAllocation hipErrorOutOfMemory
+#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
+#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
+#define cudaEventCreateWithFlags hipEventCreateWithFlags
+#define cudaEventDisableTiming hipEventDisableTiming
+#define cudaEventRecord hipEventRecord
+#define cudaEventSynchronize hipEventSynchronize
+#define cudaEvent_t hipEvent_t
+#define cudaEventDestroy hipEventDestroy
+#define cudaFree hipFree
+#define cudaFreeHost hipHostFree
+#define cudaGetDevice hipGetDevice
+#define cudaGetDeviceCount hipGetDeviceCount
+#define cudaGetDeviceProperties hipGetDeviceProperties
+#define cudaGetErrorString hipGetErrorString
+#define cudaGetLastError hipGetLastError
+#define cudaHostAlloc hipHostMalloc
+#define cudaHostAllocPortable hipHostMallocPortable
+#define cudaHostAllocMapped hipHostMallocMapped
+#define cudaHostGetDevicePointer hipHostGetDevicePointer
+#define cudaHostRegister hipHostRegister
+#define cudaHostRegisterPortable hipHostRegisterPortable
+#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
+#define cudaHostUnregister hipHostUnregister
+#define cudaLaunchCooperativeKernel hipLaunchCooperativeKernel
+#define cudaLaunchHostFunc hipLaunchHostFunc
+#define cudaMalloc hipMalloc
+#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault)
+#define cudaMallocManaged hipMallocManaged
+#define cudaMemAdvise hipMemAdvise
+#define cudaMemcpy hipMemcpy
+#define cudaMemcpyAsync hipMemcpyAsync
+#define cudaMemcpyPeerAsync hipMemcpyPeerAsync
+#define cudaMemcpy2DAsync hipMemcpy2DAsync
+#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
+#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
+#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
+#define cudaMemcpyKind hipMemcpyKind
+#define cudaMemset hipMemset
+#define cudaMemsetAsync hipMemsetAsync
+#define cudaMemGetInfo hipMemGetInfo
+#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize
+#define cudaSetDevice hipSetDevice
+#define cuDeviceGet hipDeviceGet
+#define CUdevice hipDevice_t
+#define CUdeviceptr hipDeviceptr_t
+#define cuMemUnmap hipMemUnmap
+#define CUmemAccessDesc hipMemAccessDesc
+#define cuMemAddressFree hipMemAddressFree
+#define cuMemRelease hipMemRelease
+#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t
+#define cuMemCreate hipMemCreate
+#define cuMemAddressReserve hipMemAddressReserve
+#define cuMemMap hipMemMap
+#define cuMemSetAccess hipMemSetAccess
+#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity
+#define CUmemAllocationProp hipMemAllocationProp
+#define cuDeviceGetAttribute hipDeviceGetAttribute
+#define cudaStreamCreateWithFlags hipStreamCreateWithFlags
+#define cudaStreamDestroy hipStreamDestroy
+#define cudaStreamFireAndForget hipStreamFireAndForget
+#define cudaStreamNonBlocking hipStreamNonBlocking
+#define cudaStreamPerThread hipStreamPerThread
+#define cudaStreamSynchronize hipStreamSynchronize
+#define cudaStreamWaitEvent hipStreamWaitEvent
+#define cudaGraphExec_t hipGraphExec_t
+#define cudaGraphNode_t hipGraphNode_t
+#define cudaKernelNodeParams hipKernelNodeParams
+#define cudaKernelNodeParams hipKernelNodeParams
+#define cudaGraphExecDestroy hipGraphExecDestroy
+#define cudaGraphLaunch hipGraphLaunch
+#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure
+#define cudaGraphExecUpdateResult hipGraphExecUpdateResult
+#define cudaGraphNodeType hipGraphNodeType
+#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel
+#define cudaGraphInstantiate hipGraphInstantiate
+#define cudaStreamEndCapture hipStreamEndCapture
+#define cudaGraphDestroy hipGraphDestroy
+#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams
+#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction
+#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams
+#define cudaGraphNodeGetType hipGraphNodeGetType
+#define cudaGraphGetNodes hipGraphGetNodes
+#define cudaGraphExecUpdate hipGraphExecUpdate
+#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed
+#define cudaStreamBeginCapture hipStreamBeginCapture
+#define cudaGraph_t hipGraph_t
+#define cudaStream_t hipStream_t
+#define cudaSuccess hipSuccess
+#define cudaOccupancyMaxActiveBlocksPerMultiprocessor hipOccupancyMaxActiveBlocksPerMultiprocessor
+#define cudaFuncSetAttribute hipFuncSetAttribute
+#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
+#define __trap() do { abort(); __builtin_unreachable(); } while(0)
+#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
+#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED
+#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED
+#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE
+#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH
+#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR
+#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED
+#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR
+#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED
+
+#if HIP_VERSION >= 60500000
+#define CUBLAS_COMPUTE_16F HIPBLAS_COMPUTE_16F
+#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F
+#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_COMPUTE_32F_FAST_16F
+#define cublasComputeType_t hipblasComputeType_t
+#define cudaDataType_t hipDataType
+#else
+#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F
+#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F
+#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F
+#define cublasComputeType_t hipblasDatatype_t
+#define cudaDataType_t hipblasDatatype_t
+#endif // HIP_VERSION >= 6050000
+
+#if !defined(__HIP_PLATFORM_AMD__)
+#error "The HIP backend supports only AMD targets"
+#endif // !defined(__HIP_PLATFORM_AMD__)
+
+#define __CUDA_ARCH__ 1300
+
+#if defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__)
+#define GCN5
+#endif // defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__)
+
+#if defined(__gfx803__)
+#define GCN4
+#endif // defined(__gfx803__)
+
+#if defined(GCN5) || defined(GCN4)
+#define GCN
+#endif // defined(GCN5) || defined(GCN4)
+
+#if defined(__gfx950__)
+#define CDNA4
+#endif // defined(__gfx950__)
+
+#if defined(__gfx942__)
+#define CDNA3
+#endif // defined(__gfx942__)
+
+#if defined(__gfx90a__)
+#define CDNA2
+#endif // defined(__gfx90a__)
+
+#if defined(__gfx908__)
+#define CDNA1
+#endif // defined(__gfx908__)
+
+#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1)
+#define CDNA // For the entire family
+#endif // defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1)
+
+#if defined(__GFX12__)
+#define RDNA4
+#endif // defined(__GFX12__)
+
+#if defined(__GFX11__)
+#define RDNA3
+#endif // defined(__GFX11__)
+
+#if defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__)
+#define RDNA3_5
+#endif // defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__)
+
+#if defined(RDNA3) && !defined(RDNA3_5)
+#define RDNA3_0
+#endif // defined(RDNA3) && !defined(RDNA3_5)
+
+#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \
+    defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__)
+#define RDNA2
+#endif
+
+#if defined(__gfx1010__) || defined(__gfx1012__)
+#define RDNA1
+#endif // defined(__gfx1010__) || defined(__gfx1012__)
+
+#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
+#define RDNA // For the entire family
+#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
+
+#ifndef __has_builtin
+    #define __has_builtin(x) 0
+#endif
+
+typedef __hip_bfloat16 nv_bfloat16;
+typedef __hip_bfloat162 nv_bfloat162;
+
+#if HIP_VERSION >= 60300000
+#include <hip/hip_fp8.h>
+typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
+#define FP8_AVAILABLE
+#endif // HIP_VERSION >= 60300000
+
+typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
+typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4)));
+static __device__ __forceinline__ int __vsubss4(const int a, const int b) {
+    const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
+    const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
+#if __has_builtin(__builtin_elementwise_sub_sat)
+    const int8x4_t c = __builtin_elementwise_sub_sat(va, vb);
+    return reinterpret_cast<const int &>(c);
+#else
+    int8x4_t c;
+    int16_t tmp;
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+        tmp = va[i] - vb[i];
+        if(tmp > std::numeric_limits<int8_t>::max()) tmp = std::numeric_limits<int8_t>::max();
+        if(tmp < std::numeric_limits<int8_t>::min()) tmp = std::numeric_limits<int8_t>::min();
+        c[i] = tmp;
+    }
+    return reinterpret_cast<int &>(c);
+#endif // __has_builtin(__builtin_elementwise_sub_sat)
+}
+
+static __device__ __forceinline__ int __vsub4(const int a, const int b) {
+    // do some small modifications to a and b to make the subtraction not underflow
+    const unsigned int a_large = a | 0x80808080;
+    const unsigned int b_small = b & 0x7f7f7f7f;
+    const unsigned int result_low_7bits = a_large - b_small;
+
+    // if two ops share the same high bit, we should flip the high bit of the result
+    const unsigned int should_flip_high_1bit = (a ^ ~b) & 0x80808080;
+
+    return result_low_7bits ^ should_flip_high_1bit;
+}
+
+static __device__ __forceinline__ unsigned int __vcmpeq4(unsigned int a, unsigned int b) {
+    const uint8x4_t& va = reinterpret_cast<const uint8x4_t&>(a);
+    const uint8x4_t& vb = reinterpret_cast<const uint8x4_t&>(b);
+    unsigned int c;
+    uint8x4_t& vc = reinterpret_cast<uint8x4_t&>(c);
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+        vc[i] = va[i] == vb[i] ? 0xff : 0x00;
+    }
+    return c;
+}
+
+static __device__ __forceinline__ unsigned int __vcmpne4(unsigned int a, unsigned int b) {
+    const unsigned int x = a ^ b;
+
+    // any non-equal bit in a byte will set the high bit of that byte here
+    // the addition will not overflow in the byte as op1 and op2 are both less than 0x80
+    const unsigned int ne_low_7bits = ((x & 0x7f7f7f7f) + 0x7f7f7f7f) & 0x80808080;
+    const unsigned int ne_high_1bit = x & 0x80808080;
+    const unsigned int ne_any_bit = ne_low_7bits | ne_high_1bit;
+
+    return (ne_any_bit >> 7) * 0xff;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h
new file mode 100644
index 0000000000..22126decfb
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h
@@ -0,0 +1,161 @@
+#pragma once
+
+#include <musa_runtime.h>
+#include <musa.h>
+#include <mublas.h>
+#include <musa_bf16.h>
+#include <musa_fp16.h>
+
+#ifdef __MUSA_ARCH__
+#define __CUDA_ARCH__ 1300 // GGML_CUDA_CC_RUBIN
+#endif // __MUSA_ARCH__
+
+#define CUBLAS_COMPUTE_16F CUDA_R_16F
+#define CUBLAS_COMPUTE_32F CUDA_R_32F
+#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F
+#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT
+#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT
+#define CUBLAS_OP_N MUBLAS_OP_N
+#define CUBLAS_OP_T MUBLAS_OP_T
+#define CUBLAS_DEFAULT_MATH MUBLAS_DEFAULT_MATH
+#define CUBLAS_SIDE_RIGHT MUBLAS_SIDE_RIGHT
+#define CUBLAS_FILL_MODE_UPPER MUBLAS_FILL_MODE_UPPER
+#define CUBLAS_DIAG_NON_UNIT MUBLAS_DIAG_NON_UNIT
+#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS
+#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_TENSOR_OP_MATH
+#define CUDA_R_16F  MUSA_R_16F
+#define CUDA_R_16BF MUSA_R_16BF
+#define CUDA_R_32F  MUSA_R_32F
+#define cublasStrsmBatched mublasStrsmBatched
+#define cublasComputeType_t cudaDataType_t
+#define cublasCreate mublasCreate
+#define cublasDestroy mublasDestroy
+#define cublasGemmEx mublasGemmEx
+#define cublasGemmBatchedEx mublasGemmBatchedEx
+#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx
+#define cublasHandle_t mublasHandle_t
+#define cublasSetMathMode mublasSetMathMode
+#define cublasSetStream mublasSetStream
+#define cublasSgemm mublasSgemm
+#define cublasSgemmBatched mublasSgemmBatched
+#define cublasSgemmStridedBatched mublasSgemmStridedBatched
+#define cublasStatus_t mublasStatus_t
+#define cublasOperation_t mublasOperation_t
+#define cublasGetStatusString mublasGetStatusString
+#define cudaDataType_t musaDataType_t
+#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer
+#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess
+#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess
+#define cudaDeviceGetPCIBusId musaDeviceGetPCIBusId
+#define cudaDeviceProp musaDeviceProp
+#define cudaDeviceSynchronize musaDeviceSynchronize
+#define cudaDeviceGetAttribute musaDeviceGetAttribute
+#define cudaError_t musaError_t
+#define cudaErrorMemoryAllocation musaErrorMemoryAllocation
+#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
+#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled
+#define cudaEventCreateWithFlags musaEventCreateWithFlags
+#define cudaEventDisableTiming musaEventDisableTiming
+#define cudaEventRecord musaEventRecord
+#define cudaEventSynchronize musaEventSynchronize
+#define cudaEvent_t musaEvent_t
+#define cudaEventDestroy musaEventDestroy
+#define cudaFree musaFree
+#define cudaFreeHost musaFreeHost
+#define cudaGetDevice musaGetDevice
+#define cudaGetDeviceCount musaGetDeviceCount
+#define cudaGetDeviceProperties musaGetDeviceProperties
+#define cudaGetErrorString musaGetErrorString
+#define cudaGetLastError musaGetLastError
+#define cudaHostRegister musaHostRegister
+#define cudaHostRegisterPortable musaHostRegisterPortable
+#define cudaHostRegisterReadOnly musaHostRegisterReadOnly
+#define cudaHostUnregister musaHostUnregister
+#define cudaLaunchCooperativeKernel musaLaunchCooperativeKernel
+#define cudaLaunchHostFunc musaLaunchHostFunc
+#define cudaMalloc musaMalloc
+#define cudaMallocHost musaMallocHost
+#define cudaMallocManaged musaMallocManaged
+#define cudaMemcpy musaMemcpy
+#define cudaMemcpyAsync musaMemcpyAsync
+#define cudaMemcpyPeerAsync musaMemcpyPeerAsync
+#define cudaMemcpy2DAsync musaMemcpy2DAsync
+#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice
+#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost
+#define cudaMemcpyHostToDevice musaMemcpyHostToDevice
+#define cudaMemcpyKind musaMemcpyKind
+#define cudaMemset musaMemset
+#define cudaMemsetAsync musaMemsetAsync
+#define cudaMemGetInfo musaMemGetInfo
+#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize
+#define cudaSetDevice musaSetDevice
+#define cudaStreamCreateWithFlags musaStreamCreateWithFlags
+#define cudaStreamDestroy musaStreamDestroy
+#define cudaStreamFireAndForget musaStreamFireAndForget
+#define cudaStreamNonBlocking musaStreamNonBlocking
+#define cudaStreamPerThread musaStreamPerThread
+#define cudaStreamSynchronize musaStreamSynchronize
+#define cudaStreamWaitEvent musaStreamWaitEvent
+#define cudaStream_t musaStream_t
+#define cudaSuccess musaSuccess
+
+// Additional mappings for MUSA virtual memory pool
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
+#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE
+#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED
+#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED
+#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE
+#define CUdevice MUdevice
+#define CUdeviceptr MUdeviceptr
+#define CUmemAccessDesc MUmemAccessDesc
+#define CUmemAllocationProp MUmemAllocationProp
+#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle
+#define cuDeviceGet muDeviceGet
+#define cuDeviceGetAttribute muDeviceGetAttribute
+#define cuMemAddressFree muMemAddressFree
+#define cuMemAddressReserve muMemAddressReserve
+#define cuMemCreate muMemCreate
+#define cuMemGetAllocationGranularity muMemGetAllocationGranularity
+#define cuMemMap muMemMap
+#define cuMemRelease muMemRelease
+#define cuMemSetAccess muMemSetAccess
+#define cuMemUnmap muMemUnmap
+#define cudaDevAttrCooperativeLaunch musaDevAttrCooperativeLaunch
+#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncSetAttribute musaFuncSetAttribute
+#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms
+#define make_cudaExtent make_musaExtent
+#define make_cudaPitchedPtr make_musaPitchedPtr
+
+// Additional mappings for MUSA graphs
+#define CUDA_SUCCESS MUSA_SUCCESS
+#define CUresult MUresult
+#define cuGetErrorString muGetErrorString
+#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure
+#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction
+#define cudaGraphDestroy musaGraphDestroy
+#define cudaGraphExecDestroy musaGraphExecDestroy
+#define cudaGraphExec_t musaGraphExec_t
+#define cudaGraphExecUpdate musaGraphExecUpdate
+#define cudaGraphExecUpdateResult musaGraphExecUpdateResult
+#define cudaGraphGetNodes musaGraphGetNodes
+#define cudaGraphInstantiate musaGraphInstantiate
+#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams
+#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams
+#define cudaGraphLaunch musaGraphLaunch
+#define cudaGraphNodeGetType musaGraphNodeGetType
+#define cudaGraphNode_t musaGraphNode_t
+#define cudaGraphNodeType musaGraphNodeType
+#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel
+#define cudaGraph_t musaGraph_t
+#define cudaKernelNodeParams musaKernelNodeParams
+#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed
+#define cudaStreamBeginCapture musaStreamBeginCapture
+#define cudaStreamEndCapture musaStreamEndCapture
+#define cudaStreamCaptureStatus musaStreamCaptureStatus
+#define cudaStreamCaptureStatusNone musaStreamCaptureStatusNone
+#define cudaStreamIsCapturing musaStreamIsCapturing
+#define cudaOccupancyMaxActiveBlocksPerMultiprocessor musaOccupancyMaxActiveBlocksPerMultiprocessor
+
+typedef __mt_bfloat16 nv_bfloat16;
+typedef __mt_bfloat162 nv_bfloat162;
diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h b/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h
new file mode 100644
index 0000000000..81cce2e2a6
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h
@@ -0,0 +1,87 @@
+#pragma once
+
+#include "ggml.h"
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t;
+typedef struct      ggml_backend_buffer * ggml_backend_buffer_t;
+typedef struct             ggml_backend * ggml_backend_t;
+
+// Tensor allocator
+struct ggml_tallocr {
+    ggml_backend_buffer_t buffer;
+    void * base;
+    size_t alignment;
+    size_t offset;
+};
+
+GGML_API struct ggml_tallocr ggml_tallocr_new(ggml_backend_buffer_t buffer);
+GGML_API enum ggml_status    ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_tensor * tensor);
+
+// Graph allocator
+/*
+  Example usage:
+    ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type());
+
+    // optional: create a worst-case graph and reserve the buffers to avoid reallocations
+    ggml_gallocr_reserve(galloc, build_graph(max_batch));
+
+    // allocate the graph
+    struct ggml_cgraph * graph = build_graph(batch);
+    ggml_gallocr_alloc_graph(galloc, graph);
+
+    printf("compute buffer size: %zu bytes\n", ggml_gallocr_get_buffer_size(galloc, 0));
+
+    // evaluate the graph
+    ggml_backend_graph_compute(backend, graph);
+*/
+
+// special tensor flags for use with the graph allocator:
+//   ggml_set_input(): all input tensors are allocated at the beginning of the graph in non-overlapping addresses
+//   ggml_set_output(): output tensors are never freed and never overwritten
+
+typedef struct ggml_gallocr * ggml_gallocr_t;
+
+GGML_API ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft);
+GGML_API ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs);
+GGML_API void           ggml_gallocr_free(ggml_gallocr_t galloc);
+
+// pre-allocate buffers from a measure graph - does not allocate or modify the graph
+// call with a worst-case graph to avoid buffer reallocations
+// not strictly required for single buffer usage: ggml_gallocr_alloc_graph will reallocate the buffers automatically if needed
+// returns false if the buffer allocation failed
+// ggml_gallocr_resrve_n_size writes the buffer sizes per galloc buffer that would be allocated by ggml_gallocr_reserve_n to sizes
+GGML_API bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
+GGML_API void ggml_gallocr_reserve_n_size(
+    ggml_gallocr_t galloc,
+    struct ggml_cgraph * graph,
+    const int * node_buffer_ids,
+    const int * leaf_buffer_ids,
+    size_t * sizes);
+GGML_API bool ggml_gallocr_reserve_n(
+    ggml_gallocr_t galloc,
+    struct ggml_cgraph * graph,
+    const int * node_buffer_ids,
+    const int * leaf_buffer_ids);
+
+// automatic reallocation if the topology changes when using a single buffer
+// returns false if using multiple buffers and a re-allocation is needed (call ggml_gallocr_reserve_n first to set the node buffers)
+GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
+
+GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
+
+// Utils
+// Create a buffer and allocate all the tensors in a ggml_context
+
+// returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft. returns 0 on failure
+GGML_API size_t                       ggml_backend_alloc_ctx_tensors_from_buft_size(struct ggml_context * ctx, ggml_backend_buffer_type_t buft);
+// returns NULL on failure or if all tensors in ctx are already allocated or zero-sized
+GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft);
+GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend);
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h b/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h
new file mode 100644
index 0000000000..7a65d8121b
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h
@@ -0,0 +1,463 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-alloc.h"
+
+#ifdef GGML_BACKEND_SHARED
+#    if defined(_WIN32) && !defined(__MINGW32__)
+#        ifdef GGML_BACKEND_BUILD
+#            define GGML_BACKEND_API __declspec(dllexport) extern
+#        else
+#            define GGML_BACKEND_API __declspec(dllimport) extern
+#        endif
+#    else
+#        define GGML_BACKEND_API __attribute__ ((visibility ("default"))) extern
+#    endif
+#else
+#    define GGML_BACKEND_API extern
+#endif
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t;
+    typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
+    typedef struct ggml_backend_event * ggml_backend_event_t;
+    typedef struct ggml_backend * ggml_backend_t;
+    typedef void * ggml_backend_graph_plan_t;
+    typedef struct ggml_backend_reg * ggml_backend_reg_t;
+    typedef struct ggml_backend_device * ggml_backend_dev_t;
+
+
+    //
+    // Backend buffer type
+    //
+
+    GGML_API const char *          ggml_backend_buft_name            (ggml_backend_buffer_type_t buft);
+    GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer    (ggml_backend_buffer_type_t buft, size_t size);
+    GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer_n  (ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors);
+    GGML_API size_t                ggml_backend_buft_get_alignment   (ggml_backend_buffer_type_t buft);
+    GGML_API size_t                ggml_backend_buft_get_max_size    (ggml_backend_buffer_type_t buft);
+    GGML_API size_t                ggml_backend_buft_get_alloc_size  (ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor);
+    GGML_API size_t                ggml_backend_buft_get_alloc_size_n(ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors);
+    GGML_API bool                  ggml_backend_buft_is_host         (ggml_backend_buffer_type_t buft);
+    GGML_API ggml_backend_dev_t    ggml_backend_buft_get_device      (ggml_backend_buffer_type_t buft);
+
+    //
+    // Backend buffer
+    //
+
+    enum ggml_backend_buffer_usage {
+        GGML_BACKEND_BUFFER_USAGE_ANY = 0,
+        GGML_BACKEND_BUFFER_USAGE_WEIGHTS = 1,
+        GGML_BACKEND_BUFFER_USAGE_COMPUTE = 2,
+    };
+
+    GGML_API const char *                   ggml_backend_buffer_name          (ggml_backend_buffer_t buffer);
+    GGML_API void                           ggml_backend_buffer_free          (ggml_backend_buffer_t buffer);
+    GGML_API void *                         ggml_backend_buffer_get_base      (ggml_backend_buffer_t buffer);
+    GGML_API size_t                         ggml_backend_buffer_get_size      (ggml_backend_buffer_t buffer);
+    GGML_API enum ggml_status               ggml_backend_buffer_init_tensor   (ggml_backend_buffer_t buffer, struct ggml_tensor * tensor);
+    GGML_API size_t                         ggml_backend_buffer_get_alignment (ggml_backend_buffer_t buffer);
+    GGML_API size_t                         ggml_backend_buffer_get_max_size  (ggml_backend_buffer_t buffer);
+    GGML_API size_t                         ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor);
+    GGML_API void                           ggml_backend_buffer_clear         (ggml_backend_buffer_t buffer, uint8_t value);
+    GGML_API bool                           ggml_backend_buffer_is_host       (ggml_backend_buffer_t buffer);
+    GGML_API void                           ggml_backend_buffer_set_usage     (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+    GGML_API enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage     (ggml_backend_buffer_t buffer);
+    GGML_API ggml_backend_buffer_type_t     ggml_backend_buffer_get_type      (ggml_backend_buffer_t buffer);
+    GGML_API void                           ggml_backend_buffer_reset         (ggml_backend_buffer_t buffer);
+
+    // tensor copy between different backends
+    GGML_API void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst);
+
+    //
+    // Backend (stream)
+    //
+
+    GGML_API ggml_guid_t  ggml_backend_guid(ggml_backend_t backend);
+    GGML_API const char * ggml_backend_name(ggml_backend_t backend);
+    GGML_API void         ggml_backend_free(ggml_backend_t backend);
+
+    GGML_API ggml_backend_buffer_type_t ggml_backend_get_default_buffer_type(ggml_backend_t backend);
+    GGML_API ggml_backend_buffer_t      ggml_backend_alloc_buffer(ggml_backend_t backend, size_t size);
+    GGML_API size_t                     ggml_backend_get_alignment(ggml_backend_t backend);
+    GGML_API size_t                     ggml_backend_get_max_size(ggml_backend_t backend);
+
+    GGML_API void ggml_backend_tensor_set_async   (ggml_backend_t backend,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+    GGML_API void ggml_backend_tensor_get_async   (ggml_backend_t backend, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+    GGML_API void ggml_backend_tensor_set_2d_async(ggml_backend_t backend,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+    GGML_API void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+
+    // "offset" refers to the offset in tensor->data for setting/getting data
+    GGML_API void ggml_backend_tensor_set   (      struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+    GGML_API void ggml_backend_tensor_get   (const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+    GGML_API void ggml_backend_tensor_set_2d(      struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+    GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+    GGML_API void ggml_backend_tensor_memset(      struct ggml_tensor * tensor,     uint8_t value, size_t offset, size_t size);
+
+    GGML_API void ggml_backend_synchronize(ggml_backend_t backend);
+
+    GGML_API ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph);
+    GGML_API void                      ggml_backend_graph_plan_free  (ggml_backend_t backend, ggml_backend_graph_plan_t plan);
+
+    GGML_API enum ggml_status ggml_backend_graph_plan_compute (ggml_backend_t backend, ggml_backend_graph_plan_t plan);
+    GGML_API enum ggml_status ggml_backend_graph_compute      (ggml_backend_t backend, struct ggml_cgraph * cgraph);
+    GGML_API enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph);
+
+    // NOTE: will be removed, use device version instead
+    GGML_API bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op);
+    GGML_API bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft);
+    GGML_API bool ggml_backend_offload_op(ggml_backend_t backend, const struct ggml_tensor * op);
+
+    // asynchronous copy
+    // the copy is performed after all the currently queued operations in backend_src
+    // backend_dst will wait for the copy to complete before performing other operations
+    // automatic fallback to sync copy if async is not supported
+    GGML_API void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);
+
+    GGML_API ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend);
+
+    //
+    // Events
+    //
+
+    GGML_API ggml_backend_event_t ggml_backend_event_new(ggml_backend_dev_t device);
+    GGML_API void                 ggml_backend_event_free(ggml_backend_event_t event);
+    GGML_API void                 ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend);
+    GGML_API void                 ggml_backend_event_synchronize(ggml_backend_event_t event);
+    GGML_API void                 ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event);
+
+    //
+    // Backend device
+    //
+
+    enum ggml_backend_dev_type {
+        // CPU device using system memory
+        GGML_BACKEND_DEVICE_TYPE_CPU,
+        // GPU device using dedicated memory
+        GGML_BACKEND_DEVICE_TYPE_GPU,
+        // integrated GPU device using host memory
+        GGML_BACKEND_DEVICE_TYPE_IGPU,
+        // accelerator devices intended to be used together with the CPU backend (e.g. BLAS or AMX)
+        GGML_BACKEND_DEVICE_TYPE_ACCEL,
+        // "meta" device wrapping multiple other devices for tensor parallelism
+        GGML_BACKEND_DEVICE_TYPE_META,
+    };
+
+    // functionality supported by the device
+    struct ggml_backend_dev_caps {
+        // asynchronous operations
+        bool async;
+        // pinned host buffer
+        bool host_buffer;
+        // creating buffers from host ptr
+        bool buffer_from_host_ptr;
+        // event synchronization
+        bool events;
+        // mmap is supported for loading
+        bool mmap_support;
+    };
+
+    // all the device properties
+    struct ggml_backend_dev_props {
+        // device name
+        const char * name;
+        // device description
+        const char * description;
+        // device free memory in bytes
+        size_t memory_free;
+        // device total memory in bytes
+        size_t memory_total;
+        // device type
+        enum ggml_backend_dev_type type;
+        // device id
+        //   for PCI devices, this should be the lower-case PCI bus id formatted as "domain:bus:device.function" (e.g. "0000:c1:00.0")
+        //   if the id is unknown, this should be NULL
+        const char * device_id;
+        // device capabilities
+        struct ggml_backend_dev_caps caps;
+    };
+
+    GGML_API const char *                  ggml_backend_dev_name(ggml_backend_dev_t device);
+    GGML_API const char *                  ggml_backend_dev_description(ggml_backend_dev_t device);
+    GGML_API void                          ggml_backend_dev_memory(ggml_backend_dev_t device, size_t * free, size_t * total);
+    GGML_API enum ggml_backend_dev_type    ggml_backend_dev_type(ggml_backend_dev_t device);
+    GGML_API void                          ggml_backend_dev_get_props(ggml_backend_dev_t device, struct ggml_backend_dev_props * props);
+    GGML_API ggml_backend_reg_t            ggml_backend_dev_backend_reg(ggml_backend_dev_t device);
+    GGML_API ggml_backend_t                ggml_backend_dev_init(ggml_backend_dev_t device, const char * params);
+    GGML_API ggml_backend_buffer_type_t    ggml_backend_dev_buffer_type(ggml_backend_dev_t device);
+    GGML_API ggml_backend_buffer_type_t    ggml_backend_dev_host_buffer_type(ggml_backend_dev_t device);
+    GGML_API ggml_backend_buffer_t         ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size);
+
+    GGML_API bool                          ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op);
+    GGML_API bool                          ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft);
+    GGML_API bool                          ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op);
+
+    //
+    // Backend (reg)
+    //
+
+    GGML_API const char *       ggml_backend_reg_name(ggml_backend_reg_t reg);
+    GGML_API size_t             ggml_backend_reg_dev_count(ggml_backend_reg_t reg);
+    GGML_API ggml_backend_dev_t ggml_backend_reg_dev_get(ggml_backend_reg_t reg, size_t index);
+    GGML_API void *             ggml_backend_reg_get_proc_address(ggml_backend_reg_t reg, const char * name);
+
+    // Common functions that may be obtained using ggml_backend_reg_get_proc_address
+
+    // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
+    typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+    typedef void   (*ggml_backend_comm_free_t)(void * comm_ctx);
+    typedef bool   (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);
+
+    // Split buffer type for tensor parallelism (old)
+    typedef ggml_backend_buffer_type_t   (*ggml_backend_split_buffer_type_t)(int main_device, const float * tensor_split);
+    // Set the number of threads for the backend
+    typedef void                         (*ggml_backend_set_n_threads_t)(ggml_backend_t backend, int n_threads);
+    // Get additional buffer types provided by the device (returns a NULL-terminated array)
+    typedef ggml_backend_buffer_type_t * (*ggml_backend_dev_get_extra_bufts_t)(ggml_backend_dev_t device);
+    // Set the abort callback for the backend
+    typedef void                         (*ggml_backend_set_abort_callback_t)(ggml_backend_t backend, ggml_abort_callback abort_callback, void * abort_callback_data);
+    // Get a list of feature flags supported by the backend (returns a NULL-terminated array)
+    struct ggml_backend_feature {
+        const char * name;
+        const char * value;
+    };
+    typedef struct ggml_backend_feature * (*ggml_backend_get_features_t)(ggml_backend_reg_t reg);
+
+    //
+    // Backend registry
+    //
+
+    GGML_API void ggml_backend_register(ggml_backend_reg_t reg);
+
+    GGML_API void ggml_backend_device_register(ggml_backend_dev_t device);
+
+    // Backend (reg) enumeration
+    GGML_API size_t             ggml_backend_reg_count(void);
+    GGML_API ggml_backend_reg_t ggml_backend_reg_get(size_t index);
+    GGML_API ggml_backend_reg_t ggml_backend_reg_by_name(const char * name);
+
+    // Device enumeration
+    GGML_API size_t             ggml_backend_dev_count(void);
+    GGML_API ggml_backend_dev_t ggml_backend_dev_get(size_t index);
+    GGML_API ggml_backend_dev_t ggml_backend_dev_by_name(const char * name);
+    GGML_API ggml_backend_dev_t ggml_backend_dev_by_type(enum ggml_backend_dev_type type);
+
+    // Direct backend (stream) initialization
+    // = ggml_backend_dev_init(ggml_backend_dev_by_name(name), params)
+    GGML_API ggml_backend_t ggml_backend_init_by_name(const char * name, const char * params);
+    // = ggml_backend_dev_init(ggml_backend_dev_by_type(type), params)
+    GGML_API ggml_backend_t ggml_backend_init_by_type(enum ggml_backend_dev_type type, const char * params);
+    // = ggml_backend_dev_init(ggml_backend_dev_by_type(GPU) OR ggml_backend_dev_by_type(CPU), NULL)
+    GGML_API ggml_backend_t ggml_backend_init_best(void);
+
+    // Load a backend from a dynamic library and register it
+    GGML_API ggml_backend_reg_t ggml_backend_load(const char * path);
+    // Unload a backend if loaded dynamically and unregister it
+    GGML_API void               ggml_backend_unload(ggml_backend_reg_t reg);
+    // Load all known backends from dynamic libraries
+    GGML_API void               ggml_backend_load_all(void);
+    GGML_API void               ggml_backend_load_all_from_path(const char * dir_path);
+
+    //
+    // Backend scheduler
+    //
+
+    // The backend scheduler allows for multiple backend devices to be used together
+    // Handles compute buffer allocation, assignment of tensors to backends, and copying of tensors between backends
+    // The backends are selected based on:
+    // - the backend that supports the operation
+    // - the location of the pre-allocated tensors (e.g. the weights)
+    /*
+      Example usage:
+
+        // operations that use tensors allocated in a buffer with USAGE_WEIGHTS will be assigned
+        // preferably to run on the same backend as the buffer
+        ggml_backend_buffer_set_usage(buf_weights, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        sched = ggml_backend_sched_new({backend_gpu, backend_gpu2, backend_cpu}, NULL, num_backends, GGML_DEFAULT_GRAPH_SIZE, false, true);
+
+        // initialize buffers from a max size graph (optional)
+        reserve_graph = build_graph(sched, max_batch_size);
+
+        // manually assign nodes to a backend (optional, should not be needed in most cases)
+        struct ggml_tensor * node = ggml_mul_mat(ctx, ...);
+        ggml_backend_sched_set_tensor_backend(sched, node, backend_gpu);
+
+        ggml_backend_sched_reserve(sched, reserve_graph);
+
+        // compute
+        graph = build_graph(sched); // the graph and its tensors are single-use in terms of allocation, multi-use in terms of computation
+        for (int i = 0; i < 10; ++i) {
+            ggml_backend_sched_graph_compute(sched, graph); // on the first iteration the graph is allocated automatically
+        }
+
+        // if there are graph inputs:
+        graph = build_graph(sched); // get a new graph that is not allocated (the metadata for the old graph is freed once ggml_free is called)
+        ggml_backend_sched_reset(sched); // clear the allocation of the previous graph
+        ggml_backend_sched_alloc_graph(sched, graph); // explicitly allocate the new graph but do not execute it
+        ggml_backend_tensor_set(input_tensor, ...); // copy data to the newly allocated graph tensors
+        ggml_backend_sched_graph_compute(sched, graph); // execute the graph
+
+        // as an alternative to the above it is also possible to assign the inputs to a dedicated context and
+        // allocate them statically via ggml_backend_alloc_ctx_tensors
+    }
+    */
+
+    typedef struct ggml_backend_sched * ggml_backend_sched_t;
+
+    // Evaluation callback for each node in the graph (set with ggml_backend_sched_set_eval_callback)
+    // when ask == true, the scheduler wants to know if the user wants to observe this node
+    // this allows the scheduler to batch nodes together in order to evaluate them in a single call
+    //
+    // when ask == false, the scheduler is passing the node tensor to the user for observation
+    // if the user returns false, the scheduler will cancel the graph compute
+    //
+    typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data);
+
+    // Callback while copying input weights of a split
+    // if the user returns false the scheduler simply copies the entire weight
+    // the callback is called only for input weights in host buffers
+    // the callback is called after all non-weight inputs of the split have been copied
+    // `src` is the tensor in the previous split
+    // `dst` is the copy of `src` in the split
+    // `graph` is the compute graph of the split
+    typedef bool (*ggml_backend_sched_copy_callback)(ggml_backend_t backend, const struct ggml_tensor * src, struct ggml_tensor * dst, struct ggml_cgraph * graph, void * user_data);
+
+    // Callback when the buffers reserved for a graph are not sufficient and have to be reserved again
+    // it happens on the graph allocation path and costs a synchronization of the backends
+    // `sched` is the scheduler
+    // `graph` is the compute graph being allocated
+    // `backend_ids_changed` is true if the backend assignment of the graph nodes changed since the last allocation
+    // `graph_size_changed` is true if the graph size changed since the last allocation
+    typedef void (*ggml_backend_sched_realloc_callback)(ggml_backend_sched_t sched, struct ggml_cgraph * graph, bool backend_ids_changed, bool graph_size_changed, void * user_data);
+
+    // Initialize a backend scheduler, backends with low index are given priority over backends with high index
+    GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload);
+    GGML_API void                 ggml_backend_sched_free(ggml_backend_sched_t sched);
+
+    // Initialize backend buffers from a measure graph
+    GGML_API void                 ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes);
+    GGML_API bool                 ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); // returns success
+
+    GGML_API int                  ggml_backend_sched_get_n_backends(ggml_backend_sched_t sched);
+    GGML_API ggml_backend_t       ggml_backend_sched_get_backend(ggml_backend_sched_t sched, int i);
+
+    // Get the number of splits of the last graph
+    GGML_API int                  ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched);
+    GGML_API int                  ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched);
+
+    GGML_API ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend);
+    GGML_API size_t                     ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
+
+    GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
+    GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
+
+    // Split graph without allocating it
+    GGML_API void                 ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph);
+
+    // Allocate and compute graph on the backend scheduler
+    GGML_API bool                 ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph); // returns success
+    GGML_API enum ggml_status     ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph);
+    GGML_API enum ggml_status     ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph);
+    GGML_API void                 ggml_backend_sched_synchronize(ggml_backend_sched_t sched);
+
+    // Reset all assignments and allocators - must be called before changing the node backends or allocating a new graph.
+    // This in effect deallocates all tensors that were previously allocated and leaves them with dangling pointers.
+    // The correct way to use this API is to discard the deallocated tensors and create new ones.
+    GGML_API void                 ggml_backend_sched_reset(ggml_backend_sched_t sched);
+
+    // Set a callback to be called for each resulting node during graph compute
+    GGML_API void                 ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data);
+
+    // Set a callback to be called when the inputs weights of a split are being copied
+    GGML_API void                 ggml_backend_sched_set_copy_callback(ggml_backend_sched_t sched, ggml_backend_sched_copy_callback callback, void * user_data);
+
+    // Set a callback to be called when the compute buffers of a graph are being reallocated
+    GGML_API void                 ggml_backend_sched_set_realloc_callback(ggml_backend_sched_t sched, ggml_backend_sched_realloc_callback callback, void * user_data);
+
+    //
+    // Meta backend
+    //
+
+#define GGML_BACKEND_META_MAX_DEVICES 16
+
+    enum ggml_backend_meta_split_axis {
+        // tensor split by tensor dimensions:
+        GGML_BACKEND_SPLIT_AXIS_0 = 0,
+        GGML_BACKEND_SPLIT_AXIS_1 = 1,
+        GGML_BACKEND_SPLIT_AXIS_2 = 2,
+        GGML_BACKEND_SPLIT_AXIS_3 = 3,
+
+        GGML_BACKEND_SPLIT_AXIS_MIRRORED = 10, // all values on all backends
+        GGML_BACKEND_SPLIT_AXIS_PARTIAL  = 11, // each backend has a partial sum
+
+        // for internal bookkeeping only:
+        GGML_BACKEND_SPLIT_AXIS_NONE    = 98,
+        GGML_BACKEND_SPLIT_AXIS_UNKNOWN = 99,
+    };
+    GGML_API const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis);
+
+    struct ggml_backend_meta_split_state {
+        enum ggml_backend_meta_split_axis axis;
+
+        // for tensors with axis >= 0 && axis < GGML_MAX_DIMS:
+        //   - each device has a slice of the tensor along the split axis
+        //   - most tensors have n_segments == 1 and a contiguous slice of the tensor data
+        //   - some tensors have an inhomogenenous data layout along the split axis,
+        //     those tensors are divided into segments which are each individually split across devices
+        //     (this usually happens when multiple tensors are fused into a single one)
+        //   - ne has one entry per segment and device and that segment repeats nr times,
+        //     in total when accounting for repetitions the segments add up to ggml_tensor::ne for that axis,
+        //     the outer/inner loops are over segments/devices like [seg0_dev0_r0, seg0_dev1_r0, seg0_dev0_r1, seg0_dev1_r1, seg1_dev0_r0, seg1_dev1_r0],
+        //   - for example, a transformer may have a fused QKV matrix rather than 3 matrices, those would be 3 separate segments
+        //     that each need to be split individually across devices so that each device gets a slice of Q, K, and V,
+        //     the Q matrix can be larger than the K and V matrices so this can either be expressed as 3 segments or as 2 segments
+        //     where the segment for K/V repeats twice
+        int64_t  ne[16*GGML_BACKEND_META_MAX_DEVICES];
+        uint32_t nr[16];
+        uint32_t n_segments;
+    };
+
+    // function to assign split states for statically allocated tensors, compute tensor split states will be assigned to be compatible:
+    typedef struct ggml_backend_meta_split_state(*ggml_backend_meta_get_split_state_t)(const struct ggml_tensor * tensor, void * userdata);
+
+    // create a new meta device from "simple" devices, meta buffer type/buffer/backend is then derived from this:
+    // TODO: this looks a bit strange - a backend API creates a device. I think we should try
+    //       express this as a backend registry functionality instead
+    GGML_API ggml_backend_dev_t ggml_backend_meta_device(
+        ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud);
+
+    //
+    // Utils
+    //
+
+    struct ggml_backend_graph_copy {
+        ggml_backend_buffer_t buffer;
+        struct ggml_context * ctx_allocated;
+        struct ggml_context * ctx_unallocated;
+        struct ggml_cgraph * graph;
+    };
+
+    // Copy a graph to a different backend
+    GGML_API struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph);
+    GGML_API void                           ggml_backend_graph_copy_free(struct ggml_backend_graph_copy copy);
+
+    typedef bool (*ggml_backend_eval_callback)(int node_index, struct ggml_tensor * t1, struct ggml_tensor * t2, void * user_data);
+
+    // Compare the output of two backends
+    GGML_API bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes);
+
+    // Tensor initialization
+    GGML_API enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr);
+    GGML_API enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor);
+
+    // CPU buffer types are always available
+    GGML_API ggml_backend_buffer_t      ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size);
+    GGML_API ggml_backend_buffer_type_t ggml_backend_cpu_buffer_type(void);
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h b/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h
new file mode 100644
index 0000000000..1cd81eeaeb
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h
@@ -0,0 +1,47 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+#ifdef GGML_USE_HIP
+#define GGML_CUDA_NAME "ROCm"
+#define GGML_CUBLAS_NAME "hipBLAS"
+#elif defined(GGML_USE_MUSA)
+#define GGML_CUDA_NAME "MUSA"
+#define GGML_CUBLAS_NAME "muBLAS"
+#else
+#define GGML_CUDA_NAME "CUDA"
+#define GGML_CUBLAS_NAME "cuBLAS"
+#endif
+#define GGML_CUDA_MAX_DEVICES       16
+
+// backend API
+GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device);
+
+GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend);
+
+// device buffer
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device);
+
+// conduct allreduce operation between devices
+GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends);
+
+// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void);
+
+GGML_BACKEND_API int  ggml_backend_cuda_get_device_count(void);
+GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size);
+GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total);
+
+GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size);
+GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer);
+
+GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void);
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml.h b/python/freetoken/kernel/csrc/ggml/include/ggml.h
new file mode 100644
index 0000000000..23ea6437e2
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/ggml.h
@@ -0,0 +1,3027 @@
+#pragma once
+
+//
+// GGML Tensor Library
+//
+// This documentation is still a work in progress.
+// If you wish some specific topics to be covered, feel free to drop a comment:
+//
+//   https://github.com/ggml-org/whisper.cpp/issues/40
+//
+// ## Overview
+//
+// This library implements:
+//
+//  - a set of tensor operations
+//  - automatic differentiation
+//  - basic optimization algorithms
+//
+// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes,
+// but is not limited to, the following:
+//
+//  - linear regression
+//  - support vector machines
+//  - neural networks
+//
+// The library allows the user to define a certain function using the available tensor operations. This function
+// definition is represented internally via a computation graph. Each tensor operation in the function definition
+// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the
+// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized
+// using one of the available optimization algorithms.
+//
+// For example, here we define the function: f(x) = a*x^2 + b
+//
+//   {
+//       struct ggml_init_params params = {
+//           .mem_size   = 16*1024*1024,
+//           .mem_buffer = NULL,
+//       };
+//
+//       // memory allocation happens here
+//       struct ggml_context * ctx = ggml_init(params);
+//
+//       struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//
+//       ggml_set_param(ctx, x); // x is an input variable
+//
+//       struct ggml_tensor * a  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//       struct ggml_tensor * b  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//       struct ggml_tensor * x2 = ggml_mul(ctx, x, x);
+//       struct ggml_tensor * f  = ggml_add(ctx, ggml_mul(ctx, a, x2), b);
+//
+//       ...
+//   }
+//
+// Notice that the function definition above does not involve any actual computation. The computation is performed only
+// when the user explicitly requests it. For example, to compute the function's value at x = 2.0:
+//
+//   {
+//       ...
+//
+//       struct ggml_cgraph * gf = ggml_new_graph(ctx);
+//       ggml_build_forward_expand(gf, f);
+//
+//       // set the input variable and parameter values
+//       ggml_set_f32(x, 2.0f);
+//       ggml_set_f32(a, 3.0f);
+//       ggml_set_f32(b, 4.0f);
+//
+//       ggml_graph_compute_with_ctx(ctx, &gf, n_threads);
+//
+//       printf("f = %f\n", ggml_get_f32_1d(f, 0));
+//
+//       ...
+//   }
+//
+// The actual computation is performed in the ggml_graph_compute() function.
+//
+// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the
+// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know
+// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory
+// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was
+// actually needed.
+//
+// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic
+// differentiation and optimization algorithms.
+//
+// The described approach allows to define the function graph once and then compute its forward or backward graphs
+// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way
+// the user can avoid the memory allocation overhead at runtime.
+//
+// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class
+// citizens, but in theory the library can be extended to support FP8 and integer data types.
+//
+// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary
+// and binary operations. Most of the available operations fall into one of these two categories. With time, it became
+// clear that the library needs to support more complex operations. The way to support these operations is not clear
+// yet, but a few examples are demonstrated in the following operations:
+//
+//   - ggml_permute()
+//   - ggml_conv_1d_1s()
+//   - ggml_conv_1d_2s()
+//
+// For each tensor operator, the library implements a forward and backward computation function. The forward function
+// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the
+// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a
+// calculus class, or watch the following video:
+//
+//   What is Automatic Differentiation?
+//   https://www.youtube.com/watch?v=wG_nF1awSSY
+//
+//
+// ## Tensor data (struct ggml_tensor)
+//
+// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of
+// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains
+// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example:
+//
+//   {
+//       struct ggml_tensor * c = ggml_add(ctx, a, b);
+//
+//       assert(c->src[0] == a);
+//       assert(c->src[1] == b);
+//   }
+//
+// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the
+// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows
+// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and
+// permutation. All tensor operations have to take the stride into account and not assume that the tensor is
+// contiguous in memory.
+//
+// The data of the tensor is accessed via the "data" pointer. For example:
+//
+//   {
+//       const int nx = 2;
+//       const int ny = 3;
+//
+//       struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny);
+//
+//       for (int y = 0; y < ny; y++) {
+//           for (int x = 0; x < nx; x++) {
+//               *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y;
+//           }
+//       }
+//
+//       ...
+//   }
+//
+// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used.
+//
+// ## The matrix multiplication operator (ggml_mul_mat)
+//
+// TODO
+//
+//
+// ## Multi-threading
+//
+// TODO
+//
+//
+// ## Overview of ggml.c
+//
+// TODO
+//
+//
+// ## SIMD optimizations
+//
+// TODO
+//
+//
+// ## Debugging ggml
+//
+// TODO
+//
+//
+
+#ifdef GGML_SHARED
+#    if defined(_WIN32) && !defined(__MINGW32__)
+#        ifdef GGML_BUILD
+#            define GGML_API __declspec(dllexport) extern
+#        else
+#            define GGML_API __declspec(dllimport) extern
+#        endif
+#    else
+#        define GGML_API __attribute__ ((visibility ("default"))) extern
+#    endif
+#else
+#    define GGML_API extern
+#endif
+
+// TODO: support for clang
+#ifdef __GNUC__
+#    define GGML_DEPRECATED(func, hint) func __attribute__((deprecated(hint)))
+#elif defined(_MSC_VER)
+#    define GGML_DEPRECATED(func, hint) __declspec(deprecated(hint)) func
+#else
+#    define GGML_DEPRECATED(func, hint) func
+#endif
+
+#ifndef __GNUC__
+#    define GGML_ATTRIBUTE_FORMAT(...)
+#elif defined(__MINGW32__) && !defined(__clang__)
+#    define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__)))
+#else
+#    define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__)))
+#endif
+
+#if defined(_WIN32) && !defined(_WIN32_WINNT)
+#    define _WIN32_WINNT 0x0A00
+#endif
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#define GGML_FILE_MAGIC   0x67676d6c // "ggml"
+#define GGML_FILE_VERSION 2
+
+#define GGML_QNT_VERSION        2    // bump this on quantization format changes
+#define GGML_QNT_VERSION_FACTOR 1000 // do not change this
+
+#define GGML_MAX_DIMS           4
+#define GGML_MAX_PARAMS         2048
+#define GGML_MAX_SRC            10
+#define GGML_MAX_N_THREADS      512
+#define GGML_MAX_OP_PARAMS      64
+
+#ifndef GGML_MAX_NAME
+#   define GGML_MAX_NAME        64
+#endif
+
+#define GGML_DEFAULT_N_THREADS  4
+#define GGML_DEFAULT_GRAPH_SIZE 2048
+
+#if UINTPTR_MAX == 0xFFFFFFFF
+    #define GGML_MEM_ALIGN 4
+#elif defined(__EMSCRIPTEN__)
+// emscripten uses max_align_t == 8, so we need GGML_MEM_ALIGN == 8 for 64-bit wasm.
+// (for 32-bit wasm, the first conditional is true and GGML_MEM_ALIGN stays 4.)
+// ref: https://github.com/ggml-org/llama.cpp/pull/18628
+    #define GGML_MEM_ALIGN 8
+#else
+    #define GGML_MEM_ALIGN 16
+#endif
+
+#define GGML_EXIT_SUCCESS 0
+#define GGML_EXIT_ABORTED 1
+
+// TODO: convert to enum https://github.com/ggml-org/llama.cpp/pull/16187#discussion_r2388538726
+#define GGML_ROPE_TYPE_NORMAL 0
+#define GGML_ROPE_TYPE_NEOX   2
+#define GGML_ROPE_TYPE_MROPE  8
+#define GGML_ROPE_TYPE_VISION 24
+#define GGML_ROPE_TYPE_IMROPE 40 // binary: 101000
+
+#define GGML_MROPE_SECTIONS   4
+
+#define GGML_UNUSED(x) (void)(x)
+#ifdef __CUDACC__
+template<typename... Args>
+__host__ __device__ constexpr inline void ggml_unused_vars_impl(Args&&...) noexcept {}
+#define GGML_UNUSED_VARS(...) ggml_unused_vars_impl(__VA_ARGS__)
+#else
+#define GGML_UNUSED_VARS(...) do { (void)sizeof((__VA_ARGS__, 0)); } while(0)
+#endif // __CUDACC__
+
+#define GGML_PAD(x, n) (((x) + (n) - 1) & ~((n) - 1))
+
+#ifndef NDEBUG
+#   define GGML_UNREACHABLE() do { fprintf(stderr, "statement should be unreachable\n"); abort(); } while(0)
+#elif defined(__GNUC__)
+#   define GGML_UNREACHABLE() __builtin_unreachable()
+#elif defined(_MSC_VER)
+#   define GGML_UNREACHABLE() __assume(0)
+#else
+#   define GGML_UNREACHABLE() ((void) 0)
+#endif
+
+#ifdef __cplusplus
+#   define GGML_NORETURN [[noreturn]]
+#elif defined(_MSC_VER)
+#   define GGML_NORETURN __declspec(noreturn)
+#else
+#   define GGML_NORETURN _Noreturn
+#endif
+
+#define GGML_ABORT(...) ggml_abort(__FILE__, __LINE__, __VA_ARGS__)
+#define GGML_ASSERT(x) if (!(x)) GGML_ABORT("GGML_ASSERT(%s) failed", #x)
+
+// used to copy the number of elements and stride in bytes of tensors into local variables.
+// main purpose is to reduce code duplication and improve readability.
+//
+// example:
+//
+//    GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne);
+//    GGML_TENSOR_LOCALS(size_t,  nb1, src1, nb);
+//
+#define GGML_TENSOR_LOCALS_1(type, prefix, pointer, array) \
+    const type prefix##0 = (pointer) ? (pointer)->array[0] : 0; \
+    GGML_UNUSED(prefix##0);
+#define GGML_TENSOR_LOCALS_2(type, prefix, pointer, array) \
+    GGML_TENSOR_LOCALS_1    (type, prefix, pointer, array) \
+    const type prefix##1 = (pointer) ? (pointer)->array[1] : 0; \
+    GGML_UNUSED(prefix##1);
+#define GGML_TENSOR_LOCALS_3(type, prefix, pointer, array) \
+    GGML_TENSOR_LOCALS_2    (type, prefix, pointer, array) \
+    const type prefix##2 = (pointer) ? (pointer)->array[2] : 0; \
+    GGML_UNUSED(prefix##2);
+#define GGML_TENSOR_LOCALS(type, prefix, pointer, array) \
+    GGML_TENSOR_LOCALS_3  (type, prefix, pointer, array) \
+    const type prefix##3 = (pointer) ? (pointer)->array[3] : 0; \
+    GGML_UNUSED(prefix##3);
+
+#define GGML_TENSOR_UNARY_OP_LOCALS \
+    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne,  dst,  ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb,  dst,  nb)
+
+#define GGML_TENSOR_BINARY_OP_LOCALS \
+    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb1, src1, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne,  dst,  ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb,  dst,  nb)
+
+#define GGML_TENSOR_TERNARY_OP_LOCALS \
+    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb1, src1, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne2, src2, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb2, src2, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne,  dst,  ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb,  dst,  nb)
+
+#define GGML_TENSOR_BINARY_OP_LOCALS01 \
+    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb) \
+    GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \
+    GGML_TENSOR_LOCALS(size_t,  nb1, src1, nb)
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    // Function type used in fatal error callbacks
+    typedef void (*ggml_abort_callback_t)(const char * error_message);
+
+    // Set the abort callback (passing null will restore original abort functionality: printing a message to stdout)
+    // Returns the old callback for chaining
+    GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback);
+
+    GGML_NORETURN GGML_ATTRIBUTE_FORMAT(3, 4)
+    GGML_API void ggml_abort(const char * file, int line, const char * fmt, ...);
+
+    enum ggml_status {
+        GGML_STATUS_ALLOC_FAILED = -2,
+        GGML_STATUS_FAILED = -1,
+        GGML_STATUS_SUCCESS = 0,
+        GGML_STATUS_ABORTED = 1,
+    };
+
+    // get ggml_status name string
+    GGML_API const char * ggml_status_to_string(enum ggml_status status);
+
+    // ieee 754-2008 half-precision float16
+    // todo: make this not an integral type
+    typedef uint16_t ggml_fp16_t;
+    GGML_API float       ggml_fp16_to_fp32(ggml_fp16_t);
+    GGML_API ggml_fp16_t ggml_fp32_to_fp16(float);
+    GGML_API void        ggml_fp16_to_fp32_row(const ggml_fp16_t *, float *, int64_t);
+    GGML_API void        ggml_fp32_to_fp16_row(const float *, ggml_fp16_t *, int64_t);
+
+    // google brain half-precision bfloat16
+    typedef struct { uint16_t bits; } ggml_bf16_t;
+    GGML_API ggml_bf16_t ggml_fp32_to_bf16(float);
+    GGML_API float       ggml_bf16_to_fp32(ggml_bf16_t);  // consider just doing << 16
+    GGML_API void        ggml_bf16_to_fp32_row(const ggml_bf16_t *, float *, int64_t);
+    GGML_API void        ggml_fp32_to_bf16_row_ref(const float *, ggml_bf16_t *, int64_t);
+    GGML_API void        ggml_fp32_to_bf16_row(const float *, ggml_bf16_t *, int64_t);
+
+    struct ggml_object;
+    struct ggml_context;
+    struct ggml_cgraph;
+
+    // NOTE: always add types at the end of the enum to keep backward compatibility
+    enum ggml_type {
+        GGML_TYPE_F32     = 0,
+        GGML_TYPE_F16     = 1,
+        GGML_TYPE_Q4_0    = 2,
+        GGML_TYPE_Q4_1    = 3,
+        // GGML_TYPE_Q4_2 = 4, support has been removed
+        // GGML_TYPE_Q4_3 = 5, support has been removed
+        GGML_TYPE_Q5_0    = 6,
+        GGML_TYPE_Q5_1    = 7,
+        GGML_TYPE_Q8_0    = 8,
+        GGML_TYPE_Q8_1    = 9,
+        GGML_TYPE_Q2_K    = 10,
+        GGML_TYPE_Q3_K    = 11,
+        GGML_TYPE_Q4_K    = 12,
+        GGML_TYPE_Q5_K    = 13,
+        GGML_TYPE_Q6_K    = 14,
+        GGML_TYPE_Q8_K    = 15,
+        GGML_TYPE_IQ2_XXS = 16,
+        GGML_TYPE_IQ2_XS  = 17,
+        GGML_TYPE_IQ3_XXS = 18,
+        GGML_TYPE_IQ1_S   = 19,
+        GGML_TYPE_IQ4_NL  = 20,
+        GGML_TYPE_IQ3_S   = 21,
+        GGML_TYPE_IQ2_S   = 22,
+        GGML_TYPE_IQ4_XS  = 23,
+        GGML_TYPE_I8      = 24,
+        GGML_TYPE_I16     = 25,
+        GGML_TYPE_I32     = 26,
+        GGML_TYPE_I64     = 27,
+        GGML_TYPE_F64     = 28,
+        GGML_TYPE_IQ1_M   = 29,
+        GGML_TYPE_BF16    = 30,
+        // GGML_TYPE_Q4_0_4_4 = 31, support has been removed from gguf files
+        // GGML_TYPE_Q4_0_4_8 = 32,
+        // GGML_TYPE_Q4_0_8_8 = 33,
+        GGML_TYPE_TQ1_0   = 34,
+        GGML_TYPE_TQ2_0   = 35,
+        // GGML_TYPE_IQ4_NL_4_4 = 36,
+        // GGML_TYPE_IQ4_NL_4_8 = 37,
+        // GGML_TYPE_IQ4_NL_8_8 = 38,
+        GGML_TYPE_MXFP4   = 39, // MXFP4 (1 block)
+        GGML_TYPE_NVFP4   = 40, // NVFP4 (4 blocks, E4M3 scale)
+        GGML_TYPE_Q1_0    = 41,
+        GGML_TYPE_Q2_0    = 42,
+        GGML_TYPE_COUNT   = 43,
+    };
+
+    // [TAG_GGML_PREC]
+    // this enum is used to declare the allowed numerical precision/data-types types that can be used during the compute of an op
+    // the declared types can be:
+    //  - result accumulation type
+    //  - source tensor data representation type
+    //  - etc.
+    // the precision parameters are stored as ggml_tensor.op_params to the respective ops
+    enum ggml_prec {
+        GGML_PREC_UNDEFINED = 0,
+        GGML_PREC_DEFAULT   = 0,  // note: deprecated, use GGML_PREC_UNDEFINED
+        GGML_PREC_F32       = 10,
+        GGML_PREC_BF16      = 15,
+        GGML_PREC_F16       = 20,
+        GGML_PREC_Q8        = 30,
+        GGML_PREC_Q4        = 40,
+    };
+
+    // op hint
+    enum ggml_op_hint {
+        GGML_HINT_NONE             = 0,
+        GGML_HINT_SRC0_IS_HADAMARD = 1,
+    };
+
+    // model file types
+    enum ggml_ftype {
+        GGML_FTYPE_UNKNOWN        = -1,
+        GGML_FTYPE_ALL_F32        = 0,
+        GGML_FTYPE_MOSTLY_F16     = 1,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q4_0    = 2,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q4_1    = 3,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16
+        GGML_FTYPE_MOSTLY_Q8_0    = 7,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q5_0    = 8,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q5_1    = 9,  // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q2_K    = 10, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q3_K    = 11, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q4_K    = 12, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q5_K    = 13, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q6_K    = 14, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ2_XXS = 15, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ2_XS  = 16, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ3_XXS = 17, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ1_S   = 18, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ4_NL  = 19, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ3_S   = 20, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ2_S   = 21, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ4_XS  = 22, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ1_M   = 23, // except 1d tensors
+        GGML_FTYPE_MOSTLY_BF16    = 24, // except 1d tensors
+        GGML_FTYPE_MOSTLY_MXFP4   = 25, // except 1d tensors
+        GGML_FTYPE_MOSTLY_NVFP4   = 26, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q1_0    = 27, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q2_0    = 28, // except 1d tensors
+    };
+
+    // available tensor operations:
+    enum ggml_op {
+        GGML_OP_NONE = 0,
+
+        GGML_OP_DUP,
+        GGML_OP_ADD,
+        GGML_OP_ADD_ID,
+        GGML_OP_ADD1,
+        GGML_OP_ACC,
+        GGML_OP_SUB,
+        GGML_OP_MUL,
+        GGML_OP_DIV,
+        GGML_OP_SQR,
+        GGML_OP_SQRT,
+        GGML_OP_LOG,
+        GGML_OP_SIN,
+        GGML_OP_COS,
+        GGML_OP_SUM,
+        GGML_OP_SUM_ROWS,
+        GGML_OP_CUMSUM,
+        GGML_OP_MEAN,
+        GGML_OP_ARGMAX,
+        GGML_OP_COUNT_EQUAL,
+        GGML_OP_REPEAT,
+        GGML_OP_REPEAT_BACK,
+        GGML_OP_CONCAT,
+        GGML_OP_SILU_BACK,
+        GGML_OP_NORM, // normalize
+        GGML_OP_RMS_NORM,
+        GGML_OP_RMS_NORM_BACK,
+        GGML_OP_GROUP_NORM,
+        GGML_OP_L2_NORM,
+
+        GGML_OP_MUL_MAT,
+        GGML_OP_MUL_MAT_ID,
+        GGML_OP_OUT_PROD,
+
+        GGML_OP_SCALE,
+        GGML_OP_SET,
+        GGML_OP_CPY,
+        GGML_OP_CONT,
+        GGML_OP_RESHAPE,
+        GGML_OP_VIEW,
+        GGML_OP_PERMUTE,
+        GGML_OP_TRANSPOSE,
+        GGML_OP_GET_ROWS,
+        GGML_OP_GET_ROWS_BACK,
+        GGML_OP_SET_ROWS,
+        GGML_OP_DIAG,
+        GGML_OP_DIAG_MASK_INF,
+        GGML_OP_DIAG_MASK_ZERO,
+        GGML_OP_SOFT_MAX,
+        GGML_OP_SOFT_MAX_BACK,
+        GGML_OP_ROPE,
+        GGML_OP_ROPE_BACK,
+        GGML_OP_CLAMP,
+        GGML_OP_CONV_TRANSPOSE_1D,
+        GGML_OP_IM2COL,
+        GGML_OP_IM2COL_BACK,
+        GGML_OP_IM2COL_3D,
+        GGML_OP_COL2IM_1D,
+        GGML_OP_CONV_2D,
+        GGML_OP_CONV_3D,
+        GGML_OP_CONV_2D_DW,
+        GGML_OP_CONV_TRANSPOSE_2D,
+        GGML_OP_POOL_1D,
+        GGML_OP_POOL_2D,
+        GGML_OP_POOL_2D_BACK,
+        GGML_OP_UPSCALE,
+        GGML_OP_PAD,
+        GGML_OP_PAD_REFLECT_1D,
+        GGML_OP_ROLL,
+        GGML_OP_ARANGE,
+        GGML_OP_TIMESTEP_EMBEDDING,
+        GGML_OP_ARGSORT,
+        GGML_OP_TOP_K,
+        GGML_OP_LEAKY_RELU,
+        GGML_OP_TRI,
+        GGML_OP_FILL,
+
+        GGML_OP_FLASH_ATTN_EXT,
+        GGML_OP_FLASH_ATTN_BACK,
+        GGML_OP_SSM_CONV,
+        GGML_OP_SSM_SCAN,
+        GGML_OP_WIN_PART,
+        GGML_OP_WIN_UNPART,
+        GGML_OP_GET_REL_POS,
+        GGML_OP_ADD_REL_POS,
+        GGML_OP_RWKV_WKV6,
+        GGML_OP_GATED_LINEAR_ATTN,
+        GGML_OP_RWKV_WKV7,
+        GGML_OP_SOLVE_TRI,
+        GGML_OP_GATED_DELTA_NET,
+        GGML_OP_LIGHTNING_INDEXER,
+        GGML_OP_DSV4_HC_COMB,
+        GGML_OP_DSV4_HC_PRE,
+        GGML_OP_DSV4_HC_POST,
+
+        GGML_OP_UNARY,
+
+        GGML_OP_MAP_CUSTOM1,
+        GGML_OP_MAP_CUSTOM2,
+        GGML_OP_MAP_CUSTOM3,
+
+        GGML_OP_CUSTOM,
+
+        GGML_OP_CROSS_ENTROPY_LOSS,
+        GGML_OP_CROSS_ENTROPY_LOSS_BACK,
+        GGML_OP_OPT_STEP_ADAMW,
+        GGML_OP_OPT_STEP_SGD,
+
+        GGML_OP_GLU,
+
+        GGML_OP_COUNT,
+    };
+
+    enum ggml_unary_op {
+        GGML_UNARY_OP_ABS,
+        GGML_UNARY_OP_SGN,
+        GGML_UNARY_OP_NEG,
+        GGML_UNARY_OP_STEP,
+        GGML_UNARY_OP_TANH,
+        GGML_UNARY_OP_ELU,
+        GGML_UNARY_OP_RELU,
+        GGML_UNARY_OP_SIGMOID,
+        GGML_UNARY_OP_GELU,
+        GGML_UNARY_OP_GELU_QUICK,
+        GGML_UNARY_OP_SILU,
+        GGML_UNARY_OP_HARDSWISH,
+        GGML_UNARY_OP_HARDSIGMOID,
+        GGML_UNARY_OP_EXP,
+        GGML_UNARY_OP_EXPM1,
+        GGML_UNARY_OP_SOFTPLUS,
+        GGML_UNARY_OP_GELU_ERF,
+        GGML_UNARY_OP_XIELU,
+        GGML_UNARY_OP_FLOOR,
+        GGML_UNARY_OP_CEIL,
+        GGML_UNARY_OP_ROUND,
+        GGML_UNARY_OP_TRUNC,
+
+        GGML_UNARY_OP_COUNT,
+    };
+
+    enum ggml_glu_op {
+        GGML_GLU_OP_REGLU,
+        GGML_GLU_OP_GEGLU,
+        GGML_GLU_OP_SWIGLU,
+        GGML_GLU_OP_SWIGLU_OAI,
+        GGML_GLU_OP_GEGLU_ERF,
+        GGML_GLU_OP_GEGLU_QUICK,
+        GGML_GLU_OP_SWIGLU_CLAMP,
+
+        GGML_GLU_OP_COUNT,
+    };
+
+    enum ggml_object_type {
+        GGML_OBJECT_TYPE_TENSOR,
+        GGML_OBJECT_TYPE_GRAPH,
+        GGML_OBJECT_TYPE_WORK_BUFFER
+    };
+
+    enum ggml_log_level {
+        GGML_LOG_LEVEL_NONE  = 0,
+        GGML_LOG_LEVEL_DEBUG = 1,
+        GGML_LOG_LEVEL_INFO  = 2,
+        GGML_LOG_LEVEL_WARN  = 3,
+        GGML_LOG_LEVEL_ERROR = 4,
+        GGML_LOG_LEVEL_CONT  = 5, // continue previous log
+    };
+
+    // this tensor...
+    enum ggml_tensor_flag {
+        GGML_TENSOR_FLAG_INPUT   =  1, // ...is an input for the GGML compute graph
+        GGML_TENSOR_FLAG_OUTPUT  =  2, // ...is an output for the GGML compute graph
+        GGML_TENSOR_FLAG_PARAM   =  4, // ...contains trainable parameters
+        GGML_TENSOR_FLAG_LOSS    =  8, // ...defines loss for numerical optimization (multiple loss tensors add up)
+        GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed
+    };
+
+    enum ggml_tri_type {
+        GGML_TRI_TYPE_UPPER_DIAG = 0,
+        GGML_TRI_TYPE_UPPER      = 1,
+        GGML_TRI_TYPE_LOWER_DIAG = 2,
+        GGML_TRI_TYPE_LOWER      = 3
+    };
+
+    struct ggml_init_params {
+        // memory pool
+        size_t mem_size;   // bytes
+        void * mem_buffer; // if NULL, memory will be allocated internally
+        bool   no_alloc;   // don't allocate memory for the tensor data
+    };
+
+    // n-dimensional tensor
+    struct ggml_tensor {
+        enum ggml_type type;
+
+        struct ggml_backend_buffer * buffer;
+
+        int64_t ne[GGML_MAX_DIMS]; // number of elements
+        size_t  nb[GGML_MAX_DIMS]; // stride in bytes:
+                                   // nb[0] = ggml_type_size(type)
+                                   // nb[1] = nb[0]   * (ne[0] / ggml_blck_size(type)) + padding
+                                   // nb[i] = nb[i-1] * ne[i-1]
+
+        // compute data
+        enum ggml_op op;
+
+        // op params - allocated as int32_t for alignment
+        int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
+
+        int32_t flags;
+
+        struct ggml_tensor * src[GGML_MAX_SRC];
+
+        // source tensor and offset for views
+        struct ggml_tensor * view_src;
+        size_t               view_offs;
+
+        void * data;
+
+        char name[GGML_MAX_NAME];
+
+        void * extra; // extra things e.g. for ggml-cuda.cu
+
+        char padding[8];
+    };
+
+    static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor);
+
+    // Abort callback
+    // If not NULL, called before ggml computation
+    // If it returns true, the computation is aborted
+    typedef bool (*ggml_abort_callback)(void * data);
+
+
+    //
+    // GUID
+    //
+
+    // GUID types
+    typedef uint8_t ggml_guid[16];
+    typedef ggml_guid * ggml_guid_t;
+
+    GGML_API bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b);
+
+    // misc
+
+    GGML_API const char * ggml_version(void);
+    GGML_API const char * ggml_commit(void);
+
+    GGML_API void    ggml_time_init(void); // call this once at the beginning of the program
+    GGML_API int64_t ggml_time_ms(void);
+    GGML_API int64_t ggml_time_us(void);
+    GGML_API int64_t ggml_cycles(void);
+    GGML_API int64_t ggml_cycles_per_ms(void);
+
+    // accepts a UTF-8 path, even on Windows
+    GGML_API FILE *  ggml_fopen(const char * fname, const char * mode);
+
+    GGML_API void    ggml_print_object (const struct ggml_object * obj);
+    GGML_API void    ggml_print_objects(const struct ggml_context * ctx);
+
+    GGML_API int64_t ggml_nelements (const struct ggml_tensor * tensor);
+    GGML_API int64_t ggml_nrows     (const struct ggml_tensor * tensor);
+    GGML_API size_t  ggml_nbytes    (const struct ggml_tensor * tensor);
+    GGML_API size_t  ggml_nbytes_pad(const struct ggml_tensor * tensor); // same as ggml_nbytes() but padded to GGML_MEM_ALIGN
+
+    GGML_API int64_t ggml_blck_size(enum ggml_type type);
+    GGML_API size_t  ggml_type_size(enum ggml_type type);             // size in bytes for all elements in a block
+    GGML_API size_t  ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row
+
+    GGML_DEPRECATED(
+    GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float
+    "use ggml_row_size() instead");
+
+    GGML_API const char * ggml_type_name(enum ggml_type type);
+    GGML_API const char * ggml_op_name  (enum ggml_op   op);
+    GGML_API const char * ggml_op_symbol(enum ggml_op   op);
+
+    GGML_API const char * ggml_unary_op_name(enum ggml_unary_op op);
+    GGML_API const char * ggml_glu_op_name(enum ggml_glu_op op);
+    GGML_API const char * ggml_op_desc(const struct ggml_tensor * t); // unary or op name
+
+    GGML_API size_t  ggml_element_size(const struct ggml_tensor * tensor);
+
+    GGML_API bool    ggml_is_quantized(enum ggml_type type);
+
+    // TODO: temporary until model loading of ggml examples is refactored
+    GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype);
+
+    GGML_API bool ggml_is_transposed(const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_permuted  (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_empty     (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_view      (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_scalar    (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_vector    (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_matrix    (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_3d        (const struct ggml_tensor * tensor);
+    GGML_API int  ggml_n_dims       (const struct ggml_tensor * tensor); // returns 1 for scalars
+
+    // returns whether the tensor elements can be iterated over with a flattened index (no gaps, no permutation)
+    GGML_API bool ggml_is_contiguous  (const struct ggml_tensor * tensor);
+    GGML_API bool ggml_is_contiguous_0(const struct ggml_tensor * tensor); // same as ggml_is_contiguous()
+    GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1
+    GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2
+
+    GGML_API bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor); // contiguous for dims < 1
+    GGML_API bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor); // contiguous for dims < 2
+    GGML_API bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor); // contiguous for dims < 3
+
+    // returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok)
+    GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor);
+
+    // true for tensor that is stored in memory as CxWxHxN and has been permuted to WxHxCxN
+    GGML_API bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor);
+
+    // true if the elements in dimension 0 are contiguous, or there is just 1 block of elements
+    GGML_API bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor);
+
+    GGML_API bool ggml_are_same_shape (const struct ggml_tensor * t0, const struct ggml_tensor * t1);
+    GGML_API bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1);
+
+    GGML_API bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1);
+
+    // use this to compute the memory overhead of a tensor
+    GGML_API size_t ggml_tensor_overhead(void);
+
+    GGML_API bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes);
+
+    // main
+
+    GGML_API struct ggml_context * ggml_init (struct ggml_init_params params);
+    GGML_API void                  ggml_reset(struct ggml_context * ctx);
+    GGML_API void                  ggml_free (struct ggml_context * ctx);
+
+    GGML_API size_t  ggml_used_mem(const struct ggml_context * ctx);
+
+    GGML_API bool    ggml_get_no_alloc(struct ggml_context * ctx);
+    GGML_API void    ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc);
+
+    GGML_API void *  ggml_get_mem_buffer     (const struct ggml_context * ctx);
+    GGML_API size_t  ggml_get_mem_size       (const struct ggml_context * ctx);
+    GGML_API size_t  ggml_get_max_tensor_size(const struct ggml_context * ctx);
+
+    GGML_API struct ggml_tensor * ggml_new_tensor(
+            struct ggml_context * ctx,
+            enum   ggml_type type,
+            int    n_dims,
+            const int64_t *ne);
+
+    GGML_API struct ggml_tensor * ggml_new_tensor_1d(
+            struct ggml_context * ctx,
+            enum   ggml_type type,
+            int64_t ne0);
+
+    GGML_API struct ggml_tensor * ggml_new_tensor_2d(
+            struct ggml_context * ctx,
+            enum   ggml_type type,
+            int64_t ne0,
+            int64_t ne1);
+
+    GGML_API struct ggml_tensor * ggml_new_tensor_3d(
+            struct ggml_context * ctx,
+            enum   ggml_type type,
+            int64_t ne0,
+            int64_t ne1,
+            int64_t ne2);
+
+    GGML_API struct ggml_tensor * ggml_new_tensor_4d(
+            struct ggml_context * ctx,
+            enum   ggml_type type,
+            int64_t ne0,
+            int64_t ne1,
+            int64_t ne2,
+            int64_t ne3);
+
+    GGML_API void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes);
+
+    GGML_API struct ggml_tensor * ggml_dup_tensor (struct ggml_context * ctx, const struct ggml_tensor * src);
+    GGML_API struct ggml_tensor * ggml_view_tensor(struct ggml_context * ctx, struct ggml_tensor * src);
+
+    // Context tensor enumeration and lookup
+    GGML_API struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx);
+    GGML_API struct ggml_tensor * ggml_get_next_tensor (const struct ggml_context * ctx, struct ggml_tensor * tensor);
+    GGML_API struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name);
+
+    // Converts a flat index into coordinates
+    GGML_API void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3);
+
+    GGML_API enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor);
+    GGML_API enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor);
+
+    GGML_API void *  ggml_get_data    (const struct ggml_tensor * tensor);
+    GGML_API float * ggml_get_data_f32(const struct ggml_tensor * tensor);
+
+    GGML_API const char *         ggml_get_name   (const struct ggml_tensor * tensor);
+    GGML_API struct ggml_tensor * ggml_set_name   (      struct ggml_tensor * tensor, const char * name);
+    GGML_ATTRIBUTE_FORMAT(2, 3)
+    GGML_API struct ggml_tensor * ggml_format_name(      struct ggml_tensor * tensor, const char * fmt, ...);
+
+    // Tensor flags
+    GGML_API void ggml_set_input(struct ggml_tensor * tensor);
+    GGML_API void ggml_set_output(struct ggml_tensor * tensor);
+    GGML_API void ggml_set_param(struct ggml_tensor * tensor);
+    GGML_API void ggml_set_loss(struct ggml_tensor * tensor);
+
+    //
+    // operations on tensors with backpropagation
+    //
+
+    GGML_API struct ggml_tensor * ggml_dup(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_dup_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_add(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_add_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_add_cast(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            enum   ggml_type      type);
+
+    // dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]]
+    GGML_API struct ggml_tensor * ggml_add_id(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * ids);
+
+    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b),
+        "use ggml_add instead");
+
+    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b),
+        "use ggml_add_inplace instead");
+
+    // dst = a
+    // view(dst, nb1, nb2, nb3, offset) += b
+    // return dst
+    GGML_API struct ggml_tensor * ggml_acc(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                nb2,
+            size_t                nb3,
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_acc_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                nb2,
+            size_t                nb3,
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_sub(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_sub_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_mul(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_mul_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_div(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_div_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_sqr(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sqr_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sqrt(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sqrt_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_log(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_log_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_expm1(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_expm1_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_softplus(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_softplus_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sin(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sin_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_cos(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_cos_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // return scalar
+    GGML_API struct ggml_tensor * ggml_sum(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // sums along rows, with input shape [a,b,c,d] return shape [1,b,c,d]
+    GGML_API struct ggml_tensor * ggml_sum_rows(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_cumsum(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a);
+
+    // mean along rows
+    GGML_API struct ggml_tensor * ggml_mean(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // argmax along rows
+    GGML_API struct ggml_tensor * ggml_argmax(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // count number of equal elements in a and b
+    GGML_API struct ggml_tensor * ggml_count_equal(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // if a is the same shape as b, and a is not parameter, return a
+    // otherwise, return a new tensor: repeat(a) to fit in b
+    GGML_API struct ggml_tensor * ggml_repeat(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // repeat a to the specified shape
+    GGML_API struct ggml_tensor * ggml_repeat_4d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+                       int64_t    ne0,
+                       int64_t    ne1,
+                       int64_t    ne2,
+                       int64_t    ne3);
+
+    // sums repetitions in a into shape of b
+    GGML_API struct ggml_tensor * ggml_repeat_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b); // sum up values that are adjacent in dims > 0 instead of repeated with same stride
+
+    // concat a and b along dim
+    // used in stable-diffusion
+    GGML_API struct ggml_tensor * ggml_concat(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   dim);
+
+    GGML_API struct ggml_tensor * ggml_abs(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_abs_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sgn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sgn_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_neg(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_neg_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_step(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_step_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_tanh(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_tanh_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_elu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_elu_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_relu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_leaky_relu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a, float negative_slope, bool inplace);
+
+    GGML_API struct ggml_tensor * ggml_relu_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sigmoid(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_sigmoid_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_gelu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_gelu_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // GELU using erf (error function) when possible
+    // some backends may fallback to approximation based on Abramowitz and Stegun formula
+    GGML_API struct ggml_tensor * ggml_gelu_erf(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_gelu_erf_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_gelu_quick(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_gelu_quick_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_silu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_silu_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // a - dy
+    // b - x
+    GGML_API struct ggml_tensor * ggml_silu_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // hardswish(x) = x * relu6(x + 3) / 6
+    GGML_API struct ggml_tensor * ggml_hardswish(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // hardsigmoid(x) = relu6(x + 3) / 6
+    GGML_API struct ggml_tensor * ggml_hardsigmoid(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_exp(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_exp_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_floor(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_floor_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_ceil(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_ceil_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_round(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_round_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+     /**
+     * Truncates the fractional part of each element in the tensor (towards zero).
+     * For example: trunc(3.7) = 3.0, trunc(-2.9) = -2.0
+     * Similar to std::trunc in C/C++.
+     */
+
+    GGML_API struct ggml_tensor * ggml_trunc(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_trunc_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+
+
+    // xIELU activation function
+    // x = x * (c_a(alpha_n) + c_b(alpha_p, beta) * sigmoid(beta * x)) + eps * (x > 0)
+    // where c_a = softplus and c_b(a, b) = softplus(a) + b are constraining functions
+    // that constrain the positive and negative source alpha values respectively
+    GGML_API struct ggml_tensor * ggml_xielu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float alpha_n,
+            float alpha_p,
+            float beta,
+            float eps);
+
+    // gated linear unit ops
+    // A: n columns, r rows,
+    // result is n / 2 columns, r rows,
+    // expects gate in second half of row, unless swapped is true
+    GGML_API struct ggml_tensor * ggml_glu(
+            struct ggml_context * ctx,
+             struct ggml_tensor * a,
+             enum ggml_glu_op     op,
+             bool                 swapped);
+
+    GGML_API struct ggml_tensor * ggml_reglu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_reglu_swapped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu_swapped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_swiglu(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_swiglu_swapped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu_erf(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu_erf_swapped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu_quick(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_geglu_quick_swapped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // A: n columns, r rows,
+    // B: n columns, r rows,
+    GGML_API struct ggml_tensor * ggml_glu_split(
+            struct ggml_context * ctx,
+             struct ggml_tensor * a,
+             struct ggml_tensor * b,
+             enum ggml_glu_op     op);
+
+    GGML_API struct ggml_tensor * ggml_reglu_split(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_geglu_split(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_swiglu_split(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_geglu_erf_split(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_geglu_quick_split(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    GGML_API struct ggml_tensor * ggml_swiglu_oai(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            float                 alpha,
+            float                 limit);
+
+    GGML_API struct ggml_tensor * ggml_swiglu_clamp(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            float                 limit);
+
+    // normalize along rows
+    GGML_API struct ggml_tensor * ggml_norm(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    GGML_API struct ggml_tensor * ggml_norm_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    GGML_API struct ggml_tensor * ggml_rms_norm(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    GGML_API struct ggml_tensor * ggml_rms_norm_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    // group normalize along ne0*ne1*n_groups
+    // used in stable-diffusion
+    GGML_API struct ggml_tensor * ggml_group_norm(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_groups,
+            float                 eps);
+
+    GGML_API struct ggml_tensor * ggml_group_norm_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_groups,
+            float                 eps);
+
+    // l2 normalize along rows
+    // used in rwkv v7
+    GGML_API struct ggml_tensor * ggml_l2_norm(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    GGML_API struct ggml_tensor * ggml_l2_norm_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 eps);
+
+    // a - x
+    // b - dy
+    GGML_API struct ggml_tensor * ggml_rms_norm_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            float                 eps);
+
+    // [TAG_GGML_PREC]
+    // set the minimum required accumulator type for the implementation to use during the compute
+    // for example:
+    //  - GGML_PREC_F32  - requires accumulation of the results in F32
+    //  - GGML_PREC_BF16 - can accumulate the results in BF16, F32
+    //  - GGML_PREC_F16  - can accumulate the results in F16, F32
+    //  - GGML_PREC_Q8   - not allowed
+    //  - GGML_PREC_Q4   - not allowed
+    //
+    // return false on faliure
+    GGML_API bool ggml_prec_set_acc(
+            struct ggml_tensor * a,
+            enum ggml_prec       prec);
+
+    // [TAG_GGML_PREC]
+    // set the smallest rank that the implementation can use to internally convert the src[idx] data to
+    // ranks in decreasing order:
+    //  - GGML_PREC_F32  - GGML_TYPE_F32
+    //  - GGML_PREC_BF16 - GGML_TYPE_BF16
+    //  - GGML_PREC_F16  - GGML_TYPE_F16,
+    //  - GGML_PREC_Q8   - GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, GGML_TYPE_Q8_K, etc.
+    //  - GGML_PREC_Q4   - GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_NVFP4, GGML_TYPE_MXFP4, etc.
+    //
+    // for example:
+    //   - ggml_prec_set_src(a, GGML_PREC_Q8, 1):
+    //     - allows the implementation to quantize F32, BF16, F16 data of src[1] down to GGML_TYPE_Q8_0
+    //     - cannot quantize it down to GGML_TYPE_Q4_0 or GGML_TYPE_NVFP4
+    //   - ggml_prec_set_src(a, GGML_PREC_Q4, 1):
+    //     - allows the implementation to quantize F32, BF16, F16 data of src[1] down to 4-bit datatypes such as GGML_TYPE_Q4_K, GGML_TYPE_NVFP4 etc.
+    //
+    // return false on faliure
+    GGML_API bool ggml_prec_set_src(
+            struct ggml_tensor * a,
+            enum ggml_prec       prec,
+            int                  idx);
+
+    // A: k columns, n rows => [ne03, ne02, n, k]
+    // B: k columns, m rows  (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k]
+    // result is n columns, m rows => [ne03 * x, ne02 * y, m, n]
+    GGML_API struct ggml_tensor * ggml_mul_mat(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // change the precision of a matrix multiplication
+    // set to GGML_PREC_F32 for higher precision (useful for phi-2)
+    GGML_DEPRECATED(GGML_API void ggml_mul_mat_set_prec(
+            struct ggml_tensor * a,
+            enum ggml_prec       prec),
+        "use ggml_prec_set_acc() instead");
+
+    // change the hint of a matrix multiplication
+    GGML_API void ggml_mul_mat_set_hint(
+            struct ggml_tensor * a,
+            enum ggml_op_hint    hint);
+
+    // indirect matrix multiplication
+    //
+    // [TAG_MUL_MAT_ID_DUPLICATE_ID]
+    // note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
+    //       (i.e. duplicate experts per token are not guaranteed to be supported in all backends)
+    // TODO: add backend support: https://github.com/ggml-org/llama.cpp/pull/29998#issuecomment-6153742533
+    GGML_API struct ggml_tensor * ggml_mul_mat_id(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * as,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * ids);
+
+    // A: m columns, n rows,
+    // B: p columns, n rows,
+    // result is m columns, p rows
+    GGML_API struct ggml_tensor * ggml_out_prod(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    //
+    // operations on tensors without backpropagation
+    //
+
+    GGML_API struct ggml_tensor * ggml_scale(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 s);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_scale_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 s);
+
+    // x = s * a + b
+    GGML_API struct ggml_tensor * ggml_scale_bias(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        float                 s,
+        float                 b);
+
+    GGML_API struct ggml_tensor * ggml_scale_bias_inplace(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        float                 s,
+        float                 b);
+
+    // b -> view(a,offset,nb1,nb2,3), return modified a
+    GGML_API struct ggml_tensor * ggml_set(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                nb2,
+            size_t                nb3,
+            size_t                offset); // in bytes
+
+    // b -> view(a,offset,nb1,nb2,3), return view(a)
+    GGML_API struct ggml_tensor * ggml_set_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                nb2,
+            size_t                nb3,
+            size_t                offset); // in bytes
+
+    GGML_API struct ggml_tensor * ggml_set_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                offset); // in bytes
+
+    GGML_API struct ggml_tensor * ggml_set_1d_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                offset); // in bytes
+
+    // b -> view(a,offset,nb1,nb2,3), return modified a
+    GGML_API struct ggml_tensor * ggml_set_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                offset); // in bytes
+
+    // b -> view(a,offset,nb1,nb2,3), return view(a)
+    GGML_API struct ggml_tensor * ggml_set_2d_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            size_t                nb1,
+            size_t                offset); // in bytes
+
+    // a -> b, return view(b)
+    GGML_API struct ggml_tensor * ggml_cpy(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // note: casting from f32 to i32 will discard the fractional part
+    GGML_API struct ggml_tensor * ggml_cast(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            enum   ggml_type      type);
+
+    // make contiguous
+    GGML_API struct ggml_tensor * ggml_cont(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // make contiguous, with new shape
+    GGML_API struct ggml_tensor * ggml_cont_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0);
+
+    GGML_API struct ggml_tensor * ggml_cont_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1);
+
+    GGML_API struct ggml_tensor * ggml_cont_3d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2);
+
+    GGML_API struct ggml_tensor * ggml_cont_4d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            int64_t               ne3);
+
+    // return view(a), b specifies the new shape
+    // TODO: when we start computing gradient, make a copy instead of view
+    GGML_API struct ggml_tensor * ggml_reshape(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // return view(a)
+    // TODO: when we start computing gradient, make a copy instead of view
+    GGML_API struct ggml_tensor * ggml_reshape_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0);
+
+    GGML_API struct ggml_tensor * ggml_reshape_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1);
+
+    // return view(a)
+    // TODO: when we start computing gradient, make a copy instead of view
+    GGML_API struct ggml_tensor * ggml_reshape_3d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2);
+
+    GGML_API struct ggml_tensor * ggml_reshape_4d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            int64_t               ne3);
+
+    // offset in bytes
+    GGML_API struct ggml_tensor * ggml_view_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_view_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            size_t                nb1, // row stride in bytes
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_view_3d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            size_t                nb1, // row   stride in bytes
+            size_t                nb2, // slice stride in bytes
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_view_4d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            int64_t               ne3,
+            size_t                nb1, // row   stride in bytes
+            size_t                nb2, // slice stride in bytes
+            size_t                nb3,
+            size_t                offset);
+
+    GGML_API struct ggml_tensor * ggml_permute(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   axis0,
+            int                   axis1,
+            int                   axis2,
+            int                   axis3);
+
+    // alias for ggml_permute(ctx, a, 1, 0, 2, 3)
+    GGML_API struct ggml_tensor * ggml_transpose(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // supports 4D a:
+    // a     [n_embd, ne1, ne2, ne3]
+    // b I32 [n_rows, ne2, ne3, 1]
+    //
+    // return [n_embd, n_rows, ne2, ne3]
+    GGML_API struct ggml_tensor * ggml_get_rows(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // data
+            struct ggml_tensor  * b); // row indices
+
+    GGML_API struct ggml_tensor * ggml_get_rows_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // gradients of ggml_get_rows result
+            struct ggml_tensor  * b,  // row indices
+            struct ggml_tensor  * c); // data for ggml_get_rows, only used for its shape
+
+    // a TD  [n_embd, ne1,    ne2,    ne3]
+    // b TS  [n_embd, n_rows, ne02,   ne03] | ne02 == ne2, ne03 == ne3
+    // c I64 [n_rows, ne11,   ne12,   1]    | c[i] in [0, ne1)
+    //
+    // undefined behavior if destination rows overlap
+    //
+    // broadcast:
+    //   ne2 % ne11 == 0
+    //   ne3 % ne12 == 0
+    //
+    // return view(a)
+    GGML_API struct ggml_tensor * ggml_set_rows(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // destination
+            struct ggml_tensor  * b,  // source
+            struct ggml_tensor  * c); // row indices
+
+    GGML_API struct ggml_tensor * ggml_diag(
+        struct ggml_context     * ctx,
+        struct ggml_tensor      * a);
+
+    // set elements above the diagonal to -INF
+    GGML_API struct ggml_tensor * ggml_diag_mask_inf(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_past);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_past);
+
+    // set elements above the diagonal to 0
+    GGML_API struct ggml_tensor * ggml_diag_mask_zero(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_past);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_past);
+
+    GGML_API struct ggml_tensor * ggml_clamp(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 min,
+            float                 max);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_clamp_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 min,
+            float                 max);
+
+    GGML_API struct ggml_tensor * ggml_soft_max(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_soft_max_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    // a    [ne0, ne01, ne02, ne03]
+    // mask [ne0, ne11, ne12, ne13] | ne11 >= ne01, F16 or F32, optional
+    //
+    // broadcast:
+    //   ne02 % ne12 == 0
+    //   ne03 % ne13 == 0
+    //
+    // fused soft_max(a*scale + mask*(ALiBi slope))
+    // max_bias = 0.0f for no ALiBi
+    GGML_API struct ggml_tensor * ggml_soft_max_ext(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * mask,
+            float                 scale,
+            float                 max_bias);
+
+    GGML_API struct ggml_tensor * ggml_soft_max_ext_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * mask,
+            float                 scale,
+            float                 max_bias);
+
+    GGML_API void ggml_soft_max_add_sinks(
+            struct ggml_tensor * a,
+            struct ggml_tensor * sinks);
+
+    GGML_API struct ggml_tensor * ggml_soft_max_ext_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            float                 scale,
+            float                 max_bias);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_soft_max_ext_back_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            float                 scale,
+            float                 max_bias);
+
+    // rotary position embedding
+    // if (mode & 1) - skip n_past elements (NOT SUPPORTED)
+    // if (mode & GGML_ROPE_TYPE_NEOX) - GPT-NeoX style
+    //
+    // b is an int32 vector with size a->ne[2], it contains the positions
+    GGML_API struct ggml_tensor * ggml_rope(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   n_dims,
+            int                   mode);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_rope_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   n_dims,
+            int                   mode);
+
+    // RoPE operations with extended options
+    // a is the input tensor to apply RoPE to, shape [n_embd, n_head, n_token]
+    // b is an int32 vector with size n_token
+    // c is freq factors (e.g. phi3-128k), (optional)
+    // mode can be GGML_ROPE_TYPE_NORMAL or NEOX; for MROPE and VISION mode, use ggml_rope_multi
+    //
+    // pseudo-code for computing theta:
+    //   for i in [0, n_dims/2):
+    //     theta[i] = b[i] * powf(freq_base, -2.0 * i / n_dims);
+    //     theta[i] = theta[i] / c[i];  # if c is provided, divide theta by c
+    //     theta[i] = rope_yarn(theta[i], ...);  # note: theta = theta * freq_scale is applied here
+    //
+    // other params are used by YaRN RoPE scaling, these default values will disable YaRN:
+    //   freq_scale  = 1.0f
+    //   ext_factor  = 0.0f
+    //   attn_factor = 1.0f
+    //   beta_fast   = 0.0f
+    //   beta_slow   = 0.0f
+    //
+    // example:
+    //   (marking: c = cos, s = sin, 0 = unrotated)
+    //   given a single head with size = 8 --> [00000000]
+    //   GGML_ROPE_TYPE_NORMAL  n_dims = 4 --> [cscs0000]
+    //   GGML_ROPE_TYPE_NORMAL  n_dims = 8 --> [cscscscs]
+    //   GGML_ROPE_TYPE_NEOX    n_dims = 4 --> [ccss0000]
+    //   GGML_ROPE_TYPE_NEOX    n_dims = 8 --> [ccccssss]
+    GGML_API struct ggml_tensor * ggml_rope_ext(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * c,
+            int                   n_dims,
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    // multi-dimensional RoPE, for Qwen-VL and similar vision models
+    // mode can be either VISION, MROPE, IMROPE, cannot be combined with NORMAL or NEOX
+    // sections specify how many dimensions to rotate in each section:
+    //   section length is equivalent to number of cos/sin pairs, NOT the number of dims
+    //   (i.e. sum of 4 sections are expected to be n_dims/2)
+    //   last sections can be 0, means ignored
+    // all other options are identical to ggml_rope_ext
+    //
+    // important note:
+    //   - NEOX ordering is automatically applied and cannot be disabled for MROPE and VISION
+    //     if you need normal ordering, there are 2 methods:
+    //     (1) split the tensor manually using ggml_view
+    //     (2) permute the weight upon conversion
+    //   - for VISION, n_dims must be head_size/2
+    //
+    // example M-RoPE:
+    //  given sections = [t=4, y=2, x=2, 0]
+    //  given a single head with size = 18 --> [000000000000000000]
+    //  GGML_ROPE_TYPE_MROPE   n_dims = 16 --> [ttttyyxxttttyyxx00] (cos/sin are applied in NEOX ordering)
+    //  GGML_ROPE_TYPE_IMROPE  n_dims = 16 --> [ttyxttyxttyxttyx00] (interleaved M-RoPE, still NEOX ordering)
+    //  note: the theta for each dim is computed the same way as ggml_rope_ext, no matter the section
+    //        in other words, idx used for theta: [0123456789... until n_dims/2], not reset for each section
+    //
+    // example vision RoPE:
+    //  given sections = [y=4, x=4, 0, 0] (last 2 sections are ignored)
+    //  given a single head with size = 8 --> [00000000]
+    //  GGML_ROPE_TYPE_VISION  n_dims = 4 --> [yyyyxxxx]
+    //  other values of n_dims are untested and is undefined behavior
+    //  note: unlike MROPE, the theta for each dim is computed differently for each section
+    //        in other words, idx used for theta: [0123] for y section, then [0123] for x section
+    GGML_API struct ggml_tensor * ggml_rope_multi(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * c,
+            int                   n_dims,
+            int                   sections[GGML_MROPE_SECTIONS],
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_rope_ext_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * c,
+            int                   n_dims,
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    GGML_API struct ggml_tensor * ggml_rope_multi_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * c,
+            int                   n_dims,
+            int                   sections[GGML_MROPE_SECTIONS],
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   n_dims,
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow),
+        "use ggml_rope_ext instead");
+
+    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   n_dims,
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow),
+        "use ggml_rope_ext_inplace instead");
+
+    // compute correction dims for YaRN RoPE scaling
+    GGML_API void ggml_rope_yarn_corr_dims(
+        int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]);
+
+    // rotary position embedding backward, i.e compute dx from dy
+    // a - dy
+    GGML_API struct ggml_tensor * ggml_rope_ext_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a, // gradients of ggml_rope result
+            struct ggml_tensor  * b, // positions
+            struct ggml_tensor  * c, // freq factors
+            int                   n_dims,
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    GGML_API struct ggml_tensor * ggml_rope_multi_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * c,
+            int                   n_dims,
+            int                   sections[4],
+            int                   mode,
+            int                   n_ctx_orig,
+            float                 freq_base,
+            float                 freq_scale,
+            float                 ext_factor,
+            float                 attn_factor,
+            float                 beta_fast,
+            float                 beta_slow);
+
+    // set the offset dims for RoPE
+    // a must be GGML_OP_ROPE or GGML_OP_ROPE_BACK
+    // vision RoPE is not supported
+    // example: (marking: x = rotated, 0 = unrotated)
+    //     n_embd = 10, n_dims = 4, offset = 2 --> [00xxxx0000]
+    GGML_API struct ggml_tensor * ggml_rope_set_offset(
+            struct ggml_tensor  * a,
+            int                   n_offs);
+
+    // im2col
+    // converts data into a format that effectively results in a convolution when combined with matrix multiplication
+    GGML_API struct ggml_tensor * ggml_im2col(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // convolution kernel
+            struct ggml_tensor  * b,  // data
+            int                   s0, // stride dimension 0
+            int                   s1, // stride dimension 1
+            int                   p0, // padding dimension 0
+            int                   p1, // padding dimension 1
+            int                   d0, // dilation dimension 0
+            int                   d1, // dilation dimension 1
+            bool                  is_2D,
+            enum ggml_type        dst_type);
+
+    GGML_API struct ggml_tensor * ggml_im2col_back(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,  // convolution kernel
+        struct ggml_tensor  * b,  // gradient of im2col output
+        int64_t             * ne, // shape of im2col input
+        int                   s0, // stride dimension 0
+        int                   s1, // stride dimension 1
+        int                   p0, // padding dimension 0
+        int                   p1, // padding dimension 1
+        int                   d0, // dilation dimension 0
+        int                   d1, // dilation dimension 1
+        bool                  is_2D);
+
+    // col2im_1d: scatter-add GEMM columns back to 1D signal
+    // a: [K*OC, T_in]  (columns from matmul, K = a->ne[0]/OC)
+    // result: [T_out, OC]  where T_out = (T_in - 1)*s0 + K - 2*p0
+    GGML_API struct ggml_tensor * ggml_col2im_1d(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,   // columns [K*OC, T_in]
+        int                   s0,  // stride
+        int                   oc,  // output channels
+        int                   p0); // padding to crop from both sides
+
+    GGML_API struct ggml_tensor * ggml_conv_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel
+            struct ggml_tensor  * b,   // data
+            int                   s0,  // stride
+            int                   p0,  // padding
+            int                   d0); // dilation
+
+    // conv_1d with padding = half
+    // alias for ggml_conv_1d(a, b, s, a->ne[0]/2, d)
+    GGML_API struct ggml_tensor* ggml_conv_1d_ph(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // convolution kernel
+            struct ggml_tensor  * b,  // data
+            int                   s,  // stride
+            int                   d); // dilation
+
+    // depthwise
+    // TODO: this is very likely wrong for some cases! - needs more testing
+    GGML_API struct ggml_tensor * ggml_conv_1d_dw(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel
+            struct ggml_tensor  * b,   // data
+            int                   s0,  // stride
+            int                   p0,  // padding
+            int                   d0); // dilation
+
+    GGML_API struct ggml_tensor * ggml_conv_1d_dw_ph(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel
+            struct ggml_tensor  * b,   // data
+            int                   s0,  // stride
+            int                   d0); // dilation
+
+    GGML_API struct ggml_tensor * ggml_conv_transpose_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel
+            struct ggml_tensor  * b,   // data
+            int                   s0,  // stride
+            int                   p0,  // padding
+            int                   d0); // dilation
+
+    GGML_API struct ggml_tensor * ggml_conv_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel
+            struct ggml_tensor  * b,   // data
+            int                   s0,  // stride dimension 0
+            int                   s1,  // stride dimension 1
+            int                   p0,  // padding dimension 0
+            int                   p1,  // padding dimension 1
+            int                   d0,  // dilation dimension 0
+            int                   d1); // dilation dimension 1
+
+    GGML_API struct ggml_tensor * ggml_im2col_3d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int64_t               IC,
+            int                   s0, // stride width
+            int                   s1, // stride height
+            int                   s2, // stride depth
+            int                   p0, // padding width
+            int                   p1, // padding height
+            int                   p2, // padding depth
+            int                   d0, // dilation width
+            int                   d1, // dilation height
+            int                   d2, // dilation depth
+            enum ggml_type        dst_type);
+
+    // a: [OC*IC, KD, KH, KW]
+    // b: [N*IC, ID, IH, IW]
+    // result: [N*OC, OD, OH, OW]
+    GGML_API struct ggml_tensor * ggml_conv_3d(
+                struct ggml_context * ctx,
+                struct ggml_tensor  * a,
+                struct ggml_tensor  * b,
+                int64_t               IC,
+                int                   s0, // stride width
+                int                   s1, // stride height
+                int                   s2, // stride depth
+                int                   p0, // padding width
+                int                   p1, // padding height
+                int                   p2, // padding depth
+                int                   d0, // dilation width
+                int                   d1, // dilation height
+                int                   d2  // dilation depth
+        );
+
+    // kernel size is a->ne[0] x a->ne[1]
+    // stride is equal to kernel size
+    // padding is zero
+    // example:
+    // a:     16   16    3  768
+    // b:   1024 1024    3    1
+    // res:   64   64  768    1
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_conv_2d_sk_p0(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // kernel size is a->ne[0] x a->ne[1]
+    // stride is 1
+    // padding is half
+    // example:
+    // a:      3    3    256  256
+    // b:     64   64    256    1
+    // res:   64   64    256    1
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_conv_2d_s1_ph(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
+    // depthwise (via im2col and mul_mat)
+    GGML_API struct ggml_tensor * ggml_conv_2d_dw(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // convolution kernel
+            struct ggml_tensor  * b,  // data
+            int                  s0,  // stride dimension 0
+            int                  s1,  // stride dimension 1
+            int                  p0,  // padding dimension 0
+            int                  p1,  // padding dimension 1
+            int                  d0,  // dilation dimension 0
+            int                  d1); // dilation dimension 1
+
+    // Depthwise 2D convolution
+    // may be faster than ggml_conv_2d_dw, but not available in all backends
+    // a:   KW    KH    1    C    convolution kernel
+    // b:   W     H     C    N    input data
+    // res: W_out H_out C    N
+    GGML_API struct ggml_tensor * ggml_conv_2d_dw_direct(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   stride0,
+            int                   stride1,
+            int                   pad0,
+            int                   pad1,
+            int                   dilation0,
+            int                   dilation1);
+
+    GGML_API struct ggml_tensor * ggml_conv_transpose_2d_p0(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            int                   stride);
+
+    GGML_API struct ggml_tensor * ggml_conv_2d_direct(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // convolution kernel [KW, KH, IC, OC]
+            struct ggml_tensor  * b,   // input data [W, H, C, N]
+            int                   s0,  // stride dimension 0
+            int                   s1,  // stride dimension 1
+            int                   p0,  // padding dimension 0
+            int                   p1,  // padding dimension 1
+            int                   d0,  // dilation dimension 0
+            int                   d1); // dilation dimension 1
+
+    GGML_API struct ggml_tensor * ggml_conv_3d_direct(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // kernel [KW, KH, KD, IC * OC]
+            struct ggml_tensor  * b,   // input  [W, H, D, C * N]
+            int                   s0,  // stride
+            int                   s1,
+            int                   s2,
+            int                   p0,  // padding
+            int                   p1,
+            int                   p2,
+            int                   d0,  // dilation
+            int                   d1,
+            int                   d2,
+            int                   n_channels,
+            int                   n_batch,
+            int                   n_channels_out);
+
+    enum ggml_op_pool {
+        GGML_OP_POOL_MAX,
+        GGML_OP_POOL_AVG,
+        GGML_OP_POOL_COUNT,
+    };
+
+    GGML_API struct ggml_tensor * ggml_pool_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            enum ggml_op_pool     op,
+            int                   k0, // kernel size
+            int                   s0, // stride
+            int                   p0); // padding
+
+    // the result will have 2*p0 padding for the first dimension
+    // and 2*p1 padding for the second dimension
+    GGML_API struct ggml_tensor * ggml_pool_2d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            enum ggml_op_pool     op,
+            int                   k0,
+            int                   k1,
+            int                   s0,
+            int                   s1,
+            float                 p0,
+            float                 p1);
+
+    GGML_API struct ggml_tensor * ggml_pool_2d_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * af, // "a"/input used in forward pass
+            enum ggml_op_pool     op,
+            int                   k0,
+            int                   k1,
+            int                   s0,
+            int                   s1,
+            float                 p0,
+            float                 p1);
+
+    enum ggml_scale_mode {
+        GGML_SCALE_MODE_NEAREST  = 0,
+        GGML_SCALE_MODE_BILINEAR = 1,
+        GGML_SCALE_MODE_BICUBIC  = 2,
+
+        GGML_SCALE_MODE_COUNT
+    };
+
+    enum ggml_scale_flag {
+        GGML_SCALE_FLAG_ALIGN_CORNERS = (1 << 8),
+        GGML_SCALE_FLAG_ANTIALIAS     = (1 << 9),
+    };
+
+    // interpolate
+    // multiplies ne0 and ne1 by scale factor
+    GGML_API struct ggml_tensor * ggml_upscale(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   scale_factor,
+            enum ggml_scale_mode  mode);
+
+    // interpolate
+    // interpolate scale to specified dimensions
+    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   ne0,
+            int                   ne1,
+            int                   ne2,
+            int                   ne3,
+            enum ggml_scale_mode  mode),
+        "use ggml_interpolate instead");
+
+    // Up- or downsamples the input to the specified size.
+    // 2D scale modes (eg. bilinear) are applied to the first two dimensions.
+    GGML_API struct ggml_tensor * ggml_interpolate(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            int64_t               ne3,
+            uint32_t              mode); // ggml_scale_mode [ | ggml_scale_flag...]
+
+    // pad each dimension with zeros: [x, ..., x] -> [x, ..., x, 0, ..., 0]
+    GGML_API struct ggml_tensor * ggml_pad(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                  p0,
+            int                  p1,
+            int                  p2,
+            int                  p3);
+
+    // pad each dimension with values on the other side of the torus (looping around)
+    GGML_API struct ggml_tensor * ggml_pad_circular(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   p0,
+            int                   p1,
+            int                   p2,
+            int                   p3);
+
+    GGML_API struct ggml_tensor * ggml_pad_ext(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                  lp0,
+            int                  rp0,
+            int                  lp1,
+            int                  rp1,
+            int                  lp2,
+            int                  rp2,
+            int                  lp3,
+            int                  rp3
+            );
+
+    // pad each dimension with values on the other side of the torus (looping around)
+    GGML_API struct ggml_tensor * ggml_pad_ext_circular(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   lp0,
+            int                   rp0,
+            int                   lp1,
+            int                   rp1,
+            int                   lp2,
+            int                   rp2,
+            int                   lp3,
+            int                   rp3);
+
+    // pad each dimension with reflection: [a, b, c, d] -> [b, a, b, c, d, c]
+    GGML_API struct ggml_tensor * ggml_pad_reflect_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   p0,
+            int                   p1);
+
+    // Move tensor elements by an offset given for each dimension. Elements that
+    // are shifted beyond the last position are wrapped around to the beginning.
+    GGML_API struct ggml_tensor * ggml_roll(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   shift0,
+            int                   shift1,
+            int                   shift2,
+            int                   shift3);
+
+    // Convert matrix into a triangular one (upper, strict upper, lower or strict lower) by writing
+    // zeroes everywhere outside the masked area
+    GGML_API struct ggml_tensor * ggml_tri(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            enum ggml_tri_type    type);
+
+    // Fill tensor a with constant c
+    GGML_API struct ggml_tensor * ggml_fill(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 c);
+
+    GGML_API struct ggml_tensor * ggml_fill_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 c);
+
+    // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
+    // timesteps: [N,]
+    // return: [N, dim]
+    GGML_API struct ggml_tensor * ggml_timestep_embedding(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * timesteps,
+            int                   dim,
+            int                   max_period);
+
+    // sort rows
+    enum ggml_sort_order {
+        GGML_SORT_ORDER_ASC,
+        GGML_SORT_ORDER_DESC,
+    };
+
+    GGML_API struct ggml_tensor * ggml_argsort(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            enum ggml_sort_order  order);
+
+    // similar to ggml_top_k but implemented as `argsort` + `view`
+    GGML_API struct ggml_tensor * ggml_argsort_top_k(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   k);
+
+    // top k elements per row
+    // note: the resulting top k indices are in no particular order
+    GGML_API struct ggml_tensor * ggml_top_k(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   k);
+
+    GGML_API struct ggml_tensor * ggml_arange(
+            struct ggml_context * ctx,
+            float                 start,
+            float                 stop,
+            float                 step);
+
+    // q:    [n_embd_k, n_batch, n_head,    ne3 ]
+    // k:    [n_embd_k, n_kv,    n_head_kv, ne3 ]
+    // v:    [n_embd_v, n_kv,    n_head_kv, ne3 ] !! not transposed !!
+    // mask: [n_kv,     n_batch, ne32,      ne33]
+    // res:  [n_embd_v, n_head,  n_batch,   ne3 ] !! permuted !!
+    //
+    // broadcast:
+    //   n_head % n_head_kv == 0
+    //   n_head % ne32      == 0
+    //   ne3    % ne33      == 0
+    //
+    GGML_API struct ggml_tensor * ggml_flash_attn_ext(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * mask,
+            float                 scale,
+            float                 max_bias,
+            float                 logit_softcap);
+
+    GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec(
+            struct ggml_tensor * a,
+            enum ggml_prec       prec),
+        "use ggml_prec_set_acc() instead");
+
+    GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
+            const struct ggml_tensor * a);
+
+    // Use finite mask entries as a sparse K/V set. Set 0 to disable.
+    // n_kv_max must bound the number of finite entries in every mask row.
+    GGML_API void ggml_flash_attn_ext_set_n_kv_max(
+            struct ggml_tensor * a,
+            int32_t              n_kv_max);
+
+    GGML_API void ggml_flash_attn_ext_add_sinks(
+            struct ggml_tensor * a,
+            struct ggml_tensor * sinks);
+
+    // TODO: needs to be adapted to ggml_flash_attn_ext
+    GGML_API struct ggml_tensor * ggml_flash_attn_back(
+           struct ggml_context * ctx,
+           struct ggml_tensor  * q,
+           struct ggml_tensor  * k,
+           struct ggml_tensor  * v,
+           struct ggml_tensor  * d,
+           bool                  masked);
+
+    GGML_API struct ggml_tensor * ggml_ssm_conv(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * sx,
+            struct ggml_tensor  * c);
+
+    GGML_API struct ggml_tensor * ggml_ssm_scan(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * s,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * dt,
+            struct ggml_tensor  * A,
+            struct ggml_tensor  * B,
+            struct ggml_tensor  * C,
+            struct ggml_tensor  * ids,
+            int64_t               K);
+
+    // partition into non-overlapping windows with padding if needed
+    // example:
+    // a:   768   64   64    1
+    // w:    14
+    // res: 768   14   14    25
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_win_part(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   w);
+
+    // reverse of ggml_win_part
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_win_unpart(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   w0,
+            int                   h0,
+            int                   w);
+
+    GGML_API struct ggml_tensor * ggml_unary(
+            struct ggml_context * ctx,
+             struct ggml_tensor * a,
+             enum ggml_unary_op op);
+
+    GGML_API struct ggml_tensor * ggml_unary_inplace(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        enum ggml_unary_op op);
+
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_get_rel_pos(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   qh,
+            int                   kh);
+
+    // used in sam
+    GGML_API struct ggml_tensor * ggml_add_rel_pos(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * pw,
+            struct ggml_tensor  * ph);
+
+    GGML_API struct ggml_tensor * ggml_add_rel_pos_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * pw,
+            struct ggml_tensor  * ph);
+
+    GGML_API struct ggml_tensor * ggml_rwkv_wkv6(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * r,
+            struct ggml_tensor  * tf,
+            struct ggml_tensor  * td,
+            struct ggml_tensor  * state);
+
+    GGML_API struct ggml_tensor * ggml_gated_linear_attn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * g,
+            struct ggml_tensor  * state,
+            float scale);
+
+    GGML_API struct ggml_tensor * ggml_rwkv_wkv7(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * r,
+            struct ggml_tensor  * w,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * state);
+
+    /* Solves a specific equation of the form Ax=B, where A is a triangular matrix
+    *  without zeroes on the diagonal (i.e. invertible).
+    *  B can have any number of columns, but must have the same number of rows as A
+    *  If A is [n, n] and B is [n, m], then the result will be [n, m] as well
+    *  Has O(n^3) complexity (unlike most matrix ops out there), so use on cases
+    *  where n > 100 sparingly, pre-chunk if necessary.
+    *
+    *  If left = false, solves xA=B instead
+    *  If lower = false, assumes upper triangular instead
+    *  If uni = true, assumes diagonal of A to be all ones (will override actual values)
+    *
+    *  TODO: currently only lower, right, non-unitriangular variant is implemented
+    */
+    GGML_API struct ggml_tensor * ggml_solve_tri(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        struct ggml_tensor  * b,
+        bool                  left,
+        bool                  lower,
+        bool                  uni);
+
+    // TODO: add ggml_gated_delta_net_set_bcast() to be able to configure Q, K broadcast type: tiled vs interleaved [TAG_GGML_GDN_BCAST]
+    // ref: https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306
+    //
+    // tensor shapes (S_k == S_v, H_v % H_k == 0):
+    //   q, k  : [S_k, H_k, n_tokens, n_seqs]
+    //   v     : [S_v, H_v, n_tokens, n_seqs]
+    //   g     : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA)
+    //   beta  : [1, H_v, n_tokens, n_seqs]
+    //   state : [S_v, S_v, H_v, n_seqs] -- initial recurrent state s0
+    //
+    // the output packs the attention scores [S_v, H_v, n_tokens, n_seqs] followed by K state
+    // snapshots, most-recent first (slot 0 = final state, slot s = state s tokens back). K == 1
+    // keeps only the final state; when n_tokens < K only slots 0..n_tokens-1 are written.
+    GGML_API struct ggml_tensor * ggml_gated_delta_net(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * g,
+            struct ggml_tensor  * beta,
+            struct ggml_tensor  * state,
+            int64_t               K);
+
+    // DSA lightning indexer
+    //
+    // q:       [n_embd_idx, n_head_idx, n_batch, ne3 ]
+    // k:       [n_embd_idx, 1,          n_kv,    ne3 ]
+    // weights: [n_head_idx, n_batch,    1,       ne3 ] !! prescaled !!
+    // mask:    [n_kv,       n_batch,    1,       ne33] !! f16 !!
+    // res:     [n_kv,       n_batch,    1,       ne3 ]
+    //
+    // broadcast:
+    //   ne3 % ne33 == 0
+    //
+    GGML_API struct ggml_tensor * ggml_lightning_indexer(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * q,
+        struct ggml_tensor  * k,
+        struct ggml_tensor  * weights,
+        struct ggml_tensor  * mask);
+
+    // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
+    // In short these operations are replacements for the original residual connection (x = transformer(x) + x)
+    // using a richer representation through streams.
+    //
+    // hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc]
+    //          -> [dst_hc, src_hc, n_tokens]
+    // logits[dst, src, t] = mixes[2*hc + dst + hc*src, t]*scale[2]
+    //                         + base[2*hc + dst + hc*src]
+    // Softmax over dst, add eps, normalize over src, then repeat normalization
+    // over dst followed by src for iterations 1 through n_iter - 1.
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_comb(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mixes,
+            struct ggml_tensor  * scale,
+            struct ggml_tensor  * base,
+            float                 eps,
+            int32_t               n_iter);
+
+    // hc_pre: x [n_embd, hc, n_tokens], weights [hc, n_tokens] -> [n_embd, n_tokens]
+    //   result[i, t] = sum_h x[i, h, t]*weights[h, t]
+    //
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_pre(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * weights);
+
+    // hc_pre with a per-element gate (Qwen3.8-Flash-Next): gate [n_embd, hc, n_tokens]
+    //   result[i, t] = scale*sum_h x[i, h, t]*sigmoid(gate[i, h, t])
+    //
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_pre_gated(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * gate,
+            float                 scale);
+
+    // hc_post: x [n_embd, n_tokens], residual [n_embd, hc, n_tokens],
+    //          post [hc, n_tokens], comb [dst_hc, src_hc, n_tokens]
+    //          -> [n_embd, hc, n_tokens]
+    //   result[i, dst, t] = x[i, t]*post[dst, t]
+    //                       + sum_src residual[i, src, t]*comb[dst, src, t]
+    //   comb == NULL uses the identity: result[i, dst, t] = x[i, t]*post[dst, t] + residual[i, dst, t]
+    //
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_post(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * residual,
+            struct ggml_tensor  * post,
+            struct ggml_tensor  * comb);
+
+    // custom operators
+
+    typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
+    typedef void (*ggml_custom2_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata);
+    typedef void (*ggml_custom3_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata);
+
+#define GGML_N_TASKS_MAX (-1)
+    // n_tasks == GGML_N_TASKS_MAX means to use max number of tasks
+
+    GGML_API struct ggml_tensor * ggml_map_custom1(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            ggml_custom1_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    GGML_API struct ggml_tensor * ggml_map_custom1_inplace(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            ggml_custom1_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    GGML_API struct ggml_tensor * ggml_map_custom2(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            struct ggml_tensor    * b,
+            ggml_custom2_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    GGML_API struct ggml_tensor * ggml_map_custom2_inplace(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            struct ggml_tensor    * b,
+            ggml_custom2_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    GGML_API struct ggml_tensor * ggml_map_custom3(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            struct ggml_tensor    * b,
+            struct ggml_tensor    * c,
+            ggml_custom3_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    GGML_API struct ggml_tensor * ggml_map_custom3_inplace(
+            struct ggml_context   * ctx,
+            struct ggml_tensor    * a,
+            struct ggml_tensor    * b,
+            struct ggml_tensor    * c,
+            ggml_custom3_op_t       fun,
+            int                     n_tasks,
+            void                  * userdata);
+
+    typedef void (*ggml_custom_op_t)(struct ggml_tensor * dst , int ith, int nth, void * userdata);
+
+    GGML_API struct ggml_tensor * ggml_custom_4d(
+            struct ggml_context * ctx,
+            enum ggml_type        type,
+            int64_t               ne0,
+            int64_t               ne1,
+            int64_t               ne2,
+            int64_t               ne3,
+            struct ggml_tensor ** args,
+            int                   n_args,
+            ggml_custom_op_t      fun,
+            int                   n_tasks,
+            void                * userdata);
+
+    GGML_API struct ggml_tensor * ggml_custom_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor ** args,
+            int                   n_args,
+            ggml_custom_op_t      fun,
+            int                   n_tasks,
+            void                * userdata);
+
+    // loss function
+
+    GGML_API struct ggml_tensor * ggml_cross_entropy_loss(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // logits
+            struct ggml_tensor  * b); // labels
+
+    GGML_API struct ggml_tensor * ggml_cross_entropy_loss_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,  // logits
+            struct ggml_tensor  * b,  // labels
+            struct ggml_tensor  * c); // gradients of cross_entropy_loss result
+
+    // AdamW optimizer step
+    // Paper: https://arxiv.org/pdf/1711.05101v3.pdf
+    // PyTorch: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html
+    GGML_API struct ggml_tensor * ggml_opt_step_adamw(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * grad,
+            struct ggml_tensor  * m,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * adamw_params); // parameters such as the learning rate
+
+    // stochastic gradient descent step (with weight decay)
+    GGML_API struct ggml_tensor * ggml_opt_step_sgd(
+        struct ggml_context * ctx,
+        struct ggml_tensor *  a,
+        struct ggml_tensor *  grad,
+        struct ggml_tensor *  sgd_params); // alpha, weight decay
+
+    // build forward multiple tensors and select one of them for computing
+    // this is useful for creating graphs that have constant topology but compute different things based on the input
+    // ref: https://github.com/ggml-org/llama.cpp/pull/18550
+    //
+    // nodes:
+    //   | - build forward into the graph but do not compute
+    //   c - build forward into the graph and compute
+    //
+    //    |  |  ...  c  ...  |
+    //    |  |  ...  c  ...  |
+    //    |  |  ...  c  ...  |
+    //   [0  1  ... idx ...  n-1]        <-- ggml_build_forward_select(..., n, idx)
+    //               c
+    //               c
+    //
+    // example:
+    //   struct ggml_tensor * curs[3];
+    //
+    //   curs[0]  = compute0(...);
+    //   curs[1]  = compute1(...);
+    //   curs[2]  = compute2(...);
+    //
+    //   int idx = select_branch(some_input);
+    //
+    //   struct ggml_tensor * out = ggml_build_forward_select(cgraph, curs, 3, idx);
+    //
+    GGML_API struct ggml_tensor * ggml_build_forward_select(
+            struct ggml_cgraph  * cgraph,
+            struct ggml_tensor ** tensors,
+            int                   n_tensors,
+            int                   idx);
+
+    GGML_API void ggml_build_forward_expand(
+            struct ggml_cgraph * cgraph,
+            struct ggml_tensor * tensor);
+
+    // add the tensor and its parents to the graph without marking them for compute
+    // the flag is set later, when the tensor is reached from a node that computes
+    GGML_API void ggml_build_forward_order(
+            struct ggml_cgraph * cgraph,
+            struct ggml_tensor * tensor);
+
+    GGML_API void ggml_build_backward_expand(
+        struct ggml_context *  ctx,        // context for gradient computation
+        struct ggml_cgraph  *  cgraph,
+        struct ggml_tensor  ** grad_accs);
+
+    // graph allocation in a context
+    GGML_API struct ggml_cgraph * ggml_new_graph       (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false
+    GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads);
+    GGML_API struct ggml_cgraph * ggml_graph_dup       (struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads);
+    GGML_API void                 ggml_graph_cpy       (struct ggml_cgraph * src, struct ggml_cgraph * dst);
+    GGML_API void                 ggml_graph_reset     (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1
+    GGML_API void                 ggml_graph_clear     (struct ggml_cgraph * cgraph);
+
+    GGML_API int                   ggml_graph_size   (struct ggml_cgraph * cgraph);
+    GGML_API struct ggml_tensor *  ggml_graph_node   (struct ggml_cgraph * cgraph, int i); // if i < 0, returns nodes[n_nodes + i]
+    GGML_API struct ggml_tensor ** ggml_graph_nodes  (struct ggml_cgraph * cgraph);
+    GGML_API int                   ggml_graph_n_nodes(struct ggml_cgraph * cgraph);
+
+    GGML_API void   ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor);
+
+    GGML_API size_t ggml_graph_overhead(void);
+    GGML_API size_t ggml_graph_overhead_custom(size_t size, bool grads);
+
+    GGML_API struct ggml_tensor * ggml_graph_get_tensor  (const struct ggml_cgraph * cgraph, const char * name);
+    GGML_API struct ggml_tensor * ggml_graph_get_grad    (const struct ggml_cgraph * cgraph, const struct ggml_tensor * node);
+    GGML_API struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node);
+
+    // print info and performance information for the graph
+    GGML_API void ggml_graph_print(const struct ggml_cgraph * cgraph);
+
+    // dump the graph into a file using the dot format
+    GGML_API void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename);
+
+    // TODO these functions were sandwiched in the old optimization interface, is there a better place for them?
+    typedef void (*ggml_log_callback)(enum ggml_log_level level, const char * text, void * user_data);
+
+    // Set callback for all future logging events.
+    // If this is not called, or NULL is supplied, everything is output on stderr.
+    GGML_API void ggml_log_get(ggml_log_callback * log_callback, void ** user_data);
+    GGML_API void ggml_log_set(ggml_log_callback   log_callback, void *  user_data);
+
+    GGML_API struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor);
+
+    //
+    // quantization
+    //
+
+    // - ggml_quantize_init can be called multiple times with the same type
+    //   it will only initialize the quantization tables for the first call or after ggml_quantize_free
+    //   automatically called by ggml_quantize_chunk for convenience
+    //
+    // - ggml_quantize_free will free any memory allocated by ggml_quantize_init
+    //   call this at the end of the program to avoid memory leaks
+    //
+    // note: these are thread-safe
+    //
+    GGML_API void ggml_quantize_init(enum ggml_type type);
+    GGML_API void ggml_quantize_free(void);
+
+    // some quantization type cannot be used without an importance matrix
+    GGML_API bool ggml_quantize_requires_imatrix(enum ggml_type type);
+
+    // calls ggml_quantize_init internally (i.e. can allocate memory)
+    GGML_API size_t ggml_quantize_chunk(
+            enum ggml_type   type,
+               const float * src,
+                      void * dst,
+                   int64_t   start,
+                   int64_t   nrows,
+                   int64_t   n_per_row,
+               const float * imatrix);
+
+#ifdef __cplusplus
+    // restrict not standard in C++
+#    if defined(__GNUC__)
+#        define GGML_RESTRICT __restrict__
+#    elif defined(__clang__)
+#        define GGML_RESTRICT __restrict
+#    elif defined(_MSC_VER)
+#        define GGML_RESTRICT __restrict
+#    else
+#        define GGML_RESTRICT
+#    endif
+#else
+#    if defined (_MSC_VER) && (__STDC_VERSION__ < 201112L)
+#        define GGML_RESTRICT __restrict
+#    else
+#        define GGML_RESTRICT restrict
+#    endif
+#endif
+    typedef void (*ggml_to_float_t)  (const void  * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+    typedef void (*ggml_from_float_t)(const float * GGML_RESTRICT x, void  * GGML_RESTRICT y, int64_t k);
+
+    struct ggml_type_traits {
+        const char             * type_name;
+        int64_t                  blck_size;
+        int64_t                  blck_size_interleave; // interleave elements in blocks
+        size_t                   type_size;
+        bool                     is_quantized;
+        ggml_to_float_t          to_float;
+        ggml_from_float_t        from_float_ref;
+    };
+
+    GGML_API const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type);
+
+    // ggml threadpool
+    // TODO: currently, only a few functions are in the base ggml API, while the rest are in the CPU backend
+    // the goal should be to create an API that other backends can use move everything to the ggml base
+
+    // scheduling priorities
+    enum ggml_sched_priority {
+        GGML_SCHED_PRIO_LOW = -1,
+        GGML_SCHED_PRIO_NORMAL,
+        GGML_SCHED_PRIO_MEDIUM,
+        GGML_SCHED_PRIO_HIGH,
+        GGML_SCHED_PRIO_REALTIME
+    };
+
+    // threadpool params
+    // Use ggml_threadpool_params_default() or ggml_threadpool_params_init() to populate the defaults
+    struct ggml_threadpool_params {
+        bool                cpumask[GGML_MAX_N_THREADS]; // mask of cpu cores (all-zeros means use default affinity settings)
+        int                 n_threads;                   // number of threads
+        enum ggml_sched_priority prio;                   // thread priority
+        uint32_t            poll;                        // polling level (0 - no polling, 100 - aggressive polling)
+        bool                strict_cpu;                  // strict cpu placement
+        bool                paused;                      // start in paused state
+    };
+
+    struct ggml_threadpool;     // forward declaration, see ggml.c
+
+    typedef struct ggml_threadpool * ggml_threadpool_t;
+
+    GGML_API struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads);
+    GGML_API void                          ggml_threadpool_params_init   (struct ggml_threadpool_params * p, int n_threads);
+    GGML_API bool                          ggml_threadpool_params_match  (const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1);
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/include/gguf.h b/python/freetoken/kernel/csrc/ggml/include/gguf.h
new file mode 100644
index 0000000000..b3a1e1230a
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/gguf.h
@@ -0,0 +1,211 @@
+// This file contains functionality related to "GGUF" files, the binary file format used by ggml.
+// GGUF files have the following structure:
+//
+// 1. File magic "GGUF" (4 bytes).
+// 2. File version (uint32_t).
+// 3. Number of ggml tensors in file (int64_t).
+// 4. Number of key-value-pairs in file (int64_t).
+// 5. For each KV pair:
+//   1. The key (string).
+//   2. The value type (gguf_type).
+//   3a. If the value type is GGUF_TYPE_ARRAY:
+//     1. The type of the array (gguf_type).
+//     2. The number of elements in the array (uint64_t).
+//     3. The binary representation of each element in the array.
+//   3b. Otherwise:
+//     1. The binary representation of the value.
+// 6. For each ggml tensor:
+//   1. The tensor name (string).
+//   2. The number of dimensions of the tensor (uint32_t).
+//   3. For each dimension:
+//     1. The size of the tensor in the dimension (int64_t).
+//   4. The tensor data type (ggml_type).
+//   5. The tensor data offset in the tensor data binary blob (uint64_t).
+// 7. The tensor data binary blob (optional, aligned).
+//
+// Strings are serialized as the string length (uint64_t) followed by the C string without the null terminator.
+// All enums are stored as int32_t.
+// All bool values are stored as int8_t.
+// If the special key "general.alignment" (uint32_t) is defined it is used for alignment,
+//   otherwise GGUF_DEFAULT_ALIGNMENT is used.
+//
+// Module maintainer: Johannes Gäßler (@JohannesGaessler, johannesg@5d6.de)
+
+#pragma once
+
+#include "ggml.h"
+
+#include <stdbool.h>
+#include <stdint.h>
+
+#define GGUF_MAGIC   "GGUF"
+#define GGUF_VERSION 3
+
+#define GGUF_KEY_GENERAL_ALIGNMENT "general.alignment"
+
+#define GGUF_DEFAULT_ALIGNMENT 32
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    // types that can be stored as GGUF KV data
+    enum gguf_type {
+        GGUF_TYPE_UINT8   = 0,
+        GGUF_TYPE_INT8    = 1,
+        GGUF_TYPE_UINT16  = 2,
+        GGUF_TYPE_INT16   = 3,
+        GGUF_TYPE_UINT32  = 4,
+        GGUF_TYPE_INT32   = 5,
+        GGUF_TYPE_FLOAT32 = 6,
+        GGUF_TYPE_BOOL    = 7,
+        GGUF_TYPE_STRING  = 8,
+        GGUF_TYPE_ARRAY   = 9,
+        GGUF_TYPE_UINT64  = 10,
+        GGUF_TYPE_INT64   = 11,
+        GGUF_TYPE_FLOAT64 = 12,
+        GGUF_TYPE_COUNT,       // marks the end of the enum
+    };
+
+    struct gguf_context;
+
+    struct gguf_init_params {
+        bool no_alloc;
+
+        // if not NULL, create a ggml_context and allocate the tensor data in it
+        struct ggml_context ** ctx;
+    };
+
+    // callback to simulate or wrap a FILE pointer - read up to `len` bytes at `offset` into `output` and return the number of bytes read
+    typedef size_t (*gguf_reader_callback_t)(void * userdata, void * output, uint64_t offset, size_t len);
+
+    GGML_API struct gguf_context * gguf_init_empty(void);
+    GGML_API struct gguf_context * gguf_init_from_file_ptr(FILE * file, struct gguf_init_params params);
+    GGML_API struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_params params);
+    GGML_API struct gguf_context * gguf_init_from_buffer(const void * data, size_t size, struct gguf_init_params params);
+
+    // max_chunk_read is the maximum number of bytes that the GGUF code will read at once from the callback, a value of 0 means no limit
+    GGML_API struct gguf_context * gguf_init_from_callback(gguf_reader_callback_t callback, void * userdata, size_t max_chunk_read, uint64_t max_expected_size, struct gguf_init_params params);
+
+    GGML_API void gguf_free(struct gguf_context * ctx);
+
+    GGML_API const char * gguf_type_name(enum gguf_type type);
+
+    GGML_API uint32_t gguf_get_version    (const struct gguf_context * ctx);
+    GGML_API size_t   gguf_get_alignment  (const struct gguf_context * ctx);
+    GGML_API size_t   gguf_get_data_offset(const struct gguf_context * ctx);  // padded to gguf_get_alignment if and only if the gguf_context contains at least one tensor
+
+    GGML_API int64_t      gguf_get_n_kv(const struct gguf_context * ctx);
+    GGML_API int64_t      gguf_find_key(const struct gguf_context * ctx, const char * key); // returns -1 if key is not found
+    GGML_API const char * gguf_get_key (const struct gguf_context * ctx, int64_t key_id);
+
+    GGML_API enum gguf_type gguf_get_kv_type (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id);
+
+    // will abort if the wrong type is used for the key
+    GGML_API uint8_t      gguf_get_val_u8  (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API int8_t       gguf_get_val_i8  (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API uint16_t     gguf_get_val_u16 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API int16_t      gguf_get_val_i16 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API uint32_t     gguf_get_val_u32 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API int32_t      gguf_get_val_i32 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API float        gguf_get_val_f32 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API uint64_t     gguf_get_val_u64 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API int64_t      gguf_get_val_i64 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API double       gguf_get_val_f64 (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API bool         gguf_get_val_bool(const struct gguf_context * ctx, int64_t key_id);
+    GGML_API const char * gguf_get_val_str (const struct gguf_context * ctx, int64_t key_id);
+    GGML_API const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id);
+    GGML_API size_t       gguf_get_arr_n   (const struct gguf_context * ctx, int64_t key_id);
+
+    // get raw pointer to the first element of the array with the given key_id
+    // for bool arrays, note that they are always stored as int8 on all platforms (usually this makes no difference)
+    GGML_API const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id);
+
+    // get ith C string from array with given key_id
+    GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i);
+
+    GGML_API int64_t         gguf_get_n_tensors    (const struct gguf_context * ctx);
+    GGML_API int64_t         gguf_find_tensor      (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found
+    GGML_API size_t          gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id);
+    GGML_API const char *    gguf_get_tensor_name  (const struct gguf_context * ctx, int64_t tensor_id);
+    GGML_API const int64_t * gguf_get_tensor_ne    (const struct gguf_context * ctx, int64_t tensor_id); // returns ne, an array of GGML_MAX_DIMS elements; ne[dim] is 1 for dim >= n_dims
+    GGML_API enum ggml_type  gguf_get_tensor_type  (const struct gguf_context * ctx, int64_t tensor_id);
+    GGML_API size_t          gguf_get_tensor_size  (const struct gguf_context * ctx, int64_t tensor_id);
+
+    // removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist)
+    GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key);
+
+    // overrides an existing KV pair or adds a new one, the new KV pair is always at the back
+    GGML_API void gguf_set_val_u8  (struct gguf_context * ctx, const char * key, uint8_t      val);
+    GGML_API void gguf_set_val_i8  (struct gguf_context * ctx, const char * key, int8_t       val);
+    GGML_API void gguf_set_val_u16 (struct gguf_context * ctx, const char * key, uint16_t     val);
+    GGML_API void gguf_set_val_i16 (struct gguf_context * ctx, const char * key, int16_t      val);
+    GGML_API void gguf_set_val_u32 (struct gguf_context * ctx, const char * key, uint32_t     val);
+    GGML_API void gguf_set_val_i32 (struct gguf_context * ctx, const char * key, int32_t      val);
+    GGML_API void gguf_set_val_f32 (struct gguf_context * ctx, const char * key, float        val);
+    GGML_API void gguf_set_val_u64 (struct gguf_context * ctx, const char * key, uint64_t     val);
+    GGML_API void gguf_set_val_i64 (struct gguf_context * ctx, const char * key, int64_t      val);
+    GGML_API void gguf_set_val_f64 (struct gguf_context * ctx, const char * key, double       val);
+    GGML_API void gguf_set_val_bool(struct gguf_context * ctx, const char * key, bool         val);
+    GGML_API void gguf_set_val_str (struct gguf_context * ctx, const char * key, const char * val);
+
+    // creates a new array with n elements of the given type and copies the corresponding number of bytes from data
+    GGML_API void gguf_set_arr_data(struct gguf_context * ctx, const char * key, enum gguf_type type, const void * data, size_t n);
+
+    // creates a new array with n strings and copies the corresponding strings from data
+    GGML_API void gguf_set_arr_str (struct gguf_context * ctx, const char * key, const char ** data, size_t n);
+
+    // set or add KV pairs from another context
+    GGML_API void gguf_set_kv(struct gguf_context * ctx, const struct gguf_context * src);
+
+    // add tensor to GGUF context, tensor name must be unique
+    GGML_API void gguf_add_tensor(struct gguf_context * ctx, const struct ggml_tensor * tensor);
+
+    // after changing a tensor's type, the offsets of all tensors with higher indices are immediately recalculated
+    //   in such a way that the tensor data remains as one contiguous block (except for padding)
+    GGML_API void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type);
+
+    // assumes that at least gguf_get_tensor_size bytes can be read from data
+    GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data);
+
+    // writing gguf files can be done in 3 ways:
+    //
+    // - write the entire gguf_context to a binary file in a single pass:
+    //
+    //   gguf_write_to_file(ctx, fname, /*only_meta =*/ false);
+    //
+    // - write only the meta data to a file, then re-open the file and append the tensor data:
+    //
+    //   gguf_write_to_file(ctx, fname, /*only_meta =*/ true);
+    //   FILE * f = fopen(fname, "ab");
+    //   fwrite(f, ...); // write tensor data
+    //   fclose(f);
+    //
+    // - first prepare a file with a placeholder for the meta data, write the tensor data, then write the meta data:
+    //
+    //   FILE * f = fopen(fname, "wb");
+    //   const size_t size_meta = gguf_get_meta_size(ctx);
+    //   fseek(f, size_meta, SEEK_SET);
+    //   fwrite(f, ...); // write tensor data
+    //   void * data = malloc(size_meta);
+    //   gguf_get_meta_data(ctx, data);
+    //   rewind(f);
+    //   fwrite(data, 1, data, f);
+    //   free(data);
+    //   fclose(f);
+    //
+
+    // write the entire context to a binary file
+    GGML_API bool gguf_write_to_file_ptr(const struct gguf_context * ctx, FILE * file, bool only_meta);
+    GGML_API bool gguf_write_to_file(const struct gguf_context * ctx, const char * fname, bool only_meta);
+
+    // get the size in bytes of the meta data (header, kv pairs, tensor info) including padding
+    GGML_API size_t gguf_get_meta_size(const struct gguf_context * ctx);
+
+    // writes the meta data to pointer "data"
+    GGML_API void   gguf_get_meta_data(const struct gguf_context * ctx, void * data);
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-common.h b/python/freetoken/kernel/csrc/ggml/src/ggml-common.h
new file mode 100644
index 0000000000..1dbbe326d0
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/src/ggml-common.h
@@ -0,0 +1,1911 @@
+#ifndef GGML_COMMON_DECL
+
+#if defined(GGML_COMMON_DECL_C)
+#include <stdint.h>
+
+typedef uint16_t ggml_half;
+typedef uint32_t ggml_half2;
+
+#define GGML_COMMON_AGGR_U
+#define GGML_COMMON_AGGR_S
+
+#define GGML_COMMON_DECL
+#elif defined(GGML_COMMON_DECL_CPP)
+#include <cstdint>
+
+typedef uint16_t ggml_half;
+typedef uint32_t ggml_half2;
+
+// std-c++ allow anonymous unions but some compiler warn on it
+#define GGML_COMMON_AGGR_U data
+// std-c++ do not allow it.
+#define GGML_COMMON_AGGR_S data
+
+#define GGML_COMMON_DECL
+#elif defined(GGML_COMMON_DECL_METAL)
+#include <metal_stdlib>
+
+typedef half  ggml_half;
+typedef half2 ggml_half2;
+
+#define GGML_COMMON_AGGR_U
+#define GGML_COMMON_AGGR_S
+
+#define GGML_COMMON_DECL
+#elif defined(GGML_COMMON_DECL_CUDA)
+#if defined(GGML_COMMON_DECL_MUSA)
+#include <musa_fp16.h>
+#else
+#include <cuda_fp16.h>
+#endif
+#include <cstdint>
+
+typedef half  ggml_half;
+typedef half2 ggml_half2;
+
+#define GGML_COMMON_AGGR_U
+#define GGML_COMMON_AGGR_S data
+
+#define GGML_COMMON_DECL
+#elif defined(GGML_COMMON_DECL_HIP)
+#include <hip/hip_fp16.h>
+#include <cstdint>
+
+typedef half  ggml_half;
+typedef half2 ggml_half2;
+
+#define GGML_COMMON_AGGR_U
+#define GGML_COMMON_AGGR_S data
+
+#define GGML_COMMON_DECL
+#elif defined(GGML_COMMON_DECL_SYCL)
+#include <sycl/half_type.hpp>
+#include <cstdint>
+
+typedef sycl::half  ggml_half;
+typedef sycl::half2 ggml_half2;
+
+#define GGML_COMMON_AGGR_U
+#define GGML_COMMON_AGGR_S data
+
+#define GGML_COMMON_DECL
+#endif
+
+#if defined(GGML_COMMON_DECL)
+
+#ifndef __cplusplus
+#ifndef static_assert
+#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201100L)
+#define static_assert(cond, msg) _Static_assert(cond, msg)
+#else
+#define static_assert(cond, msg) struct global_scope_noop_trick
+#endif
+#endif
+#endif // __cplusplus
+
+// QK = number of values after dequantization
+// QK_K = super-block size
+
+#define QK_K 256
+#define K_SCALE_SIZE 12
+
+#if defined(GGML_COMMON_DECL_CUDA) || defined(GGML_COMMON_DECL_HIP) || defined(GGML_COMMON_DECL_SYCL)
+// QR = QK / number of values before dequantization
+// QI = number of 32 bit integers before dequantization
+
+#define QI1_0 (QK1_0 / 32)
+#define QR1_0 1
+
+#define QI2_0 (QK2_0 / 32)
+#define QR2_0 1
+
+
+#define QI4_0 (QK4_0 / (4 * QR4_0))
+#define QR4_0 2
+
+#define QI4_1 (QK4_1 / (4 * QR4_1))
+#define QR4_1 2
+
+#define QI_MXFP4 (QK_MXFP4 / (4 * QR_MXFP4))
+#define QR_MXFP4 2
+
+#define QI_NVFP4 (QK_NVFP4 / (4 * QR_NVFP4))
+#define QR_NVFP4 2
+
+#define QI5_0 (QK5_0 / (4 * QR5_0))
+#define QR5_0 2
+
+#define QI5_1 (QK5_1 / (4 * QR5_1))
+#define QR5_1 2
+
+#define QI8_0 (QK8_0 / (4 * QR8_0))
+#define QR8_0 1
+
+#define QI8_1 (QK8_1 / (4 * QR8_1))
+#define QR8_1 1
+
+#define QI2_K (QK_K / (4*QR2_K))
+#define QR2_K 4
+
+#define QI3_K (QK_K / (4*QR3_K))
+#define QR3_K 4
+
+#define QI4_K (QK_K / (4*QR4_K))
+#define QR4_K 2
+
+#define QI5_K (QK_K / (4*QR5_K))
+#define QR5_K 2
+
+#define QI6_K (QK_K / (4*QR6_K))
+#define QR6_K 2
+
+#define QI2_XXS (QK_K / (4*QR2_XXS))
+#define QR2_XXS 4
+
+#define QI2_XS (QK_K / (4*QR2_XS))
+#define QR2_XS 4
+
+#define QI2_S (QK_K / (4*QR2_S))
+#define QR2_S 4
+
+#define QI3_XXS (QK_K / (4*QR3_XXS))
+#define QR3_XXS 4
+
+#define QI3_XS (QK_K / (4*QR3_XS))
+#define QR3_XS 4
+
+#define QI1_S (QK_K / (4*QR1_S))
+#define QR1_S 8
+
+#define QI1_M (QK_K / (4*QR1_M))
+#define QR1_M 8
+
+#define QI4_NL (QK4_NL / (4*QR4_NL))
+#define QR4_NL 2
+
+#define QI4_XS (QK_K / (4*QR4_XS))
+#define QR4_XS 2
+
+#define QI3_S (QK_K / (4*QR3_S))
+#define QR3_S 4
+
+#endif // GGML_COMMON_DECL_CUDA || GGML_COMMON_DECL_HIP
+
+#ifdef _MSC_VER
+#define GGML_EXTENSION
+#else // _MSC_VER
+#define GGML_EXTENSION __extension__
+#endif // _MSC_VER
+
+#define QK1_0 128
+typedef struct {
+    ggml_half d;           // delta
+    uint8_t qs[QK1_0 / 8]; // bits / quants
+} block_q1_0;
+static_assert(sizeof(block_q1_0) == sizeof(ggml_half) + QK1_0 / 8, "wrong q1_0 block size/padding");
+
+#define QK2_0 64
+typedef struct {
+    ggml_half d;              // delta (scale)
+    uint8_t qs[QK2_0 / 4];   // 2 bits per element
+} block_q2_0;
+static_assert(sizeof(block_q2_0) == sizeof(ggml_half) + QK2_0 / 4, "wrong q2_0 block size/padding");
+
+#define QK4_0 32
+typedef struct {
+    ggml_half d;           // delta
+    uint8_t qs[QK4_0 / 2]; // nibbles / quants
+} block_q4_0;
+static_assert(sizeof(block_q4_0) == sizeof(ggml_half) + QK4_0 / 2, "wrong q4_0 block size/padding");
+
+#define QK4_1 32
+typedef struct {
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d; // delta
+            ggml_half m; // min
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 dm;
+    } GGML_COMMON_AGGR_U;
+    uint8_t qs[QK4_1 / 2]; // nibbles / quants
+} block_q4_1;
+static_assert(sizeof(block_q4_1) == 2 * sizeof(ggml_half) + QK4_1 / 2, "wrong q4_1 block size/padding");
+
+#define QK_MXFP4 32
+typedef struct {
+    uint8_t e; // E8M0
+    uint8_t qs[QK_MXFP4/2];
+} block_mxfp4;
+static_assert(sizeof(block_mxfp4) == sizeof(uint8_t) + QK_MXFP4/2, "wrong mxfp4 block size/padding");
+
+#define QK_NVFP4 64
+#define QK_NVFP4_SUB 16  // sub-block size for per-group scales
+typedef struct {
+    uint8_t d[QK_NVFP4/QK_NVFP4_SUB]; // UE4M3 scales (4 bytes, one per 16-element sub-block)
+    uint8_t qs[QK_NVFP4/2];           // packed 4-bit E2M1 values (32 bytes)
+} block_nvfp4;
+static_assert(sizeof(block_nvfp4) == sizeof(uint8_t)*(QK_NVFP4/QK_NVFP4_SUB) + QK_NVFP4/2, "wrong nvfp4 block size/padding");
+
+#define QK5_0 32
+typedef struct {
+    ggml_half d;           // delta
+    uint8_t qh[4];         // 5-th bit of quants
+    uint8_t qs[QK5_0 / 2]; // nibbles / quants
+} block_q5_0;
+static_assert(sizeof(block_q5_0) == sizeof(ggml_half) + sizeof(uint32_t) + QK5_0 / 2, "wrong q5_0 block size/padding");
+
+#define QK5_1 32
+typedef struct {
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d; // delta
+            ggml_half m; // min
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 dm;
+    } GGML_COMMON_AGGR_U;
+    uint8_t qh[4];         // 5-th bit of quants
+    uint8_t qs[QK5_1 / 2]; // nibbles / quants
+} block_q5_1;
+static_assert(sizeof(block_q5_1) == 2 * sizeof(ggml_half) + sizeof(uint32_t) + QK5_1 / 2, "wrong q5_1 block size/padding");
+
+#define QK8_0 32
+typedef struct {
+    ggml_half d;       // delta
+    int8_t  qs[QK8_0]; // quants
+} block_q8_0;
+static_assert(sizeof(block_q8_0) == sizeof(ggml_half) + QK8_0, "wrong q8_0 block size/padding");
+
+#define QK8_1 32
+typedef struct {
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d; // delta
+            ggml_half s; // d * sum(qs[i])
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 ds;
+    } GGML_COMMON_AGGR_U;
+    int8_t qs[QK8_1]; // quants
+} block_q8_1;
+static_assert(sizeof(block_q8_1) == 2*sizeof(ggml_half) + QK8_1, "wrong q8_1 block size/padding");
+
+//
+// Ternary quantization
+//
+
+// 1.6875 bpw
+typedef struct {
+    uint8_t qs[(QK_K - 4 * QK_K / 64) / 5]; // 5 elements per byte (3^5 = 243 < 256)
+    uint8_t qh[QK_K/64]; // 4 elements per byte
+    ggml_half d;
+} block_tq1_0;
+static_assert(sizeof(block_tq1_0) == sizeof(ggml_half) + QK_K / 64 + (QK_K - 4 * QK_K / 64) / 5, "wrong tq1_0 block size/padding");
+
+// 2.0625 bpw
+typedef struct {
+    uint8_t qs[QK_K/4]; // 2 bits per element
+    ggml_half d;
+} block_tq2_0;
+static_assert(sizeof(block_tq2_0) == sizeof(ggml_half) + QK_K / 4, "wrong tq2_0 block size/padding");
+
+//
+// Super-block quantization structures
+//
+
+// 2-bit quantization
+// weight is represented as x = a * q + b
+// 16 blocks of 16 elements each
+// Effectively 2.625 bits per weight
+typedef struct {
+    uint8_t scales[QK_K/16]; // scales and mins, quantized with 4 bits
+    uint8_t qs[QK_K/4];      // quants
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d;    // super-block scale for quantized scales
+            ggml_half dmin; // super-block scale for quantized mins
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 dm;
+    } GGML_COMMON_AGGR_U;
+} block_q2_K;
+static_assert(sizeof(block_q2_K) == 2*sizeof(ggml_half) + QK_K/16 + QK_K/4, "wrong q2_K block size/padding");
+
+// 3-bit quantization
+// weight is represented as x = a * q
+// 16 blocks of 16 elements each
+// Effectively 3.4375 bits per weight
+typedef struct {
+    uint8_t hmask[QK_K/8]; // quants - high bit
+    uint8_t qs[QK_K/4];    // quants - low 2 bits
+    uint8_t scales[12];    // scales, quantized with 6 bits
+    ggml_half d;           // super-block scale
+} block_q3_K;
+static_assert(sizeof(block_q3_K) == sizeof(ggml_half) + QK_K / 4 + QK_K / 8 + 12, "wrong q3_K block size/padding");
+
+// 4-bit quantization
+// 8 blocks of 32 elements each
+// weight is represented as x = a * q + b
+// Effectively 4.5 bits per weight
+typedef struct {
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d;    // super-block scale for quantized scales
+            ggml_half dmin; // super-block scale for quantized mins
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 dm;
+    } GGML_COMMON_AGGR_U;
+    uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits
+    uint8_t qs[QK_K/2];           // 4--bit quants
+} block_q4_K;
+static_assert(sizeof(block_q4_K) == 2*sizeof(ggml_half) + K_SCALE_SIZE + QK_K/2, "wrong q4_K block size/padding");
+
+// 5-bit quantization
+// 8 blocks of 32 elements each
+// weight is represented as x = a * q + b
+// Effectively 5.5 bits per weight
+typedef struct {
+    GGML_EXTENSION union {
+        struct {
+            ggml_half d;    // super-block scale for quantized scales
+            ggml_half dmin; // super-block scale for quantized mins
+        } GGML_COMMON_AGGR_S;
+        ggml_half2 dm;
+    } GGML_COMMON_AGGR_U;
+    uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits
+    uint8_t qh[QK_K/8];           // quants, high bit
+    uint8_t qs[QK_K/2];           // quants, low 4 bits
+} block_q5_K;
+static_assert(sizeof(block_q5_K) == 2*sizeof(ggml_half) + K_SCALE_SIZE + QK_K/2 + QK_K/8, "wrong q5_K block size/padding");
+
+// 6-bit quantization
+// weight is represented as x = a * q
+// 16 blocks of 16 elements each
+// Effectively 6.5625 bits per weight
+typedef struct {
+    uint8_t ql[QK_K/2];      // quants, lower 4 bits
+    uint8_t qh[QK_K/4];      // quants, upper 2 bits
+    int8_t  scales[QK_K/16]; // scales, quantized with 8 bits
+    ggml_half d;             // super-block scale
+} block_q6_K;
+static_assert(sizeof(block_q6_K) == sizeof(ggml_half) + QK_K / 16 + 3*QK_K/4, "wrong q6_K block size/padding");
+
+// This is only used for intermediate quantization and dot products
+typedef struct {
+    float   d;              // delta
+    int8_t  qs[QK_K];       // quants
+    int16_t bsums[QK_K/16]; // sum of quants in groups of 16
+} block_q8_K;
+static_assert(sizeof(block_q8_K) == sizeof(float) + QK_K + QK_K/16*sizeof(int16_t), "wrong q8_K block size/padding");
+
+// (Almost) "true" 2-bit quantization.
+// Due to the need to use blocks as per ggml design, it ends up using
+// 2.0625 bpw because of the 16-bit scale for each block of 256.
+typedef struct {
+    ggml_half d;
+    uint16_t qs[QK_K/8];
+} block_iq2_xxs;
+static_assert(sizeof(block_iq2_xxs) == sizeof(ggml_half) + QK_K/8*sizeof(uint16_t), "wrong iq2_xxs block size/padding");
+
+// 2.3125 bpw quants
+typedef struct {
+    ggml_half d;
+    uint16_t qs[QK_K/8];
+    uint8_t  scales[QK_K/32];
+} block_iq2_xs;
+static_assert(sizeof(block_iq2_xs) == sizeof(ggml_half) + QK_K/8*sizeof(uint16_t) + QK_K/32, "wrong iq2_xs block size/padding");
+
+// 2.5625 bpw quants
+typedef struct {
+    ggml_half d;
+    uint8_t qs[QK_K/4];
+    uint8_t qh[QK_K/32];
+    uint8_t scales[QK_K/32];
+} block_iq2_s;
+static_assert(sizeof(block_iq2_s) == sizeof(ggml_half) + QK_K/4 + QK_K/16, "wrong iq2_s block size/padding");
+
+// (Almost) "true" 3-bit quantization.
+// Due to the need to use blocks as per ggml design, it ends up using
+// 3.0625 bpw because of the 16-bit scale for each block of 256.
+typedef struct {
+    ggml_half d;
+    uint8_t qs[3*QK_K/8];
+} block_iq3_xxs;
+static_assert(sizeof(block_iq3_xxs) == sizeof(ggml_half) + 3*(QK_K/8), "wrong iq3_xxs block size/padding");
+
+// 3.4375 bpw
+#define IQ3S_N_SCALE QK_K/64
+typedef struct {
+    ggml_half d;
+    uint8_t qs[QK_K/4];
+    uint8_t qh[QK_K/32];
+    uint8_t signs[QK_K/8];
+    uint8_t scales[IQ3S_N_SCALE];
+} block_iq3_s;
+static_assert(sizeof(block_iq3_s) == sizeof(ggml_half) + 13*(QK_K/32) + IQ3S_N_SCALE, "wrong iq3_s block size/padding");
+
+// 1.5625 bpw
+typedef struct {
+    ggml_half d;
+    uint8_t  qs[QK_K/8];
+    uint16_t qh[QK_K/32];
+} block_iq1_s;
+static_assert(sizeof(block_iq1_s) == sizeof(ggml_half) + QK_K/8 + QK_K/16, "wrong iq1_s block size/padding");
+
+// 1.75 bpw
+typedef struct {
+    uint8_t  qs[QK_K/8];      // grid index, low 8 bits
+    uint8_t  qh[QK_K/16];     // grid index, high 3 bits + grid shift bit (for two groups of 8)
+    uint8_t  scales[QK_K/32]; // 3-bit block scales (4-bit if QK_K == 64)
+} block_iq1_m;
+static_assert(sizeof(block_iq1_m) == QK_K/8 + QK_K/16 + QK_K/32, "wrong iq1_m block size/padding");
+
+// Used by IQ1_M quants
+typedef union {
+    ggml_half f16;
+    uint16_t  u16;
+} iq1m_scale_t;
+
+// Non-linear quants
+#define QK4_NL 32
+typedef struct {
+    ggml_half d;
+    uint8_t qs[QK4_NL/2];
+} block_iq4_nl;
+static_assert(sizeof(block_iq4_nl) == sizeof(ggml_half) + QK4_NL/2, "wrong iq4_nl block size/padding");
+
+typedef struct {
+    ggml_half d;
+    uint16_t scales_h;
+    uint8_t  scales_l[QK_K/64];
+    uint8_t  qs[QK_K/2];
+} block_iq4_xs;
+static_assert(sizeof(block_iq4_xs) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/64 + QK_K/2, "wrong iq4_xs block size/padding");
+
+#endif // GGML_COMMON_DECL
+#endif // GGML_COMMON_DECL
+
+////////////////////////////////////////////////////////////////////////////////
+
+#ifndef GGML_COMMON_IMPL
+
+#if defined(GGML_COMMON_IMPL_C)
+#include <stdint.h>
+
+#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = {
+#define GGML_TABLE_END() };
+
+#define GGML_COMMON_IMPL
+#elif defined(GGML_COMMON_IMPL_CPP)
+#include <cstdint>
+
+#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = {
+#define GGML_TABLE_END() };
+
+#define GGML_COMMON_IMPL
+#elif defined(GGML_COMMON_IMPL_METAL)
+#include <metal_stdlib>
+
+#define GGML_TABLE_BEGIN(type, name, size) static const constant type name[size] = {
+#define GGML_TABLE_END() };
+
+#define GGML_COMMON_IMPL
+#elif defined(GGML_COMMON_IMPL_CUDA) || defined(GGML_COMMON_IMPL_HIP) || defined(GGML_COMMON_IMPL_MUSA)
+#include <cstdint>
+
+#define GGML_TABLE_BEGIN(type, name, size) static const __device__ type name[size] = {
+#define GGML_TABLE_END() };
+
+#define GGML_COMMON_IMPL
+#elif defined(GGML_COMMON_IMPL_SYCL)
+
+#include <cstdint>
+
+#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = {
+#define GGML_TABLE_END() };
+
+#define GGML_COMMON_IMPL
+#endif
+
+#if defined(GGML_COMMON_IMPL)
+
+GGML_TABLE_BEGIN(uint8_t, kmask_iq2xs, 8)
+    1, 2, 4, 8, 16, 32, 64, 128
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint8_t, ksigns_iq2xs, 128)
+      0, 129, 130,   3, 132,   5,   6, 135, 136,   9,  10, 139,  12, 141, 142,  15,
+    144,  17,  18, 147,  20, 149, 150,  23,  24, 153, 154,  27, 156,  29,  30, 159,
+    160,  33,  34, 163,  36, 165, 166,  39,  40, 169, 170,  43, 172,  45,  46, 175,
+     48, 177, 178,  51, 180,  53,  54, 183, 184,  57,  58, 187,  60, 189, 190,  63,
+    192,  65,  66, 195,  68, 197, 198,  71,  72, 201, 202,  75, 204,  77,  78, 207,
+     80, 209, 210,  83, 212,  85,  86, 215, 216,  89,  90, 219,  92, 221, 222,  95,
+     96, 225, 226,  99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
+    240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint64_t, ksigns64, 128)
+    0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff,
+    0xff00000000ff0000, 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff,
+    0xff000000ff000000, 0x00000000ff0000ff, 0x00000000ff00ff00, 0xff000000ff00ffff,
+    0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, 0x00000000ffffffff,
+    0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff,
+    0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff,
+    0x000000ffff000000, 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff,
+    0xff0000ffffff0000, 0x000000ffffff00ff, 0x000000ffffffff00, 0xff0000ffffffffff,
+    0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, 0xff00ff000000ffff,
+    0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff,
+    0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff,
+    0xff00ff00ffff0000, 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff,
+    0x0000ffff00000000, 0xff00ffff000000ff, 0xff00ffff0000ff00, 0x0000ffff0000ffff,
+    0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, 0xff00ffff00ffffff,
+    0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff,
+    0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff,
+    0xffff000000000000, 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff,
+    0x00ff000000ff0000, 0xffff000000ff00ff, 0xffff000000ffff00, 0x00ff000000ffffff,
+    0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, 0x00ff0000ff00ffff,
+    0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff,
+    0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff,
+    0xffff00ff00ff0000, 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff,
+    0xffff00ffff000000, 0x00ff00ffff0000ff, 0x00ff00ffff00ff00, 0xffff00ffff00ffff,
+    0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, 0x00ff00ffffffffff,
+    0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff,
+    0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff,
+    0xffffff00ff000000, 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff,
+    0x00ffff00ffff0000, 0xffffff00ffff00ff, 0xffffff00ffffff00, 0x00ffff00ffffffff,
+    0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, 0xffffffff0000ffff,
+    0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff,
+    0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff,
+    0xffffffffffff0000, 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff,
+GGML_TABLE_END()
+
+
+GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256)
+    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
+    0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
+    0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
+    0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
+    0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
+    0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
+    0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
+    0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
+    0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
+    0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
+    0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
+    0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
+    0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
+    0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
+    0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
+    0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
+    0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
+    0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
+    0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
+    0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
+    0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
+    0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
+    0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
+    0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
+    0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
+    0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
+    0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
+    0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
+    0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
+    0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
+    0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
+    0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
+    0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
+    0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
+    0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
+    0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
+    0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
+    0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
+    0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
+    0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
+    0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
+    0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
+    0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
+    0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
+    0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
+    0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
+    0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
+    0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
+    0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
+    0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
+    0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
+    0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
+    0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
+    0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
+    0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
+    0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
+    0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
+    0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
+    0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
+    0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
+    0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
+    0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
+    0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
+    0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512)
+    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
+    0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b,
+    0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919,
+    0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b,
+    0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919,
+    0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808,
+    0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819,
+    0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819,
+    0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808,
+    0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b,
+    0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b,
+    0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908,
+    0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908,
+    0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919,
+    0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808,
+    0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919,
+    0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908,
+    0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b,
+    0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908,
+    0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08,
+    0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808,
+    0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808,
+    0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819,
+    0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908,
+    0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819,
+    0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808,
+    0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b,
+    0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819,
+    0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819,
+    0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808,
+    0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908,
+    0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19,
+    0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b,
+    0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b,
+    0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919,
+    0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808,
+    0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819,
+    0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819,
+    0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b,
+    0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908,
+    0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808,
+    0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819,
+    0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808,
+    0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919,
+    0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808,
+    0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808,
+    0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908,
+    0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908,
+    0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808,
+    0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b,
+    0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819,
+    0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919,
+    0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908,
+    0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808,
+    0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908,
+    0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919,
+    0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08,
+    0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19,
+    0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b,
+    0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b,
+    0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808,
+    0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08,
+    0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b,
+    0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908,
+    0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b,
+    0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908,
+    0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08,
+    0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808,
+    0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808,
+    0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08,
+    0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819,
+    0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919,
+    0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808,
+    0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808,
+    0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819,
+    0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819,
+    0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908,
+    0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908,
+    0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b,
+    0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908,
+    0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908,
+    0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908,
+    0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808,
+    0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819,
+    0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819,
+    0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819,
+    0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808,
+    0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b,
+    0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819,
+    0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819,
+    0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08,
+    0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808,
+    0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19,
+    0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919,
+    0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808,
+    0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19,
+    0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b,
+    0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808,
+    0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b,
+    0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b,
+    0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08,
+    0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b,
+    0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808,
+    0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819,
+    0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808,
+    0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808,
+    0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08,
+    0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b,
+    0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19,
+    0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08,
+    0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919,
+    0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08,
+    0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08,
+    0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908,
+    0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908,
+    0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b,
+    0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908,
+    0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808,
+    0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b,
+    0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808,
+    0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808,
+    0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19,
+    0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08,
+    0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808,
+    0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b,
+    0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808,
+    0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b,
+    0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024)
+    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
+    0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b,
+    0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919,
+    0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b,
+    0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919,
+    0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b,
+    0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919,
+    0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808,
+    0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908,
+    0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b,
+    0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908,
+    0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08,
+    0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19,
+    0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819,
+    0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919,
+    0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b,
+    0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908,
+    0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908,
+    0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919,
+    0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908,
+    0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b,
+    0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919,
+    0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b,
+    0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919,
+    0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908,
+    0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b,
+    0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b,
+    0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08,
+    0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b,
+    0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819,
+    0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808,
+    0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908,
+    0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b,
+    0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908,
+    0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08,
+    0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819,
+    0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808,
+    0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08,
+    0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819,
+    0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b,
+    0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908,
+    0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919,
+    0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b,
+    0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919,
+    0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808,
+    0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819,
+    0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919,
+    0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919,
+    0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808,
+    0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819,
+    0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b,
+    0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908,
+    0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b,
+    0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908,
+    0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919,
+    0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b,
+    0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919,
+    0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b,
+    0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819,
+    0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919,
+    0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908,
+    0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b,
+    0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908,
+    0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b,
+    0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908,
+    0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08,
+    0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908,
+    0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819,
+    0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819,
+    0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808,
+    0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08,
+    0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19,
+    0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819,
+    0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808,
+    0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819,
+    0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919,
+    0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808,
+    0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19,
+    0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08,
+    0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b,
+    0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908,
+    0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808,
+    0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819,
+    0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908,
+    0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819,
+    0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808,
+    0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808,
+    0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819,
+    0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908,
+    0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08,
+    0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819,
+    0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b,
+    0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b,
+    0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08,
+    0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19,
+    0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819,
+    0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919,
+    0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908,
+    0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808,
+    0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808,
+    0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908,
+    0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808,
+    0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08,
+    0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08,
+    0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908,
+    0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919,
+    0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808,
+    0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819,
+    0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908,
+    0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08,
+    0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819,
+    0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808,
+    0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808,
+    0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819,
+    0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808,
+    0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908,
+    0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b,
+    0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819,
+    0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808,
+    0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b,
+    0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808,
+    0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b,
+    0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19,
+    0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819,
+    0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08,
+    0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b,
+    0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908,
+    0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b,
+    0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b,
+    0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919,
+    0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808,
+    0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819,
+    0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908,
+    0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08,
+    0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08,
+    0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819,
+    0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919,
+    0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908,
+    0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b,
+    0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908,
+    0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b,
+    0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908,
+    0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08,
+    0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819,
+    0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808,
+    0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819,
+    0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919,
+    0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808,
+    0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808,
+    0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08,
+    0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819,
+    0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919,
+    0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808,
+    0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819,
+    0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919,
+    0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808,
+    0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b,
+    0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908,
+    0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808,
+    0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908,
+    0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b,
+    0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908,
+    0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b,
+    0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908,
+    0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b,
+    0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908,
+    0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08,
+    0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908,
+    0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b,
+    0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908,
+    0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08,
+    0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819,
+    0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919,
+    0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808,
+    0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19,
+    0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b,
+    0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919,
+    0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808,
+    0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819,
+    0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908,
+    0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919,
+    0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808,
+    0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808,
+    0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b,
+    0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919,
+    0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808,
+    0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b,
+    0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808,
+    0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919,
+    0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b,
+    0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08,
+    0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919,
+    0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808,
+    0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b,
+    0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908,
+    0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808,
+    0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808,
+    0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808,
+    0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908,
+    0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808,
+    0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808,
+    0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b,
+    0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908,
+    0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808,
+    0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808,
+    0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819,
+    0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919,
+    0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b,
+    0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808,
+    0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819,
+    0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b,
+    0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908,
+    0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08,
+    0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908,
+    0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919,
+    0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819,
+    0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908,
+    0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b,
+    0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808,
+    0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819,
+    0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908,
+    0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919,
+    0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808,
+    0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808,
+    0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808,
+    0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919,
+    0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908,
+    0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908,
+    0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08,
+    0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819,
+    0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b,
+    0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808,
+    0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819,
+    0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908,
+    0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819,
+    0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808,
+    0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808,
+    0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b,
+    0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908,
+    0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808,
+    0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908,
+    0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819,
+    0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819,
+    0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808,
+    0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b,
+    0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b,
+    0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819,
+    0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b,
+    0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b,
+    0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b,
+    0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819,
+    0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19,
+    0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819,
+    0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908,
+    0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808,
+    0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256)
+    0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414,
+    0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14,
+    0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404,
+    0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e,
+    0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c,
+    0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c,
+    0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34,
+    0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c,
+    0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c,
+    0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04,
+    0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c,
+    0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414,
+    0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434,
+    0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c,
+    0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e,
+    0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24,
+    0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24,
+    0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c,
+    0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c,
+    0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14,
+    0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414,
+    0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e,
+    0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404,
+    0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c,
+    0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c,
+    0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14,
+    0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c,
+    0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c,
+    0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14,
+    0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14,
+    0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c,
+    0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512)
+    0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305,
+    0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905,
+    0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09,
+    0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b,
+    0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b,
+    0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d,
+    0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03,
+    0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505,
+    0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03,
+    0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901,
+    0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d,
+    0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303,
+    0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501,
+    0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105,
+    0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505,
+    0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101,
+    0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707,
+    0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b,
+    0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01,
+    0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f,
+    0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305,
+    0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103,
+    0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509,
+    0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503,
+    0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b,
+    0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f,
+    0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f,
+    0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f,
+    0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109,
+    0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f,
+    0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509,
+    0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501,
+    0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303,
+    0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f,
+    0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907,
+    0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703,
+    0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03,
+    0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01,
+    0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01,
+    0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903,
+    0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505,
+    0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b,
+    0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107,
+    0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509,
+    0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303,
+    0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103,
+    0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05,
+    0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b,
+    0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f,
+    0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701,
+    0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909,
+    0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305,
+    0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d,
+    0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b,
+    0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d,
+    0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307,
+    0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09,
+    0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309,
+    0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709,
+    0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f,
+    0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303,
+    0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503,
+    0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b,
+    0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101,
+GGML_TABLE_END()
+
+// TODO: fix name to kvalues_iq4_nl
+GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16)
+    -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113,
+GGML_TABLE_END()
+
+// e2m1 values (doubled), shared by MXFP4 and NVFP4
+// ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16)
+    0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12,
+GGML_TABLE_END()
+#define kvalues_mxfp4 kvalues_fp4
+
+#define NGRID_IQ1S 2048
+#define IQ1S_DELTA 0.125f
+#define IQ1M_DELTA 0.125f
+#if defined(GGML_COMMON_IMPL_C) || defined(GGML_COMMON_IMPL_CPP)
+GGML_TABLE_BEGIN(uint64_t, iq1s_grid, NGRID_IQ1S)
+    0xffffffffffffffff, 0xffffffffffffff01, 0xffffffffffff0000, 0xffffffffffff01ff,
+    0xffffffffffff0101, 0xffffffffff00ff00, 0xffffffffff000000, 0xffffffffff01ffff,
+    0xffffffffff01ff01, 0xffffffffff0101ff, 0xffffffffff010101, 0xffffffff00ff0000,
+    0xffffffff0000ff00, 0xffffffff000000ff, 0xffffffff00000001, 0xffffffff00010000,
+    0xffffffff01ffffff, 0xffffffff01ffff01, 0xffffffff01ff01ff, 0xffffffff01ff0101,
+    0xffffffff01000000, 0xffffffff0101ffff, 0xffffffff0101ff01, 0xffffffff010101ff,
+    0xffffffff01010101, 0xffffff00ffff00ff, 0xffffff00ffff0000, 0xffffff00ff00ff00,
+    0xffffff00ff0000ff, 0xffffff00ff000001, 0xffffff00ff000100, 0xffffff00ff000101,
+    0xffffff00ff010000, 0xffffff0000ffff00, 0xffffff0000ff0001, 0xffffff0000ff0100,
+    0xffffff000000ff01, 0xffffff0000000000, 0xffffff0000000101, 0xffffff000001ff00,
+    0xffffff00000100ff, 0xffffff0000010001, 0xffffff00000101ff, 0xffffff0001ff0000,
+    0xffffff000100ff00, 0xffffff00010000ff, 0xffffff0001000001, 0xffffff0001010000,
+    0xffffff01ffffffff, 0xffffff01ffffff01, 0xffffff01ffff01ff, 0xffffff01ffff0101,
+    0xffffff01ff000000, 0xffffff01ff01ffff, 0xffffff01ff01ff01, 0xffffff01ff0101ff,
+    0xffffff01ff010101, 0xffffff0100ff0000, 0xffffff010000ff00, 0xffffff0100000100,
+    0xffffff01000100ff, 0xffffff0100010100, 0xffffff0101ffffff, 0xffffff0101ffff01,
+    0xffffff0101ff01ff, 0xffffff0101ff0101, 0xffffff010100ff00, 0xffffff0101000000,
+    0xffffff0101000100, 0xffffff010101ffff, 0xffffff010101ff01, 0xffffff01010101ff,
+    0xffffff0101010101, 0xffff00ffff00ff00, 0xffff00ffff0000ff, 0xffff00ffff000001,
+    0xffff00ffff010000, 0xffff00ff00ffff00, 0xffff00ff00ff0100, 0xffff00ff00000000,
+    0xffff00ff00000101, 0xffff00ff000100ff, 0xffff00ff00010000, 0xffff00ff0100ff00,
+    0xffff00ff01000100, 0xffff00ff01010000, 0xffff0000ffffff00, 0xffff0000ffff00ff,
+    0xffff0000ffff0000, 0xffff0000ffff0001, 0xffff0000ff000000, 0xffff0000ff0001ff,
+    0xffff0000ff000101, 0xffff0000ff010100, 0xffff000000ffffff, 0xffff000000ff0000,
+    0xffff000000ff0101, 0xffff00000000ffff, 0xffff00000000ff00, 0xffff0000000000ff,
+    0xffff000000000000, 0xffff000000000001, 0xffff000000000100, 0xffff00000001ffff,
+    0xffff00000001ff01, 0xffff000000010000, 0xffff0000000101ff, 0xffff000000010101,
+    0xffff000001ffff00, 0xffff00000100ff00, 0xffff000001000000, 0xffff0000010001ff,
+    0xffff000001000101, 0xffff00000101ff00, 0xffff0000010100ff, 0xffff000001010000,
+    0xffff000001010001, 0xffff000001010100, 0xffff0001ff0000ff, 0xffff0001ff000100,
+    0xffff000100ffff00, 0xffff000100ff00ff, 0xffff00010000ffff, 0xffff00010000ff01,
+    0xffff000100000000, 0xffff0001000001ff, 0xffff00010001ffff, 0xffff00010001ff00,
+    0xffff000100010001, 0xffff000100010100, 0xffff000101ff0000, 0xffff00010100ff00,
+    0xffff0001010000ff, 0xffff000101000100, 0xffff01ffffffffff, 0xffff01ffffffff01,
+    0xffff01ffffff01ff, 0xffff01ffffff0101, 0xffff01ffff000000, 0xffff01ffff01ffff,
+    0xffff01ffff01ff01, 0xffff01ffff0101ff, 0xffff01ffff010101, 0xffff01ff00ff0000,
+    0xffff01ff0000ff00, 0xffff01ff00000001, 0xffff01ff00010000, 0xffff01ff01ffffff,
+    0xffff01ff01ffff01, 0xffff01ff01ff01ff, 0xffff01ff01ff0101, 0xffff01ff01000000,
+    0xffff01ff0101ffff, 0xffff01ff0101ff01, 0xffff01ff010101ff, 0xffff01ff01010101,
+    0xffff0100ffff0000, 0xffff0100ff00ff00, 0xffff0100ff0000ff, 0xffff0100ff000100,
+    0xffff0100ff0100ff, 0xffff0100ff010000, 0xffff010000ffff00, 0xffff01000000ffff,
+    0xffff01000000ff00, 0xffff010000000000, 0xffff01000001ff00, 0xffff0100000100ff,
+    0xffff010000010100, 0xffff01000100ff00, 0xffff0100010000ff, 0xffff010001000001,
+    0xffff010001000100, 0xffff010001010000, 0xffff0101ffffffff, 0xffff0101ffffff01,
+    0xffff0101ffff01ff, 0xffff0101ffff0101, 0xffff0101ff000000, 0xffff0101ff01ffff,
+    0xffff0101ff01ff01, 0xffff0101ff0101ff, 0xffff0101ff010101, 0xffff010100ff0000,
+    0xffff01010000ff00, 0xffff010100000100, 0xffff01010001ff00, 0xffff010100010000,
+    0xffff010101ffffff, 0xffff010101ffff01, 0xffff010101ff0000, 0xffff010101ff01ff,
+    0xffff010101ff0101, 0xffff010101000000, 0xffff01010101ffff, 0xffff01010101ff01,
+    0xffff0101010101ff, 0xffff010101010101, 0xff00ffffff00ffff, 0xff00ffffff00ff00,
+    0xff00ffffff0000ff, 0xff00ffffff000100, 0xff00ffffff0100ff, 0xff00ffffff010000,
+    0xff00ffff00ffff00, 0xff00ffff00ff00ff, 0xff00ffff0000ffff, 0xff00ffff00000000,
+    0xff00ffff000001ff, 0xff00ffff0001ff00, 0xff00ffff000100ff, 0xff00ffff00010000,
+    0xff00ffff00010100, 0xff00ffff0100ff00, 0xff00ffff010000ff, 0xff00ffff01000001,
+    0xff00ffff0101ff00, 0xff00ffff01010000, 0xff00ff00ffffff00, 0xff00ff00ffff00ff,
+    0xff00ff00ffff0001, 0xff00ff00ffff0100, 0xff00ff00ff00ffff, 0xff00ff00ff00ff01,
+    0xff00ff00ff000000, 0xff00ff00ff0001ff, 0xff00ff00ff01ff00, 0xff00ff00ff0100ff,
+    0xff00ff00ff010100, 0xff00ff0000ff0000, 0xff00ff0000ff0101, 0xff00ff000000ffff,
+    0xff00ff000000ff00, 0xff00ff000000ff01, 0xff00ff00000000ff, 0xff00ff0000000000,
+    0xff00ff0000000001, 0xff00ff0000000100, 0xff00ff000001ffff, 0xff00ff0000010000,
+    0xff00ff0001ff00ff, 0xff00ff000100ff01, 0xff00ff0001000000, 0xff00ff000101ff00,
+    0xff00ff00010100ff, 0xff00ff01ff00ff00, 0xff00ff01ff0000ff, 0xff00ff01ff000001,
+    0xff00ff01ff010000, 0xff00ff0100ffffff, 0xff00ff0100ff0001, 0xff00ff0100ff0100,
+    0xff00ff010000ff01, 0xff00ff0100000000, 0xff00ff01000001ff, 0xff00ff0100000101,
+    0xff00ff01000100ff, 0xff00ff0100010001, 0xff00ff0101ff0000, 0xff00ff010100ff00,
+    0xff00ff01010000ff, 0xff00ff0101000001, 0xff00ff0101010000, 0xff0000ffffffff00,
+    0xff0000ffffff0001, 0xff0000ffffff0100, 0xff0000ffff0000ff, 0xff0000ffff000000,
+    0xff0000ffff0001ff, 0xff0000ffff000100, 0xff0000ffff01ff00, 0xff0000ffff010001,
+    0xff0000ff00ffff00, 0xff0000ff00ff0000, 0xff0000ff00ff0001, 0xff0000ff00ff01ff,
+    0xff0000ff00ff0101, 0xff0000ff0000ff00, 0xff0000ff000000ff, 0xff0000ff00000000,
+    0xff0000ff00000001, 0xff0000ff00000100, 0xff0000ff0001ff01, 0xff0000ff00010000,
+    0xff0000ff000101ff, 0xff0000ff01ff00ff, 0xff0000ff01ff0100, 0xff0000ff0100ffff,
+    0xff0000ff010000ff, 0xff0000ff01000000, 0xff0000ff010001ff, 0xff0000ff01000100,
+    0xff0000ff01000101, 0xff0000ff0101ff00, 0xff0000ff010100ff, 0xff0000ff01010000,
+    0xff0000ff01010100, 0xff000000ffffff01, 0xff000000ffff0000, 0xff000000ffff0101,
+    0xff000000ff00ff00, 0xff000000ff0000ff, 0xff000000ff000000, 0xff000000ff000001,
+    0xff000000ff000100, 0xff000000ff01ffff, 0xff000000ff01ff01, 0xff000000ff010000,
+    0xff000000ff0101ff, 0xff000000ff010101, 0xff00000000ffff00, 0xff00000000ff00ff,
+    0xff00000000ff0000, 0xff00000000ff0001, 0xff0000000000ff00, 0xff0000000000ff01,
+    0xff000000000000ff, 0xff00000000000000, 0xff00000000000001, 0xff00000000000100,
+    0xff00000000000101, 0xff0000000001ff00, 0xff000000000100ff, 0xff00000000010000,
+    0xff00000000010001, 0xff00000000010100, 0xff00000001ffffff, 0xff00000001ffff01,
+    0xff00000001ff00ff, 0xff00000001ff0000, 0xff00000001ff01ff, 0xff00000001ff0101,
+    0xff0000000100ffff, 0xff0000000100ff00, 0xff000000010000ff, 0xff00000001000000,
+    0xff00000001000001, 0xff00000001000100, 0xff00000001000101, 0xff0000000101ffff,
+    0xff0000000101ff01, 0xff00000001010000, 0xff000001ffffff00, 0xff000001ffff00ff,
+    0xff000001ffff0000, 0xff000001ffff0001, 0xff000001ff000000, 0xff000001ff000001,
+    0xff000001ff0001ff, 0xff000001ff000101, 0xff000001ff01ff00, 0xff000001ff010001,
+    0xff00000100ffffff, 0xff00000100ffff01, 0xff00000100ff00ff, 0xff00000100ff0000,
+    0xff00000100ff01ff, 0xff00000100ff0101, 0xff0000010000ff00, 0xff00000100000000,
+    0xff00000100000001, 0xff000001000001ff, 0xff00000100000100, 0xff0000010001ff00,
+    0xff000001000100ff, 0xff00000100010000, 0xff000001000101ff, 0xff00000100010100,
+    0xff00000100010101, 0xff00000101ff0001, 0xff00000101ff0101, 0xff0000010100ff01,
+    0xff00000101000000, 0xff000001010100ff, 0xff00000101010100, 0xff0001ffff00ff00,
+    0xff0001ffff000001, 0xff0001ffff010000, 0xff0001ff00ffff00, 0xff0001ff00ff00ff,
+    0xff0001ff00ff0001, 0xff0001ff00ff0100, 0xff0001ff0000ffff, 0xff0001ff00000000,
+    0xff0001ff000001ff, 0xff0001ff00000101, 0xff0001ff0001ffff, 0xff0001ff0001ff00,
+    0xff0001ff000100ff, 0xff0001ff00010001, 0xff0001ff00010100, 0xff0001ff01ff0000,
+    0xff0001ff0100ff00, 0xff0001ff010000ff, 0xff0001ff01010000, 0xff000100ff00ffff,
+    0xff000100ff00ff01, 0xff000100ff000000, 0xff000100ff000101, 0xff000100ff01ff00,
+    0xff000100ff010000, 0xff00010000ffff01, 0xff00010000ff00ff, 0xff00010000ff0000,
+    0xff00010000ff01ff, 0xff0001000000ff00, 0xff000100000000ff, 0xff00010000000000,
+    0xff00010000000001, 0xff00010000000100, 0xff00010000000101, 0xff0001000001ffff,
+    0xff00010000010000, 0xff00010000010101, 0xff00010001ff0100, 0xff0001000100ff00,
+    0xff0001000100ff01, 0xff00010001000000, 0xff000100010001ff, 0xff0001000101ff00,
+    0xff00010001010001, 0xff00010001010100, 0xff000101ffff0100, 0xff000101ff000001,
+    0xff000101ff0100ff, 0xff000101ff010001, 0xff00010100ff00ff, 0xff00010100ff0001,
+    0xff00010100ff0100, 0xff0001010000ffff, 0xff0001010000ff01, 0xff00010100000000,
+    0xff000101000001ff, 0xff0001010001ff00, 0xff00010100010001, 0xff00010100010100,
+    0xff00010101ff0000, 0xff0001010100ff00, 0xff00010101000001, 0xff00010101000101,
+    0xff01ffffffffffff, 0xff01ffffffffff01, 0xff01ffffffff01ff, 0xff01ffffffff0101,
+    0xff01ffffff000000, 0xff01ffffff01ffff, 0xff01ffffff01ff01, 0xff01ffffff010000,
+    0xff01ffffff0101ff, 0xff01ffffff010101, 0xff01ffff00ff0000, 0xff01ffff0000ff00,
+    0xff01ffff00000100, 0xff01ffff0001ff00, 0xff01ffff00010000, 0xff01ffff01ffffff,
+    0xff01ffff01ffff01, 0xff01ffff01ff01ff, 0xff01ffff01ff0101, 0xff01ffff01000000,
+    0xff01ffff0101ffff, 0xff01ffff0101ff01, 0xff01ffff01010000, 0xff01ffff010101ff,
+    0xff01ffff01010101, 0xff01ff00ffff0000, 0xff01ff00ff00ff00, 0xff01ff00ff0000ff,
+    0xff01ff00ff000100, 0xff01ff00ff010000, 0xff01ff0000ffff01, 0xff01ff0000ff00ff,
+    0xff01ff0000ff0100, 0xff01ff0000000000, 0xff01ff00000001ff, 0xff01ff0000000101,
+    0xff01ff000001ff00, 0xff01ff00000100ff, 0xff01ff0000010000, 0xff01ff0000010001,
+    0xff01ff0001ff0000, 0xff01ff000100ffff, 0xff01ff0001000001, 0xff01ff0001000100,
+    0xff01ff0001010000, 0xff01ff01ffffff00, 0xff01ff01ffff01ff, 0xff01ff01ffff0101,
+    0xff01ff01ff00ff00, 0xff01ff01ff000000, 0xff01ff01ff01ffff, 0xff01ff01ff01ff01,
+    0xff01ff01ff0101ff, 0xff01ff01ff010101, 0xff01ff0100ff0000, 0xff01ff010000ff00,
+    0xff01ff0100000001, 0xff01ff0100000100, 0xff01ff0100010000, 0xff01ff0101ffff00,
+    0xff01ff0101ff01ff, 0xff01ff0101ff0101, 0xff01ff010100ff00, 0xff01ff0101000000,
+    0xff01ff010101ffff, 0xff01ff010101ff01, 0xff01ff01010101ff, 0xff01ff0101010101,
+    0xff0100ffffff0000, 0xff0100ffff0000ff, 0xff0100ffff000001, 0xff0100ffff000100,
+    0xff0100ffff010000, 0xff0100ff00ff00ff, 0xff0100ff00ff0000, 0xff0100ff00ff0001,
+    0xff0100ff00ff0100, 0xff0100ff0000ff01, 0xff0100ff00000000, 0xff0100ff000001ff,
+    0xff0100ff00000101, 0xff0100ff00010001, 0xff0100ff01ff0000, 0xff0100ff0100ff00,
+    0xff0100ff010000ff, 0xff0100ff01000100, 0xff0100ff0101ff00, 0xff0100ff01010000,
+    0xff010000ffff0100, 0xff010000ff000000, 0xff010000ff01ff00, 0xff010000ff010100,
+    0xff01000000ffffff, 0xff01000000ff0000, 0xff01000000ff01ff, 0xff0100000000ff00,
+    0xff010000000000ff, 0xff01000000000000, 0xff01000000000100, 0xff0100000001ff01,
+    0xff01000000010000, 0xff010000000101ff, 0xff01000001ff0100, 0xff0100000100ffff,
+    0xff010000010000ff, 0xff01000001000000, 0xff010000010001ff, 0xff01000001000101,
+    0xff0100000101ff00, 0xff010000010100ff, 0xff01000001010001, 0xff01000001010100,
+    0xff010001ffff0000, 0xff010001ff00ffff, 0xff010001ff00ff01, 0xff010001ff000100,
+    0xff010001ff010000, 0xff01000100ffff00, 0xff01000100ff0100, 0xff01000100000000,
+    0xff0100010001ffff, 0xff0100010001ff00, 0xff01000100010100, 0xff01000101ff00ff,
+    0xff01000101ff0001, 0xff0100010100ffff, 0xff01000101000101, 0xff0101ffffffffff,
+    0xff0101ffffffff01, 0xff0101ffffff01ff, 0xff0101ffffff0101, 0xff0101ffff000000,
+    0xff0101ffff01ffff, 0xff0101ffff01ff01, 0xff0101ffff0101ff, 0xff0101ffff010101,
+    0xff0101ff00ff0000, 0xff0101ff0000ff00, 0xff0101ff000000ff, 0xff0101ff00010000,
+    0xff0101ff01ffffff, 0xff0101ff01ffff01, 0xff0101ff01ff01ff, 0xff0101ff01ff0101,
+    0xff0101ff0101ffff, 0xff0101ff0101ff01, 0xff0101ff010101ff, 0xff0101ff01010101,
+    0xff010100ffff0100, 0xff010100ff00ff00, 0xff010100ff0000ff, 0xff010100ff000100,
+    0xff010100ff010000, 0xff01010000ff0001, 0xff01010000ff0100, 0xff0101000000ff01,
+    0xff01010000000000, 0xff0101000001ff00, 0xff010100000100ff, 0xff01010000010001,
+    0xff01010000010100, 0xff01010001ff0000, 0xff0101000100ffff, 0xff01010001000001,
+    0xff01010001000100, 0xff010100010100ff, 0xff01010001010000, 0xff010101ffffffff,
+    0xff010101ffffff01, 0xff010101ffff01ff, 0xff010101ffff0101, 0xff010101ff01ffff,
+    0xff010101ff01ff01, 0xff010101ff0101ff, 0xff010101ff010101, 0xff01010100ff0000,
+    0xff0101010000ff00, 0xff01010100000001, 0xff01010100000100, 0xff01010100010000,
+    0xff01010101ffffff, 0xff01010101ffff01, 0xff01010101ff01ff, 0xff01010101ff0101,
+    0xff01010101000000, 0xff0101010101ffff, 0xff0101010101ff01, 0xff010101010101ff,
+    0xff01010101010101, 0x00ffffffffff0000, 0x00ffffffff00ff00, 0x00ffffffff000001,
+    0x00ffffffff010000, 0x00ffffff00ff0100, 0x00ffffff0000ff01, 0x00ffffff00000000,
+    0x00ffffff000001ff, 0x00ffffff00000101, 0x00ffffff0001ff00, 0x00ffffff000100ff,
+    0x00ffffff00010001, 0x00ffffff010000ff, 0x00ffffff01000100, 0x00ffffff0101ff00,
+    0x00ffffff01010001, 0x00ffff00ffffffff, 0x00ffff00ffffff00, 0x00ffff00ffff00ff,
+    0x00ffff00ffff0001, 0x00ffff00ffff0100, 0x00ffff00ff00ff01, 0x00ffff00ff000000,
+    0x00ffff00ff000001, 0x00ffff00ff0001ff, 0x00ffff00ff000101, 0x00ffff00ff01ff00,
+    0x00ffff00ff010001, 0x00ffff00ff010100, 0x00ffff0000ff0000, 0x00ffff0000ff01ff,
+    0x00ffff0000ff0101, 0x00ffff000000ff00, 0x00ffff00000000ff, 0x00ffff0000000000,
+    0x00ffff0000000001, 0x00ffff0000000100, 0x00ffff0000000101, 0x00ffff0000010000,
+    0x00ffff00000101ff, 0x00ffff0000010101, 0x00ffff0001ffff00, 0x00ffff0001ff00ff,
+    0x00ffff0001ff0001, 0x00ffff000100ffff, 0x00ffff000100ff01, 0x00ffff0001000000,
+    0x00ffff000101ffff, 0x00ffff000101ff00, 0x00ffff000101ff01, 0x00ffff01ffff0000,
+    0x00ffff01ff00ff00, 0x00ffff01ff0000ff, 0x00ffff01ff000001, 0x00ffff01ff010000,
+    0x00ffff0100ffff00, 0x00ffff010000ff01, 0x00ffff0100000000, 0x00ffff0100000101,
+    0x00ffff01000100ff, 0x00ffff0100010100, 0x00ffff0101ff0100, 0x00ffff01010000ff,
+    0x00ffff0101010000, 0x00ff00ffffffff00, 0x00ff00ffff000000, 0x00ff00ffff000100,
+    0x00ff00ffff010100, 0x00ff00ff00ff0000, 0x00ff00ff00ff01ff, 0x00ff00ff00ff0101,
+    0x00ff00ff0000ff00, 0x00ff00ff000000ff, 0x00ff00ff00000000, 0x00ff00ff00000001,
+    0x00ff00ff0001ff00, 0x00ff00ff0001ff01, 0x00ff00ff00010000, 0x00ff00ff000101ff,
+    0x00ff00ff00010101, 0x00ff00ff01ffff00, 0x00ff00ff01ff0001, 0x00ff00ff01ff0100,
+    0x00ff00ff0100ffff, 0x00ff00ff0100ff01, 0x00ff00ff01000000, 0x00ff00ff0101ffff,
+    0x00ff00ff0101ff00, 0x00ff00ff01010100, 0x00ff0000ffffff00, 0x00ff0000ffffff01,
+    0x00ff0000ffff0000, 0x00ff0000ffff0101, 0x00ff0000ff00ff00, 0x00ff0000ff0000ff,
+    0x00ff0000ff000000, 0x00ff0000ff000001, 0x00ff0000ff000100, 0x00ff0000ff01ffff,
+    0x00ff0000ff010000, 0x00ff0000ff010101, 0x00ff000000ffff00, 0x00ff000000ff00ff,
+    0x00ff000000ff0000, 0x00ff000000ff0001, 0x00ff000000ff0100, 0x00ff00000000ffff,
+    0x00ff00000000ff00, 0x00ff0000000000ff, 0x00ff000000000000, 0x00ff000000000001,
+    0x00ff0000000001ff, 0x00ff000000000100, 0x00ff00000001ff00, 0x00ff0000000100ff,
+    0x00ff000000010000, 0x00ff000000010001, 0x00ff000000010100, 0x00ff000001ffff01,
+    0x00ff000001ff00ff, 0x00ff000001ff0000, 0x00ff000001ff01ff, 0x00ff00000100ff00,
+    0x00ff0000010000ff, 0x00ff000001000000, 0x00ff000001000001, 0x00ff000001000100,
+    0x00ff000001000101, 0x00ff000001010000, 0x00ff0000010101ff, 0x00ff000001010101,
+    0x00ff0001ffffff00, 0x00ff0001ffff0000, 0x00ff0001ffff0100, 0x00ff0001ff0000ff,
+    0x00ff0001ff000000, 0x00ff0001ff0001ff, 0x00ff0001ff000101, 0x00ff0001ff01ff00,
+    0x00ff0001ff0100ff, 0x00ff0001ff010100, 0x00ff000100ffffff, 0x00ff000100ffff01,
+    0x00ff000100ff0000, 0x00ff000100ff01ff, 0x00ff00010000ffff, 0x00ff00010000ff00,
+    0x00ff00010000ff01, 0x00ff000100000000, 0x00ff000100000001, 0x00ff000100000100,
+    0x00ff00010001ff01, 0x00ff000100010000, 0x00ff0001000101ff, 0x00ff000101ffff00,
+    0x00ff000101ff0000, 0x00ff000101ff0101, 0x00ff0001010000ff, 0x00ff000101000000,
+    0x00ff00010101ff00, 0x00ff0001010100ff, 0x00ff000101010001, 0x00ff01ffffff0000,
+    0x00ff01ffff00ff00, 0x00ff01ffff000000, 0x00ff01ffff000101, 0x00ff01ffff010000,
+    0x00ff01ff00ffff01, 0x00ff01ff00ff0100, 0x00ff01ff0000ffff, 0x00ff01ff00000000,
+    0x00ff01ff000001ff, 0x00ff01ff0001ff00, 0x00ff01ff000100ff, 0x00ff01ff00010001,
+    0x00ff01ff00010100, 0x00ff01ff01ff0000, 0x00ff01ff0100ff00, 0x00ff01ff010000ff,
+    0x00ff01ff01000001, 0x00ff01ff01000100, 0x00ff01ff01010000, 0x00ff0100ffffff00,
+    0x00ff0100ffff0000, 0x00ff0100ffff0001, 0x00ff0100ffff0101, 0x00ff0100ff00ffff,
+    0x00ff0100ff0000ff, 0x00ff0100ff000000, 0x00ff0100ff0001ff, 0x00ff0100ff01ff00,
+    0x00ff0100ff0100ff, 0x00ff0100ff010001, 0x00ff010000ffffff, 0x00ff010000ff0000,
+    0x00ff010000ff0101, 0x00ff01000000ff00, 0x00ff01000000ff01, 0x00ff0100000000ff,
+    0x00ff010000000000, 0x00ff010000000001, 0x00ff010000000100, 0x00ff01000001ffff,
+    0x00ff01000001ff01, 0x00ff010000010000, 0x00ff010000010001, 0x00ff010000010101,
+    0x00ff010001ff0001, 0x00ff010001ff0100, 0x00ff01000100ff01, 0x00ff010001000000,
+    0x00ff010001000001, 0x00ff0100010001ff, 0x00ff01000101ff00, 0x00ff0100010100ff,
+    0x00ff010001010001, 0x00ff010001010100, 0x00ff0101ff000001, 0x00ff010100ff00ff,
+    0x00ff010100ff0001, 0x00ff010100ff0100, 0x00ff010100000000, 0x00ff0101000001ff,
+    0x00ff010100000101, 0x00ff0101000100ff, 0x00ff010100010100, 0x00ff0101010000ff,
+    0x00ff010101010000, 0x0000ffffffffff00, 0x0000ffffffff00ff, 0x0000ffffffff0000,
+    0x0000ffffffff0001, 0x0000ffffffff0100, 0x0000ffffff00ff01, 0x0000ffffff000000,
+    0x0000ffffff000101, 0x0000ffffff01ff00, 0x0000ffffff0100ff, 0x0000ffffff010100,
+    0x0000ffff00ffffff, 0x0000ffff00ff0000, 0x0000ffff00ff01ff, 0x0000ffff0000ff00,
+    0x0000ffff000000ff, 0x0000ffff00000000, 0x0000ffff00000001, 0x0000ffff00000100,
+    0x0000ffff00010000, 0x0000ffff000101ff, 0x0000ffff01ff0001, 0x0000ffff01ff0100,
+    0x0000ffff01000000, 0x0000ffff010001ff, 0x0000ffff0101ffff, 0x0000ffff0101ff00,
+    0x0000ffff01010001, 0x0000ffff01010100, 0x0000ff00ffff0000, 0x0000ff00ffff01ff,
+    0x0000ff00ffff0100, 0x0000ff00ffff0101, 0x0000ff00ff00ff00, 0x0000ff00ff0000ff,
+    0x0000ff00ff000000, 0x0000ff00ff000001, 0x0000ff00ff0001ff, 0x0000ff00ff000100,
+    0x0000ff00ff01ffff, 0x0000ff00ff010000, 0x0000ff00ff010001, 0x0000ff00ff0101ff,
+    0x0000ff00ff010101, 0x0000ff0000ffff00, 0x0000ff0000ff00ff, 0x0000ff0000ff0000,
+    0x0000ff0000ff0001, 0x0000ff0000ff0100, 0x0000ff000000ffff, 0x0000ff000000ff00,
+    0x0000ff000000ff01, 0x0000ff00000000ff, 0x0000ff0000000000, 0x0000ff0000000001,
+    0x0000ff00000001ff, 0x0000ff0000000100, 0x0000ff0000000101, 0x0000ff000001ff00,
+    0x0000ff00000100ff, 0x0000ff0000010000, 0x0000ff0000010001, 0x0000ff0000010100,
+    0x0000ff0001ffff01, 0x0000ff0001ff0000, 0x0000ff000100ff00, 0x0000ff00010000ff,
+    0x0000ff0001000000, 0x0000ff0001000001, 0x0000ff0001000100, 0x0000ff000101ffff,
+    0x0000ff0001010000, 0x0000ff0001010101, 0x0000ff01ffffff00, 0x0000ff01ffff0001,
+    0x0000ff01ff00ff01, 0x0000ff01ff000000, 0x0000ff01ff000101, 0x0000ff01ff01ff00,
+    0x0000ff01ff0100ff, 0x0000ff0100ffff01, 0x0000ff0100ff0000, 0x0000ff0100ff0101,
+    0x0000ff010000ff00, 0x0000ff01000000ff, 0x0000ff0100000000, 0x0000ff0100000001,
+    0x0000ff0100000100, 0x0000ff010001ff01, 0x0000ff0100010000, 0x0000ff0101ff0000,
+    0x0000ff010100ffff, 0x0000ff010100ff01, 0x0000ff0101000000, 0x0000ff0101000100,
+    0x0000ff0101000101, 0x0000ff01010100ff, 0x000000ffffff00ff, 0x000000ffffff0000,
+    0x000000ffff00ff00, 0x000000ffff0000ff, 0x000000ffff000000, 0x000000ffff000001,
+    0x000000ffff0001ff, 0x000000ffff000100, 0x000000ffff01ff00, 0x000000ffff010000,
+    0x000000ffff0101ff, 0x000000ffff010101, 0x000000ff00ffff00, 0x000000ff00ff00ff,
+    0x000000ff00ff0000, 0x000000ff00ff0001, 0x000000ff00ff0100, 0x000000ff00ff0101,
+    0x000000ff0000ffff, 0x000000ff0000ff00, 0x000000ff000000ff, 0x000000ff00000000,
+    0x000000ff00000001, 0x000000ff000001ff, 0x000000ff00000100, 0x000000ff00000101,
+    0x000000ff0001ff00, 0x000000ff0001ff01, 0x000000ff000100ff, 0x000000ff00010000,
+    0x000000ff00010001, 0x000000ff00010100, 0x000000ff01ffffff, 0x000000ff01ff01ff,
+    0x000000ff01ff0101, 0x000000ff0100ff00, 0x000000ff010000ff, 0x000000ff01000000,
+    0x000000ff01000001, 0x000000ff01000100, 0x000000ff0101ff00, 0x000000ff010100ff,
+    0x000000ff01010000, 0x000000ff01010101, 0x00000000ffffff00, 0x00000000ffffff01,
+    0x00000000ffff00ff, 0x00000000ffff0000, 0x00000000ffff0001, 0x00000000ffff0100,
+    0x00000000ff00ffff, 0x00000000ff00ff00, 0x00000000ff00ff01, 0x00000000ff0000ff,
+    0x00000000ff000000, 0x00000000ff000001, 0x00000000ff000100, 0x00000000ff000101,
+    0x00000000ff01ff00, 0x00000000ff0100ff, 0x00000000ff010000, 0x00000000ff010001,
+    0x00000000ff010100, 0x0000000000ffffff, 0x0000000000ffff00, 0x0000000000ffff01,
+    0x0000000000ff00ff, 0x0000000000ff0000, 0x0000000000ff0001, 0x0000000000ff01ff,
+    0x0000000000ff0100, 0x000000000000ffff, 0x000000000000ff00, 0x000000000000ff01,
+    0x00000000000000ff, 0x0000000000000000, 0x0000000000000001, 0x00000000000001ff,
+    0x0000000000000100, 0x0000000000000101, 0x000000000001ffff, 0x000000000001ff00,
+    0x00000000000100ff, 0x0000000000010000, 0x0000000000010001, 0x00000000000101ff,
+    0x0000000000010100, 0x0000000000010101, 0x0000000001ffff00, 0x0000000001ff00ff,
+    0x0000000001ff0000, 0x0000000001ff0100, 0x0000000001ff0101, 0x000000000100ffff,
+    0x000000000100ff00, 0x00000000010000ff, 0x0000000001000000, 0x0000000001000001,
+    0x00000000010001ff, 0x0000000001000100, 0x000000000101ff00, 0x00000000010100ff,
+    0x0000000001010000, 0x0000000001010001, 0x0000000001010100, 0x00000001ffffffff,
+    0x00000001ffffff00, 0x00000001ffffff01, 0x00000001ffff00ff, 0x00000001ffff0001,
+    0x00000001ffff01ff, 0x00000001ffff0100, 0x00000001ff00ff00, 0x00000001ff0000ff,
+    0x00000001ff000000, 0x00000001ff0001ff, 0x00000001ff000100, 0x00000001ff01ffff,
+    0x00000001ff01ff00, 0x00000001ff01ff01, 0x00000001ff0100ff, 0x00000001ff010000,
+    0x00000001ff010001, 0x00000001ff0101ff, 0x00000001ff010100, 0x0000000100ffff00,
+    0x0000000100ff0000, 0x0000000100ff0001, 0x0000000100ff01ff, 0x0000000100ff0100,
+    0x0000000100ff0101, 0x000000010000ffff, 0x000000010000ff00, 0x000000010000ff01,
+    0x00000001000000ff, 0x0000000100000000, 0x0000000100000001, 0x00000001000001ff,
+    0x0000000100000100, 0x0000000100000101, 0x000000010001ff00, 0x00000001000100ff,
+    0x0000000100010000, 0x0000000100010100, 0x0000000101ffff01, 0x0000000101ff0000,
+    0x0000000101ff0001, 0x0000000101ff01ff, 0x0000000101ff0100, 0x0000000101ff0101,
+    0x000000010100ff00, 0x0000000101000000, 0x0000000101000101, 0x000000010101ff01,
+    0x0000000101010000, 0x0000000101010001, 0x00000001010101ff, 0x0000000101010100,
+    0x000001ffffff00ff, 0x000001ffffff0000, 0x000001ffffff0001, 0x000001ffffff0100,
+    0x000001ffff00ffff, 0x000001ffff000000, 0x000001ffff0001ff, 0x000001ffff01ff00,
+    0x000001ffff010101, 0x000001ff00ff0000, 0x000001ff00ff01ff, 0x000001ff00ff0101,
+    0x000001ff0000ff00, 0x000001ff000000ff, 0x000001ff00000000, 0x000001ff00000001,
+    0x000001ff000001ff, 0x000001ff00000100, 0x000001ff0001ffff, 0x000001ff0001ff01,
+    0x000001ff000100ff, 0x000001ff00010000, 0x000001ff01ffff01, 0x000001ff01ff0100,
+    0x000001ff0100ffff, 0x000001ff0100ff01, 0x000001ff01000000, 0x000001ff010001ff,
+    0x000001ff0101ff00, 0x000001ff01010100, 0x00000100ffffff00, 0x00000100ffffff01,
+    0x00000100ffff0000, 0x00000100ffff0101, 0x00000100ff00ff00, 0x00000100ff0000ff,
+    0x00000100ff000000, 0x00000100ff000001, 0x00000100ff000100, 0x00000100ff010000,
+    0x0000010000ffff00, 0x0000010000ff00ff, 0x0000010000ff0000, 0x0000010000ff0001,
+    0x0000010000ff0100, 0x000001000000ffff, 0x000001000000ff00, 0x000001000000ff01,
+    0x00000100000000ff, 0x0000010000000000, 0x0000010000000001, 0x00000100000001ff,
+    0x0000010000000100, 0x0000010000000101, 0x000001000001ff00, 0x00000100000100ff,
+    0x0000010000010000, 0x0000010000010001, 0x0000010000010100, 0x0000010001ffff00,
+    0x0000010001ff0000, 0x0000010001ff0100, 0x000001000100ff00, 0x00000100010000ff,
+    0x0000010001000000, 0x0000010001000001, 0x00000100010001ff, 0x0000010001000100,
+    0x0000010001010000, 0x00000101ffff00ff, 0x00000101ffff01ff, 0x00000101ff000000,
+    0x00000101ff000101, 0x00000101ff01ffff, 0x00000101ff010000, 0x00000101ff010001,
+    0x00000101ff010100, 0x0000010100ff0000, 0x0000010100ff01ff, 0x0000010100ff0100,
+    0x000001010000ff00, 0x0000010100000000, 0x0000010100000001, 0x00000101000001ff,
+    0x0000010100000100, 0x000001010001ff01, 0x0000010100010000, 0x00000101000101ff,
+    0x0000010100010101, 0x0000010101ffff00, 0x0000010101ff0101, 0x000001010100ff01,
+    0x0000010101000000, 0x0000010101000001, 0x00000101010001ff, 0x0000010101000101,
+    0x000001010101ff00, 0x0001ffffffff0000, 0x0001ffffff0000ff, 0x0001ffffff000001,
+    0x0001ffffff000100, 0x0001ffffff010000, 0x0001ffff00ff00ff, 0x0001ffff0000ffff,
+    0x0001ffff00000000, 0x0001ffff00000001, 0x0001ffff000001ff, 0x0001ffff00000101,
+    0x0001ffff0001ff00, 0x0001ffff000100ff, 0x0001ffff00010001, 0x0001ffff00010100,
+    0x0001ffff01ffff00, 0x0001ffff01000001, 0x0001ffff01010000, 0x0001ff00ffffff00,
+    0x0001ff00ffff00ff, 0x0001ff00ffff0001, 0x0001ff00ffff0100, 0x0001ff00ff00ff01,
+    0x0001ff00ff000000, 0x0001ff00ff01ff00, 0x0001ff00ff01ff01, 0x0001ff00ff010001,
+    0x0001ff00ff010100, 0x0001ff0000ff0000, 0x0001ff0000ff0100, 0x0001ff000000ff00,
+    0x0001ff0000000000, 0x0001ff0000000001, 0x0001ff0000000100, 0x0001ff0000010000,
+    0x0001ff0000010001, 0x0001ff0000010101, 0x0001ff0001ff00ff, 0x0001ff0001ff0101,
+    0x0001ff000100ff01, 0x0001ff0001000000, 0x0001ff000101ff00, 0x0001ff0001010001,
+    0x0001ff0001010100, 0x0001ff01ff00ff00, 0x0001ff01ff000001, 0x0001ff01ff000100,
+    0x0001ff0100ffffff, 0x0001ff0100ffff00, 0x0001ff0100ff0001, 0x0001ff0100000000,
+    0x0001ff0100000001, 0x0001ff01000001ff, 0x0001ff010001ffff, 0x0001ff0101ff0000,
+    0x0001ff010100ff00, 0x0001ff0101000001, 0x0001ff0101010000, 0x000100ffff00ff00,
+    0x000100ffff00ff01, 0x000100ffff000000, 0x000100ffff000001, 0x000100ffff000101,
+    0x000100ffff01ff00, 0x000100ffff010001, 0x000100ffff010100, 0x000100ff00ffffff,
+    0x000100ff00ffff01, 0x000100ff00ff0000, 0x000100ff00ff01ff, 0x000100ff00ff0101,
+    0x000100ff0000ff00, 0x000100ff000000ff, 0x000100ff00000000, 0x000100ff00000001,
+    0x000100ff00000100, 0x000100ff00000101, 0x000100ff0001ffff, 0x000100ff0001ff01,
+    0x000100ff00010000, 0x000100ff01ff00ff, 0x000100ff01ff0000, 0x000100ff01ff0100,
+    0x000100ff0100ffff, 0x000100ff0100ff01, 0x000100ff010000ff, 0x000100ff01000000,
+    0x000100ff01000001, 0x000100ff010001ff, 0x000100ff01000101, 0x000100ff0101ff00,
+    0x000100ff010100ff, 0x000100ff01010100, 0x00010000ffff0000, 0x00010000ffff01ff,
+    0x00010000ffff0101, 0x00010000ff00ff00, 0x00010000ff000000, 0x00010000ff000001,
+    0x00010000ff000100, 0x0001000000ff00ff, 0x0001000000ff0000, 0x0001000000ff0001,
+    0x0001000000ff0100, 0x000100000000ffff, 0x000100000000ff00, 0x00010000000000ff,
+    0x0001000000000000, 0x0001000000000001, 0x0001000000000100, 0x000100000001ff00,
+    0x00010000000100ff, 0x0001000000010000, 0x0001000000010001, 0x0001000000010100,
+    0x0001000001ff0001, 0x0001000001ff0100, 0x0001000001ff0101, 0x000100000100ff00,
+    0x0001000001000000, 0x0001000001000001, 0x0001000001000100, 0x0001000001000101,
+    0x000100000101ff01, 0x0001000001010000, 0x0001000001010001, 0x00010000010101ff,
+    0x00010001ffffff01, 0x00010001ffff0100, 0x00010001ff000000, 0x00010001ff01ffff,
+    0x00010001ff010001, 0x00010001ff0101ff, 0x00010001ff010100, 0x0001000100ffffff,
+    0x0001000100ff0000, 0x0001000100ff01ff, 0x0001000100ff0101, 0x000100010000ff00,
+    0x00010001000000ff, 0x0001000100000000, 0x0001000100000001, 0x00010001000001ff,
+    0x0001000100000101, 0x000100010001ffff, 0x0001000100010000, 0x00010001000101ff,
+    0x0001000101ffffff, 0x0001000101ffff01, 0x0001000101ff0000, 0x0001000101ff0101,
+    0x00010001010000ff, 0x0001000101000001, 0x00010001010001ff, 0x0001000101000100,
+    0x000100010101ffff, 0x00010001010100ff, 0x0001000101010001, 0x0001000101010101,
+    0x000101ffff000001, 0x000101ffff000100, 0x000101ffff010000, 0x000101ff00ffff00,
+    0x000101ff0000ff01, 0x000101ff00000000, 0x000101ff00000101, 0x000101ff0001ff00,
+    0x000101ff00010100, 0x000101ff01ff0000, 0x000101ff0100ff00, 0x000101ff010001ff,
+    0x000101ff01010001, 0x00010100ffffff00, 0x00010100ffff00ff, 0x00010100ff00ffff,
+    0x00010100ff000000, 0x00010100ff01ff00, 0x00010100ff0100ff, 0x00010100ff010001,
+    0x00010100ff010100, 0x0001010000ffffff, 0x0001010000ffff00, 0x0001010000ff0000,
+    0x0001010000ff0001, 0x0001010000ff01ff, 0x000101000000ff00, 0x00010100000000ff,
+    0x0001010000000000, 0x0001010000000001, 0x0001010000000100, 0x000101000001ffff,
+    0x0001010000010000, 0x0001010000010101, 0x0001010001ffff01, 0x0001010001ff00ff,
+    0x0001010001ff0101, 0x0001010001000000, 0x000101000101ff00, 0x00010100010100ff,
+    0x0001010001010000, 0x0001010001010100, 0x00010101ff00ff00, 0x00010101ff000001,
+    0x00010101ff0001ff, 0x0001010100ffff00, 0x0001010100ff00ff, 0x0001010100ff0100,
+    0x000101010000ffff, 0x0001010100000000, 0x00010101000001ff, 0x0001010100000101,
+    0x00010101000100ff, 0x0001010100010000, 0x0001010100010100, 0x0001010101ff0001,
+    0x00010101010000ff, 0x00010101010001ff, 0x0001010101000101, 0x0001010101010001,
+    0x01ffffffffffffff, 0x01ffffffffffff01, 0x01ffffffffff01ff, 0x01ffffffffff0101,
+    0x01ffffffff01ffff, 0x01ffffffff01ff01, 0x01ffffffff0101ff, 0x01ffffffff010101,
+    0x01ffffff00ff0000, 0x01ffffff0000ffff, 0x01ffffff0000ff00, 0x01ffffff000000ff,
+    0x01ffffff00000001, 0x01ffffff00000100, 0x01ffffff00010000, 0x01ffffff01ffffff,
+    0x01ffffff01ffff01, 0x01ffffff01ff01ff, 0x01ffffff01ff0101, 0x01ffffff01000000,
+    0x01ffffff0101ffff, 0x01ffffff0101ff01, 0x01ffffff010101ff, 0x01ffffff01010101,
+    0x01ffff00ffff0000, 0x01ffff00ff00ff00, 0x01ffff00ff0000ff, 0x01ffff00ff000001,
+    0x01ffff00ff000100, 0x01ffff00ff010000, 0x01ffff0000ffff00, 0x01ffff0000ff00ff,
+    0x01ffff0000ff0100, 0x01ffff000000ffff, 0x01ffff000000ff01, 0x01ffff0000000000,
+    0x01ffff0000000001, 0x01ffff00000001ff, 0x01ffff0000000100, 0x01ffff00000100ff,
+    0x01ffff0000010001, 0x01ffff0000010100, 0x01ffff0001ff0000, 0x01ffff0001ff0100,
+    0x01ffff00010000ff, 0x01ffff0001000001, 0x01ffff0001000100, 0x01ffff0001010000,
+    0x01ffff01ffffffff, 0x01ffff01ffffff01, 0x01ffff01ffff01ff, 0x01ffff01ffff0101,
+    0x01ffff01ff000000, 0x01ffff01ff01ffff, 0x01ffff01ff01ff01, 0x01ffff01ff0101ff,
+    0x01ffff01ff010101, 0x01ffff010000ff00, 0x01ffff01000000ff, 0x01ffff0100000100,
+    0x01ffff0100010000, 0x01ffff0101ffffff, 0x01ffff0101ffff01, 0x01ffff0101ff01ff,
+    0x01ffff0101ff0101, 0x01ffff0101000000, 0x01ffff010101ffff, 0x01ffff010101ff01,
+    0x01ffff01010101ff, 0x01ffff0101010101, 0x01ff00ffff0000ff, 0x01ff00ffff000100,
+    0x01ff00ff00ffff00, 0x01ff00ff00ff00ff, 0x01ff00ff0000ff00, 0x01ff00ff00000000,
+    0x01ff00ff00000101, 0x01ff00ff0001ff00, 0x01ff00ff000100ff, 0x01ff00ff00010100,
+    0x01ff00ff010000ff, 0x01ff00ff01000100, 0x01ff0000ffffff00, 0x01ff0000ffff0100,
+    0x01ff0000ff00ff01, 0x01ff0000ff000000, 0x01ff0000ff000101, 0x01ff0000ff010001,
+    0x01ff0000ff010100, 0x01ff000000ffffff, 0x01ff000000ffff00, 0x01ff000000ff0000,
+    0x01ff000000ff01ff, 0x01ff00000000ff00, 0x01ff0000000000ff, 0x01ff000000000000,
+    0x01ff000000000001, 0x01ff000000000100, 0x01ff000000000101, 0x01ff000000010000,
+    0x01ff000000010001, 0x01ff0000000101ff, 0x01ff000000010101, 0x01ff000001ffff00,
+    0x01ff000001ff00ff, 0x01ff000001ff0001, 0x01ff000001ff0100, 0x01ff00000100ffff,
+    0x01ff00000100ff01, 0x01ff000001000000, 0x01ff0000010001ff, 0x01ff000001010001,
+    0x01ff0001ff00ff00, 0x01ff0001ff000001, 0x01ff0001ff000100, 0x01ff0001ff010000,
+    0x01ff000100ffff00, 0x01ff000100ff00ff, 0x01ff000100ff0100, 0x01ff000100ff0101,
+    0x01ff00010000ffff, 0x01ff000100000000, 0x01ff000100000100, 0x01ff000100000101,
+    0x01ff00010001ff00, 0x01ff000100010001, 0x01ff000100010101, 0x01ff000101ff0000,
+    0x01ff00010100ff00, 0x01ff000101000101, 0x01ff0001010100ff, 0x01ff01ffffffffff,
+    0x01ff01ffffffff01, 0x01ff01ffffff01ff, 0x01ff01ffffff0101, 0x01ff01ffff000000,
+    0x01ff01ffff01ffff, 0x01ff01ffff01ff01, 0x01ff01ffff0101ff, 0x01ff01ffff010101,
+    0x01ff01ff00ffff00, 0x01ff01ff00ff0000, 0x01ff01ff0000ff00, 0x01ff01ff000000ff,
+    0x01ff01ff00000100, 0x01ff01ff00010000, 0x01ff01ff00010100, 0x01ff01ff01ffffff,
+    0x01ff01ff01ffff01, 0x01ff01ff01ff01ff, 0x01ff01ff01ff0101, 0x01ff01ff01000000,
+    0x01ff01ff0101ffff, 0x01ff01ff0101ff01, 0x01ff01ff010101ff, 0x01ff01ff01010101,
+    0x01ff0100ffff0000, 0x01ff0100ffff0001, 0x01ff0100ff00ff00, 0x01ff0100ff0000ff,
+    0x01ff0100ff000001, 0x01ff0100ff010000, 0x01ff010000ffff00, 0x01ff010000ff00ff,
+    0x01ff010000ff0001, 0x01ff010000ff0100, 0x01ff01000000ffff, 0x01ff01000000ff01,
+    0x01ff010000000000, 0x01ff010000000101, 0x01ff01000001ff00, 0x01ff0100000100ff,
+    0x01ff010001ff0000, 0x01ff010001000001, 0x01ff010001000100, 0x01ff010001010000,
+    0x01ff0101ffffffff, 0x01ff0101ffffff01, 0x01ff0101ffff01ff, 0x01ff0101ffff0101,
+    0x01ff0101ff000000, 0x01ff0101ff01ffff, 0x01ff0101ff01ff01, 0x01ff0101ff0101ff,
+    0x01ff0101ff010101, 0x01ff010100ff0000, 0x01ff01010000ff00, 0x01ff0101000000ff,
+    0x01ff010100000001, 0x01ff010101ffffff, 0x01ff010101ffff01, 0x01ff010101ff01ff,
+    0x01ff010101ff0101, 0x01ff010101000000, 0x01ff01010101ffff, 0x01ff01010101ff01,
+    0x01ff0101010101ff, 0x01ff010101010101, 0x0100ffffffff0000, 0x0100ffffff00ff00,
+    0x0100ffffff000001, 0x0100ffffff0001ff, 0x0100ffffff000100, 0x0100ffffff010000,
+    0x0100ffff00ffff00, 0x0100ffff00ff0001, 0x0100ffff00ff0100, 0x0100ffff00000000,
+    0x0100ffff000001ff, 0x0100ffff00000101, 0x0100ffff00010100, 0x0100ffff00010101,
+    0x0100ffff01ff0000, 0x0100ffff0100ff00, 0x0100ffff010000ff, 0x0100ffff01000001,
+    0x0100ffff01000100, 0x0100ffff01010000, 0x0100ff00ffffff00, 0x0100ff00ffff00ff,
+    0x0100ff00ffff0001, 0x0100ff00ffff0100, 0x0100ff00ff00ffff, 0x0100ff00ff000000,
+    0x0100ff00ff0001ff, 0x0100ff00ff000101, 0x0100ff00ff01ff00, 0x0100ff00ff0100ff,
+    0x0100ff00ff010001, 0x0100ff00ff010100, 0x0100ff0000ffffff, 0x0100ff0000ff0000,
+    0x0100ff000000ffff, 0x0100ff000000ff00, 0x0100ff00000000ff, 0x0100ff0000000000,
+    0x0100ff0000000001, 0x0100ff0000000100, 0x0100ff000001ff01, 0x0100ff0000010000,
+    0x0100ff0001ff00ff, 0x0100ff0001ff0001, 0x0100ff000100ff01, 0x0100ff0001000000,
+    0x0100ff00010001ff, 0x0100ff000101ff00, 0x0100ff00010100ff, 0x0100ff0001010001,
+    0x0100ff0001010100, 0x0100ff01ffff0000, 0x0100ff01ff00ff00, 0x0100ff01ff0000ff,
+    0x0100ff01ff000100, 0x0100ff01ff010000, 0x0100ff0100ff00ff, 0x0100ff0100ff0001,
+    0x0100ff0100ff0100, 0x0100ff010000ffff, 0x0100ff010000ff01, 0x0100ff0100000000,
+    0x0100ff01000001ff, 0x0100ff0100010001, 0x0100ff0100010100, 0x0100ff0101ff0000,
+    0x0100ff01010000ff, 0x0100ff0101000001, 0x0100ff0101010100, 0x010000ffffffff00,
+    0x010000ffffff00ff, 0x010000ffffff0001, 0x010000ffff00ffff, 0x010000ffff000000,
+    0x010000ffff0001ff, 0x010000ffff010001, 0x010000ff00ffffff, 0x010000ff00ff0101,
+    0x010000ff0000ff00, 0x010000ff000000ff, 0x010000ff00000000, 0x010000ff00000001,
+    0x010000ff000001ff, 0x010000ff00000100, 0x010000ff0001ffff, 0x010000ff0001ff00,
+    0x010000ff0001ff01, 0x010000ff00010000, 0x010000ff01ff00ff, 0x010000ff01ff0001,
+    0x010000ff0100ff01, 0x010000ff010000ff, 0x010000ff01000000, 0x010000ff010001ff,
+    0x010000ff0101ff00, 0x010000ff01010100, 0x01000000ffffffff, 0x01000000ffff0000,
+    0x01000000ffff01ff, 0x01000000ffff0101, 0x01000000ff00ffff, 0x01000000ff00ff00,
+    0x01000000ff0000ff, 0x01000000ff000000, 0x01000000ff000001, 0x01000000ff000100,
+    0x01000000ff01ff00, 0x01000000ff010000, 0x01000000ff010100, 0x01000000ff010101,
+    0x0100000000ffff00, 0x0100000000ff00ff, 0x0100000000ff0000, 0x0100000000ff0001,
+    0x0100000000ff0100, 0x010000000000ffff, 0x010000000000ff00, 0x010000000000ff01,
+    0x01000000000000ff, 0x0100000000000000, 0x0100000000000001, 0x01000000000001ff,
+    0x0100000000000100, 0x0100000000000101, 0x010000000001ff00, 0x01000000000100ff,
+    0x0100000000010000, 0x0100000000010001, 0x0100000000010100, 0x0100000001ffff00,
+    0x0100000001ff0000, 0x0100000001ff01ff, 0x010000000100ff00, 0x010000000100ff01,
+    0x01000000010000ff, 0x0100000001000000, 0x0100000001000001, 0x0100000001000100,
+    0x0100000001000101, 0x010000000101ffff, 0x010000000101ff01, 0x0100000001010000,
+    0x01000000010101ff, 0x0100000001010101, 0x01000001ffffff00, 0x01000001ffff00ff,
+    0x01000001ff00ffff, 0x01000001ff000000, 0x01000001ff000100, 0x01000001ff01ffff,
+    0x01000001ff010001, 0x01000001ff010100, 0x0100000100ff0000, 0x0100000100ff01ff,
+    0x0100000100ff0100, 0x010000010000ff00, 0x010000010000ff01, 0x0100000100000000,
+    0x0100000100000001, 0x0100000100000100, 0x0100000100010000, 0x01000001000101ff,
+    0x0100000101ffff01, 0x0100000101ff00ff, 0x0100000101ff0100, 0x0100000101ff0101,
+    0x010000010100ff01, 0x01000001010000ff, 0x0100000101000000, 0x01000001010100ff,
+    0x0100000101010001, 0x0100000101010100, 0x010001ffffff0000, 0x010001ffff000001,
+    0x010001ffff000100, 0x010001ffff010000, 0x010001ff00ffff00, 0x010001ff00ff0001,
+    0x010001ff0000ffff, 0x010001ff0000ff01, 0x010001ff00000000, 0x010001ff00000001,
+    0x010001ff00000101, 0x010001ff000100ff, 0x010001ff00010000, 0x010001ff01ff0000,
+    0x010001ff0100ff00, 0x010001ff01000001, 0x010001ff01000100, 0x010001ff01010000,
+    0x01000100ffff00ff, 0x01000100ffff0001, 0x01000100ffff0100, 0x01000100ff00ffff,
+    0x01000100ff00ff01, 0x01000100ff000000, 0x01000100ff0001ff, 0x01000100ff000101,
+    0x01000100ff01ffff, 0x01000100ff01ff00, 0x01000100ff0100ff, 0x01000100ff010001,
+    0x0100010000ffffff, 0x0100010000ffff01, 0x0100010000ff0000, 0x0100010000ff01ff,
+    0x0100010000ff0101, 0x010001000000ff00, 0x01000100000000ff, 0x0100010000000000,
+    0x0100010000000001, 0x0100010000000100, 0x010001000001ff01, 0x0100010000010000,
+    0x0100010000010001, 0x0100010000010101, 0x0100010001ffff00, 0x0100010001ff00ff,
+    0x010001000100ffff, 0x010001000100ff01, 0x0100010001000000, 0x0100010001000101,
+    0x010001000101ff00, 0x0100010001010001, 0x01000101ffff0000, 0x01000101ff000000,
+    0x01000101ff010000, 0x0100010100ff00ff, 0x0100010100ff0001, 0x0100010100ff0100,
+    0x010001010000ffff, 0x0100010100000000, 0x01000101000001ff, 0x010001010001ff00,
+    0x0100010101ff0000, 0x010001010100ff00, 0x01000101010000ff, 0x0100010101000000,
+    0x0100010101000001, 0x0101ffffffffffff, 0x0101ffffffffff01, 0x0101ffffffff01ff,
+    0x0101ffffffff0101, 0x0101ffffff000000, 0x0101ffffff01ffff, 0x0101ffffff01ff01,
+    0x0101ffffff0101ff, 0x0101ffffff010101, 0x0101ffff00ff0000, 0x0101ffff0000ff00,
+    0x0101ffff000000ff, 0x0101ffff00000001, 0x0101ffff00000100, 0x0101ffff01ffffff,
+    0x0101ffff01ffff01, 0x0101ffff01ff01ff, 0x0101ffff01ff0101, 0x0101ffff01000000,
+    0x0101ffff0101ffff, 0x0101ffff0101ff01, 0x0101ffff010101ff, 0x0101ffff01010101,
+    0x0101ff00ffff0000, 0x0101ff00ffff0100, 0x0101ff00ff00ff00, 0x0101ff00ff0000ff,
+    0x0101ff00ff000001, 0x0101ff00ff000100, 0x0101ff00ff000101, 0x0101ff0000ff0001,
+    0x0101ff0000ff0100, 0x0101ff000000ff00, 0x0101ff0000000000, 0x0101ff00000001ff,
+    0x0101ff0000000101, 0x0101ff000001ff00, 0x0101ff00000100ff, 0x0101ff0001ff0000,
+    0x0101ff000100ffff, 0x0101ff000100ff01, 0x0101ff0001000001, 0x0101ff0001000100,
+    0x0101ff01ffffff01, 0x0101ff01ffff01ff, 0x0101ff01ffff0101, 0x0101ff01ff00ffff,
+    0x0101ff01ff000100, 0x0101ff01ff01ff01, 0x0101ff01ff0101ff, 0x0101ff01ff010101,
+    0x0101ff0100ff0000, 0x0101ff010000ff00, 0x0101ff0100000001, 0x0101ff0100000100,
+    0x0101ff0100010000, 0x0101ff0101ffffff, 0x0101ff0101ffff01, 0x0101ff0101ff01ff,
+    0x0101ff0101ff0101, 0x0101ff0101000000, 0x0101ff010101ffff, 0x0101ff010101ff01,
+    0x0101ff01010101ff, 0x0101ff0101010101, 0x010100ffff000100, 0x010100ffff010000,
+    0x010100ff00ffff00, 0x010100ff00ff00ff, 0x010100ff0000ffff, 0x010100ff000000ff,
+    0x010100ff00000000, 0x010100ff000001ff, 0x010100ff00000101, 0x010100ff0001ff00,
+    0x010100ff00010000, 0x010100ff00010001, 0x010100ff000101ff, 0x010100ff00010100,
+    0x010100ff01ff0000, 0x01010000ffff0001, 0x01010000ffff0100, 0x01010000ff00ffff,
+    0x01010000ff00ff01, 0x01010000ff000000, 0x01010000ff0001ff, 0x01010000ff010001,
+    0x01010000ff010100, 0x0101000000ffff01, 0x0101000000ff0000, 0x010100000000ff00,
+    0x01010000000000ff, 0x0101000000000000, 0x0101000000000001, 0x0101000000000100,
+    0x0101000000010000, 0x0101000000010101, 0x0101000001ffff00, 0x0101000001ff00ff,
+    0x0101000001ff0000, 0x0101000001ff0001, 0x0101000001ff0100, 0x010100000100ff01,
+    0x0101000001000000, 0x01010000010001ff, 0x01010001ffff0000, 0x01010001ff00ff00,
+    0x01010001ff000001, 0x01010001ff000101, 0x01010001ff01ff00, 0x01010001ff010000,
+    0x0101000100ff00ff, 0x0101000100ff0001, 0x0101000100ff0101, 0x010100010000ff01,
+    0x0101000100000000, 0x0101000100000001, 0x01010001000001ff, 0x010100010001ffff,
+    0x010100010001ff01, 0x0101000101ff0001, 0x010100010100ffff, 0x0101000101000000,
+    0x0101000101000001, 0x0101000101000100, 0x010100010101ff00, 0x01010001010100ff,
+    0x0101000101010001, 0x010101ffffffffff, 0x010101ffffffff01, 0x010101ffffff01ff,
+    0x010101ffffff0101, 0x010101ffff01ffff, 0x010101ffff01ff01, 0x010101ffff0101ff,
+    0x010101ffff010101, 0x010101ff0000ff00, 0x010101ff000000ff, 0x010101ff00000001,
+    0x010101ff00000100, 0x010101ff01ffffff, 0x010101ff01ffff01, 0x010101ff01ff01ff,
+    0x010101ff01ff0101, 0x010101ff01000000, 0x010101ff0101ffff, 0x010101ff0101ff01,
+    0x010101ff010101ff, 0x010101ff01010101, 0x01010100ffff0000, 0x01010100ff0000ff,
+    0x01010100ff000100, 0x01010100ff01ff00, 0x01010100ff010000, 0x0101010000ffff00,
+    0x010101000000ffff, 0x0101010000000000, 0x0101010000000101, 0x010101000001ff00,
+    0x0101010000010001, 0x0101010000010100, 0x010101000100ffff, 0x0101010001000001,
+    0x01010101ffffffff, 0x01010101ffffff01, 0x01010101ffff01ff, 0x01010101ffff0101,
+    0x01010101ff01ffff, 0x01010101ff01ff01, 0x01010101ff0101ff, 0x01010101ff010101,
+    0x010101010000ff00, 0x01010101000000ff, 0x0101010100000001, 0x0101010101ffffff,
+    0x0101010101ffff01, 0x0101010101ff01ff, 0x0101010101ff0101, 0x0101010101000000,
+    0x010101010101ffff, 0x010101010101ff01, 0x01010101010101ff, 0x0101010101010101,
+GGML_TABLE_END()
+#else
+GGML_TABLE_BEGIN(uint32_t, iq1s_grid_gpu, NGRID_IQ1S)
+    0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000,
+    0x00020002, 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101,
+    0x02000000, 0x02000002, 0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200,
+    0x02020202, 0x00000110, 0x00000111, 0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212,
+    0x00020111, 0x01000011, 0x01000112, 0x01000211, 0x01010012, 0x01010111, 0x01010212, 0x01020011,
+    0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011, 0x02010110, 0x02010112, 0x02020111,
+    0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020, 0x00020022, 0x00020220,
+    0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020, 0x02000022,
+    0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220,
+    0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101,
+    0x01011202, 0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110,
+    0x00001111, 0x00001112, 0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111,
+    0x01001212, 0x01011010, 0x01011011, 0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010,
+    0x01021012, 0x01021111, 0x01021210, 0x01021212, 0x02001011, 0x02011011, 0x02011111, 0x02011210,
+    0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112, 0x02021211, 0x00011120, 0x00011221,
+    0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220, 0x01021020, 0x01021021,
+    0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000, 0x00002002,
+    0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101,
+    0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101,
+    0x02022000, 0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211,
+    0x00022110, 0x00022111, 0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110,
+    0x01022211, 0x02012011, 0x02012110, 0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022,
+    0x00002220, 0x00002222, 0x00012121, 0x00022020, 0x00022022, 0x00022220, 0x00022222, 0x01002121,
+    0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020, 0x02002022, 0x02002121, 0x02002220,
+    0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222, 0x00110000, 0x00110001,
+    0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000, 0x01110101,
+    0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102,
+    0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012,
+    0x00110111, 0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010,
+    0x01110011, 0x01110012, 0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111,
+    0x02100110, 0x02110012, 0x02110111, 0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122,
+    0x00120121, 0x01100020, 0x01100122, 0x01100221, 0x01110022, 0x01110121, 0x01110220, 0x01110222,
+    0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120, 0x02110122, 0x02120121, 0x00101001,
+    0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201, 0x00121001, 0x00121102,
+    0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100, 0x01111101,
+    0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000,
+    0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101,
+    0x02121201, 0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112,
+    0x00111211, 0x00121010, 0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110,
+    0x01101111, 0x01101112, 0x01111011, 0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211,
+    0x01111212, 0x01121011, 0x01121110, 0x01121111, 0x01121112, 0x01121211, 0x02101010, 0x02101012,
+    0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010, 0x02111011, 0x02111110, 0x02111111,
+    0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111, 0x00101021, 0x00101120,
+    0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021, 0x00121122,
+    0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121,
+    0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221,
+    0x01121222, 0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001,
+    0x00112102, 0x00122101, 0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101,
+    0x01112200, 0x01112202, 0x01122000, 0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101,
+    0x02112001, 0x02112100, 0x02122101, 0x00112010, 0x00112012, 0x00112111, 0x00112212, 0x00122011,
+    0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210, 0x01112011, 0x01112110, 0x01112111,
+    0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212, 0x02102211, 0x02112011,
+    0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221, 0x00112122,
+    0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121,
+    0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222,
+    0x00200000, 0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101,
+    0x00220200, 0x00220202, 0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000,
+    0x02200002, 0x02200200, 0x02200202, 0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200,
+    0x02220202, 0x00200111, 0x00210011, 0x00210110, 0x00210211, 0x00220111, 0x01200012, 0x01200110,
+    0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011, 0x01220110, 0x01220111, 0x01220112,
+    0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021, 0x00200220, 0x00200222,
+    0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121, 0x01210021,
+    0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121,
+    0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201,
+    0x00221101, 0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200,
+    0x01211202, 0x01221102, 0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101,
+    0x00201211, 0x00211111, 0x00221011, 0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011,
+    0x01211110, 0x01211111, 0x01211211, 0x01221012, 0x01221111, 0x01221210, 0x02201211, 0x02211010,
+    0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011, 0x02221110, 0x02221112, 0x02221211,
+    0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021, 0x01201221, 0x01211121,
+    0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222, 0x00202000,
+    0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202,
+    0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202,
+    0x02222000, 0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211,
+    0x00222111, 0x01202112, 0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112,
+    0x01222211, 0x02202111, 0x02212010, 0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020,
+    0x00202022, 0x00202220, 0x00202222, 0x00222020, 0x00222022, 0x00222220, 0x00222222, 0x01202121,
+    0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020, 0x02202022, 0x02202220, 0x02202222,
+    0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101, 0x10010001, 0x10010102,
+    0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001, 0x11020100,
+    0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110,
+    0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011,
+    0x10020112, 0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111,
+    0x11010112, 0x11010211, 0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110,
+    0x12000112, 0x12010010, 0x12010012, 0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121,
+    0x10010021, 0x10010120, 0x10010122, 0x10020121, 0x11000021, 0x11010022, 0x11010121, 0x11010222,
+    0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121, 0x10001001, 0x10011101, 0x10011201,
+    0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100, 0x11011101, 0x11011102,
+    0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102, 0x12001201,
+    0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012,
+    0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010,
+    0x10021111, 0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010,
+    0x11011011, 0x11011110, 0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110,
+    0x11021111, 0x11021112, 0x11021211, 0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011,
+    0x12011110, 0x12011111, 0x12011112, 0x12011211, 0x12011212, 0x12021111, 0x12021210, 0x12021212,
+    0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121, 0x10011220, 0x10011222, 0x10021021,
+    0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220, 0x11011020, 0x11011021,
+    0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220, 0x12001021,
+    0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101,
+    0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101,
+    0x11012200, 0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100,
+    0x12012102, 0x12012201, 0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010,
+    0x10012110, 0x10012111, 0x10012210, 0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111,
+    0x11002212, 0x11012011, 0x11012012, 0x11012110, 0x11012111, 0x11012112, 0x11012211, 0x11022010,
+    0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112, 0x12002211, 0x12012012, 0x12012111,
+    0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211, 0x10012122, 0x11002120,
+    0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221, 0x12012120,
+    0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101,
+    0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001,
+    0x11110100, 0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201,
+    0x12110101, 0x12110200, 0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210,
+    0x10100211, 0x10100212, 0x10110011, 0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211,
+    0x10120010, 0x10120111, 0x10120112, 0x10120210, 0x10120212, 0x11100011, 0x11100110, 0x11100111,
+    0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012, 0x11110110, 0x11110111, 0x11110112,
+    0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111, 0x11120112, 0x11120211,
+    0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211, 0x12120010,
+    0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021,
+    0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122,
+    0x11110221, 0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221,
+    0x12110222, 0x12120120, 0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102,
+    0x10111200, 0x10111201, 0x10121001, 0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100,
+    0x11101101, 0x11101102, 0x11101201, 0x11101202, 0x11111000, 0x11111001, 0x11111100, 0x11111101,
+    0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001, 0x11121002, 0x11121100, 0x11121101,
+    0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001, 0x12111100, 0x12111101,
+    0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011, 0x10101012,
+    0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110,
+    0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112,
+    0x10121211, 0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210,
+    0x11101211, 0x11111010, 0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210,
+    0x11111211, 0x11111212, 0x11121010, 0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210,
+    0x11121211, 0x11121212, 0x12101011, 0x12101110, 0x12101111, 0x12101211, 0x12101212, 0x12111010,
+    0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210, 0x12111211, 0x12121011, 0x12121110,
+    0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022, 0x10101120, 0x10101122,
+    0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221, 0x10121020,
+    0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021,
+    0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022,
+    0x11111120, 0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120,
+    0x11121121, 0x11121221, 0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222,
+    0x12111021, 0x12111121, 0x12111222, 0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221,
+    0x10102100, 0x10102101, 0x10102102, 0x10102201, 0x10112000, 0x10112101, 0x10112200, 0x10122001,
+    0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001, 0x11112100, 0x11112101, 0x11112102,
+    0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101, 0x12102002, 0x12102201,
+    0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011, 0x10102012,
+    0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111,
+    0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012,
+    0x11112110, 0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110,
+    0x11122111, 0x11122112, 0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110,
+    0x12112111, 0x12112112, 0x12112210, 0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121,
+    0x10112222, 0x10122020, 0x10122121, 0x10122122, 0x10122221, 0x11102121, 0x11102220, 0x11102221,
+    0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221, 0x11122022, 0x11122121, 0x11122220,
+    0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122, 0x12112220, 0x12112222,
+    0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100, 0x11210000,
+    0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201,
+    0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012,
+    0x10210111, 0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011,
+    0x11210111, 0x11210112, 0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212,
+    0x12210012, 0x12210111, 0x12220011, 0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221,
+    0x11200020, 0x11200021, 0x11200122, 0x11210121, 0x11210122, 0x11210220, 0x11220020, 0x12200121,
+    0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002, 0x10211101, 0x10211102, 0x10211202,
+    0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101, 0x11201200, 0x11201202,
+    0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000, 0x11221002,
+    0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101,
+    0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210,
+    0x10201212, 0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112,
+    0x11201211, 0x11211010, 0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011,
+    0x11221110, 0x11221111, 0x11221112, 0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011,
+    0x12211111, 0x12211112, 0x12211211, 0x12211212, 0x12221012, 0x12221111, 0x12221112, 0x12221210,
+    0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122, 0x10221220, 0x10221221, 0x11201020,
+    0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121, 0x11211122, 0x11211220,
+    0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121, 0x12201222,
+    0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222,
+    0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001,
+    0x11222201, 0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010,
+    0x10212111, 0x10222011, 0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111,
+    0x11202112, 0x11202210, 0x11212011, 0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010,
+    0x11222111, 0x11222212, 0x12202012, 0x12202110, 0x12202212, 0x12212111, 0x12222011, 0x12222110,
+    0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220, 0x11202021, 0x11202120, 0x11202221,
+    0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121, 0x11222221, 0x12202122,
+    0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200, 0x20000202,
+    0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100,
+    0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101,
+    0x22020000, 0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112,
+    0x20010211, 0x20020111, 0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111,
+    0x21010112, 0x21010210, 0x21010211, 0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211,
+    0x22010110, 0x22010112, 0x22010211, 0x22020111, 0x20000020, 0x20000022, 0x20000220, 0x20000222,
+    0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222, 0x21010021, 0x21010120, 0x21010221,
+    0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121, 0x22020020, 0x22020022,
+    0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001, 0x21011101,
+    0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211,
+    0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111,
+    0x21001210, 0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111,
+    0x21021112, 0x21021210, 0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010,
+    0x22011012, 0x22011111, 0x22011210, 0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121,
+    0x21001021, 0x21001120, 0x21001221, 0x21001222, 0x21011020, 0x21011121, 0x21011221, 0x21011222,
+    0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021, 0x22011222, 0x22021120, 0x20002000,
+    0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002, 0x20022200, 0x20022202,
+    0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201, 0x22002000,
+    0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202,
+    0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110,
+    0x21002112, 0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110,
+    0x22002111, 0x22012112, 0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222,
+    0x20012121, 0x20022020, 0x20022022, 0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120,
+    0x21012122, 0x22002020, 0x22002022, 0x22002220, 0x22002222, 0x22012121, 0x22022020, 0x22022022,
+    0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102, 0x20110200, 0x20110201, 0x20120101,
+    0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202, 0x21120201, 0x21120202,
+    0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011, 0x20100110,
+    0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110,
+    0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111,
+    0x21110112, 0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111,
+    0x22110210, 0x22120011, 0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120,
+    0x20110221, 0x20120121, 0x21100120, 0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121,
+    0x21110220, 0x21120122, 0x21120221, 0x22100121, 0x22110120, 0x22110122, 0x22120221, 0x20101001,
+    0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200, 0x20121102, 0x21101000, 0x21101202,
+    0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201, 0x21121000, 0x21121001,
+    0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101, 0x22111200,
+    0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011,
+    0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212,
+    0x21101011, 0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012,
+    0x21111110, 0x21111111, 0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110,
+    0x21121111, 0x21121112, 0x21121211, 0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012,
+    0x22111110, 0x22111111, 0x22111112, 0x22111211, 0x22111212, 0x22121010, 0x22121012, 0x22121111,
+    0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020, 0x20111121, 0x20111221, 0x20121020,
+    0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021, 0x21111022, 0x21111121,
+    0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221, 0x22101222,
+    0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102,
+    0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102,
+    0x21112202, 0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101,
+    0x20102110, 0x20102112, 0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212,
+    0x20122010, 0x20122011, 0x20122110, 0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210,
+    0x21102212, 0x21112011, 0x21112110, 0x21112111, 0x21112112, 0x21112211, 0x21122012, 0x21122111,
+    0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010, 0x22112012, 0x22112111, 0x22112212,
+    0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120, 0x21102122, 0x21102221,
+    0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120, 0x22112121,
+    0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002,
+    0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000,
+    0x22200002, 0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202,
+    0x20200111, 0x20200211, 0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112,
+    0x21200211, 0x21210011, 0x21210111, 0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111,
+    0x22210010, 0x22210012, 0x22210112, 0x22210211, 0x20200022, 0x20200220, 0x20200222, 0x20210020,
+    0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121, 0x21210021, 0x21210122, 0x21210221,
+    0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121, 0x22220020, 0x22220022,
+    0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000, 0x21211100,
+    0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201,
+    0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112,
+    0x20221211, 0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211,
+    0x21221111, 0x21221212, 0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012,
+    0x22211111, 0x22211210, 0x20201121, 0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121,
+    0x21201120, 0x21201122, 0x21201222, 0x21211022, 0x21211121, 0x21211122, 0x21211220, 0x21221020,
+    0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122, 0x22211221, 0x22221021, 0x22221120,
+    0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000, 0x20222002, 0x20222200,
+    0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002, 0x22202200,
+    0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110,
+    0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011,
+    0x21222112, 0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222,
+    0x20222020, 0x20222022, 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020,
+    0x22202022, 0x22202220, 0x22202222, 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222,
+GGML_TABLE_END()
+#endif
+
+#endif // GGML_COMMON_IMPL
+#endif // GGML_COMMON_IMPL
diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h b/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h
new file mode 100644
index 0000000000..ae26e0c23b
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h
@@ -0,0 +1,795 @@
+#pragma once
+
+// GGML internal header
+
+#include "ggml.h"
+#include "gguf.h"
+
+#include <assert.h>
+#include <math.h>
+#include <stdlib.h> // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/
+#include <stdbool.h>
+#include <stdint.h>
+#include <string.h>
+
+#ifdef __ARM_FEATURE_SVE
+#include <arm_sve.h>
+#endif // __ARM_FEATURE_SVE
+
+#if defined(__ARM_NEON) && !defined(__CUDACC__) && !defined(__MUSACC__)
+// if YCM cannot find <arm_neon.h>, make a symbolic link to it, for example:
+//
+//   $ ln -sfn /Library/Developer/CommandLineTools/usr/lib/clang/13.1.6/include/arm_neon.h ./src/
+//
+#include <arm_neon.h>
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void ggml_print_backtrace(void);
+
+uint64_t ggml_graph_next_uid(void);
+
+#ifndef MIN
+#    define MIN(a, b) ((a) < (b) ? (a) : (b))
+#endif
+
+#ifndef MAX
+#    define MAX(a, b) ((a) > (b) ? (a) : (b))
+#endif
+
+// required for mmap as gguf only guarantees 32-byte alignment
+#define TENSOR_ALIGNMENT 32
+
+// static_assert should be a #define, but if it's not,
+// fall back to the _Static_assert C11 keyword.
+// if C99 - static_assert is noop
+// ref: https://stackoverflow.com/a/53923785/4039976
+#ifndef __cplusplus
+    #ifndef static_assert
+        #if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201100L)
+            #define static_assert(cond, msg) _Static_assert(cond, msg)
+        #else
+            #define static_assert(cond, msg) struct global_scope_noop_trick
+        #endif
+    #endif
+#endif
+
+static inline int ggml_up32(int n) {
+    return (n + 31) & ~31;
+}
+
+//static inline int ggml_up64(int n) {
+//    return (n + 63) & ~63;
+//}
+
+static inline int ggml_up(int n, int m) {
+    // assert m is a power of 2
+    GGML_ASSERT((m & (m - 1)) == 0);
+    return (n + m - 1) & ~(m - 1);
+}
+
+// TODO: move to ggml.h? (won't be able to inline)
+static bool ggml_are_same_layout(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    if (a->type != b->type) {
+        return false;
+    }
+    for (int i = 0; i < GGML_MAX_DIMS; i++) {
+        if (a->ne[i] != b->ne[i]) {
+            return false;
+        }
+        if (a->nb[i] != b->nb[i]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool ggml_op_is_empty(enum ggml_op op) {
+    switch (op) {
+        case GGML_OP_NONE:
+        case GGML_OP_RESHAPE:
+        case GGML_OP_TRANSPOSE:
+        case GGML_OP_VIEW:
+        case GGML_OP_PERMUTE:
+            return true;
+        default:
+            return false;
+    }
+}
+
+static inline bool ggml_impl_is_view(const struct ggml_tensor * t) {
+    return t->view_src != NULL;
+}
+
+static inline float ggml_compute_softplus_f32(float input) {
+    return (input > 20.0f) ? input : logf(1 + expf(input));
+}
+//
+// logging
+//
+
+GGML_ATTRIBUTE_FORMAT(2, 3)
+GGML_API void ggml_log_internal        (enum ggml_log_level level, const char * format, ...);
+GGML_API void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data);
+
+#define GGML_LOG(...)       ggml_log_internal(GGML_LOG_LEVEL_NONE , __VA_ARGS__)
+#define GGML_LOG_INFO(...)  ggml_log_internal(GGML_LOG_LEVEL_INFO , __VA_ARGS__)
+#define GGML_LOG_WARN(...)  ggml_log_internal(GGML_LOG_LEVEL_WARN , __VA_ARGS__)
+#define GGML_LOG_ERROR(...) ggml_log_internal(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)
+#define GGML_LOG_DEBUG(...) ggml_log_internal(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__)
+#define GGML_LOG_CONT(...)  ggml_log_internal(GGML_LOG_LEVEL_CONT , __VA_ARGS__)
+
+#define GGML_DEBUG 0
+
+#if (GGML_DEBUG >= 1)
+#define GGML_PRINT_DEBUG(...) GGML_LOG_DEBUG(__VA_ARGS__)
+#else
+#define GGML_PRINT_DEBUG(...)
+#endif
+
+#if (GGML_DEBUG >= 5)
+#define GGML_PRINT_DEBUG_5(...) GGML_LOG_DEBUG(__VA_ARGS__)
+#else
+#define GGML_PRINT_DEBUG_5(...)
+#endif
+
+#if (GGML_DEBUG >= 10)
+#define GGML_PRINT_DEBUG_10(...) GGML_LOG_DEBUG(__VA_ARGS__)
+#else
+#define GGML_PRINT_DEBUG_10(...)
+#endif
+
+// tensor params
+
+static void ggml_set_op_params(struct ggml_tensor * tensor, const void * params, size_t params_size) {
+    GGML_ASSERT(tensor != NULL); // silence -Warray-bounds warnings
+    assert(params_size <= GGML_MAX_OP_PARAMS);
+    memcpy(tensor->op_params, params, params_size);
+}
+
+static int32_t ggml_get_op_params_i32(const struct ggml_tensor * tensor, uint32_t i) {
+    assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t));
+    return ((const int32_t *)(tensor->op_params))[i];
+}
+
+static float ggml_get_op_params_f32(const struct ggml_tensor * tensor, uint32_t i) {
+    assert(i < GGML_MAX_OP_PARAMS / sizeof(float));
+    return ((const float *)(tensor->op_params))[i];
+}
+
+// [TAG_GGML_PREC]
+// - GGML_OP_MUL_MAT
+//   0 - acc
+//   1 - hint
+//   2 - src0 precision
+//   3 - src1 precision
+//
+// - GGML_OP_MUL_MAT_ID
+//   0 - acc
+//   1 - hint
+//   2 - src0 precision
+//   3 - src1 precision
+static void ggml_set_op_params_i32(struct ggml_tensor * tensor, uint32_t i, int32_t value) {
+    assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t));
+    ((int32_t *)(tensor->op_params))[i] = value;
+}
+
+static void ggml_set_op_params_f32(struct ggml_tensor * tensor, uint32_t i, float value) {
+    assert(i < GGML_MAX_OP_PARAMS / sizeof(float));
+    ((float *)(tensor->op_params))[i] = value;
+}
+
+struct ggml_map_custom1_op_params {
+    ggml_custom1_op_t  fun;
+    int                n_tasks;
+    void             * userdata;
+};
+
+struct ggml_map_custom2_op_params {
+    ggml_custom2_op_t   fun;
+    int                 n_tasks;
+    void              * userdata;
+};
+
+struct ggml_map_custom3_op_params {
+    ggml_custom3_op_t fun;
+    int               n_tasks;
+    void            * userdata;
+};
+
+struct ggml_custom_op_params {
+    ggml_custom_op_t fun;
+    int              n_tasks;
+    void           * userdata;
+};
+
+// bitset
+
+typedef uint32_t ggml_bitset_t;
+
+static_assert(sizeof(ggml_bitset_t) == 4, "bitset_t constants must be updated");
+#define BITSET_SHR 5 // log2(sizeof(ggml_bitset_t)*8)
+#define BITSET_MASK (sizeof(ggml_bitset_t)*8 - 1)
+
+static size_t ggml_bitset_size(size_t n) {
+    return (n + BITSET_MASK) >> BITSET_SHR;
+}
+
+static inline bool ggml_bitset_get(const ggml_bitset_t * bitset, size_t i) {
+    return !!(bitset[i >> BITSET_SHR] & (1u << (i & BITSET_MASK)));
+}
+
+static inline void ggml_bitset_set(ggml_bitset_t * bitset, size_t i) {
+    bitset[i >> BITSET_SHR] |= (1u << (i & BITSET_MASK));
+}
+
+static inline void ggml_bitset_clear(ggml_bitset_t * bitset, size_t i) {
+    bitset[i >> BITSET_SHR] &= ~(1u << (i & BITSET_MASK));
+}
+
+// hash set
+
+#define GGML_HASHSET_FULL ((size_t)-1)
+#define GGML_HASHSET_ALREADY_EXISTS ((size_t)-2)
+
+struct ggml_hash_set {
+    size_t size;
+    ggml_bitset_t * used;       // whether or not the keys are in use i.e. set
+    struct ggml_tensor ** keys; // actual tensors in the set, keys[i] is only defined if ggml_bitset_get(used, i)
+};
+
+struct ggml_hash_set ggml_hash_set_new(size_t size);
+void                 ggml_hash_set_free(struct ggml_hash_set * hash_set);
+
+// returns the minimum size for a hash set that can hold min_sz elements
+size_t ggml_hash_size(size_t min_sz);
+
+// remove all elements from the hash set
+void ggml_hash_set_reset(struct ggml_hash_set * hash_set);
+
+// returns true if key is in the hash set
+static bool ggml_hash_contains(const struct ggml_hash_set * hash_set, struct ggml_tensor * key);
+
+// returns GGML_HASHSET_FULL if table is full, otherwise the current index of the key or where it should be inserted
+static size_t ggml_hash_find(const struct ggml_hash_set * hash_set, const struct ggml_tensor * key);
+
+// returns GGML_HASHSET_ALREADY_EXISTS if key already exists, index otherwise, asserts if table is full
+static size_t ggml_hash_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key);
+
+// return index, asserts if table is full
+static size_t ggml_hash_find_or_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key);
+
+// hash function for ggml_tensor
+static inline size_t ggml_hash(const struct ggml_tensor * p) {
+    // the last 4 bits are always zero due to alignment
+    return (size_t)(uintptr_t)p >> 4;
+}
+
+static size_t ggml_hash_find(const struct ggml_hash_set * hash_set, const struct ggml_tensor * key) {
+    size_t h = ggml_hash(key) % hash_set->size;
+
+    // linear probing
+    size_t i = h;
+    while (ggml_bitset_get(hash_set->used, i) && hash_set->keys[i] != key) {
+        i = (i + 1) % hash_set->size;
+        if (i == h) {
+            // visited all hash table entries -> not found
+            return GGML_HASHSET_FULL;
+        }
+    }
+    return i;
+}
+
+static bool ggml_hash_contains(const struct ggml_hash_set * hash_set, struct ggml_tensor * key) {
+    size_t i = ggml_hash_find(hash_set, key);
+    return i != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, i);
+}
+
+static size_t ggml_hash_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key) {
+    size_t h = ggml_hash(key) % hash_set->size;
+
+    // linear probing
+    size_t i = h;
+    do {
+        if (!ggml_bitset_get(hash_set->used, i)) {
+            ggml_bitset_set(hash_set->used, i);
+            hash_set->keys[i] = key;
+            return i;
+        }
+        if (hash_set->keys[i] == key) {
+            return GGML_HASHSET_ALREADY_EXISTS;
+        }
+        i = (i + 1) % hash_set->size;
+    } while (i != h);
+
+    // visited all hash table entries -> not found
+    GGML_ABORT("fatal error");
+}
+
+static size_t ggml_hash_find_or_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key) {
+    size_t h = ggml_hash(key) % hash_set->size;
+
+    // linear probing
+    size_t i = h;
+    do {
+        if (!ggml_bitset_get(hash_set->used, i)) {
+            ggml_bitset_set(hash_set->used, i);
+            hash_set->keys[i] = key;
+            return i;
+        }
+        if (hash_set->keys[i] == key) {
+            return i;
+        }
+        i = (i + 1) % hash_set->size;
+    } while (i != h);
+
+    // visited all hash table entries -> not found
+    GGML_ABORT("fatal error");
+}
+
+// computation graph
+
+enum ggml_cgraph_eval_order {
+    GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT = 0,
+    GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT,
+    GGML_CGRAPH_EVAL_ORDER_COUNT
+};
+
+struct ggml_cgraph {
+    int size;    // maximum number of nodes/leafs/grads/grad_accs
+    int n_nodes; // number of nodes currently in use
+    int n_leafs; // number of leafs currently in use
+
+    struct ggml_tensor ** nodes;     // tensors with data that can change if the graph is evaluated
+    struct ggml_tensor ** grads;     // the outputs of these tensors are the gradients of the nodes
+    struct ggml_tensor ** grad_accs; // accumulators for node gradients
+    struct ggml_tensor ** leafs;     // tensors with constant data
+    int32_t             * use_counts;// number of uses of each tensor, indexed by hash table slot
+
+    struct ggml_hash_set visited_hash_set;
+
+    enum ggml_cgraph_eval_order order;
+
+    // an optional identifier that can be utilized to recognize same graphs if two non-zero values match
+    // a value of 0 means it is not set and should be ignored
+    uint64_t uid;
+};
+
+// returns a slice of cgraph with nodes [i0, i1)
+// the slice does not have leafs or gradients
+// if you need the gradients, get them from the original graph
+struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1);
+
+// ggml-alloc.c: true if the operation can reuse memory from its sources
+GGML_API bool ggml_op_can_inplace(enum ggml_op op);
+
+
+// Memory allocation
+
+GGML_API void * ggml_aligned_malloc(size_t size);
+GGML_API void ggml_aligned_free(void * ptr, size_t size);
+
+// FP16 <-> FP32
+// ref: https://github.com/Maratyszcza/FP16
+
+static inline float fp32_from_bits(uint32_t w) {
+    union {
+        uint32_t as_bits;
+        float as_value;
+    } fp32;
+    fp32.as_bits = w;
+    return fp32.as_value;
+}
+
+static inline uint32_t fp32_to_bits(float f) {
+    union {
+        float as_value;
+        uint32_t as_bits;
+    } fp32;
+    fp32.as_value = f;
+    return fp32.as_bits;
+}
+
+static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
+    const uint32_t w = (uint32_t) h << 16;
+    const uint32_t sign = w & UINT32_C(0x80000000);
+    const uint32_t two_w = w + w;
+
+    const uint32_t exp_offset = UINT32_C(0xE0) << 23;
+#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) || defined(__GNUC__) && !defined(__STRICT_ANSI__)) && (!defined(__cplusplus) || __cplusplus >= 201703L)
+    const float exp_scale = 0x1.0p-112f;
+#else
+    const float exp_scale = fp32_from_bits(UINT32_C(0x7800000));
+#endif
+    const float normalized_value = fp32_from_bits((two_w >> 4) + exp_offset) * exp_scale;
+
+    const uint32_t magic_mask = UINT32_C(126) << 23;
+    const float magic_bias = 0.5f;
+    const float denormalized_value = fp32_from_bits((two_w >> 17) | magic_mask) - magic_bias;
+
+    const uint32_t denormalized_cutoff = UINT32_C(1) << 27;
+    const uint32_t result = sign |
+        (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value) : fp32_to_bits(normalized_value));
+    return fp32_from_bits(result);
+}
+
+static inline ggml_fp16_t ggml_compute_fp32_to_fp16(float f) {
+#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) || defined(__GNUC__) && !defined(__STRICT_ANSI__)) && (!defined(__cplusplus) || __cplusplus >= 201703L)
+    const float scale_to_inf = 0x1.0p+112f;
+    const float scale_to_zero = 0x1.0p-110f;
+#else
+    const float scale_to_inf = fp32_from_bits(UINT32_C(0x77800000));
+    const float scale_to_zero = fp32_from_bits(UINT32_C(0x08800000));
+#endif
+    float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
+
+    const uint32_t w = fp32_to_bits(f);
+    const uint32_t shl1_w = w + w;
+    const uint32_t sign = w & UINT32_C(0x80000000);
+    uint32_t bias = shl1_w & UINT32_C(0xFF000000);
+    if (bias < UINT32_C(0x71000000)) {
+        bias = UINT32_C(0x71000000);
+    }
+
+    base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
+    const uint32_t bits = fp32_to_bits(base);
+    const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
+    const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
+    const uint32_t nonsign = exp_bits + mantissa_bits;
+    return (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign);
+}
+
+#define GGML_COMPUTE_FP16_TO_FP32(x) ggml_compute_fp16_to_fp32(x)
+#define GGML_COMPUTE_FP32_TO_FP16(x) ggml_compute_fp32_to_fp16(x)
+
+#define GGML_FP16_TO_FP32(x) GGML_COMPUTE_FP16_TO_FP32(x)
+#define GGML_FP32_TO_FP16(x) GGML_COMPUTE_FP32_TO_FP16(x)
+
+static inline float ggml_e8m0_to_fp32(uint8_t x) {
+    uint32_t bits;  // Stores the raw bit representation of the float
+
+    // Handle special case for minimum exponent (denormalized float)
+    if (x == 0) {
+        // Bit pattern for 2^(-127):
+        // - Sign bit: 0 (positive)
+        // - Exponent: 0 (denormalized number)
+        // - Mantissa: 0x400000 (0.5 in fractional form)
+        // Value = 0.5 * 2^(-126) = 2^(-127)
+        bits = 0x00400000;
+    }
+    // note: disabled as we don't need to handle NaNs
+    //// Handle special case for NaN (all bits set)
+    //else if (x == 0xFF) {
+    //    // Standard quiet NaN pattern:
+    //    // - Sign bit: 0
+    //    // - Exponent: all 1s (0xFF)
+    //    // - Mantissa: 0x400000 (quiet NaN flag)
+    //    bits = 0x7FC00000;
+    //}
+    // Normalized values (most common case)
+    else {
+        // Construct normalized float by shifting exponent into position:
+        // - Exponent field: 8 bits (positions 30-23)
+        // - Mantissa: 0 (implicit leading 1)
+        // Value = 2^(x - 127)
+        bits = (uint32_t) x << 23;
+    }
+
+    float result;  // Final float value
+                   // Safely reinterpret bit pattern as float without type-punning issues
+    memcpy(&result, &bits, sizeof(float));
+    return result;
+}
+
+// Equal to ggml_e8m0_to_fp32/2
+// Useful with MXFP4 quantization since the E0M2 values are doubled
+static inline float ggml_e8m0_to_fp32_half(uint8_t x) {
+    uint32_t bits;
+
+    // For x < 2: use precomputed denormal patterns
+    if (x < 2) {
+        // 0x00200000 = 2^(-128), 0x00400000 = 2^(-127)
+        bits = 0x00200000 << x;
+    }
+    // For x >= 2: normalized exponent adjustment
+    else {
+        // 0.5 * 2^(x-127) = 2^(x-128) = normalized with exponent (x-1)
+        bits = (uint32_t)(x - 1) << 23;
+    }
+    // Note: NaNs are not handled here
+
+    float result;
+    memcpy(&result, &bits, sizeof(float));
+    return result;
+}
+
+#define GGML_E8M0_TO_FP32(x) ggml_e8m0_to_fp32(x)
+#define GGML_E8M0_TO_FP32_HALF(x) ggml_e8m0_to_fp32_half(x)
+
+// UE4M3: unsigned, 4 exp bits (bias=7), 3 mantissa bits
+// Returns value * 0.5 to match kvalues_mxfp4 convention (kvalues = 2 * E2M1_float)
+static inline float ggml_ue4m3_to_fp32(uint8_t x) {
+    if (x == 0 || x == 0x7F) {
+        return 0.0f;
+    }
+    int   exp = (x >> 3) & 0xF;
+    int   man = x & 0x7;
+    float raw;
+    if (exp == 0) {
+        raw = ldexpf((float) man, -9);
+    } else {
+        raw = ldexpf(1.0f + (float) man / 8.0f, exp - 7);
+    }
+    return raw * 0.5f;
+}
+
+static inline uint8_t ggml_fp32_to_ue4m3(float x) {
+    if (!(x > 0.0f)) {
+        return 0;
+    }
+    if (x > 448.0f) {
+        x = 448.0f;
+    }
+    uint32_t bits;
+    memcpy(&bits, &x, 4);
+    int fp32_exp  = ((bits >> 23) & 0xFF) - 127;
+    int fp32_man  = (bits >> 20) & 0x7;
+    int ue4m3_exp = fp32_exp + 7;
+    if (ue4m3_exp <= 0) {
+        // subnormal: value = man * 2^-9, man = round(x * 2^9)
+        int man = (int) (x * 512.0f + 0.5f);
+        if (man > 7) {
+            man = 7;
+        }
+        if (man < 1) {
+            return 0;
+        }
+        return (uint8_t) man;
+    }
+    if (ue4m3_exp >= 15) {
+        return 0x7E;
+    }
+    int round_bit = (bits >> 19) & 1;
+    int ue4m3_man = fp32_man + round_bit;
+    if (ue4m3_man > 7) {
+        ue4m3_man = 0;
+        ue4m3_exp++;
+        if (ue4m3_exp >= 15) {
+            return 0x7E;
+        }
+    }
+    return (uint8_t) ((ue4m3_exp << 3) | ue4m3_man);
+}
+
+/**
+ * Converts brain16 to float32.
+ *
+ * The bfloat16 floating point format has the following structure:
+ *
+ *       ┌sign
+ *       │
+ *       │   ┌exponent
+ *       │   │
+ *       │   │      ┌mantissa
+ *       │   │      │
+ *       │┌──┴───┐┌─┴───┐
+ *     0b0000000000000000 brain16
+ *
+ * Since bf16 has the same number of exponent bits as a 32bit float,
+ * encoding and decoding numbers becomes relatively straightforward.
+ *
+ *       ┌sign
+ *       │
+ *       │   ┌exponent
+ *       │   │
+ *       │   │      ┌mantissa
+ *       │   │      │
+ *       │┌──┴───┐┌─┴───────────────────┐
+ *     0b00000000000000000000000000000000 IEEE binary32
+ *
+ * For comparison, the standard fp16 format has fewer exponent bits.
+ *
+ *       ┌sign
+ *       │
+ *       │  ┌exponent
+ *       │  │
+ *       │  │    ┌mantissa
+ *       │  │    │
+ *       │┌─┴─┐┌─┴──────┐
+ *     0b0000000000000000 IEEE binary16
+ *
+ * @see IEEE 754-2008
+ */
+static inline float ggml_compute_bf16_to_fp32(ggml_bf16_t h) {
+    union {
+        float f;
+        uint32_t i;
+    } u;
+    u.i = (uint32_t)h.bits << 16;
+    return u.f;
+}
+
+/**
+ * Converts float32 to brain16.
+ *
+ * This is binary identical with Google Brain float conversion.
+ * Floats shall round to nearest even, and NANs shall be quiet.
+ * Subnormals aren't flushed to zero, except perhaps when used.
+ * This code should vectorize nicely if using modern compilers.
+ */
+static inline ggml_bf16_t ggml_compute_fp32_to_bf16(float s) {
+    ggml_bf16_t h;
+    union {
+        float f;
+        uint32_t i;
+    } u;
+    u.f = s;
+    if ((u.i & 0x7fffffff) > 0x7f800000) { /* nan */
+        h.bits = (u.i >> 16) | 64; /* force to quiet */
+        return h;
+    }
+    h.bits = (u.i + (0x7fff + ((u.i >> 16) & 1))) >> 16;
+    return h;
+}
+
+#define GGML_FP32_TO_BF16(x) ggml_compute_fp32_to_bf16(x)
+#define GGML_BF16_TO_FP32(x) ggml_compute_bf16_to_fp32(x)
+
+static inline int32_t ggml_node_get_use_count(const struct ggml_cgraph * cgraph, int node_idx) {
+    const struct ggml_tensor * node = cgraph->nodes[node_idx];
+
+    size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node);
+    if (!ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) {
+        return 0;
+    }
+    return cgraph->use_counts[hash_pos];
+}
+
+// return true if the node's results are only used by N other nodes
+// and can be fused into their calculations.
+static inline bool ggml_node_has_n_uses(const struct ggml_cgraph * cgraph, int node_idx, int32_t n_uses) {
+    const struct ggml_tensor * node = cgraph->nodes[node_idx];
+
+    // check the use count against how many we're replacing
+    if (ggml_node_get_use_count(cgraph, node_idx) != n_uses) {
+        return false;
+    }
+
+    // if node is a view, some other node might be using the intermediate result
+    // via the view source.
+    if (node->view_src) {
+        return false;
+    }
+
+    // If the user requested output for the node, can't fuse
+    if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
+        return false;
+    }
+
+    return true;
+}
+
+// Returns true if nodes with indices { node_idxs } are the sequence of ggml_ops in ops[]
+// and are fusable. Nodes are considered fusable according to this function if:
+// - all nodes except the last have only one use and are not views/outputs (see ggml_node_has_N_uses).
+// - all nodes except the last are a src of the following node.
+// - all nodes are the same shape.
+// TODO: Consider allowing GGML_OP_NONE nodes in between
+static inline bool ggml_can_fuse_ext(const struct ggml_cgraph * cgraph, const int * node_idxs, const enum ggml_op * ops, int num_ops) {
+    for (int i = 0; i < num_ops; ++i) {
+        if (node_idxs[i] >= cgraph->n_nodes) {
+            return false;
+        }
+
+        struct ggml_tensor * node = cgraph->nodes[node_idxs[i]];
+        if (node->op != ops[i]) {
+            return false;
+        }
+        if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+            return false;
+        }
+        if (i < num_ops - 1 && !ggml_node_has_n_uses(cgraph, node_idxs[i], 1)) {
+            return false;
+        }
+        if (i > 0) {
+            struct ggml_tensor * prev = cgraph->nodes[node_idxs[i - 1]];
+            if (node->src[0] != prev && node->src[1] != prev) {
+                return false;
+            }
+            if (!ggml_are_same_shape(node, prev)) {
+                return false;
+            }
+        }
+    }
+    return true;
+}
+
+// same as above, for sequential indices starting at node_idx
+static inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, const enum ggml_op * ops, int num_ops) {
+    assert(num_ops < 32);
+
+    if (node_idx + num_ops > cgraph->n_nodes) {
+        return false;
+    }
+
+    int idxs[32];
+    for (int i = 0; i < num_ops; ++i) {
+        idxs[i] = node_idx + i;
+    }
+
+    return ggml_can_fuse_ext(cgraph, idxs, ops, num_ops);
+}
+
+GGML_API bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
+                                         const int *                node_idxs,
+                                         int                        count,
+                                         const enum ggml_op *       ops,
+                                         const int *                outputs,
+                                         int                        num_outputs);
+
+// Returns true if the subgraph formed by {node_idxs} can be fused
+// checks whethers all nodes which are not part of outputs can be elided
+// by checking if their num_uses are confined to the subgraph
+static inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph,
+                                          int                        node_idx,
+                                          int                        count,
+                                          const enum ggml_op *       ops,
+                                          const int *                outputs,
+                                          int                        num_outputs) {
+    GGML_ASSERT(count < 32);
+    if (node_idx + count > cgraph->n_nodes) {
+        return false;
+    }
+
+    int idxs[32];
+
+    for (int i = 0; i < count; ++i) {
+        idxs[i] = node_idx + i;
+    }
+
+    return ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, outputs, num_outputs);
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#ifdef __cplusplus
+#include <array>
+#include <initializer_list>
+#include <vector>
+
+// nicer C++ syntax for ggml_can_fuse
+inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops) {
+    return ggml_can_fuse(cgraph, node_idx, ops.begin(), (int)ops.size());
+}
+
+inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph *          cgraph,
+                                   int                                 start_idx,
+                                   std::initializer_list<enum ggml_op> ops,
+                                   std::initializer_list<int>          outputs = {}) {
+    return ggml_can_fuse_subgraph(cgraph, start_idx, ops.size(), ops.begin(), outputs.begin(), outputs.size());
+}
+
+// Return true if the edges in the graph match expectations.
+inline bool ggml_check_edges(const struct ggml_cgraph *                cgraph,
+                             int                                       start_idx,
+                             std::initializer_list<std::array<int, 3>> edges) {
+    for (const auto & edge : edges) {
+        int dst_node = edge[0];
+        int src_idx  = edge[1];
+        int src_node = edge[2];
+        if (cgraph->nodes[start_idx + dst_node]->src[src_idx] != cgraph->nodes[start_idx + src_node]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+// expose GGUF internals for test code
+GGML_API size_t gguf_type_size(enum gguf_type type);
+GGML_API void gguf_write_to_buf(const struct gguf_context * ctx, std::vector<int8_t> & buf, bool only_meta);
+#endif // __cplusplus
diff --git a/python/freetoken/kernel/csrc/gguf/dequantize.cuh b/python/freetoken/kernel/csrc/gguf/dequantize.cuh
deleted file mode 100644
index 57cb1a27cf..0000000000
--- a/python/freetoken/kernel/csrc/gguf/dequantize.cuh
+++ /dev/null
@@ -1,583 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/dequantize.cuh
-// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/convert.cu
-// Dequant functions
-static __device__ __forceinline__ void dequantize_q4_0(const void* vx, const int ib, const int iqs, dfloat2& v) {
-  const block_q4_0* x = (const block_q4_0*)vx;
-
-  const dfloat d = x[ib].d;
-
-  const int vui = x[ib].qs[iqs];
-
-  v.x = __int2half_rn(vui & 0xF);
-  v.y = __int2half_rn(vui >> 4);
-
-  v = __hsub2(v, __floats2half2_rn(8.0f, 8.0f));
-  v = __hmul2(v, {d, d});
-}
-
-static __device__ __forceinline__ void dequantize_q4_1(const void* vx, const int ib, const int iqs, dfloat2& v) {
-  const block_q4_1* x = (const block_q4_1*)vx;
-
-  const dfloat d = __low2half(x[ib].dm);
-  const dfloat m = __high2half(x[ib].dm);
-
-  const int vui = x[ib].qs[iqs];
-
-  v.x = __int2half_rn(vui & 0xF);
-  v.y = __int2half_rn(vui >> 4);
-
-  v = __hmul2(v, {d, d});
-  v = __hadd2(v, {m, m});
-}
-
-static __device__ __forceinline__ void dequantize_q5_0(const void* vx, const int ib, const int iqs, dfloat2& v) {
-  const block_q5_0* x = (const block_q5_0*)vx;
-
-  const dfloat d = x[ib].d;
-
-  uint32_t qh;
-  memcpy(&qh, x[ib].qh, sizeof(qh));
-
-  const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10;
-  const int xh_1 = ((qh >> (iqs + 12))) & 0x10;
-
-  v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0);
-  v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1);
-
-  v = __hsub2(v, __floats2half2_rn(16.0f, 16.0f));
-  v = __hmul2(v, {d, d});
-}
-
-static __device__ __forceinline__ void dequantize_q5_1(const void* vx, const int ib, const int iqs, dfloat2& v) {
-  const block_q5_1* x = (const block_q5_1*)vx;
-
-  const dfloat d = __low2half(x[ib].dm);
-  const dfloat m = __high2half(x[ib].dm);
-
-  uint32_t qh;
-  memcpy(&qh, x[ib].qh, sizeof(qh));
-
-  const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10;
-  const int xh_1 = ((qh >> (iqs + 12))) & 0x10;
-
-  v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0);
-  v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1);
-
-  v = __hmul2(v, {d, d});
-  v = __hadd2(v, {m, m});
-}
-
-static __device__ __forceinline__ void dequantize_q8_0(const void* vx, const int ib, const int iqs, dfloat2& v) {
-  const block_q8_0* x = (const block_q8_0*)vx;
-
-  const dfloat d = x[ib].d;
-
-  v.x = __int2half_rn(x[ib].qs[iqs + 0]);
-  v.y = __int2half_rn(x[ib].qs[iqs + 1]);
-
-  v = __hmul2(v, {d, d});
-}
-
-template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
-static __global__ void dequantize_block(const void* __restrict__ vx, dst_t* __restrict__ y, const int k) {
-  const int i = 2 * (blockDim.x * blockIdx.x + threadIdx.x);
-
-  if (i >= k) {
-    return;
-  }
-
-  const int ib = i / qk;          // block index
-  const int iqs = (i % qk) / qr;  // quant index
-  const int iybs = i - i % qk;    // y block start index
-  const int y_offset = qr == 1 ? 1 : qk / 2;
-
-  // dequantize
-  dfloat2 v;
-  dequantize_kernel(vx, ib, iqs, v);
-
-  y[iybs + iqs + 0] = convert_from_half<dst_t>(v.x);
-  y[iybs + iqs + y_offset] = convert_from_half<dst_t>(v.y);
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_q2_K(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_q2_K* x = (const block_q2_K*)vx;
-
-  const auto tid = threadIdx.x;
-  const int n = tid / 32;
-  const int l = tid - 32 * n;
-  const int is = 8 * n + l / 16;
-
-  const uint8_t q = x[i].qs[32 * n + l];
-  dst_t* y = yy + i * QK_K + 128 * n;
-
-  half dall = __low2half(x[i].dm);
-  half dmin = __high2half(x[i].dm);
-  y[l + 0] = convert_from_half<dst_t>(__hsub(
-      __hmul(dall, __int2half_rn((x[i].scales[is + 0] & 0xF) * ((q >> 0) & 3))),
-      __hmul(dmin, __int2half_rn(x[i].scales[is + 0] >> 4))));
-  y[l + 32] = convert_from_half<dst_t>(__hsub(
-      __hmul(dall, __int2half_rn((x[i].scales[is + 2] & 0xF) * ((q >> 2) & 3))),
-      __hmul(dmin, __int2half_rn(x[i].scales[is + 2] >> 4))));
-  y[l + 64] = convert_from_half<dst_t>(__hsub(
-      __hmul(dall, __int2half_rn((x[i].scales[is + 4] & 0xF) * ((q >> 4) & 3))),
-      __hmul(dmin, __int2half_rn(x[i].scales[is + 4] >> 4))));
-  y[l + 96] = convert_from_half<dst_t>(__hsub(
-      __hmul(dall, __int2half_rn((x[i].scales[is + 6] & 0xF) * ((q >> 6) & 3))),
-      __hmul(dmin, __int2half_rn(x[i].scales[is + 6] >> 4))));
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_q3_K(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_q3_K* x = (const block_q3_K*)vx;
-
-  const auto r = threadIdx.x / 4;
-  const int tid = r / 2;
-  const int is0 = r % 2;
-  const int l0 = 16 * is0 + 4 * (threadIdx.x % 4);
-  const int n = tid / 4;
-  const int j = tid - 4 * n;
-
-  uint8_t m = 1 << (4 * n + j);
-  int is = 8 * n + 2 * j + is0;
-  int shift = 2 * j;
-
-  int8_t us = is < 4    ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 8] >> 0) & 3) << 4)
-              : is < 8  ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 4] >> 2) & 3) << 4)
-              : is < 12 ? (x[i].scales[is - 8] >> 4) | (((x[i].scales[is + 0] >> 4) & 3) << 4)
-                        : (x[i].scales[is - 8] >> 4) | (((x[i].scales[is - 4] >> 6) & 3) << 4);
-  half d_all = x[i].d;
-  half dl = __hmul(d_all, __int2half_rn(us - 32));
-
-  dst_t* y = yy + i * QK_K + 128 * n + 32 * j;
-  const uint8_t* q = x[i].qs + 32 * n;
-  const uint8_t* hm = x[i].hmask;
-
-  for (int l = l0; l < l0 + 4; ++l) {
-    y[l] = convert_from_half<dst_t>(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4))));
-  }
-}
-
-static inline __device__ void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) {
-  if (j < 4) {
-    d = q[j] & 63;
-    m = q[j + 4] & 63;
-  } else {
-    d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
-    m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_q4_K(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const block_q4_K* x = (const block_q4_K*)vx;
-
-  const auto i = blockIdx.x;
-
-  // assume 32 threads
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;
-  const int ir = tid % 8;
-  const int is = 2 * il;
-  const int n = 4;
-
-  dst_t* y = yy + i * QK_K + 64 * il + n * ir;
-
-  const half dall = __low2half(x[i].dm);
-  const half dmin = __high2half(x[i].dm);
-
-  const uint8_t* q = x[i].qs + 32 * il + n * ir;
-
-  uint8_t sc, m;
-  get_scale_min_k4(is + 0, x[i].scales, sc, m);
-  const half d1 = __hmul(dall, __int2half_rn(sc));
-  const half m1 = __hmul(dmin, __int2half_rn(m));
-  get_scale_min_k4(is + 1, x[i].scales, sc, m);
-  const half d2 = __hmul(dall, __int2half_rn(sc));
-  const half m2 = __hmul(dmin, __int2half_rn(m));
-  for (int l = 0; l < n; ++l) {
-    y[l + 0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1));
-    y[l + 32] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2));
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_q5_K(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const block_q5_K* x = (const block_q5_K*)vx;
-
-  const auto i = blockIdx.x;
-
-  // assume 64 threads - this is very slightly better than the one below
-  const auto tid = threadIdx.x;
-  const int il = tid / 16;  // il is in 0...3
-  const int ir = tid % 16;  // ir is in 0...15
-  const int is = 2 * il;    // is is in 0...6
-
-  dst_t* y = yy + i * QK_K + 64 * il + 2 * ir;
-
-  const half dall = __low2half(x[i].dm);
-  const half dmin = __high2half(x[i].dm);
-
-  const uint8_t* ql = x[i].qs + 32 * il + 2 * ir;
-  const uint8_t* qh = x[i].qh + 2 * ir;
-
-  uint8_t sc, m;
-  get_scale_min_k4(is + 0, x[i].scales, sc, m);
-  const half d1 = __hmul(dall, __int2half_rn(sc));
-  const half m1 = __hmul(dmin, __int2half_rn(m));
-  get_scale_min_k4(is + 1, x[i].scales, sc, m);
-  const half d2 = __hmul(dall, __int2half_rn(sc));
-  const half m2 = __hmul(dmin, __int2half_rn(m));
-
-  uint8_t hm = 1 << (2 * il);
-  y[0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1));
-  y[1] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1));
-  hm <<= 1;
-  y[32] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2));
-  y[33] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2));
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_q6_K(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const block_q6_K* x = (const block_q6_K*)vx;
-
-  const auto i = blockIdx.x;
-
-  // assume 64 threads - this is very slightly better than the one below
-  const auto tid = threadIdx.x;
-  const int ip = tid / 32;       // ip is 0 or 1
-  const int il = tid - 32 * ip;  // 0...32
-  const int is = 8 * ip + il / 16;
-
-  dst_t* y = yy + i * QK_K + 128 * ip + il;
-
-  const half d = x[i].d;
-
-  const uint8_t* ql = x[i].ql + 64 * ip + il;
-  const uint8_t qh = x[i].qh[32 * ip + il];
-  const int8_t* sc = x[i].scales + is;
-
-  y[0] = convert_from_half<dst_t>(
-      __hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))));
-  y[32] = convert_from_half<dst_t>(
-      __hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))));
-  y[64] = convert_from_half<dst_t>(
-      __hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))));
-  y[96] = convert_from_half<dst_t>(
-      __hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))));
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq2_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq2_xxs* x = (const block_iq2_xxs*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint16_t* q2 = x[i].qs + 4 * ib;
-  const uint8_t* aux8 = (const uint8_t*)q2;
-  const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[il]);
-  const uint32_t aux32 = q2[2] | (q2[3] << 16);
-  const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f;
-  const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127];
-  for (int j = 0; j < 8; ++j)
-    y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq2_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq2_xs* x = (const block_iq2_xs*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint16_t* q2 = x[i].qs + 4 * ib;
-  const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[il] & 511));
-  const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f;
-  const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
-  for (int j = 0; j < 8; ++j)
-    y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq2_s(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq2_s* x = (const block_iq2_s*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint8_t* grid = (const uint8_t*)(iq2s_grid + (x[i].qs[4 * ib + il] | ((x[i].qh[ib] << (8 - 2 * il)) & 0x300)));
-  const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f;
-  const uint8_t signs = x[i].qs[QK_K / 8 + 4 * ib + il];
-  for (int j = 0; j < 8; ++j)
-    y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq3_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq3_xxs* x = (const block_iq3_xxs*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint8_t* q3 = x[i].qs + 8 * ib;
-  const uint16_t* gas = (const uint16_t*)(x[i].qs + QK_K / 4) + 2 * ib;
-  const uint8_t* grid1 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 0]);
-  const uint8_t* grid2 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 1]);
-  const uint32_t aux32 = gas[0] | (gas[1] << 16);
-  const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f;
-  const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127];
-  for (int j = 0; j < 4; ++j) {
-    y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f);
-    y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f);
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq3_s(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq3_s* x = (const block_iq3_s*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint8_t* qs = x[i].qs + 8 * ib;
-  const uint8_t* grid1 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 0] | ((x[i].qh[ib] << (8 - 2 * il)) & 256)));
-  const uint8_t* grid2 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 1] | ((x[i].qh[ib] << (7 - 2 * il)) & 256)));
-  const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib / 2] >> 4 * (ib % 2)) & 0xf)) * 0.5f;
-  const uint8_t signs = x[i].signs[4 * ib + il];
-  for (int j = 0; j < 4; ++j) {
-    y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f);
-    y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f);
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq1_s(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const int64_t i = blockIdx.x;
-  const block_iq1_s* x = (const block_iq1_s*)vx;
-
-  const int64_t tid = threadIdx.x;
-  const int64_t il = tid / 8;  // 0...3
-  const int64_t ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
-  const float d = __half2float(x[i].d) * (2 * ((x[i].qh[ib] >> 12) & 7) + 1);
-  uint32_t grid32[2];
-  const int8_t* q = (const int8_t*)grid32;
-  grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[ib] >> 3 * il) & 7) << 8)];
-  grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-  grid32[0] &= 0x0f0f0f0f;
-  for (int j = 0; j < 8; ++j) {
-    y[j] = d * (q[j] + delta);
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq1_m(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const int64_t i = blockIdx.x;
-  const block_iq1_m* x = (const block_iq1_m*)vx;
-
-  const int64_t tid = threadIdx.x;
-  const int64_t il = tid / 8;  // 0...3
-  const int64_t ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 8 * il;
-  const uint16_t* sc = (const uint16_t*)x[i].scales;
-  iq1m_scale_t scale;
-  scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
-  const int64_t ib16 = 2 * ib + il / 2;  // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
-  const float d = __half2float(scale.f16) * (2 * ((sc[ib16 / 4] >> 3 * (ib16 % 4)) & 0x7) + 1);
-  const float delta = x[i].qh[2 * ib + il / 2] & (0x08 << 4 * (il % 2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
-  uint32_t grid32[2];
-  const int8_t* q = (const int8_t*)grid32;
-  grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[2 * ib + il / 2] >> 4 * (il % 2)) & 7) << 8)];
-  grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-  grid32[0] &= 0x0f0f0f0f;
-  for (int j = 0; j < 8; ++j) {
-    y[j] = d * (q[j] + delta);
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq4_nl(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq4_nl* x = (const block_iq4_nl*)vx + i * (QK_K / QK4_NL);
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 4 * il;
-  const uint8_t* q4 = x[ib].qs + 4 * il;
-  const float d = __half2float(x[ib].d);
-  for (int j = 0; j < 4; ++j) {
-    y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-    y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4];
-  }
-}
-
-template <typename dst_t>
-static __global__ void dequantize_block_iq4_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) {
-  const auto i = blockIdx.x;
-  const block_iq4_xs* x = (const block_iq4_xs*)vx;
-
-  const auto tid = threadIdx.x;
-  const int il = tid / 8;  // 0...3
-  const int ib = tid % 8;  // 0...7
-  dst_t* y = yy + i * QK_K + 32 * ib + 4 * il;
-  const uint8_t* q4 = x[i].qs + 16 * ib + 4 * il;
-  const float d = __half2float(x[i].d) *
-                  ((((x[i].scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((x[i].scales_h >> 2 * ib) & 3) << 4)) - 32);
-  for (int j = 0; j < 4; ++j) {
-    y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-    y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4];
-  }
-}
-
-template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
-static void
-dequantize_block_cuda(const void* __restrict__ vx, dst_t* __restrict__ y, const int k, cudaStream_t stream) {
-  const int num_blocks = (k + 2 * CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2 * CUDA_DEQUANTIZE_BLOCK_SIZE);
-  dequantize_block<qk, qr, dequantize_kernel><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(vx, y, k);
-}
-
-template <typename dst_t>
-static void dequantize_row_q2_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_q3_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_q4_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_q4_K<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_q5_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_q5_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_q6_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_q6_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq2_xxs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq2_xxs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq2_xs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq2_xs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq2_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq2_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq3_xxs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq3_xxs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq3_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq3_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq1_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq1_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq1_m_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = k / QK_K;
-  dequantize_block_iq1_m<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq4_nl_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = (k + QK_K - 1) / QK_K;
-  dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static void dequantize_row_iq4_xs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) {
-  const int nb = (k + QK_K - 1) / QK_K;
-  dequantize_block_iq4_xs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template <typename dst_t>
-static to_cuda_ggml_t<dst_t> ggml_get_to_cuda(int64_t type) {
-  switch (type) {
-    case 2:
-      return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
-    case 3:
-      return dequantize_block_cuda<QK4_1, QR4_1, dequantize_q4_1>;
-    case 6:
-      return dequantize_block_cuda<QK5_0, QR5_0, dequantize_q5_0>;
-    case 7:
-      return dequantize_block_cuda<QK5_1, QR5_1, dequantize_q5_1>;
-    case 8:
-      return dequantize_block_cuda<QK8_0, QR8_0, dequantize_q8_0>;
-    case 10:
-      return dequantize_row_q2_K_cuda;
-    case 11:
-      return dequantize_row_q3_K_cuda;
-    case 12:
-      return dequantize_row_q4_K_cuda;
-    case 13:
-      return dequantize_row_q5_K_cuda;
-    case 14:
-      return dequantize_row_q6_K_cuda;
-    case 16:
-      return dequantize_row_iq2_xxs_cuda;
-    case 17:
-      return dequantize_row_iq2_xs_cuda;
-    case 18:
-      return dequantize_row_iq3_xxs_cuda;
-    case 19:
-      return dequantize_row_iq1_s_cuda;
-    case 20:
-      return dequantize_row_iq4_nl_cuda;
-    case 21:
-      return dequantize_row_iq3_s_cuda;
-    case 22:
-      return dequantize_row_iq2_s_cuda;
-    case 23:
-      return dequantize_row_iq4_xs_cuda;
-    case 29:
-      return dequantize_row_iq1_m_cuda;
-    default:
-      return nullptr;
-  }
-}
diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h
deleted file mode 100644
index f42a216332..0000000000
--- a/python/freetoken/kernel/csrc/gguf/dispatch.h
+++ /dev/null
@@ -1,28 +0,0 @@
-// Minimal AT_DISPATCH helper for the vendored GGUF kernels (borrowed from
-// sgl-kernel csrc/quantization/gguf, which are ports of llama.cpp). The donor
-// pulls these macros from its large include/utils.h; we only need the float
-// dispatch, so vendor just that to keep the JIT compile self-contained.
-#pragma once
-
-#include <ATen/Dispatch.h>
-
-#ifndef WARP_SIZE
-#define WARP_SIZE 32
-#endif
-
-// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants).
-#ifndef SGLANG_SHFL_XOR_SYNC
-#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask))
-#endif
-#ifndef SGLANG_SHFL_XOR_SYNC_WIDTH
-#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
-  __shfl_xor_sync((mask), (var), (lane_mask), (width))
-#endif
-
-#define DISPATCH_CASE_FLOAT_TYPES(...)                 \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
-
-#define DISPATCH_FLOAT_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOAT_TYPES(__VA_ARGS__))
diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common.h b/python/freetoken/kernel/csrc/gguf/ggml-common.h
deleted file mode 100644
index 88c21a4ab3..0000000000
--- a/python/freetoken/kernel/csrc/gguf/ggml-common.h
+++ /dev/null
@@ -1,1029 +0,0 @@
-// adapted from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/ggml-common.h
-// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-common.h
-#define QK_K 256
-#define K_QUANTS_PER_ITERATION 2
-#define WARP_SIZE_GGUF 32
-#define K_SCALE_SIZE 12
-#define CUDA_DEQUANTIZE_BLOCK_SIZE 256
-#define CUDA_QUANTIZE_BLOCK_SIZE 256
-#define GGML_CUDA_DMMV_X 32
-#define GGML_CUDA_MMV_Y 1
-
-// Data Structures
-// QK = number of values after dequantization
-// QR = QK / number of values before dequantization
-// QI = number of 32 bit integers before dequantization
-
-#define QK4_0 32
-#define QR4_0 2
-#define QI4_0 (QK4_0 / (4 * QR4_0))
-typedef struct {
-  half d;                 // delta
-  uint8_t qs[QK4_0 / 2];  // nibbles / quants
-} block_q4_0;
-
-#define QK4_1 32
-#define QR4_1 2
-#define QI4_1 (QK4_1 / (4 * QR4_1))
-typedef struct {
-  half2 dm;               // dm.x = delta, dm.y = min
-  uint8_t qs[QK4_1 / 2];  // nibbles / quants
-} block_q4_1;
-
-#define QK5_0 32
-#define QR5_0 2
-#define QI5_0 (QK5_0 / (4 * QR5_0))
-typedef struct {
-  half d;                 // delta
-  uint8_t qh[4];          // 5-th bit of quants
-  uint8_t qs[QK5_0 / 2];  // nibbles / quants
-} block_q5_0;
-
-#define QK5_1 32
-#define QR5_1 2
-#define QI5_1 (QK5_1 / (4 * QR5_1))
-typedef struct {
-  half2 dm;               // dm.x = delta, dm.y = min
-  uint8_t qh[4];          // 5-th bit of quants
-  uint8_t qs[QK5_1 / 2];  // nibbles / quants
-} block_q5_1;
-
-#define QK8_0 32
-#define QR8_0 1
-#define QI8_0 (QK8_0 / (4 * QR8_0))
-typedef struct {
-  half d;            // delta
-  int8_t qs[QK8_0];  // quants
-} block_q8_0;
-
-#define QK8_1 32
-#define QR8_1 1
-#define QI8_1 (QK8_1 / (4 * QR8_1))
-typedef struct {
-  half2 ds;          // ds.x = delta, ds.y = sum
-  int8_t qs[QK8_0];  // quants
-} block_q8_1;
-
-#define QR2_K 4
-#define QI2_K (QK_K / (4 * QR2_K))
-typedef struct {
-  uint8_t scales[QK_K / 16];  // scales and mins, quantized with 4 bits
-  uint8_t qs[QK_K / 4];       // quants
-  half2 dm;                   // super-block scale for quantized scales/mins
-} block_q2_K;
-
-#define QR3_K 4
-#define QI3_K (QK_K / (4 * QR3_K))
-typedef struct {
-  uint8_t hmask[QK_K / 8];       // quants - high bit
-  uint8_t qs[QK_K / 4];          // quants - low 2 bits
-  uint8_t scales[K_SCALE_SIZE];  // scales, quantized with 6 bits
-  half d;                        // super-block scale
-} block_q3_K;
-
-#define QR4_K 2
-#define QI4_K (QK_K / (4 * QR4_K))
-typedef struct {
-  half2 dm;                       // super-block scale for quantized scales/mins
-  uint8_t scales[3 * QK_K / 64];  // scales, quantized with 6 bits
-  uint8_t qs[QK_K / 2];           // 4--bit quants
-} block_q4_K;
-
-#define QR5_K 2
-#define QI5_K (QK_K / (4 * QR5_K))
-typedef struct {
-  half2 dm;                      // super-block scale for quantized scales/mins
-  uint8_t scales[K_SCALE_SIZE];  // scales and mins, quantized with 6 bits
-  uint8_t qh[QK_K / 8];          // quants, high bit
-  uint8_t qs[QK_K / 2];          // quants, low 4 bits
-} block_q5_K;
-
-#define QR6_K 2
-#define QI6_K (QK_K / (4 * QR6_K))
-typedef struct {
-  uint8_t ql[QK_K / 2];      // quants, lower 4 bits
-  uint8_t qh[QK_K / 4];      // quants, upper 2 bits
-  int8_t scales[QK_K / 16];  // scales
-  half d;                    // delta
-} block_q6_K;
-
-#define QR2_XXS 8
-#define QI2_XXS (QK_K / (4 * QR2_XXS))
-typedef struct {
-  half d;
-  uint16_t qs[QK_K / 8];
-} block_iq2_xxs;
-
-#define QR2_XS 8
-#define QI2_XS (QK_K / (4 * QR2_XS))
-typedef struct {
-  half d;
-  uint16_t qs[QK_K / 8];
-  uint8_t scales[QK_K / 32];
-} block_iq2_xs;
-
-#define QR2_S 8
-#define QI2_S (QK_K / (4 * QR2_S))
-typedef struct {
-  half d;
-  uint8_t qs[QK_K / 4];
-  uint8_t qh[QK_K / 32];
-  uint8_t scales[QK_K / 32];
-} block_iq2_s;
-
-#define QR3_XXS 8
-#define QI3_XXS (QK_K / (4 * QR3_XXS))
-typedef struct {
-  half d;
-  uint8_t qs[3 * (QK_K / 8)];
-} block_iq3_xxs;
-
-#define QR3_XS 8
-#define QI3_XS (QK_K / (4 * QR3_XS))
-#define IQ3S_N_SCALE QK_K / 64
-typedef struct {
-  half d;
-  uint8_t qs[QK_K / 4];
-  uint8_t qh[QK_K / 32];
-  uint8_t signs[QK_K / 8];
-  uint8_t scales[IQ3S_N_SCALE];
-} block_iq3_s;
-
-// 1.5625 bpw
-#define QR1_S 8
-#define QI1_S (QK_K / (4 * QR1_S))
-typedef struct {
-  half d;
-  uint8_t qs[QK_K / 8];
-  uint16_t qh[QK_K / 32];
-} block_iq1_s;
-
-// 1.75 bpw
-#define QR1_M 8
-#define QI1_M (QK_K / (4 * QR1_M))
-typedef struct {
-  uint8_t qs[QK_K / 8];       // grid index, low 8 bits
-  uint8_t qh[QK_K / 16];      // grid index, high 3 bits + grid shift bit (for two groups of 8)
-  uint8_t scales[QK_K / 32];  // 3-bit block scales (4-bit if QK_K == 64)
-} block_iq1_m;
-
-// Used by IQ1_M quants
-typedef union {
-  half f16;
-  uint16_t u16;
-} iq1m_scale_t;
-
-#define QK4_NL 32
-#define QR4_NL 2
-#define QI4_NL (QK4_NL / (4 * QR4_NL))
-typedef struct {
-  half d;
-  uint8_t qs[QK4_NL / 2];
-} block_iq4_nl;
-
-#define QR4_XS 8
-#define QI4_XS (QK_K / (4 * QR4_XS))
-typedef struct {
-  half d;
-  uint16_t scales_h;
-  uint8_t scales_l[QK_K / 64];
-  uint8_t qs[QK_K / 2];
-} block_iq4_xs;
-
-static const __device__ uint64_t iq2xxs_grid[256] = {
-    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b,
-    0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08,
-    0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08,
-    0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
-    0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, 0x0808081908191919,
-    0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, 0x0808082b08080808, 0x0808082b0808082b,
-    0x0808082b082b082b, 0x0808082b2b08082b, 0x0808190808080819, 0x0808190808081908, 0x0808190808190808,
-    0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
-    0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x080819082b2b1908,
-    0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, 0x08081919082b0808, 0x080819191908192b,
-    0x08081919192b2b19, 0x080819192b080808, 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808,
-    0x0808192b19080808, 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
-    0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, 0x08082b0819081908,
-    0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, 0x08082b1908081908, 0x08082b1919080808,
-    0x08082b2b0808082b, 0x08082b2b08191908, 0x0819080808080819, 0x0819080808081908, 0x0819080808190808,
-    0x08190808082b0819, 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
-    0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, 0x0819081919190808,
-    0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, 0x0819082b19081919, 0x0819190808080808,
-    0x0819190808082b08, 0x08191908082b0808, 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808,
-    0x0819191908192b08, 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
-    0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, 0x08192b1908080808,
-    0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, 0x082b080808080808, 0x082b08080808082b,
-    0x082b080808082b2b, 0x082b080819081908, 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b,
-    0x082b0819082b2b19, 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
-    0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, 0x082b191908080808,
-    0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, 0x082b2b0808082b08, 0x082b2b08082b0808,
-    0x082b2b082b191908, 0x082b2b2b19081908, 0x1908080808080819, 0x1908080808081908, 0x1908080808190808,
-    0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
-    0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808,
-    0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, 0x190808192b080808, 0x190808192b081919,
-    0x1908082b08080819, 0x1908082b08190808, 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08,
-    0x1908190808080808, 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
-    0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, 0x19082b0808081908,
-    0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b1908080808, 0x19082b1919192b08,
-    0x19082b19192b0819, 0x19082b192b08082b, 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808,
-    0x1919080808082b08, 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
-    0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, 0x1919082b2b190819,
-    0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, 0x1919192b08080819, 0x1919192b19191908,
-    0x19192b0808080808, 0x19192b0808190819, 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808,
-    0x19192b2b08082b08, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
-    0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, 0x192b190808080808,
-    0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, 0x192b19192b081908, 0x192b2b081908082b,
-    0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b,
-    0x2b08081908081908, 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
-    0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808,
-    0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, 0x2b08192b08082b19, 0x2b08192b19080808,
-    0x2b08192b192b0808, 0x2b082b080808082b, 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908,
-    0x2b19080808190808, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
-    0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, 0x2b19190819081908,
-    0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808,
-    0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808,
-    0x2b2b2b1908081908,
-};
-
-static const __device__ uint64_t iq2xs_grid[512] = {
-    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b,
-    0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808,
-    0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908,
-    0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919,
-    0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b,
-    0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b192b19,
-    0x080808082b2b0808, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19,
-    0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b,
-    0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919,
-    0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, 0x08080819192b0808, 0x08080819192b2b08,
-    0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b,
-    0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808,
-    0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808,
-    0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19,
-    0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819,
-    0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08,
-    0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, 0x080819082b080819,
-    0x080819082b081908, 0x080819082b190808, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919,
-    0x0808191908082b08, 0x0808191908190819, 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819,
-    0x0808191919081908, 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819,
-    0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, 0x0808192b1908082b,
-    0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08,
-    0x08082b0808082b2b, 0x08082b0808190819, 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919,
-    0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808,
-    0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808,
-    0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, 0x08082b2b08080808, 0x08082b2b082b0808,
-    0x08082b2b082b2b08, 0x08082b2b2b19192b, 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908,
-    0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919,
-    0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, 0x081908081908082b,
-    0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x08190808192b0808,
-    0x08190808192b2b2b, 0x081908082b080819, 0x081908082b081908, 0x081908082b190808, 0x0819081908080808,
-    0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908,
-    0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, 0x081908192b080808,
-    0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b0808192b,
-    0x0819082b08190808, 0x0819082b19080808, 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b,
-    0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808,
-    0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, 0x08191908192b1908,
-    0x081919082b080808, 0x0819191908080819, 0x0819191908081908, 0x0819191908190808, 0x0819191919080808,
-    0x0819192b08080808, 0x0819192b08191908, 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908,
-    0x08192b0808190808, 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b,
-    0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, 0x08192b2b2b2b2b19,
-    0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, 0x082b080808082b2b,
-    0x082b080808190819, 0x082b080808191908, 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908,
-    0x082b080819190808, 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908,
-    0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, 0x082b082b08080808,
-    0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908,
-    0x082b190808190808, 0x082b1908082b2b19, 0x082b190819080808, 0x082b191908080808, 0x082b191919080819,
-    0x082b19191919082b, 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b,
-    0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, 0x082b2b0819191919,
-    0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, 0x082b2b192b190808, 0x082b2b2b08082b08,
-    0x082b2b2b082b0808, 0x082b2b2b2b08082b, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819,
-    0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b,
-    0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808,
-    0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819,
-    0x1908080819191908, 0x19080808192b0808, 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908,
-    0x190808082b190808, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08,
-    0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, 0x1908081919081908,
-    0x1908081919190808, 0x190808192b080808, 0x190808192b081919, 0x190808192b2b082b, 0x1908082b08080819,
-    0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808,
-    0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819,
-    0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, 0x1908190819081908,
-    0x1908190819190808, 0x190819082b080808, 0x190819082b191908, 0x1908191908080819, 0x1908191908081908,
-    0x1908191908190808, 0x19081919082b1908, 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808,
-    0x1908192b08082b2b, 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908,
-    0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, 0x19082b08192b082b,
-    0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b19192b2b19,
-    0x19082b2b08081908, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08,
-    0x1919080808190819, 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819,
-    0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, 0x1919081908081908,
-    0x1919081908190808, 0x1919081908191919, 0x1919081919080808, 0x191908191908082b, 0x1919082b08080808,
-    0x1919082b19081908, 0x1919082b2b2b2b2b, 0x1919190808080819, 0x1919190808081908, 0x1919190808190808,
-    0x19191908082b0819, 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819,
-    0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, 0x1919192b082b0819,
-    0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, 0x19192b0808191908, 0x19192b0819080819,
-    0x19192b0819190808, 0x19192b082b192b19, 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b,
-    0x19192b2b2b081919, 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808,
-    0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, 0x192b081908080808,
-    0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, 0x192b190808080808, 0x192b19080819192b,
-    0x192b191908190808, 0x192b191919080808, 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819,
-    0x192b2b08192b2b2b, 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b,
-    0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819,
-    0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, 0x2b08080819080819, 0x2b08080819081908,
-    0x2b08080819190808, 0x2b0808082b080808, 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b,
-    0x2b08081908080819, 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808,
-    0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, 0x2b08082b2b080808,
-    0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, 0x2b08190808080819, 0x2b08190808081908,
-    0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808,
-    0x2b0819082b082b19, 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08,
-    0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, 0x2b082b0819192b2b,
-    0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, 0x2b082b190808192b, 0x2b082b2b082b082b,
-    0x2b082b2b2b080808, 0x2b082b2b2b082b08, 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819,
-    0x2b19080808081908, 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908,
-    0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, 0x2b19082b2b082b19,
-    0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, 0x2b19190819190808, 0x2b19190819192b08,
-    0x2b191919082b2b19, 0x2b1919192b190808, 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819,
-    0x2b192b082b2b192b, 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808,
-    0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, 0x2b2b0808082b2b2b,
-    0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, 0x2b2b08192b2b192b, 0x2b2b082b08080808,
-    0x2b2b082b0808082b, 0x2b2b082b08082b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808,
-    0x2b2b190819080808, 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b,
-    0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, 0x2b2b2b082b2b2b08,
-    0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b,
-    0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b,
-};
-
-static const __device__ uint64_t iq2s_grid[1024] = {
-    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b,
-    0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808,
-    0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908,
-    0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919,
-    0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, 0x08080808192b2b19,
-    0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819,
-    0x080808082b191908, 0x080808082b2b0808, 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819,
-    0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b,
-    0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808,
-    0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908,
-    0x080808191919192b, 0x0808081919192b19, 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08,
-    0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919,
-    0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919,
-    0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b082b2b2b,
-    0x0808082b19080819, 0x0808082b19081908, 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808,
-    0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908,
-    0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19,
-    0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819,
-    0x08081908082b1908, 0x08081908082b192b, 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b,
-    0x0808190819081919, 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908,
-    0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, 0x08081908192b1919,
-    0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, 0x080819082b082b19, 0x080819082b190808,
-    0x080819082b191919, 0x080819082b192b08, 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808,
-    0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819,
-    0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, 0x08081919082b1919,
-    0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, 0x080819191908192b, 0x0808191919082b19,
-    0x0808191919190808, 0x080819191919082b, 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819,
-    0x08081919192b1908, 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08,
-    0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, 0x0808192b08081908,
-    0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b08191919, 0x0808192b19080808,
-    0x0808192b19081919, 0x0808192b19082b08, 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808,
-    0x0808192b2b080819, 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b,
-    0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, 0x08082b080819192b,
-    0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b08082b2b2b, 0x08082b0819080819,
-    0x08082b0819081908, 0x08082b081908192b, 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b,
-    0x08082b0819191919, 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808,
-    0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908,
-    0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, 0x08082b1908192b08, 0x08082b19082b0819,
-    0x08082b1919080808, 0x08082b1919081919, 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908,
-    0x08082b19192b0808, 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819,
-    0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, 0x08082b2b19190808,
-    0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, 0x0819080808082b19,
-    0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819,
-    0x08190808082b1908, 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919,
-    0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, 0x0819080819192b19,
-    0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, 0x08190808192b2b08, 0x081908082b080819,
-    0x081908082b081908, 0x081908082b08192b, 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08,
-    0x081908082b2b0819, 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919,
-    0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, 0x081908190819192b,
-    0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, 0x08190819082b1919, 0x08190819082b2b08,
-    0x0819081919080819, 0x0819081919081908, 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808,
-    0x081908191919082b, 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908,
-    0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, 0x081908192b190819,
-    0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b08082b19, 0x0819082b08190808,
-    0x0819082b08191919, 0x0819082b082b0819, 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919,
-    0x0819082b19190819, 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808,
-    0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819,
-    0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, 0x08191908082b0808, 0x08191908082b1919,
-    0x08191908082b2b08, 0x0819190819080819, 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19,
-    0x0819190819190808, 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819,
-    0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, 0x081919082b082b08,
-    0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, 0x0819191908080819, 0x0819191908081908,
-    0x081919190808192b, 0x0819191908082b19, 0x0819191908190808, 0x081919190819082b, 0x0819191908191919,
-    0x0819191908192b08, 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b,
-    0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, 0x08191919192b0808,
-    0x081919192b080819, 0x081919192b081908, 0x081919192b190808, 0x0819192b08080808, 0x0819192b08081919,
-    0x0819192b08082b08, 0x0819192b08190819, 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819,
-    0x0819192b19081908, 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819,
-    0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, 0x08192b0808191919,
-    0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, 0x08192b081908082b, 0x08192b0819081919,
-    0x08192b0819082b08, 0x08192b0819190819, 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819,
-    0x08192b082b081908, 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08,
-    0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, 0x08192b1919081908,
-    0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, 0x08192b2b08081908, 0x08192b2b08190808,
-    0x08192b2b19080808, 0x08192b2b1919192b, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919,
-    0x082b080808082b08, 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19,
-    0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, 0x082b080819081908,
-    0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, 0x082b0808192b1908, 0x082b08082b080808,
-    0x082b08082b082b2b, 0x082b08082b191908, 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908,
-    0x082b081908190808, 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808,
-    0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, 0x082b0819192b0808,
-    0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, 0x082b082b08080808, 0x082b082b08082b2b,
-    0x082b082b082b082b, 0x082b082b082b2b08, 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808,
-    0x082b082b2b082b08, 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908,
-    0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, 0x082b190808192b08,
-    0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, 0x082b19081908082b, 0x082b190819081919,
-    0x082b190819082b08, 0x082b190819190819, 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819,
-    0x082b19082b081908, 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08,
-    0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, 0x082b191919081908,
-    0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, 0x082b192b08080819, 0x082b192b08081908,
-    0x082b192b08190808, 0x082b192b19080808, 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919,
-    0x082b2b0808190819, 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808,
-    0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, 0x082b2b1908190808,
-    0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, 0x082b2b2b192b1908, 0x082b2b2b2b082b08,
-    0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19,
-    0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b,
-    0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, 0x190808081908082b,
-    0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908,
-    0x190808081919192b, 0x1908080819192b19, 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919,
-    0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08,
-    0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919,
-    0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x190808190819192b, 0x1908081908192b19,
-    0x19080819082b0808, 0x19080819082b082b, 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908,
-    0x190808191908192b, 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919,
-    0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, 0x190808192b08082b,
-    0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, 0x190808192b191908, 0x190808192b2b0808,
-    0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919,
-    0x1908082b08192b08, 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08,
-    0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, 0x1908082b2b081908,
-    0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808082b2b,
-    0x1908190808190819, 0x1908190808191908, 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808,
-    0x19081908082b082b, 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908,
-    0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, 0x1908190819191919,
-    0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, 0x190819082b080808, 0x190819082b08082b,
-    0x190819082b081919, 0x190819082b082b08, 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808,
-    0x1908191908080819, 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808,
-    0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, 0x19081919082b1908,
-    0x1908191919080808, 0x190819191908082b, 0x1908191919081919, 0x1908191919082b08, 0x1908191919190819,
-    0x1908191919191908, 0x19081919192b0808, 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908,
-    0x190819192b190808, 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08,
-    0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, 0x1908192b19081908,
-    0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, 0x19082b0808080819, 0x19082b0808081908,
-    0x19082b0808082b19, 0x19082b0808190808, 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08,
-    0x19082b08082b0819, 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919,
-    0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, 0x19082b082b081908,
-    0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, 0x19082b1908081919, 0x19082b1908082b08,
-    0x19082b1908190819, 0x19082b1908191908, 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908,
-    0x19082b1919190808, 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908,
-    0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919,
-    0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, 0x191908080819192b, 0x1919080808192b19,
-    0x19190808082b0808, 0x19190808082b082b, 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819,
-    0x1919080819081908, 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b,
-    0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, 0x191908082b080808,
-    0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, 0x191908082b190819, 0x191908082b191908,
-    0x1919081908080819, 0x1919081908081908, 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808,
-    0x191908190819082b, 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908,
-    0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, 0x1919081919190819,
-    0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, 0x191908192b081908, 0x191908192b190808,
-    0x1919082b08080808, 0x1919082b08081919, 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908,
-    0x1919082b082b0808, 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19,
-    0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, 0x1919190808082b19,
-    0x1919190808190808, 0x191919080819082b, 0x1919190808191919, 0x1919190808192b08, 0x19191908082b0819,
-    0x19191908082b1908, 0x1919190819080808, 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08,
-    0x1919190819190819, 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908,
-    0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, 0x1919191908082b08,
-    0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, 0x1919191919080819, 0x1919191919081908,
-    0x1919191919190808, 0x191919192b080808, 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808,
-    0x1919192b082b192b, 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919,
-    0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, 0x19192b0819080819,
-    0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, 0x19192b082b080808, 0x19192b1908080819,
-    0x19192b1908081908, 0x19192b1908190808, 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19,
-    0x19192b2b2b081919, 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b,
-    0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, 0x192b0808082b0819,
-    0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, 0x192b080819082b08, 0x192b080819190819,
-    0x192b080819191908, 0x192b0808192b0808, 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808,
-    0x192b08190808082b, 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908,
-    0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, 0x192b08192b080808,
-    0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, 0x192b082b19080808, 0x192b082b1919192b,
-    0x192b082b2b2b0819, 0x192b190808080808, 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819,
-    0x192b190808191908, 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808,
-    0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, 0x192b191919080808,
-    0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, 0x192b192b08080808, 0x192b192b2b191908,
-    0x192b2b0808080819, 0x192b2b0808081908, 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08,
-    0x192b2b1908080808, 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808,
-    0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908,
-    0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, 0x2b08080819080819, 0x2b08080819081908,
-    0x2b08080819190808, 0x2b0808081919082b, 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819,
-    0x2b0808082b080808, 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819,
-    0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, 0x2b08081908191919,
-    0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, 0x2b08081919080808, 0x2b0808191908082b,
-    0x2b08081919081919, 0x2b08081919082b08, 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819,
-    0x2b0808192b081908, 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919,
-    0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, 0x2b08082b19081908,
-    0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, 0x2b0819080808192b, 0x2b08190808082b19,
-    0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819,
-    0x2b08190819080808, 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819,
-    0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, 0x2b0819082b190808,
-    0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, 0x2b08191908082b08, 0x2b08191908190819,
-    0x2b08191908191908, 0x2b081919082b0808, 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808,
-    0x2b0819192b080808, 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808,
-    0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, 0x2b082b0808190819,
-    0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, 0x2b082b0819190808, 0x2b082b082b2b082b,
-    0x2b082b1908080819, 0x2b082b1908081908, 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b,
-    0x2b082b2b19192b08, 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819,
-    0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, 0x2b19080808191919,
-    0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908081908082b, 0x2b19080819081919,
-    0x2b19080819082b08, 0x2b19080819190819, 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819,
-    0x2b1908082b081908, 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819,
-    0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, 0x2b19081919192b2b,
-    0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, 0x2b19082b19080808, 0x2b19082b2b2b192b,
-    0x2b19190808080808, 0x2b1919080808082b, 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819,
-    0x2b19190808191908, 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808,
-    0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, 0x2b19191908190808,
-    0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, 0x2b19192b08080808, 0x2b19192b1908192b,
-    0x2b19192b192b1908, 0x2b192b0808080819, 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b,
-    0x2b192b0819080808, 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b,
-    0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, 0x2b2b080808191908,
-    0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, 0x2b2b080819081908, 0x2b2b080819190808,
-    0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b,
-    0x2b2b082b08082b2b, 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b,
-    0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, 0x2b2b190808081908,
-    0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, 0x2b2b19082b2b1908, 0x2b2b191908080808,
-    0x2b2b191908192b19, 0x2b2b192b19190819, 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b,
-    0x2b2b2b1919191908, 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808,
-    0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b,
-};
-
-static const __device__ uint32_t iq3xxs_grid[256] = {
-    0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, 0x04041c0c,
-    0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, 0x040c140c, 0x040c142c,
-    0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, 0x04140414, 0x04140424, 0x04140c0c,
-    0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c,
-    0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e,
-    0x04243e1c, 0x04243e2c, 0x042c040c, 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04,
-    0x043e0c24, 0x043e0c34, 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c,
-    0x0c04141c, 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c,
-    0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, 0x0c143e14,
-    0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, 0x0c24042c, 0x0c242c04,
-    0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, 0x0c3e2404, 0x14040404, 0x14040414,
-    0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c,
-    0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404,
-    0x14140414, 0x14140c0c, 0x14140c3e, 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c,
-    0x141c0c04, 0x141c0c24, 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c,
-    0x142c3e24, 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c,
-    0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, 0x1c0c2424,
-    0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, 0x1c1c0c0c, 0x1c1c1c1c,
-    0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, 0x1c2c2c2c, 0x1c340c24, 0x1c341c34,
-    0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e,
-    0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424,
-    0x24242c0c, 0x24243424, 0x242c142c, 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04,
-    0x2c040c14, 0x2c04240c, 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14,
-    0x2c143e14, 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c,
-    0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, 0x340c340c,
-    0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, 0x34341c1c, 0x343e041c,
-    0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, 0x3e042c14, 0x3e0c1434, 0x3e0c2404,
-    0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c,
-    0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04,
-};
-
-static const __device__ uint32_t iq3xs_grid[512] = {
-    0x04040404, 0x0404040c, 0x04040414, 0x0404042c, 0x0404043e, 0x04040c04, 0x04040c0c, 0x04040c14, 0x04040c24,
-    0x04040c34, 0x04041404, 0x0404140c, 0x0404142c, 0x04041c1c, 0x04042404, 0x04042414, 0x0404242c, 0x0404243e,
-    0x04042c0c, 0x04042c1c, 0x04043404, 0x04043414, 0x04043e0c, 0x04043e24, 0x04043e3e, 0x040c0404, 0x040c040c,
-    0x040c0414, 0x040c0424, 0x040c0c04, 0x040c0c0c, 0x040c0c2c, 0x040c1404, 0x040c141c, 0x040c143e, 0x040c1c0c,
-    0x040c1c2c, 0x040c2424, 0x040c340c, 0x040c342c, 0x040c3e14, 0x04140404, 0x0414040c, 0x0414042c, 0x0414043e,
-    0x04140c04, 0x04140c1c, 0x04140c34, 0x0414140c, 0x0414142c, 0x04141c04, 0x04141c24, 0x04142414, 0x0414242c,
-    0x0414243e, 0x04142c0c, 0x04142c1c, 0x04143e04, 0x04143e1c, 0x041c041c, 0x041c0c0c, 0x041c0c2c, 0x041c1404,
-    0x041c1414, 0x041c1c0c, 0x041c1c1c, 0x041c1c34, 0x041c2424, 0x041c2c04, 0x041c2c14, 0x041c343e, 0x041c3e0c,
-    0x041c3e2c, 0x04240404, 0x04240c1c, 0x04240c3e, 0x0424140c, 0x04241424, 0x04241c14, 0x04242404, 0x0424241c,
-    0x04242c0c, 0x04243e04, 0x042c0414, 0x042c0424, 0x042c1404, 0x042c1414, 0x042c1434, 0x042c1c1c, 0x042c240c,
-    0x042c242c, 0x042c243e, 0x042c3434, 0x042c3e1c, 0x04340434, 0x04340c0c, 0x04340c1c, 0x04341c0c, 0x04342c14,
-    0x04343e0c, 0x043e0404, 0x043e0414, 0x043e0424, 0x043e1404, 0x043e1414, 0x043e1434, 0x043e1c1c, 0x043e2c04,
-    0x043e2c24, 0x0c040404, 0x0c04040c, 0x0c040414, 0x0c040424, 0x0c040c04, 0x0c040c0c, 0x0c040c1c, 0x0c040c2c,
-    0x0c040c3e, 0x0c041404, 0x0c041414, 0x0c041c0c, 0x0c041c24, 0x0c041c34, 0x0c042c24, 0x0c042c34, 0x0c04340c,
-    0x0c043e14, 0x0c0c0404, 0x0c0c040c, 0x0c0c041c, 0x0c0c0434, 0x0c0c0c04, 0x0c0c0c24, 0x0c0c140c, 0x0c0c1c04,
-    0x0c0c1c1c, 0x0c0c240c, 0x0c0c2c04, 0x0c0c2c14, 0x0c0c3e04, 0x0c0c3e34, 0x0c140404, 0x0c140c14, 0x0c140c2c,
-    0x0c140c3e, 0x0c141404, 0x0c141424, 0x0c141c14, 0x0c142404, 0x0c14241c, 0x0c142c2c, 0x0c143404, 0x0c143e14,
-    0x0c1c040c, 0x0c1c0424, 0x0c1c043e, 0x0c1c0c04, 0x0c1c0c1c, 0x0c1c140c, 0x0c1c143e, 0x0c1c1c04, 0x0c1c1c24,
-    0x0c1c240c, 0x0c1c3414, 0x0c1c3e04, 0x0c24041c, 0x0c24042c, 0x0c240c14, 0x0c240c24, 0x0c241c0c, 0x0c241c1c,
-    0x0c242414, 0x0c242434, 0x0c242c04, 0x0c242c24, 0x0c2c040c, 0x0c2c0c04, 0x0c2c0c1c, 0x0c2c140c, 0x0c2c1c04,
-    0x0c2c1c14, 0x0c2c2c0c, 0x0c341404, 0x0c341424, 0x0c34143e, 0x0c342424, 0x0c342434, 0x0c3e040c, 0x0c3e041c,
-    0x0c3e0c04, 0x0c3e0c14, 0x0c3e140c, 0x0c3e1c2c, 0x0c3e240c, 0x0c3e3414, 0x0c3e3e04, 0x14040404, 0x1404040c,
-    0x1404041c, 0x1404042c, 0x1404043e, 0x14040c04, 0x14040c14, 0x14040c24, 0x14040c34, 0x1404140c, 0x1404141c,
-    0x1404143e, 0x14041c04, 0x14041c14, 0x1404240c, 0x1404241c, 0x1404242c, 0x14042c04, 0x14042c14, 0x1404343e,
-    0x14043e04, 0x14043e1c, 0x14043e2c, 0x140c0404, 0x140c0414, 0x140c0c04, 0x140c0c1c, 0x140c0c3e, 0x140c1414,
-    0x140c142c, 0x140c1c0c, 0x140c1c24, 0x140c2414, 0x140c2c0c, 0x1414040c, 0x14140424, 0x1414043e, 0x1414140c,
-    0x1414141c, 0x14141c04, 0x14141c3e, 0x1414240c, 0x14142c1c, 0x14142c3e, 0x14143e0c, 0x14143e24, 0x141c0404,
-    0x141c0414, 0x141c042c, 0x141c0c0c, 0x141c1414, 0x141c1424, 0x141c1c0c, 0x141c1c1c, 0x141c2414, 0x141c2c04,
-    0x141c3434, 0x1424040c, 0x1424043e, 0x14241404, 0x1424141c, 0x14241c14, 0x14241c2c, 0x1424240c, 0x14243e14,
-    0x14243e2c, 0x142c0424, 0x142c0c0c, 0x142c1414, 0x142c1c3e, 0x142c2404, 0x142c2c1c, 0x142c3e04, 0x14340404,
-    0x14340414, 0x1434043e, 0x1434140c, 0x14342c2c, 0x1434340c, 0x143e042c, 0x143e0c0c, 0x143e1434, 0x143e1c04,
-    0x143e241c, 0x143e2c04, 0x1c040414, 0x1c040c0c, 0x1c040c1c, 0x1c040c2c, 0x1c040c3e, 0x1c041414, 0x1c041c0c,
-    0x1c041c1c, 0x1c041c2c, 0x1c042414, 0x1c042424, 0x1c04243e, 0x1c042c0c, 0x1c04341c, 0x1c043e0c, 0x1c0c040c,
-    0x1c0c041c, 0x1c0c042c, 0x1c0c0c24, 0x1c0c140c, 0x1c0c141c, 0x1c0c2404, 0x1c0c3404, 0x1c0c3e14, 0x1c0c3e34,
-    0x1c140404, 0x1c140c14, 0x1c141404, 0x1c141c14, 0x1c141c24, 0x1c142c04, 0x1c1c040c, 0x1c1c0c04, 0x1c1c0c24,
-    0x1c1c140c, 0x1c1c141c, 0x1c1c143e, 0x1c1c1c04, 0x1c1c240c, 0x1c1c241c, 0x1c1c243e, 0x1c1c2c2c, 0x1c1c3e1c,
-    0x1c24041c, 0x1c240c0c, 0x1c240c34, 0x1c241414, 0x1c241c0c, 0x1c242c14, 0x1c243404, 0x1c243424, 0x1c2c040c,
-    0x1c2c0c04, 0x1c2c0c14, 0x1c2c142c, 0x1c2c1c14, 0x1c2c2424, 0x1c2c2c34, 0x1c2c3e1c, 0x1c340c34, 0x1c34240c,
-    0x1c3e040c, 0x1c3e041c, 0x1c3e1404, 0x1c3e1414, 0x1c3e1c2c, 0x24040404, 0x24040424, 0x24040c14, 0x24041404,
-    0x24041424, 0x2404143e, 0x24041c14, 0x2404240c, 0x24042c04, 0x24043e04, 0x240c0414, 0x240c043e, 0x240c0c0c,
-    0x240c0c1c, 0x240c1414, 0x240c1c04, 0x240c1c2c, 0x240c241c, 0x240c2c0c, 0x240c2c2c, 0x2414040c, 0x2414041c,
-    0x24140c04, 0x24140c2c, 0x2414140c, 0x24141c1c, 0x24142404, 0x24142c3e, 0x24143414, 0x24143e04, 0x241c0424,
-    0x241c0c0c, 0x241c0c1c, 0x241c1404, 0x241c1414, 0x241c1c0c, 0x241c1c2c, 0x24240404, 0x24240414, 0x24241424,
-    0x24241c3e, 0x24242404, 0x24243e0c, 0x242c042c, 0x242c043e, 0x242c140c, 0x242c3414, 0x24340c1c, 0x24341c24,
-    0x24343404, 0x243e0c04, 0x243e0c2c, 0x243e1c04, 0x243e241c, 0x243e2c0c, 0x2c040414, 0x2c040c04, 0x2c040c24,
-    0x2c041414, 0x2c042404, 0x2c042424, 0x2c04243e, 0x2c042c14, 0x2c043434, 0x2c043e24, 0x2c0c040c, 0x2c0c041c,
-    0x2c0c042c, 0x2c0c0c14, 0x2c0c140c, 0x2c0c1c14, 0x2c0c3e14, 0x2c140404, 0x2c140c0c, 0x2c14141c, 0x2c141c04,
-    0x2c141c34, 0x2c142c1c, 0x2c1c0414, 0x2c1c043e, 0x2c1c0c04, 0x2c1c143e, 0x2c1c2424, 0x2c1c2c0c, 0x2c1c342c,
-    0x2c1c3e1c, 0x2c24040c, 0x2c240424, 0x2c241404, 0x2c241c14, 0x2c242434, 0x2c2c0c14, 0x2c2c1434, 0x2c2c2c0c,
-    0x2c2c2c1c, 0x2c342414, 0x2c3e0414, 0x2c3e0424, 0x2c3e1414, 0x34040c0c, 0x34040c1c, 0x34040c2c, 0x34041c0c,
-    0x34041c1c, 0x34043404, 0x340c0404, 0x340c1404, 0x340c143e, 0x340c3424, 0x34140c14, 0x34141c24, 0x34142414,
-    0x34142c2c, 0x34143414, 0x34143e04, 0x341c0404, 0x341c0c24, 0x341c140c, 0x341c2404, 0x3424142c, 0x3424241c,
-    0x34243414, 0x342c0404, 0x342c041c, 0x342c1c24, 0x342c3404, 0x3434042c, 0x34342404, 0x343e0c0c, 0x343e0c1c,
-    0x3e040404, 0x3e040424, 0x3e04043e, 0x3e041404, 0x3e041414, 0x3e041c34, 0x3e042404, 0x3e042c24, 0x3e043414,
-    0x3e0c0414, 0x3e0c0c0c, 0x3e0c1424, 0x3e0c241c, 0x3e0c242c, 0x3e14040c, 0x3e140424, 0x3e140c04, 0x3e140c34,
-    0x3e14140c, 0x3e141c04, 0x3e142c0c, 0x3e1c0414, 0x3e1c1c14, 0x3e1c1c2c, 0x3e1c2c1c, 0x3e24040c, 0x3e24042c,
-    0x3e240c1c, 0x3e241404, 0x3e242c04, 0x3e2c1414, 0x3e2c2414, 0x3e340414, 0x3e341c0c, 0x3e3e0404,
-};
-
-#define IQ1S_DELTA 0.125f
-#define IQ1M_DELTA 0.125f
-static const __device__ uint64_t iq1s_grid_gpu[2048] = {
-    0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, 0x00020002,
-    0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, 0x02000000, 0x02000002,
-    0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200, 0x02020202, 0x00000110, 0x00000111,
-    0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212, 0x00020111, 0x01000011, 0x01000112, 0x01000211,
-    0x01010012, 0x01010111, 0x01010212, 0x01020011, 0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011,
-    0x02010110, 0x02010112, 0x02020111, 0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020,
-    0x00020022, 0x00020220, 0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020,
-    0x02000022, 0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220,
-    0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101, 0x01011202,
-    0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110, 0x00001111, 0x00001112,
-    0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111, 0x01001212, 0x01011010, 0x01011011,
-    0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010, 0x01021012, 0x01021111, 0x01021210, 0x01021212,
-    0x02001011, 0x02011011, 0x02011111, 0x02011210, 0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112,
-    0x02021211, 0x00011120, 0x00011221, 0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220,
-    0x01021020, 0x01021021, 0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000,
-    0x00002002, 0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101,
-    0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101, 0x02022000,
-    0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211, 0x00022110, 0x00022111,
-    0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110, 0x01022211, 0x02012011, 0x02012110,
-    0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022, 0x00002220, 0x00002222, 0x00012121, 0x00022020,
-    0x00022022, 0x00022220, 0x00022222, 0x01002121, 0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020,
-    0x02002022, 0x02002121, 0x02002220, 0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222,
-    0x00110000, 0x00110001, 0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000,
-    0x01110101, 0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102,
-    0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012, 0x00110111,
-    0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010, 0x01110011, 0x01110012,
-    0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111, 0x02100110, 0x02110012, 0x02110111,
-    0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122, 0x00120121, 0x01100020, 0x01100122, 0x01100221,
-    0x01110022, 0x01110121, 0x01110220, 0x01110222, 0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120,
-    0x02110122, 0x02120121, 0x00101001, 0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201,
-    0x00121001, 0x00121102, 0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100,
-    0x01111101, 0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000,
-    0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101, 0x02121201,
-    0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112, 0x00111211, 0x00121010,
-    0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110, 0x01101111, 0x01101112, 0x01111011,
-    0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211, 0x01111212, 0x01121011, 0x01121110, 0x01121111,
-    0x01121112, 0x01121211, 0x02101010, 0x02101012, 0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010,
-    0x02111011, 0x02111110, 0x02111111, 0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111,
-    0x00101021, 0x00101120, 0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021,
-    0x00121122, 0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121,
-    0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221, 0x01121222,
-    0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001, 0x00112102, 0x00122101,
-    0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101, 0x01112200, 0x01112202, 0x01122000,
-    0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101, 0x02112001, 0x02112100, 0x02122101, 0x00112010,
-    0x00112012, 0x00112111, 0x00112212, 0x00122011, 0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210,
-    0x01112011, 0x01112110, 0x01112111, 0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212,
-    0x02102211, 0x02112011, 0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221,
-    0x00112122, 0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121,
-    0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222, 0x00200000,
-    0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101, 0x00220200, 0x00220202,
-    0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000, 0x02200002, 0x02200200, 0x02200202,
-    0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200, 0x02220202, 0x00200111, 0x00210011, 0x00210110,
-    0x00210211, 0x00220111, 0x01200012, 0x01200110, 0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011,
-    0x01220110, 0x01220111, 0x01220112, 0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021,
-    0x00200220, 0x00200222, 0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121,
-    0x01210021, 0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121,
-    0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201, 0x00221101,
-    0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200, 0x01211202, 0x01221102,
-    0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101, 0x00201211, 0x00211111, 0x00221011,
-    0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011, 0x01211110, 0x01211111, 0x01211211, 0x01221012,
-    0x01221111, 0x01221210, 0x02201211, 0x02211010, 0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011,
-    0x02221110, 0x02221112, 0x02221211, 0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021,
-    0x01201221, 0x01211121, 0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222,
-    0x00202000, 0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202,
-    0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202, 0x02222000,
-    0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211, 0x00222111, 0x01202112,
-    0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112, 0x01222211, 0x02202111, 0x02212010,
-    0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020, 0x00202022, 0x00202220, 0x00202222, 0x00222020,
-    0x00222022, 0x00222220, 0x00222222, 0x01202121, 0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020,
-    0x02202022, 0x02202220, 0x02202222, 0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101,
-    0x10010001, 0x10010102, 0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001,
-    0x11020100, 0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110,
-    0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011, 0x10020112,
-    0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111, 0x11010112, 0x11010211,
-    0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110, 0x12000112, 0x12010010, 0x12010012,
-    0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121, 0x10010021, 0x10010120, 0x10010122, 0x10020121,
-    0x11000021, 0x11010022, 0x11010121, 0x11010222, 0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121,
-    0x10001001, 0x10011101, 0x10011201, 0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100,
-    0x11011101, 0x11011102, 0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102,
-    0x12001201, 0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012,
-    0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010, 0x10021111,
-    0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010, 0x11011011, 0x11011110,
-    0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110, 0x11021111, 0x11021112, 0x11021211,
-    0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011, 0x12011110, 0x12011111, 0x12011112, 0x12011211,
-    0x12011212, 0x12021111, 0x12021210, 0x12021212, 0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121,
-    0x10011220, 0x10011222, 0x10021021, 0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220,
-    0x11011020, 0x11011021, 0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220,
-    0x12001021, 0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101,
-    0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101, 0x11012200,
-    0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100, 0x12012102, 0x12012201,
-    0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010, 0x10012110, 0x10012111, 0x10012210,
-    0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111, 0x11002212, 0x11012011, 0x11012012, 0x11012110,
-    0x11012111, 0x11012112, 0x11012211, 0x11022010, 0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112,
-    0x12002211, 0x12012012, 0x12012111, 0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211,
-    0x10012122, 0x11002120, 0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221,
-    0x12012120, 0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101,
-    0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001, 0x11110100,
-    0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201, 0x12110101, 0x12110200,
-    0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210, 0x10100211, 0x10100212, 0x10110011,
-    0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211, 0x10120010, 0x10120111, 0x10120112, 0x10120210,
-    0x10120212, 0x11100011, 0x11100110, 0x11100111, 0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012,
-    0x11110110, 0x11110111, 0x11110112, 0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111,
-    0x11120112, 0x11120211, 0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211,
-    0x12120010, 0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021,
-    0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122, 0x11110221,
-    0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221, 0x12110222, 0x12120120,
-    0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102, 0x10111200, 0x10111201, 0x10121001,
-    0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100, 0x11101101, 0x11101102, 0x11101201, 0x11101202,
-    0x11111000, 0x11111001, 0x11111100, 0x11111101, 0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001,
-    0x11121002, 0x11121100, 0x11121101, 0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001,
-    0x12111100, 0x12111101, 0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011,
-    0x10101012, 0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110,
-    0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112, 0x10121211,
-    0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210, 0x11101211, 0x11111010,
-    0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210, 0x11111211, 0x11111212, 0x11121010,
-    0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210, 0x11121211, 0x11121212, 0x12101011, 0x12101110,
-    0x12101111, 0x12101211, 0x12101212, 0x12111010, 0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210,
-    0x12111211, 0x12121011, 0x12121110, 0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022,
-    0x10101120, 0x10101122, 0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221,
-    0x10121020, 0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021,
-    0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022, 0x11111120,
-    0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120, 0x11121121, 0x11121221,
-    0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222, 0x12111021, 0x12111121, 0x12111222,
-    0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221, 0x10102100, 0x10102101, 0x10102102, 0x10102201,
-    0x10112000, 0x10112101, 0x10112200, 0x10122001, 0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001,
-    0x11112100, 0x11112101, 0x11112102, 0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101,
-    0x12102002, 0x12102201, 0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011,
-    0x10102012, 0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111,
-    0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012, 0x11112110,
-    0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110, 0x11122111, 0x11122112,
-    0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110, 0x12112111, 0x12112112, 0x12112210,
-    0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121, 0x10112222, 0x10122020, 0x10122121, 0x10122122,
-    0x10122221, 0x11102121, 0x11102220, 0x11102221, 0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221,
-    0x11122022, 0x11122121, 0x11122220, 0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122,
-    0x12112220, 0x12112222, 0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100,
-    0x11210000, 0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201,
-    0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012, 0x10210111,
-    0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011, 0x11210111, 0x11210112,
-    0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212, 0x12210012, 0x12210111, 0x12220011,
-    0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221, 0x11200020, 0x11200021, 0x11200122, 0x11210121,
-    0x11210122, 0x11210220, 0x11220020, 0x12200121, 0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002,
-    0x10211101, 0x10211102, 0x10211202, 0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101,
-    0x11201200, 0x11201202, 0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000,
-    0x11221002, 0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101,
-    0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210, 0x10201212,
-    0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112, 0x11201211, 0x11211010,
-    0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011, 0x11221110, 0x11221111, 0x11221112,
-    0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011, 0x12211111, 0x12211112, 0x12211211, 0x12211212,
-    0x12221012, 0x12221111, 0x12221112, 0x12221210, 0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122,
-    0x10221220, 0x10221221, 0x11201020, 0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121,
-    0x11211122, 0x11211220, 0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121,
-    0x12201222, 0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222,
-    0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001, 0x11222201,
-    0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010, 0x10212111, 0x10222011,
-    0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111, 0x11202112, 0x11202210, 0x11212011,
-    0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010, 0x11222111, 0x11222212, 0x12202012, 0x12202110,
-    0x12202212, 0x12212111, 0x12222011, 0x12222110, 0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220,
-    0x11202021, 0x11202120, 0x11202221, 0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121,
-    0x11222221, 0x12202122, 0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200,
-    0x20000202, 0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100,
-    0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101, 0x22020000,
-    0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112, 0x20010211, 0x20020111,
-    0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111, 0x21010112, 0x21010210, 0x21010211,
-    0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211, 0x22010110, 0x22010112, 0x22010211, 0x22020111,
-    0x20000020, 0x20000022, 0x20000220, 0x20000222, 0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222,
-    0x21010021, 0x21010120, 0x21010221, 0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121,
-    0x22020020, 0x22020022, 0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001,
-    0x21011101, 0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211,
-    0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111, 0x21001210,
-    0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111, 0x21021112, 0x21021210,
-    0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010, 0x22011012, 0x22011111, 0x22011210,
-    0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121, 0x21001021, 0x21001120, 0x21001221, 0x21001222,
-    0x21011020, 0x21011121, 0x21011221, 0x21011222, 0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021,
-    0x22011222, 0x22021120, 0x20002000, 0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002,
-    0x20022200, 0x20022202, 0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201,
-    0x22002000, 0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202,
-    0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110, 0x21002112,
-    0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110, 0x22002111, 0x22012112,
-    0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222, 0x20012121, 0x20022020, 0x20022022,
-    0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120, 0x21012122, 0x22002020, 0x22002022, 0x22002220,
-    0x22002222, 0x22012121, 0x22022020, 0x22022022, 0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102,
-    0x20110200, 0x20110201, 0x20120101, 0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202,
-    0x21120201, 0x21120202, 0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011,
-    0x20100110, 0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110,
-    0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111, 0x21110112,
-    0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111, 0x22110210, 0x22120011,
-    0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120, 0x20110221, 0x20120121, 0x21100120,
-    0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121, 0x21110220, 0x21120122, 0x21120221, 0x22100121,
-    0x22110120, 0x22110122, 0x22120221, 0x20101001, 0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200,
-    0x20121102, 0x21101000, 0x21101202, 0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201,
-    0x21121000, 0x21121001, 0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101,
-    0x22111200, 0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011,
-    0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212, 0x21101011,
-    0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012, 0x21111110, 0x21111111,
-    0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110, 0x21121111, 0x21121112, 0x21121211,
-    0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012, 0x22111110, 0x22111111, 0x22111112, 0x22111211,
-    0x22111212, 0x22121010, 0x22121012, 0x22121111, 0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020,
-    0x20111121, 0x20111221, 0x20121020, 0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021,
-    0x21111022, 0x21111121, 0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221,
-    0x22101222, 0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102,
-    0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102, 0x21112202,
-    0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101, 0x20102110, 0x20102112,
-    0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212, 0x20122010, 0x20122011, 0x20122110,
-    0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210, 0x21102212, 0x21112011, 0x21112110, 0x21112111,
-    0x21112112, 0x21112211, 0x21122012, 0x21122111, 0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010,
-    0x22112012, 0x22112111, 0x22112212, 0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120,
-    0x21102122, 0x21102221, 0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120,
-    0x22112121, 0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002,
-    0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000, 0x22200002,
-    0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202, 0x20200111, 0x20200211,
-    0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112, 0x21200211, 0x21210011, 0x21210111,
-    0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111, 0x22210010, 0x22210012, 0x22210112, 0x22210211,
-    0x20200022, 0x20200220, 0x20200222, 0x20210020, 0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121,
-    0x21210021, 0x21210122, 0x21210221, 0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121,
-    0x22220020, 0x22220022, 0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000,
-    0x21211100, 0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201,
-    0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112, 0x20221211,
-    0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211, 0x21221111, 0x21221212,
-    0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012, 0x22211111, 0x22211210, 0x20201121,
-    0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121, 0x21201120, 0x21201122, 0x21201222, 0x21211022,
-    0x21211121, 0x21211122, 0x21211220, 0x21221020, 0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122,
-    0x22211221, 0x22221021, 0x22221120, 0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000,
-    0x20222002, 0x20222200, 0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002,
-    0x22202200, 0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110,
-    0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011, 0x21222112,
-    0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222, 0x20222020, 0x20222022,
-    0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, 0x22202022, 0x22202220, 0x22202222,
-    0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222,
-};
-
-static const __device__ uint8_t ksigns_iq2xs[128] = {
-    0,   129, 130, 3,   132, 5,   6,   135, 136, 9,   10,  139, 12,  141, 142, 15,  144, 17,  18,  147, 20,  149,
-    150, 23,  24,  153, 154, 27,  156, 29,  30,  159, 160, 33,  34,  163, 36,  165, 166, 39,  40,  169, 170, 43,
-    172, 45,  46,  175, 48,  177, 178, 51,  180, 53,  54,  183, 184, 57,  58,  187, 60,  189, 190, 63,  192, 65,
-    66,  195, 68,  197, 198, 71,  72,  201, 202, 75,  204, 77,  78,  207, 80,  209, 210, 83,  212, 85,  86,  215,
-    216, 89,  90,  219, 92,  221, 222, 95,  96,  225, 226, 99,  228, 101, 102, 231, 232, 105, 106, 235, 108, 237,
-    238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
-};
-
-static const __device__ uint64_t ksigns64[128] = {
-    0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, 0xff00000000ff0000,
-    0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, 0xff000000ff000000, 0x00000000ff0000ff,
-    0x00000000ff00ff00, 0xff000000ff00ffff, 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00,
-    0x00000000ffffffff, 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff,
-    0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, 0x000000ffff000000,
-    0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, 0xff0000ffffff0000, 0x000000ffffff00ff,
-    0x000000ffffffff00, 0xff0000ffffffffff, 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00,
-    0xff00ff000000ffff, 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff,
-    0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, 0xff00ff00ffff0000,
-    0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, 0x0000ffff00000000, 0xff00ffff000000ff,
-    0xff00ffff0000ff00, 0x0000ffff0000ffff, 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00,
-    0xff00ffff00ffffff, 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff,
-    0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, 0xffff000000000000,
-    0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, 0x00ff000000ff0000, 0xffff000000ff00ff,
-    0xffff000000ffff00, 0x00ff000000ffffff, 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00,
-    0x00ff0000ff00ffff, 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff,
-    0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, 0xffff00ff00ff0000,
-    0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, 0xffff00ffff000000, 0x00ff00ffff0000ff,
-    0x00ff00ffff00ff00, 0xffff00ffff00ffff, 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00,
-    0x00ff00ffffffffff, 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff,
-    0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, 0xffffff00ff000000,
-    0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, 0x00ffff00ffff0000, 0xffffff00ffff00ff,
-    0xffffff00ffffff00, 0x00ffff00ffffffff, 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00,
-    0xffffffff0000ffff, 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff,
-    0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, 0xffffffffffff0000,
-    0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff,
-};
-
-static const __device__ uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128};
-static const __device__ int8_t kvalues_iq4nl[16] = {
-    -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};
-
-typedef half dfloat;  // dequantize float
-typedef half2 dfloat2;
-typedef void (*dequantize_kernel_t)(const void* vx, const int ib, const int iqs, dfloat2& v);
-template <typename dst_t>
-using to_cuda_ggml_t = void (*)(const void* __restrict__ x, dst_t* __restrict__ y, int k, cudaStream_t stream);
-typedef float (*vec_dot_q_cuda_t)(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs);
-typedef void (*allocate_tiles_cuda_t)(int** x_ql, half2** x_dm, int** x_qh, int** x_sc);
-typedef void (*load_tiles_cuda_t)(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row);
-typedef float (*vec_dot_q_mul_mat_cuda_t)(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ms,
-    const int& i,
-    const int& j,
-    const int& k);
-
-// Utility function
-
-template <typename dst_t>
-static __device__ __forceinline__ dst_t convert_from_half(half val) {
-  return val;
-}
-
-template <>
-__device__ __forceinline__ c10::BFloat16 convert_from_half<c10::BFloat16>(half val) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA)
-  return __float2bfloat16(__half2float(val));
-#else
-  return __half2float(val);
-#endif  // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA)
-}
-
-template <>
-__device__ __forceinline__ float convert_from_half<float>(half val) {
-  return __half2float(val);
-}
-
-#if defined(USE_ROCM)
-
-#ifndef __has_builtin
-#define __has_builtin(x) 0
-#endif
-
-typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
-static __device__ __forceinline__ int __vsubss4(const int a, const int b) {
-  const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
-  const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
-#if __has_builtin(__builtin_elementwise_sub_sat)
-  const int8x4_t c = __builtin_elementwise_sub_sat(va, vb);
-  return reinterpret_cast<const int&>(c);
-#else
-  int8x4_t c;
-  int16_t tmp;
-#pragma unroll
-  for (int i = 0; i < 4; i++) {
-    tmp = va[i] - vb[i];
-    if (tmp > std::numeric_limits<int8_t>::max()) tmp = std::numeric_limits<int8_t>::max();
-    if (tmp < std::numeric_limits<int8_t>::min()) tmp = std::numeric_limits<int8_t>::min();
-    c[i] = tmp;
-  }
-  return reinterpret_cast<int&>(c);
-#endif  // __has_builtin(__builtin_elementwise_sub_sat)
-}
-
-static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) {
-#if __has_builtin(__builtin_amdgcn_sdot4)
-  c = __builtin_amdgcn_sdot4(a, b, c, false);
-#else
-  const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
-  const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
-  c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3];
-#endif
-  return c;
-}
-
-static __device__ __forceinline__ uint32_t __vcmpeq4(const uint32_t a, const uint32_t b) {
-  uint32_t neq = a ^ b;
-  return !(neq & 0xff000000) * 0xff000000 | !(neq & 0x00ff0000) * 0x00ff0000 | !(neq & 0x0000ff00) * 0x0000ff00 |
-         !(neq & 0x000000ff) * 0x000000ff;
-}
-
-static __device__ __forceinline__ uint32_t __vsub4(const uint32_t a, const uint32_t b) {
-  return (static_cast<uint8_t>(((a & 0xff000000) >> 24) - ((b & 0xff000000) >> 24)) << 24) +
-         (static_cast<uint8_t>(((a & 0x00ff0000) >> 16) - ((b & 0x00ff0000) >> 16)) << 16) +
-         (static_cast<uint8_t>(((a & 0x0000ff00) >> 8) - ((b & 0x0000ff00) >> 8)) << 8) +
-         (static_cast<uint8_t>(((a & 0x000000ff) >> 0) - ((b & 0x000000ff) >> 0)) << 0);
-}
-#endif  // defined(USE_ROCM)
diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu
deleted file mode 100644
index d88960d5fd..0000000000
--- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu
+++ /dev/null
@@ -1,848 +0,0 @@
-// Adatped from
-// https://github.com/vllm-project/vllm/blob/755ed7b05be4743237d3339c4ff8c22bcaae04f4/csrc/quantization/gguf/gguf_kernel.cu
-#include <c10/cuda/CUDAGuard.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
-
-// dont use clang-format here, it breaks the include order
-// clang-format off
-#include "dispatch.h"
-
-#include "ggml-common.h"
-#include "vecdotq.cuh"
-#include "dequantize.cuh"
-#include "mmvq.cuh"
-#include "mmq.cuh"
-#include "moe.cuh"
-#include "moe_vec.cuh"
-// clang-format off
-
-// Q8 gemv
-template <typename scalar_t>
-static __global__ void
-quantize_q8_1(const scalar_t* __restrict__ x, void* __restrict__ vy, const int kx, const int kx_padded) {
-  const auto ix = blockDim.x * blockIdx.x + threadIdx.x;
-  if (ix >= kx_padded) {
-    return;
-  }
-  const auto iy = blockDim.y * blockIdx.y + threadIdx.y;
-  const int i_padded = iy * kx_padded + ix;
-
-  block_q8_1* y = (block_q8_1*)vy;
-
-  const int ib = i_padded / QK8_1;   // block index
-  const int iqs = i_padded % QK8_1;  // quant index
-
-  const float xi = ix < kx ? static_cast<float>(x[iy * kx + ix]) : 0.0f;
-  float amax = fabsf(xi);
-  float sum = xi;
-
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1) {
-    amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), amax, mask, 32));
-    sum += SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), sum, mask, 32);
-  }
-
-  const float d = amax / 127;
-  const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
-
-  y[ib].qs[iqs] = q;
-
-  if (iqs > 0) {
-    return;
-  }
-
-  y[ib].ds.x = __float2half(d);
-  y[ib].ds.y = __float2half(sum);
-}
-
-template <typename scalar_t>
-static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx, const int ky, cudaStream_t stream) {
-  const int64_t kx_padded = (kx + 512 - 1) / 512 * 512;
-  const int block_num_x = (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
-  constexpr int MAX_BLOCK_SIZE = 65535;
-  for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) {
-    const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off;
-    const dim3 num_blocks(block_num_x, num_blocks_y, 1);
-    const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1);
-    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(
-        &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded);
-  }
-}
-
-torch::Tensor ggml_dequantize(
-    torch::Tensor W,  // quant weight
-    int64_t type,
-    int64_t m,
-    int64_t n,
-    std::optional<at::ScalarType> const& dtype) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(W));
-  auto dtype_ = dtype.value_or(torch::kFloat16);
-  auto options = torch::TensorOptions().dtype(dtype_).device(W.device());
-  at::Tensor DW = torch::empty({m, n}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-
-  DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
-    auto to_cuda = ggml_get_to_cuda<scalar_t>(type);
-    to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream);
-  });
-
-  return DW;
-}
-
-torch::Tensor ggml_mul_mat_vec_a8(
-    torch::Tensor W,  // quant weight
-    torch::Tensor X,  // input
-    int64_t type,
-    int64_t row) {
-  int col = X.sizes()[1];
-  int vecs = X.sizes()[0];
-  const int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({vecs, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options);
-  DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
-    quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream);
-    switch (type) {
-      case 2:
-        mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 3:
-        mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 6:
-        mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 7:
-        mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 8:
-        mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 10:
-        mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 11:
-        mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 12:
-        mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 13:
-        mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 14:
-        mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 16:
-        mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 17:
-        mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 18:
-        mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 19:
-        mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 20:
-        mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 21:
-        mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 22:
-        mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 23:
-        mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 29:
-        mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-torch::Tensor ggml_mul_mat_a8(
-    torch::Tensor W,  // quant weight
-    torch::Tensor X,  // input
-    int64_t type,
-    int64_t row) {
-  int col = X.sizes()[1];
-  int padded = (col + 512 - 1) / 512 * 512;
-  int batch = X.sizes()[0];
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({batch, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options);
-  DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] {
-    quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, batch, stream);
-
-    switch (type) {
-      case 2:
-        ggml_mul_mat_q4_0_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 3:
-        ggml_mul_mat_q4_1_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 6:
-        ggml_mul_mat_q5_0_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 7:
-        ggml_mul_mat_q5_1_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 8:
-        ggml_mul_mat_q8_0_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 10:
-        ggml_mul_mat_q2_K_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 11:
-        ggml_mul_mat_q3_K_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 12:
-        ggml_mul_mat_q4_K_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 13:
-        ggml_mul_mat_q5_K_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-      case 14:
-        ggml_mul_mat_q6_K_q8_1_cuda(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            col,
-            row,
-            batch,
-            padded,
-            row,
-            stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-torch::Tensor ggml_moe_a8(
-    torch::Tensor X,  // input
-    torch::Tensor W,  // expert weights
-    torch::Tensor sorted_token_ids,
-    torch::Tensor expert_ids,
-    torch::Tensor num_tokens_post_padded,
-    int64_t type,
-    int64_t row,
-    int64_t top_k,
-    int64_t tokens) {
-  int col = X.sizes()[1];
-  int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({tokens * top_k, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
-  DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
-    quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream);
-    switch (type) {
-      case 2:
-        ggml_moe_q4_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 3:
-        ggml_moe_q4_1_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 6:
-        ggml_moe_q5_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 7:
-        ggml_moe_q5_1_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 8:
-        ggml_moe_q8_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 10:
-        ggml_moe_q2_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 11:
-        ggml_moe_q3_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 12:
-        ggml_moe_q4_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 13:
-        ggml_moe_q5_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-      case 14:
-        ggml_moe_q6_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(),
-            (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(),
-            W.stride(0),
-            col,
-            row,
-            tokens,
-            padded,
-            row,
-            top_k,
-            sorted_token_ids.sizes()[0],
-            stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-torch::Tensor ggml_moe_a8_vec(
-    torch::Tensor X,  // input
-    torch::Tensor W,  // expert weights
-    torch::Tensor topk_ids,
-    int64_t top_k,
-    int64_t type,
-    int64_t row,
-    int64_t tokens) {
-  int col = X.sizes()[1];
-  const int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::zeros({tokens * top_k, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
-  DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
-    quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream);
-    switch (type) {
-      case 2:
-        moe_vec_q4_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 3:
-        moe_vec_q4_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 6:
-        moe_vec_q5_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 7:
-        moe_vec_q5_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 8:
-        moe_vec_q8_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 10:
-        moe_vec_q2_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 11:
-        moe_vec_q3_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 12:
-        moe_vec_q4_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 13:
-        moe_vec_q5_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 14:
-        moe_vec_q6_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 16:
-        moe_vec_iq2_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 17:
-        moe_vec_iq2_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 18:
-        moe_vec_iq3_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 19:
-        moe_vec_iq1_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 20:
-        moe_vec_iq4_nl_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 21:
-        moe_vec_iq3_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 22:
-        moe_vec_iq2_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 23:
-        moe_vec_iq4_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-      case 29:
-        moe_vec_iq1_m_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(),
-            (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(),
-            (int*)topk_ids.data_ptr(),
-            top_k,
-            tokens,
-            col,
-            row,
-            quant_X.stride(0),
-            stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-int64_t ggml_moe_get_block_size(int64_t type) {
-  switch (type) {
-    case 2:
-      return MOE_X_Q4_0;
-    case 3:
-      return MOE_X_Q4_1;
-    case 6:
-      return MOE_X_Q5_0;
-    case 7:
-      return MOE_X_Q5_1;
-    case 8:
-      return MOE_X_Q8_0;
-    case 10:
-      return MOE_X_Q2_K;
-    case 11:
-      return MOE_X_Q3_K;
-    case 12:
-      return MOE_X_Q4_K;
-    case 13:
-      return MOE_X_Q5_K;
-    case 14:
-      return MOE_X_Q6_K;
-  }
-  return 0;
-}
-
-// ---- FreeToken pybind bindings (donor registers these via TORCH_LIBRARY; we
-// expose them through torch.utils.cpp_extension.load's pybind module instead) ----
-#include <torch/extension.h>
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("ggml_dequantize", &ggml_dequantize, "");
-  m.def("ggml_mul_mat_vec_a8", &ggml_mul_mat_vec_a8, "");
-  m.def("ggml_mul_mat_a8", &ggml_mul_mat_a8, "");
-  m.def("ggml_moe_a8", &ggml_moe_a8, "");
-  m.def("ggml_moe_a8_vec", &ggml_moe_a8_vec, "");
-  m.def("ggml_moe_get_block_size", &ggml_moe_get_block_size, "");
-}
diff --git a/python/freetoken/kernel/csrc/gguf/mmq.cuh b/python/freetoken/kernel/csrc/gguf/mmq.cuh
deleted file mode 100644
index 9838c14f01..0000000000
--- a/python/freetoken/kernel/csrc/gguf/mmq.cuh
+++ /dev/null
@@ -1,881 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmq.cuh
-// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu
-template <
-    typename scalar_t,
-    int qk,
-    int qr,
-    int qi,
-    bool need_sum,
-    typename block_q_t,
-    int mmq_x,
-    int mmq_y,
-    int nwarps,
-    allocate_tiles_cuda_t allocate_tiles,
-    load_tiles_cuda_t load_tiles,
-    int vdr,
-    vec_dot_q_mul_mat_cuda_t vec_dot>
-static __device__ __forceinline__ void mul_mat_q(
-    const void* __restrict__ vx,
-    const void* __restrict__ vy,
-    scalar_t* __restrict__ dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst) {
-  const block_q_t* x = (const block_q_t*)vx;
-  const block_q8_1* y = (const block_q8_1*)vy;
-
-  const int blocks_per_row_x = ncols_x / qk;
-  const int blocks_per_col_y = nrows_y / QK8_1;
-  const int blocks_per_warp = WARP_SIZE_GGUF / qi;
-
-  const int& ncols_dst = ncols_y;
-
-  const auto row_dst_0 = blockIdx.x * mmq_y;
-  const int& row_x_0 = row_dst_0;
-
-  const auto col_dst_0 = blockIdx.y * mmq_x;
-  const int& col_y_0 = col_dst_0;
-
-  int* tile_x_ql = nullptr;
-  half2* tile_x_dm = nullptr;
-  int* tile_x_qh = nullptr;
-  int* tile_x_sc = nullptr;
-
-  allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc);
-
-  __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF];
-  __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1];
-
-  float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}};
-
-  for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) {
-    load_tiles(
-        x + row_x_0 * blocks_per_row_x + ib0,
-        tile_x_ql,
-        tile_x_dm,
-        tile_x_qh,
-        tile_x_sc,
-        threadIdx.y,
-        nrows_x - row_x_0 - 1,
-        threadIdx.x,
-        blocks_per_row_x);
-
-#pragma unroll
-    for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp / qr < blocks_per_row_x; ++ir) {
-      const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x;
-      const int kbxd = kqs / QI8_1;
-
-#pragma unroll
-      for (int i = 0; i < mmq_x; i += nwarps) {
-        const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y - 1);  // to prevent out-of-bounds memory accesses
-        const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + kbxd];
-        const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF;
-        tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1);
-      }
-
-#pragma unroll
-      for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) {
-        const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF / QI8_1)) % mmq_x;
-        const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1);
-        const int col_y_eff = min(col_y_0 + ids, ncols_y - 1);
-
-        // if the sum is not needed it's faster to transform the scale to f32 ahead of time
-        const half2* dsi_src =
-            &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby].ds;
-        half2* dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF / QI8_1) + kby];
-        if (need_sum) {
-          *dsi_dst = *dsi_src;
-        } else {
-          float* dfi_dst = (float*)dsi_dst;
-          *dfi_dst = __low2float(*dsi_src);
-        }
-      }
-
-      __syncthreads();
-
-      // #pragma unroll // unrolling this loop causes too much register pressure
-      for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) {
-#pragma unroll
-        for (int j = 0; j < mmq_x; j += nwarps) {
-#pragma unroll
-          for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-            sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot(
-                tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k);
-          }
-        }
-      }
-      __syncthreads();
-    }
-  }
-
-#pragma unroll
-  for (int j = 0; j < mmq_x; j += nwarps) {
-    const auto col_dst = col_dst_0 + j + threadIdx.y;
-    if (col_dst >= ncols_dst) {
-      return;
-    }
-
-#pragma unroll
-    for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-      const auto row_dst = row_dst_0 + threadIdx.x + i;
-      if (row_dst >= nrows_dst) {
-        continue;
-      }
-      dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps];
-    }
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q4_0 64
-#define MMQ_Y_Q4_0 128
-#define NWARPS_Q4_0 8
-#else
-#define MMQ_X_Q4_0 4
-#define MMQ_Y_Q4_0 32
-#define NWARPS_Q4_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2)
-#endif
-    mul_mat_q4_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q4_0;
-  const int mmq_y = MMQ_Y_Q4_0;
-  const int nwarps = NWARPS_Q4_0;
-
-  mul_mat_q<
-      scalar_t,
-      QK4_0,
-      QR4_0,
-      QI4_0,
-      true,
-      block_q4_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_0<mmq_y>,
-      load_tiles_q4_0<mmq_y, nwarps, need_check>,
-      VDR_Q4_0_Q8_1_MMQ,
-      vec_dot_q4_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q4_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  int mmq_x = MMQ_X_Q4_0;
-  int mmq_y = MMQ_Y_Q4_0;
-  int nwarps = NWARPS_Q4_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q4_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q4_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q4_1 64
-#define MMQ_Y_Q4_1 128
-#define NWARPS_Q4_1 8
-#else
-#define MMQ_X_Q4_1 4
-#define MMQ_Y_Q4_1 32
-#define NWARPS_Q4_1 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2)
-#endif
-    mul_mat_q4_1(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q4_1;
-  const int mmq_y = MMQ_Y_Q4_1;
-  const int nwarps = NWARPS_Q4_1;
-
-  mul_mat_q<
-      scalar_t,
-      QK4_1,
-      QR4_1,
-      QI4_1,
-      true,
-      block_q4_1,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_1<mmq_y>,
-      load_tiles_q4_1<mmq_y, nwarps, need_check>,
-      VDR_Q4_1_Q8_1_MMQ,
-      vec_dot_q4_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q4_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  int mmq_x = MMQ_X_Q4_1;
-  int mmq_y = MMQ_Y_Q4_1;
-  int nwarps = NWARPS_Q4_1;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q4_1<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q4_1<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q5_0 64
-#define MMQ_Y_Q5_0 128
-#define NWARPS_Q5_0 8
-#else
-#define MMQ_X_Q5_0 4
-#define MMQ_Y_Q5_0 32
-#define NWARPS_Q5_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2)
-#endif
-    mul_mat_q5_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q5_0;
-  const int mmq_y = MMQ_Y_Q5_0;
-  const int nwarps = NWARPS_Q5_0;
-
-  mul_mat_q<
-      scalar_t,
-      QK5_0,
-      QR5_0,
-      QI5_0,
-      false,
-      block_q5_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_0<mmq_y>,
-      load_tiles_q5_0<mmq_y, nwarps, need_check>,
-      VDR_Q5_0_Q8_1_MMQ,
-      vec_dot_q5_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q5_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q5_0;
-  const int mmq_y = MMQ_Y_Q5_0;
-  const int nwarps = NWARPS_Q5_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q5_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q5_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q5_1 64
-#define MMQ_Y_Q5_1 128
-#define NWARPS_Q5_1 8
-#else
-#define MMQ_X_Q5_1 4
-#define MMQ_Y_Q5_1 32
-#define NWARPS_Q5_1 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2)
-#endif
-    mul_mat_q5_1(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q5_1;
-  const int mmq_y = MMQ_Y_Q5_1;
-  const int nwarps = NWARPS_Q5_1;
-
-  mul_mat_q<
-      scalar_t,
-      QK5_1,
-      QR5_1,
-      QI5_1,
-      true,
-      block_q5_1,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_1<mmq_y>,
-      load_tiles_q5_1<mmq_y, nwarps, need_check>,
-      VDR_Q5_1_Q8_1_MMQ,
-      vec_dot_q5_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q5_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q5_1;
-  const int mmq_y = MMQ_Y_Q5_1;
-  const int nwarps = NWARPS_Q5_1;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q5_1<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q5_1<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q8_0 64
-#define MMQ_Y_Q8_0 128
-#define NWARPS_Q8_0 8
-#else
-#define MMQ_X_Q8_0 4
-#define MMQ_Y_Q8_0 32
-#define NWARPS_Q8_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2)
-#endif
-    mul_mat_q8_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q8_0;
-  const int mmq_y = MMQ_Y_Q8_0;
-  const int nwarps = NWARPS_Q8_0;
-
-  mul_mat_q<
-      scalar_t,
-      QK8_0,
-      QR8_0,
-      QI8_0,
-      false,
-      block_q8_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q8_0<mmq_y>,
-      load_tiles_q8_0<mmq_y, nwarps, need_check>,
-      VDR_Q8_0_Q8_1_MMQ,
-      vec_dot_q8_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q8_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q8_0;
-  const int mmq_y = MMQ_Y_Q8_0;
-  const int nwarps = NWARPS_Q8_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q8_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q8_0<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q2_K 64
-#define MMQ_Y_Q2_K 128
-#define NWARPS_Q2_K 8
-#else
-#define MMQ_X_Q2_K 4
-#define MMQ_Y_Q2_K 32
-#define NWARPS_Q2_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2)
-#endif
-    mul_mat_q2_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q2_K;
-  const int mmq_y = MMQ_Y_Q2_K;
-  const int nwarps = NWARPS_Q2_K;
-
-  mul_mat_q<
-      scalar_t,
-      QK_K,
-      QR2_K,
-      QI2_K,
-      false,
-      block_q2_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q2_K<mmq_y>,
-      load_tiles_q2_K<mmq_y, nwarps, need_check>,
-      VDR_Q2_K_Q8_1_MMQ,
-      vec_dot_q2_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q2_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q2_K;
-  const int mmq_y = MMQ_Y_Q2_K;
-  const int nwarps = NWARPS_Q2_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q2_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q2_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q3_K 64
-#define MMQ_Y_Q3_K 128
-#define NWARPS_Q3_K 8
-#else
-#define MMQ_X_Q3_K 4
-#define MMQ_Y_Q3_K 32
-#define NWARPS_Q3_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2)
-#endif
-    mul_mat_q3_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-
-  const int mmq_x = MMQ_X_Q3_K;
-  const int mmq_y = MMQ_Y_Q3_K;
-  const int nwarps = NWARPS_Q3_K;
-
-  mul_mat_q<
-      scalar_t,
-      QK_K,
-      QR3_K,
-      QI3_K,
-      false,
-      block_q3_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q3_K<mmq_y>,
-      load_tiles_q3_K<mmq_y, nwarps, need_check>,
-      VDR_Q3_K_Q8_1_MMQ,
-      vec_dot_q3_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q3_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q3_K;
-  const int mmq_y = MMQ_Y_Q3_K;
-  const int nwarps = NWARPS_Q3_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q3_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q3_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q4_K 64
-#define MMQ_Y_Q4_K 128
-#define NWARPS_Q4_K 8
-#else
-#define MMQ_X_Q4_K 4
-#define MMQ_Y_Q4_K 32
-#define NWARPS_Q4_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2)
-#endif
-    mul_mat_q4_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q4_K;
-  const int mmq_y = MMQ_Y_Q4_K;
-  const int nwarps = NWARPS_Q4_K;
-
-  mul_mat_q<
-      scalar_t,
-      QK_K,
-      QR4_K,
-      QI4_K,
-      true,
-      block_q4_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_K<mmq_y>,
-      load_tiles_q4_K<mmq_y, nwarps, need_check>,
-      VDR_Q4_K_Q8_1_MMQ,
-      vec_dot_q4_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q4_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q4_K;
-  const int mmq_y = MMQ_Y_Q4_K;
-  const int nwarps = NWARPS_Q4_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q4_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q4_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q5_K 64
-#define MMQ_Y_Q5_K 128
-#define NWARPS_Q5_K 8
-#else
-#define MMQ_X_Q5_K 4
-#define MMQ_Y_Q5_K 32
-#define NWARPS_Q5_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2)
-#endif
-    mul_mat_q5_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q5_K;
-  const int mmq_y = MMQ_Y_Q5_K;
-  const int nwarps = NWARPS_Q5_K;
-
-  mul_mat_q<
-      scalar_t,
-      QK_K,
-      QR5_K,
-      QI5_K,
-      true,
-      block_q5_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_K<mmq_y>,
-      load_tiles_q5_K<mmq_y, nwarps, need_check>,
-      VDR_Q5_K_Q8_1_MMQ,
-      vec_dot_q5_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q5_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q5_K;
-  const int mmq_y = MMQ_Y_Q5_K;
-  const int nwarps = NWARPS_Q5_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q5_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q5_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MMQ_X_Q6_K 64
-#define MMQ_Y_Q6_K 128
-#define NWARPS_Q6_K 8
-#else
-#define MMQ_X_Q6_K 4
-#define MMQ_Y_Q6_K 32
-#define NWARPS_Q6_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2)
-#endif
-    mul_mat_q6_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst) {
-  const int mmq_x = MMQ_X_Q6_K;
-  const int mmq_y = MMQ_Y_Q6_K;
-  const int nwarps = NWARPS_Q6_K;
-
-  mul_mat_q<
-      scalar_t,
-      QK_K,
-      QR6_K,
-      QI6_K,
-      false,
-      block_q6_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q6_K<mmq_y>,
-      load_tiles_q6_K<mmq_y, nwarps, need_check>,
-      VDR_Q6_K_Q8_1_MMQ,
-      vec_dot_q6_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template <typename scalar_t>
-static void ggml_mul_mat_q6_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    cudaStream_t stream) {
-  const int mmq_x = MMQ_X_Q6_K;
-  const int mmq_y = MMQ_Y_Q6_K;
-  const int nwarps = NWARPS_Q6_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    const bool need_check = false;
-    mul_mat_q6_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  } else {
-    const bool need_check = true;
-    mul_mat_q6_K<scalar_t, need_check>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-  }
-}
diff --git a/python/freetoken/kernel/csrc/gguf/mmvq.cuh b/python/freetoken/kernel/csrc/gguf/mmvq.cuh
deleted file mode 100644
index 7331731ace..0000000000
--- a/python/freetoken/kernel/csrc/gguf/mmvq.cuh
+++ /dev/null
@@ -1,352 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmvq.cuh
-// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
-template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
-static __global__ void mul_mat_vec_q(
-    const void* __restrict__ vx,
-    const void* __restrict__ vy,
-    scalar_t* __restrict__ dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs) {
-  const auto row = blockIdx.x * blockDim.y + threadIdx.y;
-  const auto vec = blockIdx.y;
-
-  if (row >= nrows || vec >= nvecs) {
-    return;
-  }
-
-  const int blocks_per_row = ncols / qk;
-  const int blocks_per_warp = vdr * WARP_SIZE / qi;
-  const int nrows_y = (ncols + 512 - 1) / 512 * 512;
-
-  // partial sum for each thread
-  float tmp = 0.0f;
-
-  const block_q_t* x = (const block_q_t*)vx;
-  const block_q8_1* y = (const block_q8_1*)vy;
-
-  for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) {
-    const int ibx = row * blocks_per_row + i;  // x block index
-
-    const int iby = vec * (nrows_y / QK8_1) + i * (qk / QK8_1);  // y block index that aligns with ibx
-
-    const int iqs = vdr * (threadIdx.x % (qi / vdr));  // x block quant index when casting the quants to int
-
-    tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);
-  }
-
-  // sum up partial sums and write back result
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
-    tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask);
-  }
-
-  if (threadIdx.x == 0) {
-    dst[vec * nrows + row] = tmp;
-  }
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q4_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q4_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q5_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q5_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q8_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q2_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q3_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q4_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q5_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_q6_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq2_xxs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq2_xs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq2_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq3_xxs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq1_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq1_m_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq4_nl_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq4_xs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template <typename scalar_t>
-static void mul_mat_vec_iq3_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int ncols,
-    const int nrows,
-    const int nvecs,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, nvecs, 1);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  mul_mat_vec_q<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
diff --git a/python/freetoken/kernel/csrc/gguf/moe.cuh b/python/freetoken/kernel/csrc/gguf/moe.cuh
deleted file mode 100644
index 91e434fa60..0000000000
--- a/python/freetoken/kernel/csrc/gguf/moe.cuh
+++ /dev/null
@@ -1,1379 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe.cuh
-#include <cstdint>
-
-/* Adapted from ./csrc/quantization/gguf/mmq.cuh
- */
-template <
-    typename scalar_t,
-    int qk,
-    int qr,
-    int qi,
-    bool need_sum,
-    typename block_q_t,
-    int mmq_x,
-    int mmq_y,
-    int nwarps,
-    allocate_tiles_cuda_t allocate_tiles,
-    load_tiles_cuda_t load_tiles,
-    int vdr,
-    vec_dot_q_mul_mat_cuda_t vec_dot>
-static __device__ __forceinline__ void moe_q(
-    const void* __restrict__ vx,
-    const void* __restrict__ vy,
-    scalar_t* __restrict__ dst,
-    const int* __restrict__ sorted_token_ids,
-    const int* __restrict__ expert_ids,
-    const int* __restrict__ num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k) {
-  const int blocks_per_row_x = ncols_x / qk;
-  const int blocks_per_col_y = nrows_y / QK8_1;
-  const int blocks_per_warp = WARP_SIZE_GGUF / qi;
-
-  const int ncols_dst = ncols_y * top_k;
-
-  const auto row_dst_0 = blockIdx.x * mmq_y;
-  const int& row_x_0 = row_dst_0;
-
-  const auto col_dst_0 = blockIdx.y * mmq_x;
-
-  int token_offs[mmq_x / nwarps];
-  for (int i = 0; i < mmq_x; i += nwarps) {
-    token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i];
-  }
-
-  const int exp_idx = expert_ids[blockIdx.y];
-  if (exp_idx > 255 || exp_idx < 0) return;
-  if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return;
-
-  const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride);
-  const block_q8_1* y = (const block_q8_1*)(vy);
-
-  int* tile_x_ql = nullptr;
-  half2* tile_x_dm = nullptr;
-  int* tile_x_qh = nullptr;
-  int* tile_x_sc = nullptr;
-
-  allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc);
-
-  __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF];
-  __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1];
-
-  float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}};
-
-  for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) {
-    load_tiles(
-        x + row_x_0 * blocks_per_row_x + ib0,
-        tile_x_ql,
-        tile_x_dm,
-        tile_x_qh,
-        tile_x_sc,
-        threadIdx.y,
-        nrows_x - row_x_0 - 1,
-        threadIdx.x,
-        blocks_per_row_x);
-
-    const int n_per_r = ((qk * blocks_per_warp) / qr);
-#pragma unroll
-    for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) {
-      const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x;
-      const int kbxd = kqs / QI8_1;
-
-#pragma unroll
-      for (int i = 0; i < mmq_x; i += nwarps) {
-        const int col_y_eff = token_offs[i / nwarps] / top_k;
-        const int block_x = ib0 * (qk / QK8_1) + kbxd;
-        if (col_y_eff < ncols_y && block_x < blocks_per_col_y) {
-          const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + block_x];
-          const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF;
-          tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1);
-        }
-      }
-
-      if (threadIdx.x < n_per_r / QK8_1) {
-        const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1);
-        const int col_y_eff = token_offs[threadIdx.y] / top_k;
-        const int block_x = ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby;
-
-        if (col_y_eff < ncols_y && block_x < blocks_per_col_y) {
-          const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds;
-          half2* dsi_dst = &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby];
-
-          if (need_sum) {
-            *dsi_dst = *dsi_src;
-          } else {
-            float* dfi_dst = (float*)dsi_dst;
-            *dfi_dst = __low2float(*dsi_src);
-          }
-        }
-      }
-      __syncthreads();
-
-      // #pragma unroll // unrolling this loop causes too much register pressure
-      for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) {
-#pragma unroll
-        for (int j = 0; j < mmq_x; j += nwarps) {
-#pragma unroll
-          for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-            sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot(
-                tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k);
-          }
-        }
-      }
-      __syncthreads();
-    }
-  }
-
-#pragma unroll
-  for (int j = 0; j < mmq_x; j += nwarps) {
-    const int col_dst = token_offs[j / nwarps];
-    if (col_dst >= ncols_dst) {
-      return;
-    }
-
-#pragma unroll
-    for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-      const auto row_dst = row_dst_0 + threadIdx.x + i;
-      if (row_dst >= nrows_dst) {
-        continue;
-      }
-      dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps];
-    }
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q4_0 8
-#define MOE_Y_Q4_0 128
-#define NWARPS_Q4_0 8
-#else
-#define MOE_X_Q4_0 4
-#define MOE_Y_Q4_0 32
-#define NWARPS_Q4_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2)
-#endif
-    moe_q4_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q4_0;
-  const int mmq_y = MOE_Y_Q4_0;
-  const int nwarps = NWARPS_Q4_0;
-
-  moe_q<
-      scalar_t,
-      QK4_0,
-      QR4_0,
-      QI4_0,
-      true,
-      block_q4_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_0<mmq_y>,
-      load_tiles_q4_0<mmq_y, nwarps, need_check>,
-      VDR_Q4_0_Q8_1_MMQ,
-      vec_dot_q4_0_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q4_0_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  int mmq_x = MOE_X_Q4_0;
-  int mmq_y = MOE_Y_Q4_0;
-  int nwarps = NWARPS_Q4_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q4_1 8
-#define MOE_Y_Q4_1 128
-#define NWARPS_Q4_1 8
-#else
-#define MOE_X_Q4_1 4
-#define MOE_Y_Q4_1 32
-#define NWARPS_Q4_1 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2)
-#endif
-    moe_q4_1(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q4_1;
-  const int mmq_y = MOE_Y_Q4_1;
-  const int nwarps = NWARPS_Q4_1;
-
-  moe_q<
-      scalar_t,
-      QK4_1,
-      QR4_1,
-      QI4_1,
-      true,
-      block_q4_1,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_1<mmq_y>,
-      load_tiles_q4_1<mmq_y, nwarps, need_check>,
-      VDR_Q4_1_Q8_1_MMQ,
-      vec_dot_q4_1_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q4_1_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  int mmq_x = MOE_X_Q4_1;
-  int mmq_y = MOE_Y_Q4_1;
-  int nwarps = NWARPS_Q4_1;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q5_0 8
-#define MOE_Y_Q5_0 128
-#define NWARPS_Q5_0 8
-#else
-#define MOE_X_Q5_0 4
-#define MOE_Y_Q5_0 32
-#define NWARPS_Q5_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2)
-#endif
-    moe_q5_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q5_0;
-  const int mmq_y = MOE_Y_Q5_0;
-  const int nwarps = NWARPS_Q5_0;
-
-  moe_q<
-      scalar_t,
-      QK5_0,
-      QR5_0,
-      QI5_0,
-      false,
-      block_q5_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_0<mmq_y>,
-      load_tiles_q5_0<mmq_y, nwarps, need_check>,
-      VDR_Q5_0_Q8_1_MMQ,
-      vec_dot_q5_0_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q5_0_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q5_0;
-  const int mmq_y = MOE_Y_Q5_0;
-  const int nwarps = NWARPS_Q5_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q5_1 8
-#define MOE_Y_Q5_1 128
-#define NWARPS_Q5_1 8
-#else
-#define MOE_X_Q5_1 4
-#define MOE_Y_Q5_1 32
-#define NWARPS_Q5_1 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2)
-#endif
-    moe_q5_1(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q5_1;
-  const int mmq_y = MOE_Y_Q5_1;
-  const int nwarps = NWARPS_Q5_1;
-
-  moe_q<
-      scalar_t,
-      QK5_1,
-      QR5_1,
-      QI5_1,
-      true,
-      block_q5_1,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_1<mmq_y>,
-      load_tiles_q5_1<mmq_y, nwarps, need_check>,
-      VDR_Q5_1_Q8_1_MMQ,
-      vec_dot_q5_1_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q5_1_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q5_1;
-  const int mmq_y = MOE_Y_Q5_1;
-  const int nwarps = NWARPS_Q5_1;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q8_0 8
-#define MOE_Y_Q8_0 128
-#define NWARPS_Q8_0 8
-#else
-#define MOE_X_Q8_0 4
-#define MOE_Y_Q8_0 32
-#define NWARPS_Q8_0 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2)
-#endif
-    moe_q8_0(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q8_0;
-  const int mmq_y = MOE_Y_Q8_0;
-  const int nwarps = NWARPS_Q8_0;
-
-  moe_q<
-      scalar_t,
-      QK8_0,
-      QR8_0,
-      QI8_0,
-      false,
-      block_q8_0,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q8_0<mmq_y>,
-      load_tiles_q8_0<mmq_y, nwarps, need_check>,
-      VDR_Q8_0_Q8_1_MMQ,
-      vec_dot_q8_0_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q8_0_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q8_0;
-  const int mmq_y = MOE_Y_Q8_0;
-  const int nwarps = NWARPS_Q8_0;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q2_K 8
-#define MOE_Y_Q2_K 128
-#define NWARPS_Q2_K 8
-#else
-#define MOE_X_Q2_K 4
-#define MOE_Y_Q2_K 32
-#define NWARPS_Q2_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2)
-#endif
-    moe_q2_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q2_K;
-  const int mmq_y = MOE_Y_Q2_K;
-  const int nwarps = NWARPS_Q2_K;
-
-  moe_q<
-      scalar_t,
-      QK_K,
-      QR2_K,
-      QI2_K,
-      false,
-      block_q2_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q2_K<mmq_y>,
-      load_tiles_q2_K<mmq_y, nwarps, need_check>,
-      VDR_Q2_K_Q8_1_MMQ,
-      vec_dot_q2_K_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q2_K_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q2_K;
-  const int mmq_y = MOE_Y_Q2_K;
-  const int nwarps = NWARPS_Q2_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q3_K 8
-#define MOE_Y_Q3_K 128
-#define NWARPS_Q3_K 8
-#else
-#define MOE_X_Q3_K 4
-#define MOE_Y_Q3_K 32
-#define NWARPS_Q3_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2)
-#endif
-    moe_q3_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-
-  const int mmq_x = MOE_X_Q3_K;
-  const int mmq_y = MOE_Y_Q3_K;
-  const int nwarps = NWARPS_Q3_K;
-
-  moe_q<
-      scalar_t,
-      QK_K,
-      QR3_K,
-      QI3_K,
-      false,
-      block_q3_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q3_K<mmq_y>,
-      load_tiles_q3_K<mmq_y, nwarps, need_check>,
-      VDR_Q3_K_Q8_1_MMQ,
-      vec_dot_q3_K_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-template <typename scalar_t>
-static void ggml_moe_q3_K_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q3_K;
-  const int mmq_y = MOE_Y_Q3_K;
-  const int nwarps = NWARPS_Q3_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q4_K 8
-#define MOE_Y_Q4_K 128
-#define NWARPS_Q4_K 8
-#else
-#define MOE_X_Q4_K 4
-#define MOE_Y_Q4_K 32
-#define NWARPS_Q4_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2)
-#endif
-    moe_q4_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q4_K;
-  const int mmq_y = MOE_Y_Q4_K;
-  const int nwarps = NWARPS_Q4_K;
-
-  moe_q<
-      scalar_t,
-      QK_K,
-      QR4_K,
-      QI4_K,
-      true,
-      block_q4_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q4_K<mmq_y>,
-      load_tiles_q4_K<mmq_y, nwarps, need_check>,
-      VDR_Q4_K_Q8_1_MMQ,
-      vec_dot_q4_K_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q4_K_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q4_K;
-  const int mmq_y = MOE_Y_Q4_K;
-  const int nwarps = NWARPS_Q4_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q5_K 8
-#define MOE_Y_Q5_K 128
-#define NWARPS_Q5_K 8
-#else
-#define MOE_X_Q5_K 4
-#define MOE_Y_Q5_K 32
-#define NWARPS_Q5_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2)
-#endif
-    moe_q5_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q5_K;
-  const int mmq_y = MOE_Y_Q5_K;
-  const int nwarps = NWARPS_Q5_K;
-
-  moe_q<
-      scalar_t,
-      QK_K,
-      QR5_K,
-      QI5_K,
-      true,
-      block_q5_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q5_K<mmq_y>,
-      load_tiles_q5_K<mmq_y, nwarps, need_check>,
-      VDR_Q5_K_Q8_1_MMQ,
-      vec_dot_q5_K_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q5_K_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q5_K;
-  const int mmq_y = MOE_Y_Q5_K;
-  const int nwarps = NWARPS_Q5_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
-
-#if defined(USE_ROCM)
-#define MOE_X_Q6_K 8
-#define MOE_Y_Q6_K 128
-#define NWARPS_Q6_K 8
-#else
-#define MOE_X_Q6_K 4
-#define MOE_Y_Q6_K 32
-#define NWARPS_Q6_K 4
-#endif
-
-template <typename scalar_t, bool need_check>
-static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2)
-#endif
-    moe_q6_K(
-        const void* __restrict__ vx,
-        const void* __restrict__ vy,
-        scalar_t* __restrict__ dst,
-        const int* sorted_token_ids,
-        const int* expert_ids,
-        const int* num_tokens_post_padded,
-        const int exp_stride,
-        const int ncols_x,
-        const int nrows_x,
-        const int ncols_y,
-        const int nrows_y,
-        const int nrows_dst,
-        const int top_k) {
-  const int mmq_x = MOE_X_Q6_K;
-  const int mmq_y = MOE_Y_Q6_K;
-  const int nwarps = NWARPS_Q6_K;
-
-  moe_q<
-      scalar_t,
-      QK_K,
-      QR6_K,
-      QI6_K,
-      false,
-      block_q6_K,
-      mmq_x,
-      mmq_y,
-      nwarps,
-      allocate_tiles_q6_K<mmq_y>,
-      load_tiles_q6_K<mmq_y, nwarps, need_check>,
-      VDR_Q6_K_Q8_1_MMQ,
-      vec_dot_q6_K_q8_1_mul_mat>(
-      vx,
-      vy,
-      dst,
-      sorted_token_ids,
-      expert_ids,
-      num_tokens_post_padded,
-      exp_stride,
-      ncols_x,
-      nrows_x,
-      ncols_y,
-      nrows_y,
-      nrows_dst,
-      top_k);
-}
-
-template <typename scalar_t>
-static void ggml_moe_q6_K_q8_1_cuda(
-    const void* inp,
-    const void* w,
-    scalar_t* dst,
-    const int* sorted_token_ids,
-    const int* expert_ids,
-    const int* num_tokens_post_padded,
-    const int exp_stride,
-    const int ncols_x,
-    const int nrows_x,
-    const int ncols_y,
-    const int nrows_y,
-    const int nrows_dst,
-    const int top_k,
-    const int tokens_post_padded,
-    cudaStream_t stream) {
-  const int mmq_x = MOE_X_Q6_K;
-  const int mmq_y = MOE_Y_Q6_K;
-  const int nwarps = NWARPS_Q6_K;
-
-  const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-  const int block_num_y = (tokens_post_padded) / mmq_x;
-  const dim3 block_nums(block_num_x, block_num_y, 1);
-  const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-  if (nrows_x % mmq_y == 0) {
-    constexpr bool need_check = false;
-    moe_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  } else {
-    constexpr bool need_check = true;
-    moe_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
-        w,
-        inp,
-        dst,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        exp_stride,
-        ncols_x,
-        nrows_x,
-        ncols_y,
-        nrows_y,
-        nrows_dst,
-        top_k);
-  }
-}
diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh
deleted file mode 100644
index 8cef9e080a..0000000000
--- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh
+++ /dev/null
@@ -1,413 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe_vec.cuh
-// copied and adapted from
-// https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
-template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
-static __global__ void moe_vec_q(
-    const void* __restrict__ vx,
-    const void* __restrict__ vy,
-    scalar_t* __restrict__ dst,
-    const int* topk_ids,
-    const int topk,
-    const int ncols,
-    const int nrows,
-    const int token_stride) {
-  const auto row = blockIdx.x * blockDim.y + threadIdx.y;
-
-  const auto token = blockIdx.z / topk;
-  const auto expert = (topk_ids)[blockIdx.z];
-
-  if (row >= nrows) {
-    return;
-  }
-
-  const int blocks_per_row = ncols / qk;
-  const int blocks_per_warp = vdr * WARP_SIZE / qi;
-
-  // partial sum for each thread
-  float tmp = 0.0f;
-
-  const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row;
-  const block_q8_1* y = (const block_q8_1*)(((const int*)vy) + token * token_stride);
-
-  for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) {
-    const int ibx = row * blocks_per_row + i;  // x block index
-
-    const int iby = i * (qk / QK8_1);  // y block index that aligns with ibx
-
-    const int iqs = vdr * (threadIdx.x % (qi / vdr));  // x block quant index when casting the quants to int
-
-    tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);
-  }
-
-  // sum up partial sums and write back result
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
-    tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask);
-  }
-
-  if (threadIdx.x == 0) {
-    dst[blockIdx.z * nrows + row] = tmp;
-  }
-}
-
-template <typename scalar_t>
-static void moe_vec_q4_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q4_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q5_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q5_1_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q8_0_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q2_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q3_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q4_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q5_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_q6_K_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq2_xxs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq2_xs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq2_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq3_xxs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq1_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq1_m_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq4_nl_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq4_xs_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
-
-template <typename scalar_t>
-static void moe_vec_iq3_s_q8_1_cuda(
-    const void* vx,
-    const void* vy,
-    scalar_t* dst,
-    const int* topk_ids,
-    const int top_k,
-    const int tokens,
-    const int ncols,
-    const int nrows,
-    const int token_stride,
-    cudaStream_t stream) {
-  const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-  const dim3 block_nums(block_num_y, 1, tokens * top_k);
-  const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-  moe_vec_q<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
-      <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
-}
diff --git a/python/freetoken/kernel/csrc/gguf/vecdotq.cuh b/python/freetoken/kernel/csrc/gguf/vecdotq.cuh
deleted file mode 100644
index 08b4cd269c..0000000000
--- a/python/freetoken/kernel/csrc/gguf/vecdotq.cuh
+++ /dev/null
@@ -1,2037 +0,0 @@
-// copied from
-// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/vecdotq.cuh
-// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/vecdotq.cuh
-// and https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu
-static __device__ __forceinline__ int get_int_b2(const void* x, const int& i32) {
-  const uint16_t* x16 = (const uint16_t*)x;  // assume at least 2 byte alignment
-
-  int x32 = x16[2 * i32 + 0] << 0;
-  x32 |= x16[2 * i32 + 1] << 16;
-
-  return x32;
-}
-
-static __device__ __forceinline__ int get_int_b4(const void* x, const int& i32) {
-  return ((const int*)x)[i32];  // assume at least 4 byte alignment
-}
-
-static __device__ __forceinline__ int get_int_from_int8(const int8_t* x8, const int& i32) {
-  const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32);  // assume at least 2 byte alignment
-  int x32 = 0;
-  x32 |= x16[0] << 0;
-  x32 |= x16[1] << 16;
-  return x32;
-}
-
-static __device__ __forceinline__ int get_int_from_uint8(const uint8_t* x8, const int& i32) {
-  const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32);  // assume at least 2 byte alignment
-  int x32 = 0;
-  x32 |= x16[0] << 0;
-  x32 |= x16[1] << 16;
-  return x32;
-}
-
-static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t* x8, const int& i32) {
-  return *((const int*)(x8 + sizeof(int) * i32));  // assume at least 4 byte alignment
-}
-
-static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t* x8, const int& i32) {
-  return *((const int*)(x8 + sizeof(int) * i32));  // assume at least 4 byte alignment
-}
-
-// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called
-// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q
-
-#define VDR_Q4_0_Q8_1_MMVQ 2
-#define VDR_Q4_0_Q8_1_MMQ 4
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q4_0_q8_1_impl(const int* v, const int* u, const float& d4, const half2& ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    const int vi0 = (v[i] >> 0) & 0x0F0F0F0F;
-    const int vi1 = (v[i] >> 4) & 0x0F0F0F0F;
-
-    // SIMD dot product of quantized values
-    sumi = __dp4a(vi0, u[2 * i + 0], sumi);
-    sumi = __dp4a(vi1, u[2 * i + 1], sumi);
-  }
-
-  const float2 ds8f = __half22float2(ds8);
-
-  // second part effectively subtracts 8 from each quant value
-  return d4 * (sumi * ds8f.x - (8 * vdr / QI4_0) * ds8f.y);
-#endif
-}
-
-#define VDR_Q4_1_Q8_1_MMVQ 2
-#define VDR_Q4_1_Q8_1_MMQ 4
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q4_1_q8_1_impl(const int* v, const int* u, const half2& dm4, const half2& ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    const int vi0 = (v[i] >> 0) & 0x0F0F0F0F;
-    const int vi1 = (v[i] >> 4) & 0x0F0F0F0F;
-
-    // SIMD dot product of quantized values
-    sumi = __dp4a(vi0, u[2 * i + 0], sumi);
-    sumi = __dp4a(vi1, u[2 * i + 1], sumi);
-  }
-
-  const float2 tmp = __half22float2(__hmul2(dm4, ds8));
-  const float d4d8 = tmp.x;
-  const float m4s8 = tmp.y;
-
-  // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it
-  return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1));
-#endif
-}
-
-#define VDR_Q5_0_Q8_1_MMVQ 2
-#define VDR_Q5_0_Q8_1_MMQ 4
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q5_0_q8_1_impl(const int* vl, const int* vh, const int* u, const float& d5, const half2& ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    int vi0 = (vl[i] >> 0) & 0x0F0F0F0F;     // lower 4 qs bits, still need qh as 5th bits
-    vi0 |= (vh[i] << 4) & 0x00000010;        // 0 ->  4
-    vi0 |= (vh[i] << 11) & 0x00001000;       // 1 -> 12
-    vi0 |= (vh[i] << 18) & 0x00100000;       // 2 -> 20
-    vi0 |= (vh[i] << 25) & 0x10000000;       // 3 -> 28
-    sumi = __dp4a(vi0, u[2 * i + 0], sumi);  // SIMD dot product of quantized values
-
-    int vi1 = (vl[i] >> 4) & 0x0F0F0F0F;     // upper 4 qs bits, still need qh as 5th bits
-    vi1 |= (vh[i] >> 12) & 0x00000010;       // 16 ->  4
-    vi1 |= (vh[i] >> 5) & 0x00001000;        // 17 -> 12
-    vi1 |= (vh[i] << 2) & 0x00100000;        // 18 -> 20
-    vi1 |= (vh[i] << 9) & 0x10000000;        // 19 -> 28
-    sumi = __dp4a(vi1, u[2 * i + 1], sumi);  // SIMD dot product of quantized values
-  }
-
-  const float2 ds8f = __half22float2(ds8);
-
-  // second part effectively subtracts 16 from each quant value
-  return d5 * (sumi * ds8f.x - (16 * vdr / QI5_0) * ds8f.y);
-#endif
-}
-
-#define VDR_Q5_1_Q8_1_MMVQ 2
-#define VDR_Q5_1_Q8_1_MMQ 4
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q5_1_q8_1_impl(const int* vl, const int* vh, const int* u, const half2& dm5, const half2& ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    int vi0 = (vl[i] >> 0) & 0x0F0F0F0F;     // lower 4 qs bits, still need qh as 5th bits
-    vi0 |= (vh[i] << 4) & 0x00000010;        // 0 ->  4
-    vi0 |= (vh[i] << 11) & 0x00001000;       // 1 -> 12
-    vi0 |= (vh[i] << 18) & 0x00100000;       // 2 -> 20
-    vi0 |= (vh[i] << 25) & 0x10000000;       // 3 -> 28
-    sumi = __dp4a(vi0, u[2 * i + 0], sumi);  // SIMD dot product of quantized values
-
-    int vi1 = (vl[i] >> 4) & 0x0F0F0F0F;     // upper 4 qs bits, still need qh as 5th bits
-    vi1 |= (vh[i] >> 12) & 0x00000010;       // 16 ->  4
-    vi1 |= (vh[i] >> 5) & 0x00001000;        // 17 -> 12
-    vi1 |= (vh[i] << 2) & 0x00100000;        // 18 -> 20
-    vi1 |= (vh[i] << 9) & 0x10000000;        // 19 -> 28
-    sumi = __dp4a(vi1, u[2 * i + 1], sumi);  // SIMD dot product of quantized values
-  }
-
-  const float2 tmp = __half22float2(__hmul2(dm5, ds8));
-  const float d5d8 = tmp.x;
-  const float m5s8 = tmp.y;
-
-  // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it
-  return sumi * d5d8 + m5s8 / (QI5_1 / vdr);
-#endif
-}
-
-#define VDR_Q8_0_Q8_1_MMVQ 2
-#define VDR_Q8_0_Q8_1_MMQ 8
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q8_0_q8_1_impl(const int* v, const int* u, const float& d8_0, const float& d8_1) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    // SIMD dot product of quantized values
-    sumi = __dp4a(v[i], u[i], sumi);
-  }
-  return d8_0 * d8_1 * sumi;
-#endif
-}
-
-template <int vdr>
-static __device__ __forceinline__ float
-vec_dot_q8_1_q8_1_impl(const int* v, const int* u, const half2& dm8, const half2& ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  int sumi = 0;
-
-#pragma unroll
-  for (int i = 0; i < vdr; ++i) {
-    // SIMD dot product of quantized values
-    sumi = __dp4a(v[i], u[i], sumi);
-  }
-
-  const float2 tmp = __half22float2(__hmul2(dm8, ds8));
-  const float d8d8 = tmp.x;
-  const float m8s8 = tmp.y;
-
-  // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it
-  return sumi * d8d8 + m8s8 / (QI8_1 / vdr);
-#endif
-}
-
-#define VDR_Q2_K_Q8_1_MMVQ 1
-#define VDR_Q2_K_Q8_1_MMQ 2
-
-// contiguous v/x values
-static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
-    const int& v,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ scales,
-    const half2& dm2,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  float sumf_d = 0.0f;
-  float sumf_m = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR2_K; ++i) {
-    const int sc = scales[2 * i];
-
-    const int vi = (v >> (2 * i)) & 0x03030303;
-
-    sumf_d += d8[i] * (__dp4a(vi, u[i], 0) * (sc & 0xF));  // SIMD dot product
-
-    // fill int with 4x m
-    int m = sc >> 4;
-    m |= m << 8;
-    m |= m << 16;
-    sumf_m += d8[i] * __dp4a(m, u[i], 0);  // multiply constant q2_K part with sum of q8_1 values
-  }
-
-  const float2 dm2f = __half22float2(dm2);
-
-  return dm2f.x * sumf_d - dm2f.y * sumf_m;
-#endif
-}
-
-static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ scales,
-    const half2& dm2,
-    const float& d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi_d = 0;
-  int sumi_m = 0;
-
-#pragma unroll
-  for (int i0 = 0; i0 < QI8_1; i0 += QI8_1 / 2) {
-    int sumi_d_sc = 0;
-
-    const int sc = scales[i0 / (QI8_1 / 2)];
-
-    // fill int with 4x m
-    int m = sc >> 4;
-    m |= m << 8;
-    m |= m << 16;
-
-#pragma unroll
-    for (int i = i0; i < i0 + QI8_1 / 2; ++i) {
-      sumi_d_sc = __dp4a(v[i], u[i], sumi_d_sc);  // SIMD dot product
-      sumi_m = __dp4a(m, u[i], sumi_m);           // multiply sum of q8_1 values with m
-    }
-
-    sumi_d += sumi_d_sc * (sc & 0xF);
-  }
-
-  const float2 dm2f = __half22float2(dm2);
-
-  return d8 * (dm2f.x * sumi_d - dm2f.y * sumi_m);
-#endif
-}
-
-#define VDR_Q3_K_Q8_1_MMVQ 1
-#define VDR_Q3_K_Q8_1_MMQ 2
-
-// contiguous v/x values
-static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq(
-    const int& vl,
-    const int& vh,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ scales,
-    const int& scale_offset,
-    const float& d3,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  float sumf = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR3_K; ++i) {
-    const int isc = scale_offset + 2 * i;
-
-    const int isc_low = isc % (QK_K / 32);
-    const int sc_shift_low = 4 * (isc / (QK_K / 32));
-    const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF;
-
-    const int isc_high = isc % (QK_K / 64);
-    const int sc_shift_high = 2 * (isc / (QK_K / 64));
-    const int sc_high = ((scales[(QK_K / 32) + isc_high] >> sc_shift_high) & 3) << 4;
-
-    const int sc = (sc_low | sc_high) - 32;
-
-    const int vil = (vl >> (2 * i)) & 0x03030303;
-
-    const int vih = ((vh >> i) << 2) & 0x04040404;
-
-    const int vi = __vsubss4(vil, vih);
-
-    sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc);  // SIMD dot product
-  }
-
-  return d3 * sumf;
-#endif
-}
-
-static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const int8_t* __restrict__ scales,
-    const float& d3,
-    const float& d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  int sumi = 0;
-
-#pragma unroll
-  for (int i0 = 0; i0 < QR3_K * VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1 / 2) {
-    int sumi_sc = 0;
-
-    for (int i = i0; i < i0 + QI8_1 / 2; ++i) {
-      sumi_sc = __dp4a(v[i], u[i], sumi_sc);  // SIMD dot product
-    }
-
-    sumi += sumi_sc * scales[i0 / (QI8_1 / 2)];
-  }
-
-  return d3 * d8 * sumi;
-#endif
-}
-
-#define VDR_Q4_K_Q8_1_MMVQ 2
-#define VDR_Q4_K_Q8_1_MMQ 8
-
-// contiguous v/x values
-static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ sc,
-    const uint8_t* __restrict__ m,
-    const half2& dm4,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  float sumf_d = 0.0f;
-  float sumf_m = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR4_K; ++i) {
-    const int v0i = (v[0] >> (4 * i)) & 0x0F0F0F0F;
-    const int v1i = (v[1] >> (4 * i)) & 0x0F0F0F0F;
-
-    const int dot1 = __dp4a(v1i, u[2 * i + 1], __dp4a(v0i, u[2 * i + 0], 0));                // SIMD dot product
-    const int dot2 = __dp4a(0x01010101, u[2 * i + 1], __dp4a(0x01010101, u[2 * i + 0], 0));  // sum of u
-
-    sumf_d += d8[i] * (dot1 * sc[i]);
-    sumf_m += d8[i] * (dot2 * m[i]);  // multiply constant part of q4_K with sum of q8_1 values
-  }
-
-  const float2 dm4f = __half22float2(dm4);
-  return dm4f.x * sumf_d - dm4f.y * sumf_m;
-#endif
-}
-
-static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ sc,
-    const uint8_t* __restrict__ m,
-    const half2& dm4,
-    const half2* __restrict__ ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  float sumf_d = 0.0f;
-  float sumf_m = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR4_K * VDR_Q4_K_Q8_1_MMQ / QI8_1; ++i) {
-    int sumi_d = 0;
-
-#pragma unroll
-    for (int j = 0; j < QI8_1; ++j) {
-      sumi_d = __dp4a((v[j] >> (4 * i)) & 0x0F0F0F0F, u[i * QI8_1 + j], sumi_d);  // SIMD dot product
-    }
-
-    const float2 ds8f = __half22float2(ds8[i]);
-
-    sumf_d += ds8f.x * (sc[i] * sumi_d);
-    sumf_m += ds8f.y * m[i];  // sum of q8_1 block * q4_K min val
-  }
-
-  const float2 dm4f = __half22float2(dm4);
-
-  return dm4f.x * sumf_d - dm4f.y * sumf_m;
-#endif
-}
-
-#define VDR_Q5_K_Q8_1_MMVQ 2
-#define VDR_Q5_K_Q8_1_MMQ 8
-
-static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq(
-    const int* __restrict__ vl,
-    const int* __restrict__ vh,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ sc,
-    const uint8_t* __restrict__ m,
-    const half2& dm5,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  float sumf_d = 0.0f;
-  float sumf_m = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR5_K; ++i) {
-    const int vl0i = (vl[0] >> (4 * i)) & 0x0F0F0F0F;
-    const int vl1i = (vl[1] >> (4 * i)) & 0x0F0F0F0F;
-
-    const int vh0i = ((vh[0] >> i) << 4) & 0x10101010;
-    const int vh1i = ((vh[1] >> i) << 4) & 0x10101010;
-
-    const int v0i = vl0i | vh0i;
-    const int v1i = vl1i | vh1i;
-
-    const int dot1 = __dp4a(v0i, u[2 * i + 0], __dp4a(v1i, u[2 * i + 1], 0));                // SIMD dot product
-    const int dot2 = __dp4a(0x01010101, u[2 * i + 0], __dp4a(0x01010101, u[2 * i + 1], 0));  // sum of u
-
-    sumf_d += d8[i] * (dot1 * sc[i]);
-    sumf_m += d8[i] * (dot2 * m[i]);
-  }
-
-  const float2 dm5f = __half22float2(dm5);
-  return dm5f.x * sumf_d - dm5f.y * sumf_m;
-#endif
-}
-
-static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const uint8_t* __restrict__ sc,
-    const uint8_t* __restrict__ m,
-    const half2& dm4,
-    const half2* __restrict__ ds8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  float sumf_d = 0.0f;
-  float sumf_m = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR5_K * VDR_Q5_K_Q8_1_MMQ / QI8_1; ++i) {
-    int sumi_d = 0;
-
-#pragma unroll
-    for (int j = 0; j < QI8_1; ++j) {
-      sumi_d = __dp4a(v[i * QI8_1 + j], u[i * QI8_1 + j], sumi_d);  // SIMD dot product
-    }
-
-    const float2 ds8f = __half22float2(ds8[i]);
-
-    sumf_d += ds8f.x * (sc[i] * sumi_d);
-    sumf_m += ds8f.y * m[i];  // sum of q8_1 block * q4_K min val
-  }
-
-  const float2 dm4f = __half22float2(dm4);
-
-  return dm4f.x * sumf_d - dm4f.y * sumf_m;
-#endif
-}
-
-#define VDR_Q6_K_Q8_1_MMVQ 1
-#define VDR_Q6_K_Q8_1_MMQ 8
-
-// contiguous v/x values
-static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq(
-    const int& vl,
-    const int& vh,
-    const int* __restrict__ u,
-    const int8_t* __restrict__ scales,
-    const float& d,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  float sumf = 0.0f;
-
-#pragma unroll
-  for (int i = 0; i < QR6_K; ++i) {
-    const int sc = scales[4 * i];
-    const int vil = (vl >> (4 * i)) & 0x0F0F0F0F;
-    const int vih = ((vh >> (4 * i)) << 4) & 0x30303030;
-    const int vi = __vsubss4((vil | vih), 0x20202020);  // vi = (vil | vih) - 32
-
-    sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc);  // SIMD dot product
-  }
-
-  return d * sumf;
-#endif
-}
-
-static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq(
-    const int* __restrict__ v,
-    const int* __restrict__ u,
-    const int8_t* __restrict__ sc,
-    const float& d6,
-    const float* __restrict__ d8) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  float sumf_d = 0.0f;
-
-#pragma unroll
-  for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) {
-    int2 sumi_d = {0, 0};  // 2 q6_K scales per q8_1 scale
-
-#pragma unroll
-    for (int i = i0; i < i0 + 2; ++i) {
-      sumi_d.x = __dp4a(v[2 * i + 0], u[2 * i + 0], sumi_d.x);  // SIMD dot product
-      sumi_d.x = __dp4a(v[2 * i + 1], u[2 * i + 1], sumi_d.x);  // SIMD dot product
-
-      sumi_d.y = __dp4a(v[2 * i + 4], u[2 * i + 4], sumi_d.y);  // SIMD dot product
-      sumi_d.y = __dp4a(v[2 * i + 5], u[2 * i + 5], sumi_d.y);  // SIMD dot product
-    }
-
-    sumf_d += d8[i0 / 4] * (sc[i0 / 2 + 0] * sumi_d.x + sc[i0 / 2 + 1] * sumi_d.y);
-  }
-
-  return d6 * sumf_d;
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_q4_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q4_0* bq4_0 = (const block_q4_0*)vbq;
-
-  int v[VDR_Q4_0_Q8_1_MMVQ];
-  int u[2 * VDR_Q4_0_Q8_1_MMVQ];
-
-#pragma unroll
-  for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) {
-    v[i] = get_int_from_uint8(bq4_0->qs, iqs + i);
-    u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i);
-    u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0);
-  }
-
-  return vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMVQ>(v, u, __half2float(bq4_0->d), bq8_1->ds);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q4_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y];
-  __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI4_0) + mmq_y / QI4_0];
-  *x_ql = tile_x_qs;
-  *x_dm = (half2*)tile_x_d;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q4_0(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI4_0;
-  const int kqsx = k % QI4_0;
-
-  const block_q4_0* bx0 = (const block_q4_0*)vx;
-  float* x_dmf = (float*)x_dm;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_0* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx);
-    // x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i / QI4_0 + kbx] = bxi->d;
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_0;
-  const int kbxd = k % blocks_per_tile_x_row;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_0) {
-    int i = i0 + i_offset * QI4_0 + k / blocks_per_tile_x_row;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_0* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + kbxd] = __half2float(bxi->d);
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q4_0_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  (void)x_qh;
-  (void)x_sc;
-
-  const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2));
-  const float* x_dmf = (const float*)x_dm;
-
-  int u[2 * VDR_Q4_0_Q8_1_MMQ];
-
-#pragma unroll
-  for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) {
-    u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF];
-    u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_0) % WARP_SIZE_GGUF];
-  }
-
-  return vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMQ>(
-      &x_ql[i * (WARP_SIZE_GGUF + 1) + k],
-      u,
-      x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + k / QI4_0],
-      y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q4_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q4_1* bq4_1 = (const block_q4_1*)vbq;
-
-  int v[VDR_Q4_1_Q8_1_MMVQ];
-  int u[2 * VDR_Q4_1_Q8_1_MMVQ];
-
-#pragma unroll
-  for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) {
-    v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i);
-    u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i);
-    u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1);
-  }
-
-  return vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMVQ>(v, u, bq4_1->dm, bq8_1->ds);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q4_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + +mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_1) + mmq_y / QI4_1];
-  *x_ql = tile_x_qs;
-  *x_dm = tile_x_dm;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q4_1(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI4_1;
-  const int kqsx = k % QI4_1;
-
-  const block_q4_1* bx0 = (const block_q4_1*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_1* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_1;
-  const int kbxd = k % blocks_per_tile_x_row;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_1) {
-    int i = i0 + i_offset * QI4_1 + k / blocks_per_tile_x_row;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_1* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + kbxd] = bxi->dm;
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q4_1_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2));
-
-  int u[2 * VDR_Q4_1_Q8_1_MMQ];
-
-#pragma unroll
-  for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) {
-    u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF];
-    u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_1) % WARP_SIZE_GGUF];
-  }
-
-  return vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMQ>(
-      &x_ql[i * (WARP_SIZE_GGUF + 1) + k],
-      u,
-      x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + k / QI4_1],
-      y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q5_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q5_0* bq5_0 = (const block_q5_0*)vbq;
-
-  int vl[VDR_Q5_0_Q8_1_MMVQ];
-  int vh[VDR_Q5_0_Q8_1_MMVQ];
-  int u[2 * VDR_Q5_0_Q8_1_MMVQ];
-
-#pragma unroll
-  for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) {
-    vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i);
-    vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i));
-    u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i);
-    u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0);
-  }
-
-  return vec_dot_q5_0_q8_1_impl<VDR_Q5_0_Q8_1_MMVQ>(vl, vh, u, __half2float(bq5_0->d), bq8_1->ds);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q5_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y];
-  __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI5_0) + mmq_y / QI5_0];
-
-  *x_ql = tile_x_ql;
-  *x_dm = (half2*)tile_x_d;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q5_0(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI5_0;
-  const int kqsx = k % QI5_0;
-
-  const block_q5_0* bx0 = (const block_q5_0*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q5_0* bxi = bx0 + i * blocks_per_row + kbx;
-    const int ql = get_int_from_uint8(bxi->qs, kqsx);
-    const int qh = get_int_from_uint8(bxi->qh, 0) >> (4 * (k % QI5_0));
-
-    int qs0 = (ql >> 0) & 0x0F0F0F0F;
-    qs0 |= (qh << 4) & 0x00000010;     // 0 ->  4
-    qs0 |= (qh << 11) & 0x00001000;    // 1 -> 12
-    qs0 |= (qh << 18) & 0x00100000;    // 2 -> 20
-    qs0 |= (qh << 25) & 0x10000000;    // 3 -> 28
-    qs0 = __vsubss4(qs0, 0x10101010);  // subtract 16
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0;
-
-    int qs1 = (ql >> 4) & 0x0F0F0F0F;
-    qs1 |= (qh >> 12) & 0x00000010;    // 16 ->  4
-    qs1 |= (qh >> 5) & 0x00001000;     // 17 -> 12
-    qs1 |= (qh << 2) & 0x00100000;     // 18 -> 20
-    qs1 |= (qh << 9) & 0x10000000;     // 19 -> 28
-    qs1 = __vsubss4(qs1, 0x10101010);  // subtract 16
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1;
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_0;
-  const int kbxd = k % blocks_per_tile_x_row;
-  float* x_dmf = (float*)x_dm;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_0) {
-    int i = i0 + i_offset * QI5_0 + k / blocks_per_tile_x_row;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_0* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dmf[i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + kbxd] = __half2float(bxi->d);
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q5_0_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2));
-  const int index_bx = i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + k / QI5_0;
-  const float* x_dmf = (const float*)x_dm;
-  const float* y_df = (const float*)y_ds;
-
-  int u[2 * VDR_Q5_0_Q8_1_MMQ];
-
-#pragma unroll
-  for (int l = 0; l < VDR_Q5_0_Q8_1_MMQ; ++l) {
-    u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF];
-    u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_0) % WARP_SIZE_GGUF];
-  }
-
-  return vec_dot_q8_0_q8_1_impl<QR5_0 * VDR_Q5_0_Q8_1_MMQ>(
-      &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k],
-      u,
-      x_dmf[index_bx],
-      y_df[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q5_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q5_1* bq5_1 = (const block_q5_1*)vbq;
-
-  int vl[VDR_Q5_1_Q8_1_MMVQ];
-  int vh[VDR_Q5_1_Q8_1_MMVQ];
-  int u[2 * VDR_Q5_1_Q8_1_MMVQ];
-
-#pragma unroll
-  for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) {
-    vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i);
-    vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i));
-    u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i);
-    u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1);
-  }
-
-  return vec_dot_q5_1_q8_1_impl<VDR_Q5_1_Q8_1_MMVQ>(vl, vh, u, bq5_1->dm, bq8_1->ds);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q5_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_1) + mmq_y / QI5_1];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q5_1(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI5_1;
-  const int kqsx = k % QI5_1;
-
-  const block_q5_1* bx0 = (const block_q5_1*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_1* bxi = bx0 + i * blocks_per_row + kbx;
-
-    const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx);
-    const int qh = get_int_from_uint8_aligned(bxi->qh, 0) >> (4 * (k % QI5_1));
-
-    int qs0 = (ql >> 0) & 0x0F0F0F0F;
-    qs0 |= (qh << 4) & 0x00000010;   // 0 ->  4
-    qs0 |= (qh << 11) & 0x00001000;  // 1 -> 12
-    qs0 |= (qh << 18) & 0x00100000;  // 2 -> 20
-    qs0 |= (qh << 25) & 0x10000000;  // 3 -> 28
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0;
-
-    int qs1 = (ql >> 4) & 0x0F0F0F0F;
-    qs1 |= (qh >> 12) & 0x00000010;  // 16 ->  4
-    qs1 |= (qh >> 5) & 0x00001000;   // 17 -> 12
-    qs1 |= (qh << 2) & 0x00100000;   // 18 -> 20
-    qs1 |= (qh << 9) & 0x10000000;   // 19 -> 28
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1;
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_1;
-  const int kbxd = k % blocks_per_tile_x_row;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_1) {
-    int i = i0 + i_offset * QI5_1 + k / blocks_per_tile_x_row;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_1* bxi = bx0 + i * blocks_per_row + kbxd;
-
-    x_dm[i * (WARP_SIZE_GGUF / QI5_1) + i / QI5_1 + kbxd] = bxi->dm;
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q5_1_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2));
-  const int index_bx = i * (WARP_SIZE_GGUF / QI5_1) + +i / QI5_1 + k / QI5_1;
-
-  int u[2 * VDR_Q5_1_Q8_1_MMQ];
-
-#pragma unroll
-  for (int l = 0; l < VDR_Q5_1_Q8_1_MMQ; ++l) {
-    u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF];
-    u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_1) % WARP_SIZE_GGUF];
-  }
-
-  return vec_dot_q8_1_q8_1_impl<QR5_1 * VDR_Q5_1_Q8_1_MMQ>(
-      &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k],
-      u,
-      x_dm[index_bx],
-      y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q8_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q8_0* bq8_0 = (const block_q8_0*)vbq;
-
-  int v[VDR_Q8_0_Q8_1_MMVQ];
-  int u[VDR_Q8_0_Q8_1_MMVQ];
-
-#pragma unroll
-  for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) {
-    v[i] = get_int_from_int8(bq8_0->qs, iqs + i);
-    u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i);
-  }
-
-  return vec_dot_q8_0_q8_1_impl<VDR_Q8_0_Q8_1_MMVQ>(v, u, __half2float(bq8_0->d), __low2float(bq8_1->ds));
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q8_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y];
-  __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI8_0) + mmq_y / QI8_0];
-
-  *x_ql = tile_x_qs;
-  *x_dm = (half2*)tile_x_d;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q8_0(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI8_0;
-  const int kqsx = k % QI8_0;
-  float* x_dmf = (float*)x_dm;
-
-  const block_q8_0* bx0 = (const block_q8_0*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q8_0* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_int8(bxi->qs, kqsx);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI8_0;
-  const int kbxd = k % blocks_per_tile_x_row;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI8_0) {
-    int i = i0 + i_offset * QI8_0 + k / blocks_per_tile_x_row;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q8_0* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + kbxd] = __half2float(bxi->d);
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q8_0_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const float* x_dmf = (const float*)x_dm;
-  const float* y_df = (const float*)y_ds;
-
-  return vec_dot_q8_0_q8_1_impl<VDR_Q8_0_Q8_1_MMQ>(
-      &x_ql[i * (WARP_SIZE_GGUF + 1) + k],
-      &y_qs[j * WARP_SIZE_GGUF + k],
-      x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + k / QI8_0],
-      y_df[j * (WARP_SIZE_GGUF / QI8_1) + k / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q2_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q2_K* bq2_K = (const block_q2_K*)vbq;
-
-  const int bq8_offset = QR2_K * (iqs / QI8_1);
-  const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2);
-
-  const uint8_t* scales = bq2_K->scales + scale_offset;
-
-  const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs);
-  int u[QR2_K];
-  float d8[QR2_K];
-
-#pragma unroll
-  for (int i = 0; i < QR2_K; ++i) {
-    u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
-    d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
-  }
-
-  return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q2_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI2_K) + mmq_y / QI2_K];
-  __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-  *x_sc = tile_x_sc;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q2_K(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI2_K;
-  const int kqsx = k % QI2_K;
-
-  const block_q2_K* bx0 = (const block_q2_K*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q2_K* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI2_K;
-  const int kbxd = k % blocks_per_tile_x_row;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI2_K) {
-    int i = (i0 + i_offset * QI2_K + k / blocks_per_tile_x_row) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q2_K* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbxd] = bxi->dm;
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) {
-    int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4);
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q2_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI2_K / 4);
-    x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] =
-        get_int_from_uint8_aligned(bxi->scales, k % (QI2_K / 4));
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q2_K_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const int kbx = k / QI2_K;
-  const int ky = (k % QI2_K) * QR2_K;
-  const float* y_df = (const float*)y_ds;
-
-  int v[QR2_K * VDR_Q2_K_Q8_1_MMQ];
-
-  const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI2_K + (QI2_K / 2) * (ky / (2 * QI2_K)) + ky % (QI2_K / 2);
-  const int shift = 2 * ((ky % (2 * QI2_K)) / (QI2_K / 2));
-
-#pragma unroll
-  for (int l = 0; l < QR2_K * VDR_Q2_K_Q8_1_MMQ; ++l) {
-    v[l] = (x_ql[kqsx + l] >> shift) & 0x03030303;
-  }
-
-  const uint8_t* scales = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4]) + ky / 4;
-
-  const int index_y = j * WARP_SIZE_GGUF + (QR2_K * k) % WARP_SIZE_GGUF;
-  return vec_dot_q2_K_q8_1_impl_mmq(
-      v, &y_qs[index_y], scales, x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbx], y_df[index_y / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q3_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q3_K* bq3_K = (const block_q3_K*)vbq;
-
-  const int bq8_offset = QR3_K * (iqs / (QI3_K / 2));
-  const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2);
-
-  const float d = __half2float(bq3_K->d);
-
-  const int vl = get_int_from_uint8(bq3_K->qs, iqs);
-
-  // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted
-  const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K / 2)) >> bq8_offset;
-
-  int u[QR3_K];
-  float d8[QR3_K];
-
-#pragma unroll
-  for (int i = 0; i < QR3_K; ++i) {
-    u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
-    d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
-  }
-
-  return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q3_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI3_K) + mmq_y / QI3_K];
-  __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 2) + mmq_y / 2];
-  __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-  *x_qh = tile_x_qh;
-  *x_sc = tile_x_sc;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q3_K(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI3_K;
-  const int kqsx = k % QI3_K;
-
-  const block_q3_K* bx0 = (const block_q3_K*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q3_K* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI3_K;
-  const int kbxd = k % blocks_per_tile_x_row;
-  float* x_dmf = (float*)x_dm;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI3_K) {
-    int i = (i0 + i_offset * QI3_K + k / blocks_per_tile_x_row) % mmq_y;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q3_K* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbxd] = __half2float(bxi->d);
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 2) {
-    int i = i0 + i_offset * 2 + k / (WARP_SIZE_GGUF / 2);
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 2)) / (QI3_K / 2);
-    // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted
-    x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + k % (WARP_SIZE_GGUF / 2)] =
-        ~get_int_from_uint8(bxi->hmask, k % (QI3_K / 2));
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) {
-    int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4);
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI3_K / 4);
-
-    const int ksc = k % (QI3_K / 4);
-
-    const int ksc_low = ksc % (QI3_K / 8);
-    const int shift_low = 4 * (ksc / (QI3_K / 8));
-    const int sc_low = (get_int_from_uint8(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F;
-
-    const int ksc_high = QI3_K / 8;
-    const int shift_high = 2 * ksc;
-    const int sc_high = ((get_int_from_uint8(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;
-
-    const int sc = __vsubss4(sc_low | sc_high, 0x20202020);
-
-    x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] = sc;
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q3_K_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const int kbx = k / QI3_K;
-  const int ky = (k % QI3_K) * QR3_K;
-  const float* x_dmf = (const float*)x_dm;
-  const float* y_df = (const float*)y_ds;
-
-  const int8_t* scales = ((const int8_t*)(x_sc + i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4)) + ky / 4;
-
-  int v[QR3_K * VDR_Q3_K_Q8_1_MMQ];
-
-#pragma unroll
-  for (int l = 0; l < QR3_K * VDR_Q3_K_Q8_1_MMQ; ++l) {
-    const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI3_K + (QI3_K / 2) * (ky / (2 * QI3_K)) + ky % (QI3_K / 2);
-    const int shift = 2 * ((ky % 32) / 8);
-    const int vll = (x_ql[kqsx + l] >> shift) & 0x03030303;
-
-    const int vh = x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + kbx * (QI3_K / 2) + (ky + l) % 8] >> ((ky + l) / 8);
-    const int vlh = (vh << 2) & 0x04040404;
-
-    v[l] = __vsubss4(vll, vlh);
-  }
-
-  const int index_y = j * WARP_SIZE_GGUF + (k * QR3_K) % WARP_SIZE_GGUF;
-  return vec_dot_q3_K_q8_1_impl_mmq(
-      v, &y_qs[index_y], scales, x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbx], y_df[index_y / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q4_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q4_K* bq4_K = (const block_q4_K*)vbq;
-
-  int v[2];
-  int u[2 * QR4_K];
-  float d8[QR4_K];
-
-  // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6
-  const int bq8_offset = QR4_K * ((iqs / 2) / (QI8_1 / 2));
-
-  // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12
-  // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44
-  // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76
-  // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108
-
-  const int* q4 = (const int*)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4));
-  v[0] = q4[0];
-  v[1] = q4[4];
-
-  const uint16_t* scales = (const uint16_t*)bq4_K->scales;
-  uint16_t aux[2];
-  const int j = bq8_offset / 2;
-  if (j < 2) {
-    aux[0] = scales[j + 0] & 0x3f3f;
-    aux[1] = scales[j + 2] & 0x3f3f;
-  } else {
-    aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2);
-    aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2);
-  }
-  const uint8_t* sc = (const uint8_t*)aux;
-  const uint8_t* m = sc + 2;
-
-  for (int i = 0; i < QR4_K; ++i) {
-    const block_q8_1* bq8i = bq8_1 + bq8_offset + i;
-    d8[i] = __low2float(bq8i->ds);
-
-    const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4);
-    u[2 * i + 0] = q8[0];
-    u[2 * i + 1] = q8[4];
-  }
-
-  return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q4_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_K) + mmq_y / QI4_K];
-  __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-  *x_sc = tile_x_sc;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q4_K(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI4_K;   // == 0 if QK_K == 256
-  const int kqsx = k % QI4_K;  // == k if QK_K == 256
-
-  const block_q4_K* bx0 = (const block_q4_K*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_K* bxi = bx0 + i * blocks_per_row + kbx;
-    x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_K;  // == 1 if QK_K == 256
-  const int kbxd = k % blocks_per_tile_x_row;                // == 0 if QK_K == 256
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_K) {
-    int i = (i0 + i_offset * QI4_K + k / blocks_per_tile_x_row) % mmq_y;
-    if (need_check) {
-      i = min(i, i_max);
-    }
-    const block_q4_K* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K + kbxd] = bxi->dm;
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) {
-    int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q4_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI4_K / 8);
-
-    const int* scales = (const int*)bxi->scales;
-
-    const int ksc = k % (WARP_SIZE_GGUF / 8);
-    // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8
-    int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F;  // lower 4 bits
-    scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030;                            // upper 2 bits
-
-    x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8;
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q4_K_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  (void)x_qh;
-
-  const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8);
-
-  const int index_y = j * WARP_SIZE_GGUF + (QR4_K * k) % WARP_SIZE_GGUF;
-  return vec_dot_q4_K_q8_1_impl_mmq(
-      &x_ql[i * (WARP_SIZE_GGUF + 1) + k],
-      &y_qs[index_y],
-      sc,
-      sc + 8,
-      x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K],
-      &y_ds[index_y / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q5_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q5_K* bq5_K = (const block_q5_K*)vbq;
-
-  int vl[2];
-  int vh[2];
-  int u[2 * QR5_K];
-  float d8[QR5_K];
-
-  const int bq8_offset = QR5_K * ((iqs / 2) / (QI8_1 / 2));
-  const int* ql = (const int*)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4));
-  const int* qh = (const int*)(bq5_K->qh + 4 * ((iqs / 2) % 4));
-
-  vl[0] = ql[0];
-  vl[1] = ql[4];
-
-  vh[0] = qh[0] >> bq8_offset;
-  vh[1] = qh[4] >> bq8_offset;
-
-  const uint16_t* scales = (const uint16_t*)bq5_K->scales;
-  uint16_t aux[2];
-  const int j = bq8_offset / 2;
-  if (j < 2) {
-    aux[0] = scales[j + 0] & 0x3f3f;
-    aux[1] = scales[j + 2] & 0x3f3f;
-  } else {
-    aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2);
-    aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2);
-  }
-  const uint8_t* sc = (const uint8_t*)aux;
-  const uint8_t* m = sc + 2;
-
-#pragma unroll
-  for (int i = 0; i < QR5_K; ++i) {
-    const block_q8_1* bq8i = bq8_1 + bq8_offset + i;
-    d8[i] = __low2float(bq8i->ds);
-
-    const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4);
-    u[2 * i + 0] = q8[0];
-    u[2 * i + 1] = q8[4];
-  }
-
-  return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q5_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_K) + mmq_y / QI5_K];
-  __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-  *x_sc = tile_x_sc;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q5_K(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI5_K;   // == 0 if QK_K == 256
-  const int kqsx = k % QI5_K;  // == k if QK_K == 256
-
-  const block_q5_K* bx0 = (const block_q5_K*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_K* bxi = bx0 + i * blocks_per_row + kbx;
-    const int ky = QR5_K * kqsx;
-
-    const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx);
-    const int ql0 = (ql >> 0) & 0x0F0F0F0F;
-    const int ql1 = (ql >> 4) & 0x0F0F0F0F;
-
-    const int qh = get_int_from_uint8_aligned(bxi->qh, kqsx % (QI5_K / 4));
-    const int qh0 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 0)) << 4) & 0x10101010;
-    const int qh1 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 1)) << 4) & 0x10101010;
-
-    const int kq0 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + 0;
-    const int kq1 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + (QI5_K / 4);
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = ql0 | qh0;
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = ql1 | qh1;
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_K;  // == 1 if QK_K == 256
-  const int kbxd = k % blocks_per_tile_x_row;                // == 0 if QK_K == 256
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_K) {
-    int i = (i0 + i_offset * QI5_K + k / blocks_per_tile_x_row) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_K* bxi = bx0 + i * blocks_per_row + kbxd;
-    x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K + kbxd] = bxi->dm;
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) {
-    int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q5_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI5_K / 8);
-
-    const int* scales = (const int*)bxi->scales;
-
-    const int ksc = k % (WARP_SIZE_GGUF / 8);
-
-    // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8
-    int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F;  // lower 4 bits
-    scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030;                            // upper 2 bits
-
-    x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8;
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q5_K_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8);
-
-  const int index_x = i * (QR5_K * WARP_SIZE_GGUF + 1) + QR5_K * k;
-  const int index_y = j * WARP_SIZE_GGUF + (QR5_K * k) % WARP_SIZE_GGUF;
-  return vec_dot_q5_K_q8_1_impl_mmq(
-      &x_ql[index_x],
-      &y_qs[index_y],
-      sc,
-      sc + 8,
-      x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K],
-      &y_ds[index_y / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_q6_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_q6_K* bq6_K = (const block_q6_K*)vbq;
-
-  const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 4);
-  const int scale_offset = (QI6_K / 4) * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 8);
-  const int vh_shift = 2 * ((iqs % (QI6_K / 2)) / (QI6_K / 4));
-
-  const int vl = get_int_from_uint8(bq6_K->ql, iqs);
-  const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K / 4) * (iqs / (QI6_K / 2)) + iqs % (QI6_K / 4)) >> vh_shift;
-
-  const int8_t* scales = bq6_K->scales + scale_offset;
-
-  int u[QR6_K];
-  float d8[QR6_K];
-
-#pragma unroll
-  for (int i = 0; i < QR6_K; ++i) {
-    u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2 * i].qs, iqs % QI8_1);
-    d8[i] = __low2float(bq8_1[bq8_offset + 2 * i].ds);
-  }
-
-  return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, __half2float(bq6_K->d), d8);
-}
-
-template <int mmq_y>
-static __device__ __forceinline__ void allocate_tiles_q6_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) {
-  __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y];
-  __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI6_K) + mmq_y / QI6_K];
-  __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8];
-
-  *x_ql = tile_x_ql;
-  *x_dm = tile_x_dm;
-  *x_sc = tile_x_sc;
-}
-
-template <int mmq_y, int nwarps, bool need_check>
-static __device__ __forceinline__ void load_tiles_q6_K(
-    const void* __restrict__ vx,
-    int* __restrict__ x_ql,
-    half2* __restrict__ x_dm,
-    int* __restrict__ x_qh,
-    int* __restrict__ x_sc,
-    const int& i_offset,
-    const int& i_max,
-    const int& k,
-    const int& blocks_per_row) {
-  const int kbx = k / QI6_K;   // == 0 if QK_K == 256
-  const int kqsx = k % QI6_K;  // == k if QK_K == 256
-
-  const block_q6_K* bx0 = (const block_q6_K*)vx;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
-    int i = i0 + i_offset;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q6_K* bxi = bx0 + i * blocks_per_row + kbx;
-    const int ky = QR6_K * kqsx;
-
-    const int ql = get_int_from_uint8(bxi->ql, kqsx);
-    const int ql0 = (ql >> 0) & 0x0F0F0F0F;
-    const int ql1 = (ql >> 4) & 0x0F0F0F0F;
-
-    const int qh = get_int_from_uint8(bxi->qh, (QI6_K / 4) * (kqsx / (QI6_K / 2)) + kqsx % (QI6_K / 4));
-    const int qh0 = ((qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) << 4) & 0x30303030;
-    const int qh1 = (qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) & 0x30303030;
-
-    const int kq0 = ky - ky % QI6_K + k % (QI6_K / 2) + 0;
-    const int kq1 = ky - ky % QI6_K + k % (QI6_K / 2) + (QI6_K / 2);
-
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
-    x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
-  }
-
-  const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI6_K;  // == 1 if QK_K == 256
-  const int kbxd = k % blocks_per_tile_x_row;                // == 0 if QK_K == 256
-  float* x_dmf = (float*)x_dm;
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI6_K) {
-    int i = (i0 + i_offset * QI6_K + k / blocks_per_tile_x_row) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q6_K* bxi = bx0 + i * blocks_per_row + kbxd;
-
-    x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K + kbxd] = __half2float(bxi->d);
-  }
-
-#pragma unroll
-  for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) {
-    int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y;
-
-    if (need_check) {
-      i = min(i, i_max);
-    }
-
-    const block_q6_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / 4;
-
-    x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k % (WARP_SIZE_GGUF / 8)] = get_int_from_int8(bxi->scales, k % (QI6_K / 8));
-  }
-}
-
-static __device__ __forceinline__ float vec_dot_q6_K_q8_1_mul_mat(
-    const int* __restrict__ x_ql,
-    const half2* __restrict__ x_dm,
-    const int* __restrict__ x_qh,
-    const int* __restrict__ x_sc,
-    const int* __restrict__ y_qs,
-    const half2* __restrict__ y_ds,
-    const int& i,
-    const int& j,
-    const int& k) {
-  const float* x_dmf = (const float*)x_dm;
-  const float* y_df = (const float*)y_ds;
-
-  const int8_t* sc = ((const int8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 8]);
-
-  const int index_x = i * (QR6_K * WARP_SIZE_GGUF + 1) + QR6_K * k;
-  const int index_y = j * WARP_SIZE_GGUF + (QR6_K * k) % WARP_SIZE_GGUF;
-  return vec_dot_q6_K_q8_1_impl_mmq(
-      &x_ql[index_x], &y_qs[index_y], sc, x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K], &y_df[index_y / QI8_1]);
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_iq2_xxs* bq2 = (const block_iq2_xxs*)vbq;
-
-  const int ib32 = iqs;
-  const uint16_t* q2 = bq2->qs + 4 * ib32;
-  const uint8_t* aux8 = (const uint8_t*)q2;
-  const int8_t* q8 = bq8_1[ib32].qs;
-  uint32_t aux32 = q2[2] | (q2[3] << 16);
-  int sumi = 0;
-  for (int l = 0; l < 4; ++l) {
-    const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]);
-    const uint8_t signs = ksigns_iq2xs[aux32 & 127];
-    for (int j = 0; j < 8; ++j) {
-      sumi += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
-    }
-    q8 += 8;
-    aux32 >>= 7;
-  }
-  const float d = __half2float(bq2->d) * (0.5f + aux32) * __half2float(bq8_1[ib32].ds.x) * 0.25f;
-  return d * sumi;
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-  const block_iq2_xs* bq2 = (const block_iq2_xs*)vbq;
-
-  const int ib32 = iqs;
-  const uint16_t* q2 = bq2->qs + 4 * ib32;
-  const int8_t* q8 = bq8_1[ib32].qs;
-  const uint8_t ls1 = bq2->scales[ib32] & 0xf;
-  const uint8_t ls2 = bq2->scales[ib32] >> 4;
-  int sumi1 = 0;
-  for (int l = 0; l < 2; ++l) {
-    const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511));
-    const uint8_t signs = ksigns_iq2xs[q2[l] >> 9];
-    for (int j = 0; j < 8; ++j) {
-      sumi1 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
-    }
-    q8 += 8;
-  }
-  int sumi2 = 0;
-  for (int l = 2; l < 4; ++l) {
-    const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511));
-    const uint8_t signs = ksigns_iq2xs[q2[l] >> 9];
-    for (int j = 0; j < 8; ++j) {
-      sumi2 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
-    }
-    q8 += 8;
-  }
-  const float d = __half2float(bq2->d) * __half2float(bq8_1[ib32].ds.x) * 0.25f;
-  return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2);
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  const block_iq2_s* bq2 = (const block_iq2_s*)vbq;
-
-  const int ib32 = iqs;
-  const int8_t* q8 = bq8_1[ib32].qs;
-  const uint8_t* signs = bq2->qs + QK_K / 8 + 4 * ib32;
-  const uint8_t ls1 = bq2->scales[ib32] & 0xf;
-  const uint8_t ls2 = bq2->scales[ib32] >> 4;
-  int sumi1 = 0;
-  for (int l = 0; l < 2; ++l) {
-    const uint32_t* grid =
-        (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300)));
-    const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201);
-    const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201);
-    const int grid_l = __vsub4(grid[0] ^ signs0, signs0);
-    const int grid_h = __vsub4(grid[1] ^ signs1, signs1);
-    sumi1 = __dp4a(grid_l, *((const int*)q8 + 0), sumi1);
-    sumi1 = __dp4a(grid_h, *((const int*)q8 + 1), sumi1);
-    q8 += 8;
-  }
-  int sumi2 = 0;
-  for (int l = 2; l < 4; ++l) {
-    const uint32_t* grid =
-        (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300)));
-    const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201);
-    const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201);
-    const int grid_l = __vsub4(grid[0] ^ signs0, signs0);
-    const int grid_h = __vsub4(grid[1] ^ signs1, signs1);
-    sumi2 = __dp4a(grid_l, *((const int*)q8 + 0), sumi2);
-    sumi2 = __dp4a(grid_h, *((const int*)q8 + 1), sumi2);
-    q8 += 8;
-  }
-  const float d = __half2float(bq2->d) * __low2float(bq8_1[ib32].ds) * 0.25f;
-  return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2);
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  const block_iq3_xxs* bq2 = (const block_iq3_xxs*)vbq;
-
-  const int ib32 = iqs;
-  const uint8_t* q3 = bq2->qs + 8 * ib32;
-  const uint16_t* gas = (const uint16_t*)(bq2->qs + QK_K / 4) + 2 * ib32;
-  const int8_t* q8 = bq8_1[ib32].qs;
-  uint32_t aux32 = gas[0] | (gas[1] << 16);
-  int sumi = 0;
-  for (int l = 0; l < 4; ++l) {
-    const uint32_t* grid1 = iq3xxs_grid + q3[2 * l + 0];
-    const uint32_t* grid2 = iq3xxs_grid + q3[2 * l + 1];
-    const uint32_t* signs = (const uint32_t*)(ksigns64 + (aux32 & 127));
-    const int grid_l = __vsub4(grid1[0] ^ signs[0], signs[0]);
-    const int grid_h = __vsub4(grid2[0] ^ signs[1], signs[1]);
-    sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi);
-    sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi);
-    q8 += 8;
-    aux32 >>= 7;
-  }
-  const float d = __half2float(bq2->d) * (0.5f + aux32) * __low2float(bq8_1[ib32].ds) * 0.5f;
-  return d * sumi;
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  const block_iq3_s* bq2 = (const block_iq3_s*)vbq;
-
-  const int ib32 = iqs;
-  const uint8_t* qs = bq2->qs + 8 * ib32;
-  const int8_t* q8 = bq8_1[ib32].qs;
-  int sumi = 0;
-  for (int l = 0; l < 4; ++l) {
-    const uint32_t* grid1 = iq3xs_grid + (qs[2 * l + 0] | ((bq2->qh[ib32] << (8 - 2 * l)) & 256));
-    const uint32_t* grid2 = iq3xs_grid + (qs[2 * l + 1] | ((bq2->qh[ib32] << (7 - 2 * l)) & 256));
-    uint32_t signs0 = __vcmpeq4(((bq2->signs[4 * ib32 + l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201);
-    uint32_t signs1 = __vcmpeq4(((bq2->signs[4 * ib32 + l] >> 4) * 0x01010101) & 0x08040201, 0x08040201);
-    const int grid_l = __vsub4(grid1[0] ^ signs0, signs0);
-    const int grid_h = __vsub4(grid2[0] ^ signs1, signs1);
-    sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi);
-    sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi);
-    q8 += 8;
-  }
-  const float d = __half2float(bq2->d) * (0.5f + ((bq2->scales[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf)) *
-                  __low2float(bq8_1[ib32].ds) * 0.5f;
-  return d * sumi;
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq1_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  const block_iq1_s* bq1 = (const block_iq1_s*)vbq;
-
-  const int qs_packed = get_int_b2(bq1->qs, iqs);
-  const uint8_t* qs = (const uint8_t*)&qs_packed;
-
-  const int qh = bq1->qh[iqs];
-
-  int sumi = 0;
-#pragma unroll
-  for (int l0 = 0; l0 < 8; l0 += 2) {
-    const int grid = iq1s_grid_gpu[qs[l0 / 2] | (((qh >> 3 * (l0 / 2)) & 0x07) << 8)];
-
-    const int grid0 = (grid >> 0) & 0x0F0F0F0F;
-    const int grid1 = (grid >> 4) & 0x0F0F0F0F;
-
-    const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0);
-    const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1);
-
-    sumi = __dp4a(grid0, u0, sumi);
-    sumi = __dp4a(grid1, u1, sumi);
-  }
-
-  const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1);
-  const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f * IQ1S_DELTA / 0x8000);
-  const float2 ds = __half22float2(bq8_1[iqs].ds);
-  return d1q * (ds.x * sumi + ds.y * delta);
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq1_m_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  const block_iq1_m* bq1 = (const block_iq1_m*)vbq;
-
-  const int qs_packed = get_int_b4(bq1->qs, iqs);
-  const uint8_t* qs = (const uint8_t*)&qs_packed;
-
-  int sumi[2] = {0};
-  float sumf[2] = {0.0f};
-#pragma unroll
-  for (int l0 = 0; l0 < 8; l0 += 2) {
-    const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2));
-
-    const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)];
-
-    const int grid0 = (grid >> 0) & 0x0F0F0F0F;
-    const int grid1 = (grid >> 4) & 0x0F0F0F0F;
-
-    const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0);
-    const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1);
-
-    sumi[l0 / 4] = __dp4a(grid0, u0, sumi[l0 / 4]);
-    sumi[l0 / 4] = __dp4a(grid1, u1, sumi[l0 / 4]);
-
-    const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08);
-    int sumy = 0;
-    sumy = __dp4a(u0, 0x01010101, sumy);
-    sumy = __dp4a(u1, 0x01010101, sumy);
-    sumf[l0 / 4] += delta * sumy;
-  }
-
-  const uint16_t* sc = (const uint16_t*)bq1->scales;
-
-  iq1m_scale_t scale;
-  scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000);
-  const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds);
-
-  const int tmp = sc[iqs / 2] >> (6 * (iqs % 2));
-  const int sc0 = 2 * ((tmp >> 0) & 0x07) + 1;
-  const int sc1 = 2 * ((tmp >> 3) & 0x07) + 1;
-  return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1);
-#endif
-}
-
-static __device__ __forceinline__ void
-get_int_from_table_16(const uint32_t& q4, const uint8_t* values, int& val1, int& val2) {
-  uint32_t aux32;
-  const uint8_t* q8 = (const uint8_t*)&aux32;
-  aux32 = q4 & 0x0f0f0f0f;
-  uint16_t v1 = values[q8[0]] | (values[q8[1]] << 8);
-  uint16_t v2 = values[q8[2]] | (values[q8[3]] << 8);
-  val1 = v1 | (v2 << 16);
-  aux32 = (q4 >> 4) & 0x0f0f0f0f;
-  v1 = values[q8[0]] | (values[q8[1]] << 8);
-  v2 = values[q8[2]] | (values[q8[3]] << 8);
-  val2 = v1 | (v2 << 16);
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-
-  const block_iq4_nl* bq = (const block_iq4_nl*)vbq;
-
-  const uint16_t* q4 = (const uint16_t*)bq->qs + 2 * iqs;
-  const int32_t* q8 = (const int32_t*)bq8_1->qs + iqs;
-
-  const uint8_t* values = (const uint8_t*)kvalues_iq4nl;
-
-  int v1, v2;
-  int sumi1 = 0, sumi2 = 0;
-  for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) {
-    const uint32_t aux = q4[2 * l] | (q4[2 * l + 1] << 16);
-    get_int_from_table_16(aux, values, v1, v2);
-    sumi1 = __dp4a(v1, q8[l + 0], sumi1);
-    sumi2 = __dp4a(v2, q8[l + 4], sumi2);
-  }
-  const float d = __half2float(bq->d) * __low2float(bq8_1->ds);
-  return d * (sumi1 + sumi2);
-#endif
-}
-
-static __device__ __forceinline__ float
-vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) {
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA
-  const block_iq4_xs* bq4 = (const block_iq4_xs*)vbq;
-  const uint8_t* values = (const uint8_t*)kvalues_iq4nl;
-
-  // iqs is 0...7
-  const int ib32 = iqs;
-  const int32_t* q8 = (const int*)bq8_1[ib32].qs;
-  const uint32_t* q4 = (const uint32_t*)bq4->qs + 4 * ib32;
-  const int8_t ls = ((bq4->scales_l[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf) | (((bq4->scales_h >> 2 * ib32) & 3) << 4);
-  const float d = __half2float(bq4->d) * (ls - 32) * __low2float(bq8_1[ib32].ds);
-  int v1, v2;
-  int sumi1 = 0, sumi2 = 0;
-  for (int j = 0; j < 4; ++j) {
-    get_int_from_table_16(q4[j], values, v1, v2);
-    sumi1 = __dp4a(v1, q8[j + 0], sumi1);
-    sumi2 = __dp4a(v2, q8[j + 4], sumi2);
-  }
-  return d * (sumi1 + sumi2);
-#endif
-}
diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py
new file mode 100644
index 0000000000..b672d58cf2
--- /dev/null
+++ b/python/freetoken/kernel/ggml.py
@@ -0,0 +1,155 @@
+"""llama.cpp's ggml-cuda kernels (vendored under ``csrc/ggml``) as a tvm-ffi module: prebuilt in the kernel cache, else compiled on first use.
+
+Weights stay in their packed ggml blocks. ``mul_mat`` / ``mul_mat_id`` pick MMVQ or MMQ the way llama.cpp's mul_mat / mul_mat_id do.
+Activations are fp32, bf16 or fp16 and are quantized to q8_1 first; results come out in the activation dtype and are accumulated in fp32.
+"""
+
+from __future__ import annotations
+
+import functools
+import pathlib
+from typing import TYPE_CHECKING
+
+import torch
+
+from .utils import _is_rocm, load_aot
+
+if TYPE_CHECKING:
+    from tvm_ffi import Module
+
+_ROOT = pathlib.Path(__file__).parent / "csrc" / "ggml"
+
+
+def load_module(build_directory: str | None = None) -> Module:
+    cuda_dir = _ROOT / "ggml-cuda"
+    sources = [
+        _ROOT / "freetoken" / "ggml_shim.cu",
+        _ROOT / "freetoken" / "ggml_ops.cu",
+        *(cuda_dir / f for f in ("mmvq.cu", "mmq.cu", "mmid.cu", "quantize.cu", "convert.cu", "getrows.cu")),
+        *sorted((cuda_dir / "template-instances").glob("mmq-instance-*.cu")),
+    ]
+    if _is_rocm():
+        flags = ["-ffast-math", "-DGGML_USE_HIP"]
+    else:
+        # a multi-arch kernel-cache build is ~6x smaller with compressed SASS
+        flags = ["-use_fast_math", "--extended-lambda", "--compress-mode=size"]
+    return load_aot(
+        "ggml",
+        cuda_files=[str(s) for s in sources],
+        extra_cuda_cflags=flags,
+        extra_include_paths=[str(_ROOT / "include"), str(_ROOT / "src"), str(cuda_dir)],
+        build_directory=build_directory,
+    )
+
+
+@functools.cache
+def _module() -> Module:
+    return load_module()
+
+
+@functools.cache
+def _mmq_tail_bytes(ggml_type: int, k: int) -> int:
+    return _module().mmq_tail_bytes(ggml_type, k)
+
+
+def _check_mmq_tail(w: torch.Tensor, ggml_type: int, k: int) -> None:
+    """MMQ tiles read up to MATRIX_ROW_PADDING values past the last row of ``w``, so its storage must hold them."""
+    tail = _mmq_tail_bytes(ggml_type, k)
+    if tail == 0 or w.numel() == 0:
+        return
+    end = w.storage_offset() + sum((n - 1) * s for n, s in zip(w.shape, w.stride())) + 1
+    if end + tail > w.untyped_storage().nbytes():
+        raise ValueError(f"ggml type {ggml_type} K={k} weight needs {tail} bytes of zero padding after its last row for MMQ")
+
+
+def mul_mat(w: torch.Tensor, ggml_type: int, x: torch.Tensor, *, out: torch.Tensor | None = None, xq: torch.Tensor | None = None) -> torch.Tensor:
+    """``x[M, K] @ w[N, K]^T`` in ``x``'s dtype with ``w`` packed as ``[N, row_bytes]`` uint8, or written into the ``[M, N]`` view ``out``.
+
+    ``xq`` is ``quantize_q8_1(x)``: the vector kernel reads it instead of quantizing ``x`` again.
+    """
+    m, k = x.shape
+    if out is None:
+        out = torch.empty(m, w.shape[0], dtype=x.dtype, device=x.device)
+    if not uses_mmvq(ggml_type, m):
+        _check_mmq_tail(w, ggml_type, k)
+    _module().mul_mat(w, ggml_type, x.contiguous(), out, xq)
+    return out
+
+
+def quantize_q8_1(x: torch.Tensor) -> torch.Tensor:
+    """``x[M, K]`` as the vector kernel's q8_1 blocks, to share among GEMVs that read the same ``x``."""
+    m, k = x.shape
+    out = torch.empty(m, _q8_1_row_bytes(k), dtype=torch.uint8, device=x.device)
+    _module().quantize_q8_1(x.contiguous(), out)
+    return out
+
+
+@functools.cache
+def _q8_1_row_bytes(k: int) -> int:
+    return _module().q8_1_row_bytes(k)
+
+
+@functools.cache
+def uses_mmvq(ggml_type: int, m: int) -> bool:
+    """Whether ``mul_mat`` of ``m`` rows on this GPU runs the vector kernel, the one that reads ``xq``."""
+    return _module().uses_mmvq(ggml_type, m)
+
+
+def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
+    """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``."""
+    t, top_k = ids.shape
+    out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device)
+    if t > mmvq_max_tokens(ggml_type):
+        _check_mmq_tail(w, ggml_type, x.shape[-1])
+    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out)
+    return out
+
+
+def mul_mat_id_glu(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, glu: str, limit: float = 0.0) -> torch.Tensor:
+    """``glu(gate, up)`` per route in one pass for small ``T``: ``w[E, 2I, row_bytes]`` holds gate rows then up rows; returns ``[T, k, I]`` in ``x``'s dtype.
+
+    ``glu`` is one of ``swiglu``, ``geglu`` (tanh GELU), ``swiglu_oai`` (alpha 1.702, limit 7) and ``swiglu_clamp`` (alpha 1, clamped at ``limit``).
+    """
+    t, top_k = ids.shape
+    out = torch.empty(t, top_k, w.shape[1] // 2, dtype=x.dtype, device=x.device)
+    _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit)
+    return out
+
+
+def moe_weighted_sum(y: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+    """``sum_j weights[t, j] * y[t, j, :]`` over ``y[T, k, N]``, accumulated in fp32 and returned in ``y``'s dtype."""
+    out = torch.empty(y.shape[0], y.shape[2], dtype=y.dtype, device=y.device)
+    _module().moe_weighted_sum(y.contiguous(), weights.float().contiguous(), out)
+    return out
+
+
+@functools.cache
+def mmvq_max_tokens(ggml_type: int) -> int:
+    """Largest token count the small-batch MoE kernels take on this GPU."""
+    return _module().mmvq_max_tokens(ggml_type)
+
+
+def dequantize(w: torch.Tensor, ggml_type: int, k: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+    out = torch.empty(w.numel() // w.shape[-1], k, dtype=dtype, device=w.device)
+    _module().dequantize(w, ggml_type, k, out)
+    return out
+
+
+def get_rows(w: torch.Tensor, ggml_type: int, k: int, ids: torch.Tensor, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+    out = torch.empty(ids.shape[0], k, dtype=dtype, device=w.device)
+    _module().get_rows(w, ggml_type, k, ids, out)
+    return out
+
+
+__all__ = [
+    "dequantize",
+    "get_rows",
+    "load_module",
+    "mmvq_max_tokens",
+    "moe_weighted_sum",
+    "mul_mat",
+    "mul_mat_id",
+    "mul_mat_id_glu",
+    "quantize_q8_1",
+    "uses_mmvq",
+]
diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py
deleted file mode 100644
index 04a1656098..0000000000
--- a/python/freetoken/kernel/gguf.py
+++ /dev/null
@@ -1,142 +0,0 @@
-"""Borrowed llama.cpp GGUF dequant/GEMM CUDA kernels, JIT-compiled on first use.
-
-The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored verbatim from sgl-kernel
-(``csrc/quantization/gguf/``), which are themselves ports of llama.cpp. We compile
-them through ``torch.utils.cpp_extension.load`` (the same toolchain sglang/vllm use)
-into a torch-op module and expose the handful of ops the GGUF path needs. This is a
-separate, torch-native extension that sits alongside FreeToken's tvm-ffi kernels.
-
-All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) and
-dequantize *inside* the kernel -- no bf16 copy of the weight is ever materialized.
-"""
-
-from __future__ import annotations
-
-import functools
-import os
-import pathlib
-import shutil
-
-import torch
-
-_CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf"
-
-
-def _host_compiler() -> str | None:
-    """A host compiler nvcc + libtorch headers accept.
-
-    The system default gcc can be too new for the torch headers (gcc 16 hard-errors),
-    and on this toolchain even nvcc+gcc-13 trips a non-conformant ``typename
-    decltype`` in ``List_inl.h`` once ``torch::Tensor`` is instantiated -- but nvcc
-    with ``clang++`` as host compiles it cleanly. So prefer clang++, then fall back
-    to an older gcc. Override with ``FREETOKEN_GGUF_HOST_CXX``.
-    """
-    override = os.environ.get("FREETOKEN_GGUF_HOST_CXX")
-    if override:
-        return override
-    for cxx in ("clang++", "g++-13", "g++-14", "g++-15"):
-        if shutil.which(cxx):
-            return cxx
-    return None
-
-
-def _c_compiler_for(cxx: str) -> str:
-    base = os.path.basename(cxx)
-    if "clang" in base:
-        return shutil.which("clang") or "clang"
-    cc = base.replace("g++", "gcc")
-    return shutil.which(cc) or cc
-
-@functools.cache
-def _module():
-    from torch.utils.cpp_extension import load
-
-    extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"]
-    host_cxx = _host_compiler()
-    if host_cxx is not None:
-        # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a
-        # libtorch/nvcc-compatible compiler. Force (not setdefault): the system
-        # default (CXX unset -> g++) can be a gcc too new for the torch headers.
-        cxx_path = shutil.which(host_cxx) or host_cxx
-        extra_cuda_cflags += ["-ccbin", cxx_path]
-        os.environ["CXX"] = cxx_path
-        os.environ["CC"] = _c_compiler_for(cxx_path)
-
-    # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a
-    # plain `load` of the single source compiles + binds the ggml_* ops.
-    return load(
-        name="freetoken_gguf_kernels",
-        sources=[str(_CSRC / "gguf_kernel.cu")],
-        extra_include_paths=[str(_CSRC)],
-        extra_cuda_cflags=extra_cuda_cflags,
-        verbose=True,
-    )
-
-
-# ---- thin typed wrappers (signatures mirror sgl_kernel.quantization.gguf) ----
-
-
-def ggml_dequantize(
-    weight: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None = None
-) -> torch.Tensor:
-    """Dequantize a packed GGUF weight ``[m, row_bytes]`` to a dense ``[m, n]`` tensor."""
-    return _module().ggml_dequantize(weight, quant_type, m, n, dtype)
-
-
-def ggml_mul_mat_vec_a8(
-    weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int
-) -> torch.Tensor:
-    """MMVQ: small-batch GEMV with on-the-fly dequant. ``row`` = output features."""
-    return _module().ggml_mul_mat_vec_a8(weight, x, quant_type, row)
-
-
-def ggml_mul_mat_a8(
-    weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int
-) -> torch.Tensor:
-    """MMQ: large-batch quantized matmul. ``row`` = output features."""
-    return _module().ggml_mul_mat_a8(weight, x, quant_type, row)
-
-
-def ggml_moe_a8(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    sorted_token_ids: torch.Tensor,
-    expert_ids: torch.Tensor,
-    num_tokens_post_padded: torch.Tensor,
-    quant_type: int,
-    row: int,
-    top_k: int,
-    tokens: int,
-) -> torch.Tensor:
-    """MMQ grouped expert matmul over stacked experts ``weight[E, row, *]``."""
-    return _module().ggml_moe_a8(
-        x, weight, sorted_token_ids, expert_ids, num_tokens_post_padded,
-        quant_type, row, top_k, tokens,
-    )
-
-
-def ggml_moe_a8_vec(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    topk_ids: torch.Tensor,
-    top_k: int,
-    quant_type: int,
-    row: int,
-    tokens: int,
-) -> torch.Tensor:
-    """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``."""
-    return _module().ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, row, tokens)
-
-
-def ggml_moe_get_block_size(quant_type: int) -> int:
-    return _module().ggml_moe_get_block_size(quant_type)
-
-
-__all__ = [
-    "ggml_dequantize",
-    "ggml_mul_mat_vec_a8",
-    "ggml_mul_mat_a8",
-    "ggml_moe_a8",
-    "ggml_moe_a8_vec",
-    "ggml_moe_get_block_size",
-]
diff --git a/python/freetoken/layers/activation.py b/python/freetoken/layers/activation.py
index cd7a402a65..e38a6caebc 100644
--- a/python/freetoken/layers/activation.py
+++ b/python/freetoken/layers/activation.py
@@ -6,10 +6,17 @@
     import torch
 
 
-def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
+def _use_flashinfer(x: torch.Tensor, out: torch.Tensor | None) -> bool:
+    """flashinfer's *_and_mul kernels take fp16 / bf16 in and the same dtype out; the Triton kernels take any float dtypes."""
+    import torch
+
     from freetoken.kernel.backend import is_flashinfer_installed
 
-    if is_flashinfer_installed():
+    return x.dtype in (torch.float16, torch.bfloat16) and (out is None or out.dtype == x.dtype) and is_flashinfer_installed()
+
+
+def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
+    if _use_flashinfer(x, out):
         from flashinfer import silu_and_mul
     else:
         from freetoken.kernel.triton.activation import silu_and_mul
@@ -18,9 +25,7 @@ def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
 
 
 def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
-    from freetoken.kernel.backend import is_flashinfer_installed
-
-    if is_flashinfer_installed():
+    if _use_flashinfer(x, out):
         from flashinfer import gelu_and_mul
     else:
         from freetoken.kernel.triton.activation import gelu_and_mul
@@ -30,9 +35,7 @@ def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
 
 def gelu_tanh_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
     """tanh-approximate GELU gate (`gelu_pytorch_tanh`) followed by elementwise mul."""
-    from freetoken.kernel.backend import is_flashinfer_installed
-
-    if is_flashinfer_installed():
+    if _use_flashinfer(x, out):
         from flashinfer import gelu_tanh_and_mul
     else:
         from freetoken.kernel.triton.activation import gelu_tanh_and_mul
diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index a22216548a..f86670b14c 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -3,7 +3,6 @@
 from typing import Dict
 
 import torch
-import torch.nn.functional as F
 from freetoken.core import get_global_ctx
 from freetoken.distributed import DistributedCommunicator, get_tp_info
 from freetoken.utils import div_ceil, nvtx_annotate
@@ -13,11 +12,18 @@
 
 
 class VocabParallelEmbedding(BaseOP):
+    """The vocab table, sharded over TP: its ``quant_method`` owns the table, looks rows up, and serves a tied LM head's product."""
+
+    quant_layer_kind = LayerKind.EMBEDDING
+
     def __init__(
         self,
         num_embeddings: int,
         embedding_dim: int,
         embed_scale: float | None = None,
+        *,
+        quant_config: QuantConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
         tp_info = get_tp_info()
@@ -28,7 +34,12 @@ def __init__(
         start_idx = self.num_embeddings_tp * tp_rank
         finish_idx = min(start_idx + self.num_embeddings_tp, num_embeddings)
         self.vocab_range = (start_idx, finish_idx - start_idx)
-        self.weight = torch.empty(self.num_embeddings_tp, embedding_dim)
+        self.embedding_dim = embedding_dim
+        # the LM head subclass is a linear and builds its own method
+        self.quant_method = None
+        if self.quant_layer_kind is LayerKind.EMBEDDING:
+            self.quant_method = quant_method_for(quant_config, self, prefix)
+            self.quant_method.create_weights(self)
         # Gemma scales embeddings by sqrt(hidden_size). The scale is materialized in
         # the weight dtype (bf16) to match HF, which downcasts the scalar. The GPU
         # scalar is built lazily (model __init__ runs on the meta device) and cached
@@ -39,14 +50,7 @@ def __init__(
 
     @nvtx_annotate("Embedding")
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        from freetoken.kernel import indexing
-
-        y = indexing(
-            weights=self.weight,
-            indices=x,
-            vocab_range=self.vocab_range if self.tp_size > 1 else None,
-        )
-
+        y = self.quant_method.embed(self, x)
         if self.tp_size > 1:
             y = self._comm.all_reduce(y)
         if self._embed_scale is not None:
@@ -83,7 +87,6 @@ def __init__(
         self.in_features = embedding_dim
         self.out_features = self.num_embeddings_tp
         self.output_sizes = (self.num_embeddings_tp,)
-        self.quant_method = None
         if tied_embedding is None:
             self.quant_method = quant_method_for(quant_config, self, prefix)
             self.quant_method.create_weights(self)
@@ -132,7 +135,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             del indices
 
         if self.tied_embedding is not None:
-            logits = F.linear(x, self.tied_embedding.weight, self.bias)
+            logits = self.tied_embedding.quant_method.logits(self.tied_embedding, x, self.bias)
         else:
             logits = self.quant_method.apply(self, x)
         if self.tp_size == 1:
diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py
deleted file mode 100644
index ac49b1a5be..0000000000
--- a/python/freetoken/layers/gguf.py
+++ /dev/null
@@ -1,128 +0,0 @@
-"""Native-GGUF quantized layers: weights stay in their packed block layout and are
-dequantized *inside* the borrowed llama.cpp kernels (no bf16 copy ever materialized).
-
-Mirrors vLLM/sglang's ``GGUFLinearMethod`` / ``GGUFEmbeddingMethod`` dispatch, ported
-onto FreeToken's ``BaseOP``. FreeToken keeps fused projections (qkv, gate_up) as a
-single tensor: because Q4_0/K-quants pack each *output row* independently over the
-input dim, the loader can concatenate the per-shard packed rows along dim 0 (they
-share an input dim, hence the same ``row_bytes``), so a fused layer is still one
-``[out, row_bytes]`` qweight -- no per-shard padding bookkeeping needed.
-
-TP is assumed to be 1 (the gemma4 GGUF path restricts to TP=1, like the HF path).
-"""
-
-from __future__ import annotations
-
-import torch
-
-from freetoken.models.gguf.dequant import (
-    BLOCK_SHAPE,
-    GGML_BF16,
-    GGML_F16,
-    GGML_F32,
-    GGML_NAME,
-    GGML_Q4_0,
-    GGML_Q6_K,
-    GGML_Q8_0,
-    row_bytes,
-)
-
-from .base import BaseOP
-
-# ggml type groups for kernel dispatch (subset we build kernels for).
-_UNQUANTIZED = {GGML_F32, GGML_F16, GGML_BF16}
-# standard + k-quants: both an MMVQ (small-batch GEMV) and MMQ (large-batch) kernel exist.
-_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K}
-_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K}
-_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K}
-
-# Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic).
-_MMVQ_SAFE = 6
-
-
-def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
-    """y = x @ dequant(qweight).T, dispatched by batch size and quant type."""
-    from freetoken.kernel.gguf import (
-        ggml_dequantize,
-        ggml_mul_mat_a8,
-        ggml_mul_mat_vec_a8,
-    )
-
-    out_features = qweight.shape[0]
-    if x.shape[0] == 0:
-        return x.new_empty((0, out_features))
-    if qweight_type in _UNQUANTIZED:
-        return x @ qweight.T
-    if x.shape[0] <= _MMVQ_SAFE and qweight_type in _MMVQ:
-        return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features)
-    if qweight_type in _MMQ:
-        return ggml_mul_mat_a8(qweight, x, qweight_type, out_features)
-    if qweight_type in _DEQUANT:
-        block, type_size = BLOCK_SHAPE[qweight_type]
-        in_features = qweight.shape[1] // type_size * block
-        weight = ggml_dequantize(qweight, qweight_type, out_features, in_features, x.dtype)
-        return x @ weight.T
-    raise NotImplementedError(f"unsupported GGUF type {GGML_NAME.get(qweight_type, qweight_type)}")
-
-
-class GGUFLinear(BaseOP):
-    """Linear whose weight is a native GGUF block-quantized ``[out, row_bytes]`` tensor."""
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        quant_type: int,
-        has_bias: bool = False,
-    ):
-        self.in_features = in_features
-        self.out_features = out_features
-        self._quant_type = quant_type
-        self.qweight = torch.empty(out_features, row_bytes(in_features, quant_type), dtype=torch.uint8)
-        self.bias = torch.empty(out_features) if has_bias else None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = fused_mul_mat_gguf(x, self.qweight, self._quant_type)
-        if self.bias is not None:
-            out = out + self.bias
-        return out
-
-
-class GGUFEmbedding(BaseOP):
-    """Vocab embedding stored as a native GGUF block-quantized table.
-
-    The full table is never dequantized: only the looked-up rows are gathered (in
-    packed form) and dequantized per lookup, matching vLLM's ``_apply_gguf_embedding``.
-    """
-
-    def __init__(
-        self,
-        num_embeddings: int,
-        embedding_dim: int,
-        quant_type: int,
-        embed_scale: float | None = None,
-    ):
-        self.num_embeddings = num_embeddings
-        self.embedding_dim = embedding_dim
-        self._quant_type = quant_type
-        self.qweight = torch.empty(
-            num_embeddings, row_bytes(embedding_dim, quant_type), dtype=torch.uint8
-        )
-        self._embed_scale = embed_scale
-        self._embed_scale_t: torch.Tensor | None = None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        from freetoken.kernel.gguf import ggml_dequantize
-
-        flat = x.flatten()
-        rows = self.qweight.index_select(0, flat)  # [n, row_bytes] packed
-        y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16)
-        y = y.view(*x.shape, self.embedding_dim)
-        if self._embed_scale is not None:
-            if self._embed_scale_t is None:
-                self._embed_scale_t = torch.tensor(self._embed_scale, dtype=y.dtype, device=y.device)
-            y = y * self._embed_scale_t
-        return y
-
-
-__all__ = ["GGUFLinear", "GGUFEmbedding", "fused_mul_mat_gguf"]
diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py
index d8a3e8619f..599ca04906 100644
--- a/python/freetoken/layers/linear.py
+++ b/python/freetoken/layers/linear.py
@@ -63,6 +63,7 @@ def __init__(
         output_size: int,
         has_bias: bool,
         *,
+        output_sizes: List[int] | None = None,
         quant_config: QuantConfig | None = None,
         prefix: str = "",
     ):
@@ -72,6 +73,7 @@ def __init__(
             local_isize=input_size,
             local_osize=output_size,
             has_bias=has_bias,
+            output_sizes=output_sizes,
             quant_config=quant_config,
             prefix=prefix,
         )
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 79ad9e3a5f..598bd1193f 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -81,16 +81,12 @@ def __init__(
         self.strategy = strategy
         self.decode_target = decode_target
         self.prefix = prefix
-        # offload layers without a quant config stay on the format-tag banks (GGUF q4_0)
-        self.quant_method = None
-        if quant_config is not None or allocate_experts:
-            self.quant_method = quant_method_for(quant_config, self, prefix)
-            if allocate_experts:
-                self.quant_method.create_weights(self)
+        self.quant_method = quant_method_for(quant_config, self, prefix)
+        if allocate_experts:
+            self.quant_method.create_weights(self)
 
     def finalize(self) -> None:
-        if self.quant_method is not None:
-            self.quant_method.finalize(self)
+        self.quant_method.finalize(self)
 
     def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if self.tp_size > 1:
@@ -103,7 +99,6 @@ def _resident_gemm(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> torch.Tensor:
-        assert self.quant_method is not None
         return self.quant_method.apply(
             hidden_states, topk_weights, topk_ids, self.quant_method.resident_view(self),
             layer=self, is_prefill=get_global_ctx().batch.is_prefill,
@@ -191,7 +186,7 @@ def __init__(
 
     @property
     def skips_inactive_slots(self) -> bool:
-        return self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots
+        return self.quant_method.kernel.supports_inactive_slots
 
     def forward(
         self,
@@ -321,7 +316,7 @@ def _decode_hybrid(
         assert executor is not None, "CPU MoE executor was not initialized"
         # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
         # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
-        # with several quants this holds only because the cache zeroes staged slot 0 rows (zero_inactive_route_slots)
+        # with several quants this holds only because the cache zeroes its staging area after prefill
         clamp = not self.skips_inactive_slots
         gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp)
         pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids)
@@ -402,7 +397,6 @@ def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, .
 
     # ------------------------------------------------------------------
     # Kernel dispatch: ``views`` are the bank tensors the movement step produced (in bank registration order) and ``topk_ids`` already index their rows.
-    # GGUF q4_0 experts still dispatch on the cache's format tag until they get a method.
     # ------------------------------------------------------------------
 
     def _expert_gemm(
@@ -417,27 +411,15 @@ def _expert_gemm(
         alphas: tuple[torch.Tensor, torch.Tensor] | None,
         is_prefill: bool,
     ) -> torch.Tensor:
-        if self.quant_method is not None:
-            from freetoken.moe.legacy_format import canonical_role  # legacy_format imports this package
+        from freetoken.moe.legacy_format import canonical_role  # legacy_format imports this package
 
-            view = ExpertView(
-                {canonical_role(name): t for name, t in zip(cache.bank_schema, views)},
-                slots=None if n is not None else topk_ids, n=n, alphas=alphas,
-            )
-            return self.quant_method.apply(
-                hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill
-            )
-        fmt = cache.quant_format
-        if fmt == "q4_0":
-            # Native GGUF Q4_0 experts: dequant-in-kernel grouped GEMV (MMVQ) over the
-            # streamed packed banks; topk_ids already index the cache slots / layer.
-            from freetoken.moe.fused_q4_0 import fused_experts_gguf_q4_0
-
-            gate_up, down = views
-            return fused_experts_gguf_q4_0(
-                hidden_states, gate_up, down, topk_weights, topk_ids, self.activation
-            )
-        raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}")
+        view = ExpertView(
+            {canonical_role(name): t for name, t in zip(cache.bank_schema, views)},
+            slots=None if n is not None else topk_ids, n=n, alphas=alphas,
+        )
+        return self.quant_method.apply(
+            hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill
+        )
 
 
 def iter_moe_layers(model) -> Iterator[MoELayer]:
diff --git a/python/freetoken/layers/quantization/__init__.py b/python/freetoken/layers/quantization/__init__.py
index 62019fc57d..905b80ef5d 100644
--- a/python/freetoken/layers/quantization/__init__.py
+++ b/python/freetoken/layers/quantization/__init__.py
@@ -1,6 +1,7 @@
 from .scheme import GroupShape, QuantKind, QuantScheme, WeightDesc
 from .registry import LayerKind, dialects, method_class, methods_for, register_dialect, register_method
 from .method import KernelSelectionError, QuantMethod, finalize_quant, select_kernel
+from .embedding import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod, GGUFEmbeddingMethod, UnquantizedEmbeddingMethod
 from .linear import (
     Fp8BlockLinearMethod,
     Fp8TensorLinearMethod,
@@ -28,6 +29,7 @@
 from .configs import (
     CompressedTensorsConfig,
     Fp8BlockConfig,
+    GGUFConfig,
     ModelOptConfig,
     Mxfp4Config,
     NoQuantConfig,
@@ -42,11 +44,12 @@
     "QuantKind", "GroupShape", "QuantScheme", "WeightDesc",
     "LayerKind", "dialects", "method_class", "methods_for", "register_dialect", "register_method",
     "KernelSelectionError", "select_kernel", "QuantMethod",
+    "EmbeddingConfig", "EmbeddingKernel", "EmbeddingMethod", "UnquantizedEmbeddingMethod", "GGUFEmbeddingMethod",
     "LinearConfig", "LinearKernel", "LinearMethod",
     "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod",
     "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod",
     "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod",
     "QuantBackend", "set_quant_backend", "get_quant_backend", "NameMap",
-    "QuantConfig", "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config",
+    "QuantConfig", "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "GGUFConfig", "Mxfp4Config",
     "quant_method_for", "quantization_config_of", "set_quant_config", "get_quant_config", "finalize_quant",
 ]
diff --git a/python/freetoken/layers/quantization/configs/__init__.py b/python/freetoken/layers/quantization/configs/__init__.py
index 4bd7a6a9a7..358e6fecd6 100644
--- a/python/freetoken/layers/quantization/configs/__init__.py
+++ b/python/freetoken/layers/quantization/configs/__init__.py
@@ -1,11 +1,12 @@
 """One QuantConfig per checkpoint dialect; importing the package registers them all."""
 
 from .base import NoQuantConfig, QuantConfig, quantization_config_of
-from . import compressed_tensors, fp8, modelopt, mxfp4
+from . import compressed_tensors, fp8, gguf, modelopt, mxfp4
 from .compressed_tensors import CompressedTensorsConfig
 from .fp8 import Fp8BlockConfig
 from .modelopt import ModelOptConfig
 from .mxfp4 import Mxfp4Config
+from .gguf import GGUFConfig
 
 _NO_QUANT = NoQuantConfig()
 
@@ -33,6 +34,6 @@ def get_quant_config() -> QuantConfig | None:
 
 __all__ = [
     "QuantConfig", "quantization_config_of", "quant_method_for", "set_quant_config", "get_quant_config",
-    "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config",
-    "compressed_tensors", "fp8", "modelopt", "mxfp4",
+    "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config", "GGUFConfig",
+    "compressed_tensors", "fp8", "gguf", "modelopt", "mxfp4",
 ]
diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py
index a64136103b..78818e8091 100644
--- a/python/freetoken/layers/quantization/configs/base.py
+++ b/python/freetoken/layers/quantization/configs/base.py
@@ -6,14 +6,15 @@
 from dataclasses import dataclass
 from typing import Any, ClassVar
 
+from ..embedding import EmbeddingConfig
 from ..linear import LinearConfig
 from ..moe import MoEConfig
 from ..names import Matcher, NameMap, name_set
 from ..quant_backend import get_quant_backend
-from ..registry import LayerKind, dialects, method_class, register_dialect
+from ..registry import LayerKind, dialects, method_class, methods_for, register_dialect
 from ..scheme import QuantKind, QuantScheme
 
-LAYER_CONFIGS = {LayerKind.LINEAR: LinearConfig, LayerKind.MOE: MoEConfig}
+LAYER_CONFIGS = {LayerKind.LINEAR: LinearConfig, LayerKind.MOE: MoEConfig, LayerKind.EMBEDDING: EmbeddingConfig}
 
 
 def cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
@@ -88,6 +89,9 @@ def get_quant_method(self, layer: Any, prefix: str):
         scheme = self.scheme_for(prefix)
         layer_kind = layer.quant_layer_kind
         kind = scheme.kind if scheme else QuantKind.NONE
+        if layer_kind is LayerKind.EMBEDDING and kind not in methods_for(layer_kind):
+            # a checkpoint that quantizes its linears stores the vocab table dense
+            scheme, kind = None, QuantKind.NONE
         cls = method_class(kind, layer_kind)
         cfg = self.layer_config(layer, layer_kind, scheme)
         return cls(cfg, get_quant_backend().select(layer_kind, kind))
diff --git a/python/freetoken/layers/quantization/configs/gguf.py b/python/freetoken/layers/quantization/configs/gguf.py
new file mode 100644
index 0000000000..bbc24f137a
--- /dev/null
+++ b/python/freetoken/layers/quantization/configs/gguf.py
@@ -0,0 +1,82 @@
+"""GGUF: every tensor carries its own ggml type, read from the file's tensor table."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, ClassVar
+
+from ..ggml_types import is_quantized, type_name
+from ..names import NameMap
+from ..quant_backend import get_quant_backend
+from ..registry import LayerKind, method_class, register_dialect
+from ..scheme import QuantKind, QuantScheme, gguf_scheme
+from .base import QuantConfig, Stored
+
+
+@register_dialect
+class GGUFConfig(QuantConfig):
+    """``tensor_types`` is the file's name -> ggml type table; ``to_gguf(name, tensor_types)`` (the family's adapter) maps a checkpoint module name to the GGUF tensors that store it."""
+
+    dialect = "gguf"
+    STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {QuantKind.GGUF: {"weight": "weight"}}
+
+    def __init__(
+        self,
+        tensor_types: dict[str, int],
+        to_gguf: Callable[[str, dict[str, int]], tuple[str, ...]],
+        name_map: NameMap | None = None,
+        unquantized: tuple[str, ...] = (),
+    ):
+        super().__init__(name_map, unquantized)
+        self.tensor_types = tensor_types
+        self.to_gguf = to_gguf
+
+    @classmethod
+    def claims(cls, q: dict[str, Any]) -> bool:
+        return False
+
+    def _part_types(self, name: str) -> tuple[int, ...]:
+        types = []
+        for tensor in self.to_gguf(name, self.tensor_types):
+            key = f"{tensor}.weight" if f"{tensor}.weight" in self.tensor_types else tensor
+            if key not in self.tensor_types:
+                raise KeyError(f"GGUF has no tensor {key!r} for {name!r}")
+            types.append(self.tensor_types[key])
+        return tuple(types)
+
+    def scheme_for_name(self, name: str) -> QuantScheme | None:
+        types = self._part_types(name)
+        if not any(is_quantized(t) for t in types):
+            return None
+        if len(types) != 1:
+            raise NotImplementedError(f"{name} is stored as {len(types)} quantized GGUF tensors")
+        return gguf_scheme(types)
+
+    def scheme_for(self, prefix: str) -> QuantScheme | None:
+        if prefix in self._schemes:
+            return self._schemes[prefix]
+        names = self.name_map.to_checkpoint(prefix)
+        parts = [() if self.unquantized(n) else self._part_types(n) for n in names]
+        quantized = [any(is_quantized(t) for t in p) for p in parts]
+        if not any(quantized):
+            scheme = None
+        elif any(len(p) != 1 for p in parts):
+            found = {n: [type_name(t) for t in p] for n, p in zip(names, parts)}
+            raise NotImplementedError(f"fused module {prefix!r} needs one GGUF tensor per part, got {found}")
+        else:
+            scheme = gguf_scheme(p[0] for p in parts)
+        self._schemes[prefix] = scheme
+        return scheme
+
+    def get_quant_method(self, layer: Any, prefix: str):
+        if layer.quant_layer_kind is not LayerKind.MOE:
+            return super().get_quant_method(layer, prefix)
+        # every expert layer loads its GGUF banks, float experts included; the GGUF MoE kernel picks how to multiply them
+        parts = [self._part_types(n) for n in self.name_map.to_checkpoint(prefix)]
+        if any(len(p) != 1 for p in parts):
+            raise NotImplementedError(f"experts {prefix!r} are stored as {parts} GGUF tensors per projection")
+        scheme = gguf_scheme(p[0] for p in parts)
+        cfg = self.layer_config(layer, LayerKind.MOE, scheme)
+        return method_class(QuantKind.GGUF, LayerKind.MOE)(cfg, get_quant_backend().select(LayerKind.MOE, QuantKind.GGUF))
+
+
+__all__ = ["GGUFConfig"]
diff --git a/python/freetoken/layers/quantization/embedding/__init__.py b/python/freetoken/layers/quantization/embedding/__init__.py
new file mode 100644
index 0000000000..ccdbece04e
--- /dev/null
+++ b/python/freetoken/layers/quantization/embedding/__init__.py
@@ -0,0 +1,10 @@
+from .base import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod
+from . import gguf, unquantized
+from .gguf import GGUFEmbeddingMethod
+from .unquantized import UnquantizedEmbeddingMethod
+
+__all__ = [
+    "EmbeddingConfig", "EmbeddingKernel", "EmbeddingMethod",
+    "UnquantizedEmbeddingMethod", "GGUFEmbeddingMethod",
+    "gguf", "unquantized",
+]
diff --git a/python/freetoken/layers/quantization/embedding/base.py b/python/freetoken/layers/quantization/embedding/base.py
new file mode 100644
index 0000000000..caa873f4f9
--- /dev/null
+++ b/python/freetoken/layers/quantization/embedding/base.py
@@ -0,0 +1,64 @@
+"""The embedding side: per-layer config, the kernel ABC and the Method base. The method owns the vocab table: it looks rows up, and a tied LM head multiplies by the same table through it."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import torch
+
+from ..method import QuantMethod
+from ..scheme import QuantScheme
+
+
+@dataclass(frozen=True)
+class EmbeddingConfig:
+    """The TP-local table shape, the dtype a lookup returns, and the table's scheme."""
+
+    num_embeddings: int
+    embedding_dim: int
+    tp_size: int
+    dtype: torch.dtype
+    scheme: QuantScheme | None = None
+
+    @classmethod
+    def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "EmbeddingConfig":
+        return cls(layer.num_embeddings_tp, layer.embedding_dim, layer.tp_size, torch.get_default_dtype(), scheme)
+
+
+class EmbeddingKernel(ABC):
+    """One backend for one table kind: the row lookup, and ``x @ table^T`` for a tied LM head."""
+
+    name: ClassVar[str]
+
+    def unusable_reason(self, cfg: EmbeddingConfig) -> str | None:
+        return None
+
+    def worth_it(self, cfg: EmbeddingConfig) -> bool:
+        return True
+
+    def finalize(self, layer: Any) -> None:
+        pass
+
+    @abstractmethod
+    def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: ...
+
+    @abstractmethod
+    def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: ...
+
+
+class EmbeddingMethod(QuantMethod):
+    """Declares the table for its kind; the lookup and the tied head's product go to the kernel."""
+
+    @abstractmethod
+    def create_weights(self, layer: Any) -> None: ...
+
+    def finalize(self, layer: Any) -> None:
+        self.kernel.finalize(layer)
+
+    def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor:
+        return self.kernel.embed(layer, ids)
+
+    def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+        return self.kernel.logits(layer, x, bias)
diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py
new file mode 100644
index 0000000000..f41b7b6a4b
--- /dev/null
+++ b/python/freetoken/layers/quantization/embedding/gguf.py
@@ -0,0 +1,45 @@
+"""GGUF tables: packed ggml rows, decoded per looked-up row; a tied head runs the GGUF linear kernel on the same rows."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+
+from ..ggml_types import Segment, pad_rows, row_bytes
+from ..linear.gguf import GgmlLinearKernel
+from ..registry import LayerKind, register_method
+from ..scheme import QuantKind, gguf_types
+from .base import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod
+
+
+class GgmlEmbeddingKernel(EmbeddingKernel):
+    name = "ggml"
+
+    def unusable_reason(self, cfg: EmbeddingConfig) -> str | None:
+        return "GGUF tables are not sharded over TP" if cfg.tp_size > 1 else None
+
+    def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel import ggml
+
+        cfg = layer.quant_method.cfg
+        (ggml_type,) = gguf_types(cfg.scheme)
+        rows = ggml.get_rows(layer.weight, ggml_type, cfg.embedding_dim, ids.reshape(-1).to(torch.int32), cfg.dtype)
+        return rows.view(*ids.shape, cfg.embedding_dim)
+
+    def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
+        cfg = layer.quant_method.cfg
+        (ggml_type,) = gguf_types(cfg.scheme)
+        table = SimpleNamespace(_ggml_segments=[Segment(ggml_type, 0, 1, 0, cfg.num_embeddings)], weight=layer.weight, bias=bias)
+        return GgmlLinearKernel().apply(table, x)
+
+
+@register_method(QuantKind.GGUF, LayerKind.EMBEDDING)
+class GGUFEmbeddingMethod(EmbeddingMethod):
+    candidates = (GgmlEmbeddingKernel,)
+
+    def create_weights(self, layer: Any) -> None:
+        (ggml_type,) = gguf_types(self.cfg.scheme)
+        h = self.cfg.embedding_dim
+        layer.weight = torch.empty(self.cfg.num_embeddings + pad_rows(h), row_bytes(ggml_type, h), dtype=torch.uint8)
diff --git a/python/freetoken/layers/quantization/embedding/unquantized.py b/python/freetoken/layers/quantization/embedding/unquantized.py
new file mode 100644
index 0000000000..5b755baf19
--- /dev/null
+++ b/python/freetoken/layers/quantization/embedding/unquantized.py
@@ -0,0 +1,33 @@
+"""Dense tables: the indexing kernel looks rows up; a tied head runs the unquantized linear kernel on the same table."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+
+from ..linear.unquantized import TorchLinearKernel
+from ..registry import LayerKind, register_method
+from ..scheme import QuantKind
+from .base import EmbeddingKernel, EmbeddingMethod
+
+
+class IndexingEmbeddingKernel(EmbeddingKernel):
+    name = "indexing"
+
+    def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel import indexing
+
+        return indexing(weights=layer.weight, indices=ids, vocab_range=layer.vocab_range if layer.tp_size > 1 else None)
+
+    def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
+        return TorchLinearKernel().apply(SimpleNamespace(weight=layer.weight, bias=bias), x)
+
+
+@register_method(QuantKind.NONE, LayerKind.EMBEDDING)
+class UnquantizedEmbeddingMethod(EmbeddingMethod):
+    candidates = (IndexingEmbeddingKernel,)
+
+    def create_weights(self, layer: Any) -> None:
+        layer.weight = torch.empty(self.cfg.num_embeddings, self.cfg.embedding_dim)
diff --git a/python/freetoken/layers/quantization/ggml_types.py b/python/freetoken/layers/quantization/ggml_types.py
new file mode 100644
index 0000000000..8b257cd8c7
--- /dev/null
+++ b/python/freetoken/layers/quantization/ggml_types.py
@@ -0,0 +1,88 @@
+"""ggml type facts shared by the GGUF dialect, its methods and the GGUF weight readers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from freetoken.third_party import gguf
+
+# ggml reads up to this many values past the end of a packed matrix (llama.cpp MATRIX_ROW_PADDING)
+MATRIX_ROW_PADDING = 512
+
+UNQUANTIZED = frozenset(
+    int(t) for t in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16)
+)
+
+
+_FLOAT_DTYPES = {
+    int(gguf.GGMLQuantizationType.F32): torch.float32,
+    int(gguf.GGMLQuantizationType.F16): torch.float16,
+    int(gguf.GGMLQuantizationType.BF16): torch.bfloat16,
+}
+
+
+def float_dtype(ggml_type: int) -> torch.dtype | None:
+    """The torch dtype of an unquantized ggml type, None for a quantized one."""
+    return _FLOAT_DTYPES.get(int(ggml_type))
+
+
+def type_name(ggml_type: int) -> str:
+    return gguf.GGMLQuantizationType(ggml_type).name
+
+
+def is_quantized(ggml_type: int) -> bool:
+    return int(ggml_type) not in UNQUANTIZED
+
+
+def row_bytes(ggml_type: int, k: int) -> int:
+    block, size = gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(ggml_type)]
+    if k % block:
+        raise ValueError(f"{type_name(ggml_type)} rows need a multiple of {block} values, got {k}")
+    return k // block * size
+
+
+def pad_rows(k: int) -> int:
+    """Zero rows to append after a packed [rows, K] matrix so ggml's tail reads stay inside it."""
+    return 0 if k % MATRIX_ROW_PADDING == 0 else 1
+
+
+def mmq_tail_bytes(ggml_type: int, k: int) -> int:
+    """Bytes past the last row of a packed [rows, K] matrix that MMQ reads and that must read as zeros."""
+    rem = k % MATRIX_ROW_PADDING
+    return 0 if rem == 0 or not is_quantized(ggml_type) else row_bytes(ggml_type, MATRIX_ROW_PADDING - rem)
+
+
+@dataclass(frozen=True)
+class Segment:
+    """A run of consecutive fused parts that share one ggml type: one packed matrix and one GEMM."""
+
+    ggml_type: int
+    first_part: int
+    num_parts: int
+    out_start: int
+    rows: int
+
+
+def segments(types: tuple[int, ...], output_sizes: tuple[int, ...]) -> list[Segment]:
+    if len(types) != len(output_sizes):
+        raise ValueError(f"{len(types)} ggml types for {len(output_sizes)} fused parts")
+    out: list[Segment] = []
+    start = 0
+    for i, (t, n) in enumerate(zip(types, output_sizes)):
+        if out and out[-1].ggml_type == t:
+            last = out[-1]
+            out[-1] = Segment(t, last.first_part, last.num_parts + 1, last.out_start, last.rows + n)
+        else:
+            out.append(Segment(int(t), i, 1, start, n))
+        start += n
+    return out
+
+
+def segment_param(index: int) -> str:
+    """State-dict leaf of segment ``index``: ``weight`` for the first, then ``weight_1``, ``weight_2``."""
+    return "weight" if index == 0 else f"weight_{index}"
+
+
+__all__ = ["MATRIX_ROW_PADDING", "Segment", "float_dtype", "is_quantized", "mmq_tail_bytes", "pad_rows", "row_bytes", "segment_param", "segments", "type_name"]
diff --git a/python/freetoken/layers/quantization/linear/__init__.py b/python/freetoken/layers/quantization/linear/__init__.py
index 8236abe018..781f2b458a 100644
--- a/python/freetoken/layers/quantization/linear/__init__.py
+++ b/python/freetoken/layers/quantization/linear/__init__.py
@@ -1,13 +1,14 @@
 from .base import LinearConfig, LinearKernel, LinearMethod
-from . import fp8_block, fp8_tensor, mxfp8, nvfp4, unquantized
+from . import fp8_block, fp8_tensor, gguf, mxfp8, nvfp4, unquantized
 from .fp8_block import Fp8BlockLinearMethod
 from .fp8_tensor import Fp8TensorLinearMethod
 from .mxfp8 import Mxfp8LinearMethod
 from .nvfp4 import Nvfp4LinearMethod
 from .unquantized import UnquantizedLinearMethod
+from .gguf import GGUFLinearMethod
 
 __all__ = [
     "LinearConfig", "LinearKernel", "LinearMethod",
-    "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod",
-    "fp8_block", "fp8_tensor", "mxfp8", "nvfp4", "unquantized",
+    "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod", "GGUFLinearMethod",
+    "fp8_block", "fp8_tensor", "gguf", "mxfp8", "nvfp4", "unquantized",
 ]
diff --git a/python/freetoken/layers/quantization/linear/gguf.py b/python/freetoken/layers/quantization/linear/gguf.py
new file mode 100644
index 0000000000..ac08dbe287
--- /dev/null
+++ b/python/freetoken/layers/quantization/linear/gguf.py
@@ -0,0 +1,56 @@
+"""GGUF Linear: packed ggml rows multiplied by llama.cpp's MMVQ / MMQ kernels; a fused module runs one GEMM per ggml type, float parts as plain bf16 rows."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ..ggml_types import is_quantized, pad_rows, row_bytes, segment_param, segments
+from ..registry import LayerKind, register_method
+from ..scheme import QuantKind, gguf_types
+from .base import LinearConfig, LinearKernel, LinearMethod
+
+
+
+class GgmlLinearKernel(LinearKernel):
+    name = "ggml"
+
+    def unusable_reason(self, cfg: LinearConfig) -> str | None:
+        return None
+
+    def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
+        from freetoken.kernel import ggml
+
+        x2 = x.reshape(-1, x.shape[-1])
+        segs = layer._ggml_segments
+        m = x2.shape[0]
+        y = torch.empty(m, sum(s.rows for s in segs), dtype=x.dtype, device=x.device)
+        # every segment reads the same x: quantize it once for all the vector-kernel segments
+        xq = ggml.quantize_q8_1(x2) if sum(is_quantized(s.ggml_type) and ggml.uses_mmvq(s.ggml_type, m) for s in segs) > 1 else None
+        for i, s in enumerate(segs):
+            w = getattr(layer, segment_param(i))[: s.rows]
+            out = y[:, s.out_start : s.out_start + s.rows]
+            if is_quantized(s.ggml_type):
+                ggml.mul_mat(w, s.ggml_type, x2, out=out, xq=xq)
+            else:
+                torch.mm(x2, w.t(), out=out)
+        if layer.bias is not None:
+            y = y + layer.bias
+        return y.reshape(*x.shape[:-1], y.shape[-1])
+
+
+@register_method(QuantKind.GGUF, LayerKind.LINEAR)
+class GGUFLinearMethod(LinearMethod):
+    candidates = (GgmlLinearKernel,)
+
+    def create_weights(self, layer: Any) -> None:
+        g = self.cfg
+        segs = segments(gguf_types(g.scheme), g.output_sizes)
+        layer._ggml_segments = segs
+        k = g.in_features
+        for i, s in enumerate(segs):
+            if is_quantized(s.ggml_type):
+                setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(k), row_bytes(s.ggml_type, k), dtype=torch.uint8))
+            else:
+                setattr(layer, segment_param(i), torch.empty(s.rows, k))
diff --git a/python/freetoken/layers/quantization/moe/__init__.py b/python/freetoken/layers/quantization/moe/__init__.py
index a5f2dc6645..1230f36f08 100644
--- a/python/freetoken/layers/quantization/moe/__init__.py
+++ b/python/freetoken/layers/quantization/moe/__init__.py
@@ -1,13 +1,14 @@
 from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod
-from . import fp8_block, mxfp4, mxfp8, nvfp4, unquantized
+from . import fp8_block, gguf, mxfp4, mxfp8, nvfp4, unquantized
 from .fp8_block import Fp8BlockMoEMethod
 from .mxfp4 import Mxfp4MoEMethod
 from .mxfp8 import Mxfp8MoEMethod
 from .nvfp4 import Nvfp4MoEMethod
 from .unquantized import UnquantizedMoEMethod
+from .gguf import GGUFMoEMethod
 
 __all__ = [
     "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod",
-    "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod",
-    "fp8_block", "mxfp4", "mxfp8", "nvfp4", "unquantized",
+    "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod", "GGUFMoEMethod",
+    "fp8_block", "gguf", "mxfp4", "mxfp8", "nvfp4", "unquantized",
 ]
diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py
new file mode 100644
index 0000000000..78e051da67
--- /dev/null
+++ b/python/freetoken/layers/quantization/moe/gguf.py
@@ -0,0 +1,118 @@
+"""GGUF experts: packed ggml rows in two banks (gate|up and down), multiplied by llama.cpp's mul_mat_id kernels; float experts run on the bf16 fused MoE."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+from ..ggml_types import float_dtype, is_quantized, row_bytes, type_name
+from ..registry import LayerKind, register_method
+from ..scheme import QuantKind, gguf_types
+from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod, gated_epilogue_reason, is_resident, limit_or_inf
+
+
+
+def fused_glu(layer) -> tuple[str, float] | None:
+    """The MMVQ epilogue (``ggml.mul_mat_id_glu``) that computes this layer's gated activation, or None to run it as a separate kernel."""
+    act, alpha, limit = layer.activation, float(layer.alpha), limit_or_inf(layer)
+    if act == "silu" and alpha == 1.0 and limit == math.inf:
+        return "swiglu", 0.0
+    if act == "gelu_tanh":
+        return "geglu", 0.0
+    if act == "swiglu_clamp" and alpha == 1.0 and limit != math.inf:
+        return "swiglu_clamp", limit
+    # ggml's fused swiglu_oai keeps its default alpha and limit
+    if act == "swigluoai" and math.isclose(alpha, 1.702, rel_tol=1e-6) and math.isclose(limit, 7.0, rel_tol=1e-6):
+        return "swiglu_oai", 0.0
+    return None
+
+
+def expert_types(cfg: MoEConfig) -> tuple[int, int]:
+    gate, up, down = gguf_types(cfg.scheme)
+    if gate != up:
+        raise NotImplementedError(f"GGUF experts with gate {type_name(gate)} and up {type_name(up)} need one gate_up bank type")
+    if is_quantized(gate) != is_quantized(down):
+        raise NotImplementedError(f"GGUF experts with gate_up {type_name(gate)} and down {type_name(down)} mix float and quantized rows")
+    return gate, down
+
+
+
+class GgmlMoEKernel(MoEKernel):
+    name = "ggml"
+    cpu_format = "gguf"
+
+    def unusable_reason(self, cfg: MoEConfig) -> str | None:
+        if not is_quantized(expert_types(cfg)[0]):
+            reason = gated_epilogue_reason(cfg)
+            if reason:
+                return f"float GGUF experts run on the bf16 fused MoE: {reason}"
+        return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+
+    def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
+        gate_up, down = expert_types(cfg)
+        i, h = cfg.local_intermediate, cfg.hidden
+        dtype = float_dtype(gate_up)
+        if dtype is not None:
+            return {"gate_up": BankSpec((2 * i, h), dtype), "down": BankSpec((h, i), dtype)}
+        return {
+            "gate_up": BankSpec((2 * i, row_bytes(gate_up, h)), torch.uint8),
+            "down": BankSpec((h, row_bytes(down, i)), torch.uint8),
+        }
+
+    def pack(self, pieces, cfg: MoEConfig, out):
+        i = cfg.local_intermediate
+        out["gate_up"][:, :i].copy_(pieces["gate"])
+        out["gate_up"][:, i:].copy_(pieces["up"])
+        out["down"].copy_(pieces["down"])
+        return {}
+
+    def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefill: bool):
+        from freetoken.kernel import ggml
+        from freetoken.layers.activation import gated_act_and_mul
+
+        gate_up_type, down_type = expert_types(layer.quant_method.cfg)
+        if not is_quantized(gate_up_type):
+            return self._apply_float(layer, x, topk_weights, topk_ids, view, gate_up_type, is_prefill=is_prefill)
+        assert not layer.apply_router_weight_on_input, "GGUF experts apply the router weights on the output"
+        ids = topk_ids.to(torch.int32)
+        gate_up = view.tensors["gate_up"]
+        glu = fused_glu(layer)
+        small = x.shape[0] <= min(ggml.mmvq_max_tokens(gate_up_type), ggml.mmvq_max_tokens(down_type))
+        if glu is not None and small:
+            inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu)
+        else:
+            h = ggml.mul_mat_id(gate_up, gate_up_type, x, ids)
+            t, k, two_i = h.shape
+            inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device)
+            gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer))
+            del h  # free the gate_up output before the down GEMM allocates its own
+        y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids)
+        return ggml.moe_weighted_sum(y, topk_weights)
+
+    @staticmethod
+    def _apply_float(layer, x, topk_weights, topk_ids, view: ExpertView, ggml_type: int, *, is_prefill: bool):
+        """Float experts: their banks in their own dtype, through the bf16 fused MoE (the unquantized kernel's path)."""
+        from freetoken.moe.fused import fused_experts_decode_impl, fused_experts_impl
+
+        gate_up, down = view.tensors["gate_up"], view.tensors["down"]
+        if gate_up.dtype != x.dtype:
+            raise NotImplementedError(f"GGUF {type_name(ggml_type)} experts need {gate_up.dtype} activations, got {x.dtype}")
+        impl = fused_experts_impl if is_prefill or is_resident(layer) else fused_experts_decode_impl
+        return impl(x, gate_up, down, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, float(layer.alpha), limit_or_inf(layer))
+
+
+@register_method(QuantKind.GGUF, LayerKind.MOE)
+class GGUFMoEMethod(MoEMethod):
+    candidates = (GgmlMoEKernel,)
+
+    @property
+    def cpu_format(self) -> str | None:
+        # the CPU executor reads Q4_0 experts in their GGUF rows and BF16 experts as its bf16 banks; other ggml types have no CPU kernel yet
+        gate_up, down = expert_types(self.cfg)
+        names = {type_name(gate_up), type_name(down)}
+        if names == {"Q4_0"}:
+            return "q4_0"
+        if names == {"BF16"}:
+            return "bf16"
+        return self.kernel.cpu_format
diff --git a/python/freetoken/layers/quantization/registry.py b/python/freetoken/layers/quantization/registry.py
index 2d13cf50cc..c561e55edc 100644
--- a/python/freetoken/layers/quantization/registry.py
+++ b/python/freetoken/layers/quantization/registry.py
@@ -11,6 +11,7 @@
 class LayerKind(Enum):
     LINEAR = "linear"
     MOE = "moe"
+    EMBEDDING = "embedding"
 
     def __str__(self) -> str:
         return self.value
diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py
index 89f968afc8..ab43814e6a 100644
--- a/python/freetoken/layers/quantization/scheme.py
+++ b/python/freetoken/layers/quantization/scheme.py
@@ -20,6 +20,7 @@ class QuantKind(Enum):
     MXFP8 = "mxfp8"
     NVFP4 = "nvfp4"
     MXFP4 = "mxfp4"
+    GGUF = "gguf"
 
     def __str__(self) -> str:
         return self.value
@@ -105,3 +106,13 @@ def nvfp4_scheme(*, input_scale: bool) -> QuantScheme:
 
 def mxfp4_scheme() -> QuantScheme:
     return QuantScheme(QuantKind.MXFP4, WeightDesc("e2m1", (1, MX_GROUP), "e8m0"), {"weight", "weight_scale"})
+
+
+def gguf_scheme(types: Iterable[int]) -> QuantScheme:
+    """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order)."""
+    return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, None), {"weight"})
+
+
+def gguf_types(scheme: QuantScheme) -> tuple[int, ...]:
+    assert scheme.kind is QuantKind.GGUF, scheme
+    return tuple(int(t) for t in scheme.weight.elem.removeprefix("ggml:").split(","))
diff --git a/python/freetoken/models/gemma4/__init__.py b/python/freetoken/models/gemma4/__init__.py
index 9f41f080fa..4fcf15a373 100644
--- a/python/freetoken/models/gemma4/__init__.py
+++ b/python/freetoken/models/gemma4/__init__.py
@@ -1,11 +1,6 @@
 from .attention import Gemma4Attention
 from .config import parse_config
-from .gguf import (
-    dummy_q4_0_expert_sources,
-    iter_gguf_weights,
-    load_q4_0_expert_sources,
-    parse_gguf_config,
-)
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Gemma4ForCausalLM, Gemma4ForConditionalGeneration, Gemma4UnifiedForConditionalGeneration
 from .moe import Gemma4Router
 from .vision import Gemma4MultimodalEmbedder, Gemma4UnifiedVisionEmbedder, Gemma4VisionModel
@@ -32,6 +27,5 @@
     "iter_weights_parallel",
     "parse_gguf_config",
     "iter_gguf_weights",
-    "load_q4_0_expert_sources",
-    "dummy_q4_0_expert_sources",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py
index 437822b519..95f89c3ca4 100644
--- a/python/freetoken/models/gemma4/gguf.py
+++ b/python/freetoken/models/gemma4/gguf.py
@@ -6,12 +6,13 @@
 ``gemma4.config.parse_config`` -- only the source is GGUF KV metadata instead of a
 HF config object. transformers' own GGUF->config conversion is rejected by the
 gemma4 strict dataclass (per-layer ``num_key_value_heads`` array), so we read the
-metadata directly. ``expert_quant`` is set to ``"q4_0"`` to route the routed experts
-through the native-Q4_0 offload-cache path.
+metadata directly. Every tensor keeps its packed ggml type (``QuantKind.GGUF``); the
+routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
 """
 
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING, Iterator
 
 import torch
@@ -22,7 +23,8 @@
     RotaryConfig,
     SWAAttentionGroupConfig,
 )
-from freetoken.models.gguf.dequant import GGML_Q4_0, GGML_Q6_K, dequantize, row_bytes
+from freetoken.models.gguf.reader import gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
 
 if TYPE_CHECKING:
     from freetoken.models.gguf.config import GgufConfigShim
@@ -36,12 +38,9 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int:
     which drives the unrotated tail's frequencies to zero (ggml ops.cpp ``theta / ff``).
     Count the untouched entries to get the rotated width back.
     """
-    from freetoken.models.gguf.reader import iter_gguf_tensors
-
-    for t in iter_gguf_tensors(shim.model_path):
-        if t.name == "rope_freqs.weight":  # F32, so the packed bytes ARE the values
-            freqs = t.packed().reshape(-1).view(torch.float32)
-            return int((freqs == 1.0).sum().item()) * 2
+    info = gguf_tensor_index(shim.model_path).get("rope_freqs.weight")
+    if info is not None:
+        return int((float_tensor(info) == 1.0).sum().item()) * 2
     # A metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table. Every
     # released gemma4 sets partial_rotary_factor 0.25 on the full layers, and the converter
     # asserts rope_type == "proportional" (llama.cpp conversion/gemma.py).
@@ -112,8 +111,7 @@ def g(key: str):
         model_type="gemma4",
         architectures=list(shim.architectures),
         moe_enabled=True,
-        expert_quant="q4_0",
-        moe_weight_format="q4_0",
+        expert_quant="gguf",
         use_qk_norm=True,
         attn_sm_scale=1.0,
         final_logit_softcapping=float(g("final_logit_softcapping")),
@@ -144,8 +142,7 @@ def g(key: str):
 # Weight loading: GGUF tensor names -> FreeToken gemma4 module params.
 # --------------------------------------------------------------------------------------
 
-# Per-layer 1:1 norm/scale tensors (gguf suffix -> freetoken module-relative name). All
-# are tiny F32 tensors dequantized to bf16.
+# Per-layer 1:1 norm/scale tensors (gguf suffix -> freetoken module-relative name), all F32.
 _LAYER_SCALAR_MAP = {
     "attn_norm.weight": "input_layernorm.weight",
     "attn_q_norm.weight": "self_attn.q_norm.weight",
@@ -161,14 +158,39 @@ def g(key: str):
     "ffn_gate_inp.scale": "feed_forward.router.scale",
     "ffn_down_exps.scale": "feed_forward.router.per_expert_scale",
 }
-# Expert weight tensors (handled by the offload bank loader, not iter_gguf_weights).
-_EXPERT_SUFFIXES = ("ffn_gate_up_exps.weight", "ffn_down_exps.weight")
-
 
-def _to_bf16(t) -> torch.Tensor:
-    """Dequantize a GgufTensor (F32/F16/Q*) to a dense bf16 tensor of its torch shape."""
-    flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16)
-    return flat.reshape(t.shape)
+# checkpoint module leaf (gemma4 HF naming) -> GGUF tensor
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+    "experts.0.gate_proj": "ffn_gate_exps",
+    "experts.0.up_proj": "ffn_up_exps",
+    "experts.0.down_proj": "ffn_down_exps",
+}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name == "model.embed_tokens":
+        return ("token_embd",)
+    if name == "lm_head":
+        return ("output",) if "output.weight" in tensors else ()
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf == "self_attn.v_proj":
+        # the full-attention layers store no attn_v: k is reused as v
+        return (f"{blk}.attn_v",) if f"{blk}.attn_v.weight" in tensors else (f"{blk}.attn_k",)
+    if leaf in ("experts.0.gate_proj", "experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors:
+        return (f"{blk}.ffn_gate_up_exps",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
 
 
 def _require_tp1(what: str) -> None:
@@ -191,284 +213,43 @@ def iter_gguf_weights(
     include_moe_experts: bool,
     include_non_moe: bool,
 ) -> Iterator[tuple[str, torch.Tensor]]:
-    """Yield (param_name, tensor) for every non-expert gemma4 param.
-
-    Quantized projections (attention q/k/v/o, shared-MLP gate/up/down) and the
-    embedding stay in their native packed block layout and are yielded as ``.qweight``
-    (uint8); norms / router / per-layer scalars dequantize to bf16. q/k/v and gate/up
-    are fused by concatenating packed rows along the output dim (same input dim ->
-    same ``row_bytes``). Routed experts are served from the offload cache, so they are
-    skipped here (asserts the offload contract like the other MoE models).
-    """
-    from freetoken.models.gguf.reader import iter_gguf_tensors
+    """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
     from freetoken.utils import cached_load_hf_config
 
-    assert not include_moe_experts, (
-        "gemma4 GGUF stores experts as Q4_0 and only supports the offload backend; "
-        "experts are loaded into the offload cache via load_q4_0_expert_sources()."
-    )
-    assert include_non_moe
     _require_tp1("weight loading")
-
-    # Full-attention layers ship no attn_v (k reused as v); SWA layers do. Knowing which
-    # is which lets us emit the fused qkv as soon as its parts are present.
     config = parse_gguf_config(cached_load_hf_config(model_path))
-    k_eq_v_layers = {
-        lid
-        for lid in range(config.num_layers)
-        if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig)
-        and config.attention_group_for_layer(lid).k_eq_v
-    }
-
-    # Per-layer fusion buffers: layer -> {slot: packed[out, row_bytes]}.
-    qkv_buf: dict[int, dict[str, torch.Tensor]] = {}
-    gate_up_buf: dict[int, dict[str, torch.Tensor]] = {}
-
-    def layer_of(name: str) -> int:
-        return int(name.split(".")[1])
-
-    for t in iter_gguf_tensors(model_path):
-        name = t.name
-        if name == "token_embd.weight":
-            yield "model.embed_tokens.qweight", t.packed()  # Q6_K packed table
-            continue
-        if name == "output_norm.weight":
-            yield "model.norm.weight", _to_bf16(t)
-            continue
-        if name == "rope_freqs.weight":
-            continue  # rope frequencies recomputed in-engine
-        if not name.startswith("blk."):
-            continue
-        if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES):
-            continue  # routed experts -> offload banks
-
-        layer = layer_of(name)
-        suffix = name.split(".", 2)[2]  # after "blk.N."
-        base = f"model.layers.{layer}"
-
-        if suffix in _LAYER_SCALAR_MAP:
-            rel = _LAYER_SCALAR_MAP[suffix]
-            tensor = _to_bf16(t)
-            if rel.endswith("layer_scalar"):
-                tensor = tensor.reshape(1)
-            yield f"{base}.{rel}", tensor
-            continue
-
-        # Quantized projections: keep packed; fuse q/k/v and gate/up.
-        if suffix == "attn_q.weight":
-            qkv_buf.setdefault(layer, {})["q"] = t.packed()
-        elif suffix == "attn_k.weight":
-            qkv_buf.setdefault(layer, {})["k"] = t.packed()
-        elif suffix == "attn_v.weight":
-            qkv_buf.setdefault(layer, {})["v"] = t.packed()
-        elif suffix == "ffn_gate.weight":
-            gate_up_buf.setdefault(layer, {})["gate"] = t.packed()
-        elif suffix == "ffn_up.weight":
-            gate_up_buf.setdefault(layer, {})["up"] = t.packed()
-        elif suffix == "attn_output.weight":
-            yield f"{base}.self_attn.o_proj.qweight", t.packed()
-        elif suffix == "ffn_down.weight":
-            yield f"{base}.feed_forward.shared_mlp.down_proj.qweight", t.packed()
-        else:
-            raise ValueError(f"unmapped gemma4 GGUF tensor: {name}")
-
-        # Emit fused qkv once all parts are present: k_eq_v full layers reuse k as v
-        # (no attn_v tensor), SWA layers wait for their separate attn_v.
-        slots = qkv_buf.get(layer)
-        if slots is not None and "q" in slots and "k" in slots:
-            if layer in k_eq_v_layers:
-                v = slots["k"]
-            elif "v" in slots:
-                v = slots["v"]
-            else:
-                v = None
-            if v is not None:
-                yield f"{base}.self_attn.qkv_proj.qweight", torch.cat(
-                    [slots["q"], slots["k"], v], dim=0
-                )
-                del qkv_buf[layer]
-        gu = gate_up_buf.get(layer)
-        if gu is not None and "gate" in gu and "up" in gu:
-            yield f"{base}.feed_forward.shared_mlp.gate_up_proj.qweight", torch.cat(
-                [gu["gate"], gu["up"]], dim=0
-            )
-            del gate_up_buf[layer]
-
-    assert not qkv_buf, f"incomplete qkv groups: {sorted(qkv_buf)}"
-    assert not gate_up_buf, f"incomplete gate_up groups: {sorted(gate_up_buf)}"
-
-
-# --------------------------------------------------------------------------------------
-# Model layer swap: dense bf16 Linear/Embedding -> native GGUF-quant ops.
-# --------------------------------------------------------------------------------------
-
-
-def is_gguf_model(config: ModelConfig) -> bool:
-    """True when the model was parsed from a GGUF checkpoint (native-quant path)."""
-    return getattr(config, "moe_weight_format", None) == "q4_0"
-
-
-class GGUFTiedLMHead:
-    """Tied LM head over a native Q6_K embedding table (logits via ggml matmul).
-
-    Holds only a reference to the GGUF embedding (no params of its own -> empty
-    state_dict), mirroring ``ParallelLMHead`` with ``tie_word_embeddings``. TP=1 only.
-    """
-
-    def __init__(self, embedding, quant_type: int):
-        self._embedding = embedding
-        self._quant_type = quant_type
-
-    def state_dict(self, *, prefix: str = "", result=None):
-        return result if result is not None else {}
-
-    def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False):
-        state_dict.pop(f"{prefix}.weight", None)
-        state_dict.pop(f"{prefix}.bias", None)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        from freetoken.core import get_global_ctx
-        from freetoken.layers.gguf import fused_mul_mat_gguf
-
-        batch = get_global_ctx().batch
-        if batch.is_prefill:
-            indices = batch.attn_metadata.get_last_indices(batch.size)
-            x = x[indices].contiguous()
-        return fused_mul_mat_gguf(x, self._embedding.qweight, self._quant_type)
-
-
-def convert_gemma4_to_gguf(model, config: ModelConfig) -> None:
-    """In place: replace gemma4's dense projections + embedding with native GGUF ops.
-
-    Quantized in the checkpoint -> swapped: attention qkv/o, shared-MLP gate_up/down
-    (all Q4_0) and the token embedding (Q6_K, also the tied LM head). Left as dense
-    bf16 (F32 in the GGUF): the router gate, all RMSNorms, the per-layer scalars, and
-    the routed experts (served from the offload cache).
-    """
-    from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear
-
-    def swap_linear(owner, attr, quant_type=GGML_Q4_0):
-        lin = getattr(owner, attr)
-        out_features, in_features = lin.weight.shape
-        setattr(
-            owner,
-            attr,
-            GGUFLinear(in_features, out_features, quant_type, has_bias=lin.bias is not None),
-        )
-
-    inner = model.model
-    embed = GGUFEmbedding(
-        num_embeddings=config.vocab_size,
-        embedding_dim=config.hidden_size,
-        quant_type=GGML_Q6_K,
-        embed_scale=config.embedding_scale,
-    )
-    inner.embed_tokens = embed
-
-    for layer in inner.layers.op_list:
-        swap_linear(layer.self_attn, "qkv_proj")
-        swap_linear(layer.self_attn, "o_proj")
-        swap_linear(layer.feed_forward.shared_mlp, "gate_up_proj")
-        swap_linear(layer.feed_forward.shared_mlp, "down_proj")
-
-    if config.tie_word_embeddings:
-        model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K)
-
-
-# --------------------------------------------------------------------------------------
-# Routed-expert host banks (native Q4_0) for the offload cache.
-# --------------------------------------------------------------------------------------
-
-def _q4_0_expert_specs(config: ModelConfig) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
-    E = config.num_experts
-    H, I = config.hidden_size, config.moe_intermediate_size
-    return {
-        "gate_up": ((E, 2 * I, row_bytes(H, GGML_Q4_0)), torch.uint8),
-        "down": ((E, H, row_bytes(I, GGML_Q4_0)), torch.uint8),
-    }
-
-
-def load_q4_0_expert_sources(
-    model_path: str, config: ModelConfig, *, layer_sink=None
-) -> dict[str, list[torch.Tensor]]:
-    """Per-layer host banks of the routed experts' native Q4_0 block bytes.
-
-    ``gate_up`` is one ``[E, 2I, H//32*18]`` tensor per layer and ``down`` one
-    ``[E, H, I//32*18]`` per layer (independent :class:`HostBank` allocations) -- each
-    expert's packed rows verbatim from the GGUF (no dequant), whole layers arriving in
-    one shot (gate_up + down = 2 writes/layer), so the offload cache streams whole
-    experts to the ggml MoE kernels.
-
-    ``layer_sink=None`` (serving): pin each layer's two banks as they complete via an
-    internally-owned :class:`PinPipeline` (or, on a CUDA-less host, allocate the mmap
-    banks but never pin -- the CPU executor reads them pageable). ``layer_sink`` given
-    (converter): the completion tracker fires into it instead -- nothing here is pinned,
-    and the sink may release banks it has written out, so the returned tensors are only
-    valid until then (the caller owns that tradeoff).
-    """
-    from freetoken.models.gguf.reader import iter_gguf_tensors
-    from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks
-
-    _require_tp1("expert banks")
-    L, E = config.num_layers, config.num_experts
-    H, I = config.hidden_size, config.moe_intermediate_size
-    h_bytes, i_bytes = row_bytes(H, GGML_Q4_0), row_bytes(I, GGML_Q4_0)
-    hb = alloc_layer_banks(_q4_0_expert_specs(config), L)  # lazy anon mmaps (unpinned)
-    banks = {name: [b.tensor for b in hb[name]] for name in hb}
-    seen_gu, seen_dn = set(), set()
-
-    def _load(sink) -> None:
-        tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None  # gate_up + down
-        for t in iter_gguf_tensors(model_path):
-            if not t.name.startswith("blk."):
-                continue
-            layer = int(t.name.split(".")[1])
-            if t.name.endswith("ffn_gate_up_exps.weight"):
-                banks["gate_up"][layer].copy_(t.packed().reshape(E, 2 * I, h_bytes))
-                seen_gu.add(layer)
-            elif t.name.endswith("ffn_down_exps.weight"):
-                banks["down"][layer].copy_(t.packed().reshape(E, H, i_bytes))
-                seen_dn.add(layer)
-            else:
-                continue
-            if tracker is not None:
-                tracker.note(layer)
-
-    if layer_sink is not None:
-        _load(layer_sink)
-    elif torch.cuda.is_available():
-        with PinPipeline() as pins:
-            _load(pins)
-    else:
-        _load(None)  # CUDA-less: mmap banks stay pageable, never pinned
-
-    want = set(range(L))
-    assert seen_gu == want and seen_dn == want, (
-        f"missing Q4_0 expert layers: gate_up {sorted(want - seen_gu)}, "
-        f"down {sorted(want - seen_dn)}"
-    )
-    return banks
-
-
-def dummy_q4_0_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tensor]]:
-    """Random Q4_0 expert banks shaped like ``load_q4_0_expert_sources`` output."""
-    from freetoken.moe.host_banks import alloc_layer_banks, pin_banks
-
-    L = config.num_layers
-    hb = alloc_layer_banks(_q4_0_expert_specs(config), L)
-    banks = {name: [b.tensor for b in hb[name]] for name in hb}
-    for t in banks["gate_up"] + banks["down"]:
-        t.random_(0, 256)
-    if torch.cuda.is_available():
-        pin_banks(hb)  # match the other dummies: pin-after-fill (no-op mmap fill on CPU-only)
-    return banks
+    index = gguf_tensor_index(model_path)
+    hidden = config.hidden_size
+
+    def t(name: str):
+        return index[name if name in index else f"{name}.weight"]
+
+    embed = t("token_embd")
+    yield from fused("model.embed_tokens", [embed], [matrix(embed)], hidden)
+    if "output.weight" in index:
+        out = t("output")
+        yield from fused("lm_head", [out], [matrix(out)], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(config.num_layers):
+        blk, base = f"blk.{layer}", f"model.layers.{layer}"
+        for suffix, rel in _LAYER_SCALAR_MAP.items():
+            tensor = float_tensor(index[f"{blk}.{suffix}"])
+            yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor
+        q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k")
+        qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)]
+        yield from fused(f"{base}.self_attn.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        o = t(f"{blk}.attn_output")
+        yield from fused(f"{base}.self_attn.o_proj", [o], [matrix(o)], int(o.ne[0]))
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{base}.feed_forward.shared_mlp.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{base}.feed_forward.shared_mlp.down_proj", [down], [matrix(down)], int(down.ne[0]))
 
 
 __all__ = [
     "parse_gguf_config",
     "iter_gguf_weights",
-    "convert_gemma4_to_gguf",
-    "is_gguf_model",
-    "load_q4_0_expert_sources",
-    "dummy_q4_0_expert_sources",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/gemma4/model.py b/python/freetoken/models/gemma4/model.py
index 5c91b90571..1b9d0da7d9 100644
--- a/python/freetoken/models/gemma4/model.py
+++ b/python/freetoken/models/gemma4/model.py
@@ -61,6 +61,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
             embed_scale=config.embedding_scale,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
@@ -91,13 +93,6 @@ def __init__(self, config: ModelConfig):
         self._final_logit_softcapping = config.final_logit_softcapping
         super().__init__()
 
-        # GGUF checkpoints carry native block-quantized weights: swap the dense
-        # projections + embedding for GGUF-quant ops (experts stay on the offload cache).
-        from .gguf import convert_gemma4_to_gguf, is_gguf_model
-
-        if is_gguf_model(config):
-            convert_gemma4_to_gguf(self, config)
-
     def forward(self) -> torch.Tensor:
         output = self.model.forward(get_global_ctx().batch.input_ids)
         logits = self.lm_head.forward(output)
diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py
index ba990c007b..08293c7eab 100644
--- a/python/freetoken/models/gemma4/weight.py
+++ b/python/freetoken/models/gemma4/weight.py
@@ -196,7 +196,7 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None:
                         continue
 
                     # Per-expert NVFP4 tensors go to the offload cache (expert pieces),
-                    # not this dense pass; fused bf16/q4_0 experts lack ".experts.<int>." so are unaffected.
+                    # not this dense pass; fused bf16 experts lack ".experts.<int>." so are unaffected.
                     if _NVFP4_EXPERT_RE.search(raw_name):
                         continue
 
diff --git a/python/freetoken/models/gguf/__init__.py b/python/freetoken/models/gguf/__init__.py
index 75e40e4faa..20a97b4610 100644
--- a/python/freetoken/models/gguf/__init__.py
+++ b/python/freetoken/models/gguf/__init__.py
@@ -1,26 +1,19 @@
-from .dequant import GGML_NAME, dequantize
 from .reader import (
     FTW_METADATA_GGUF,
-    GgufTensor,
     gguf_architecture,
     gguf_config_source,
     gguf_tensor_names,
     is_gguf_path,
-    iter_gguf_tensors,
     load_gguf_metadata,
     write_metadata_gguf,
 )
 
 __all__ = [
-    "GGML_NAME",
-    "dequantize",
     "FTW_METADATA_GGUF",
-    "GgufTensor",
     "gguf_architecture",
     "gguf_config_source",
     "gguf_tensor_names",
     "is_gguf_path",
-    "iter_gguf_tensors",
     "load_gguf_metadata",
     "write_metadata_gguf",
 ]
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 63b1a18b97..4eb2a999db 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -18,6 +18,7 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "qwen4exp": "Qwen4ExpGGUFForCausalLM",
 }
 
 
diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py
deleted file mode 100644
index 77c3ea0102..0000000000
--- a/python/freetoken/models/gguf/dequant.py
+++ /dev/null
@@ -1,153 +0,0 @@
-"""GGML block-quant dequantization in pure torch (the formats this repo's GGUF
-checkpoints use: Q4_0, Q6_K, plus trivial F32/F16/BF16).
-
-This is the *reference / CPU* path, NOT the engine's hot path: GGUF weights stay
-packed and are dequantized inside the borrowed ggml CUDA kernels (see
-``freetoken.kernel.gguf``). These routines are used only to (a) materialize the few
-dense F32/F16 tensors at load (norms, scales, router) via :func:`dequantize`, and
-(b) cross-check the CUDA kernels in tests. The ``BLOCK_SHAPE`` table and
-:func:`row_bytes` are the type metadata the packed (kernel) path also relies on.
-
-Each ``dequant_*`` takes the raw little-endian bytes as a ``uint8`` tensor whose
-final axis spans whole blocks, and returns the values in *storage order* (ggml's
-fastest axis first); the caller reshapes to the torch shape (``dims[::-1]``). The
-math mirrors ``ggml-quants.c``.
-"""
-
-from __future__ import annotations
-
-import torch
-
-# ggml_type enum values (subset present in these checkpoints).
-GGML_F32 = 0
-GGML_F16 = 1
-GGML_Q4_0 = 2
-GGML_Q8_0 = 8
-GGML_Q6_K = 14
-GGML_BF16 = 30
-
-# (block numel, bytes per block) per ggml type.
-BLOCK_SHAPE: dict[int, tuple[int, int]] = {
-    GGML_F32: (1, 4),
-    GGML_F16: (1, 2),
-    GGML_BF16: (1, 2),
-    GGML_Q4_0: (32, 18),
-    GGML_Q8_0: (32, 34),
-    GGML_Q6_K: (256, 210),
-}
-
-GGML_NAME = {
-    GGML_F32: "F32",
-    GGML_F16: "F16",
-    GGML_BF16: "BF16",
-    GGML_Q4_0: "Q4_0",
-    GGML_Q8_0: "Q8_0",
-    GGML_Q6_K: "Q6_K",
-}
-
-
-def row_bytes(numel: int, ggml_type: int) -> int:
-    """Packed byte length of one row of ``numel`` elements in ``ggml_type`` blocks.
-
-    Single source of truth for the ``numel // block * type_size`` math shared by the
-    packed-weight ops (``GGUFLinear``/``GGUFEmbedding``) and the expert bank loaders.
-    """
-    block, type_size = BLOCK_SHAPE[ggml_type]
-    assert numel % block == 0, (
-        f"{numel} not a multiple of block {block} for {GGML_NAME.get(ggml_type, ggml_type)}"
-    )
-    return numel // block * type_size
-
-
-def _f16_scales(raw: torch.Tensor, lo: int, hi: int) -> torch.Tensor:
-    """Reinterpret bytes ``[lo:hi]`` (2 per block) of each block row as fp16 -> fp32 [N,1]."""
-    return raw[:, lo:hi].contiguous().view(torch.float16).to(torch.float32)
-
-
-def dequant_q4_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
-    """Q4_0: per 32-elem block = fp16 scale ``d`` + 16 packed nibbles; ``w = d*(q-8)``.
-
-    Byte ``j`` of the 16 holds element ``j`` in its low nibble and ``j+16`` in its high
-    nibble, so storage order within the block is ``[lo0..lo15, hi0..hi15]``.
-    """
-    raw = raw.reshape(-1, 18)
-    d = _f16_scales(raw, 0, 2)  # [N,1]
-    qs = raw[:, 2:18]  # [N,16] uint8
-    lo = (qs & 0x0F).to(torch.float32)
-    hi = (qs >> 4).to(torch.float32)
-    q = torch.cat([lo, hi], dim=1)  # [N,32]
-    return ((q - 8.0) * d).reshape(-1).to(out_dtype)
-
-
-def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
-    """Q6_K: 256-elem super-block = 128B low nibbles + 64B high 2-bits + 16 int8
-    sub-scales + fp16 ``d``. Direct vectorization of ggml's two-half loop."""
-    raw = raw.reshape(-1, 210)
-    n = raw.shape[0]
-    ql = raw[:, 0:128]  # [n,128]
-    qh = raw[:, 128:192]  # [n,64]
-    sc = raw[:, 192:208].view(torch.int8).to(torch.float32)  # [n,16]
-    d = _f16_scales(raw, 208, 210)  # [n,1]
-
-    y = torch.empty((n, 256), dtype=torch.float32, device=raw.device)
-    # l in 0..15 -> is=0; l in 16..31 -> is=1 (per ggml: is = l/16).
-    is_idx = (torch.arange(32, device=raw.device) // 16)  # [32] in {0,1}
-    for h in range(2):  # two 128-elem halves of the super-block
-        qlh = ql[:, h * 64:(h + 1) * 64]  # [n,64]
-        qhh = qh[:, h * 32:(h + 1) * 32]  # [n,32]
-        sch = sc[:, h * 8:(h + 1) * 8]  # [n,8]
-        a = qlh[:, 0:32].to(torch.int32)  # ql[l]
-        b = qlh[:, 32:64].to(torch.int32)  # ql[l+32]
-        hb = qhh.to(torch.int32)  # qh[l]
-        q1 = ((a & 0x0F) | (((hb >> 0) & 3) << 4)) - 32
-        q2 = ((b & 0x0F) | (((hb >> 2) & 3) << 4)) - 32
-        q3 = ((a >> 4) | (((hb >> 4) & 3) << 4)) - 32
-        q4 = ((b >> 4) | (((hb >> 6) & 3) << 4)) - 32
-        s1 = sch.index_select(1, is_idx + 0).to(torch.float32)
-        s2 = sch.index_select(1, is_idx + 2).to(torch.float32)
-        s3 = sch.index_select(1, is_idx + 4).to(torch.float32)
-        s4 = sch.index_select(1, is_idx + 6).to(torch.float32)
-        base = h * 128
-        y[:, base + 0:base + 32] = d * s1 * q1.to(torch.float32)
-        y[:, base + 32:base + 64] = d * s2 * q2.to(torch.float32)
-        y[:, base + 64:base + 96] = d * s3 * q3.to(torch.float32)
-        y[:, base + 96:base + 128] = d * s4 * q4.to(torch.float32)
-    return y.reshape(-1).to(out_dtype)
-
-
-_DEQUANT = {
-    GGML_Q4_0: dequant_q4_0,
-    GGML_Q6_K: dequant_q6_k,
-}
-
-
-def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> torch.Tensor:
-    """Dequantize ``raw`` (uint8) of any supported ggml type to flat ``out_dtype``."""
-    if ggml_type == GGML_F32:
-        return raw.view(torch.float32).to(out_dtype)
-    if ggml_type == GGML_F16:
-        return raw.view(torch.float16).to(out_dtype)
-    if ggml_type == GGML_BF16:
-        return raw.view(torch.bfloat16).to(out_dtype)
-    fn = _DEQUANT.get(ggml_type)
-    if fn is None:
-        raise NotImplementedError(
-            f"dequant for ggml type {GGML_NAME.get(ggml_type, ggml_type)} not implemented"
-        )
-    return fn(raw, out_dtype)
-
-
-__all__ = [
-    "GGML_F32",
-    "GGML_F16",
-    "GGML_BF16",
-    "GGML_Q4_0",
-    "GGML_Q8_0",
-    "GGML_Q6_K",
-    "GGML_NAME",
-    "BLOCK_SHAPE",
-    "row_bytes",
-    "dequant_q4_0",
-    "dequant_q6_k",
-    "dequantize",
-]
diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py
new file mode 100644
index 0000000000..0da56ceb4f
--- /dev/null
+++ b/python/freetoken/models/gguf/experts.py
@@ -0,0 +1,104 @@
+"""Routed-expert banks straight from a GGUF's ``blk.N.ffn_{gate,up,down}_exps`` tensors, in their packed ggml rows: pinned host banks for the offload cache, or device banks for resident experts.
+
+Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts keep their dtype ([E, 2I, H] and [E, H, I]).
+A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+
+import torch
+
+from freetoken.layers.quantization import QuantKind
+from freetoken.layers.quantization.ggml_types import float_dtype, mmq_tail_bytes, pad_rows, type_name
+from freetoken.utils import init_logger
+
+from .reader import GgufTensorInfo, gguf_tensor_index
+
+logger = init_logger(__name__)
+
+
+def _expert_rows(info: GgufTensorInfo) -> torch.Tensor:
+    """[E, rows, row_bytes] packed ggml rows, or [E, rows, K] for a float tensor."""
+    k, rows, experts = info.ne
+    raw = torch.from_numpy(info.data()).view(experts, rows, info.nbytes // (experts * rows))
+    dtype = float_dtype(info.ggml_type)
+    return raw if dtype is None else raw.view(dtype)
+
+
+def _drop_page_cache(info: GgufTensorInfo) -> None:
+    try:
+        fd = os.open(info.path, os.O_RDONLY)
+        try:
+            os.posix_fadvise(fd, info.offset, info.nbytes, os.POSIX_FADV_DONTNEED)
+        finally:
+            os.close(fd)
+    except OSError:
+        pass
+
+
+def _device_bank(shape: tuple[int, ...], k: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    """An [E, rows, row] device bank; a zero row follows the last one when ggml's MMQ tiles read past it."""
+    experts, rows, row = shape
+    flat = torch.empty(experts * rows + pad_rows(k), row, dtype=dtype, device=device)
+    flat[experts * rows :].zero_()
+    return flat[: experts * rows].view(shape)
+
+
+def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int = 0, dummy: bool = False, device: torch.device | None = None):
+    """``ExpertBanks`` with one bank per (layer, role): pinned host banks for the offload cache, filled from the file and pinned while later layers load, or resident banks on ``device``."""
+    from freetoken.moe.expert_banks import ExpertBanks
+    from freetoken.moe.host_banks import HostBank, PinPipeline
+
+    index = gguf_tensor_index(model_path)
+    sources: dict[str, list[torch.Tensor]] = {"gate_up": [], "down": []}
+    tails: dict[str, list[int]] = {"gate_up": [], "down": []}
+    tags = []
+    with PinPipeline() if device is None else contextlib.nullcontext() as pins:
+        for layer in range(num_layers):
+            blk = f"blk.{first_layer + layer}"
+            down = index[f"{blk}.ffn_down_exps.weight"]
+            if f"{blk}.ffn_gate_up_exps.weight" in index:
+                gate = index[f"{blk}.ffn_gate_up_exps.weight"]
+                parts = [gate]
+            else:
+                gate, up = index[f"{blk}.ffn_gate_exps.weight"], index[f"{blk}.ffn_up_exps.weight"]
+                if gate.ggml_type != up.ggml_type:
+                    raise NotImplementedError(f"{blk}: gate and up experts have different ggml types")
+                parts = [gate, up]
+            tags.append(("ggml", type_name(gate.ggml_type), type_name(down.ggml_type)))
+            tails["gate_up"].append(mmq_tail_bytes(gate.ggml_type, gate.ne[0]))
+            tails["down"].append(mmq_tail_bytes(down.ggml_type, down.ne[0]))
+            rows = [_expert_rows(info) for info in parts]
+            d = _expert_rows(down)
+            gate_up_shape = (rows[0].shape[0], sum(r.shape[1] for r in rows), rows[0].shape[2])
+            if device is None:
+                host = {"gate_up": HostBank(gate_up_shape, rows[0].dtype), "down": HostBank(tuple(d.shape), d.dtype)}
+                banks = {role: bank.tensor for role, bank in host.items()}
+            else:
+                banks = {
+                    "gate_up": _device_bank(gate_up_shape, gate.ne[0], rows[0].dtype, device),
+                    "down": _device_bank(tuple(d.shape), down.ne[0], d.dtype, device),
+                }
+            if dummy:
+                banks["gate_up"].zero_()
+                banks["down"].zero_()
+            else:
+                start = 0
+                for r in rows:
+                    banks["gate_up"][:, start : start + r.shape[1]].copy_(r)
+                    start += r.shape[1]
+                banks["down"].copy_(d)
+                for info in (*parts, down):
+                    _drop_page_cache(info)
+            if device is None:
+                pins(layer, host)
+            for role, bank in banks.items():
+                sources[role].append(bank)
+    total = sum(t.numel() * t.element_size() for ts in sources.values() for t in ts)
+    logger.info_rank0(f"GGUF expert banks: {num_layers} layers, {total / 2**30:.1f} GiB {'on ' + str(device) if device is not None else 'in host RAM'}")
+    return ExpertBanks("gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails)
+
+
+__all__ = ["load_gguf_expert_banks"]
diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py
index b950d92981..2715f12909 100644
--- a/python/freetoken/models/gguf/reader.py
+++ b/python/freetoken/models/gguf/reader.py
@@ -2,22 +2,20 @@
 
 Thin layer over ``gguf.GGUFReader`` (gguf-py). Metadata is read into a plain dict
 keyed by the GGUF field name (``general.architecture``, ``gemma4.block_count`` ...);
-tensors are exposed as ``GgufTensor`` records carrying the *torch* shape (ggml dims
-reversed), the ggml quant type, and a zero-copy ``uint8`` view of the packed block
-bytes laid out as ``[rows, row_bytes]`` (rows = product of all but the fastest ggml
-dim; row_bytes spans whole quant blocks of the fastest dim).
+tensors are indexed as ``GgufTensorInfo`` records (ggml dims, ggml type, and where the
+packed bytes sit in the memory-mapped file).
 """
 
 from __future__ import annotations
 
 import functools
 import os
+import re
 import struct
 from dataclasses import dataclass
-from typing import Any, Iterator
+from typing import Any
 
 import numpy as np
-import torch
 
 
 def is_gguf_path(model_path: str) -> bool:
@@ -65,7 +63,7 @@ def write_metadata_gguf(source_gguf: str, dest_path: str) -> None:
     Validates by re-parsing: the copy must list zero tensors and expose the identical KV
     key set (the KV *bytes* are copied verbatim, so identical keys imply identical values).
     """
-    import gguf
+    from freetoken.third_party import gguf
 
     reader = gguf.GGUFReader(source_gguf)
     assert reader.tensors, f"{source_gguf}: no tensors to bound the KV section"
@@ -98,20 +96,6 @@ def write_metadata_gguf(source_gguf: str, dest_path: str) -> None:
     )
 
 
-@dataclass(frozen=True)
-class GgufTensor:
-    name: str
-    shape: tuple[int, ...]  # torch order (ggml dims reversed)
-    ggml_type: int
-    rows: int  # product of shape[:-1] over the *ggml* layout = blocks-major rows
-    row_bytes: int  # packed bytes per row (whole quant blocks of the fastest dim)
-    _raw: np.ndarray  # uint8 view, shape [rows, row_bytes]
-
-    def packed(self) -> torch.Tensor:
-        """Zero-copy ``[rows, row_bytes]`` uint8 tensor of the native block bytes."""
-        return torch.from_numpy(self._raw)
-
-
 def _field_value(reader, name: str) -> Any:
     field = reader.fields.get(name)
     if field is None:
@@ -121,7 +105,7 @@ def _field_value(reader, name: str) -> Any:
 
 @functools.cache
 def _reader(model_path: str):
-    import gguf
+    from freetoken.third_party import gguf
 
     return gguf.GGUFReader(model_path)
 
@@ -133,6 +117,55 @@ def load_gguf_metadata(model_path: str) -> dict[str, Any]:
     return {name: field.contents() for name, field in reader.fields.items()}
 
 
+_SPLIT_RE = re.compile(r"-(\d{5})-of-(\d{5})\.gguf$")
+
+
+@functools.cache
+def gguf_split_paths(model_path: str) -> tuple[str, ...]:
+    """Every file of a split GGUF (``name-00001-of-00002.gguf``), in order; a single file is its own list."""
+    count = int(_field_value(_reader(model_path), "split.count") or 1)
+    if count == 1:
+        return (model_path,)
+    match = _SPLIT_RE.search(model_path)
+    if match is None:
+        raise ValueError(f"{model_path} is one of {count} GGUF splits but is not named like name-00001-of-{count:05d}.gguf")
+    stem = model_path[: match.start()]
+    paths = tuple(f"{stem}-{i:05d}-of-{count:05d}.gguf" for i in range(1, count + 1))
+    missing = [p for p in paths if not os.path.isfile(p)]
+    if missing:
+        raise FileNotFoundError(f"GGUF split files missing: {missing}")
+    return paths
+
+
+@dataclass(frozen=True)
+class GgufTensorInfo:
+    """Where one tensor's packed bytes sit: ``ne`` is in ggml order (fastest dim first)."""
+
+    name: str
+    path: str
+    offset: int
+    nbytes: int
+    ggml_type: int
+    ne: tuple[int, ...]
+
+    def data(self) -> np.ndarray:
+        """Zero-copy uint8 view of the bytes (memory-mapped)."""
+        reader = _reader(self.path)
+        return reader.data[self.offset : self.offset + self.nbytes]
+
+
+@functools.cache
+def gguf_tensor_index(model_path: str) -> dict[str, GgufTensorInfo]:
+    """Every tensor of every split, by name."""
+    index: dict[str, GgufTensorInfo] = {}
+    for path in gguf_split_paths(model_path):
+        for t in _reader(path).tensors:
+            index[t.name] = GgufTensorInfo(
+                t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape)
+            )
+    return index
+
+
 def gguf_architecture(model_path: str) -> str:
     arch = _field_value(_reader(model_path), "general.architecture")
     if arch is None:
@@ -140,39 +173,8 @@ def gguf_architecture(model_path: str) -> str:
     return str(arch)
 
 
-def iter_gguf_tensors(model_path: str) -> Iterator[GgufTensor]:
-    """Yield every tensor with its torch shape, ggml type, and packed block bytes."""
-    import gguf
-
-    reader = _reader(model_path)
-    for t in reader.tensors:
-        ne = [int(s) for s in t.shape]  # ggml order, fastest dim first
-        torch_shape = tuple(reversed(ne))
-        block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type]
-        n_fast = ne[0]
-        if n_fast % block != 0:
-            raise ValueError(
-                f"{t.name}: fastest dim {n_fast} not a multiple of block {block} "
-                f"for {t.tensor_type.name}"
-            )
-        row_bytes = n_fast // block * type_size
-        rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1
-        # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays;
-        # normalize everything to a flat byte view before shaping into [rows, row_bytes].
-        flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8)
-        raw = flat.reshape(rows, row_bytes)
-        yield GgufTensor(
-            name=t.name,
-            shape=torch_shape,
-            ggml_type=int(t.tensor_type),
-            rows=rows,
-            row_bytes=row_bytes,
-            _raw=raw,
-        )
-
-
 def gguf_tensor_names(model_path: str) -> set[str]:
-    return {t.name for t in _reader(model_path).tensors}
+    return {t.name for path in gguf_split_paths(model_path) for t in _reader(path).tensors}
 
 
 __all__ = [
@@ -181,9 +183,10 @@ def gguf_tensor_names(model_path: str) -> set[str]:
     "OUTPUT_WEIGHT_PRESENT_KV",
     "gguf_config_source",
     "write_metadata_gguf",
-    "GgufTensor",
+    "GgufTensorInfo",
+    "gguf_split_paths",
+    "gguf_tensor_index",
     "load_gguf_metadata",
     "gguf_architecture",
-    "iter_gguf_tensors",
     "gguf_tensor_names",
 ]
diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py
new file mode 100644
index 0000000000..4cba1d082e
--- /dev/null
+++ b/python/freetoken/models/gguf/tensors.py
@@ -0,0 +1,56 @@
+"""GGUF tensors as state-dict entries: float tensors in torch order, quantized ones as packed ggml rows grouped into the GGUF Linear's per-type segments."""
+
+from __future__ import annotations
+
+from typing import Iterator
+
+import numpy as np
+import torch
+
+from freetoken.layers.quantization.ggml_types import is_quantized, pad_rows, segment_param, segments
+
+from .reader import GgufTensorInfo
+
+_FLOAT_DTYPES = {0: torch.float32, 1: torch.float16, 30: torch.bfloat16}
+
+
+def float_tensor(info: GgufTensorInfo) -> torch.Tensor:
+    """A float tensor in torch order (ggml dims reversed)."""
+    return torch.from_numpy(np.array(info.data())).view(_FLOAT_DTYPES[info.ggml_type]).view(*reversed(info.ne))
+
+
+def packed_rows(info: GgufTensorInfo) -> torch.Tensor:
+    """Packed ggml rows ``[rows, row_bytes]`` of a quantized tensor."""
+    rows = int(np.prod(info.ne[1:]))
+    return torch.from_numpy(info.data()).view(rows, info.nbytes // rows)
+
+
+def matrix(info: GgufTensorInfo) -> torch.Tensor:
+    return packed_rows(info) if is_quantized(info.ggml_type) else float_tensor(info)
+
+
+def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k: int, pad_to: int = 0) -> Iterator[tuple[str, torch.Tensor]]:
+    """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed)."""
+    types = [i.ggml_type for i in infos]
+    pad = (-sum(p.shape[0] for p in parts)) % pad_to if pad_to else 0
+    if not any(is_quantized(t) for t in types):
+        rows = [p.to(torch.bfloat16) for p in parts]
+        if pad:
+            rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=torch.bfloat16))
+        yield f"{prefix}.weight", torch.cat(rows, dim=0)
+        return
+    segs = segments(tuple(types), tuple(p.shape[0] for p in parts))
+    for i, seg in enumerate(segs):
+        rows = parts[seg.first_part : seg.first_part + seg.num_parts]
+        zeros = pad if i == len(segs) - 1 else 0
+        if is_quantized(seg.ggml_type):
+            zeros += pad_rows(k)
+            dtype = torch.uint8
+        else:
+            rows = [r.to(torch.bfloat16) for r in rows]
+            dtype = torch.bfloat16
+        tail = [torch.zeros(zeros, rows[0].shape[1], dtype=dtype)] if zeros else []
+        yield f"{prefix}.{segment_param(i)}", torch.cat([*rows, *tail], dim=0)
+
+
+__all__ = ["float_tensor", "fused", "matrix", "packed_rows"]
diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index 6d5481c177..e6a5d4aca0 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -13,7 +13,7 @@
 from .reader import gguf_architecture, load_gguf_metadata
 
 # GGUF architecture -> transformers GGUF tokenizer-converter key.
-_TOKENIZER_ARCH = {"gemma4": "gemma4_text"}
+_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen4exp": "qwen3"}
 
 
 def load_gguf_tokenizer(model_path: str):
diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py
index 0fe4a4ffa8..15f866cad2 100644
--- a/python/freetoken/models/qwen4_exp/__init__.py
+++ b/python/freetoken/models/qwen4_exp/__init__.py
@@ -10,6 +10,7 @@
 """
 
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration
 from .weight import (
     ftw_side_files,
@@ -33,4 +34,7 @@
     "load_ple_table",
     "parse_config",
     "iter_expert_pieces",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index e413786266..b1e5f9aa1d 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -43,6 +43,8 @@ class Qwen4ExpArgs:
     index_budget: int
     index_ratio: int
     image_token_id: int | None = None
+    # GGUF stores ssm_out with its value-head columns in llama.cpp's tiled order
+    gdn_out_heads_tiled: bool = False
 
     @property
     def index_topk_blocks(self) -> int:
diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py
index ed49dd42f0..48c92a8e8a 100644
--- a/python/freetoken/models/qwen4_exp/gdn.py
+++ b/python/freetoken/models/qwen4_exp/gdn.py
@@ -32,9 +32,10 @@ class Qwen4ExpGatedDeltaNet(BaseOP):
     def __init__(
         self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim,
         conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "sigmoid",
-        *, quant_config: QuantConfig | None = None, prefix: str = "",
+        *, quant_config: QuantConfig | None = None, prefix: str = "", out_heads_tiled: bool = False,
     ):
         self.layer_id = layer_id
+        self._out_heads_tiled = out_heads_tiled
         # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as
         # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the
         # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head
@@ -197,6 +198,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         core_out = core_out.reshape(-1, self.head_v_dim)
         z = z.reshape(-1, self.head_v_dim)
         out = self.norm.forward(core_out, z).reshape(total, -1)
+        if self._out_heads_tiled:
+            # grouped value heads (k-head major) -> tiled (k-head minor), the column order of a GGUF ssm_out
+            r = self.num_v_heads // self.num_k_heads
+            out = out.view(total, self.num_k_heads, r, self.head_v_dim).transpose(1, 2).reshape(total, -1)
         return self.out_proj.forward(out)
 
 
diff --git a/python/freetoken/models/qwen4_exp/gguf.py b/python/freetoken/models/qwen4_exp/gguf.py
new file mode 100644
index 0000000000..7909324b2c
--- /dev/null
+++ b/python/freetoken/models/qwen4_exp/gguf.py
@@ -0,0 +1,264 @@
+"""Qwen3.8-Flash-Next from a llama.cpp ``qwen4exp`` GGUF: config from the GGUF metadata, tensor names, and the converter's transforms undone.
+
+The llama.cpp converter (conversion/qwen.py, conversion/qwen4exp.py) stores every zero-centred norm as ``w + 1``, ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, the indexer projection split into q and k, and the GDN value heads in tiled order (k-head minor). Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead (``Qwen4ExpArgs.gdn_out_heads_tiled``).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.layers.quantization.ggml_types import is_quantized
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen4exp"
+
+# HF config fields the GGUF metadata does not carry (Qwen3.8-Flash-Next's own config.json values)
+_FAMILY_DEFAULTS = dict(
+    hidden_act="silu",
+    output_gate_type="sigmoid",
+    indexer_kv_heads=1,
+    ngram_vocab_size_base=20000000,
+    make_ngram_vocab_size_divisible_by=128,
+    split_ngram_parts=128,
+    norm_topk_prob=True,
+)
+
+_LAYER_LEAVES: dict[str, tuple[str, ...]] = {
+    "linear_attn.in_proj_qkv": ("attn_qkv",),
+    "linear_attn.in_proj_z": ("attn_gate",),
+    "linear_attn.in_proj_b": ("ssm_beta",),
+    "linear_attn.in_proj_a": ("ssm_alpha",),
+    "linear_attn.out_proj": ("ssm_out",),
+    "self_attn.q_proj": ("attn_q",),
+    "self_attn.k_proj": ("attn_k",),
+    "self_attn.v_proj": ("attn_v",),
+    "self_attn.o_proj": ("attn_output",),
+    "self_attn.indexer.index_qk_proj": ("indexer.q_proj", "indexer.k_proj"),
+    "mlp.gate": ("ffn_gate_inp",),
+    "mlp.shared_expert_gate": ("ffn_gate_inp_shexp",),
+    "mlp.shared_expert.gate_proj": ("ffn_gate_shexp",),
+    "mlp.shared_expert.up_proj": ("ffn_up_shexp",),
+    "mlp.shared_expert.down_proj": ("ffn_down_shexp",),
+    "mlp.experts.0.gate_proj": ("ffn_gate_exps",),
+    "mlp.experts.0.up_proj": ("ffn_up_exps",),
+    "mlp.experts.0.down_proj": ("ffn_down_exps",),
+    "attn_hyper_connection.input_mix_weight_down": ("hc_attn_down",),
+    "attn_hyper_connection.block_inject_weight": ("hc_attn_inject",),
+    "attn_hyper_connection.input_mix_weight_up": ("hc_attn_up",),
+    "mlp_hyper_connection.input_mix_weight_down": ("hc_ffn_down",),
+    "mlp_hyper_connection.block_inject_weight": ("hc_ffn_inject",),
+    "mlp_hyper_connection.input_mix_weight_up": ("hc_ffn_up",),
+    "ple.key_proj": ("ple_key",),
+    "ple.value_proj": ("ple_value",),
+}
+_TOP_LEAVES: dict[str, tuple[str, ...]] = {
+    "model.language_model.embed_tokens": ("token_embd",),
+    "lm_head": ("output",),
+    "model.language_model.hyper_connection_mixer.input_mix_weight_down": ("output_hc_down",),
+    "model.language_model.hyper_connection_mixer.input_mix_weight_up": ("output_hc_up",),
+}
+_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return _TOP_LEAVES[name]
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ()))
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    head_dim = int(kv("attention.key_length"))
+    sections = [int(s) for s in kv("rope.dimension_sections")][:3]
+    ratios = [int(r) for r in kv("attention.compress_ratios")]
+    n_v = int(kv("ssm.time_step_rank"))
+    text = SimpleNamespace(
+        **_FAMILY_DEFAULTS,
+        hidden_size=int(kv("embedding_length")),
+        # block_count also counts the MTP layers a GGUF may carry after the decoder
+        num_hidden_layers=int(kv("block_count")) - int(md.get(f"{_ARCH}.nextn_predict_layers", 0)),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=head_dim,
+        rope_parameters={
+            "rope_type": "default",
+            "rope_theta": float(kv("rope.freq_base")),
+            "partial_rotary_factor": int(kv("rope.dimension_count")) / head_dim,
+            "mrope_section": sections,
+            "mrope_interleaved": True,
+        },
+        max_position_embeddings=int(kv("context_length")),
+        full_attention_interval=int(kv("full_attention_interval")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        num_experts=int(kv("expert_count")),
+        num_experts_per_tok=int(kv("expert_used_count")),
+        moe_intermediate_size=int(kv("expert_feed_forward_length")),
+        shared_expert_intermediate_size=int(kv("expert_shared_feed_forward_length")),
+        linear_num_key_heads=int(kv("ssm.group_count")),
+        linear_num_value_heads=n_v,
+        linear_key_head_dim=int(kv("ssm.state_size")),
+        linear_value_head_dim=int(kv("ssm.inner_size")) // n_v,
+        linear_conv_kernel_dim=int(kv("ssm.conv_kernel")),
+        hc_count=int(kv("hyper_connection.count")),
+        hc_lowrank=int(kv("hyper_connection.low_rank")),
+        indexer_n_heads=int(kv("attention.indexer.head_count")),
+        indexer_head_dim=int(kv("attention.indexer.key_length")),
+        indexer_budget=int(kv("attention.indexer.top_k")),
+        indexer_compress_ratio=max(ratios),
+        ple_layer_ids=[int(i) + 1 for i in kv("ple.layers")],
+        # the GGUF key is one n-gram head's row; HF counts all (ngram_size - 1) * heads_per_ngram heads
+        ple_embed_dim=int(kv("embedding_length_per_layer_input")) * (int(kv("ple.ngram_size")) - 1) * int(kv("ple.heads_per_ngram")),
+        ple_conv_kernel_size=int(kv("ple.conv_kernel")),
+        ngram_size=int(kv("ple.ngram_size")),
+        heads_per_ngram=int(kv("ple.heads_per_ngram")),
+        eos_token_id=int(kv("ple.eos_token_id")),
+    )
+    image_token_id = md.get(f"{_ARCH}.ple.image_token_id")
+    hf_like = SimpleNamespace(
+        text_config=text,
+        model_type="qwen4_exp",
+        architectures=list(shim.architectures),
+        image_token_id=None if image_token_id is None else int(image_token_id),
+        vision_config=None,
+    )
+    config = parse_config(hf_like)
+    return replace(config, expert_quant="gguf", qwen4_args=replace(config.qwen4_args, gdn_out_heads_tiled=True))
+
+
+# ----------------------------------------------------------------------------------------------
+# Weights
+# ----------------------------------------------------------------------------------------------
+
+def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> torch.Tensor:
+    """Tiled value heads (k-head minor) back to grouped (k-head major) along ``dim``."""
+    n = t.shape[dim]
+    r = n // (num_k_heads * head_dim)
+    idx = torch.arange(n).view(r, num_k_heads, head_dim).transpose(0, 1).reshape(n)
+    return t.index_select(dim, idx)
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+    include_vision: bool = True,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    config = parse_gguf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    md = cached_load_hf_config(model_path).metadata
+    g = config.linear_attention_group()
+    hidden = config.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    def norm(name: str) -> torch.Tensor:
+        return float_tensor(t(name)).float() - 1.0
+
+    def hc(prefix: str, src: str) -> Iterator[tuple[str, torch.Tensor]]:
+        down, inject, up = (t(f"{src}_{p}") for p in ("down", "inject", "up"))
+        yield from fused(f"{prefix}.input_mix_weight_down_block_inject", [down, inject], [matrix(down), matrix(inject)],
+                         hidden * config.qwen4_args.hc_count, pad_to=16)
+        yield from fused(f"{prefix}.input_mix_weight_up", [up], [matrix(up)], config.qwen4_args.hc_lowrank)
+        yield f"{prefix}.hc_norm.weight", norm(f"{src}_norm")
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    for proj, k in (("down", hidden * config.qwen4_args.hc_count), ("up", config.qwen4_args.hc_lowrank)):
+        info = t(f"output_hc_{proj}")
+        yield from fused(f"model.hyper_connection_mixer.input_mix_weight_{proj}", [info], [matrix(info)], k)
+    yield "model.hyper_connection_mixer.hc_norm.weight", norm("output_hc_norm")
+
+    for layer in range(config.num_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        if config.is_linear_layer(layer):
+            la = f"{pre}.linear_attn"
+            nk, dk = g.num_key_heads, g.key_head_dim
+            qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate")
+            qkv_rows = matrix(qkv)
+            v_start = 2 * nk * dk
+            qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)])
+            z_rows = _untile(matrix(gate), nk, g.value_head_dim)
+            beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha")
+            beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info))
+            if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type):
+                yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden)
+                yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden)
+            else:
+                yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden)
+            conv = float_tensor(t(f"{blk}.ssm_conv1d")).float()
+            conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)])
+            yield f"{la}.conv1d.weight", conv.unsqueeze(1)
+            yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1)
+            yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1)
+            yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm"))
+            yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim)
+        else:
+            sa = f"{pre}.self_attn"
+            parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+            yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden)
+            out = t(f"{blk}.attn_output")
+            yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+            yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm")
+            yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm")
+            iq, ik = t(f"{blk}.indexer.q_proj"), t(f"{blk}.indexer.k_proj")
+            yield from fused(f"{sa}.indexer.index_qk_proj", [iq, ik], [matrix(iq), matrix(ik)], hidden)
+            yield f"{sa}.indexer.q_layernorm.weight", norm(f"{blk}.indexer.q_norm")
+            yield f"{sa}.indexer.k_layernorm.weight", norm(f"{blk}.indexer.k_norm")
+
+        mlp = f"{pre}.mlp"
+        yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+        sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden)
+        yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size)
+        yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1)
+        yield from hc(f"{pre}.attn_hyper_connection", f"{blk}.hc_attn")
+        yield from hc(f"{pre}.mlp_hyper_connection", f"{blk}.hc_ffn")
+
+        if layer in config.qwen4_args.ple_layer_ids:
+            ple = f"{pre}.ple"
+            for proj in ("key", "value"):
+                info = t(f"{blk}.ple_{proj}")
+                yield from fused(f"{ple}.{proj}_proj", [info], [matrix(info)], config.qwen4_args.ple_embed_dim)
+            for n in ("key", "query", "conv"):
+                yield f"{ple}.norm_{n}.weight", norm(f"{blk}.ple_norm_{n}")
+            yield f"{ple}.conv1d.weight", float_tensor(t(f"{blk}.ple_conv1d")).float().unsqueeze(1)
+            yield f"{ple}.ple_embedding.layer_multipliers", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.layer_multipliers"]], dtype=torch.int64)
+            yield f"{ple}.ple_embedding.ngram_heads_vocab_sizes", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.head_vocab_sizes"]], dtype=torch.int64)
+            yield f"{ple}.ple_embedding.ngram_heads_offsets", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.head_offsets"]], dtype=torch.int64)
+
+
+def ple_table_info(model_path: str) -> GgufTensorInfo:
+    return gguf_tensor_index(model_path)["per_layer_token_embd.weight"]
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "ple_table_info", "to_gguf"]
diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py
index f13808dccd..69aaa58302 100644
--- a/python/freetoken/models/qwen4_exp/hc.py
+++ b/python/freetoken/models/qwen4_exp/hc.py
@@ -98,6 +98,7 @@ def __init__(self, config: ModelConfig, use_combine: bool = True, *, prefix: str
             self.pad_size = (-(self.lowrank + self.hc_count)) % 16
             self.input_mix_weight_down_block_inject = LinearReplicated(
                 width, self.lowrank + self.hc_count + self.pad_size, has_bias=False,
+                output_sizes=[self.lowrank, self.hc_count + self.pad_size],
                 quant_config=config.quant, prefix=f"{prefix}.input_mix_weight_down_block_inject",
             )
         else:
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index 258badac42..4eb69165d5 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -52,6 +52,7 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseO
         output_gate=g.output_gate,
         quant_config=config.quant,
         prefix=prefix,
+        out_heads_tiled=config.qwen4_args.gdn_out_heads_tiled,
     )
 
 
@@ -92,6 +93,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None:
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
@@ -168,11 +171,8 @@ def load_host_tables(self, engine_config) -> int:
             return 0
 
         if engine_config.ple_backend == "disk":
-            from freetoken.utils import download_hf_weight
-
             from .ple_disk import DiskRowTable, resolve_row_source
 
-            folder = download_hf_weight(engine_config.model_path)
             # one WAIT node per captured graph: the flag protocol supports a single consume
             assert len(ple_layers) == 1, "disk PLE backend expects exactly one PLE layer"
             emb, args = ple_layers[0].ple_embedding, ple_layers[0].args
@@ -186,7 +186,7 @@ def load_host_tables(self, engine_config) -> int:
                 "image_token_id": args.image_token_id,
             }
             disk_table = DiskRowTable(
-                resolve_row_source(folder),
+                resolve_row_source(engine_config.model_path),
                 constants,
                 max_graph_rows=max(256, engine_config.cuda_graph_max_bs or 0),
                 max_extend_tokens=engine_config.max_extend_tokens,
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index 69ee9fa93a..bb71da2ff7 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -50,6 +50,9 @@ class PleRowSource:
     row_bytes: int
     row_stride: int
     scale: float
+    # packed ggml rows (a GGUF table) of ``row_values`` values each; None for one fp8 byte per value
+    ggml_type: int | None = None
+    row_values: int | None = None
 
     @property
     def total_rows(self) -> int:
@@ -95,9 +98,27 @@ def source_from_safetensors(folder: str) -> PleRowSource:
     return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, cols, cols, float(scale))
 
 
-def resolve_row_source(folder: str) -> PleRowSource:
+def source_from_gguf(model_path: str) -> PleRowSource:
+    """The GGUF ``per_layer_token_embd`` table in place: one extent of packed ggml rows."""
+    from freetoken.layers.quantization.ggml_types import row_bytes
+
+    from .gguf import ple_table_info
+
+    info = ple_table_info(model_path)
+    values, rows = info.ne
+    nbytes = row_bytes(info.ggml_type, values)
+    return PleRowSource([info.path], [0], [info.offset], rows, nbytes, nbytes, 1.0, info.ggml_type, values)
+
+
+def resolve_row_source(model_path: str) -> PleRowSource:
     """Pick the row source for a checkpoint; the seam where a repacked format would plug in."""
-    return source_from_safetensors(folder)
+    from freetoken.models.gguf.reader import gguf_config_source
+
+    if (gguf_src := gguf_config_source(model_path)) is not None:
+        return source_from_gguf(gguf_src)
+    from freetoken.utils import download_hf_weight
+
+    return source_from_safetensors(download_hf_weight(model_path))
 
 
 class DiskRowTable:
@@ -113,9 +134,14 @@ def __init__(
         dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         from freetoken.kernel.row_store import PleStore
+        from freetoken.layers.quantization.ggml_types import float_dtype
 
         self.num_rows = source.total_rows
-        self.head_dim = source.row_bytes  # fp8: one byte per element
+        self.head_dim = source.row_values or source.row_bytes  # fp8: one byte per element
+        self._ggml_type = source.ggml_type
+        self._row_bytes = source.row_bytes
+        # float rows (fp8, or F32 / F16 / BF16 in a GGUF) are only cast; quantized ggml rows go through ggml
+        self._row_dtype = torch.float8_e4m3fn if source.ggml_type is None else float_dtype(source.ggml_type)
         self.dtype = dtype
         self.heads = int(hash_constants["num_ngram_heads"])
         self.scale = source.scale
@@ -142,7 +168,7 @@ def __init__(
             use_io_uring=os.getenv(_IO_URING_ENV, "1") != "0",
         )
         self._device = torch.device("cuda", torch.cuda.current_device())
-        self._token_bytes = self.heads * self.head_dim
+        self._token_bytes = self.heads * self._row_bytes
         # allocated up front: pinned alloc inside stream capture is illegal; one replay consumes it at a time
         self._graph_pinned = alloc_pinned_tensor(max_graph_rows * self._token_bytes, dtype=torch.uint8)
         self._graph_pinned.zero_()  # padded decode lanes read whatever sits here
@@ -272,7 +298,13 @@ def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torc
         dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True)
         if not capturing:
             self._eager_read[self._eager_slot].record(stream)
-        values = dev[:nbytes].view(torch.float8_e4m3fn).to(self.dtype)
+        if self._row_dtype is not None:
+            # copy=True: rows already in self.dtype must not alias the staging buffer the next fill reuses
+            values = dev[:nbytes].view(self._row_dtype).to(self.dtype, copy=True)
+        else:
+            from freetoken.kernel import ggml
+
+            values = ggml.dequantize(dev[:nbytes].view(-1, self._row_bytes), self._ggml_type, self.head_dim, self.dtype)
         if self.scale != 1.0:
             values = values * self.scale
         values = values.view(*row_ids.shape[:-1], -1)
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 678424ff05..4b0cea25a9 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -213,6 +213,15 @@ class ModelSpec:
         encoders=_QWEN_VL_ENCODERS,
         packed_modules_mapping=_QWEN4_EXP_PACKED,
     ),
+    # llama.cpp qwen4exp GGUF (text only): the same model classes with every tensor in its packed ggml type
+    "Qwen4ExpGGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen4_exp",
+        "Qwen4ExpForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        checkpoint_roots=_LANGUAGE_MODEL_ROOT,
+        packed_modules_mapping=_QWEN4_EXP_PACKED,
+    ),
     # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the
     # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the
     # loader handles the compressed-tensors NVFP4 layout.
@@ -282,12 +291,14 @@ class ModelSpec:
         checkpoint_segments=_GEMMA4_SEGMENTS,
         packed_modules_mapping=_GEMMA4_PACKED,
     ),
-    # GGUF (native Q4_0/Q6_K) gemma4: same model classes, GGUF config + weight loaders.
+    # llama.cpp gemma4 GGUF (text only): the same model classes with every tensor in its packed ggml type
     "Gemma4GGUFForCausalLM": ModelSpec(
         "freetoken.models.gemma4",
         "Gemma4ForCausalLM",
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
+        checkpoint_segments=_GEMMA4_SEGMENTS,
+        packed_modules_mapping=_GEMMA4_PACKED,
     ),
     "GptOssForCausalLM": ModelSpec(
         "freetoken.models.gpt_oss",
@@ -342,16 +353,19 @@ def _load_attr(module_path: str, attr_name: str) -> Any:
 
 
 def checkpoint_quant_config(model_path: str, hf_config: Any, spec: ModelSpec):
-    """The checkpoint's QuantConfig under the family's naming, or None for GGUF, whose native-quant ops the shared parser does not model yet."""
-    from freetoken.layers.quantization import NameMap, QuantConfig
+    """The checkpoint's QuantConfig under the family's naming.
 
+    A GGUF family gets the GGUF dialect over the file's tensor table, named through the ``to_gguf`` its module exports."""
+    from freetoken.layers.quantization import GGUFConfig, NameMap, QuantConfig
+
+    name_map = NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping)
     if spec.parse_config == "parse_gguf_config":
-        return None
-    return QuantConfig.from_hf(
-        hf_config,
-        name_map=NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping),
-        unquantized=spec.unquantized_modules,
-    )
+        to_gguf = _load_attr(spec.module, "to_gguf")
+        from freetoken.models.gguf.reader import gguf_tensor_index
+
+        types = {name: info.ggml_type for name, info in gguf_tensor_index(hf_config.model_path).items()}
+        return GGUFConfig(types, to_gguf, name_map=name_map, unquantized=spec.unquantized_modules)
+    return QuantConfig.from_hf(hf_config, name_map=name_map, unquantized=spec.unquantized_modules)
 
 
 def get_model_class(model_architecture: str, model_config: ModelConfig):
diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py
index e3382b1b59..b03266fb67 100644
--- a/python/freetoken/models/weight.py
+++ b/python/freetoken/models/weight.py
@@ -262,28 +262,9 @@ def ftw_lacks_vision(model_path: str) -> bool:
     return not any(name.startswith(VISION_KEY_PREFIXES) for name in ftw_tensor_names(model_path, "weight"))
 
 
-def load_q4_0_moe_expert_sources(
-    model_path: str,
-    model_config,
-    *,
-    dummy: bool = False,
-    layer_sink=None,
-) -> dict:
-    """Load (or fabricate, with ``dummy=True``) packed GGUF Q4_0 expert source banks.
-    ``layer_sink`` (converter) streams each completed layer's banks; ignored for dummy."""
-    _config, spec = _spec_for_model_path(model_path)
-    if dummy:
-        builder = _model_override(spec, "dummy_q4_0_expert_sources")
-        assert builder is not None, "model defines no dummy_q4_0_expert_sources"
-        return builder(model_config)
-    loader = _load_attr(spec.module, "load_q4_0_expert_sources")
-    return loader(model_path, model_config, layer_sink=layer_sink)
-
-
 __all__ = [
     "load_weight",
     "load_vision_weight",
     "ftw_lacks_vision",
-    "load_q4_0_moe_expert_sources",
     "iter_expert_tensors_parallel",
 ]
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 5521c34e1f..03ce4b750b 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -2,8 +2,8 @@
 
 The expert kernel (``QuantMethod.kernel``) owns the bank layout and the pack step; the
 checkpoint side delivers pieces (``moe.expert_pieces``) and this module fills the pinned host
-banks from them (``build_expert_banks``). The GGUF q4_0 experts still
-use their own providers until they get a method.
+banks from them (``build_expert_banks``). GGUF experts keep their packed ggml rows
+(``models.gguf.experts``).
 """
 
 from __future__ import annotations
@@ -19,7 +19,7 @@
 from freetoken.utils import init_logger
 
 from .host_banks import alloc_layer_banks
-from .offload_cache import _BANK_BYTES_PER_EXPERT, _BANK_SCHEMAS
+from .offload_cache import _BANK_BYTES_PER_EXPERT
 
 logger = init_logger(__name__)
 
@@ -50,6 +50,8 @@ class ExpertBanks:
     layout: dict | None = None
     # per-layer quant tag (e.g. the gate_up and down GGUF types), for layers whose banks match in dtype and row shape but not in quant; None = one quant
     layer_tags: list | None = None
+    # per bank, per layer: bytes past an expert matrix's last row that the kernel reads and that must read as zeros (ggml MMQ); None = no such reads
+    row_tail_bytes: dict[str, list[int]] | None = None
 
 
 def _dummy_fill(role: str, tensor: torch.Tensor) -> None:
@@ -161,46 +163,6 @@ def _fill(sink) -> None:
 _PARALLEL_CHUNK = 8 << 20  # default O_DIRECT chunk for the parallel reader
 
 
-def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks:
-    if parallel:
-        raise NotImplementedError(
-            "parallel reader not implemented for q4_0: GGUF is a single packed file "
-            "(not safetensors), so the common reader doesn't apply -- it needs a GGUF-native "
-            "parallel reader (parse the tensor table, chunked O_DIRECT over the one file)"
-        )
-    from freetoken.models.weight import load_q4_0_moe_expert_sources
-
-    # Native GGUF Q4_0 routed experts: packed block bytes streamed to the GPU and
-    # dequantized inside the borrowed ggml MoE kernels (no bf16 expert copy). Banks are
-    # per-layer HostBanks (pin-after-fill), so conversion streams each completed layer's
-    # gate_up + down straight through the sink (dummy fabricates in one shot -> not streamed).
-    sink = None if dummy else layer_sink
-    sources = load_q4_0_moe_expert_sources(model_path, model_config, dummy=dummy, layer_sink=sink)
-    return ExpertBanks(
-        "q4_0", {name: sources[name] for name in _BANK_SCHEMAS["q4_0"]}, streamed=sink is not None
-    )
-
-
-# expert formats that still load through their own provider (GGUF)
-_PROVIDERS = {
-    "q4_0": _q4_0_banks,
-}
-
-
-def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks:
-    expert_quant = model_config.expert_quant
-    if expert_quant not in _PROVIDERS:
-        raise ValueError(
-            f"{expert_quant!r} experts load through their MoE quant method; "
-            f"only {sorted(_PROVIDERS)} still have a format provider"
-        )
-    return _PROVIDERS[expert_quant](
-        model_path, model_config, device, dtype, dummy,
-        parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target,
-        layer_sink=layer_sink,
-    )
-
-
 def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None, resident=False) -> ExpertBanks:
     from freetoken.moe.expert_pieces import iter_expert_pieces
 
@@ -258,7 +220,7 @@ def bank_bytes_estimate(model_config, method=None) -> int | None:
     """Estimated total expert-bank bytes of a raw checkpoint before loading it.
 
     With a bound expert ``method`` the kernel's layout gives the exact host bytes; otherwise the
-    format-tag table sizes the GGUF format. ``None`` for unknown formats or missing dims
+    format-tag table sizes the format. ``None`` for unknown formats or missing dims
     (callers then skip the pre-load sizing)."""
     layers = getattr(model_config, "num_moe_layers", None)
     if method is not None and layers:
@@ -287,12 +249,10 @@ def load_expert_banks(
     *,
     method=None,
     device: torch.device,
-    dtype: torch.dtype,
     dummy: bool = False,
     parallel: bool | None = None,
     workers: int = 8,
     chunk: int = _PARALLEL_CHUNK,
-    decode_target: str = "gpu",
     layer_sink=None,
     layer_residency: list[str] | None = None,
     resident: bool = False,
@@ -315,8 +275,8 @@ def load_expert_banks(
     streamable formats, so callers must check ``streamed`` rather than assume it fired.
 
     ``method`` (the bound expert quant method of the model's offload layers) selects
-    the generic path: the family's pieces packed by the method's kernel. Without it only the
-    GGUF q4_0 format loads, through its own provider.
+    the generic path: the family's pieces packed by the method's kernel. Without it only an
+    FTW checkpoint loads.
 
     ``layer_residency``: per-layer ``HostResidency`` labels applied at settle time -- explicitly on the FTW fast path, ambiently (``requested_residency``) in the slow-path providers.
     Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota).
@@ -370,9 +330,14 @@ def load_expert_banks(
     from freetoken.moe.host_banks import requested_residency
 
     def _build(par: bool) -> ExpertBanks:
-        if method is not None:
-            return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident)
-        return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink)
+        if method is not None and method.kind is QuantKind.GGUF:
+            from freetoken.models.gguf.experts import load_gguf_expert_banks
+
+            # each layer keeps its own ggml types, so the banks come from the file, not from one method's layout
+            return load_gguf_expert_banks(model_path, model_config.num_moe_layers, dummy=dummy, device=device if resident else None)
+        if method is None:
+            raise ValueError("the experts of a raw checkpoint are packed by the offload layers' MoE quant method; pass method=")
+        return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident)
 
     with requested_residency(layer_residency) as residency_plan:
         try:
diff --git a/python/freetoken/moe/fused_q4_0.py b/python/freetoken/moe/fused_q4_0.py
deleted file mode 100644
index cdab82bf99..0000000000
--- a/python/freetoken/moe/fused_q4_0.py
+++ /dev/null
@@ -1,53 +0,0 @@
-"""Grouped expert GEMM over native GGUF Q4_0 banks (borrowed ggml MoE kernels).
-
-Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache
-interface: the experts are streamed to the GPU as packed Q4_0 block bytes and
-dequantized *inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. We
-use the MMVQ (vector) kernel for both prefill and decode: it consumes ``topk_ids``
-directly (no ``moe_align_block_size`` needed) and on small batches it is the right
-choice anyway. ``topk_ids`` already index the streamed cache slots (decode) or the
-materialized layer positions (prefill).
-"""
-
-from __future__ import annotations
-
-import torch
-
-from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
-from freetoken.models.gguf.dequant import GGML_Q4_0
-
-_ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul}
-
-
-def fused_experts_gguf_q4_0(
-    hidden_states: torch.Tensor,
-    gate_up_q: torch.Tensor,  # [num_slots, 2I, H//32*18] uint8
-    down_q: torch.Tensor,  # [num_slots, H, I//32*18] uint8
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    activation: str,
-) -> torch.Tensor:
-    from freetoken.kernel.gguf import ggml_moe_a8_vec
-
-    act_fn = _ACT.get(activation)
-    if act_fn is None:
-        raise ValueError(f"unsupported MoE activation {activation!r}")
-
-    num_tokens = hidden_states.shape[0]
-    n2 = gate_up_q.shape[1]  # 2 * intermediate
-    h = down_q.shape[1]  # hidden
-    top_k = topk_ids.shape[1]
-    qt = int(GGML_Q4_0)
-
-    # gate_up: [num_tokens*top_k, 2I] -> activation -> [num_tokens*top_k, I]
-    gate_up = ggml_moe_a8_vec(hidden_states, gate_up_q, topk_ids, top_k, qt, n2, num_tokens)
-    inter = act_fn(gate_up)
-    # down: each of the num_tokens*top_k intermediate rows uses its own expert id.
-    out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, qt, h, num_tokens * top_k)
-    out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to(
-        out.dtype
-    )
-    return out.sum(dim=1)
-
-
-__all__ = ["fused_experts_gguf_q4_0"]
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 701524e0a1..a209f9a924 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -46,9 +46,8 @@
     # Half the host/cache footprint of bf16; the grouped GEMM (kernel/triton/fp8_blockscale_moe)
     # reads the routed fp8 rows directly and dequantizes in the K-loop (no bf16 materialization).
     "fp8_block": ("gate_up", "gate_up_scale", "down", "down_scale"),
-    # native GGUF Q4_0 experts: packed block bytes per output row, dequantized inside
-    # the borrowed ggml MoE kernels. gate_up [L*E, 2I, H//32*18], down [L*E, H, I//32*18].
-    "q4_0": ("gate_up", "down"),
+    # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts keep their dtype
+    "gguf": ("gate_up", "down"),
     # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes +
     # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct
     # globals, and folding them into the e4m3 block scales would underflow)
@@ -93,7 +92,6 @@
         (2 * I // 128) * fp8_block_scale_pad(2 * I // 128, H // 128)
         + (H // 128) * fp8_block_scale_pad(H // 128, I // 128)
     ) * 2,
-    "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18,
     "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2),
     "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2),
     "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32),
@@ -182,13 +180,12 @@ def __post_init__(self) -> None:
         # all layers = the plain --moe-strategy cpu case).
         self.cpu_layer_ids: frozenset = frozenset()
         # Bank row layouts, set by set_bank_sources: layouts[g][k] is (dtype, row shape) of bank k (layouts of two quant tags may be equal), layer_layout[l] is layer l's g.
-        # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several it is a byte segment each pool views with its own layout.
+        # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several, or with pads, it is a byte segment each pool views with its own layout.
         self.layouts: list[tuple[tuple[torch.dtype, tuple[int, ...]], ...]] = []
         self.layer_layout: list[int] = [0] * self.num_layers
+        # _pad_bytes[g][k]: zeroed bytes after the last row of every pool and staging buffer of layout g in bank k, for kernels that read past it (16-aligned)
+        self._pad_bytes: list[list[int]] = []
         self._raw_banks = False
-        # Set by the engine when hybrid decode hands CPU routes to a kernel that reads slot 0 with weight 0 (no supports_inactive_slots).
-        # With several quants a staging buffer can leave another quant's bytes in a pool's slot 0, which may decode to NaN, so the last prefill layer zeroes those rows.
-        self.zero_inactive_route_slots = False
         # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path.
         self.validate_rebuild(self.cache_size, self.pool_sizes)
         assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, (
@@ -302,6 +299,8 @@ def __post_init__(self) -> None:
         # _stage_views[g][b] are buffer b's per-bank [num_experts, ...] views for layout g; _inval[b] are the global slots it overwrites.
         self._stage_bytes: list[int] = []
         self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = []
+        # _stage_tails[g][b]: the pad after buffer b's rows in layout g
+        self._stage_tails: list[list[list[torch.Tensor]]] = []
         self._inval: list[slice | torch.Tensor] = []
         self.prefill_copy_stream: torch.cuda.Stream | None = None
         self.prefill_begin_event: torch.cuda.Event | None = None
@@ -337,6 +336,12 @@ def _bank_rows(self, k: int, layout: int, offset: int, nrows: int) -> torch.Tens
         dtype, shape = self.layouts[layout][k]
         return cache[offset : offset + nrows * row].view(dtype).view(nrows, *shape)
 
+    def _bank_bytes(self, k: int, pool_sizes: list[int]) -> int:
+        """Bytes of raw bank ``k``: each pool's slots and pad, plus the largest pad, so two staging buffers fit wherever two layers of slots do."""
+        layouts = [self._pool_layout(p) for p in range(len(pool_sizes))]
+        slots = sum(s * self._row_bytes[g][k] + self._pad_bytes[g][k] for g, s in zip(layouts, pool_sizes))
+        return slots + max(pads[k] for pads in self._pad_bytes)
+
     def _stage_fits(self, buffers: int) -> bool:
         return all(c.numel() * c.element_size() >= buffers * s for (_, c), s in zip(self.banks, self._stage_bytes))
 
@@ -356,14 +361,15 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
         # rows[k][p]: bytes of one slot of pool p in bank k; before the banks exist, one unit per slot with a stage of E
         if self.banks:
             rows = [[self._row_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))]
+            pads = [[self._pad_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))]
             stage = self._stage_bytes
         else:
-            rows, stage = [[1] * len(spans)], [E]
+            rows, pads, stage = [[1] * len(spans)], [[0] * len(spans)], [E]
         offsets = []
-        for row in rows:
+        for row, pad in zip(rows, pads):
             offsets.append([0])
-            for s, r in zip(self.pool_sizes[:-1], row):
-                offsets[-1].append(offsets[-1][-1] + s * r)
+            for s, r, d in zip(self.pool_sizes[:-1], row, pad):
+                offsets[-1].append(offsets[-1][-1] + s * r + d)
         self._pool_offsets = offsets
         # local slots below this index overlap the prefill staging area in some bank
         self._pool_stable_from = [
@@ -371,10 +377,8 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
             for p, s in enumerate(self.pool_sizes)
         ]
         self._inval = []
-        masks = []
         for b in range(2):
             mask = torch.zeros(self.cache_size, dtype=torch.bool)
-            masks.append(mask)
             for k in range(len(rows)):
                 lo, hi = b * stage[k], (b + 1) * stage[k]
                 for p, (base, s) in enumerate(spans):
@@ -387,9 +391,8 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
             else:
                 self._inval.append(idx.to(self.device))
         self._stage_dirty = [False, False]
-        self._stage_slot0_rows = [[], []]
         if not self.banks:
-            self._pool_bank_views, self._stage_views = [], []
+            self._pool_bank_views, self._stage_views, self._stage_tails = [], [], []
             return
         self._pool_bank_views = [
             tuple(self._bank_rows(k, self._pool_layout(p), offsets[k][p], s) for k in range(len(self.banks)))
@@ -402,11 +405,10 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
             ]
             for g in range(len(self.layouts))
         ]
-        if len(self.layouts) > 1:
-            self._stage_slot0_rows = [
-                [self._bank_rows(k, self._pool_layout(p), offsets[k][p], 1) for p, base in enumerate(bases) if masks[b][base] for k in range(len(self.banks))]
-                for b in range(2)
-            ]
+        self._stage_tails = [
+            [[self.banks[k][1][b * stage[k] + E * row[k] :][: pad[k]] for k in range(len(self.banks)) if pad[k] and self._stage_fits(b + 1)] for b in range(2)]
+            for row, pad in zip(self._row_bytes, self._pad_bytes)
+        ]
 
     def stage_views(self, layer_id: int, buffer_id: int) -> tuple[torch.Tensor, ...]:
         """Per-bank ``[num_experts, ...]`` views of prefill staging buffer ``buffer_id`` in ``layer_id``'s layout; row == expert id."""
@@ -432,6 +434,7 @@ def set_bank_sources(
         sources: dict[str, list[torch.Tensor]],
         layer_residency: list[str] | None = None,
         layer_tags: list | None = None,
+        row_tail_bytes: dict[str, list[int]] | None = None,
     ) -> None:
         """Attach the host (CPU pinned) expert source banks and allocate a GPU slot
         cache per bank, following the format's bank schema.
@@ -447,6 +450,8 @@ def set_bank_sources(
         Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them.
 
         ``layer_tags`` names each layer's quant (see :func:`bank_layouts`); layers of different tags never share a pool.
+
+        ``row_tail_bytes`` gives, per bank and layer, how many bytes past the last row the kernel reads; every pool and staging buffer is then followed by that many zero bytes.
         """
         from freetoken.moe.legacy_format import canonical_role
         from freetoken.moe.host_banks import HostResidency
@@ -498,12 +503,22 @@ def set_bank_sources(
         for p in range(len(self.pool_sizes)):
             mixed = {self.layer_layout[l] for l, q in enumerate(self.pool_of_layer) if q == p}
             assert len(mixed) == 1, f"slot pool {p} holds layers of {len(mixed)} layouts; a pool needs one"
-        self._raw_banks = len(set(self.layouts)) > 1
+        self._pad_bytes = [[0] * len(self.bank_schema) for _ in self.layouts]
+        if row_tail_bytes is not None:
+            tails = {canonical_role(name): per_layer for name, per_layer in row_tail_bytes.items()}
+            for k, name in enumerate(self.bank_schema):
+                for layer_id, tail in enumerate(tails[canonical_role(name)]):
+                    g = self.layer_layout[layer_id]
+                    self._pad_bytes[g][k] = max(self._pad_bytes[g][k], -(-tail // 16) * 16)
+        self._raw_banks = len(set(self.layouts)) > 1 or any(map(any, self._pad_bytes))
         self._row_bytes = [[math.prod(shape) * dtype.itemsize for dtype, shape in layout] for layout in self.layouts]
         if self._raw_banks:
             # every pool view and staging buffer starts on a row boundary of its own layout
             assert all(r % 16 == 0 for rows in self._row_bytes for r in rows), self._row_bytes
-        self._stage_bytes = [self.num_experts * max(rows[k] for rows in self._row_bytes) for k in range(len(self.bank_schema))]
+        self._stage_bytes = [
+            self.num_experts * max(rows[k] for rows in self._row_bytes) + max(pads[k] for pads in self._pad_bytes)
+            for k in range(len(self.bank_schema))
+        ]
         self._alloc_banks()
         if not self._stage_fits(1):
             raise ValueError(f"MoE cache of {self.cache_size} slots cannot stage one whole layer of experts")
@@ -518,8 +533,7 @@ def set_bank_sources(
     def _alloc_banks(self) -> None:
         for k, name in enumerate(self.bank_schema):
             if self._raw_banks:
-                nbytes = sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(self.pool_sizes))
-                self.bank_caches[name] = torch.zeros(nbytes, dtype=torch.uint8, device=self.device)
+                self.bank_caches[name] = torch.zeros(self._bank_bytes(k, self.pool_sizes), dtype=torch.uint8, device=self.device)
             else:
                 dtype, shape = self.layouts[0][k]
                 self.bank_caches[name] = torch.zeros((self.cache_size, *shape), dtype=dtype, device=self.device)
@@ -634,7 +648,7 @@ def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None)
             )
         if self._raw_banks and pool_sizes is not None:
             for k, stage in enumerate(self._stage_bytes):
-                if sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(pool_sizes)) < stage:
+                if self._bank_bytes(k, pool_sizes) < stage:
                     raise ValueError(f"moe_cache_size={cache_size} cannot stage one whole layer of experts")
 
     def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
@@ -663,7 +677,7 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         self.bank_caches = {}
         self._pool_bank_views = []
         self._stage_views = []
-        self._stage_slot0_rows = [[], []]
+        self._stage_tails = []
         self.cache_size = cache_size
         self.pool_sizes = sizes
         if self.device.type == "cuda":
@@ -796,7 +810,8 @@ def _init_prefill_overlap_buffers(self) -> None:
             )
             self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device)
 
-    def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
+    def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int) -> None:
+        """Take buffer ``buffer_id``'s slots for staging ``layer_id``."""
         slots = self._inval[buffer_id]
         old_ids = self.id_of_slot[slots]
         self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
@@ -805,6 +820,9 @@ def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
         # ensure_experts evicts them first.
         self.usage[slots] = 0
         self._stage_dirty[buffer_id] = True
+        # the pad may hold rows of a wider layout or of a pool
+        for tail in self._stage_tails[self.layer_layout[layer_id]][buffer_id]:
+            tail.zero_()
 
     def begin_prefill(self) -> None:
         if not self.prefill_overlap:
@@ -846,7 +864,7 @@ def prefetch_prefill_layer(self, layer_id: int) -> None:
             )
 
         def copy() -> None:
-            self._invalidate_prefill_buffer(buffer_id)
+            self._invalidate_prefill_buffer(buffer_id, layer_id)
             for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)):
                 view.copy_(per_layer[layer_id], non_blocking=True)
 
@@ -954,7 +972,7 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
         with torch.cuda.stream(self.prefill_copy_stream):
             if self._prefill_buffer_has_release_event[buffer_id]:
                 self.prefill_copy_stream.wait_event(self.prefill_release_events[buffer_id])
-            self._invalidate_prefill_buffer(buffer_id)
+            self._invalidate_prefill_buffer(buffer_id, layer_id)
             if miss.size:
                 run_starts = np.concatenate(([0], np.nonzero(np.diff(miss) != 1)[0] + 1))
                 starts = miss[run_starts]
@@ -997,12 +1015,13 @@ def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]:
         return self.stage_views(layer_id, buffer_id)
 
     def release_prefill_layer(self, layer_id: int) -> None:
-        if self.zero_inactive_route_slots and layer_id == self.num_layers - 1:
+        if layer_id == self.num_layers - 1:
             # the staged slots were invalidated and nothing registers them within the forward, so their bytes are free to clear
             for b in range(2):
                 if self._stage_dirty[b]:
-                    for row in self._stage_slot0_rows[b]:
-                        row.zero_()
+                    # staged rows left there may be shifted or of another quant, and kernels still read them: MMQ past a slot's last row, weight-0 routes on slot 0
+                    for (_, cache), stage in zip(self.banks, self._stage_bytes):
+                        cache.view(-1).view(torch.uint8)[b * stage : (b + 1) * stage].zero_()
                     self._stage_dirty[b] = False
         if not self.prefill_overlap:
             return
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index 1d863a2f54..8c4ad8f393 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -108,7 +108,7 @@ def materialize_layer(cache, layer_id: int) -> None:
         return
     # staging buffer 0 may overlap several pools, so it only holds the layer for this forward and registers nothing
     E = cache.num_experts
-    cache._invalidate_prefill_buffer(0)
+    cache._invalidate_prefill_buffer(0, layer_id)
     rows = torch.arange(E, dtype=torch.int32, device=cache.device)
     cache.evict_slots[:E].copy_(rows)
     cache.src_indices[:E].copy_(rows)
diff --git a/python/freetoken/third_party/__init__.py b/python/freetoken/third_party/__init__.py
new file mode 100644
index 0000000000..633cc8a798
--- /dev/null
+++ b/python/freetoken/third_party/__init__.py
@@ -0,0 +1 @@
+"""Vendored third-party Python code; each package carries its LICENSE and VERSION."""
diff --git a/python/freetoken/third_party/gguf/LICENSE b/python/freetoken/third_party/gguf/LICENSE
new file mode 100644
index 0000000000..76f67efdc6
--- /dev/null
+++ b/python/freetoken/third_party/gguf/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Georgi Gerganov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/python/freetoken/third_party/gguf/VERSION b/python/freetoken/third_party/gguf/VERSION
new file mode 100644
index 0000000000..f3ec79b970
--- /dev/null
+++ b/python/freetoken/third_party/gguf/VERSION
@@ -0,0 +1,10 @@
+llama.cpp 8c2134a8755d8a9792af513dab20282affe903e4 (2026-10-14), the same pin as kernel/csrc/ggml
+
+Only the gguf-py modules that the GGUF reader and the ggml type tables need: gguf-py/gguf/{constants,gguf_reader,quants,lazy}.py.
+The PyPI gguf release lags the ggml types in this tree (gguf 0.19.0 cannot read Q2_0 tensors), so FreeToken reads GGUF files with this copy.
+
+Local changes, kept in freetoken.patch (already applied):
+- __init__.py: imports only the vendored modules (upstream also pulls in the writer, vocab and metadata modules and their dependencies).
+- gguf_reader.py: imports constants relatively, so the installed gguf package is never picked up.
+
+To update: copy the same files from the llama.cpp commit that kernel/csrc/ggml pins, apply freetoken.patch, and update this file.
diff --git a/python/freetoken/third_party/gguf/__init__.py b/python/freetoken/third_party/gguf/__init__.py
new file mode 100644
index 0000000000..f4c74ef9e5
--- /dev/null
+++ b/python/freetoken/third_party/gguf/__init__.py
@@ -0,0 +1,4 @@
+from .constants import *
+from .lazy import *
+from .gguf_reader import *
+from .quants import *
diff --git a/python/freetoken/third_party/gguf/constants.py b/python/freetoken/third_party/gguf/constants.py
new file mode 100644
index 0000000000..c31ab4a42e
--- /dev/null
+++ b/python/freetoken/third_party/gguf/constants.py
@@ -0,0 +1,6332 @@
+from __future__ import annotations
+
+from enum import Enum, IntEnum, auto
+from typing import Any
+
+#
+# constants
+#
+
+GGUF_MAGIC             = 0x46554747  # "GGUF"
+GGUF_VERSION           = 3
+GGUF_DEFAULT_ALIGNMENT = 32
+GGML_QUANT_VERSION     = 2  # GGML_QNT_VERSION from ggml.h
+GGML_MAX_DIMS          = 4  # GGML_MAX_DIMS from ggml.h
+
+#
+# metadata keys
+#
+
+
+class Keys:
+    class General:
+        TYPE                       = "general.type"
+        ARCHITECTURE               = "general.architecture"
+        QUANTIZATION_VERSION       = "general.quantization_version"
+        ALIGNMENT                  = "general.alignment"
+        FILE_TYPE                  = "general.file_type"
+
+        # Per-tensor extra options (tensor name array + parallel option arrays, e.g. prec_a4).
+        TENSOR_EXTRA_NAME          = "general.tensor_extra.name"
+        TENSOR_EXTRA_PREC_A4       = "general.tensor_extra.prec_a4"
+
+        # Recommended Sampler Parameters
+        SAMPLING_SEQUENCE           = "general.sampling.sequence"
+        SAMPLING_TOP_K              = "general.sampling.top_k"
+        SAMPLING_TOP_P              = "general.sampling.top_p"
+        SAMPLING_MIN_P              = "general.sampling.min_p"
+        SAMPLING_XTC_PROBABILITY    = "general.sampling.xtc_probability"
+        SAMPLING_XTC_THRESHOLD      = "general.sampling.xtc_threshold"
+        SAMPLING_TEMP               = "general.sampling.temp"
+        SAMPLING_PENALTY_LAST_N     = "general.sampling.penalty_last_n"
+        SAMPLING_PENALTY_REPEAT     = "general.sampling.penalty_repeat"
+        SAMPLING_MIROSTAT           = "general.sampling.mirostat"
+        SAMPLING_MIROSTAT_TAU       = "general.sampling.mirostat_tau"
+        SAMPLING_MIROSTAT_ETA       = "general.sampling.mirostat_eta"
+
+        # Authorship Metadata
+        NAME                       = "general.name"
+        AUTHOR                     = "general.author"
+        VERSION                    = "general.version"
+        ORGANIZATION               = "general.organization"
+
+        FINETUNE                   = "general.finetune"
+        BASENAME                   = "general.basename"
+
+        DESCRIPTION                = "general.description"
+        QUANTIZED_BY               = "general.quantized_by"
+
+        SIZE_LABEL                 = "general.size_label"
+
+        # Licensing details
+        LICENSE                    = "general.license"
+        LICENSE_NAME               = "general.license.name"
+        LICENSE_LINK               = "general.license.link"
+
+        # Typically represents the converted GGUF repo (Unless native)
+        URL                        = "general.url" # Model Website/Paper
+        DOI                        = "general.doi"
+        UUID                       = "general.uuid"
+        REPO_URL                   = "general.repo_url" # Model Source Repository (git/svn/etc...)
+
+        # Model Source during conversion
+        SOURCE_URL                 = "general.source.url" # Model Website/Paper
+        SOURCE_DOI                 = "general.source.doi"
+        SOURCE_UUID                = "general.source.uuid"
+        SOURCE_REPO_URL            = "general.source.repo_url" # Model Source Repository (git/svn/etc...)
+
+        # Base Model Source. There can be more than one source if it's a merged
+        # model like with 'Mistral-7B-Merge-14-v0.1'. This will assist in
+        # tracing linage of models as it is finetuned or merged over time.
+        BASE_MODEL_COUNT           = "general.base_model.count"
+        BASE_MODEL_NAME            = "general.base_model.{id}.name"
+        BASE_MODEL_AUTHOR          = "general.base_model.{id}.author"
+        BASE_MODEL_VERSION         = "general.base_model.{id}.version"
+        BASE_MODEL_ORGANIZATION    = "general.base_model.{id}.organization"
+        BASE_MODEL_DESCRIPTION     = "general.base_model.{id}.description"
+        BASE_MODEL_URL             = "general.base_model.{id}.url" # Model Website/Paper
+        BASE_MODEL_DOI             = "general.base_model.{id}.doi"
+        BASE_MODEL_UUID            = "general.base_model.{id}.uuid"
+        BASE_MODEL_REPO_URL        = "general.base_model.{id}.repo_url" # Model Source Repository (git/svn/etc...)
+
+        # Dataset Source
+        DATASET_COUNT           = "general.dataset.count"
+        DATASET_NAME            = "general.dataset.{id}.name"
+        DATASET_AUTHOR          = "general.dataset.{id}.author"
+        DATASET_VERSION         = "general.dataset.{id}.version"
+        DATASET_ORGANIZATION    = "general.dataset.{id}.organization"
+        DATASET_DESCRIPTION     = "general.dataset.{id}.description"
+        DATASET_URL             = "general.dataset.{id}.url" # Model Website/Paper
+        DATASET_DOI             = "general.dataset.{id}.doi"
+        DATASET_UUID            = "general.dataset.{id}.uuid"
+        DATASET_REPO_URL        = "general.dataset.{id}.repo_url" # Model Source Repository (git/svn/etc...)
+
+        # Array based KV stores
+        TAGS                       = "general.tags"
+        LANGUAGES                  = "general.languages"
+
+    class LLM:
+        VOCAB_SIZE                        = "{arch}.vocab_size"
+        CONTEXT_LENGTH                    = "{arch}.context_length"
+        EMBEDDING_LENGTH                  = "{arch}.embedding_length"
+        EMBEDDING_LENGTH_OUT              = "{arch}.embedding_length_out"
+        FEATURES_LENGTH                   = "{arch}.features_length"
+        BLOCK_COUNT                       = "{arch}.block_count"
+        LEADING_DENSE_BLOCK_COUNT         = "{arch}.leading_dense_block_count"
+        FEED_FORWARD_LENGTH               = "{arch}.feed_forward_length"
+        EXPERT_FEED_FORWARD_LENGTH        = "{arch}.expert_feed_forward_length"
+        EXPERT_SHARED_FEED_FORWARD_LENGTH = "{arch}.expert_shared_feed_forward_length"
+        EXPERT_CHUNK_FEED_FORWARD_LENGTH  = "{arch}.expert_chunk_feed_forward_length"
+        USE_PARALLEL_RESIDUAL             = "{arch}.use_parallel_residual"
+        TENSOR_DATA_LAYOUT                = "{arch}.tensor_data_layout"
+        EXPERT_COUNT                      = "{arch}.expert_count"
+        EXPERT_USED_COUNT                 = "{arch}.expert_used_count"
+        EXPERT_SHARED_COUNT               = "{arch}.expert_shared_count"
+        EXPERT_GROUP_COUNT                = "{arch}.expert_group_count"
+        EXPERT_GROUP_USED_COUNT           = "{arch}.expert_group_used_count"
+        EXPERT_WEIGHTS_SCALE              = "{arch}.expert_weights_scale"
+        EXPERT_WEIGHTS_NORM               = "{arch}.expert_weights_norm"
+        EXPERT_GATING_FUNC                = "{arch}.expert_gating_func"
+        EXPERT_GROUP_SCALE                = "{arch}.expert_group_scale"
+        EXPERT_LATENT_LENGTH              = "{arch}.expert_latent_length"
+        EXPERTS_PER_GROUP                 = "{arch}.experts_per_group"
+        MOE_EVERY_N_LAYERS                = "{arch}.moe_every_n_layers"
+        MOE_LATENT_SIZE                   = "{arch}.moe_latent_size"
+        NEXTN_PREDICT_LAYERS              = "{arch}.nextn_predict_layers"
+        NUM_DEEPSTACK_LAYERS              = "{arch}.n_deepstack_layers"
+        DEEPSTACK_MAPPING                 = "{arch}.deepstack_mapping"
+        POOLING_TYPE                      = "{arch}.pooling_type"
+        LOGIT_SCALE                       = "{arch}.logit_scale"
+        DECODER_START_TOKEN_ID            = "{arch}.decoder_start_token_id"
+        DECODER_BLOCK_COUNT               = "{arch}.decoder_block_count"
+        ATTN_LOGIT_SOFTCAPPING            = "{arch}.attn_logit_softcapping"
+        ROUTER_LOGIT_SOFTCAPPING          = "{arch}.router_logit_softcapping"
+        FINAL_LOGIT_SOFTCAPPING           = "{arch}.final_logit_softcapping"
+        SWIN_NORM                         = "{arch}.swin_norm"
+        RESCALE_EVERY_N_LAYERS            = "{arch}.rescale_every_n_layers"
+        TIME_MIX_EXTRA_DIM                = "{arch}.time_mix_extra_dim"
+        TIME_DECAY_EXTRA_DIM              = "{arch}.time_decay_extra_dim"
+        RESIDUAL_SCALE                    = "{arch}.residual_scale"
+        EMBEDDING_SCALE                   = "{arch}.embedding_scale"
+        TOKEN_SHIFT_COUNT                 = "{arch}.token_shift_count"
+        INTERLEAVE_MOE_LAYER_STEP         = "{arch}.interleave_moe_layer_step"
+        FULL_ATTENTION_INTERVAL           = "{arch}.full_attention_interval"
+        NUM_LOOPS                         = "{arch}.num_loops"
+        SKIP_LOOP_FINAL_NORM              = "{arch}.skip_loop_final_norm"
+        HASH_LAYER_COUNT                  = "{arch}.hash_layer_count"
+        ACTIVATION_SPARSITY_SCALE         = "{arch}.activation_sparsity_scale"
+        ALTUP_ACTIVE_IDX                  = "{arch}.altup.active_idx"
+        ALTUP_NUM_INPUTS                  = "{arch}.altup.num_inputs"
+        EMBD_LENGTH_PER_LAYER_INP         = "{arch}.embedding_length_per_layer_input"
+        SWIGLU_CLAMP_EXP                  = "{arch}.swiglu_clamp_exp"
+        SWIGLU_CLAMP_SHEXP                = "{arch}.swiglu_clamp_shexp"
+        HIDDEN_ACT                        = "{arch}.hidden_activation"
+        DENSE_FEAT_IN_SIZE                = "{arch}.{dense}_feat_in"
+        DENSE_FEAT_OUT_SIZE               = "{arch}.{dense}_feat_out"
+        TARGET_LAYERS                     = "{arch}.target_layers"
+        TARGET_HIDDEN_SIZE                = "{arch}.target_hidden_size"
+        BLOCK_SIZE                        = "{arch}.block_size"
+        CONV_KERNEL_SIZE                  = "{arch}.conv_kernel_size"
+        CONV_GROUP_SIZE                   = "{arch}.conv_group_size"
+        SELECTOR_RANK                     = "{arch}.selector_rank"
+        SELECTOR_TOP_K                    = "{arch}.selector_top_k"
+        SAMPLE_FROM_ANCHOR                = "{arch}.sample_from_anchor"
+        HAS_CONFIDENCE_HEAD               = "{arch}.has_confidence_head"
+        NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
+        NORM_BEFORE_FC                    = "{arch}.norm_before_fc"
+
+    class Adapters:
+        COUNT                = "{arch}.adapters.count"
+        TOKEN_IDS_ACTIVATE   = "{arch}.adapters.token_ids_activate"
+        TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute"
+        LORA_RANK            = "{arch}.adapters.lora_rank"
+        ROUTER_GAIN          = "{arch}.adapters.router_gain"
+
+    class Attention:
+        HEAD_COUNT                   = "{arch}.attention.head_count"
+        HEAD_COUNT_KV                = "{arch}.attention.head_count_kv"
+        MAX_ALIBI_BIAS               = "{arch}.attention.max_alibi_bias"
+        CLAMP_KQV                    = "{arch}.attention.clamp_kqv"
+        KEY_LENGTH                   = "{arch}.attention.key_length"
+        VALUE_LENGTH                 = "{arch}.attention.value_length"
+        LAYERNORM_EPS                = "{arch}.attention.layer_norm_epsilon"
+        LAYERNORM_RMS_EPS            = "{arch}.attention.layer_norm_rms_epsilon"
+        GROUPNORM_EPS                = "{arch}.attention.group_norm_epsilon"
+        GROUPNORM_GROUPS             = "{arch}.attention.group_norm_groups"
+        CAUSAL                       = "{arch}.attention.causal"
+        Q_LORA_RANK                  = "{arch}.attention.q_lora_rank"
+        KV_LORA_RANK                 = "{arch}.attention.kv_lora_rank"
+        DECAY_LORA_RANK              = "{arch}.attention.decay_lora_rank"
+        ICLR_LORA_RANK               = "{arch}.attention.iclr_lora_rank"
+        VALUE_RESIDUAL_MIX_LORA_RANK = "{arch}.attention.value_residual_mix_lora_rank"
+        GATE_LORA_RANK               = "{arch}.attention.gate_lora_rank"
+        REL_BUCKETS_COUNT            = "{arch}.attention.relative_buckets_count"
+        SLIDING_WINDOW               = "{arch}.attention.sliding_window"
+        SCALE                        = "{arch}.attention.scale"
+        OUTPUT_GROUP_COUNT           = "{arch}.attention.output_group_count"
+        OUTPUT_LORA_RANK             = "{arch}.attention.output_lora_rank"
+        OUTPUT_SCALE                 = "{arch}.attention.output_scale"
+        VALUE_SCALE                  = "{arch}.attention.value_scale"
+        COMPRESS_RATIOS              = "{arch}.attention.compress_ratios"
+        COMPRESS_ROPE_FREQ_BASE      = "{arch}.attention.compress_rope_freq_base"
+        TEMPERATURE_LENGTH           = "{arch}.attention.temperature_length"
+        KEY_LENGTH_MLA               = "{arch}.attention.key_length_mla"
+        VALUE_LENGTH_MLA             = "{arch}.attention.value_length_mla"
+        KEY_LENGTH_SWA               = "{arch}.attention.key_length_swa"
+        VALUE_LENGTH_SWA             = "{arch}.attention.value_length_swa"
+        KEY_LENGTH_MLA_SWA           = "{arch}.attention.key_length_mla_swa"
+        VALUE_LENGTH_MLA_SWA         = "{arch}.attention.value_length_mla_swa"
+        KV_LORA_RANK_SWA             = "{arch}.attention.kv_lora_rank_swa"
+        SHARED_KV_LAYERS             = "{arch}.attention.shared_kv_layers"
+        SLIDING_WINDOW_PATTERN       = "{arch}.attention.sliding_window_pattern"
+        RECURRENT_LAYERS             = "{arch}.attention.recurrent_layers"
+        TEMPERATURE_SCALE            = "{arch}.attention.temperature_scale"
+        ROPE_PATTERN                 = "{arch}.attention.rope_pattern"
+        VALUE_EXPERT_COUNT           = "{arch}.attention.value_expert_count"
+        VALUE_EXPERT_USED_COUNT      = "{arch}.attention.value_expert_used_count"
+
+        class Indexer:
+            HEAD_COUNT = "{arch}.attention.indexer.head_count"
+            KEY_LENGTH = "{arch}.attention.indexer.key_length"
+            TOP_K      = "{arch}.attention.indexer.top_k"
+            BLOCK_SIZE   = "{arch}.attention.indexer.block_size"    # MSA
+            LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks"  # MSA
+            TYPES      = "{arch}.attention.indexer.types"
+            KPOOL             = "{arch}.attention.indexer.kpool"              # GLM5-Next
+            KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail"  # GLM5-Next
+
+    class HyperConnection:
+        COUNT                = "{arch}.hyper_connection.count"
+        SINKHORN_ITERATIONS  = "{arch}.hyper_connection.sinkhorn_iterations"
+        EPSILON              = "{arch}.hyper_connection.epsilon"
+        # scale of the post gate (DeepSeek-V4 hardcodes 2.0)
+        MAGNITUDE            = "{arch}.hyper_connection.magnitude"
+        # absent means the mix projection is full rank (DeepSeek-V4 behaviour)
+        LOW_RANK             = "{arch}.hyper_connection.low_rank"
+
+    class PerLayerEmbedding:
+        LAYERS             = "{arch}.ple.layers"
+        NGRAM_SIZE         = "{arch}.ple.ngram_size"
+        HEADS_PER_NGRAM    = "{arch}.ple.heads_per_ngram"
+        CONV_KERNEL        = "{arch}.ple.conv_kernel"
+        LAYER_MULTIPLIERS  = "{arch}.ple.layer_multipliers"
+        HEAD_OFFSETS       = "{arch}.ple.head_offsets"
+        HEAD_VOCAB_SIZES   = "{arch}.ple.head_vocab_sizes"
+        EOS_TOKEN_ID       = "{arch}.ple.eos_token_id"
+        IMAGE_TOKEN_ID     = "{arch}.ple.image_token_id"
+
+    class Rope:
+        DIMENSION_COUNT           = "{arch}.rope.dimension_count"
+        DIMENSION_COUNT_SWA       = "{arch}.rope.dimension_count_swa"
+        DIMENSION_SECTIONS        = "{arch}.rope.dimension_sections"
+        SECTION_ORDER             = "{arch}.rope.section_order"
+        FREQ_BASE                 = "{arch}.rope.freq_base"
+        FREQ_BASE_SWA             = "{arch}.rope.freq_base_swa"
+        SCALING_TYPE              = "{arch}.rope.scaling.type"
+        SCALING_FACTOR            = "{arch}.rope.scaling.factor"
+        SCALING_ALPHA             = "{arch}.rope.scaling.alpha"
+        SCALING_ATTN_FACTOR       = "{arch}.rope.scaling.attn_factor"
+        SCALING_ORIG_CTX_LEN      = "{arch}.rope.scaling.original_context_length"
+        SCALING_FINETUNED         = "{arch}.rope.scaling.finetuned"
+        SCALING_YARN_TRUNCATE     = "{arch}.rope.scaling.yarn_truncate"
+        SCALING_YARN_LOG_MUL      = "{arch}.rope.scaling.yarn_log_multiplier"
+        SCALING_YARN_EXT_FACTOR   = "{arch}.rope.scaling.yarn_ext_factor"
+        SCALING_YARN_ATTN_FACTOR  = "{arch}.rope.scaling.yarn_attn_factor"
+        SCALING_YARN_BETA_FAST    = "{arch}.rope.scaling.yarn_beta_fast"
+        SCALING_YARN_BETA_SLOW    = "{arch}.rope.scaling.yarn_beta_slow"
+
+    class Activation:
+        SITU_BETA        = "{arch}.activation.situ_beta"
+        SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta"
+
+    class AttnRes:
+        BLOCK_SIZE = "{arch}.attn_res.block_size"
+
+    class Split:
+        LLM_KV_SPLIT_NO            = "split.no"
+        LLM_KV_SPLIT_COUNT         = "split.count"
+        LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"
+
+    class HRM:
+        LAYERS_PER_STACK = "{arch}.hrm.layers_per_stack"
+        H_CYCLES         = "{arch}.hrm.h_cycles"
+        L_CYCLES         = "{arch}.hrm.l_cycles"
+        PREFIX_LM        = "{arch}.hrm.prefix_lm"
+
+    class SSM:
+        CONV_KERNEL    = "{arch}.ssm.conv_kernel"
+        INNER_SIZE     = "{arch}.ssm.inner_size"
+        STATE_SIZE     = "{arch}.ssm.state_size"
+        TIME_STEP_RANK = "{arch}.ssm.time_step_rank"
+        GROUP_COUNT    = "{arch}.ssm.group_count"
+        DT_B_C_RMS     = "{arch}.ssm.dt_b_c_rms"
+
+    class KDA:
+        HEAD_DIM         = "{arch}.kda.head_dim"
+        SAFE_GATE        = "{arch}.kda.safe_gate"
+        GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound"
+
+    class WKV:
+        HEAD_SIZE = "{arch}.wkv.head_size"
+
+    class PosNet:
+        EMBEDDING_LENGTH = "{arch}.posnet.embedding_length"
+        BLOCK_COUNT      = "{arch}.posnet.block_count"
+
+    class ConvNext:
+        EMBEDDING_LENGTH = "{arch}.convnext.embedding_length"
+        BLOCK_COUNT      = "{arch}.convnext.block_count"
+
+    class Classifier:
+        OUTPUT_LABELS = "{arch}.classifier.output_labels"
+        POOLING_TYPE  = "{arch}.classifier.pooling_type"
+        ACTIVATION    = "{arch}.classifier.activation"
+
+    class ShortConv:
+        L_CACHE = "{arch}.shortconv.l_cache"
+
+    class Decision:
+        TYPE               = "{arch}.decision.type"
+        # note: single-use-case keys can be hard-coded in cpp code
+        BLOCK_COUNT        = "{arch}.decision.block_count"
+        ROUTING_BLOCK_COUNT = "{arch}.decision.routing_block_count"
+        HEAD_COUNT         = "{arch}.decision.head_count"
+        MAX_HEAD_TOKENS    = "{arch}.decision.max_head_tokens"
+        TEMPERATURE        = "{arch}.decision.temperature.{name}"  # name: "<type>" or "<type>.<n_opt bucket>"
+        SCORE_BIAS         = "{arch}.decision.score_bias.{n}.{i}"  # added to the logit of level i of a score question with n levels
+
+    class Tokenizer:
+        MODEL                = "tokenizer.ggml.model"
+        PRE                  = "tokenizer.ggml.pre"
+        LIST                 = "tokenizer.ggml.tokens"
+        TOKEN_TYPE           = "tokenizer.ggml.token_type"
+        TOKEN_TYPE_COUNT     = "tokenizer.ggml.token_type_count"  # for BERT-style token types
+        SCORES               = "tokenizer.ggml.scores"
+        MERGES               = "tokenizer.ggml.merges"
+        BOS_ID               = "tokenizer.ggml.bos_token_id"
+        EOS_ID               = "tokenizer.ggml.eos_token_id"
+        EOT_ID               = "tokenizer.ggml.eot_token_id"
+        EOM_ID               = "tokenizer.ggml.eom_token_id"
+        UNK_ID               = "tokenizer.ggml.unknown_token_id"
+        SEP_ID               = "tokenizer.ggml.seperator_token_id"
+        PAD_ID               = "tokenizer.ggml.padding_token_id"
+        MASK_ID              = "tokenizer.ggml.mask_token_id"
+        ADD_BOS              = "tokenizer.ggml.add_bos_token"
+        ADD_EOS              = "tokenizer.ggml.add_eos_token"
+        ADD_SEP              = "tokenizer.ggml.add_sep_token"
+        ADD_PREFIX           = "tokenizer.ggml.add_space_prefix"
+        REMOVE_EXTRA_WS      = "tokenizer.ggml.remove_extra_whitespaces"
+        PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap"
+        SUPPRESS_TOKENS      = "tokenizer.ggml.suppress_tokens"
+        HF_JSON              = "tokenizer.huggingface.json"
+        RWKV                 = "tokenizer.rwkv.world"
+        CHAT_TEMPLATE        = "tokenizer.chat_template"
+        CHAT_TEMPLATE_N      = "tokenizer.chat_template.{name}"
+        CHAT_TEMPLATES       = "tokenizer.chat_templates"
+        # Normalizer constants
+        NORMALIZER_NFC           = "tokenizer.ggml.normalizer.nfc"
+        NORMALIZER_LOWERCASE     = "tokenizer.ggml.normalizer.lowercase"
+        NORMALIZER_STRIP_ACCENTS = "tokenizer.ggml.normalizer.strip_accents"
+        # FIM/Infill special tokens constants
+        FIM_PRE_ID           = "tokenizer.ggml.fim_pre_token_id"
+        FIM_SUF_ID           = "tokenizer.ggml.fim_suf_token_id"
+        FIM_MID_ID           = "tokenizer.ggml.fim_mid_token_id"
+        FIM_PAD_ID           = "tokenizer.ggml.fim_pad_token_id"
+        FIM_REP_ID           = "tokenizer.ggml.fim_rep_token_id"
+        FIM_SEP_ID           = "tokenizer.ggml.fim_sep_token_id"
+        # deprecated:
+        PREFIX_ID            = "tokenizer.ggml.prefix_token_id"
+        SUFFIX_ID            = "tokenizer.ggml.suffix_token_id"
+        MIDDLE_ID            = "tokenizer.ggml.middle_token_id"
+
+    class Adapter:
+        TYPE                    = "adapter.type"
+        LORA_ALPHA              = "adapter.lora.alpha"
+        LORA_TASK_NAME          = "adapter.lora.task_name"
+        LORA_PROMPT_PREFIX      = "adapter.lora.prompt_prefix"
+        ALORA_INVOCATION_TOKENS = "adapter.alora.invocation_tokens"
+
+    class IMatrix:
+        CHUNK_COUNT = "imatrix.chunk_count"
+        CHUNK_SIZE  = "imatrix.chunk_size"
+        DATASETS    = "imatrix.datasets"
+
+    class Clip:
+        PROJECTOR_TYPE        = "clip.projector_type"
+        HAS_VISION_ENCODER    = "clip.has_vision_encoder"
+        HAS_AUDIO_ENCODER     = "clip.has_audio_encoder"
+        HAS_GEN_AUDIO_ENCODER = "clip.has_gen_audio_encoder"
+        HAS_LLAVA_PROJECTOR   = "clip.has_llava_projector"
+
+    class ClipVision:
+        PROJECTOR_TYPE        = "clip.vision.projector_type" # for mixed modality models
+        IMAGE_SIZE            = "clip.vision.image_size"
+        IMAGE_MIN_PIXELS      = "clip.vision.image_min_pixels"
+        IMAGE_MAX_PIXELS      = "clip.vision.image_max_pixels"
+        PREPROC_MIN_TILES     = "clip.vision.preproc_min_tiles"
+        PREPROC_MAX_TILES     = "clip.vision.preproc_max_tiles"
+        PREPROC_IMAGE_SIZE    = "clip.vision.preproc_image_size"
+        PATCH_SIZE            = "clip.vision.patch_size"
+        EMBEDDING_LENGTH      = "clip.vision.embedding_length"
+        FEED_FORWARD_LENGTH   = "clip.vision.feed_forward_length"
+        PROJECTION_DIM        = "clip.vision.projection_dim"
+        BLOCK_COUNT           = "clip.vision.block_count"
+        IMAGE_MEAN            = "clip.vision.image_mean"
+        IMAGE_STD             = "clip.vision.image_std"
+        MAX_SLICE_NUMS        = "clip.vision.max_slice_nums"
+        IMAGE_RESIZE_ALGO     = "clip.vision.image_resize_algo"
+        SPATIAL_MERGE_SIZE    = "clip.vision.spatial_merge_size"
+        SWIGLU_CLAMP          = "clip.vision.swiglu_clamp"
+        EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer
+        EXPERT_USED_COUNT     = "clip.vision.expert_used_count"
+        USE_GELU              = "clip.use_gelu"
+        USE_SILU              = "clip.use_silu"
+        N_WA_PATTERN          = "clip.vision.n_wa_pattern"     # used by qwen2.5vl
+        WA_LAYER_INDEXES      = "clip.vision.wa_layer_indexes" # used by youtuvl
+        WA_PATTERN_MODE       = "clip.vision.wa_pattern_mode"  # used by mimovl, per-layer -1/0/1
+        IS_DEEPSTACK_LAYERS   = "clip.vision.is_deepstack_layers"
+        WINDOW_SIZE           = "clip.vision.window_size"
+        FEATURE_LAYERS        = "clip.vision.feature_layer" # Granite4 Vision
+        IMAGE_GRID_PINPOINTS  = "clip.vision.image_grid_pinpoints" # Granite4 Vision
+
+        class Attention:
+            HEAD_COUNT      = "clip.vision.attention.head_count"
+            HEAD_COUNT_KV   = "clip.vision.attention.head_count_kv" # used by mimovl (GQA)
+            HEAD_DIM        = "clip.vision.attention.head_dim" # set when qkv width != n_embd
+            LAYERNORM_EPS   = "clip.vision.attention.layer_norm_epsilon"
+
+        class Projector:
+            SCALE_FACTOR    = "clip.vision.projector.scale_factor"
+            QUERY_SIDE      = "clip.vision.projector.query_side"
+            WINDOW_SIDE     = "clip.vision.projector.window_side"
+            SPATIAL_OFFSETS = "clip.vision.projector.spatial_offsets"
+
+        class SAM:
+            BLOCK_COUNT         = "clip.vision.sam.block_count"
+            EMBEDDING_LENGTH    = "clip.vision.sam.embedding_length"
+            HEAD_COUNT          = "clip.vision.sam.head_count"
+
+    class ClipAudio:
+        PROJECTOR_TYPE      = "clip.audio.projector_type" # for mixed modality models
+        NUM_MEL_BINS        = "clip.audio.num_mel_bins"
+        EMBEDDING_LENGTH    = "clip.audio.embedding_length"
+        FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length"
+        PROJECTION_DIM      = "clip.audio.projection_dim"
+        BLOCK_COUNT         = "clip.audio.block_count"
+        SUBSAMPLING_FACTOR  = "clip.audio.subsampling_factor"
+        CHUNK_SIZE          = "clip.audio.chunk_size"
+        CONV_KERNEL_SIZE    = "clip.audio.conv_kernel_size"
+        MAX_POS_EMB         = "clip.audio.max_pos_emb"
+        FEATURE_LAYERS      = "clip.audio.feature_layer" # Granite Speech Plus
+        RVQ_NUM_QUANTIZERS  = "clip.audio.rvq.num_quantizers"
+        RVQ_CODEBOOK_SIZE   = "clip.audio.rvq.codebook_size"
+        WA_PATTERN_MODE     = "clip.audio.wa_pattern_mode" # per-layer -1 (full) / 0 (windowed)
+        WINDOW_SIZE         = "clip.audio.window_size"
+        LOCAL_BLOCK_COUNT   = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count
+        LOCAL_GROUP_SIZE    = "clip.audio.local_group_size"  # mimo-v2.5: input_local_transformer grouping size
+
+        class Attention:
+            HEAD_COUNT      = "clip.audio.attention.head_count"
+            LAYERNORM_EPS   = "clip.audio.attention.layer_norm_epsilon"
+
+        class Projector:
+            STACK_FACTOR    = "clip.audio.projector.stack_factor"
+            WINDOW_SIZE     = "clip.audio.projector.window_size"
+            DOWNSAMPLE_RATE = "clip.audio.projector.downsample_rate"
+            HEAD_COUNT      = "clip.audio.projector.head_count"
+
+    class ClipGenAudio:
+        PROJECTOR_TYPE      = "clip.gen.audio.projector_type" # for mixed modality models
+        # name of the weight variant, for settings that are not in the checkpoint
+        MODEL_VARIANT       = "clip.gen.audio.model_variant"
+        EMBEDDING_LENGTH    = "clip.gen.audio.embedding_length"
+        FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length"
+        BLOCK_COUNT         = "clip.gen.audio.block_count"
+        PROJECTION_DIM      = "clip.gen.audio.projection_dim"
+
+        class Attention:
+            HEAD_COUNT      = "clip.gen.audio.attention.head_count"
+            HEAD_COUNT_KV   = "clip.gen.audio.attention.head_count_kv"
+            LAYERNORM_EPS   = "clip.gen.audio.attention.layer_norm_epsilon"
+
+    class Diffusion:
+        SHIFT_LOGITS        = "diffusion.shift_logits"
+
+    class xIELU:
+        ALPHA_P             = "xielu.alpha_p"
+        ALPHA_N             = "xielu.alpha_n"
+        BETA                = "xielu.beta"
+        EPS                 = "xielu.eps"
+
+    class PrismHadamard:
+        VERSION              = "prism.hadamard.version"
+        TIED_OUTPUT          = "prism.hadamard.tied_output"
+        BLOCK_SIZE           = "prism.hadamard.block_size"
+        TRANSFORM            = "prism.hadamard.transform"
+        AXIS                 = "prism.hadamard.axis"
+        SIGN_MODE            = "prism.hadamard.sign_mode"
+        SIGN_WIDTHS          = "prism.hadamard.sign_widths"
+        SIGN_VALUES          = "prism.hadamard.sign_values"
+        WEIGHT_NAMES         = "prism.hadamard.weight_names"
+        INVERSE_WEIGHT_NAMES = "prism.hadamard.inverse_weight_names"
+        GDN_V_GROUPED        = "prism.hadamard.gdn_v_grouped"
+
+
+#
+# recommended mapping of model tensor names for storage in gguf
+#
+
+
+class GGUFType:
+    MODEL   = "model"
+    ADAPTER = "adapter"
+    IMATRIX = "imatrix"
+    MMPROJ  = "mmproj" # dummy, unused for now
+
+
+class MODEL_ARCH(IntEnum):
+    MMPROJ           = auto() # dummy arch for clip.cpp
+    LLAMA            = auto()
+    LLAMA4           = auto()
+    DECI             = auto()
+    FALCON           = auto()
+    FALCON_H1        = auto()
+    BAICHUAN         = auto()
+    GROK             = auto()
+    GPT2             = auto()
+    GPTJ             = auto()
+    GPTNEOX          = auto()
+    MPT              = auto()
+    STARCODER        = auto()
+    REFACT           = auto()
+    BERT             = auto()
+    MODERN_BERT      = auto()
+    NOMIC_BERT       = auto()
+    NOMIC_BERT_MOE   = auto()
+    NEO_BERT         = auto()
+    JINA_BERT_V2     = auto()
+    JINA_BERT_V3     = auto()
+    EUROBERT         = auto()
+    BLOOM            = auto()
+    STABLELM         = auto()
+    QWEN             = auto()
+    QWEN2            = auto()
+    QWEN2MOE         = auto()
+    QWEN2VL          = auto()
+    QWEN3            = auto()
+    QWEN3MOE         = auto()
+    QWEN3NEXT        = auto()
+    HRM_TEXT         = auto()
+    QWEN3VL          = auto()
+    QWEN3VLMOE       = auto()
+    QWEN35           = auto()
+    QWEN35MOE        = auto()
+    CLEF             = auto()
+    VLLM_DECISION2   = auto()
+    QWEN4EXP         = auto()
+    PHI2             = auto()
+    PHI3             = auto()
+    PHIMOE           = auto()
+    PLAMO            = auto()
+    PLAMO2           = auto()
+    PLAMO3           = auto()
+    CODESHELL        = auto()
+    ORION            = auto()
+    INTERNLM2        = auto()
+    MINICPM          = auto()
+    MINICPM3         = auto()
+    GEMMA            = auto()
+    GEMMA2           = auto()
+    GEMMA3           = auto()
+    GEMMA3N          = auto()
+    GEMMA4           = auto()
+    GEMMA4_ASSISTANT = auto()
+    GEMMA_EMBEDDING  = auto()
+    GEMMA_EMBEDDING2 = auto()
+    STARCODER2       = auto()
+    RWKV6            = auto()
+    RWKV6QWEN2       = auto()
+    RWKV7            = auto()
+    ARWKV7           = auto()
+    MAMBA            = auto()
+    MAMBA2           = auto()
+    MAPLE            = auto()
+    JAMBA            = auto()
+    XVERSE           = auto()
+    COMMAND_R        = auto()
+    COHERE2          = auto()
+    COHERE2MOE       = auto()
+    DBRX             = auto()
+    OLMO             = auto()
+    OLMO2            = auto()
+    OLMOE            = auto()
+    MUSE_GLIMMER     = auto()
+    OPENELM          = auto()
+    ARCTIC           = auto()
+    DEEPSEEK         = auto()
+    DEEPSEEK2        = auto()
+    DEEPSEEK2OCR     = auto()
+    DEEPSEEK32       = auto()
+    DEEPSEEK4        = auto()
+    CHATGLM          = auto()
+    GLM4             = auto()
+    GLM4_MOE         = auto()
+    GLM_DSA          = auto()
+    GLM5_NEXT        = auto()
+    BITNET           = auto()
+    T5               = auto()
+    T5ENCODER        = auto()
+    JAIS             = auto()
+    JAIS2            = auto()
+    NEMOTRON         = auto()
+    NEMOTRON_H       = auto()
+    NEMOTRON_H_MOE   = auto()
+    EXAONE           = auto()
+    EXAONE4          = auto()
+    EXAONE_MOE       = auto()
+    GRANITE          = auto()
+    GRANITE_MOE      = auto()
+    GRANITE_HYBRID   = auto()
+    GRANITE_SWITCH   = auto()
+    GRANITE_SWA      = auto()
+    CHAMELEON        = auto()
+    WAVTOKENIZER_DEC = auto()
+    PLM              = auto()
+    BAILINGMOE       = auto()
+    BAILINGMOE2      = auto()
+    BAILINGMOE3      = auto()
+    DOTS1            = auto()
+    DOTS3NOTE        = auto()
+    ARCEE            = auto()
+    AFMOE            = auto()
+    LAGUNA           = auto()
+    ERNIE4_5         = auto()
+    ERNIE4_5_MOE     = auto()
+    HUNYUAN_MOE      = auto()
+    HUNYUAN_DENSE    = auto()
+    HUNYUAN_VL       = auto()
+    HY_V3            = auto()
+    HY_V4            = auto()
+    SMOLLM3          = auto()
+    GPT_OSS          = auto()
+    LFM2             = auto()
+    LFM2MOE          = auto()
+    DREAM            = auto()
+    SMALLTHINKER     = auto()
+    LLADA            = auto()
+    LLADA_MOE        = auto()
+    SEED_OSS         = auto()
+    GROVEMOE         = auto()
+    APERTUS          = auto()
+    COGVLM           = auto()
+    MINIMAX01        = auto()
+    MINIMAXM2        = auto()
+    MINIMAXM3        = auto()
+    RND1             = auto()
+    PANGU_EMBED      = auto()
+    MISTRAL3         = auto()
+    EAGLE3           = auto()
+    DFLASH           = auto()
+    MISTRAL4         = auto()
+    PADDLEOCR        = auto()
+    MIMO2            = auto()
+    STEP35           = auto()
+    SPARK2_5           = auto()
+    LLAMA_EMBED      = auto()
+    MAINCODER        = auto()
+    KIMI_LINEAR      = auto()
+    KIMI_K3          = auto()
+    TALKIE           = auto()
+    MELLUM           = auto()
+    NANBEIGE         = auto()
+    QWEN3TTS         = auto()
+    POCKETTTS        = auto()
+    K2HORIZON        = auto()
+
+
+class VISION_PROJECTOR_TYPE(IntEnum):
+    MLP       = auto()
+    LDP       = auto()
+    LDPV2     = auto()
+    RESAMPLER = auto()
+    GLM_EDGE  = auto()
+    MERGER    = auto()
+    GEMMA3N   = auto()
+    GEMMA3    = auto()
+    QWEN3VL   = auto()
+    LING3VL   = auto()
+    STEP3VL   = auto()
+    COGVLM    = auto()
+
+
+class MODEL_TENSOR(IntEnum):
+    TOKEN_EMBD           = auto()
+    TOKEN_EMBD_NORM      = auto()
+    MASKED_EMBD_CENTROIDS= auto()
+    MASKED_EMBD_ORDERING = auto()
+    TOKEN_TYPES          = auto()
+    POS_EMBD             = auto()
+    OUTPUT               = auto()
+    HRM_Z_L_INIT         = auto()
+    DENSE_2_OUT          = auto() # embeddinggemma 2_Dense
+    DENSE_3_OUT          = auto() # embeddinggemma 3_Dense
+    OUTPUT_NORM          = auto()
+    HC_HEAD_FN           = auto()
+    HC_HEAD_BASE         = auto()
+    HC_HEAD_SCALE        = auto()
+    HC_HEAD_NORM         = auto() # qwen4exp
+    HC_HEAD_DOWN         = auto() # qwen4exp
+    HC_HEAD_UP           = auto() # qwen4exp
+    ROPE_FREQS           = auto()
+    ROPE_FACTORS_LONG    = auto()
+    ROPE_FACTORS_SHORT   = auto()
+    ATTN_Q               = auto()
+    ATTN_K               = auto()
+    ATTN_V               = auto()
+    ATTN_QKV             = auto()
+    ATTN_OUT             = auto()
+    ATTN_NORM            = auto()
+    ATTN_NORM_2          = auto()
+    ATTN_OUT_NORM        = auto()
+    ATTN_POST_NORM       = auto()
+    ATTN_ROT_EMBD        = auto()
+    ATTN_SINKS           = auto()
+    ATTN_GATE            = auto()
+    FFN_GATE_INP         = auto()
+    FFN_GATE_INP_SHEXP   = auto()
+    FFN_NORM             = auto()
+    FFN_PRE_NORM         = auto() # alias of FFN_NORM
+    FFN_PRE_NORM_2       = auto() # gemma4
+    FFN_POST_NORM        = auto()
+    FFN_POST_NORM_1      = auto() # gemma4
+    FFN_POST_NORM_2      = auto() # gemma4
+    FFN_GATE             = auto()
+    FFN_DOWN             = auto()
+    FFN_UP               = auto()
+    FFN_ACT              = auto()
+    FFN_NORM_EXP         = auto()
+    FFN_GATE_EXP         = auto()
+    FFN_DOWN_EXP         = auto()
+    FFN_UP_EXP           = auto()
+    FFN_GATE_UP_EXP      = auto()
+    FFN_GATE_SHEXP       = auto()
+    FFN_DOWN_SHEXP       = auto()
+    FFN_UP_SHEXP         = auto()
+    FFN_GATE_CHEXP       = auto()
+    FFN_DOWN_CHEXP       = auto()
+    FFN_UP_CHEXP         = auto()
+    FFN_EXP_PROBS_B      = auto()
+    FFN_EXP_PROBS_B_VL   = auto() # deepseek4 vision (bias for image tokens)
+    FFN_GATE_TID2EID     = auto()
+    MOE_LATENT_DOWN      = auto() # nemotron 3 super
+    MOE_LATENT_UP        = auto() # nemotron 3 super
+    ATTN_Q_NORM          = auto()
+    ATTN_K_NORM          = auto()
+    LAYER_OUT_NORM       = auto()
+    LAYER_OUT_SCALE      = auto()
+    PER_LAYER_TOKEN_EMBD = auto() # gemma3n
+    PER_LAYER_MODEL_PROJ = auto() # gemma3n
+    PER_LAYER_INP_GATE   = auto() # gemma3n
+    PER_LAYER_PROJ       = auto() # gemma3n
+    PER_LAYER_PROJ_NORM  = auto() # gemma3n
+    PER_LAYER_POST_NORM  = auto() # gemma3n
+    ALTUP_PROJ           = auto() # gemma3n
+    ALTUP_UNEMBD_PROJ    = auto() # gemma3n
+    ALTUP_CORRECT_COEF   = auto() # gemma3n
+    ALTUP_CORRECT_SCALE  = auto() # gemma3n
+    ALTUP_PREDICT_COEF   = auto() # gemma3n
+    ALTUP_ROUTER         = auto() # gemma3n
+    ALTUP_ROUTER_NORM    = auto() # gemma3n
+    LAUREL_L             = auto() # gemma3n
+    LAUREL_R             = auto() # gemma3n
+    LAUREL_POST_NORM     = auto() # gemma3n
+    SSM_IN               = auto()
+    SSM_CONV1D           = auto()
+    SSM_X                = auto()
+    SSM_DT               = auto()
+    SSM_DT_NORM          = auto()
+    SSM_A                = auto()
+    SSM_B_NORM           = auto()
+    SSM_C_NORM           = auto()
+    SSM_D                = auto()
+    SSM_NORM             = auto()
+    SSM_OUT              = auto()
+    SSM_ALPHA            = auto() # qwen3.5
+    SSM_BETA_ALPHA       = auto() # qwen3next
+    SSM_CONV1D_Q         = auto() # Kimi Linear
+    SSM_CONV1D_K         = auto() # Kimi Linear
+    SSM_CONV1D_V         = auto() # Kimi Linear
+    SSM_F_A              = auto() # Kimi Linear
+    SSM_F_B              = auto() # Kimi Linear
+    SSM_BETA             = auto() # Kimi Linear qwen3.5
+    SSM_G_A              = auto() # Kimi Linear
+    SSM_G_B              = auto() # Kimi Linear
+    SSM_G                = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B)
+    ATTN_RES_SCORE       = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention)
+    FFN_RES_SCORE        = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN)
+    OUTPUT_RES_SCORE     = auto() # Kimi K3 (fused res_norm * res_proj, final)
+    FFN_ROUTED_DOWN      = auto() # Kimi K3 (latent MoE: hidden -> latent)
+    FFN_ROUTED_UP        = auto() # Kimi K3 (latent MoE: latent -> hidden)
+    FFN_ROUTED_NORM      = auto() # Kimi K3 (latent MoE: norm on expert output)
+    TIME_MIX_W0          = auto()
+    TIME_MIX_W1          = auto()
+    TIME_MIX_W2          = auto()
+    TIME_MIX_A0          = auto()
+    TIME_MIX_A1          = auto()
+    TIME_MIX_A2          = auto()
+    TIME_MIX_V0          = auto()
+    TIME_MIX_V1          = auto()
+    TIME_MIX_V2          = auto()
+    TIME_MIX_G1          = auto()
+    TIME_MIX_G2          = auto()
+    TIME_MIX_K_K         = auto()
+    TIME_MIX_K_A         = auto()
+    TIME_MIX_R_K         = auto()
+    TIME_MIX_LERP_X      = auto()
+    TIME_MIX_LERP_K      = auto()
+    TIME_MIX_LERP_V      = auto()
+    TIME_MIX_LERP_R      = auto()
+    TIME_MIX_LERP_G      = auto()
+    TIME_MIX_LERP_FUSED  = auto()
+    TIME_MIX_LERP_W      = auto()
+    TIME_MIX_FIRST       = auto()
+    TIME_MIX_DECAY       = auto()
+    TIME_MIX_DECAY_W1    = auto()
+    TIME_MIX_DECAY_W2    = auto()
+    TIME_MIX_KEY         = auto()
+    TIME_MIX_VALUE       = auto()
+    TIME_MIX_RECEPTANCE  = auto()
+    TIME_MIX_GATE        = auto()
+    TIME_MIX_LN          = auto()
+    TIME_MIX_OUTPUT      = auto()
+    CHANNEL_MIX_LERP_K   = auto()
+    CHANNEL_MIX_LERP_R   = auto()
+    CHANNEL_MIX_KEY      = auto()
+    CHANNEL_MIX_RECEPTANCE = auto()
+    CHANNEL_MIX_VALUE    = auto()
+    ATTN_Q_A             = auto()
+    ATTN_Q_B             = auto()
+    ATTN_KV_A_MQA        = auto()
+    ATTN_KV_B            = auto()
+    ATTN_K_B             = auto()
+    ATTN_V_B             = auto()
+    ATTN_Q_A_NORM        = auto()
+    ATTN_KV_A_NORM       = auto()
+    ATTN_KV              = auto()
+    ATTN_KV_NORM         = auto()
+    ATTN_OUT_A           = auto()
+    ATTN_OUT_B           = auto()
+    HC_ATTN_FN           = auto()
+    HC_ATTN_BASE         = auto()
+    HC_ATTN_SCALE        = auto()
+    HC_FFN_FN            = auto()
+    HC_FFN_BASE          = auto()
+    HC_FFN_SCALE         = auto()
+    HC_ATTN_NORM         = auto() # qwen4exp
+    HC_ATTN_DOWN         = auto() # qwen4exp
+    HC_ATTN_UP           = auto() # qwen4exp
+    HC_ATTN_INJECT       = auto() # qwen4exp
+    HC_FFN_NORM          = auto() # qwen4exp
+    HC_FFN_DOWN          = auto() # qwen4exp
+    HC_FFN_UP            = auto() # qwen4exp
+    HC_FFN_INJECT        = auto() # qwen4exp
+    PLE_KEY              = auto() # qwen4exp
+    PLE_VALUE            = auto() # qwen4exp
+    PLE_NORM_KEY         = auto() # qwen4exp
+    PLE_NORM_QUERY       = auto() # qwen4exp
+    PLE_NORM_CONV        = auto() # qwen4exp
+    PLE_CONV1D           = auto() # qwen4exp
+    ATTN_COMPRESSOR_WKV  = auto()
+    ATTN_COMPRESSOR_WGATE = auto()
+    ATTN_COMPRESSOR_APE  = auto()
+    ATTN_COMPRESSOR_NORM = auto()
+    FFN_SUB_NORM         = auto()
+    ATTN_SUB_NORM        = auto()
+    DEC_ATTN_NORM        = auto()
+    DEC_ATTN_Q           = auto()
+    DEC_ATTN_K           = auto()
+    DEC_ATTN_V           = auto()
+    DEC_ATTN_OUT         = auto()
+    DEC_ATTN_REL_B       = auto()
+    DEC_CROSS_ATTN_NORM  = auto()
+    DEC_CROSS_ATTN_NORM_KV = auto()
+    DEC_CROSS_ATTN_Q     = auto()
+    DEC_CROSS_ATTN_K     = auto()
+    DEC_CROSS_ATTN_V     = auto()
+    DEC_CROSS_ATTN_OUT   = auto()
+    DEC_CROSS_ATTN_REL_B = auto()
+    DEC_FFN_NORM         = auto()
+    DEC_FFN_GATE         = auto()
+    DEC_FFN_DOWN         = auto()
+    DEC_FFN_UP           = auto()
+    DEC_OUTPUT_NORM      = auto()
+    ENC_ATTN_NORM        = auto()
+    ENC_ATTN_Q           = auto()
+    ENC_ATTN_K           = auto()
+    ENC_ATTN_V           = auto()
+    ENC_ATTN_OUT         = auto()
+    ENC_ATTN_REL_B       = auto()
+    ENC_FFN_NORM         = auto()
+    ENC_FFN_GATE         = auto()
+    ENC_FFN_DOWN         = auto()
+    ENC_FFN_UP           = auto()
+    ENC_OUTPUT_NORM      = auto()
+    CLS                  = auto() # classifier
+    CLS_OUT              = auto() # classifier output projection
+    CLS_NORM             = auto()
+    DECISION_HIDDEN_NORM          = auto()
+    DECISION_PROJ_MEMORY          = auto()
+    DECISION_PROJ_QUESTION        = auto()
+    DECISION_PROJ_OPTION_QUESTION = auto()
+    DECISION_PROJ_GLOBAL          = auto()
+    DECISION_PROJ_OPTION_CONTEXT  = auto()
+    DECISION_PROJ_OPTION_LEXICAL  = auto()
+    DECISION_OPTION_SUMMARY_NORM  = auto()
+    DECISION_FIELD_NORM           = auto()
+    DECISION_OPTION_NORM          = auto()
+    DECISION_SCALES               = auto()
+    DECISION_SCORER               = auto()
+    DECISION_SCORER_OUT           = auto()
+    DECISION_QUERY_NORM           = auto()
+    DECISION_KEY                  = auto()
+    DECISION_QUERY                = auto()
+    DECISION_OPTION_MLP           = auto()
+    DECISION_QUERY_MLP            = auto()
+    CONV1D               = auto()
+    CONVNEXT_DW          = auto()
+    CONVNEXT_NORM        = auto()
+    CONVNEXT_PW1         = auto()
+    CONVNEXT_PW2         = auto()
+    CONVNEXT_GAMMA       = auto()
+    POSNET_CONV1         = auto()
+    POSNET_CONV2         = auto()
+    POSNET_NORM          = auto()
+    POSNET_NORM1         = auto()
+    POSNET_NORM2         = auto()
+    POSNET_ATTN_NORM     = auto()
+    POSNET_ATTN_Q        = auto()
+    POSNET_ATTN_K        = auto()
+    POSNET_ATTN_V        = auto()
+    POSNET_ATTN_OUT      = auto()
+    SHORTCONV_CONV       = auto()
+    SHORTCONV_INPROJ     = auto()
+    SHORTCONV_OUTPROJ    = auto()
+    VISEXP_ATTN_QKV      = auto()
+    VISEXP_ATTN_OUT      = auto()
+    VISEXP_GATE          = auto()
+    VISEXP_DOWN          = auto()
+    VISEXP_UP            = auto()
+    INDEXER_K_NORM       = auto()
+    INDEXER_PROJ         = auto()
+    INDEXER_ATTN_K       = auto()
+    INDEXER_ATTN_Q_B     = auto()
+    INDEXER_Q_PROJ       = auto()
+    INDEXER_K_PROJ       = auto()
+    INDEXER_Q_NORM       = auto()
+    INDEXER_COMPRESSOR_WKV = auto()
+    INDEXER_COMPRESSOR_WGATE = auto()
+    INDEXER_COMPRESSOR_APE = auto()
+    INDEXER_COMPRESSOR_NORM = auto()
+    INDEXER_KPOOL_GATE   = auto()
+    INDEXER_KPOOL_APE    = auto()
+    ATTN_V_GATE          = auto() # k2-horizon MoVA router
+    ATTN_V_EXP           = auto() # k2-horizon MoVA value experts
+    # vision
+    V_MMPROJ             = auto()
+    V_MMPROJ_FC          = auto()
+    V_MMPROJ_MLP         = auto()
+    V_MMPROJ_PEG         = auto()
+    V_ENC_EMBD_CLS       = auto()
+    V_ENC_EMBD_PATCH     = auto()
+    V_ENC_EMBD_NORM      = auto()
+    V_ENC_EMBD_PATCH_NORM = auto() # allow multiple norms in the same embd, e.g. for gemma4u
+    V_ENC_EMBD_POS       = auto()
+    V_ENC_INPUT_NORM     = auto()
+    V_ENC_ATTN_QKV       = auto()
+    V_ENC_ATTN_Q         = auto()
+    V_ENC_ATTN_Q_NORM    = auto()
+    V_ENC_ATTN_K         = auto()
+    V_ENC_ATTN_K_NORM    = auto()
+    V_ENC_ATTN_V         = auto()
+    V_ENC_ATTN_O         = auto()
+    V_ENC_ATTN_O_NORM    = auto()
+    V_ENC_ATTN_SINKS     = auto() # mimovl
+    V_ENC_POST_ATTN_NORM = auto()
+    V_ENC_FFN_UP         = auto()
+    V_ENC_FFN_GATE       = auto()
+    V_ENC_FFN_DOWN       = auto()
+    V_ENC_FFN_GATE_INP   = auto() # dots3note vision MoE router
+    V_ENC_FFN_GATE_EXPS  = auto()
+    V_ENC_FFN_UP_EXPS    = auto()
+    V_ENC_FFN_DOWN_EXPS  = auto()
+    V_ENC_FFN_EXP_PROBS_B = auto()
+    V_ENC_ATTN_POST_NORM = auto() # gemma4
+    V_ENC_FFN_POST_NORM  = auto()
+    V_LAYER_SCALE_1      = auto()
+    V_LAYER_SCALE_2      = auto()
+    V_LAYER_OUT_SCALE    = auto()
+    V_PRE_NORM           = auto()
+    V_POST_NORM          = auto()
+    V_MM_PRE_NORM        = auto() # hunyuanvl
+    V_MM_POST_NORM       = auto()
+    V_MM_INP_NORM        = auto()
+    V_MM_INP_PROJ        = auto() # gemma3
+    V_MM_SOFT_EMB_NORM   = auto() # gemma3
+    V_MM_EMBEDDING       = auto() # gemma3n
+    V_MM_HARD_EMB_NORM   = auto() # gemma3n
+    V_ENC_CONV_STEM      = auto() # gemma3n
+    V_ENC_CONV_STEM_NORM = auto() # gemma3n
+    V_ENC_MSFA_EXP       = auto() # gemma3n
+    V_ENC_MSFA_EXP_NORM  = auto() # gemma3n
+    V_ENC_MSFA_PROJ      = auto() # gemma3n
+    V_ENC_MSFA_PROJ_NORM = auto() # gemma3n
+    V_ENC_MSFA_NORM      = auto() # gemma3n
+    V_RESMPL_POS_EMBD_K  = auto() # minicpmv
+    V_RESMPL_ATTN_Q      = auto() # minicpmv
+    V_RESMPL_ATTN_K      = auto() # minicpmv
+    V_RESMPL_ATTN_V      = auto() # minicpmv
+    V_RESMPL_ATTN_OUT    = auto() # minicpmv
+    V_RESMPL_KV          = auto() # minicpmv
+    V_RESMPL_KV_NORM     = auto() # minicpmv
+    V_RESMPL_POST_NORM   = auto() # minicpmv
+    V_RESMPL_Q_NORM      = auto() # minicpmv
+    V_RESMPL_PROJ        = auto() # minicpmv
+    V_RESMPL_QUERY       = auto() # minicpmv
+    V_TOK_EMBD_IMG_BREAK = auto() # pixtral
+    V_TOK_EMBD_IMG_START = auto() # deepseek4v
+    V_TOK_EMBD_IMG_END   = auto() # deepseek4v
+    V_TOK_EMBD_IMG_PAD   = auto() # deepseek4v
+    V_MM_PATCH_MERGER    = auto() # mistral small 3.1
+    V_DS_NORM            = auto() # qwen3vl
+    V_DS_FC1             = auto() # qwen3vl
+    V_DS_FC2             = auto() # qwen3vl
+    V_MERGER_LN1         = auto() # minicpmv4_6
+    V_MERGER_ATTN_Q      = auto() # minicpmv4_6
+    V_MERGER_ATTN_K      = auto() # minicpmv4_6
+    V_MERGER_ATTN_V      = auto() # minicpmv4_6
+    V_MERGER_ATTN_O      = auto() # minicpmv4_6
+    V_MERGER_DS_LN       = auto() # minicpmv4_6
+    V_MERGER_DS_UP       = auto() # minicpmv4_6
+    V_MERGER_DS_DOWN     = auto() # minicpmv4_6
+    V_MM_POST_FC_NORM    = auto() # cogvlm
+    V_MM_UP              = auto() # cogvlm
+    V_MM_DOWN            = auto() # cogvlm
+    V_MM_GATE            = auto() # cogvlm
+    V_MM_MERGER_FC1      = auto() # minimax-m3 (patch-merge MLP)
+    V_MM_MERGER_FC2      = auto() # minimax-m3 (patch-merge MLP)
+    V_TOK_BOI            = auto() # cogvlm
+    V_TOK_EOI            = auto() # cogvlm
+    V_TOK_IMG_BEGIN      = auto() # hunyuanvl
+    V_TOK_IMG_END        = auto() # hunyuanvl
+    V_STD_BIAS           = auto() # gemma4
+    V_STD_SCALE          = auto() # gemma4
+    V_SAM_POS_EMBD       = auto() # Deepseek-OCR
+    V_SAM_PATCH_EMBD     = auto() # Deepseek-OCR
+    V_SAM_PRE_NORM       = auto() # Deepseek-OCR
+    V_SAM_POST_NORM      = auto() # Deepseek-OCR
+    V_SAM_ATTN_POS_H     = auto() # Deepseek-OCR
+    V_SAM_ATTN_POS_W     = auto() # Deepseek-OCR
+    V_SAM_ATTN_QKV       = auto() # Deepseek-OCR
+    V_SAM_ATTN_OUT       = auto() # Deepseek-OCR
+    V_SAM_MLP_LIN_1      = auto() # Deepseek-OCR
+    V_SAM_MLP_LIN_2      = auto() # Deepseek-OCR
+    V_SAM_NECK           = auto() # Deepseek-OCR
+    V_SAM_NET_2          = auto() # Deepseek-OCR
+    V_SAM_NET_3          = auto() # Deepseek-OCR
+    V_ENC_EMBD_IMGNL     = auto() # Deepseek-OCR
+    V_ENC_EMBD_VSEP      = auto() # Deepseek-OCR
+    V_TOK_EMBD_SEP       = auto() # MiniCPM-V 4.7
+    V_RESMPL_QUERY_768   = auto() # Deepseek-OCR-2
+    V_RESMPL_QUERY_1024  = auto() # Deepseek-OCR-2
+
+    # qformer projector (vision) - Granite4 Vision
+    V_QF_PROJ_QUERY      = auto()
+    V_QF_PROJ_NORM       = auto()
+    V_QF_PROJ_LINEAR     = auto()
+    V_QF_SELF_ATTN_Q     = auto()
+    V_QF_SELF_ATTN_K     = auto()
+    V_QF_SELF_ATTN_V     = auto()
+    V_QF_SELF_ATTN_O     = auto()
+    V_QF_SELF_ATTN_NORM  = auto()
+    V_QF_CROSS_ATTN_Q    = auto()
+    V_QF_CROSS_ATTN_K    = auto()
+    V_QF_CROSS_ATTN_V    = auto()
+    V_QF_CROSS_ATTN_O    = auto()
+    V_QF_CROSS_ATTN_NORM = auto()
+    V_QF_FFN_UP          = auto()
+    V_QF_FFN_DOWN        = auto()
+    V_QF_FFN_NORM        = auto()
+    V_PROJ_NORM          = auto()
+    # multi-projector (bid => projector id) - Granite4 vision
+    V_MULTI_PROJ_IMG_POS   = auto()
+    V_MULTI_PROJ_QUERY     = auto()
+    V_MULTI_PROJ_NORM      = auto()
+    V_MULTI_PROJ_LINEAR    = auto()
+    V_MULTI_PROJ_POST_NORM = auto()
+
+    # audio (mtmd)
+    A_ENC_EMBD_POS        = auto()
+    A_ENC_EMBD_NORM       = auto()
+    A_ENC_EMBD_TO_LOGITS  = auto() # lfm2
+    A_ENC_INP_PROJ        = auto() # gemma4
+    A_ENC_CONV1D          = auto()
+    A_ENC_CONV1D_NORM     = auto() # gemma3n
+    A_ENC_CONV2D          = auto()
+    A_ENC_CONV_OUT        = auto()
+    A_PRE_NORM            = auto()
+    A_POST_NORM           = auto()
+    A_ENC_LAYER_PRE_NORM  = auto() # gemma3n
+    A_ENC_ATTN_Q          = auto()
+    A_ENC_ATTN_K          = auto()
+    A_ENC_ATTN_V          = auto()
+    A_ENC_ATTN_POST_NORM  = auto()
+    A_ENC_ATTN_PRE_NORM   = auto()
+    A_ENC_ATTN_K_REL      = auto() # gemma4
+    A_ENC_PER_DIM_SCALE   = auto() # gemma3n
+    A_ENC_INPUT_NORM      = auto()
+    A_ENC_OUTPUT          = auto() # TODO @ngxson: rename to ATTN_OUT
+    A_ENC_OUTPUT_NORM     = auto() # TODO @ngxson: rename to ATTN_OUT
+    A_ENC_FFN_UP          = auto()
+    A_ENC_FFN_NORM        = auto()
+    A_ENC_FFN_POST_NORM   = auto() # gemma3n
+    A_ENC_FFN_SCALE       = auto() # gemma3n
+    A_ENC_FFN_GATE        = auto()
+    A_ENC_FFN_DOWN        = auto()
+    A_ENC_FFN_UP_1        = auto() # lfm2, gemma3n
+    A_ENC_FFN_NORM_1      = auto() # lfm2, gemma3n (pre-norm)
+    A_ENC_FFN_POST_NORM_1 = auto() # gemma3n
+    A_ENC_FFN_SCALE_1     = auto() # gemma3n
+    A_ENC_FFN_GATE_1      = auto() # lfm2, gemma3n
+    A_ENC_FFN_DOWN_1      = auto() # lfm2, gemma3n
+    A_ENC_DOWNSAMPLE_CONV = auto() # mimo-audio-tokenizer: post-transformer downsample conv
+    A_ENC_DOWNSAMPLE_NORM = auto() # mimo-audio-tokenizer: post-transformer downsample norm
+    A_ENC_RVQ_CODEBOOK    = auto() # mimo-audio-tokenizer: residual vector quantizer codebook, per quantizer index
+    A_ENC_CONV_RES2       = auto() # qwen3tts
+    A_ENC_SE_CONV1        = auto() # qwen3tts
+    A_ENC_SE_CONV2        = auto() # qwen3tts
+    A_ENC_ASP_ATTN        = auto() # qwen3tts
+    A_ENC_ASP_TDNN        = auto() # qwen3tts
+    # qwen3tts code_predictor: predicts the remaining RVQ codebooks
+    A_GEN_CODE_PROJ_IN     = auto() # small_to_mtp_projection
+    A_GEN_CODE_EMBD        = auto() # per-codebook embedding table, merged 3D [n_codebooks, vocab, dim]
+    A_GEN_CODE_HEAD        = auto() # per-codebook output head, merged 3D [n_codebooks, vocab, dim]
+    A_GEN_CODE_OUT_EMBD    = auto() # codebook-0 embedding, re-fed into the talker backbone (talker.model.codec_embedding)
+    A_GEN_CODE_ATTN_NORM   = auto()
+    A_GEN_CODE_ATTN_Q      = auto()
+    A_GEN_CODE_ATTN_Q_NORM = auto()
+    A_GEN_CODE_ATTN_K      = auto()
+    A_GEN_CODE_ATTN_K_NORM = auto()
+    A_GEN_CODE_ATTN_V      = auto()
+    A_GEN_CODE_ATTN_OUT    = auto()
+    A_GEN_CODE_FFN_NORM    = auto()
+    A_GEN_CODE_FFN_GATE    = auto()
+    A_GEN_CODE_FFN_UP      = auto()
+    A_GEN_CODE_FFN_DOWN    = auto()
+    A_GEN_CODE_OUTPUT_NORM = auto()
+    # qwen3tts code2wav: RVQ codes -> raw PCM
+    A_GEN_WAV_QUANT_FIRST_IN       = auto() # semantic RVQ, in_proj (1x1 conv, loaded as 2D)
+    A_GEN_WAV_QUANT_FIRST_OUT      = auto() # semantic RVQ, out_proj
+    A_GEN_WAV_QUANT_FIRST_CB = auto() # semantic RVQ codebook (1 layer), folded from embedding_sum/cluster_usage
+    A_GEN_WAV_QUANT_REST_IN        = auto() # acoustic RVQ, in_proj
+    A_GEN_WAV_QUANT_REST_OUT       = auto() # acoustic RVQ, out_proj
+    A_GEN_WAV_QUANT_REST_CB  = auto() # acoustic RVQ codebooks, merged 3D [15, vocab, dim]
+    A_GEN_WAV_PRE_CONV             = auto()
+    A_GEN_WAV_TFM_IN_PROJ          = auto()
+    A_GEN_WAV_TFM_OUT_PROJ         = auto()
+    A_GEN_WAV_TFM_OUTPUT_NORM      = auto()
+    A_GEN_WAV_TFM_ATTN_NORM        = auto()
+    A_GEN_WAV_TFM_ATTN_Q           = auto()
+    A_GEN_WAV_TFM_ATTN_K           = auto()
+    A_GEN_WAV_TFM_ATTN_V           = auto()
+    A_GEN_WAV_TFM_ATTN_OUT         = auto()
+    A_GEN_WAV_TFM_ATTN_SCALE       = auto() # layer scale (gamma) on the attn output
+    A_GEN_WAV_TFM_FFN_NORM         = auto()
+    A_GEN_WAV_TFM_FFN_GATE         = auto()
+    A_GEN_WAV_TFM_FFN_UP           = auto()
+    A_GEN_WAV_TFM_FFN_DOWN         = auto()
+    A_GEN_WAV_TFM_FFN_SCALE        = auto() # layer scale (gamma) on the FFN output
+    A_GEN_WAV_UP_CONV              = auto() # causal ConvTranspose1d, 2x upsample
+    A_GEN_WAV_UP_DWCONV            = auto() # ConvNeXt depthwise conv
+    A_GEN_WAV_UP_NORM              = auto() # ConvNeXt LayerNorm
+    A_GEN_WAV_UP_PW1               = auto() # ConvNeXt pointwise conv 1 (expand)
+    A_GEN_WAV_UP_PW2               = auto() # ConvNeXt pointwise conv 2 (project)
+    A_GEN_WAV_UP_GAMMA             = auto() # ConvNeXt layer scale
+    A_GEN_WAV_DAC_ENTRY            = auto() # DAC conv_pre
+    A_GEN_WAV_DAC_UP_SNAKE         = auto() # DAC per-block SnakeBeta before the upsample conv
+    A_GEN_WAV_DAC_UP_CONV          = auto() # DAC per-block causal ConvTranspose1d
+    A_GEN_WAV_DAC_RES_ACT1         = auto() # DAC residual unit, SnakeBeta before conv1
+    A_GEN_WAV_DAC_RES_CONV1        = auto() # DAC residual unit, dilated causal conv
+    A_GEN_WAV_DAC_RES_ACT2         = auto() # DAC residual unit, SnakeBeta before conv2
+    A_GEN_WAV_DAC_RES_CONV2        = auto() # DAC residual unit, pointwise causal conv
+    A_GEN_WAV_DAC_POST_SNAKE       = auto() # DAC final SnakeBeta
+    A_GEN_WAV_DAC_POST_CONV        = auto() # DAC conv_post -> 1-channel PCM
+    # pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path)
+    A_ENC_SEANET_CONV_IN     = auto()
+    A_ENC_SEANET_CONV_OUT    = auto()
+    A_ENC_SEANET_RES_CONV1   = auto() # residual unit, dilated conv
+    A_ENC_SEANET_RES_CONV2   = auto() # residual unit, pointwise conv
+    A_ENC_SEANET_SCALE_CONV  = auto() # strided downsample conv
+    A_ENC_ATTN_SCALE         = auto() # layer scale (gamma) on the attn output
+    A_ENC_FFN_SCALE_LS       = auto() # layer scale (gamma) on the FFN output
+    A_ENC_SPEAKER_PROJ       = auto() # voice latent -> backbone embd
+    A_GEN_FLOW_INPUT_PROJ    = auto()
+    A_GEN_FLOW_COND_EMBD     = auto()
+    A_GEN_FLOW_TIME_FREQS    = auto() # timestep embedder, stored cos/sin frequencies
+    A_GEN_FLOW_TIME_UP       = auto()
+    A_GEN_FLOW_TIME_DOWN     = auto()
+    A_GEN_FLOW_TIME_NORM     = auto() # RMSNorm alpha
+    A_GEN_FLOW_BLK_NORM      = auto() # AdaLN res block, in_ln
+    A_GEN_FLOW_BLK_UP        = auto()
+    A_GEN_FLOW_BLK_DOWN      = auto()
+    A_GEN_FLOW_BLK_ADA       = auto() # AdaLN modulation, -> shift/scale/gate
+    A_GEN_FLOW_FINAL_ADA     = auto() # final layer AdaLN modulation, -> shift/scale
+    A_GEN_FLOW_FINAL_PROJ    = auto()
+    A_GEN_OUT_EOS            = auto() # end-of-speech head on the backbone hidden state
+    A_GEN_INPUT_LINEAR       = auto() # generated latent -> backbone embd
+    A_GEN_EMB_MEAN           = auto() # latent denormalization stats
+    A_GEN_EMB_STD            = auto()
+    A_GEN_WAV_QUANT_OUT      = auto() # DummyQuantizer output_proj, latent -> decoder dim
+    A_GEN_WAV_UPSAMPLE       = auto() # frame rate -> encoder frame rate, depthwise convtr
+    A_GEN_WAV_SEANET_CONV_IN   = auto()
+    A_GEN_WAV_SEANET_CONV_OUT  = auto() # -> 1-channel PCM
+    A_GEN_WAV_SEANET_RES_CONV1 = auto()
+    A_GEN_WAV_SEANET_RES_CONV2 = auto()
+    A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr
+    A_MMPROJ              = auto()
+    A_MMPROJ_FC           = auto()
+    A_MM_NORM_PRE         = auto()
+    A_MM_NORM_MID         = auto()
+    A_MM_EMBEDDING        = auto() # gemma3n
+    A_MM_HARD_EMB_NORM    = auto() # gemma3n
+    A_MM_SOFT_EMB_NORM    = auto() # gemma3n
+    A_MM_INP_PROJ         = auto() # gemma3n
+    A_MM_CODE_EMBD        = auto() # mimo: text-side RVQ code embedding table ("text codebook"), merged 3D [n_channels, vocab, dim]
+    A_MM_LOCAL_ATTN_Q     = auto() # mimo: input_local_transformer (LLM-side connector)
+    A_MM_LOCAL_ATTN_K     = auto()
+    A_MM_LOCAL_ATTN_V     = auto()
+    A_MM_LOCAL_ATTN_OUT   = auto()
+    A_MM_LOCAL_FFN_GATE   = auto()
+    A_MM_LOCAL_FFN_UP     = auto()
+    A_MM_LOCAL_FFN_DOWN   = auto()
+    A_MM_LOCAL_LN1        = auto()
+    A_MM_LOCAL_LN2        = auto()
+    A_MM_LOCAL_NORM       = auto() # final norm after all input_local_transformer layers
+    A_PER_DIM_K_SCALE     = auto() # gemma4
+    A_PER_DIM_SCALE       = auto() # gemma4
+    # nextn/mtp
+    NEXTN_PROJ_PRE         = auto()
+    NEXTN_PROJ_POST        = auto()
+    NEXTN_EH_PROJ          = auto()
+    NEXTN_EMBED_TOKENS     = auto()
+    NEXTN_ENORM            = auto()
+    NEXTN_HNORM            = auto()
+    NEXTN_SHARED_HEAD_HEAD = auto()
+    NEXTN_SHARED_HEAD_NORM = auto()
+    NEXTN_HC_HEAD_NORM     = auto()
+    NEXTN_HC_HEAD_DOWN     = auto()
+    NEXTN_HC_HEAD_UP       = auto()
+    # eagle3
+    FC                     = auto()  # feature fusion layer
+    D2T                    = auto()  # draft to target vocabulary mapping
+    # dspark
+    DSPARK_MARKOV_W1       = auto()  # markov head: prev-token embed
+    DSPARK_MARKOV_W2       = auto()  # markov head: bias projection
+    DSPARK_CONF_PROJ       = auto()  # confidence head
+    DFLASH_ATTN_CONV_BASE  = auto()
+    DFLASH_ATTN_CONV_PROJ  = auto()
+    DFLASH_FFN_CONV_BASE   = auto()
+    DFLASH_FFN_CONV_PROJ   = auto()
+    DFLASH_SELECTOR_PREV   = auto()
+    DFLASH_SELECTOR_NEXT   = auto()
+    DFLASH_SELECTOR_HIDDEN = auto()
+    # lfm2 audio
+    A_ENC_NORM_CONV        = auto()
+    A_ENC_LINEAR_POS       = auto()
+    A_ENC_POS_BIAS_U       = auto()
+    A_ENC_POS_BIAS_V       = auto()
+    A_ENC_OUT              = auto()
+    A_ENC_CONV_DW          = auto() # SSM conv
+    A_ENC_CONV_NORM        = auto() # SSM conv
+    A_ENC_CONV_PW1         = auto()
+    A_ENC_CONV_PW2         = auto()
+    A_ENC_CONV_NORM_MEAN   = auto() # parakeet
+    A_ENC_CONV_NORM_VAR    = auto() # parakeet
+    A_ENC_MEL_FILTERS      = auto() # parakeet
+    A_ENC_WINDOW           = auto() # parakeet
+    A_CTC_OUT              = auto()
+    A_CTC_OUT_MID          = auto()
+    A_ENC_ATTN_REL_POS_EMB = auto()
+    # audio qformer projector
+    A_QF_PROJ_QUERY        = auto()
+    A_QF_PROJ_NORM         = auto()
+    A_QF_PROJ_LINEAR       = auto()
+    A_QF_SELF_ATTN_Q       = auto()
+    A_QF_SELF_ATTN_K       = auto()
+    A_QF_SELF_ATTN_V       = auto()
+    A_QF_SELF_ATTN_O       = auto()
+    A_QF_SELF_ATTN_NORM    = auto()
+    A_QF_CROSS_ATTN_Q      = auto()
+    A_QF_CROSS_ATTN_K      = auto()
+    A_QF_CROSS_ATTN_V      = auto()
+    A_QF_CROSS_ATTN_O      = auto()
+    A_QF_CROSS_ATTN_NORM   = auto()
+    A_QF_FFN_UP            = auto()
+    A_QF_FFN_DOWN          = auto()
+    A_QF_FFN_NORM          = auto()
+
+
+MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
+    MODEL_ARCH.MMPROJ:           "clip", # dummy arch for clip.cpp
+    MODEL_ARCH.LLAMA:            "llama",
+    MODEL_ARCH.LLAMA4:           "llama4",
+    MODEL_ARCH.DECI:             "deci",
+    MODEL_ARCH.FALCON:           "falcon",
+    MODEL_ARCH.BAICHUAN:         "baichuan",
+    MODEL_ARCH.GROK:             "grok",
+    MODEL_ARCH.GPT2:             "gpt2",
+    MODEL_ARCH.GPTJ:             "gptj",
+    MODEL_ARCH.GPTNEOX:          "gptneox",
+    MODEL_ARCH.MPT:              "mpt",
+    MODEL_ARCH.STARCODER:        "starcoder",
+    MODEL_ARCH.REFACT:           "refact",
+    MODEL_ARCH.BERT:             "bert",
+    MODEL_ARCH.MODERN_BERT:      "modern-bert",
+    MODEL_ARCH.NOMIC_BERT:       "nomic-bert",
+    MODEL_ARCH.NOMIC_BERT_MOE:   "nomic-bert-moe",
+    MODEL_ARCH.NEO_BERT:         "neo-bert",
+    MODEL_ARCH.JINA_BERT_V2:     "jina-bert-v2",
+    MODEL_ARCH.JINA_BERT_V3:     "jina-bert-v3",
+    MODEL_ARCH.EUROBERT:         "eurobert",
+    MODEL_ARCH.BLOOM:            "bloom",
+    MODEL_ARCH.STABLELM:         "stablelm",
+    MODEL_ARCH.QWEN:             "qwen",
+    MODEL_ARCH.QWEN2:            "qwen2",
+    MODEL_ARCH.QWEN2MOE:         "qwen2moe",
+    MODEL_ARCH.QWEN2VL:          "qwen2vl",
+    MODEL_ARCH.QWEN3:            "qwen3",
+    MODEL_ARCH.QWEN3MOE:         "qwen3moe",
+    MODEL_ARCH.QWEN3NEXT:        "qwen3next",
+    MODEL_ARCH.HRM_TEXT:         "hrm_text",
+    MODEL_ARCH.QWEN3VL:          "qwen3vl",
+    MODEL_ARCH.QWEN3VLMOE:       "qwen3vlmoe",
+    MODEL_ARCH.QWEN35:           "qwen35",
+    MODEL_ARCH.QWEN35MOE:        "qwen35moe",
+    MODEL_ARCH.CLEF:             "clef",
+    MODEL_ARCH.VLLM_DECISION2:   "vllm-decision2",
+    MODEL_ARCH.QWEN4EXP:         "qwen4exp",
+    MODEL_ARCH.PHI2:             "phi2",
+    MODEL_ARCH.PHI3:             "phi3",
+    MODEL_ARCH.PHIMOE:           "phimoe",
+    MODEL_ARCH.PLAMO:            "plamo",
+    MODEL_ARCH.PLAMO2:           "plamo2",
+    MODEL_ARCH.PLAMO3:           "plamo3",
+    MODEL_ARCH.CODESHELL:        "codeshell",
+    MODEL_ARCH.ORION:            "orion",
+    MODEL_ARCH.INTERNLM2:        "internlm2",
+    MODEL_ARCH.MINICPM:          "minicpm",
+    MODEL_ARCH.MINICPM3:         "minicpm3",
+    MODEL_ARCH.GEMMA:            "gemma",
+    MODEL_ARCH.GEMMA2:           "gemma2",
+    MODEL_ARCH.GEMMA3:           "gemma3",
+    MODEL_ARCH.GEMMA3N:          "gemma3n",
+    MODEL_ARCH.GEMMA4:           "gemma4",
+    MODEL_ARCH.GEMMA4_ASSISTANT: "gemma4-assistant",
+    MODEL_ARCH.GEMMA_EMBEDDING:  "gemma-embedding",
+    MODEL_ARCH.GEMMA_EMBEDDING2: "gemma-embedding2",
+    MODEL_ARCH.STARCODER2:       "starcoder2",
+    MODEL_ARCH.RWKV6:            "rwkv6",
+    MODEL_ARCH.RWKV6QWEN2:       "rwkv6qwen2",
+    MODEL_ARCH.RWKV7:            "rwkv7",
+    MODEL_ARCH.ARWKV7:           "arwkv7",
+    MODEL_ARCH.MAMBA:            "mamba",
+    MODEL_ARCH.MAMBA2:           "mamba2",
+    MODEL_ARCH.MAPLE:            "maple",
+    MODEL_ARCH.JAMBA:            "jamba",
+    MODEL_ARCH.XVERSE:           "xverse",
+    MODEL_ARCH.COMMAND_R:        "command-r",
+    MODEL_ARCH.COHERE2:          "cohere2",
+    MODEL_ARCH.COHERE2MOE:       "cohere2moe",
+    MODEL_ARCH.DBRX:             "dbrx",
+    MODEL_ARCH.OLMO:             "olmo",
+    MODEL_ARCH.OLMO2:            "olmo2",
+    MODEL_ARCH.OLMOE:            "olmoe",
+    MODEL_ARCH.MUSE_GLIMMER:     "muse-glimmer",
+    MODEL_ARCH.OPENELM:          "openelm",
+    MODEL_ARCH.ARCTIC:           "arctic",
+    MODEL_ARCH.DEEPSEEK:         "deepseek",
+    MODEL_ARCH.DEEPSEEK2:        "deepseek2",
+    MODEL_ARCH.DEEPSEEK2OCR:     "deepseek2-ocr",
+    MODEL_ARCH.DEEPSEEK32:       "deepseek32",
+    MODEL_ARCH.DEEPSEEK4:        "deepseek4",
+    MODEL_ARCH.CHATGLM:          "chatglm",
+    MODEL_ARCH.GLM4:             "glm4",
+    MODEL_ARCH.GLM4_MOE:         "glm4moe",
+    MODEL_ARCH.GLM_DSA:          "glm-dsa",
+    MODEL_ARCH.GLM5_NEXT:        "glm5-next",
+    MODEL_ARCH.BITNET:           "bitnet",
+    MODEL_ARCH.T5:               "t5",
+    MODEL_ARCH.T5ENCODER:        "t5encoder",
+    MODEL_ARCH.JAIS:             "jais",
+    MODEL_ARCH.JAIS2:            "jais2",
+    MODEL_ARCH.NEMOTRON:         "nemotron",
+    MODEL_ARCH.NEMOTRON_H:       "nemotron_h",
+    MODEL_ARCH.NEMOTRON_H_MOE:   "nemotron_h_moe",
+    MODEL_ARCH.EXAONE:           "exaone",
+    MODEL_ARCH.EXAONE4:          "exaone4",
+    MODEL_ARCH.EXAONE_MOE:       "exaone-moe",
+    MODEL_ARCH.GRANITE:          "granite",
+    MODEL_ARCH.GRANITE_MOE:      "granitemoe",
+    MODEL_ARCH.GRANITE_HYBRID:   "granitehybrid",
+    MODEL_ARCH.GRANITE_SWITCH:   "graniteswitch",
+    MODEL_ARCH.GRANITE_SWA:      "granite_swa",
+    MODEL_ARCH.CHAMELEON:        "chameleon",
+    MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
+    MODEL_ARCH.PLM:              "plm",
+    MODEL_ARCH.BAILINGMOE:       "bailingmoe",
+    MODEL_ARCH.BAILINGMOE2:      "bailingmoe2",
+    MODEL_ARCH.BAILINGMOE3:      "bailingmoe3",
+    MODEL_ARCH.DOTS1:            "dots1",
+    MODEL_ARCH.DOTS3NOTE:        "dots3note",
+    MODEL_ARCH.ARCEE:            "arcee",
+    MODEL_ARCH.AFMOE:            "afmoe",
+    MODEL_ARCH.LAGUNA:           "laguna",
+    MODEL_ARCH.ERNIE4_5:         "ernie4_5",
+    MODEL_ARCH.ERNIE4_5_MOE:     "ernie4_5-moe",
+    MODEL_ARCH.FALCON_H1:        "falcon-h1",
+    MODEL_ARCH.HUNYUAN_MOE:      "hunyuan-moe",
+    MODEL_ARCH.HUNYUAN_DENSE:    "hunyuan-dense",
+    MODEL_ARCH.HUNYUAN_VL:       "hunyuan_vl",
+    MODEL_ARCH.HY_V3:            "hy_v3",
+    MODEL_ARCH.HY_V4:            "hy_v4",
+    MODEL_ARCH.SMOLLM3:          "smollm3",
+    MODEL_ARCH.GPT_OSS:          "gpt-oss",
+    MODEL_ARCH.LFM2:             "lfm2",
+    MODEL_ARCH.LFM2MOE:          "lfm2moe",
+    MODEL_ARCH.DREAM:            "dream",
+    MODEL_ARCH.SMALLTHINKER:     "smallthinker",
+    MODEL_ARCH.LLADA:            "llada",
+    MODEL_ARCH.LLADA_MOE:        "llada-moe",
+    MODEL_ARCH.SEED_OSS:         "seed_oss",
+    MODEL_ARCH.GROVEMOE:         "grovemoe",
+    MODEL_ARCH.APERTUS:          "apertus",
+    MODEL_ARCH.MINIMAX01:        "minimax-01",
+    MODEL_ARCH.MINIMAXM2:        "minimax-m2",
+    MODEL_ARCH.MINIMAXM3:        "minimax-m3",
+    MODEL_ARCH.COGVLM:           "cogvlm",
+    MODEL_ARCH.RND1:             "rnd1",
+    MODEL_ARCH.PANGU_EMBED:      "pangu-embedded",
+    MODEL_ARCH.MISTRAL3:         "mistral3",
+    MODEL_ARCH.EAGLE3:           "eagle3",
+    MODEL_ARCH.DFLASH:           "dflash",
+    MODEL_ARCH.MISTRAL4:         "mistral4",
+    MODEL_ARCH.PADDLEOCR:        "paddleocr",
+    MODEL_ARCH.MIMO2:            "mimo2",
+    MODEL_ARCH.STEP35:           "step35",
+    MODEL_ARCH.SPARK2_5:         "spark2_5",
+    MODEL_ARCH.LLAMA_EMBED:      "llama-embed",
+    MODEL_ARCH.MAINCODER:        "maincoder",
+    MODEL_ARCH.KIMI_LINEAR:      "kimi-linear",
+    MODEL_ARCH.KIMI_K3:          "kimi-k3",
+    MODEL_ARCH.TALKIE:           "talkie",
+    MODEL_ARCH.MELLUM:           "mellum",
+    MODEL_ARCH.NANBEIGE:         "nanbeige",
+    MODEL_ARCH.QWEN3TTS:         "qwen3tts",
+    MODEL_ARCH.POCKETTTS:        "pockettts",
+    MODEL_ARCH.K2HORIZON:        "k2-horizon",
+}
+
+VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
+    VISION_PROJECTOR_TYPE.MLP:       "mlp",
+    VISION_PROJECTOR_TYPE.LDP:       "ldp",
+    VISION_PROJECTOR_TYPE.LDPV2:     "ldpv2",
+    VISION_PROJECTOR_TYPE.RESAMPLER: "resampler",
+    VISION_PROJECTOR_TYPE.GLM_EDGE:  "adapter",
+    VISION_PROJECTOR_TYPE.MERGER:    "qwen2vl_merger",
+    VISION_PROJECTOR_TYPE.GEMMA3:    "gemma3",
+    VISION_PROJECTOR_TYPE.QWEN3VL:   "qwen3vl_merger",
+    VISION_PROJECTOR_TYPE.LING3VL:   "ling3vl",
+    VISION_PROJECTOR_TYPE.STEP3VL:   "step3vl",
+}
+
+TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
+    MODEL_TENSOR.TOKEN_EMBD:                "token_embd",
+    MODEL_TENSOR.TOKEN_EMBD_NORM:           "token_embd_norm",
+    MODEL_TENSOR.TOKEN_TYPES:               "token_types",
+    MODEL_TENSOR.MASKED_EMBD_CENTROIDS:     "masked_embd_centroids",
+    MODEL_TENSOR.MASKED_EMBD_ORDERING:      "masked_embd_ordering",
+    MODEL_TENSOR.POS_EMBD:                  "position_embd",
+    MODEL_TENSOR.OUTPUT_NORM:               "output_norm",
+    MODEL_TENSOR.OUTPUT:                    "output",
+    MODEL_TENSOR.HRM_Z_L_INIT:              "hrm.z_l_init",
+    MODEL_TENSOR.DENSE_2_OUT:               "dense_2", # embeddinggemma 2_Dense
+    MODEL_TENSOR.DENSE_3_OUT:               "dense_3", # embeddinggemma 2_Dense
+    MODEL_TENSOR.HC_HEAD_FN:                "output_hc_fn",
+    MODEL_TENSOR.HC_HEAD_BASE:              "output_hc_base",
+    MODEL_TENSOR.HC_HEAD_SCALE:             "output_hc_scale",
+    MODEL_TENSOR.HC_HEAD_NORM:              "output_hc_norm",                 # qwen4exp
+    MODEL_TENSOR.HC_HEAD_DOWN:              "output_hc_down",                 # qwen4exp
+    MODEL_TENSOR.HC_HEAD_UP:                "output_hc_up",                   # qwen4exp
+    MODEL_TENSOR.ROPE_FREQS:                "rope_freqs",
+    MODEL_TENSOR.ROPE_FACTORS_LONG:         "rope_factors_long",
+    MODEL_TENSOR.ROPE_FACTORS_SHORT:        "rope_factors_short",
+    MODEL_TENSOR.ATTN_NORM:                 "blk.{bid}.attn_norm",
+    MODEL_TENSOR.ATTN_NORM_2:               "blk.{bid}.attn_norm_2",
+    MODEL_TENSOR.ATTN_QKV:                  "blk.{bid}.attn_qkv",
+    MODEL_TENSOR.ATTN_Q:                    "blk.{bid}.attn_q",
+    MODEL_TENSOR.ATTN_K:                    "blk.{bid}.attn_k",
+    MODEL_TENSOR.ATTN_V:                    "blk.{bid}.attn_v",
+    MODEL_TENSOR.ATTN_OUT:                  "blk.{bid}.attn_output",
+    MODEL_TENSOR.ATTN_ROT_EMBD:             "blk.{bid}.attn_rot_embd",
+    MODEL_TENSOR.ATTN_SINKS:                "blk.{bid}.attn_sinks",
+    MODEL_TENSOR.ATTN_GATE:                 "blk.{bid}.attn_gate",
+    MODEL_TENSOR.ATTN_Q_NORM:               "blk.{bid}.attn_q_norm",
+    MODEL_TENSOR.ATTN_K_NORM:               "blk.{bid}.attn_k_norm",
+    MODEL_TENSOR.ATTN_OUT_NORM:             "blk.{bid}.attn_output_norm",
+    MODEL_TENSOR.ATTN_POST_NORM:            "blk.{bid}.post_attention_norm",
+    MODEL_TENSOR.FFN_GATE_INP:              "blk.{bid}.ffn_gate_inp",
+    MODEL_TENSOR.FFN_GATE_INP_SHEXP:        "blk.{bid}.ffn_gate_inp_shexp",
+    MODEL_TENSOR.FFN_NORM:                  "blk.{bid}.ffn_norm",
+    MODEL_TENSOR.FFN_PRE_NORM:              "blk.{bid}.ffn_norm",
+    MODEL_TENSOR.FFN_POST_NORM:             "blk.{bid}.post_ffw_norm",
+    MODEL_TENSOR.FFN_PRE_NORM_2:            "blk.{bid}.pre_ffw_norm_2",  # gemma4
+    MODEL_TENSOR.FFN_POST_NORM_1:           "blk.{bid}.post_ffw_norm_1", # gemma4
+    MODEL_TENSOR.FFN_POST_NORM_2:           "blk.{bid}.post_ffw_norm_2", # gemma4
+    MODEL_TENSOR.FFN_GATE:                  "blk.{bid}.ffn_gate",
+    MODEL_TENSOR.FFN_DOWN:                  "blk.{bid}.ffn_down",
+    MODEL_TENSOR.FFN_UP:                    "blk.{bid}.ffn_up",
+    MODEL_TENSOR.FFN_GATE_SHEXP:            "blk.{bid}.ffn_gate_shexp",
+    MODEL_TENSOR.FFN_DOWN_SHEXP:            "blk.{bid}.ffn_down_shexp",
+    MODEL_TENSOR.FFN_UP_SHEXP:              "blk.{bid}.ffn_up_shexp",
+    MODEL_TENSOR.FFN_GATE_CHEXP:            "blk.{bid}.ffn_gate_chexps",
+    MODEL_TENSOR.FFN_DOWN_CHEXP:            "blk.{bid}.ffn_down_chexps",
+    MODEL_TENSOR.FFN_UP_CHEXP:              "blk.{bid}.ffn_up_chexps",
+    MODEL_TENSOR.FFN_ACT:                   "blk.{bid}.ffn",
+    MODEL_TENSOR.FFN_NORM_EXP:              "blk.{bid}.ffn_norm_exps",
+    MODEL_TENSOR.FFN_GATE_EXP:              "blk.{bid}.ffn_gate_exps",
+    MODEL_TENSOR.FFN_DOWN_EXP:              "blk.{bid}.ffn_down_exps",
+    MODEL_TENSOR.FFN_UP_EXP:                "blk.{bid}.ffn_up_exps",
+    MODEL_TENSOR.FFN_GATE_UP_EXP:           "blk.{bid}.ffn_gate_up_exps",
+    MODEL_TENSOR.FFN_EXP_PROBS_B:           "blk.{bid}.exp_probs_b",
+    MODEL_TENSOR.FFN_EXP_PROBS_B_VL:        "blk.{bid}.exp_probs_b_vl",
+    MODEL_TENSOR.FFN_GATE_TID2EID:          "blk.{bid}.ffn_gate_tid2eid",
+    MODEL_TENSOR.MOE_LATENT_DOWN:           "blk.{bid}.ffn_latent_down",      # nemotron 3 super
+    MODEL_TENSOR.MOE_LATENT_UP:             "blk.{bid}.ffn_latent_up",        # nemotron 3 super
+    MODEL_TENSOR.LAYER_OUT_NORM:            "blk.{bid}.layer_output_norm",
+    MODEL_TENSOR.LAYER_OUT_SCALE:           "blk.{bid}.layer_output_scale",
+    MODEL_TENSOR.PER_LAYER_TOKEN_EMBD:      "per_layer_token_embd",           # gemma3n
+    MODEL_TENSOR.PER_LAYER_MODEL_PROJ:      "per_layer_model_proj",           # gemma3n
+    MODEL_TENSOR.PER_LAYER_PROJ_NORM:       "per_layer_proj_norm",            # gemma3n
+    MODEL_TENSOR.ALTUP_UNEMBD_PROJ:         "altup_unembd_proj",              # gemma3n
+    MODEL_TENSOR.ALTUP_PROJ:                "altup_proj",                     # gemma3n
+    MODEL_TENSOR.PER_LAYER_INP_GATE:        "blk.{bid}.inp_gate",             # gemma3n
+    MODEL_TENSOR.PER_LAYER_PROJ:            "blk.{bid}.proj",                 # gemma3n
+    MODEL_TENSOR.PER_LAYER_POST_NORM:       "blk.{bid}.post_norm",            # gemma3n
+    MODEL_TENSOR.ALTUP_CORRECT_COEF:        "blk.{bid}.altup_correct_coef",   # gemma3n
+    MODEL_TENSOR.ALTUP_CORRECT_SCALE:       "blk.{bid}.altup_correct_scale",  # gemma3n
+    MODEL_TENSOR.ALTUP_PREDICT_COEF:        "blk.{bid}.altup_predict_coef",   # gemma3n
+    MODEL_TENSOR.ALTUP_ROUTER:              "blk.{bid}.altup_router",         # gemma3n
+    MODEL_TENSOR.ALTUP_ROUTER_NORM:         "blk.{bid}.altup_router_norm",    # gemma3n
+    MODEL_TENSOR.LAUREL_L:                  "blk.{bid}.laurel_l",             # gemma3n
+    MODEL_TENSOR.LAUREL_R:                  "blk.{bid}.laurel_r",             # gemma3n
+    MODEL_TENSOR.LAUREL_POST_NORM:          "blk.{bid}.laurel_post_norm",     # gemma3n
+    MODEL_TENSOR.SSM_IN:                    "blk.{bid}.ssm_in",
+    MODEL_TENSOR.SSM_CONV1D:                "blk.{bid}.ssm_conv1d",
+    MODEL_TENSOR.SSM_X:                     "blk.{bid}.ssm_x",
+    MODEL_TENSOR.SSM_DT:                    "blk.{bid}.ssm_dt",
+    MODEL_TENSOR.SSM_DT_NORM:               "blk.{bid}.ssm_dt_norm",
+    MODEL_TENSOR.SSM_A:                     "blk.{bid}.ssm_a",
+    MODEL_TENSOR.SSM_B_NORM:                "blk.{bid}.ssm_b_norm",
+    MODEL_TENSOR.SSM_C_NORM:                "blk.{bid}.ssm_c_norm",
+    MODEL_TENSOR.SSM_D:                     "blk.{bid}.ssm_d",
+    MODEL_TENSOR.SSM_NORM:                  "blk.{bid}.ssm_norm",
+    MODEL_TENSOR.SSM_OUT:                   "blk.{bid}.ssm_out",
+    MODEL_TENSOR.SSM_ALPHA:                 "blk.{bid}.ssm_alpha",            # qwen3.5
+    MODEL_TENSOR.SSM_BETA_ALPHA:            "blk.{bid}.ssm_ba",
+    MODEL_TENSOR.SSM_CONV1D_Q:              "blk.{bid}.ssm_conv1d_q",         # Kimi Linear
+    MODEL_TENSOR.SSM_CONV1D_K:              "blk.{bid}.ssm_conv1d_k",         # Kimi Linear
+    MODEL_TENSOR.SSM_CONV1D_V:              "blk.{bid}.ssm_conv1d_v",         # Kimi Linear
+    MODEL_TENSOR.SSM_F_A:                   "blk.{bid}.ssm_f_a",              # Kimi Linear
+    MODEL_TENSOR.SSM_F_B:                   "blk.{bid}.ssm_f_b",              # Kimi Linear
+    MODEL_TENSOR.SSM_BETA:                  "blk.{bid}.ssm_beta",             # Kimi Linear qwen3.5
+    MODEL_TENSOR.SSM_G_A:                   "blk.{bid}.ssm_g_a",              # Kimi Linear
+    MODEL_TENSOR.SSM_G_B:                   "blk.{bid}.ssm_g_b",              # Kimi Linear
+    MODEL_TENSOR.SSM_G:                     "blk.{bid}.ssm_g",                # Kimi K3
+    MODEL_TENSOR.ATTN_RES_SCORE:            "blk.{bid}.attn_res_score",       # Kimi K3
+    MODEL_TENSOR.FFN_RES_SCORE:             "blk.{bid}.ffn_res_score",        # Kimi K3
+    MODEL_TENSOR.OUTPUT_RES_SCORE:          "output_res_score",               # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_DOWN:           "blk.{bid}.ffn_routed_down",      # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_UP:             "blk.{bid}.ffn_routed_up",        # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_NORM:           "blk.{bid}.ffn_routed_norm",      # Kimi K3
+    MODEL_TENSOR.TIME_MIX_W0:               "blk.{bid}.time_mix_w0",
+    MODEL_TENSOR.TIME_MIX_W1:               "blk.{bid}.time_mix_w1",
+    MODEL_TENSOR.TIME_MIX_W2:               "blk.{bid}.time_mix_w2",
+    MODEL_TENSOR.TIME_MIX_A0:               "blk.{bid}.time_mix_a0",
+    MODEL_TENSOR.TIME_MIX_A1:               "blk.{bid}.time_mix_a1",
+    MODEL_TENSOR.TIME_MIX_A2:               "blk.{bid}.time_mix_a2",
+    MODEL_TENSOR.TIME_MIX_V0:               "blk.{bid}.time_mix_v0",
+    MODEL_TENSOR.TIME_MIX_V1:               "blk.{bid}.time_mix_v1",
+    MODEL_TENSOR.TIME_MIX_V2:               "blk.{bid}.time_mix_v2",
+    MODEL_TENSOR.TIME_MIX_G1:               "blk.{bid}.time_mix_g1",
+    MODEL_TENSOR.TIME_MIX_G2:               "blk.{bid}.time_mix_g2",
+    MODEL_TENSOR.TIME_MIX_K_K:              "blk.{bid}.time_mix_k_k",
+    MODEL_TENSOR.TIME_MIX_K_A:              "blk.{bid}.time_mix_k_a",
+    MODEL_TENSOR.TIME_MIX_R_K:              "blk.{bid}.time_mix_r_k",
+    MODEL_TENSOR.TIME_MIX_LERP_X:           "blk.{bid}.time_mix_lerp_x",
+    MODEL_TENSOR.TIME_MIX_LERP_K:           "blk.{bid}.time_mix_lerp_k",
+    MODEL_TENSOR.TIME_MIX_LERP_V:           "blk.{bid}.time_mix_lerp_v",
+    MODEL_TENSOR.TIME_MIX_LERP_R:           "blk.{bid}.time_mix_lerp_r",
+    MODEL_TENSOR.TIME_MIX_LERP_G:           "blk.{bid}.time_mix_lerp_g",
+    MODEL_TENSOR.TIME_MIX_LERP_FUSED:       "blk.{bid}.time_mix_lerp_fused",
+    MODEL_TENSOR.TIME_MIX_LERP_W:           "blk.{bid}.time_mix_lerp_w",
+    MODEL_TENSOR.TIME_MIX_FIRST:            "blk.{bid}.time_mix_first",
+    MODEL_TENSOR.TIME_MIX_DECAY:            "blk.{bid}.time_mix_decay",
+    MODEL_TENSOR.TIME_MIX_DECAY_W1:         "blk.{bid}.time_mix_decay_w1",
+    MODEL_TENSOR.TIME_MIX_DECAY_W2:         "blk.{bid}.time_mix_decay_w2",
+    MODEL_TENSOR.TIME_MIX_KEY:              "blk.{bid}.time_mix_key",
+    MODEL_TENSOR.TIME_MIX_VALUE:            "blk.{bid}.time_mix_value",
+    MODEL_TENSOR.TIME_MIX_RECEPTANCE:       "blk.{bid}.time_mix_receptance",
+    MODEL_TENSOR.TIME_MIX_GATE:             "blk.{bid}.time_mix_gate",
+    MODEL_TENSOR.TIME_MIX_LN:               "blk.{bid}.time_mix_ln",
+    MODEL_TENSOR.TIME_MIX_OUTPUT:           "blk.{bid}.time_mix_output",
+    MODEL_TENSOR.CHANNEL_MIX_LERP_K:        "blk.{bid}.channel_mix_lerp_k",
+    MODEL_TENSOR.CHANNEL_MIX_LERP_R:        "blk.{bid}.channel_mix_lerp_r",
+    MODEL_TENSOR.CHANNEL_MIX_KEY:           "blk.{bid}.channel_mix_key",
+    MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE:    "blk.{bid}.channel_mix_receptance",
+    MODEL_TENSOR.CHANNEL_MIX_VALUE:         "blk.{bid}.channel_mix_value",
+    MODEL_TENSOR.ATTN_Q_A:                  "blk.{bid}.attn_q_a",
+    MODEL_TENSOR.ATTN_Q_B:                  "blk.{bid}.attn_q_b",
+    MODEL_TENSOR.ATTN_KV_A_MQA:             "blk.{bid}.attn_kv_a_mqa",
+    MODEL_TENSOR.ATTN_KV_B:                 "blk.{bid}.attn_kv_b",
+    MODEL_TENSOR.ATTN_K_B:                  "blk.{bid}.attn_k_b",
+    MODEL_TENSOR.ATTN_V_B:                  "blk.{bid}.attn_v_b",
+    MODEL_TENSOR.ATTN_Q_A_NORM:             "blk.{bid}.attn_q_a_norm",
+    MODEL_TENSOR.ATTN_KV_A_NORM:            "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_KV:                   "blk.{bid}.attn_kv",
+    MODEL_TENSOR.ATTN_KV_NORM:              "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_OUT_A:                "blk.{bid}.attn_output_a",
+    MODEL_TENSOR.ATTN_OUT_B:                "blk.{bid}.attn_output_b",
+    MODEL_TENSOR.HC_ATTN_FN:                "blk.{bid}.hc_attn_fn",
+    MODEL_TENSOR.HC_ATTN_BASE:              "blk.{bid}.hc_attn_base",
+    MODEL_TENSOR.HC_ATTN_SCALE:             "blk.{bid}.hc_attn_scale",
+    MODEL_TENSOR.HC_FFN_FN:                 "blk.{bid}.hc_ffn_fn",
+    MODEL_TENSOR.HC_FFN_BASE:               "blk.{bid}.hc_ffn_base",
+    MODEL_TENSOR.HC_FFN_SCALE:              "blk.{bid}.hc_ffn_scale",
+    MODEL_TENSOR.HC_ATTN_NORM:              "blk.{bid}.hc_attn_norm",         # qwen4exp
+    MODEL_TENSOR.HC_ATTN_DOWN:              "blk.{bid}.hc_attn_down",         # qwen4exp
+    MODEL_TENSOR.HC_ATTN_UP:                "blk.{bid}.hc_attn_up",           # qwen4exp
+    MODEL_TENSOR.HC_ATTN_INJECT:            "blk.{bid}.hc_attn_inject",       # qwen4exp
+    MODEL_TENSOR.HC_FFN_NORM:               "blk.{bid}.hc_ffn_norm",          # qwen4exp
+    MODEL_TENSOR.HC_FFN_DOWN:               "blk.{bid}.hc_ffn_down",          # qwen4exp
+    MODEL_TENSOR.HC_FFN_UP:                 "blk.{bid}.hc_ffn_up",            # qwen4exp
+    MODEL_TENSOR.HC_FFN_INJECT:             "blk.{bid}.hc_ffn_inject",        # qwen4exp
+    MODEL_TENSOR.PLE_KEY:                   "blk.{bid}.ple_key",              # qwen4exp
+    MODEL_TENSOR.PLE_VALUE:                 "blk.{bid}.ple_value",            # qwen4exp
+    MODEL_TENSOR.PLE_NORM_KEY:              "blk.{bid}.ple_norm_key",         # qwen4exp
+    MODEL_TENSOR.PLE_NORM_QUERY:            "blk.{bid}.ple_norm_query",       # qwen4exp
+    MODEL_TENSOR.PLE_NORM_CONV:             "blk.{bid}.ple_norm_conv",        # qwen4exp
+    MODEL_TENSOR.PLE_CONV1D:                "blk.{bid}.ple_conv1d",           # qwen4exp
+    MODEL_TENSOR.ATTN_COMPRESSOR_WKV:       "blk.{bid}.attn_compressor_kv",
+    MODEL_TENSOR.ATTN_COMPRESSOR_WGATE:     "blk.{bid}.attn_compressor_gate",
+    MODEL_TENSOR.ATTN_COMPRESSOR_APE:       "blk.{bid}.attn_compressor_ape",
+    MODEL_TENSOR.ATTN_COMPRESSOR_NORM:      "blk.{bid}.attn_compressor_norm",
+    MODEL_TENSOR.ATTN_SUB_NORM:             "blk.{bid}.attn_sub_norm",
+    MODEL_TENSOR.FFN_SUB_NORM:              "blk.{bid}.ffn_sub_norm",
+    MODEL_TENSOR.DEC_ATTN_NORM:             "dec.blk.{bid}.attn_norm",
+    MODEL_TENSOR.DEC_ATTN_Q:                "dec.blk.{bid}.attn_q",
+    MODEL_TENSOR.DEC_ATTN_K:                "dec.blk.{bid}.attn_k",
+    MODEL_TENSOR.DEC_ATTN_V:                "dec.blk.{bid}.attn_v",
+    MODEL_TENSOR.DEC_ATTN_OUT:              "dec.blk.{bid}.attn_o",
+    MODEL_TENSOR.DEC_ATTN_REL_B:            "dec.blk.{bid}.attn_rel_b",
+    MODEL_TENSOR.DEC_CROSS_ATTN_NORM:       "dec.blk.{bid}.cross_attn_norm",
+    MODEL_TENSOR.DEC_CROSS_ATTN_NORM_KV:    "dec.blk.{bid}.cross_attn_norm_kv",
+    MODEL_TENSOR.DEC_CROSS_ATTN_Q:          "dec.blk.{bid}.cross_attn_q",
+    MODEL_TENSOR.DEC_CROSS_ATTN_K:          "dec.blk.{bid}.cross_attn_k",
+    MODEL_TENSOR.DEC_CROSS_ATTN_V:          "dec.blk.{bid}.cross_attn_v",
+    MODEL_TENSOR.DEC_CROSS_ATTN_OUT:        "dec.blk.{bid}.cross_attn_o",
+    MODEL_TENSOR.DEC_CROSS_ATTN_REL_B:      "dec.blk.{bid}.cross_attn_rel_b",
+    MODEL_TENSOR.DEC_FFN_NORM:              "dec.blk.{bid}.ffn_norm",
+    MODEL_TENSOR.DEC_FFN_GATE:              "dec.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.DEC_FFN_DOWN:              "dec.blk.{bid}.ffn_down",
+    MODEL_TENSOR.DEC_FFN_UP:                "dec.blk.{bid}.ffn_up",
+    MODEL_TENSOR.DEC_OUTPUT_NORM:           "dec.output_norm",
+    MODEL_TENSOR.ENC_ATTN_NORM:             "enc.blk.{bid}.attn_norm",
+    MODEL_TENSOR.ENC_ATTN_Q:                "enc.blk.{bid}.attn_q",
+    MODEL_TENSOR.ENC_ATTN_K:                "enc.blk.{bid}.attn_k",
+    MODEL_TENSOR.ENC_ATTN_V:                "enc.blk.{bid}.attn_v",
+    MODEL_TENSOR.ENC_ATTN_OUT:              "enc.blk.{bid}.attn_o",
+    MODEL_TENSOR.ENC_ATTN_REL_B:            "enc.blk.{bid}.attn_rel_b",
+    MODEL_TENSOR.ENC_FFN_NORM:              "enc.blk.{bid}.ffn_norm",
+    MODEL_TENSOR.ENC_FFN_GATE:              "enc.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.ENC_FFN_DOWN:              "enc.blk.{bid}.ffn_down",
+    MODEL_TENSOR.ENC_FFN_UP:                "enc.blk.{bid}.ffn_up",
+    MODEL_TENSOR.ENC_OUTPUT_NORM:           "enc.output_norm",
+    MODEL_TENSOR.CLS:                       "cls",
+    MODEL_TENSOR.CLS_OUT:                   "cls.output",
+    MODEL_TENSOR.CLS_NORM:                  "cls.norm",
+    MODEL_TENSOR.DECISION_HIDDEN_NORM:          "decision.hidden_norm",
+    MODEL_TENSOR.DECISION_PROJ_MEMORY:          "decision.proj_memory",
+    MODEL_TENSOR.DECISION_PROJ_QUESTION:        "decision.proj_question",
+    MODEL_TENSOR.DECISION_PROJ_OPTION_QUESTION: "decision.proj_option_question",
+    MODEL_TENSOR.DECISION_PROJ_GLOBAL:          "decision.proj_global",
+    MODEL_TENSOR.DECISION_PROJ_OPTION_CONTEXT:  "decision.proj_option_context",
+    MODEL_TENSOR.DECISION_PROJ_OPTION_LEXICAL:  "decision.proj_option_lexical",
+    MODEL_TENSOR.DECISION_OPTION_SUMMARY_NORM:  "decision.option_summary_norm",
+    MODEL_TENSOR.DECISION_FIELD_NORM:           "decision.field_norm",
+    MODEL_TENSOR.DECISION_OPTION_NORM:          "decision.option_norm",
+    MODEL_TENSOR.DECISION_SCALES:               "decision.scales",
+    MODEL_TENSOR.DECISION_SCORER:               "decision.scorer",
+    MODEL_TENSOR.DECISION_SCORER_OUT:           "decision.scorer_out",
+    MODEL_TENSOR.DECISION_QUERY_NORM:           "decision.query_norm",
+    MODEL_TENSOR.DECISION_KEY:                  "decision.key",
+    MODEL_TENSOR.DECISION_QUERY:                "decision.query",
+    MODEL_TENSOR.DECISION_OPTION_MLP:           "decision.option_mlp",
+    MODEL_TENSOR.DECISION_QUERY_MLP:            "decision.query_mlp",
+    MODEL_TENSOR.CONV1D:                    "conv1d",
+    MODEL_TENSOR.CONVNEXT_DW:               "convnext.{bid}.dw",
+    MODEL_TENSOR.CONVNEXT_NORM:             "convnext.{bid}.norm",
+    MODEL_TENSOR.CONVNEXT_PW1:              "convnext.{bid}.pw1",
+    MODEL_TENSOR.CONVNEXT_PW2:              "convnext.{bid}.pw2",
+    MODEL_TENSOR.CONVNEXT_GAMMA:            "convnext.{bid}.gamma",
+    MODEL_TENSOR.POSNET_CONV1:              "posnet.{bid}.conv1",
+    MODEL_TENSOR.POSNET_CONV2:              "posnet.{bid}.conv2",
+    MODEL_TENSOR.POSNET_NORM:               "posnet.{bid}.norm",
+    MODEL_TENSOR.POSNET_NORM1:              "posnet.{bid}.norm1",
+    MODEL_TENSOR.POSNET_NORM2:              "posnet.{bid}.norm2",
+    MODEL_TENSOR.POSNET_ATTN_NORM:          "posnet.{bid}.attn_norm",
+    MODEL_TENSOR.POSNET_ATTN_Q:             "posnet.{bid}.attn_q",
+    MODEL_TENSOR.POSNET_ATTN_K:             "posnet.{bid}.attn_k",
+    MODEL_TENSOR.POSNET_ATTN_V:             "posnet.{bid}.attn_v",
+    MODEL_TENSOR.POSNET_ATTN_OUT:           "posnet.{bid}.attn_output",
+    MODEL_TENSOR.SHORTCONV_CONV:            "blk.{bid}.shortconv.conv",
+    MODEL_TENSOR.SHORTCONV_INPROJ:          "blk.{bid}.shortconv.in_proj",
+    MODEL_TENSOR.SHORTCONV_OUTPROJ:         "blk.{bid}.shortconv.out_proj",
+    MODEL_TENSOR.VISEXP_ATTN_QKV:           "blk.{bid}.vis_attn_qkv",
+    MODEL_TENSOR.VISEXP_ATTN_OUT:           "blk.{bid}.vis_attn_output",
+    MODEL_TENSOR.VISEXP_GATE:               "blk.{bid}.vis_gate",
+    MODEL_TENSOR.VISEXP_DOWN:               "blk.{bid}.vis_down",
+    MODEL_TENSOR.VISEXP_UP:                 "blk.{bid}.vis_up",
+    MODEL_TENSOR.INDEXER_K_NORM:            "blk.{bid}.indexer.k_norm",
+    MODEL_TENSOR.INDEXER_PROJ:              "blk.{bid}.indexer.proj",
+    MODEL_TENSOR.INDEXER_ATTN_K:            "blk.{bid}.indexer.attn_k",
+    MODEL_TENSOR.INDEXER_ATTN_Q_B:          "blk.{bid}.indexer.attn_q_b",
+    MODEL_TENSOR.INDEXER_Q_PROJ:            "blk.{bid}.indexer.q_proj",
+    MODEL_TENSOR.INDEXER_K_PROJ:            "blk.{bid}.indexer.k_proj",
+    MODEL_TENSOR.INDEXER_Q_NORM:            "blk.{bid}.indexer.q_norm",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_WKV:    "blk.{bid}.indexer_compressor_kv",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE:  "blk.{bid}.indexer_compressor_gate",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_APE:    "blk.{bid}.indexer_compressor_ape",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_NORM:   "blk.{bid}.indexer_compressor_norm",
+    MODEL_TENSOR.INDEXER_KPOOL_GATE:        "blk.{bid}.indexer_compressor_gate",
+    MODEL_TENSOR.INDEXER_KPOOL_APE:         "blk.{bid}.indexer_compressor_ape",
+    # vision
+    MODEL_TENSOR.V_MMPROJ:                  "mm.{bid}",
+    MODEL_TENSOR.V_MMPROJ_FC:               "mm.model.fc",
+    MODEL_TENSOR.V_MMPROJ_MLP:              "mm.model.mlp.{bid}",
+    MODEL_TENSOR.V_MMPROJ_PEG:              "mm.model.peg.{bid}",
+    MODEL_TENSOR.V_ENC_EMBD_CLS:            "v.class_embd",
+    MODEL_TENSOR.V_ENC_EMBD_PATCH:          "v.patch_embd",
+    MODEL_TENSOR.V_ENC_EMBD_NORM:           "v.norm_embd",
+    MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM:     "v.patch_norm.{bid}",
+    MODEL_TENSOR.V_ENC_EMBD_POS:            "v.position_embd",
+    MODEL_TENSOR.V_ENC_ATTN_QKV:            "v.blk.{bid}.attn_qkv",
+    MODEL_TENSOR.V_ENC_ATTN_Q:              "v.blk.{bid}.attn_q",
+    MODEL_TENSOR.V_ENC_ATTN_Q_NORM:         "v.blk.{bid}.attn_q_norm",
+    MODEL_TENSOR.V_ENC_ATTN_K:              "v.blk.{bid}.attn_k",
+    MODEL_TENSOR.V_ENC_ATTN_K_NORM:         "v.blk.{bid}.attn_k_norm",
+    MODEL_TENSOR.V_ENC_ATTN_V:              "v.blk.{bid}.attn_v",
+    MODEL_TENSOR.V_ENC_INPUT_NORM:          "v.blk.{bid}.ln1",
+    MODEL_TENSOR.V_ENC_ATTN_O:              "v.blk.{bid}.attn_out",
+    MODEL_TENSOR.V_ENC_ATTN_O_NORM:         "v.blk.{bid}.attn_out_norm",
+    MODEL_TENSOR.V_ENC_ATTN_SINKS:          "v.blk.{bid}.attn_sinks",
+    MODEL_TENSOR.V_ENC_POST_ATTN_NORM:      "v.blk.{bid}.ln2",
+    MODEL_TENSOR.V_ENC_FFN_UP:              "v.blk.{bid}.ffn_up",
+    MODEL_TENSOR.V_ENC_FFN_GATE:            "v.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.V_ENC_FFN_DOWN:            "v.blk.{bid}.ffn_down",
+    MODEL_TENSOR.V_ENC_FFN_GATE_INP:        "v.blk.{bid}.ffn_gate_inp",
+    MODEL_TENSOR.V_ENC_FFN_GATE_EXPS:       "v.blk.{bid}.ffn_gate_exps",
+    MODEL_TENSOR.V_ENC_FFN_UP_EXPS:         "v.blk.{bid}.ffn_up_exps",
+    MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS:       "v.blk.{bid}.ffn_down_exps",
+    MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B:     "v.blk.{bid}.exp_probs_b",
+    MODEL_TENSOR.V_ENC_ATTN_POST_NORM:      "v.blk.{bid}.attn_post_norm",
+    MODEL_TENSOR.V_ENC_FFN_POST_NORM:       "v.blk.{bid}.ffn_post_norm",
+    MODEL_TENSOR.V_LAYER_SCALE_1:           "v.blk.{bid}.ls1",
+    MODEL_TENSOR.V_LAYER_SCALE_2:           "v.blk.{bid}.ls2",
+    MODEL_TENSOR.V_LAYER_OUT_SCALE:         "v.blk.{bid}.out_scale",
+    MODEL_TENSOR.V_PRE_NORM:                "v.pre_ln",
+    MODEL_TENSOR.V_POST_NORM:               "v.post_ln",
+    MODEL_TENSOR.V_MM_POST_NORM:            "mm.post_norm",
+    MODEL_TENSOR.V_MM_INP_PROJ:             "mm.input_projection",
+    MODEL_TENSOR.V_MM_INP_NORM:             "mm.input_norm",
+    MODEL_TENSOR.V_MM_SOFT_EMB_NORM:        "mm.soft_emb_norm",         # gemma3n
+    MODEL_TENSOR.V_MM_EMBEDDING:            "mm.embedding",             # gemma3n
+    MODEL_TENSOR.V_MM_HARD_EMB_NORM:        "mm.hard_emb_norm",         # gemma3n
+    MODEL_TENSOR.V_ENC_CONV_STEM:           "v.conv_stem.conv",         # gemma3n
+    MODEL_TENSOR.V_ENC_CONV_STEM_NORM:      "v.conv_stem.bn",           # gemma3n
+    MODEL_TENSOR.V_ENC_MSFA_EXP:            "v.msfa.ffn.pw_exp.conv",   # gemma3n
+    MODEL_TENSOR.V_ENC_MSFA_EXP_NORM:       "v.msfa.ffn.pw_exp.bn",     # gemma3n
+    MODEL_TENSOR.V_ENC_MSFA_PROJ:           "v.msfa.ffn.pw_proj.conv",  # gemma3n
+    MODEL_TENSOR.V_ENC_MSFA_PROJ_NORM:      "v.msfa.ffn.pw_proj.bn",    # gemma3n
+    MODEL_TENSOR.V_ENC_MSFA_NORM:           "v.msfa.norm",              # gemma3n
+    MODEL_TENSOR.V_RESMPL_POS_EMBD_K:       "resampler.pos_embd_k",
+    MODEL_TENSOR.V_RESMPL_ATTN_Q:           "resampler.attn.q",
+    MODEL_TENSOR.V_RESMPL_ATTN_K:           "resampler.attn.k",
+    MODEL_TENSOR.V_RESMPL_ATTN_V:           "resampler.attn.v",
+    MODEL_TENSOR.V_RESMPL_ATTN_OUT:         "resampler.attn.out",
+    MODEL_TENSOR.V_RESMPL_KV:               "resampler.kv",
+    MODEL_TENSOR.V_RESMPL_KV_NORM:          "resampler.ln_kv",
+    MODEL_TENSOR.V_RESMPL_POST_NORM:        "resampler.ln_post",
+    MODEL_TENSOR.V_RESMPL_Q_NORM:           "resampler.ln_q",
+    MODEL_TENSOR.V_RESMPL_PROJ:             "resampler.proj",
+    MODEL_TENSOR.V_RESMPL_QUERY:            "resampler.query",
+    MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK:      "v.token_embd.img_break", # pixtral
+    MODEL_TENSOR.V_TOK_EMBD_IMG_START:      "v.token_embd.img_start", # deepseek4v
+    MODEL_TENSOR.V_TOK_EMBD_IMG_END:        "v.token_embd.img_end",   # deepseek4v
+    MODEL_TENSOR.V_TOK_EMBD_IMG_PAD:        "v.token_embd.img_pad",   # deepseek4v
+    MODEL_TENSOR.V_MM_PATCH_MERGER:         "mm.patch_merger", # mistral small 3.1
+    MODEL_TENSOR.V_DS_NORM:                 "v.deepstack.{bid}.norm",
+    MODEL_TENSOR.V_DS_FC1:                  "v.deepstack.{bid}.fc1",
+    MODEL_TENSOR.V_DS_FC2:                  "v.deepstack.{bid}.fc2",
+    MODEL_TENSOR.V_MERGER_LN1:              "v.vit_merger.ln1",
+    MODEL_TENSOR.V_MERGER_ATTN_Q:           "v.vit_merger.attn_q",
+    MODEL_TENSOR.V_MERGER_ATTN_K:           "v.vit_merger.attn_k",
+    MODEL_TENSOR.V_MERGER_ATTN_V:           "v.vit_merger.attn_v",
+    MODEL_TENSOR.V_MERGER_ATTN_O:           "v.vit_merger.attn_out",
+    MODEL_TENSOR.V_MERGER_DS_LN:            "v.vit_merger.ds_ln",
+    MODEL_TENSOR.V_MERGER_DS_UP:            "v.vit_merger.ds_ffn_up",
+    MODEL_TENSOR.V_MERGER_DS_DOWN:          "v.vit_merger.ds_ffn_down",
+    MODEL_TENSOR.V_MM_POST_FC_NORM:         "mm.post_fc_norm", # cogvlm
+    MODEL_TENSOR.V_MM_UP:                   "mm.up",
+    MODEL_TENSOR.V_MM_DOWN:                 "mm.down",
+    MODEL_TENSOR.V_MM_GATE:                 "mm.gate",
+    MODEL_TENSOR.V_MM_MERGER_FC1:           "mm.merger.fc1",
+    MODEL_TENSOR.V_MM_MERGER_FC2:           "mm.merger.fc2",
+    MODEL_TENSOR.V_TOK_BOI:                 "v.boi",
+    MODEL_TENSOR.V_TOK_EOI:                 "v.eoi",
+    MODEL_TENSOR.V_MM_PRE_NORM:             "mm.pre_norm",
+    MODEL_TENSOR.V_TOK_IMG_BEGIN:           "mm.image_begin",
+    MODEL_TENSOR.V_TOK_IMG_END:             "mm.image_end",
+    MODEL_TENSOR.V_STD_BIAS:                "v.std_bias", # gemma4
+    MODEL_TENSOR.V_STD_SCALE:               "v.std_scale", # gemma4
+    # DeepSeek-OCR SAM
+    MODEL_TENSOR.V_SAM_POS_EMBD:            "v.sam.pos_embd",
+    MODEL_TENSOR.V_SAM_PATCH_EMBD:          "v.sam.patch_embd",
+    MODEL_TENSOR.V_SAM_PRE_NORM:            "v.sam.blk.{bid}.pre_ln",
+    MODEL_TENSOR.V_SAM_POST_NORM:           "v.sam.blk.{bid}.post_ln",
+    MODEL_TENSOR.V_SAM_ATTN_POS_H:          "v.sam.blk.{bid}.attn.pos_h",
+    MODEL_TENSOR.V_SAM_ATTN_POS_W:          "v.sam.blk.{bid}.attn.pos_w",
+    MODEL_TENSOR.V_SAM_ATTN_QKV:            "v.sam.blk.{bid}.attn.qkv",
+    MODEL_TENSOR.V_SAM_ATTN_OUT:            "v.sam.blk.{bid}.attn.out",
+    MODEL_TENSOR.V_SAM_MLP_LIN_1:           "v.sam.blk.{bid}.mlp.lin1",
+    MODEL_TENSOR.V_SAM_MLP_LIN_2:           "v.sam.blk.{bid}.mlp.lin2",
+    MODEL_TENSOR.V_SAM_NECK:                "v.sam.neck.{bid}",
+    MODEL_TENSOR.V_SAM_NET_2:               "v.sam.net_2",
+    MODEL_TENSOR.V_SAM_NET_3:               "v.sam.net_3",
+    MODEL_TENSOR.V_ENC_EMBD_IMGNL:          "v.image_newline", # Deepseek-OCR, Granite4Vision
+    MODEL_TENSOR.V_ENC_EMBD_VSEP:           "v.view_seperator", # Deepseek-OCR
+    MODEL_TENSOR.V_TOK_EMBD_SEP:            "v.tok_embd_sep", # MiniCPM-V 4.7
+    MODEL_TENSOR.V_RESMPL_QUERY_768:        "v.resample_query_768", # Deepseek-OCR-2 qwen2
+    MODEL_TENSOR.V_RESMPL_QUERY_1024:       "v.resample_query_1024", # Deepseek-OCR-2 qwen2
+    # Granite4 Vision
+    # qformer layers (bid => proj_id)
+    # NOTE: Names align with A_QF_*
+    MODEL_TENSOR.V_QF_SELF_ATTN_Q:          "v.proj_blk.{bid}.self_attn_q",
+    MODEL_TENSOR.V_QF_SELF_ATTN_K:          "v.proj_blk.{bid}.self_attn_k",
+    MODEL_TENSOR.V_QF_SELF_ATTN_V:          "v.proj_blk.{bid}.self_attn_v",
+    MODEL_TENSOR.V_QF_SELF_ATTN_O:          "v.proj_blk.{bid}.self_attn_out",
+    MODEL_TENSOR.V_QF_SELF_ATTN_NORM:       "v.proj_blk.{bid}.self_attn_norm",
+    MODEL_TENSOR.V_QF_CROSS_ATTN_Q:         "v.proj_blk.{bid}.cross_attn_q",
+    MODEL_TENSOR.V_QF_CROSS_ATTN_K:         "v.proj_blk.{bid}.cross_attn_k",
+    MODEL_TENSOR.V_QF_CROSS_ATTN_V:         "v.proj_blk.{bid}.cross_attn_v",
+    MODEL_TENSOR.V_QF_CROSS_ATTN_O:         "v.proj_blk.{bid}.cross_attn_out",
+    MODEL_TENSOR.V_QF_CROSS_ATTN_NORM:      "v.proj_blk.{bid}.cross_attn_norm",
+    MODEL_TENSOR.V_QF_FFN_UP:               "v.proj_blk.{bid}.ffn_up",
+    MODEL_TENSOR.V_QF_FFN_DOWN:             "v.proj_blk.{bid}.ffn_down",
+    MODEL_TENSOR.V_QF_FFN_NORM:             "v.proj_blk.{bid}.ffn_norm",
+    # multi-projector (bid => projector ID)
+    MODEL_TENSOR.V_MULTI_PROJ_IMG_POS:   "v.proj_blk.{bid}.img_pos",
+    MODEL_TENSOR.V_MULTI_PROJ_QUERY:     "v.proj_blk.{bid}.query",
+    MODEL_TENSOR.V_MULTI_PROJ_NORM:      "v.proj_blk.{bid}.norm",
+    MODEL_TENSOR.V_MULTI_PROJ_LINEAR:    "v.proj_blk.{bid}.linear",
+    MODEL_TENSOR.V_MULTI_PROJ_POST_NORM: "v.proj_blk.{bid}.post_norm",
+
+    # audio (mtmd)
+    # note: all audio tensor names must use prefix "a." or "mm.a."
+    MODEL_TENSOR.A_ENC_EMBD_POS:            "a.position_embd",
+    MODEL_TENSOR.A_ENC_EMBD_NORM:           "a.position_embd_norm",
+    MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS:      "a.embd_to_logits",
+    MODEL_TENSOR.A_ENC_INP_PROJ:            "a.input_projection",
+    MODEL_TENSOR.A_ENC_CONV1D:              "a.conv1d.{bid}",
+    MODEL_TENSOR.A_ENC_CONV2D:              "a.conv2d.{bid}",
+    MODEL_TENSOR.A_ENC_CONV_OUT:            "a.conv_out",
+    MODEL_TENSOR.A_ENC_CONV1D_NORM:         "a.conv1d.{bid}.norm",
+    MODEL_TENSOR.A_PRE_NORM:                "a.pre_ln",
+    MODEL_TENSOR.A_POST_NORM:               "a.post_ln",
+    MODEL_TENSOR.A_ENC_LAYER_PRE_NORM:      "a.blk.{bid}.layer_pre_norm",
+    MODEL_TENSOR.A_ENC_ATTN_Q:              "a.blk.{bid}.attn_q",
+    MODEL_TENSOR.A_ENC_ATTN_K:              "a.blk.{bid}.attn_k",
+    MODEL_TENSOR.A_ENC_ATTN_V:              "a.blk.{bid}.attn_v",
+    MODEL_TENSOR.A_ENC_ATTN_POST_NORM:      "a.blk.{bid}.attn_post_norm",
+    MODEL_TENSOR.A_ENC_ATTN_PRE_NORM:       "a.blk.{bid}.attn_pre_norm",
+    MODEL_TENSOR.A_ENC_ATTN_K_REL:          "a.blk.{bid}.attn_k_rel",
+    MODEL_TENSOR.A_ENC_PER_DIM_SCALE:       "a.blk.{bid}.per_dim_scale",
+    MODEL_TENSOR.A_ENC_INPUT_NORM:          "a.blk.{bid}.ln1",
+    MODEL_TENSOR.A_ENC_OUTPUT:              "a.blk.{bid}.attn_out",
+    MODEL_TENSOR.A_ENC_OUTPUT_NORM:         "a.blk.{bid}.ln2",
+    MODEL_TENSOR.A_ENC_FFN_NORM:            "a.blk.{bid}.ffn_norm",
+    MODEL_TENSOR.A_ENC_FFN_POST_NORM:       "a.blk.{bid}.ffn_post_norm",
+    MODEL_TENSOR.A_ENC_FFN_SCALE:           "a.blk.{bid}.ffn_scale",
+    MODEL_TENSOR.A_ENC_FFN_UP:              "a.blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_ENC_FFN_GATE:            "a.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.A_ENC_FFN_DOWN:            "a.blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_ENC_FFN_NORM_1:          "a.blk.{bid}.ffn_norm_1",
+    MODEL_TENSOR.A_ENC_FFN_POST_NORM_1:     "a.blk.{bid}.ffn_post_norm_1",
+    MODEL_TENSOR.A_ENC_FFN_SCALE_1:         "a.blk.{bid}.ffn_scale_1",
+    MODEL_TENSOR.A_ENC_FFN_UP_1:            "a.blk.{bid}.ffn_up_1",
+    MODEL_TENSOR.A_ENC_FFN_GATE_1:          "a.blk.{bid}.ffn_gate_1",
+    MODEL_TENSOR.A_ENC_FFN_DOWN_1:          "a.blk.{bid}.ffn_down_1",
+    MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV:     "a.downsample.conv",
+    MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM:     "a.downsample.norm",
+    MODEL_TENSOR.A_ENC_RVQ_CODEBOOK:        "a.rvq.codebook",
+    MODEL_TENSOR.A_ENC_CONV_RES2:           "a.blk.{bid}.res2.{xid}",
+    MODEL_TENSOR.A_ENC_SE_CONV1:            "a.blk.{bid}.se_conv1",
+    MODEL_TENSOR.A_ENC_SE_CONV2:            "a.blk.{bid}.se_conv2",
+    MODEL_TENSOR.A_ENC_ASP_ATTN:            "a.asp_attn",
+    MODEL_TENSOR.A_ENC_ASP_TDNN:            "a.asp_tdnn",
+    MODEL_TENSOR.A_GEN_CODE_PROJ_IN:        "a.gen.code.proj_in",
+    MODEL_TENSOR.A_GEN_CODE_EMBD:           "a.gen.code.embd",
+    MODEL_TENSOR.A_GEN_CODE_HEAD:           "a.gen.code.head",
+    MODEL_TENSOR.A_GEN_CODE_OUT_EMBD:       "a.gen.code.out_embd",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_NORM:      "a.gen.code.blk.{bid}.ln1", # reuses the generic clip.cpp block loader (TN_LN_1)
+    MODEL_TENSOR.A_GEN_CODE_ATTN_Q:         "a.gen.code.blk.{bid}.attn_q",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM:    "a.gen.code.blk.{bid}.attn_q_norm",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_K:         "a.gen.code.blk.{bid}.attn_k",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM:    "a.gen.code.blk.{bid}.attn_k_norm",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_V:         "a.gen.code.blk.{bid}.attn_v",
+    MODEL_TENSOR.A_GEN_CODE_ATTN_OUT:       "a.gen.code.blk.{bid}.attn_out",
+    MODEL_TENSOR.A_GEN_CODE_FFN_NORM:       "a.gen.code.blk.{bid}.ln2", # reuses the generic clip.cpp block loader (TN_LN_2)
+    MODEL_TENSOR.A_GEN_CODE_FFN_GATE:       "a.gen.code.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.A_GEN_CODE_FFN_UP:         "a.gen.code.blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_GEN_CODE_FFN_DOWN:       "a.gen.code.blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM:    "a.gen.code.output_norm",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN:  "a.gen.wav.quant.first.in_proj",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT: "a.gen.wav.quant.first.out_proj",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB:  "a.gen.wav.quant.first.codebook",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN:   "a.gen.wav.quant.rest.in_proj",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT:  "a.gen.wav.quant.rest.out_proj",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB:   "a.gen.wav.quant.rest.codebook",
+    MODEL_TENSOR.A_GEN_WAV_PRE_CONV:        "a.gen.wav.pre_conv",
+    MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ:     "a.gen.wav.tfm.in_proj",
+    MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ:    "a.gen.wav.tfm.out_proj",
+    MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM: "a.gen.wav.tfm.output_norm",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM:   "a.gen.wav.tfm.blk.{bid}.ln1",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q:      "a.gen.wav.tfm.blk.{bid}.attn_q",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K:      "a.gen.wav.tfm.blk.{bid}.attn_k",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V:      "a.gen.wav.tfm.blk.{bid}.attn_v",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT:    "a.gen.wav.tfm.blk.{bid}.attn_out",
+    MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE:  "a.gen.wav.tfm.blk.{bid}.ls1",
+    MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM:    "a.gen.wav.tfm.blk.{bid}.ln2",
+    MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE:    "a.gen.wav.tfm.blk.{bid}.ffn_gate",
+    MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP:      "a.gen.wav.tfm.blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN:    "a.gen.wav.tfm.blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE:   "a.gen.wav.tfm.blk.{bid}.ls2",
+    MODEL_TENSOR.A_GEN_WAV_UP_CONV:         "a.gen.wav.up.blk.{bid}.conv",
+    MODEL_TENSOR.A_GEN_WAV_UP_DWCONV:       "a.gen.wav.up.blk.{bid}.dwconv",
+    MODEL_TENSOR.A_GEN_WAV_UP_NORM:         "a.gen.wav.up.blk.{bid}.norm",
+    MODEL_TENSOR.A_GEN_WAV_UP_PW1:          "a.gen.wav.up.blk.{bid}.pw1",
+    MODEL_TENSOR.A_GEN_WAV_UP_PW2:          "a.gen.wav.up.blk.{bid}.pw2",
+    MODEL_TENSOR.A_GEN_WAV_UP_GAMMA:        "a.gen.wav.up.blk.{bid}.gamma",
+    MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY:       "a.gen.wav.dac.entry",
+    MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE:    "a.gen.wav.dac.blk.{bid}.snake",
+    MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV:     "a.gen.wav.dac.blk.{bid}.conv",
+    MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1:    "a.gen.wav.dac.blk.{bid}.res.{xid}.act1",
+    MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1:   "a.gen.wav.dac.blk.{bid}.res.{xid}.conv1",
+    MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2:    "a.gen.wav.dac.blk.{bid}.res.{xid}.act2",
+    MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2:   "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2",
+    MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE:  "a.gen.wav.dac.post_snake",
+    MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV:   "a.gen.wav.dac.post_conv",
+    MODEL_TENSOR.A_ENC_SEANET_CONV_IN:      "a.seanet.conv_in",
+    MODEL_TENSOR.A_ENC_SEANET_CONV_OUT:     "a.seanet.conv_out",
+    MODEL_TENSOR.A_ENC_SEANET_RES_CONV1:    "a.seanet.blk.{bid}.res_conv1",
+    MODEL_TENSOR.A_ENC_SEANET_RES_CONV2:    "a.seanet.blk.{bid}.res_conv2",
+    MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV:   "a.seanet.blk.{bid}.scale_conv",
+    MODEL_TENSOR.A_ENC_ATTN_SCALE:          "a.blk.{bid}.ls1",
+    MODEL_TENSOR.A_ENC_FFN_SCALE_LS:        "a.blk.{bid}.ls2",
+    MODEL_TENSOR.A_ENC_SPEAKER_PROJ:        "a.speaker_proj",
+    MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ:     "a.gen.flow.input_proj",
+    MODEL_TENSOR.A_GEN_FLOW_COND_EMBD:      "a.gen.flow.cond_embd",
+    MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS:     "a.gen.flow.time.{bid}.freqs",
+    MODEL_TENSOR.A_GEN_FLOW_TIME_UP:        "a.gen.flow.time.{bid}.up",
+    MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN:      "a.gen.flow.time.{bid}.down",
+    MODEL_TENSOR.A_GEN_FLOW_TIME_NORM:      "a.gen.flow.time.{bid}.norm",
+    MODEL_TENSOR.A_GEN_FLOW_BLK_NORM:       "a.gen.flow.blk.{bid}.norm",
+    MODEL_TENSOR.A_GEN_FLOW_BLK_UP:         "a.gen.flow.blk.{bid}.up",
+    MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN:       "a.gen.flow.blk.{bid}.down",
+    MODEL_TENSOR.A_GEN_FLOW_BLK_ADA:        "a.gen.flow.blk.{bid}.ada",
+    MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA:      "a.gen.flow.final.ada",
+    MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ:     "a.gen.flow.final.proj",
+    MODEL_TENSOR.A_GEN_OUT_EOS:             "a.gen.out_eos",
+    MODEL_TENSOR.A_GEN_INPUT_LINEAR:        "a.gen.input_linear",
+    MODEL_TENSOR.A_GEN_EMB_MEAN:            "a.gen.emb_mean",
+    MODEL_TENSOR.A_GEN_EMB_STD:             "a.gen.emb_std",
+    MODEL_TENSOR.A_GEN_WAV_QUANT_OUT:       "a.gen.wav.quant_out",
+    MODEL_TENSOR.A_GEN_WAV_UPSAMPLE:        "a.gen.wav.upsample",
+    MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN:  "a.gen.wav.seanet.conv_in",
+    MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out",
+    MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1",
+    MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2",
+    MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv",
+    MODEL_TENSOR.A_MMPROJ:                  "mm.a.mlp.{bid}",
+    MODEL_TENSOR.A_MMPROJ_FC:               "mm.a.fc",
+    MODEL_TENSOR.A_MM_NORM_PRE:             "mm.a.norm_pre",
+    MODEL_TENSOR.A_MM_NORM_MID:             "mm.a.norm_mid",
+    MODEL_TENSOR.A_MM_INP_PROJ:             "mm.a.input_projection",      # gemma3n
+    MODEL_TENSOR.A_MM_SOFT_EMB_NORM:        "mm.a.soft_emb_norm",         # gemma3n
+    MODEL_TENSOR.A_MM_EMBEDDING:            "mm.a.embedding",             # gemma3n
+    MODEL_TENSOR.A_MM_HARD_EMB_NORM:        "mm.a.hard_emb_norm",         # gemma3n
+    MODEL_TENSOR.A_MM_CODE_EMBD:            "mm.a.code_embd",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_Q:         "mm.a.local_blk.{bid}.attn_q",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_K:         "mm.a.local_blk.{bid}.attn_k",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_V:         "mm.a.local_blk.{bid}.attn_v",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT:       "mm.a.local_blk.{bid}.attn_out",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_GATE:       "mm.a.local_blk.{bid}.ffn_gate",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_UP:         "mm.a.local_blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN:       "mm.a.local_blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_MM_LOCAL_LN1:            "mm.a.local_blk.{bid}.ln1",
+    MODEL_TENSOR.A_MM_LOCAL_LN2:            "mm.a.local_blk.{bid}.ln2",
+    MODEL_TENSOR.A_MM_LOCAL_NORM:           "mm.a.local_norm",
+    MODEL_TENSOR.A_PER_DIM_K_SCALE:         "a.blk.{bid}.per_dim_k_scale", # gemma4
+    MODEL_TENSOR.A_PER_DIM_SCALE:           "a.blk.{bid}.per_dim_scale",   # gemma4
+    # lfm2 audio
+    MODEL_TENSOR.A_ENC_NORM_CONV:           "a.blk.{bid}.norm_conv",
+    MODEL_TENSOR.A_ENC_LINEAR_POS:          "a.blk.{bid}.linear_pos",
+    MODEL_TENSOR.A_ENC_POS_BIAS_U:          "a.blk.{bid}.pos_bias_u",
+    MODEL_TENSOR.A_ENC_POS_BIAS_V:          "a.blk.{bid}.pos_bias_v",
+    MODEL_TENSOR.A_ENC_OUT:                 "a.pre_encode.out",
+    MODEL_TENSOR.A_ENC_CONV_DW:             "a.blk.{bid}.conv_dw",
+    MODEL_TENSOR.A_ENC_CONV_NORM:           "a.blk.{bid}.conv_norm",
+    MODEL_TENSOR.A_ENC_CONV_PW1:            "a.blk.{bid}.conv_pw1",
+    MODEL_TENSOR.A_ENC_CONV_PW2:            "a.blk.{bid}.conv_pw2",
+    MODEL_TENSOR.A_ENC_CONV_NORM_MEAN:      "a.blk.{bid}.conv_norm_mean",
+    MODEL_TENSOR.A_ENC_CONV_NORM_VAR:       "a.blk.{bid}.conv_norm_var",
+    MODEL_TENSOR.A_ENC_MEL_FILTERS:         "a.mel_filters",
+    MODEL_TENSOR.A_ENC_WINDOW:              "a.window",
+    MODEL_TENSOR.A_CTC_OUT:                 "a.enc_ctc_out",
+    MODEL_TENSOR.A_CTC_OUT_MID:             "a.enc_ctc_out_mid",
+    MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB:    "a.blk.{bid}.attn_rel_pos_emb",
+    # qformer projector
+    MODEL_TENSOR.A_QF_PROJ_QUERY:           "a.proj_query",
+    MODEL_TENSOR.A_QF_PROJ_NORM:            "a.proj_norm",
+    MODEL_TENSOR.A_QF_PROJ_LINEAR:          "a.proj_linear",
+    MODEL_TENSOR.A_QF_SELF_ATTN_Q:          "a.proj_blk.{bid}.self_attn_q",
+    MODEL_TENSOR.A_QF_SELF_ATTN_K:          "a.proj_blk.{bid}.self_attn_k",
+    MODEL_TENSOR.A_QF_SELF_ATTN_V:          "a.proj_blk.{bid}.self_attn_v",
+    MODEL_TENSOR.A_QF_SELF_ATTN_O:          "a.proj_blk.{bid}.self_attn_out",
+    MODEL_TENSOR.A_QF_SELF_ATTN_NORM:       "a.proj_blk.{bid}.self_attn_norm",
+    MODEL_TENSOR.A_QF_CROSS_ATTN_Q:         "a.proj_blk.{bid}.cross_attn_q",
+    MODEL_TENSOR.A_QF_CROSS_ATTN_K:         "a.proj_blk.{bid}.cross_attn_k",
+    MODEL_TENSOR.A_QF_CROSS_ATTN_V:         "a.proj_blk.{bid}.cross_attn_v",
+    MODEL_TENSOR.A_QF_CROSS_ATTN_O:         "a.proj_blk.{bid}.cross_attn_out",
+    MODEL_TENSOR.A_QF_CROSS_ATTN_NORM:      "a.proj_blk.{bid}.cross_attn_norm",
+    MODEL_TENSOR.A_QF_FFN_UP:               "a.proj_blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_QF_FFN_DOWN:             "a.proj_blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_QF_FFN_NORM:             "a.proj_blk.{bid}.ffn_norm",
+    # NextN/MTP
+    MODEL_TENSOR.NEXTN_PROJ_PRE:            "nextn.pre_projection",
+    MODEL_TENSOR.NEXTN_PROJ_POST:           "nextn.post_projection",
+    MODEL_TENSOR.NEXTN_EH_PROJ:             "blk.{bid}.nextn.eh_proj",
+    MODEL_TENSOR.NEXTN_EMBED_TOKENS:        "blk.{bid}.nextn.embed_tokens",
+    MODEL_TENSOR.NEXTN_ENORM:               "blk.{bid}.nextn.enorm",
+    MODEL_TENSOR.NEXTN_HNORM:               "blk.{bid}.nextn.hnorm",
+    MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD:    "blk.{bid}.nextn.shared_head_head",
+    MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM:    "blk.{bid}.nextn.shared_head_norm",
+    MODEL_TENSOR.NEXTN_HC_HEAD_NORM:        "blk.{bid}.nextn.hc_head_norm",
+    MODEL_TENSOR.NEXTN_HC_HEAD_DOWN:        "blk.{bid}.nextn.hc_head_down",
+    MODEL_TENSOR.NEXTN_HC_HEAD_UP:          "blk.{bid}.nextn.hc_head_up",
+    MODEL_TENSOR.FC:                        "fc",
+    MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
+    MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
+    MODEL_TENSOR.DSPARK_CONF_PROJ:          "conf_proj",
+    MODEL_TENSOR.DFLASH_ATTN_CONV_BASE:     "blk.{bid}.attn_conv_base",
+    MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ:     "blk.{bid}.attn_conv_proj",
+    MODEL_TENSOR.DFLASH_FFN_CONV_BASE:      "blk.{bid}.ffn_conv_base",
+    MODEL_TENSOR.DFLASH_FFN_CONV_PROJ:      "blk.{bid}.ffn_conv_proj",
+    MODEL_TENSOR.DFLASH_SELECTOR_PREV:      "selector_predecessor",
+    MODEL_TENSOR.DFLASH_SELECTOR_NEXT:      "selector_successor",
+    MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN:    "selector_hidden",
+    MODEL_TENSOR.D2T:                       "d2t",
+    MODEL_TENSOR.ATTN_V_GATE:               "blk.{bid}.attn_v_gate",
+    MODEL_TENSOR.ATTN_V_EXP:                "blk.{bid}.attn_v_exps",
+}
+
+MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
+    MODEL_ARCH.MMPROJ: [
+        MODEL_TENSOR.V_MMPROJ,
+        MODEL_TENSOR.V_MMPROJ_FC,
+        MODEL_TENSOR.V_MMPROJ_MLP,
+        MODEL_TENSOR.V_MMPROJ_PEG,
+        MODEL_TENSOR.V_ENC_EMBD_CLS,
+        MODEL_TENSOR.V_ENC_EMBD_PATCH,
+        MODEL_TENSOR.V_ENC_EMBD_NORM,
+        MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM,
+        MODEL_TENSOR.V_ENC_EMBD_POS,
+        MODEL_TENSOR.V_ENC_EMBD_IMGNL,
+        MODEL_TENSOR.V_ENC_EMBD_VSEP,
+        MODEL_TENSOR.V_TOK_EMBD_SEP,
+        MODEL_TENSOR.V_ENC_INPUT_NORM,
+        MODEL_TENSOR.V_ENC_ATTN_QKV,
+        MODEL_TENSOR.V_ENC_ATTN_Q,
+        MODEL_TENSOR.V_ENC_ATTN_Q_NORM,
+        MODEL_TENSOR.V_ENC_ATTN_K,
+        MODEL_TENSOR.V_ENC_ATTN_K_NORM,
+        MODEL_TENSOR.V_ENC_ATTN_V,
+        MODEL_TENSOR.V_ENC_ATTN_O,
+        MODEL_TENSOR.V_ENC_ATTN_O_NORM,
+        MODEL_TENSOR.V_ENC_ATTN_SINKS,
+        MODEL_TENSOR.V_ENC_POST_ATTN_NORM,
+        MODEL_TENSOR.V_ENC_FFN_UP,
+        MODEL_TENSOR.V_ENC_FFN_GATE,
+        MODEL_TENSOR.V_ENC_FFN_DOWN,
+        MODEL_TENSOR.V_ENC_FFN_GATE_INP,
+        MODEL_TENSOR.V_ENC_FFN_GATE_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_UP_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B,
+        MODEL_TENSOR.V_ENC_ATTN_POST_NORM,
+        MODEL_TENSOR.V_ENC_FFN_POST_NORM,
+        MODEL_TENSOR.V_LAYER_SCALE_1,
+        MODEL_TENSOR.V_LAYER_SCALE_2,
+        MODEL_TENSOR.V_LAYER_OUT_SCALE,
+        MODEL_TENSOR.V_PRE_NORM,
+        MODEL_TENSOR.V_POST_NORM,
+        MODEL_TENSOR.V_MM_POST_NORM,
+        MODEL_TENSOR.V_MM_INP_PROJ,
+        MODEL_TENSOR.V_MM_INP_NORM,
+        MODEL_TENSOR.V_MM_SOFT_EMB_NORM,
+        MODEL_TENSOR.V_MM_EMBEDDING,
+        MODEL_TENSOR.V_MM_HARD_EMB_NORM,
+        MODEL_TENSOR.V_ENC_CONV_STEM,
+        MODEL_TENSOR.V_ENC_CONV_STEM_NORM,
+        MODEL_TENSOR.V_ENC_MSFA_EXP,
+        MODEL_TENSOR.V_ENC_MSFA_EXP_NORM,
+        MODEL_TENSOR.V_ENC_MSFA_PROJ,
+        MODEL_TENSOR.V_ENC_MSFA_PROJ_NORM,
+        MODEL_TENSOR.V_ENC_MSFA_NORM,
+        MODEL_TENSOR.V_RESMPL_POS_EMBD_K,
+        MODEL_TENSOR.V_RESMPL_ATTN_Q,
+        MODEL_TENSOR.V_RESMPL_ATTN_K,
+        MODEL_TENSOR.V_RESMPL_ATTN_V,
+        MODEL_TENSOR.V_RESMPL_ATTN_OUT,
+        MODEL_TENSOR.V_RESMPL_KV,
+        MODEL_TENSOR.V_RESMPL_KV_NORM,
+        MODEL_TENSOR.V_RESMPL_POST_NORM,
+        MODEL_TENSOR.V_RESMPL_Q_NORM,
+        MODEL_TENSOR.V_RESMPL_PROJ,
+        MODEL_TENSOR.V_RESMPL_QUERY,
+        MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK,
+        MODEL_TENSOR.V_TOK_EMBD_IMG_START,
+        MODEL_TENSOR.V_TOK_EMBD_IMG_END,
+        MODEL_TENSOR.V_TOK_EMBD_IMG_PAD,
+        MODEL_TENSOR.V_MM_PATCH_MERGER,
+        MODEL_TENSOR.V_MM_MERGER_FC1,
+        MODEL_TENSOR.V_MM_MERGER_FC2,
+        MODEL_TENSOR.V_DS_NORM,
+        MODEL_TENSOR.V_DS_FC1,
+        MODEL_TENSOR.V_DS_FC2,
+        MODEL_TENSOR.V_MERGER_LN1,
+        MODEL_TENSOR.V_MERGER_ATTN_Q,
+        MODEL_TENSOR.V_MERGER_ATTN_K,
+        MODEL_TENSOR.V_MERGER_ATTN_V,
+        MODEL_TENSOR.V_MERGER_ATTN_O,
+        MODEL_TENSOR.V_MERGER_DS_LN,
+        MODEL_TENSOR.V_MERGER_DS_UP,
+        MODEL_TENSOR.V_MERGER_DS_DOWN,
+        MODEL_TENSOR.V_MM_POST_FC_NORM,
+        MODEL_TENSOR.V_MM_UP,
+        MODEL_TENSOR.V_MM_DOWN,
+        MODEL_TENSOR.V_MM_GATE,
+        MODEL_TENSOR.V_TOK_BOI,
+        MODEL_TENSOR.V_TOK_EOI,
+        MODEL_TENSOR.V_MM_PRE_NORM,
+        MODEL_TENSOR.V_TOK_IMG_BEGIN,
+        MODEL_TENSOR.V_TOK_IMG_END,
+        MODEL_TENSOR.V_STD_BIAS,
+        MODEL_TENSOR.V_STD_SCALE,
+        MODEL_TENSOR.V_SAM_POS_EMBD,
+        MODEL_TENSOR.V_SAM_PATCH_EMBD,
+        MODEL_TENSOR.V_SAM_PRE_NORM,
+        MODEL_TENSOR.V_SAM_POST_NORM,
+        MODEL_TENSOR.V_SAM_ATTN_POS_H,
+        MODEL_TENSOR.V_SAM_ATTN_POS_W,
+        MODEL_TENSOR.V_SAM_ATTN_QKV,
+        MODEL_TENSOR.V_SAM_ATTN_OUT,
+        MODEL_TENSOR.V_SAM_MLP_LIN_1,
+        MODEL_TENSOR.V_SAM_MLP_LIN_2,
+        MODEL_TENSOR.V_SAM_NECK,
+        MODEL_TENSOR.V_SAM_NET_2,
+        MODEL_TENSOR.V_SAM_NET_3,
+        MODEL_TENSOR.V_RESMPL_QUERY_768,
+        MODEL_TENSOR.V_RESMPL_QUERY_1024,
+        MODEL_TENSOR.V_PROJ_NORM,
+        MODEL_TENSOR.V_QF_PROJ_QUERY,
+        MODEL_TENSOR.V_QF_PROJ_NORM,
+        MODEL_TENSOR.V_QF_PROJ_LINEAR,
+        MODEL_TENSOR.V_QF_SELF_ATTN_Q,
+        MODEL_TENSOR.V_QF_SELF_ATTN_K,
+        MODEL_TENSOR.V_QF_SELF_ATTN_V,
+        MODEL_TENSOR.V_QF_SELF_ATTN_O,
+        MODEL_TENSOR.V_QF_SELF_ATTN_NORM,
+        MODEL_TENSOR.V_QF_CROSS_ATTN_Q,
+        MODEL_TENSOR.V_QF_CROSS_ATTN_K,
+        MODEL_TENSOR.V_QF_CROSS_ATTN_V,
+        MODEL_TENSOR.V_QF_CROSS_ATTN_O,
+        MODEL_TENSOR.V_QF_CROSS_ATTN_NORM,
+        MODEL_TENSOR.V_QF_FFN_UP,
+        MODEL_TENSOR.V_QF_FFN_DOWN,
+        MODEL_TENSOR.V_QF_FFN_NORM,
+        MODEL_TENSOR.V_QF_PROJ_NORM,
+        MODEL_TENSOR.V_MULTI_PROJ_IMG_POS,
+        MODEL_TENSOR.V_MULTI_PROJ_QUERY,
+        MODEL_TENSOR.V_MULTI_PROJ_LINEAR,
+        MODEL_TENSOR.V_MULTI_PROJ_NORM,
+        MODEL_TENSOR.V_MULTI_PROJ_POST_NORM,
+        # audio
+        MODEL_TENSOR.A_ENC_EMBD_POS,
+        MODEL_TENSOR.A_ENC_EMBD_NORM,
+        MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS,
+        MODEL_TENSOR.A_ENC_INP_PROJ,
+        MODEL_TENSOR.A_ENC_CONV1D,
+        MODEL_TENSOR.A_ENC_CONV2D,
+        MODEL_TENSOR.A_ENC_CONV_OUT,
+        MODEL_TENSOR.A_ENC_CONV1D_NORM,
+        MODEL_TENSOR.A_PRE_NORM,
+        MODEL_TENSOR.A_POST_NORM,
+        MODEL_TENSOR.A_ENC_LAYER_PRE_NORM,
+        MODEL_TENSOR.A_ENC_ATTN_Q,
+        MODEL_TENSOR.A_ENC_ATTN_K,
+        MODEL_TENSOR.A_ENC_ATTN_V,
+        MODEL_TENSOR.A_ENC_ATTN_POST_NORM,
+        MODEL_TENSOR.A_ENC_ATTN_PRE_NORM,
+        MODEL_TENSOR.A_ENC_ATTN_K_REL,
+        MODEL_TENSOR.A_ENC_PER_DIM_SCALE,
+        MODEL_TENSOR.A_ENC_INPUT_NORM,
+        MODEL_TENSOR.A_ENC_OUTPUT,
+        MODEL_TENSOR.A_ENC_OUTPUT_NORM,
+        MODEL_TENSOR.A_ENC_FFN_NORM,
+        MODEL_TENSOR.A_ENC_FFN_POST_NORM,
+        MODEL_TENSOR.A_ENC_FFN_SCALE,
+        MODEL_TENSOR.A_ENC_FFN_UP,
+        MODEL_TENSOR.A_ENC_FFN_GATE,
+        MODEL_TENSOR.A_ENC_FFN_DOWN,
+        MODEL_TENSOR.A_ENC_FFN_NORM_1,
+        MODEL_TENSOR.A_ENC_FFN_POST_NORM_1,
+        MODEL_TENSOR.A_ENC_FFN_SCALE_1,
+        MODEL_TENSOR.A_ENC_FFN_UP_1,
+        MODEL_TENSOR.A_ENC_FFN_GATE_1,
+        MODEL_TENSOR.A_ENC_FFN_DOWN_1,
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV,
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM,
+        MODEL_TENSOR.A_ENC_RVQ_CODEBOOK,
+        MODEL_TENSOR.A_MMPROJ,
+        MODEL_TENSOR.A_MMPROJ_FC,
+        MODEL_TENSOR.A_MM_NORM_PRE,
+        MODEL_TENSOR.A_MM_NORM_MID,
+        MODEL_TENSOR.A_MM_CODE_EMBD,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_Q,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_K,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_V,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_GATE,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_UP,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN,
+        MODEL_TENSOR.A_MM_LOCAL_LN1,
+        MODEL_TENSOR.A_MM_LOCAL_LN2,
+        MODEL_TENSOR.A_MM_LOCAL_NORM,
+        MODEL_TENSOR.A_ENC_NORM_CONV,
+        MODEL_TENSOR.A_ENC_LINEAR_POS,
+        MODEL_TENSOR.A_ENC_POS_BIAS_U,
+        MODEL_TENSOR.A_ENC_POS_BIAS_V,
+        MODEL_TENSOR.A_ENC_OUT,
+        MODEL_TENSOR.A_ENC_CONV_DW,
+        MODEL_TENSOR.A_ENC_CONV_NORM,
+        MODEL_TENSOR.A_ENC_CONV_PW1,
+        MODEL_TENSOR.A_ENC_CONV_PW2,
+        MODEL_TENSOR.A_ENC_CONV_RES2,
+        MODEL_TENSOR.A_ENC_SE_CONV1,
+        MODEL_TENSOR.A_ENC_SE_CONV2,
+        MODEL_TENSOR.A_ENC_ASP_ATTN,
+        MODEL_TENSOR.A_ENC_ASP_TDNN,
+        MODEL_TENSOR.A_GEN_CODE_PROJ_IN,
+        MODEL_TENSOR.A_GEN_CODE_EMBD,
+        MODEL_TENSOR.A_GEN_CODE_HEAD,
+        MODEL_TENSOR.A_GEN_CODE_OUT_EMBD,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_NORM,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_Q,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_K,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_V,
+        MODEL_TENSOR.A_GEN_CODE_ATTN_OUT,
+        MODEL_TENSOR.A_GEN_CODE_FFN_NORM,
+        MODEL_TENSOR.A_GEN_CODE_FFN_GATE,
+        MODEL_TENSOR.A_GEN_CODE_FFN_UP,
+        MODEL_TENSOR.A_GEN_CODE_FFN_DOWN,
+        MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB,
+        MODEL_TENSOR.A_GEN_WAV_PRE_CONV,
+        MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ,
+        MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ,
+        MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT,
+        MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE,
+        MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM,
+        MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE,
+        MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP,
+        MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN,
+        MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE,
+        MODEL_TENSOR.A_GEN_WAV_UP_CONV,
+        MODEL_TENSOR.A_GEN_WAV_UP_DWCONV,
+        MODEL_TENSOR.A_GEN_WAV_UP_NORM,
+        MODEL_TENSOR.A_GEN_WAV_UP_PW1,
+        MODEL_TENSOR.A_GEN_WAV_UP_PW2,
+        MODEL_TENSOR.A_GEN_WAV_UP_GAMMA,
+        MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY,
+        MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE,
+        MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV,
+        MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1,
+        MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1,
+        MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2,
+        MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2,
+        MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE,
+        MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV,
+        MODEL_TENSOR.A_ENC_SEANET_CONV_IN,
+        MODEL_TENSOR.A_ENC_SEANET_CONV_OUT,
+        MODEL_TENSOR.A_ENC_SEANET_RES_CONV1,
+        MODEL_TENSOR.A_ENC_SEANET_RES_CONV2,
+        MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV,
+        MODEL_TENSOR.A_ENC_ATTN_SCALE,
+        MODEL_TENSOR.A_ENC_FFN_SCALE_LS,
+        MODEL_TENSOR.A_ENC_SPEAKER_PROJ,
+        MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ,
+        MODEL_TENSOR.A_GEN_FLOW_COND_EMBD,
+        MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS,
+        MODEL_TENSOR.A_GEN_FLOW_TIME_UP,
+        MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN,
+        MODEL_TENSOR.A_GEN_FLOW_TIME_NORM,
+        MODEL_TENSOR.A_GEN_FLOW_BLK_NORM,
+        MODEL_TENSOR.A_GEN_FLOW_BLK_UP,
+        MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN,
+        MODEL_TENSOR.A_GEN_FLOW_BLK_ADA,
+        MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA,
+        MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ,
+        MODEL_TENSOR.A_GEN_OUT_EOS,
+        MODEL_TENSOR.A_GEN_INPUT_LINEAR,
+        MODEL_TENSOR.A_GEN_EMB_MEAN,
+        MODEL_TENSOR.A_GEN_EMB_STD,
+        MODEL_TENSOR.A_GEN_WAV_QUANT_OUT,
+        MODEL_TENSOR.A_GEN_WAV_UPSAMPLE,
+        MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN,
+        MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT,
+        MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1,
+        MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2,
+        MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV,
+        MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
+        MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
+        MODEL_TENSOR.A_ENC_MEL_FILTERS,
+        MODEL_TENSOR.A_ENC_WINDOW,
+        MODEL_TENSOR.A_MM_INP_PROJ,
+        MODEL_TENSOR.A_MM_SOFT_EMB_NORM,
+        MODEL_TENSOR.A_MM_EMBEDDING,
+        MODEL_TENSOR.A_MM_HARD_EMB_NORM,
+        MODEL_TENSOR.A_PER_DIM_K_SCALE,
+        MODEL_TENSOR.A_PER_DIM_SCALE,
+        MODEL_TENSOR.A_CTC_OUT,
+        MODEL_TENSOR.A_CTC_OUT_MID,
+        MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB,
+        # qformer projector
+        MODEL_TENSOR.A_QF_PROJ_QUERY,
+        MODEL_TENSOR.A_QF_PROJ_NORM,
+        MODEL_TENSOR.A_QF_PROJ_LINEAR,
+        MODEL_TENSOR.A_QF_SELF_ATTN_Q,
+        MODEL_TENSOR.A_QF_SELF_ATTN_K,
+        MODEL_TENSOR.A_QF_SELF_ATTN_V,
+        MODEL_TENSOR.A_QF_SELF_ATTN_O,
+        MODEL_TENSOR.A_QF_SELF_ATTN_NORM,
+        MODEL_TENSOR.A_QF_CROSS_ATTN_Q,
+        MODEL_TENSOR.A_QF_CROSS_ATTN_K,
+        MODEL_TENSOR.A_QF_CROSS_ATTN_V,
+        MODEL_TENSOR.A_QF_CROSS_ATTN_O,
+        MODEL_TENSOR.A_QF_CROSS_ATTN_NORM,
+        MODEL_TENSOR.A_QF_FFN_UP,
+        MODEL_TENSOR.A_QF_FFN_DOWN,
+        MODEL_TENSOR.A_QF_FFN_NORM,
+    ],
+    MODEL_ARCH.LLAMA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.LLAMA4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.DECI: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.GROK: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.GPTNEOX: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.FALCON: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.BAICHUAN: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.STARCODER: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.BERT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+        MODEL_TENSOR.CLS,
+        MODEL_TENSOR.CLS_OUT,
+    ],
+    MODEL_ARCH.MODERN_BERT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.CLS,
+        MODEL_TENSOR.CLS_OUT,
+        MODEL_TENSOR.CLS_NORM,
+    ],
+    MODEL_ARCH.NOMIC_BERT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.NOMIC_BERT_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.NEO_BERT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+        MODEL_TENSOR.CLS,
+        MODEL_TENSOR.CLS_OUT,
+    ],
+    MODEL_ARCH.JINA_BERT_V2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+        MODEL_TENSOR.CLS,
+    ],
+    MODEL_ARCH.JINA_BERT_V3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.EUROBERT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_DOWN,
+    ],
+    MODEL_ARCH.MPT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_ACT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.POS_EMBD,
+    ],
+    MODEL_ARCH.GPTJ: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.REFACT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.BLOOM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.STABLELM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+    ],
+    MODEL_ARCH.QWEN: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.DREAM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.LLADA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN2VL: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN2MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.QWEN3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN3MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.QWEN3NEXT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_BETA_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.HRM_TEXT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.HRM_Z_L_INIT,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN3VL: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN3VLMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.QWEN35: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.CLS_OUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.CLEF: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        # decision head
+        MODEL_TENSOR.TOKEN_TYPES,
+        MODEL_TENSOR.DEC_ATTN_NORM,
+        MODEL_TENSOR.DEC_ATTN_Q,
+        MODEL_TENSOR.DEC_ATTN_K,
+        MODEL_TENSOR.DEC_ATTN_V,
+        MODEL_TENSOR.DEC_ATTN_OUT,
+        MODEL_TENSOR.DEC_CROSS_ATTN_NORM,
+        MODEL_TENSOR.DEC_CROSS_ATTN_NORM_KV,
+        MODEL_TENSOR.DEC_CROSS_ATTN_Q,
+        MODEL_TENSOR.DEC_CROSS_ATTN_K,
+        MODEL_TENSOR.DEC_CROSS_ATTN_V,
+        MODEL_TENSOR.DEC_CROSS_ATTN_OUT,
+        MODEL_TENSOR.DEC_FFN_NORM,
+        MODEL_TENSOR.DEC_FFN_DOWN,
+        MODEL_TENSOR.DEC_FFN_UP,
+        MODEL_TENSOR.DECISION_HIDDEN_NORM,
+        MODEL_TENSOR.DECISION_PROJ_MEMORY,
+        MODEL_TENSOR.DECISION_PROJ_QUESTION,
+        MODEL_TENSOR.DECISION_PROJ_OPTION_QUESTION,
+        MODEL_TENSOR.DECISION_PROJ_GLOBAL,
+        MODEL_TENSOR.DECISION_PROJ_OPTION_CONTEXT,
+        MODEL_TENSOR.DECISION_PROJ_OPTION_LEXICAL,
+        MODEL_TENSOR.DECISION_OPTION_SUMMARY_NORM,
+        MODEL_TENSOR.DECISION_FIELD_NORM,
+        MODEL_TENSOR.DECISION_OPTION_NORM,
+        MODEL_TENSOR.DECISION_SCALES,
+        MODEL_TENSOR.DECISION_SCORER,
+        MODEL_TENSOR.DECISION_SCORER_OUT,
+    ],
+    MODEL_ARCH.VLLM_DECISION2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        # decision head
+        MODEL_TENSOR.DECISION_OPTION_NORM,
+        MODEL_TENSOR.DECISION_QUERY_NORM,
+        MODEL_TENSOR.DECISION_KEY,
+        MODEL_TENSOR.DECISION_QUERY,
+        MODEL_TENSOR.DECISION_OPTION_MLP,
+        MODEL_TENSOR.DECISION_QUERY_MLP,
+        MODEL_TENSOR.DECISION_SCORER_OUT,
+    ],
+    MODEL_ARCH.QWEN35MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.QWEN4EXP: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        # no OUTPUT_NORM / ATTN_NORM / ATTN_POST_NORM: hyper-connections replace every layer norm
+        MODEL_TENSOR.HC_HEAD_NORM,
+        MODEL_TENSOR.HC_HEAD_DOWN,
+        MODEL_TENSOR.HC_HEAD_UP,
+        MODEL_TENSOR.HC_ATTN_NORM,
+        MODEL_TENSOR.HC_ATTN_DOWN,
+        MODEL_TENSOR.HC_ATTN_UP,
+        MODEL_TENSOR.HC_ATTN_INJECT,
+        MODEL_TENSOR.HC_FFN_NORM,
+        MODEL_TENSOR.HC_FFN_DOWN,
+        MODEL_TENSOR.HC_FFN_UP,
+        MODEL_TENSOR.HC_FFN_INJECT,
+        # full attention layers: ATTN_Q holds [q|gate] interleaved per head
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.INDEXER_Q_PROJ,
+        MODEL_TENSOR.INDEXER_K_PROJ,
+        MODEL_TENSOR.INDEXER_Q_NORM,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
+        MODEL_TENSOR.PLE_KEY,
+        MODEL_TENSOR.PLE_VALUE,
+        MODEL_TENSOR.PLE_NORM_KEY,
+        MODEL_TENSOR.PLE_NORM_QUERY,
+        MODEL_TENSOR.PLE_NORM_CONV,
+        MODEL_TENSOR.PLE_CONV1D,
+        # MTP block: [fc_embedding | fc_hidden] as eh_proj, its own hyper-connection mixer as the head
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_HC_HEAD_NORM,
+        MODEL_TENSOR.NEXTN_HC_HEAD_DOWN,
+        MODEL_TENSOR.NEXTN_HC_HEAD_UP,
+    ],
+    MODEL_ARCH.PLAMO: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.PLAMO2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_X,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.SSM_DT_NORM,
+        MODEL_TENSOR.SSM_B_NORM,
+        MODEL_TENSOR.SSM_C_NORM,
+    ],
+    MODEL_ARCH.PLAMO3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_POST_NORM,
+    ],
+    MODEL_ARCH.GPT2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.PHI2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.PHI3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FACTORS_LONG,
+        MODEL_TENSOR.ROPE_FACTORS_SHORT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.PHIMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FACTORS_LONG,
+        MODEL_TENSOR.ROPE_FACTORS_SHORT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.CODESHELL: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.POS_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.ORION: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.INTERNLM2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.MINICPM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ROPE_FACTORS_LONG,
+        MODEL_TENSOR.ROPE_FACTORS_SHORT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.MINICPM3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FACTORS_LONG,
+        MODEL_TENSOR.ROPE_FACTORS_SHORT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GEMMA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_NORM,
+    ],
+    MODEL_ARCH.GEMMA2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+    ],
+    MODEL_ARCH.GEMMA3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+    ],
+    MODEL_ARCH.GEMMA3N: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        # altup / laurel
+        MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
+        MODEL_TENSOR.PER_LAYER_MODEL_PROJ,
+        MODEL_TENSOR.PER_LAYER_INP_GATE,
+        MODEL_TENSOR.PER_LAYER_PROJ,
+        MODEL_TENSOR.PER_LAYER_PROJ_NORM,
+        MODEL_TENSOR.PER_LAYER_POST_NORM,
+        MODEL_TENSOR.ALTUP_PROJ,
+        MODEL_TENSOR.ALTUP_UNEMBD_PROJ,
+        MODEL_TENSOR.ALTUP_CORRECT_COEF,
+        MODEL_TENSOR.ALTUP_CORRECT_SCALE,
+        MODEL_TENSOR.ALTUP_PREDICT_COEF,
+        MODEL_TENSOR.ALTUP_ROUTER,
+        MODEL_TENSOR.ALTUP_ROUTER_NORM,
+        MODEL_TENSOR.LAUREL_L,
+        MODEL_TENSOR.LAUREL_R,
+        MODEL_TENSOR.LAUREL_POST_NORM,
+    ],
+    MODEL_ARCH.GEMMA4: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM_2,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.FFN_POST_NORM_1,
+        MODEL_TENSOR.FFN_POST_NORM_2,
+        MODEL_TENSOR.LAYER_OUT_SCALE,
+        MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
+        MODEL_TENSOR.PER_LAYER_MODEL_PROJ,
+        MODEL_TENSOR.PER_LAYER_INP_GATE,
+        MODEL_TENSOR.PER_LAYER_PROJ,
+        MODEL_TENSOR.PER_LAYER_PROJ_NORM,
+        MODEL_TENSOR.PER_LAYER_POST_NORM,
+    ],
+    MODEL_ARCH.GEMMA4_ASSISTANT: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.MASKED_EMBD_CENTROIDS,
+        MODEL_TENSOR.MASKED_EMBD_ORDERING,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.NEXTN_PROJ_PRE,
+        MODEL_TENSOR.NEXTN_PROJ_POST,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.LAYER_OUT_SCALE,
+    ],
+    MODEL_ARCH.GEMMA_EMBEDDING: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.DENSE_2_OUT,
+        MODEL_TENSOR.DENSE_3_OUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+    ],
+    MODEL_ARCH.GEMMA_EMBEDDING2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.LAYER_OUT_SCALE,
+        MODEL_TENSOR.PER_LAYER_MODEL_PROJ,
+        MODEL_TENSOR.PER_LAYER_INP_GATE,
+        MODEL_TENSOR.PER_LAYER_PROJ,
+        MODEL_TENSOR.PER_LAYER_PROJ_NORM,
+        MODEL_TENSOR.PER_LAYER_POST_NORM,
+    ],
+    MODEL_ARCH.STARCODER2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.RWKV6: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.TIME_MIX_W1,
+        MODEL_TENSOR.TIME_MIX_W2,
+        MODEL_TENSOR.TIME_MIX_LERP_X,
+        MODEL_TENSOR.TIME_MIX_LERP_K,
+        MODEL_TENSOR.TIME_MIX_LERP_V,
+        MODEL_TENSOR.TIME_MIX_LERP_R,
+        MODEL_TENSOR.TIME_MIX_LERP_G,
+        MODEL_TENSOR.TIME_MIX_LERP_W,
+        MODEL_TENSOR.TIME_MIX_LERP_FUSED,
+        MODEL_TENSOR.TIME_MIX_FIRST,
+        MODEL_TENSOR.TIME_MIX_DECAY,
+        MODEL_TENSOR.TIME_MIX_DECAY_W1,
+        MODEL_TENSOR.TIME_MIX_DECAY_W2,
+        MODEL_TENSOR.TIME_MIX_KEY,
+        MODEL_TENSOR.TIME_MIX_VALUE,
+        MODEL_TENSOR.TIME_MIX_RECEPTANCE,
+        MODEL_TENSOR.TIME_MIX_GATE,
+        MODEL_TENSOR.TIME_MIX_LN,
+        MODEL_TENSOR.TIME_MIX_OUTPUT,
+        MODEL_TENSOR.CHANNEL_MIX_LERP_K,
+        MODEL_TENSOR.CHANNEL_MIX_LERP_R,
+        MODEL_TENSOR.CHANNEL_MIX_KEY,
+        MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE,
+        MODEL_TENSOR.CHANNEL_MIX_VALUE,
+    ],
+    MODEL_ARCH.RWKV6QWEN2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.TIME_MIX_W1,
+        MODEL_TENSOR.TIME_MIX_W2,
+        MODEL_TENSOR.TIME_MIX_LERP_X,
+        MODEL_TENSOR.TIME_MIX_LERP_K,
+        MODEL_TENSOR.TIME_MIX_LERP_V,
+        MODEL_TENSOR.TIME_MIX_LERP_R,
+        MODEL_TENSOR.TIME_MIX_LERP_G,
+        MODEL_TENSOR.TIME_MIX_LERP_W,
+        MODEL_TENSOR.TIME_MIX_LERP_FUSED,
+        MODEL_TENSOR.TIME_MIX_FIRST,
+        MODEL_TENSOR.TIME_MIX_DECAY,
+        MODEL_TENSOR.TIME_MIX_DECAY_W1,
+        MODEL_TENSOR.TIME_MIX_DECAY_W2,
+        MODEL_TENSOR.TIME_MIX_KEY,
+        MODEL_TENSOR.TIME_MIX_VALUE,
+        MODEL_TENSOR.TIME_MIX_RECEPTANCE,
+        MODEL_TENSOR.TIME_MIX_GATE,
+        MODEL_TENSOR.TIME_MIX_LN,
+        MODEL_TENSOR.TIME_MIX_OUTPUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.RWKV7: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.TIME_MIX_LERP_FUSED,
+        MODEL_TENSOR.TIME_MIX_W0,
+        MODEL_TENSOR.TIME_MIX_W1,
+        MODEL_TENSOR.TIME_MIX_W2,
+        MODEL_TENSOR.TIME_MIX_A0,
+        MODEL_TENSOR.TIME_MIX_A1,
+        MODEL_TENSOR.TIME_MIX_A2,
+        MODEL_TENSOR.TIME_MIX_V0,
+        MODEL_TENSOR.TIME_MIX_V1,
+        MODEL_TENSOR.TIME_MIX_V2,
+        MODEL_TENSOR.TIME_MIX_G1,
+        MODEL_TENSOR.TIME_MIX_G2,
+        MODEL_TENSOR.TIME_MIX_K_K,
+        MODEL_TENSOR.TIME_MIX_K_A,
+        MODEL_TENSOR.TIME_MIX_R_K,
+        MODEL_TENSOR.TIME_MIX_KEY,
+        MODEL_TENSOR.TIME_MIX_VALUE,
+        MODEL_TENSOR.TIME_MIX_RECEPTANCE,
+        MODEL_TENSOR.TIME_MIX_LN,
+        MODEL_TENSOR.TIME_MIX_OUTPUT,
+        MODEL_TENSOR.CHANNEL_MIX_LERP_K,
+        MODEL_TENSOR.CHANNEL_MIX_KEY,
+        MODEL_TENSOR.CHANNEL_MIX_VALUE,
+    ],
+    MODEL_ARCH.ARWKV7: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.TIME_MIX_LERP_FUSED,
+        MODEL_TENSOR.TIME_MIX_W0,
+        MODEL_TENSOR.TIME_MIX_W1,
+        MODEL_TENSOR.TIME_MIX_W2,
+        MODEL_TENSOR.TIME_MIX_A0,
+        MODEL_TENSOR.TIME_MIX_A1,
+        MODEL_TENSOR.TIME_MIX_A2,
+        MODEL_TENSOR.TIME_MIX_V0,
+        MODEL_TENSOR.TIME_MIX_V1,
+        MODEL_TENSOR.TIME_MIX_V2,
+        MODEL_TENSOR.TIME_MIX_G1,
+        MODEL_TENSOR.TIME_MIX_G2,
+        MODEL_TENSOR.TIME_MIX_K_K,
+        MODEL_TENSOR.TIME_MIX_K_A,
+        MODEL_TENSOR.TIME_MIX_R_K,
+        MODEL_TENSOR.TIME_MIX_KEY,
+        MODEL_TENSOR.TIME_MIX_VALUE,
+        MODEL_TENSOR.TIME_MIX_RECEPTANCE,
+        MODEL_TENSOR.TIME_MIX_LN,
+        MODEL_TENSOR.TIME_MIX_OUTPUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.MAMBA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_X,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_OUT,
+    ],
+    MODEL_ARCH.MAMBA2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_OUT,
+    ],
+    MODEL_ARCH.MAPLE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.JAMBA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_X,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_DT_NORM,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_B_NORM,
+        MODEL_TENSOR.SSM_C_NORM,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.XVERSE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.COMMAND_R: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_Q_NORM,
+    ],
+    MODEL_ARCH.COHERE2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.COHERE2MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.DBRX: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_OUT_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.OLMO: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.OLMO2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.SEED_OSS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+    ],
+    MODEL_ARCH.OLMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+    ],
+    MODEL_ARCH.MUSE_GLIMMER: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+    ],
+    MODEL_ARCH.OPENELM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.ARCTIC: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_NORM_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.DEEPSEEK: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.DEEPSEEK2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP tensors
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.DEEPSEEK2OCR: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.DEEPSEEK32: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.DEEPSEEK4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV,
+        MODEL_TENSOR.ATTN_KV_NORM,
+        MODEL_TENSOR.ATTN_OUT_A,
+        MODEL_TENSOR.ATTN_OUT_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WKV,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WGATE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_APE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_WKV,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_APE,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_TID2EID,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_EXP_PROBS_B_VL,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.ERNIE4_5_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.PLM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_DOWN,
+    ],
+    MODEL_ARCH.CHATGLM : [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GLM4 : [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.GLM4_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP tensors
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.GLM_DSA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.GLM5_NEXT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        # mHC
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        # KDA (linear attention) layers
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.SSM_CONV1D_Q,
+        MODEL_TENSOR.SSM_CONV1D_K,
+        MODEL_TENSOR.SSM_CONV1D_V,
+        MODEL_TENSOR.SSM_F_A,
+        MODEL_TENSOR.SSM_F_B,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_G_A,
+        MODEL_TENSOR.SSM_G_B,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        # MLA (nope) + DSA layers
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        MODEL_TENSOR.INDEXER_KPOOL_GATE,
+        MODEL_TENSOR.INDEXER_KPOOL_APE,
+        # FFN
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.BITNET: [
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.ATTN_SUB_NORM,
+        MODEL_TENSOR.FFN_SUB_NORM,
+    ],
+    MODEL_ARCH.T5: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.DEC_ATTN_NORM,
+        MODEL_TENSOR.DEC_ATTN_Q,
+        MODEL_TENSOR.DEC_ATTN_K,
+        MODEL_TENSOR.DEC_ATTN_V,
+        MODEL_TENSOR.DEC_ATTN_OUT,
+        MODEL_TENSOR.DEC_ATTN_REL_B,
+        MODEL_TENSOR.DEC_CROSS_ATTN_NORM,
+        MODEL_TENSOR.DEC_CROSS_ATTN_Q,
+        MODEL_TENSOR.DEC_CROSS_ATTN_K,
+        MODEL_TENSOR.DEC_CROSS_ATTN_V,
+        MODEL_TENSOR.DEC_CROSS_ATTN_OUT,
+        MODEL_TENSOR.DEC_CROSS_ATTN_REL_B,
+        MODEL_TENSOR.DEC_FFN_NORM,
+        MODEL_TENSOR.DEC_FFN_GATE,
+        MODEL_TENSOR.DEC_FFN_DOWN,
+        MODEL_TENSOR.DEC_FFN_UP,
+        MODEL_TENSOR.DEC_OUTPUT_NORM,
+        MODEL_TENSOR.ENC_ATTN_NORM,
+        MODEL_TENSOR.ENC_ATTN_Q,
+        MODEL_TENSOR.ENC_ATTN_K,
+        MODEL_TENSOR.ENC_ATTN_V,
+        MODEL_TENSOR.ENC_ATTN_OUT,
+        MODEL_TENSOR.ENC_ATTN_REL_B,
+        MODEL_TENSOR.ENC_FFN_NORM,
+        MODEL_TENSOR.ENC_FFN_GATE,
+        MODEL_TENSOR.ENC_FFN_DOWN,
+        MODEL_TENSOR.ENC_FFN_UP,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+    ],
+    MODEL_ARCH.T5ENCODER: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ENC_ATTN_NORM,
+        MODEL_TENSOR.ENC_ATTN_Q,
+        MODEL_TENSOR.ENC_ATTN_K,
+        MODEL_TENSOR.ENC_ATTN_V,
+        MODEL_TENSOR.ENC_ATTN_OUT,
+        MODEL_TENSOR.ENC_ATTN_REL_B,
+        MODEL_TENSOR.ENC_FFN_NORM,
+        MODEL_TENSOR.ENC_FFN_GATE,
+        MODEL_TENSOR.ENC_FFN_DOWN,
+        MODEL_TENSOR.ENC_FFN_UP,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+    ],
+    MODEL_ARCH.JAIS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.JAIS2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.NEMOTRON: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.NEMOTRON_H: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.NEMOTRON_H_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        # experts
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        # expert latent
+        MODEL_TENSOR.MOE_LATENT_DOWN,
+        MODEL_TENSOR.MOE_LATENT_UP,
+        # shared expert
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP (draft head)
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.EXAONE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.EXAONE4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_POST_NORM,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.EXAONE_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.GRANITE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GRANITE_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
+    MODEL_ARCH.GRANITE_HYBRID: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        # MoE
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        # Dense
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GRANITE_SWITCH: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GRANITE_SWA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        # MoE (GraniteMoeSWA)
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        # Shared expert - gate+up kept fused in FFN_UP_SHEXP (LLM_FFN_SWIGLU)
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
+    MODEL_ARCH.CHAMELEON: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.WAVTOKENIZER_DEC: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.CONV1D,
+        MODEL_TENSOR.CONVNEXT_DW,
+        MODEL_TENSOR.CONVNEXT_NORM,
+        MODEL_TENSOR.CONVNEXT_PW1,
+        MODEL_TENSOR.CONVNEXT_PW2,
+        MODEL_TENSOR.CONVNEXT_GAMMA,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.POSNET_CONV1,
+        MODEL_TENSOR.POSNET_CONV2,
+        MODEL_TENSOR.POSNET_NORM,
+        MODEL_TENSOR.POSNET_NORM1,
+        MODEL_TENSOR.POSNET_NORM2,
+        MODEL_TENSOR.POSNET_ATTN_NORM,
+        MODEL_TENSOR.POSNET_ATTN_Q,
+        MODEL_TENSOR.POSNET_ATTN_K,
+        MODEL_TENSOR.POSNET_ATTN_V,
+        MODEL_TENSOR.POSNET_ATTN_OUT,
+    ],
+    MODEL_ARCH.BAILINGMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.BAILINGMOE2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.BAILINGMOE3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.SSM_CONV1D_Q,
+        MODEL_TENSOR.SSM_CONV1D_K,
+        MODEL_TENSOR.SSM_CONV1D_V,
+        MODEL_TENSOR.SSM_F_A,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_G_A,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+    ],
+    MODEL_ARCH.DOTS1: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.DOTS3NOTE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        # NextN/MTP tensors - preserved but unused
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.ARCEE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.AFMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.LAGUNA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
+    MODEL_ARCH.ERNIE4_5: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.PADDLEOCR: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.FALCON_H1: [
+        # Token embedding
+        MODEL_TENSOR.TOKEN_EMBD,
+
+        # Input layernorm
+        MODEL_TENSOR.ATTN_NORM,
+
+        # Attention components
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,         # Query projection
+        MODEL_TENSOR.ATTN_K,         # Key projection
+        MODEL_TENSOR.ATTN_V,         # Value projection
+        MODEL_TENSOR.ATTN_OUT,       # Output projection
+
+        # SSM components (Mamba2 specific)
+        MODEL_TENSOR.SSM_IN,         # Input projection for SSM
+        MODEL_TENSOR.SSM_CONV1D,     # Convolution layer
+        MODEL_TENSOR.SSM_DT,         # Delta time projection
+        MODEL_TENSOR.SSM_A,          # A parameter (log form)
+        MODEL_TENSOR.SSM_D,          # D parameter
+        MODEL_TENSOR.SSM_NORM,       # Normalization in SSM
+        MODEL_TENSOR.SSM_OUT,        # Output projection
+
+        # Pre-feedforward layernorm
+        MODEL_TENSOR.FFN_PRE_NORM,
+
+        # Feed-forward network components
+        MODEL_TENSOR.FFN_GATE,       # Gate projection (SwiGLU)
+        MODEL_TENSOR.FFN_DOWN,       # Down projection
+        MODEL_TENSOR.FFN_UP,         # Up projection
+
+        # Post-feedforward layernorm
+        MODEL_TENSOR.OUTPUT_NORM,    # Final layer norm
+        MODEL_TENSOR.OUTPUT,         # Output projection (lm_head)
+    ],
+    MODEL_ARCH.HUNYUAN_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.HUNYUAN_DENSE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.HUNYUAN_VL: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.HY_V3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        # NextN/MTP tensors (draft head)
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.HY_V4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.SMOLLM3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.GPT_OSS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.LFM2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.SHORTCONV_CONV,
+        MODEL_TENSOR.SHORTCONV_INPROJ,
+        MODEL_TENSOR.SHORTCONV_OUTPROJ,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.ATTN_NORM, # operator_norm
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.DENSE_2_OUT, # LFM2-ColBert-350M
+        MODEL_TENSOR.TOKEN_TYPES, # decision head
+        MODEL_TENSOR.CLS,
+        MODEL_TENSOR.CLS_NORM,
+        MODEL_TENSOR.CLS_OUT,
+    ],
+    MODEL_ARCH.LFM2MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.TOKEN_EMBD_NORM,
+        MODEL_TENSOR.SHORTCONV_CONV,
+        MODEL_TENSOR.SHORTCONV_INPROJ,
+        MODEL_TENSOR.SHORTCONV_OUTPROJ,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.ATTN_NORM, # operator_norm
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.SMALLTHINKER: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.APERTUS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.LLADA_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+    ],
+    MODEL_ARCH.GROVEMOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_CHEXP,
+        MODEL_TENSOR.FFN_DOWN_CHEXP,
+        MODEL_TENSOR.FFN_UP_CHEXP,
+    ],
+    MODEL_ARCH.MINIMAX01: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.MINIMAXM2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.MINIMAXM3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.INDEXER_Q_PROJ,
+        MODEL_TENSOR.INDEXER_K_PROJ,
+        MODEL_TENSOR.INDEXER_Q_NORM,
+        MODEL_TENSOR.INDEXER_K_NORM,
+    ],
+    MODEL_ARCH.COGVLM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.VISEXP_ATTN_QKV,
+        MODEL_TENSOR.VISEXP_ATTN_OUT,
+        MODEL_TENSOR.VISEXP_GATE,
+        MODEL_TENSOR.VISEXP_UP,
+        MODEL_TENSOR.VISEXP_DOWN,
+    ],
+    MODEL_ARCH.RND1: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.PANGU_EMBED: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.MISTRAL3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.EAGLE3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FC,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+        MODEL_TENSOR.D2T,
+    ],
+    MODEL_ARCH.DFLASH: [
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
+        MODEL_TENSOR.LAYER_OUT_SCALE,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV,
+        MODEL_TENSOR.ATTN_KV_NORM,
+        MODEL_TENSOR.ATTN_OUT_A,
+        MODEL_TENSOR.ATTN_OUT_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FC,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+        MODEL_TENSOR.D2T,
+        # optional DSpark heads
+        MODEL_TENSOR.DSPARK_MARKOV_W1,
+        MODEL_TENSOR.DSPARK_MARKOV_W2,
+        MODEL_TENSOR.DSPARK_CONF_PROJ,
+        MODEL_TENSOR.DFLASH_ATTN_CONV_BASE,
+        MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ,
+        MODEL_TENSOR.DFLASH_FFN_CONV_BASE,
+        MODEL_TENSOR.DFLASH_FFN_CONV_PROJ,
+        MODEL_TENSOR.DFLASH_SELECTOR_PREV,
+        MODEL_TENSOR.DFLASH_SELECTOR_NEXT,
+        MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN,
+    ],
+    MODEL_ARCH.MISTRAL4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
+    MODEL_ARCH.MIMO2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.LAYER_OUT_NORM,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.STEP35: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP tensors (Step3p5 draft head)
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+    ],
+    MODEL_ARCH.SPARK2_5: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.LLAMA_EMBED: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.MAINCODER: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.KIMI_LINEAR: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.SSM_CONV1D_Q,
+        MODEL_TENSOR.SSM_CONV1D_K,
+        MODEL_TENSOR.SSM_CONV1D_V,
+        MODEL_TENSOR.SSM_F_A,
+        MODEL_TENSOR.SSM_F_B,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_G_A,
+        MODEL_TENSOR.SSM_G_B,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
+    MODEL_ARCH.KIMI_K3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_RES_SCORE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_RES_SCORE,
+        MODEL_TENSOR.FFN_RES_SCORE,
+        # MLA (full-attention layers)
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        # KDA (linear-attention layers)
+        MODEL_TENSOR.SSM_CONV1D_Q,
+        MODEL_TENSOR.SSM_CONV1D_K,
+        MODEL_TENSOR.SSM_CONV1D_V,
+        MODEL_TENSOR.SSM_F_A,
+        MODEL_TENSOR.SSM_F_B,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_G,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        # FFN
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_ROUTED_DOWN,
+        MODEL_TENSOR.FFN_ROUTED_UP,
+        MODEL_TENSOR.FFN_ROUTED_NORM,
+    ],
+    MODEL_ARCH.TALKIE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.LAYER_OUT_SCALE,
+    ],
+    MODEL_ARCH.MELLUM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
+    MODEL_ARCH.NANBEIGE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.QWEN3TTS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.POCKETTTS: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
+    MODEL_ARCH.K2HORIZON: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_V_GATE,
+        MODEL_TENSOR.ATTN_V_EXP,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
+}
+
+# tensors that will not be serialized
+MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
+    MODEL_ARCH.LLAMA: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.DECI: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.BAICHUAN: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.QWEN: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.CODESHELL: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.ORION: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.STARCODER2: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.XVERSE: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.DEEPSEEK: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.DEEPSEEK2: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.DEEPSEEK2OCR: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.DEEPSEEK32: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.HY_V4: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.CHATGLM: [
+        MODEL_TENSOR.ROPE_FREQS,
+    ],
+    MODEL_ARCH.NEMOTRON: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.BAILINGMOE: [
+        MODEL_TENSOR.ROPE_FREQS,
+    ],
+    MODEL_ARCH.PANGU_EMBED: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+    MODEL_ARCH.NANBEIGE: [
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+    ],
+}
+
+#
+# types
+#
+
+
+class TokenType(IntEnum):
+    NORMAL       = 1
+    UNKNOWN      = 2
+    CONTROL      = 3
+    USER_DEFINED = 4
+    UNUSED       = 5
+    BYTE         = 6
+
+
+class RopeScalingType(Enum):
+    NONE     = 'none'
+    LINEAR   = 'linear'
+    YARN     = 'yarn'
+    LONGROPE = 'longrope'
+
+
+# M-RoPE: input position slot (t, y, x, z) that feeds each RoPE section, in section order
+class RopeSectionOrder(Enum):
+    TYXZ = 'tyxz' # default
+    ZYXT = 'zyxt'
+
+
+class PoolingType(IntEnum):
+    NONE = 0
+    MEAN = 1
+    CLS  = 2
+    LAST = 3
+    RANK = 4
+
+
+class GGMLQuantizationType(IntEnum):
+    F32     = 0
+    F16     = 1
+    Q4_0    = 2
+    Q4_1    = 3
+    Q5_0    = 6
+    Q5_1    = 7
+    Q8_0    = 8
+    Q8_1    = 9
+    Q2_K    = 10
+    Q3_K    = 11
+    Q4_K    = 12
+    Q5_K    = 13
+    Q6_K    = 14
+    Q8_K    = 15
+    IQ2_XXS = 16
+    IQ2_XS  = 17
+    IQ3_XXS = 18
+    IQ1_S   = 19
+    IQ4_NL  = 20
+    IQ3_S   = 21
+    IQ2_S   = 22
+    IQ4_XS  = 23
+    I8      = 24
+    I16     = 25
+    I32     = 26
+    I64     = 27
+    F64     = 28
+    IQ1_M   = 29
+    BF16    = 30
+    TQ1_0   = 34
+    TQ2_0   = 35
+    MXFP4   = 39
+    NVFP4   = 40
+    Q1_0    = 41
+    Q2_0    = 42
+
+
+class ExpertGatingFuncType(IntEnum):
+    SOFTMAX       = 1
+    SIGMOID       = 2
+    SQRTSOFTPLUS  = 4
+
+
+# TODO: add GGMLFileType from ggml_ftype in ggml.h
+
+
+# from llama_ftype in llama.h
+# ALL VALUES SHOULD BE THE SAME HERE AS THEY ARE OVER THERE.
+class LlamaFileType(IntEnum):
+    ALL_F32              = 0
+    MOSTLY_F16           = 1   # except 1d tensors
+    MOSTLY_Q4_0          = 2   # except 1d tensors
+    MOSTLY_Q4_1          = 3   # except 1d tensors
+    # MOSTLY_Q4_1_SOME_F16 = 4   # tok_embeddings.weight and output.weight are F16
+    # MOSTLY_Q4_2        = 5   # support has been removed
+    # MOSTLY_Q4_3        = 6   # support has been removed
+    MOSTLY_Q8_0          = 7   # except 1d tensors
+    MOSTLY_Q5_0          = 8   # except 1d tensors
+    MOSTLY_Q5_1          = 9   # except 1d tensors
+    MOSTLY_Q2_K          = 10  # except 1d tensors
+    MOSTLY_Q3_K_S        = 11  # except 1d tensors
+    MOSTLY_Q3_K_M        = 12  # except 1d tensors
+    MOSTLY_Q3_K_L        = 13  # except 1d tensors
+    MOSTLY_Q4_K_S        = 14  # except 1d tensors
+    MOSTLY_Q4_K_M        = 15  # except 1d tensors
+    MOSTLY_Q5_K_S        = 16  # except 1d tensors
+    MOSTLY_Q5_K_M        = 17  # except 1d tensors
+    MOSTLY_Q6_K          = 18  # except 1d tensors
+    MOSTLY_IQ2_XXS       = 19  # except 1d tensors
+    MOSTLY_IQ2_XS        = 20  # except 1d tensors
+    MOSTLY_Q2_K_S        = 21  # except 1d tensors
+    MOSTLY_IQ3_XS        = 22  # except 1d tensors
+    MOSTLY_IQ3_XXS       = 23  # except 1d tensors
+    MOSTLY_IQ1_S         = 24  # except 1d tensors
+    MOSTLY_IQ4_NL        = 25  # except 1d tensors
+    MOSTLY_IQ3_S         = 26  # except 1d tensors
+    MOSTLY_IQ3_M         = 27  # except 1d tensors
+    MOSTLY_IQ2_S         = 28  # except 1d tensors
+    MOSTLY_IQ2_M         = 29  # except 1d tensors
+    MOSTLY_IQ4_XS        = 30  # except 1d tensors
+    MOSTLY_IQ1_M         = 31  # except 1d tensors
+    MOSTLY_BF16          = 32  # except 1d tensors
+    # MOSTLY_Q4_0_4_4      = 33  # removed from gguf files, use Q4_0 and runtime repack
+    # MOSTLY_Q4_0_4_8      = 34  # removed from gguf files, use Q4_0 and runtime repack
+    # MOSTLY_Q4_0_8_8      = 35  # removed from gguf files, use Q4_0 and runtime repack
+    MOSTLY_TQ1_0         = 36  # except 1d tensors
+    MOSTLY_TQ2_0         = 37  # except 1d tensors
+    MOSTLY_MXFP4_MOE     = 38  # except 1d tensors
+    MOSTLY_NVFP4         = 39  # except 1d tensors
+    MOSTLY_Q1_0          = 40  # except 1d tensors
+    MOSTLY_Q2_0          = 41  # except 1d tensors
+
+    GUESSED              = 1024  # not specified in the model file
+
+
+class GGUFEndian(IntEnum):
+    LITTLE = 0
+    BIG = 1
+
+
+class GGUFValueType(IntEnum):
+    UINT8   = 0
+    INT8    = 1
+    UINT16  = 2
+    INT16   = 3
+    UINT32  = 4
+    INT32   = 5
+    FLOAT32 = 6
+    BOOL    = 7
+    STRING  = 8
+    ARRAY   = 9
+    UINT64  = 10
+    INT64   = 11
+    FLOAT64 = 12
+
+    @staticmethod
+    def get_type(val: Any) -> GGUFValueType:
+        if isinstance(val, (str, bytes, bytearray)):
+            return GGUFValueType.STRING
+        elif isinstance(val, list):
+            return GGUFValueType.ARRAY
+        elif isinstance(val, float):
+            return GGUFValueType.FLOAT32
+        elif isinstance(val, bool):
+            return GGUFValueType.BOOL
+        elif isinstance(val, int):
+            return GGUFValueType.INT32
+        # TODO: need help with 64-bit types in Python
+        else:
+            raise ValueError(f"Unknown type: {type(val)}")
+
+
+class DecisionType:
+    LAYA    = "laya"     # head blocks + scorer on the hidden state of one marker token per option
+    OPENJEV = "openjev"  # logits of one label token per option
+    LEV     = "lev"      # same as openjev, noul is read from a rating scale
+    KEV     = "kev"      # dot product of the hidden states of the last token and of one end token per option
+    NIMBLE  = "nimble"   # same as openjev, the prompt lists all the questions of the request
+    CLEF    = "clef"     # joint head over all questions, one score per option
+    PPLX_DECIDER = "pplx-decider"  # same as openjev, label codes of 1 or 2 letters
+    VLLM_DECISION2 = "vllm-decision2"  # head over the end token of each option and the last token, one score per option
+    VLLM_D3 = "vllm-d3"  # same as pplx-decider, non-causal attention
+    LFM2_D1 = "lfm2-d1"  # same as openjev, the labels depend on the question type
+    LFM2_D1_OMNI = "lfm2-d1-omni"  # same head as laya on a bidirectional LFM2 trunk, other prompt layout
+
+
+class VisionProjectorType:
+    GEMMA3 = "gemma3"
+    GEMMA3NV = "gemma3nv"
+    GEMMA3NA = "gemma3na"
+    GEMMA4V = "gemma4v"
+    GEMMA4A = "gemma4a"
+    GEMMA4UV = "gemma4uv" # "unified" variant
+    GEMMA4UA = "gemma4ua" # "unified" variant
+    PHI4 = "phi4"
+    IDEFICS3 = "idefics3"
+    PIXTRAL = "pixtral"
+    LLAMA4 = "llama4"
+    QWEN2VL = "qwen2vl_merger"
+    QWEN25VL = "qwen2.5vl_merger"
+    EXAONE4_5 = "exaone4_5"
+    QWEN3VL = "qwen3vl_merger"
+    LING3VL = "ling3vl"
+    STEP3VL = "step3vl"
+    ULTRAVOX = "ultravox"
+    INTERNVL = "internvl"
+    QWEN2A = "qwen2a" # audio
+    QWEN3A = "qwen3a" # audio
+    GLMA = "glma" # audio
+    QWEN25O = "qwen2.5o" # omni
+    VOXTRAL = "voxtral"
+    MERALION = "meralion"  # audio: Whisper + gated MLP adaptor
+    LFM2 = "lfm2"
+    KIMIVL = "kimivl"
+    PADDLEOCR = "paddleocr"
+    KIMIK25 = "kimik25"
+    LIGHTONOCR = "lightonocr"
+    COGVLM = "cogvlm"
+    JANUS_PRO = "janus_pro"
+    DOTSOCR = "dots_ocr"
+    DOTS3NOTE_V = "dots3note_v"
+    DOTS3NOTE_A = "dots3note_a" # audio
+    DEEPSEEKOCR = "deepseekocr"
+    DEEPSEEKOCR2 = "deepseekocr2"
+    DEEPSEEK4V = "deepseek4v"
+    LFM2A = "lfm2a" # audio
+    MUSIC_FLAMINGO = "musicflamingo" # audio
+    GLM4V = "glm4v"
+    GLM5V = "glm5v"
+    YOUTUVL = "youtuvl"
+    NEMOTRON_V2_VL = "nemotron_v2_vl"
+    QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder
+    QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
+    POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
+    POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
+    HUNYUANVL      = "hunyuanvl"
+    PARAKEET       = "parakeet"  # audio
+    MINIMAXM3      = "minimax_m3"
+    MINICPMV4_6    = "minicpmv4_6"
+    MINICPMV4_7    = "minicpmv4_7"
+    GRANITE_SPEECH = "granite_speech"  # audio
+    MIMOVL         = "mimovl"
+    MIMO_AUDIO     = "mimo_audio"
+    GRANITE4_VISION = "granite4_vision"
+    MUSE_GLIMMER   = "muse-glimmer"
+    COHERE2V       = "cohere2v"
+    D1OMNI_V       = "d1omni_v"  # lfm2 vision, without separator tokens
+    D1OMNI_A       = "d1omni_a"  # lfm2a audio, with a residual block after the projector
+
+
+# Items here are (block size, type size)
+QK_K = 256
+GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = {
+    GGMLQuantizationType.F32:     (1, 4),
+    GGMLQuantizationType.F16:     (1, 2),
+    GGMLQuantizationType.Q4_0:    (32, 2 + 16),
+    GGMLQuantizationType.Q4_1:    (32, 2 + 2 + 16),
+    GGMLQuantizationType.Q5_0:    (32, 2 + 4 + 16),
+    GGMLQuantizationType.Q5_1:    (32, 2 + 2 + 4 + 16),
+    GGMLQuantizationType.Q8_0:    (32, 2 + 32),
+    GGMLQuantizationType.Q8_1:    (32, 2 + 2 + 32),
+    GGMLQuantizationType.Q2_K:    (256, 2 + 2 + QK_K // 16 + QK_K // 4),
+    GGMLQuantizationType.Q3_K:    (256, 2 + QK_K // 4 + QK_K // 8 + 12),
+    GGMLQuantizationType.Q4_K:    (256, 2 + 2 + QK_K // 2 + 12),
+    GGMLQuantizationType.Q5_K:    (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12),
+    GGMLQuantizationType.Q6_K:    (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16),
+    GGMLQuantizationType.Q8_K:    (256, 4 + QK_K + QK_K // 8),
+    GGMLQuantizationType.IQ2_XXS: (256, 2 + QK_K // 4),
+    GGMLQuantizationType.IQ2_XS:  (256, 2 + QK_K // 4 + QK_K // 32),
+    GGMLQuantizationType.IQ3_XXS: (256, 2 + QK_K // 4 + QK_K // 8),
+    GGMLQuantizationType.IQ1_S:   (256, 2 + QK_K // 8 + QK_K // 16),
+    GGMLQuantizationType.IQ4_NL:  (32, 2 + 16),
+    GGMLQuantizationType.IQ3_S:   (256, 2 + QK_K // 4 + QK_K // 8 + QK_K // 32 + 4),
+    GGMLQuantizationType.IQ2_S:   (256, 2 + QK_K // 4 + QK_K // 16),
+    GGMLQuantizationType.IQ4_XS:  (256, 2 + 2 + QK_K // 2 + QK_K // 64),
+    GGMLQuantizationType.I8:      (1, 1),
+    GGMLQuantizationType.I16:     (1, 2),
+    GGMLQuantizationType.I32:     (1, 4),
+    GGMLQuantizationType.I64:     (1, 8),
+    GGMLQuantizationType.F64:     (1, 8),
+    GGMLQuantizationType.IQ1_M:   (256, QK_K // 8 + QK_K // 16  + QK_K // 32),
+    GGMLQuantizationType.BF16:    (1, 2),
+    GGMLQuantizationType.TQ1_0:   (256, 2 + 4 * 13),
+    GGMLQuantizationType.TQ2_0:   (256, 2 + 64),
+    GGMLQuantizationType.MXFP4:   (32, 1 + 16),
+    GGMLQuantizationType.NVFP4:   (64, 4 + 32),
+    GGMLQuantizationType.Q1_0:    (128, 2 + 16),
+    GGMLQuantizationType.Q2_0:    (64, 2 + 16),
+}
+
+
+# Aliases for backward compatibility.
+
+# general
+KEY_GENERAL_ARCHITECTURE         = Keys.General.ARCHITECTURE
+KEY_GENERAL_QUANTIZATION_VERSION = Keys.General.QUANTIZATION_VERSION
+KEY_GENERAL_ALIGNMENT            = Keys.General.ALIGNMENT
+KEY_GENERAL_NAME                 = Keys.General.NAME
+KEY_GENERAL_AUTHOR               = Keys.General.AUTHOR
+KEY_GENERAL_URL                  = Keys.General.URL
+KEY_GENERAL_DESCRIPTION          = Keys.General.DESCRIPTION
+KEY_GENERAL_LICENSE              = Keys.General.LICENSE
+KEY_GENERAL_SOURCE_URL           = Keys.General.SOURCE_URL
+KEY_GENERAL_FILE_TYPE            = Keys.General.FILE_TYPE
+
+# LLM
+KEY_VOCAB_SIZE            = Keys.LLM.VOCAB_SIZE
+KEY_CONTEXT_LENGTH        = Keys.LLM.CONTEXT_LENGTH
+KEY_EMBEDDING_LENGTH      = Keys.LLM.EMBEDDING_LENGTH
+KEY_BLOCK_COUNT           = Keys.LLM.BLOCK_COUNT
+KEY_FEED_FORWARD_LENGTH   = Keys.LLM.FEED_FORWARD_LENGTH
+KEY_USE_PARALLEL_RESIDUAL = Keys.LLM.USE_PARALLEL_RESIDUAL
+KEY_TENSOR_DATA_LAYOUT    = Keys.LLM.TENSOR_DATA_LAYOUT
+
+# attention
+KEY_ATTENTION_HEAD_COUNT        = Keys.Attention.HEAD_COUNT
+KEY_ATTENTION_HEAD_COUNT_KV     = Keys.Attention.HEAD_COUNT_KV
+KEY_ATTENTION_MAX_ALIBI_BIAS    = Keys.Attention.MAX_ALIBI_BIAS
+KEY_ATTENTION_CLAMP_KQV         = Keys.Attention.CLAMP_KQV
+KEY_ATTENTION_LAYERNORM_EPS     = Keys.Attention.LAYERNORM_EPS
+KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS
+
+# RoPE
+KEY_ROPE_DIMENSION_COUNT           = Keys.Rope.DIMENSION_COUNT
+KEY_ROPE_FREQ_BASE                 = Keys.Rope.FREQ_BASE
+KEY_ROPE_SCALING_TYPE              = Keys.Rope.SCALING_TYPE
+KEY_ROPE_SCALING_FACTOR            = Keys.Rope.SCALING_FACTOR
+KEY_ROPE_SCALING_ORIG_CTX_LEN      = Keys.Rope.SCALING_ORIG_CTX_LEN
+KEY_ROPE_SCALING_FINETUNED         = Keys.Rope.SCALING_FINETUNED
+
+# SSM
+KEY_SSM_CONV_KERNEL    = Keys.SSM.CONV_KERNEL
+KEY_SSM_INNER_SIZE     = Keys.SSM.INNER_SIZE
+KEY_SSM_STATE_SIZE     = Keys.SSM.STATE_SIZE
+KEY_SSM_TIME_STEP_RANK = Keys.SSM.TIME_STEP_RANK
+KEY_SSM_GROUP_COUNT    = Keys.SSM.GROUP_COUNT
+KEY_SSM_DT_B_C_RMS     = Keys.SSM.DT_B_C_RMS
+
+# KDA
+KEY_KDA_HEAD_DIM         = Keys.KDA.HEAD_DIM
+KEY_KDA_SAFE_GATE        = Keys.KDA.SAFE_GATE
+KEY_KDA_GATE_LOWER_BOUND = Keys.KDA.GATE_LOWER_BOUND
+
+# tokenization
+KEY_TOKENIZER_MODEL      = Keys.Tokenizer.MODEL
+KEY_TOKENIZER_PRE        = Keys.Tokenizer.PRE
+KEY_TOKENIZER_LIST       = Keys.Tokenizer.LIST
+KEY_TOKENIZER_TOKEN_TYPE = Keys.Tokenizer.TOKEN_TYPE
+KEY_TOKENIZER_SCORES     = Keys.Tokenizer.SCORES
+KEY_TOKENIZER_MERGES     = Keys.Tokenizer.MERGES
+KEY_TOKENIZER_BOS_ID     = Keys.Tokenizer.BOS_ID
+KEY_TOKENIZER_EOS_ID     = Keys.Tokenizer.EOS_ID
+KEY_TOKENIZER_EOT_ID     = Keys.Tokenizer.EOT_ID
+KEY_TOKENIZER_EOM_ID     = Keys.Tokenizer.EOM_ID
+KEY_TOKENIZER_UNK_ID     = Keys.Tokenizer.UNK_ID
+KEY_TOKENIZER_SEP_ID     = Keys.Tokenizer.SEP_ID
+KEY_TOKENIZER_PAD_ID     = Keys.Tokenizer.PAD_ID
+KEY_TOKENIZER_MASK_ID    = Keys.Tokenizer.MASK_ID
+KEY_TOKENIZER_HF_JSON    = Keys.Tokenizer.HF_JSON
+KEY_TOKENIZER_RWKV       = Keys.Tokenizer.RWKV
+
+KEY_TOKENIZER_FIM_PRE_ID = Keys.Tokenizer.FIM_PRE_ID
+KEY_TOKENIZER_FIM_SUF_ID = Keys.Tokenizer.FIM_SUF_ID
+KEY_TOKENIZER_FIM_MID_ID = Keys.Tokenizer.FIM_MID_ID
+KEY_TOKENIZER_FIM_PAD_ID = Keys.Tokenizer.FIM_PAD_ID
+KEY_TOKENIZER_FIM_REP_ID = Keys.Tokenizer.FIM_REP_ID
+KEY_TOKENIZER_FIM_SEP_ID = Keys.Tokenizer.FIM_SEP_ID
+
+# deprecated
+KEY_TOKENIZER_PREFIX_ID  = Keys.Tokenizer.PREFIX_ID
+KEY_TOKENIZER_SUFFIX_ID  = Keys.Tokenizer.SUFFIX_ID
+KEY_TOKENIZER_MIDDLE_ID  = Keys.Tokenizer.MIDDLE_ID
diff --git a/python/freetoken/third_party/gguf/freetoken.patch b/python/freetoken/third_party/gguf/freetoken.patch
new file mode 100644
index 0000000000..b8016ff554
--- /dev/null
+++ b/python/freetoken/third_party/gguf/freetoken.patch
@@ -0,0 +1,23 @@
+--- a/gguf-py/gguf/__init__.py
++++ b/gguf-py/gguf/__init__.py
+@@ -1,9 +1,4 @@
+ from .constants import *
+ from .lazy import *
+ from .gguf_reader import *
+-from .gguf_writer import *
+ from .quants import *
+-from .tensor_mapping import *
+-from .vocab import *
+-from .utility import *
+-from .metadata import *
+--- a/gguf-py/gguf/gguf_reader.py
++++ b/gguf-py/gguf/gguf_reader.py
+@@ -21,7 +21,7 @@
+     # Allow running file in package as a script.
+     sys.path.insert(0, str(Path(__file__).parent.parent))
+ 
+-from gguf.constants import (
++from .constants import (
+     GGML_MAX_DIMS,
+     GGML_QUANT_SIZES,
+     GGUF_DEFAULT_ALIGNMENT,
diff --git a/python/freetoken/third_party/gguf/gguf_reader.py b/python/freetoken/third_party/gguf/gguf_reader.py
new file mode 100644
index 0000000000..41eb6521ea
--- /dev/null
+++ b/python/freetoken/third_party/gguf/gguf_reader.py
@@ -0,0 +1,390 @@
+#
+# GGUF file reading/modification support. For API usage information,
+# please see the files scripts/ for some fairly simple examples.
+#
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from collections import OrderedDict
+from typing import Any, Literal, NamedTuple, TypeVar, Union
+
+import numpy as np
+import numpy.typing as npt
+
+from .quants import quant_shape_to_byte_shape
+
+if __name__ == "__main__":
+    from pathlib import Path
+
+    # Allow running file in package as a script.
+    sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from .constants import (
+    GGML_MAX_DIMS,
+    GGML_QUANT_SIZES,
+    GGUF_DEFAULT_ALIGNMENT,
+    GGUF_MAGIC,
+    GGUF_VERSION,
+    GGMLQuantizationType,
+    GGUFValueType,
+    GGUFEndian,
+)
+
+# limits mirroring ggml/src/gguf.cpp (not part of gguf.h)
+GGUF_MAX_STRING_LENGTH  = 1024 * 1024 * 1024
+GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024
+
+logger = logging.getLogger(__name__)
+
+READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION]
+
+
+class ReaderField(NamedTuple):
+    # Offset to start of this field.
+    offset: int
+
+    # Name of the field (not necessarily from file data).
+    name: str
+
+    # Data parts. Some types have multiple components, such as strings
+    # that consist of a length followed by the string data.
+    parts: list[npt.NDArray[Any]] = []
+
+    # Indexes into parts that we can call the actual data. For example
+    # an array of strings will be populated with indexes to the actual
+    # string data.
+    data: list[int] = [-1]
+
+    types: list[GGUFValueType] = []
+
+    def contents(self, index_or_slice: int | slice = slice(None)) -> Any:
+        if self.types:
+            to_string = lambda x: str(x.tobytes(), encoding='utf-8') # noqa: E731
+            main_type = self.types[0]
+
+            if main_type == GGUFValueType.ARRAY:
+                sub_type = self.types[-1]
+
+                if sub_type == GGUFValueType.STRING:
+                    indices = self.data[index_or_slice]
+
+                    if isinstance(index_or_slice, int):
+                        return to_string(self.parts[indices]) # type: ignore
+                    else:
+                        return [to_string(self.parts[idx]) for idx in indices] # type: ignore
+                else:
+                    # FIXME: When/if _get_field_parts() support multi-dimensional arrays, this must do so too
+
+                    # Check if it's unsafe to perform slice optimization on data
+                    # if any(True for idx in self.data if len(self.parts[idx]) != 1):
+                    #     optim_slice = slice(None)
+                    # else:
+                    #     optim_slice = index_or_slice
+                    #     index_or_slice = slice(None)
+
+                    # if isinstance(optim_slice, int):
+                    #     return self.parts[self.data[optim_slice]].tolist()[0]
+                    # else:
+                    #     return [pv for idx in self.data[optim_slice] for pv in self.parts[idx].tolist()][index_or_slice]
+
+                    if isinstance(index_or_slice, int):
+                        return self.parts[self.data[index_or_slice]].tolist()[0]
+                    else:
+                        return [pv for idx in self.data[index_or_slice] for pv in self.parts[idx].tolist()]
+
+            if main_type == GGUFValueType.STRING:
+                return to_string(self.parts[-1])
+            else:
+                return self.parts[-1].tolist()[0]
+
+        return None
+
+
+class ReaderTensor(NamedTuple):
+    name: str
+    tensor_type: GGMLQuantizationType
+    shape: npt.NDArray[np.uint32]
+    n_elements: int
+    n_bytes: int
+    data_offset: int
+    data: npt.NDArray[Any]
+    field: ReaderField
+
+
+class GGUFReader:
+    # I - same as host, S - swapped
+    byte_order: Literal['I', 'S'] = 'I'
+    alignment: int = GGUF_DEFAULT_ALIGNMENT
+    data_offset: int
+
+    # Note: Internal helper, API may change.
+    gguf_scalar_to_np: dict[GGUFValueType, type[np.generic]] = {
+        GGUFValueType.UINT8:   np.uint8,
+        GGUFValueType.INT8:    np.int8,
+        GGUFValueType.UINT16:  np.uint16,
+        GGUFValueType.INT16:   np.int16,
+        GGUFValueType.UINT32:  np.uint32,
+        GGUFValueType.INT32:   np.int32,
+        GGUFValueType.FLOAT32: np.float32,
+        GGUFValueType.UINT64:  np.uint64,
+        GGUFValueType.INT64:   np.int64,
+        GGUFValueType.FLOAT64: np.float64,
+        GGUFValueType.BOOL:    np.bool_,
+    }
+
+    def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'):
+        self.data = np.memmap(path, mode = mode)
+        offs = 0
+
+        # Check for GGUF magic
+        if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC:
+            raise ValueError('GGUF magic invalid')
+        offs += 4
+
+        # Check GGUF version
+        temp_version = self._get(offs, np.uint32)
+        if temp_version[0] & 65535 == 0:
+            # If we get 0 here that means it's (probably) a GGUF file created for
+            # the opposite byte order of the machine this script is running on.
+            self.byte_order = 'S'
+            temp_version = temp_version.view(temp_version.dtype.newbyteorder(self.byte_order))
+        version = temp_version[0]
+        if version not in READER_SUPPORTED_VERSIONS:
+            raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle')
+        if sys.byteorder == "little":
+            # Host is little endian
+            host_endian = GGUFEndian.LITTLE
+            swapped_endian = GGUFEndian.BIG
+        else:
+            # Sorry PDP or other weird systems that don't use BE or LE.
+            host_endian = GGUFEndian.BIG
+            swapped_endian = GGUFEndian.LITTLE
+        self.endianess = swapped_endian if self.byte_order == "S" else host_endian
+        self.fields: OrderedDict[str, ReaderField] = OrderedDict()
+        self.tensors: list[ReaderTensor] = []
+        offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32]))
+
+        # Check tensor count and kv count
+        temp_counts = self._get(offs, np.uint64, 2)
+        offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64]))
+        offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64]))
+        tensor_count, kv_count = temp_counts
+        if tensor_count > GGUF_MAX_ARRAY_ELEMENTS:
+            raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
+        if kv_count > GGUF_MAX_ARRAY_ELEMENTS:
+            raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
+        offs = self._build_fields(offs, kv_count)
+
+        # Build Tensor Info Fields
+        offs, tensors_fields = self._build_tensor_info(offs, tensor_count)
+        new_align = self.fields.get('general.alignment')
+        if new_align is not None:
+            if new_align.types != [GGUFValueType.UINT32]:
+                raise ValueError('Bad type for general.alignment field')
+            self.alignment = new_align.parts[-1][0]
+            # Ensure alignment is a non-zero power of two
+            if self.alignment == 0 or (self.alignment & (self.alignment - 1)) != 0:
+                raise ValueError('Invalid alignment: must be a non-zero power of two')
+        padding = offs % self.alignment
+        if padding != 0:
+            offs += self.alignment - padding
+        self.data_offset = offs
+        self._build_tensors(offs, tensors_fields)
+
+    _DT = TypeVar('_DT', bound = npt.DTypeLike)
+
+    # Fetch a key/value metadata field by key.
+    def get_field(self, key: str) -> Union[ReaderField, None]:
+        return self.fields.get(key, None)
+
+    # Fetch a tensor from the list by index.
+    def get_tensor(self, idx: int) -> ReaderTensor:
+        return self.tensors[idx]
+
+    def _get(
+        self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None,
+    ) -> npt.NDArray[Any]:
+        count = int(count)
+        itemsize = int(np.empty([], dtype = dtype).itemsize)
+        end_offs = offset + itemsize * count
+        arr = self.data[offset:end_offs].view(dtype=dtype)[:count]
+        return arr.view(arr.dtype.newbyteorder(self.byte_order if override_order is None else override_order))
+
+    def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int:
+        if field.name in self.fields:
+            # TODO: add option to make this a warning and accept duplicate keys like below
+            raise KeyError(f'Duplicate {field.name} already in list at offset {field.offset}')
+
+            # logger.warning(f'Duplicate key {field.name} at offset {field.offset}')
+            # self.fields[field.name + '_{}'.format(field.offset)] = field
+        else:
+            self.fields[field.name] = field
+        return 0 if skip_sum else sum(int(part.nbytes) for part in field.parts)
+
+    def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]:
+        slen = self._get(offset, np.uint64)
+        if int(slen[0]) > GGUF_MAX_STRING_LENGTH:
+            raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}')
+        if offset + 8 + int(slen[0]) > self.data.nbytes:
+            raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}')
+        return slen, self._get(offset + 8, np.uint8, slen[0])
+
+    def _get_field_parts(
+        self, orig_offs: int, raw_type: int,
+    ) -> tuple[int, list[npt.NDArray[Any]], list[int], list[GGUFValueType]]:
+        offs = orig_offs
+        types: list[GGUFValueType] = []
+        gtype = GGUFValueType(raw_type)
+        types.append(gtype)
+        # Handle strings.
+        if gtype == GGUFValueType.STRING:
+            sparts: list[npt.NDArray[Any]] = list(self._get_str(offs))
+            size = sum(int(part.nbytes) for part in sparts)
+            return size, sparts, [1], types
+        # Check if it's a simple scalar type.
+        nptype = self.gguf_scalar_to_np.get(gtype)
+        if nptype is not None:
+            val = self._get(offs, nptype)
+            return int(val.nbytes), [val], [0], types
+        # Handle arrays.
+        if gtype == GGUFValueType.ARRAY:
+            raw_itype = self._get(offs, np.uint32)
+            offs += int(raw_itype.nbytes)
+            alen = self._get(offs, np.uint64)
+            if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS:
+                raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
+            offs += int(alen.nbytes)
+            aparts: list[npt.NDArray[Any]] = [raw_itype, alen]
+            data_idxs: list[int] = []
+            # FIXME: Handle multi-dimensional arrays properly instead of flattening
+            for idx in range(alen[0]):
+                curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
+                if idx == 0:
+                    types += curr_types
+                idxs_offs = len(aparts)
+                aparts += curr_parts
+                data_idxs += (idx + idxs_offs for idx in curr_idxs)
+                offs += curr_size
+            return offs - orig_offs, aparts, data_idxs, types
+        # We can't deal with this one.
+        raise ValueError(f'Unknown/unhandled field type {gtype}')
+
+    def _get_tensor_info_field(self, orig_offs: int) -> ReaderField:
+        offs = orig_offs
+
+        # Get Tensor Name
+        name_len, name_data = self._get_str(offs)
+        offs += int(name_len.nbytes + name_data.nbytes)
+
+        # Get Tensor Dimensions Count
+        n_dims = self._get(offs, np.uint32)
+        offs += int(n_dims.nbytes)
+        if n_dims[0] > GGML_MAX_DIMS:
+            raise ValueError(f'Tensor dimensions count {n_dims[0]} exceeds GGML_MAX_DIMS ({GGML_MAX_DIMS})')
+
+        # Get Tensor Dimension Array
+        dims = self._get(offs, np.uint64, n_dims[0])
+        offs += int(dims.nbytes)
+
+        # Get Tensor Encoding Scheme Type
+        raw_dtype = self._get(offs, np.uint32)
+        offs += int(raw_dtype.nbytes)
+
+        # Get Tensor Offset
+        offset_tensor = self._get(offs, np.uint64)
+        offs += int(offset_tensor.nbytes)
+
+        return ReaderField(
+            orig_offs,
+            str(bytes(name_data), encoding = 'utf-8'),
+            [name_len, name_data, n_dims, dims, raw_dtype, offset_tensor],
+            [1, 3, 4, 5],
+        )
+
+    def _build_fields(self, offs: int, count: int) -> int:
+        for _ in range(count):
+            orig_offs = offs
+            kv_klen, kv_kdata = self._get_str(offs)
+            offs += int(kv_klen.nbytes + kv_kdata.nbytes)
+            raw_kv_type = self._get(offs, np.uint32)
+            offs += int(raw_kv_type.nbytes)
+            parts: list[npt.NDArray[Any]] = [kv_klen, kv_kdata, raw_kv_type]
+            idxs_offs = len(parts)
+            field_size, field_parts, field_idxs, field_types = self._get_field_parts(offs, raw_kv_type[0])
+            parts += field_parts
+            self._push_field(ReaderField(
+                orig_offs,
+                str(bytes(kv_kdata), encoding = 'utf-8'),
+                parts,
+                [idx + idxs_offs for idx in field_idxs],
+                field_types,
+            ), skip_sum = True)
+            offs += field_size
+        return offs
+
+    def _build_tensor_info(self, offs: int, count: int) -> tuple[int, list[ReaderField]]:
+        tensor_fields = []
+        for _ in range(count):
+            field = self._get_tensor_info_field(offs)
+            offs += sum(int(part.nbytes) for part in field.parts)
+            tensor_fields.append(field)
+        return offs, tensor_fields
+
+    def _build_tensors(self, start_offs: int, fields: list[ReaderField]) -> None:
+        tensors = []
+        tensor_names = set() # keep track of name to prevent duplicated tensors
+        for field in fields:
+            _name_len, name_data, _n_dims, dims, raw_dtype, offset_tensor = field.parts
+            # check if there's any tensor having same name already in the list
+            tensor_name = str(bytes(name_data), encoding = 'utf-8')
+            if tensor_name in tensor_names:
+                raise ValueError(f'Found duplicated tensor with name {tensor_name}')
+            tensor_names.add(tensor_name)
+            ggml_type = GGMLQuantizationType(raw_dtype[0])
+            # use Python ints: np.prod on uint64 wraps silently on overflow
+            n_elems = 1
+            for dim in dims.tolist():
+                n_elems *= int(dim)
+            np_dims = tuple(reversed(dims.tolist()))
+            block_size, type_size = GGML_QUANT_SIZES[ggml_type]
+            n_bytes = n_elems * type_size // block_size
+            data_offs = int(start_offs + offset_tensor[0])
+            item_type: npt.DTypeLike
+            if ggml_type == GGMLQuantizationType.F16:
+                item_count = n_elems
+                item_type = np.float16
+            elif ggml_type == GGMLQuantizationType.F32:
+                item_count = n_elems
+                item_type = np.float32
+            elif ggml_type == GGMLQuantizationType.F64:
+                item_count = n_elems
+                item_type = np.float64
+            elif ggml_type == GGMLQuantizationType.I8:
+                item_count = n_elems
+                item_type = np.int8
+            elif ggml_type == GGMLQuantizationType.I16:
+                item_count = n_elems
+                item_type = np.int16
+            elif ggml_type == GGMLQuantizationType.I32:
+                item_count = n_elems
+                item_type = np.int32
+            elif ggml_type == GGMLQuantizationType.I64:
+                item_count = n_elems
+                item_type = np.int64
+            else:
+                item_count = n_bytes
+                item_type = np.uint8
+                np_dims = quant_shape_to_byte_shape(np_dims, ggml_type)
+            tensors.append(ReaderTensor(
+                name = tensor_name,
+                tensor_type = ggml_type,
+                shape = dims,
+                n_elements = n_elems,
+                n_bytes = n_bytes,
+                data_offset = data_offs,
+                data = self._get(data_offs, item_type, item_count).reshape(np_dims),
+                field = field,
+            ))
+        self.tensors = tensors
diff --git a/python/freetoken/third_party/gguf/lazy.py b/python/freetoken/third_party/gguf/lazy.py
new file mode 100644
index 0000000000..a39f223215
--- /dev/null
+++ b/python/freetoken/third_party/gguf/lazy.py
@@ -0,0 +1,293 @@
+from __future__ import annotations
+from abc import ABC, ABCMeta, abstractmethod
+
+import logging
+from typing import Any, Callable
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+
+logger = logging.getLogger(__name__)
+
+
+class LazyMeta(ABCMeta):
+
+    def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any], **kwargs):
+        def __getattr__(self, name: str) -> Any:
+            meta_attr = getattr(self._meta, name)
+            if callable(meta_attr):
+                return type(self)._wrap_fn(
+                    (lambda s, *args, **kwargs: getattr(s, name)(*args, **kwargs)),
+                    use_self=self,
+                )
+            elif isinstance(meta_attr, self._tensor_type):
+                # e.g. self.T with torch.Tensor should still be wrapped
+                return type(self)._wrap_fn(lambda s: getattr(s, name))(self)
+            else:
+                # no need to wrap non-tensor properties,
+                # and they likely don't depend on the actual contents of the tensor
+                return meta_attr
+
+        namespace["__getattr__"] = __getattr__
+
+        # need to make a builder for the wrapped wrapper to copy the name,
+        # or else it fails with very cryptic error messages,
+        # because somehow the same string would end up in every closures
+        def mk_wrap(op_name: str, *, meta_noop: bool = False):
+            # need to wrap the wrapper to get self
+            def wrapped_special_op(self, *args, **kwargs):
+                return type(self)._wrap_fn(
+                    getattr(type(self)._tensor_type, op_name),
+                    meta_noop=meta_noop,
+                )(self, *args, **kwargs)
+            return wrapped_special_op
+
+        # special methods bypass __getattr__, so they need to be added manually
+        # ref: https://docs.python.org/3/reference/datamodel.html#special-lookup
+        # NOTE: doing this from a metaclass is very convenient
+        # TODO: make this even more comprehensive
+        for binary_op in (
+            "lt", "le", "eq", "ne", "ge", "gt",
+            "add", "and", "floordiv", "lshift", "mod", "mul", "matmul",
+            "or", "pow", "rshift", "sub", "truediv", "xor",
+            "iadd", "iand", "ifloordiv", "ilshift", "imod", "imul", "ior", "irshift", "isub", "ixor",
+            "radd", "rand", "rfloordiv", "rmul", "ror", "rpow", "rsub", "rtruediv", "rxor",
+        ):
+            attr_name = f"__{binary_op}__"
+            # evaluation on the meta tensor is needed in case there's broadcasting
+            namespace[attr_name] = mk_wrap(attr_name, meta_noop=False)
+
+        for unary_op in ("not", "abs", "invert", "neg", "pos"):
+            attr_name = f"__{unary_op}__"
+            # the result of these operators usually has the same shape and dtype as the input,
+            # so evaluation on the meta tensor can be skipped.
+            namespace[attr_name] = mk_wrap(attr_name, meta_noop=True)
+
+        for special_op in (
+            "getitem", "setitem", "len",
+        ):
+            attr_name = f"__{special_op}__"
+            namespace[attr_name] = mk_wrap(attr_name, meta_noop=False)
+
+        return super().__new__(cls, name, bases, namespace, **kwargs)
+
+
+# Tree of lazy tensors
+class LazyBase(ABC, metaclass=LazyMeta):
+    _tensor_type: type
+    _meta: Any
+    _data: Any | None
+    _args: tuple
+    _kwargs: dict[str, Any]
+    _func: Callable[[Any], Any] | None
+
+    def __init__(self, *, meta: Any, data: Any | None = None, args: tuple = (), kwargs: dict[str, Any] | None = None, func: Callable[[Any], Any] | None = None):
+        super().__init__()
+        self._meta = meta
+        self._data = data
+        self._args = args
+        self._kwargs = kwargs if kwargs is not None else {}
+        self._func = func
+        assert self._func is not None or self._data is not None
+
+    def __init_subclass__(cls) -> None:
+        if "_tensor_type" not in cls.__dict__:
+            raise TypeError(f"property '_tensor_type' must be defined for {cls!r}")
+        return super().__init_subclass__()
+
+    @staticmethod
+    def _recurse_apply(o: Any, fn: Callable[[Any], Any]) -> Any:
+        # TODO: dict and set
+        if isinstance(o, (list, tuple)):
+            L = []
+            for item in o:
+                L.append(LazyBase._recurse_apply(item, fn))
+            if isinstance(o, tuple):
+                L = tuple(L)
+            return L
+        elif isinstance(o, LazyBase):
+            return fn(o)
+        else:
+            return o
+
+    @classmethod
+    def _wrap_fn(cls, fn: Callable, *, use_self: LazyBase | None = None, meta_noop: bool | DTypeLike | tuple[DTypeLike, Callable[[tuple[int, ...]], tuple[int, ...]]] = False) -> Callable[[Any], Any]:
+        def wrapped_fn(*args, **kwargs):
+            if kwargs is None:
+                kwargs = {}
+            args = ((use_self,) if use_self is not None else ()) + args
+
+            meta_args = LazyBase._recurse_apply(args, lambda t: t._meta)
+            # TODO: maybe handle tensors in kwargs too
+
+            if isinstance(meta_noop, bool) and not meta_noop:
+                try:
+                    res = fn(*meta_args, **kwargs)
+                except NotImplementedError:
+                    # running some operations on PyTorch's Meta tensors can cause this exception
+                    res = None
+            else:
+                # some operators don't need to actually run on the meta tensors
+                assert len(args) > 0
+                res = args[0]
+                assert isinstance(res, cls)
+                res = res._meta
+                # allow operations to override the dtype and shape
+                if meta_noop is not True:
+                    if isinstance(meta_noop, tuple):
+                        dtype, shape = meta_noop
+                        assert callable(shape)
+                        res = cls.meta_with_dtype_and_shape(dtype, shape(res.shape))  # ty: ignore[call-top-callable]
+                    else:
+                        res = cls.meta_with_dtype_and_shape(meta_noop, res.shape)
+
+            if isinstance(res, cls._tensor_type):
+                return cls(meta=cls.eager_to_meta(res), args=args, kwargs=kwargs, func=fn)
+            elif isinstance(res, tuple) and all(isinstance(t, cls._tensor_type) for t in res):
+                # share the evaluation between lazy tuple elements
+                shared_args: list = [args, None]
+
+                def eager_tuple_element(a: list[Any], i: int = 0, /, **kw) -> LazyBase:
+                    assert len(a) == 2
+                    if a[1] is None:
+                        a[1] = fn(*a[0], **kw)
+                    return a[1][i]
+                return tuple(cls(meta=cls.eager_to_meta(res[i]), args=(shared_args, i), kwargs=kwargs, func=eager_tuple_element) for i in range(len(res)))
+            else:
+                del res  # not needed
+                # non-tensor return likely relies on the contents of the args
+                # (e.g. the result of torch.equal)
+                eager_args = cls.to_eager(args)
+                return fn(*eager_args, **kwargs)
+        return wrapped_fn
+
+    @classmethod
+    def to_eager(cls, t: Any) -> Any:
+        def simple_to_eager(_t: LazyBase) -> Any:
+            if _t._data is not None:
+                return _t._data
+
+            # NOTE: there's a recursion limit in Python (usually 1000)
+
+            assert _t._func is not None
+            _t._args = cls._recurse_apply(_t._args, simple_to_eager)
+            _t._data = _t._func(*_t._args, **_t._kwargs)
+            # sanity check
+            assert _t._data is not None
+            assert _t._data.dtype == _t._meta.dtype
+            assert _t._data.shape == _t._meta.shape
+
+            return _t._data
+
+        # recurse into lists and/or tuples, keeping their structure
+        return cls._recurse_apply(t, simple_to_eager)
+
+    @classmethod
+    def eager_to_meta(cls, t: Any) -> Any:
+        return cls.meta_with_dtype_and_shape(t.dtype, t.shape)
+
+    # must be overridden, meta tensor init is backend-specific
+    @classmethod
+    @abstractmethod
+    def meta_with_dtype_and_shape(cls, dtype: Any, shape: Any) -> Any: pass
+
+    @classmethod
+    def from_eager(cls, t: Any) -> Any:
+        if type(t) is cls:
+            # already lazy
+            return t
+        elif isinstance(t, cls._tensor_type):
+            return cls(meta=cls.eager_to_meta(t), data=t)
+        else:
+            return TypeError(f"{type(t)!r} is not compatible with {cls._tensor_type!r}")
+
+
+class LazyNumpyTensor(LazyBase):
+    _tensor_type = np.ndarray
+
+    shape: tuple[int, ...]  # Makes the type checker happy in quants.py
+
+    @classmethod
+    def meta_with_dtype_and_shape(cls, dtype: DTypeLike, shape: tuple[int, ...]) -> np.ndarray[Any, Any]:
+        # The initial idea was to use np.nan as the fill value,
+        # but non-float types like np.int16 can't use that.
+        # So zero it is.
+        cheat = np.zeros(1, dtype)
+        return np.lib.stride_tricks.as_strided(cheat, shape, (0 for _ in shape))
+
+    def astype(self, dtype, *args, **kwargs):
+        meta = type(self).meta_with_dtype_and_shape(dtype, self._meta.shape)
+        full_args = (self, dtype,) + args
+        return type(self)(meta=meta, args=full_args, kwargs=kwargs, func=(lambda a, *args, **kwargs: a.astype(*args, **kwargs)))
+
+    def tofile(self, *args, **kwargs):
+        eager = LazyNumpyTensor.to_eager(self)
+        return eager.tofile(*args, **kwargs)
+
+    # TODO: __array_function__
+
+
+# Tensor written to file one row-chunk at a time
+class LazyChunkedTensor:
+
+    def __init__(
+        self, chunks: list[Callable[[], np.ndarray]], shape: tuple[int, ...], dtype: DTypeLike,
+        qtype: Any = None, byteswap: bool = False,
+    ):
+        self._chunks = chunks
+        self._qtype = qtype
+        self._byteswap = byteswap
+        self.shape = tuple(shape)
+        self.dtype = np.dtype(dtype)
+
+    @property
+    def nbytes(self) -> int:
+        n = self.dtype.itemsize
+        for d in self.shape:
+            n *= d
+        return n
+
+    def numpy(self) -> LazyChunkedTensor:
+        return self
+
+    def __array__(self, *args, **kwargs):
+        # numpy would otherwise make a 1-element object array of self, and write 8 bytes
+        raise TypeError("LazyChunkedTensor cannot become an ndarray, it is written in chunks")
+
+    def quantize(self, qtype: Any) -> LazyChunkedTensor:
+        from .constants import GGMLQuantizationType
+        from .quants import QuantError, quant_shape_to_byte_shape
+
+        if qtype == GGMLQuantizationType.F32:
+            shape, dtype = self.shape, np.dtype(np.float32)
+        elif qtype == GGMLQuantizationType.F16:
+            shape, dtype = self.shape, np.dtype(np.float16)
+        else:
+            try:
+                shape, dtype = quant_shape_to_byte_shape(self.shape, qtype), np.dtype(np.uint8)
+            except ValueError as e:
+                # raised here and not per chunk, so callers can still fall back to F16
+                raise QuantError(str(e)) from e
+        return LazyChunkedTensor(self._chunks, shape, dtype, qtype, self._byteswap)
+
+    def byteswap(self, inplace: bool = False) -> LazyChunkedTensor:
+        if inplace:
+            raise NotImplementedError("a chunked tensor cannot be byteswapped in place")
+        return LazyChunkedTensor(self._chunks, self.shape, self.dtype, self._qtype, not self._byteswap)
+
+    def tofile(self, *args, **kwargs) -> None:
+        from .quants import quantize
+
+        written = 0
+        for load_chunk in self._chunks:
+            chunk = load_chunk()
+            if self._qtype is not None:
+                # exact only because chunks split on rows, and blocks never cross one
+                chunk = quantize(chunk, self._qtype)
+            if self._byteswap:
+                chunk = chunk.byteswap(inplace=False)
+            chunk.tofile(*args, **kwargs)
+            written += chunk.nbytes
+            del chunk
+        assert written == self.nbytes, f"chunked tensor wrote {written} bytes, expected {self.nbytes}"
diff --git a/python/freetoken/third_party/gguf/quants.py b/python/freetoken/third_party/gguf/quants.py
new file mode 100644
index 0000000000..80966b6ef1
--- /dev/null
+++ b/python/freetoken/third_party/gguf/quants.py
@@ -0,0 +1,1378 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Sequence
+from math import log2, ceil
+
+from numpy.typing import DTypeLike
+
+from .constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K
+from .lazy import LazyNumpyTensor
+
+import numpy as np
+
+
+def quant_shape_to_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
+    block_size, type_size = GGML_QUANT_SIZES[quant_type]
+    if shape[-1] % block_size != 0:
+        raise ValueError(f"Quantized tensor row size ({shape[-1]}) is not a multiple of {quant_type.name} block size ({block_size})")
+    return (*shape[:-1], shape[-1] // block_size * type_size)
+
+
+def quant_shape_from_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
+    block_size, type_size = GGML_QUANT_SIZES[quant_type]
+    if shape[-1] % type_size != 0:
+        raise ValueError(f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of {quant_type.name} type size ({type_size})")
+    return (*shape[:-1], shape[-1] // type_size * block_size)
+
+
+# This is faster than np.vectorize and np.apply_along_axis because it works on more than one row at a time
+def _apply_over_grouped_rows(func: Callable[[np.ndarray], np.ndarray], arr: np.ndarray, otype: DTypeLike, oshape: tuple[int, ...]) -> np.ndarray:
+    rows = arr.reshape((-1, arr.shape[-1]))
+    assert len(rows.shape)
+    osize = 1
+    for dim in oshape:
+        osize *= dim
+    out = np.empty(shape=osize, dtype=otype)
+    # compute over groups of 16 rows (arbitrary, but seems good for performance)
+    n_groups = (rows.shape[0] // 16) or 1
+    np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out)
+    return out.reshape(oshape)
+
+
+# round away from zero
+# ref: https://stackoverflow.com/a/59143326/22827863
+def np_roundf(n: np.ndarray) -> np.ndarray:
+    a = abs(n)
+    floored = np.floor(a)
+    b = floored + np.floor(2 * (a - floored))
+    return np.sign(n) * b
+
+
+class QuantError(Exception): ...
+
+
+_type_traits: dict[GGMLQuantizationType, type[__Quant]] = {}
+
+
+def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
+    if qtype == GGMLQuantizationType.F32:
+        return data.astype(np.float32, copy=False)
+    elif qtype == GGMLQuantizationType.F16:
+        return data.astype(np.float16, copy=False)
+    elif (q := _type_traits.get(qtype)) is not None:
+        return q.quantize(data)
+    else:
+        raise NotImplementedError(f"Quantization for {qtype.name} is not yet implemented")
+
+
+def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
+    if qtype == GGMLQuantizationType.F32:
+        return data.view(np.float32)
+    elif qtype == GGMLQuantizationType.F16:
+        return data.view(np.float16).astype(np.float32)
+    elif (q := _type_traits.get(qtype)) is not None:
+        return q.dequantize(data)
+    else:
+        raise NotImplementedError(f"Dequantization for {qtype.name} is not yet implemented")
+
+
+class __Quant(ABC):
+    qtype: GGMLQuantizationType
+    block_size: int
+    type_size: int
+
+    grid: np.ndarray[Any, np.dtype[np.float32]] | None = None
+    grid_shape: tuple[int, int] = (0, 0)
+    grid_map: tuple[int | float, ...] = ()
+    grid_hex: bytes | None = None
+
+    def __init__(self):
+        return TypeError("Quant conversion classes can't have instances")
+
+    def __init_subclass__(cls, qtype: GGMLQuantizationType) -> None:
+        cls.qtype = qtype
+        cls.block_size, cls.type_size = GGML_QUANT_SIZES[qtype]
+        cls.__quantize_lazy: Any = LazyNumpyTensor._wrap_fn(
+            cls.__quantize_array,
+            meta_noop=(np.uint8, cls.__shape_to_bytes)
+        )
+        cls.__dequantize_lazy: Any = LazyNumpyTensor._wrap_fn(
+            cls.__dequantize_array,
+            meta_noop=(np.float32, cls.__shape_from_bytes)
+        )
+        assert qtype not in _type_traits
+        _type_traits[qtype] = cls
+
+    @classmethod
+    def init_grid(cls):
+        if cls.grid is not None or cls.grid_hex is None:
+            return
+
+        bits_per_elem = ceil(log2(len(cls.grid_map)))
+        assert bits_per_elem != 0, cls.qtype.name
+        elems_per_byte = 8 // bits_per_elem
+
+        grid = np.frombuffer(cls.grid_hex, dtype=np.uint8)
+        # decode hexadecimal chars from grid
+        grid = grid.reshape((-1, 2))
+        grid = (np.where(grid > 0x40, grid + 9, grid) & 0x0F) << np.array([4, 0], dtype=np.uint8).reshape((1, 2))
+        grid = grid[..., 0] | grid[..., 1]
+        # unpack the grid values
+        grid = grid.reshape((-1, 1)) >> np.array([i for i in range(0, 8, 8 // elems_per_byte)], dtype=np.uint8).reshape((1, elems_per_byte))
+        grid = (grid & ((1 << bits_per_elem) - 1)).reshape((-1, 1))
+        grid_map = np.array(cls.grid_map, dtype=np.float32).reshape((1, -1))
+        grid = np.take_along_axis(grid_map, grid, axis=-1)
+        cls.grid = grid.reshape((1, 1, *cls.grid_shape))
+
+    @classmethod
+    @abstractmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        raise NotImplementedError
+
+    @classmethod
+    def quantize_rows(cls, rows: np.ndarray) -> np.ndarray:
+        rows = rows.astype(np.float32, copy=False)
+        shape = rows.shape
+        n_blocks = rows.size // cls.block_size
+        blocks = rows.reshape((n_blocks, cls.block_size))
+        blocks = cls.quantize_blocks(blocks)
+        assert blocks.dtype == np.uint8
+        assert blocks.shape[-1] == cls.type_size
+        return blocks.reshape(cls.__shape_to_bytes(shape))
+
+    @classmethod
+    def dequantize_rows(cls, rows: np.ndarray) -> np.ndarray:
+        rows = rows.view(np.uint8)
+        shape = rows.shape
+        n_blocks = rows.size // cls.type_size
+        blocks = rows.reshape((n_blocks, cls.type_size))
+        blocks = cls.dequantize_blocks(blocks)
+        assert blocks.dtype == np.float32
+        assert blocks.shape[-1] == cls.block_size
+        return blocks.reshape(cls.__shape_from_bytes(shape))
+
+    @classmethod
+    def __shape_to_bytes(cls, shape: Sequence[int]):
+        return quant_shape_to_byte_shape(shape, cls.qtype)
+
+    @classmethod
+    def __shape_from_bytes(cls, shape: Sequence[int]):
+        return quant_shape_from_byte_shape(shape, cls.qtype)
+
+    @classmethod
+    def __quantize_array(cls, array: np.ndarray) -> np.ndarray:
+        return _apply_over_grouped_rows(cls.quantize_rows, arr=array, otype=np.uint8, oshape=cls.__shape_to_bytes(array.shape))
+
+    @classmethod
+    def __dequantize_array(cls, array: np.ndarray) -> np.ndarray:
+        cls.init_grid()
+        return _apply_over_grouped_rows(cls.dequantize_rows, arr=array, otype=np.float32, oshape=cls.__shape_from_bytes(array.shape))
+
+    @classmethod
+    def __quantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
+        pass
+
+    @classmethod
+    def __dequantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
+        pass
+
+    @classmethod
+    def can_quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> bool:
+        return tensor.shape[-1] % cls.block_size == 0
+
+    @classmethod
+    def quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
+        if not cls.can_quantize(tensor):
+            raise QuantError(f"Can't quantize tensor with shape {tensor.shape} to {cls.qtype.name}")
+        if isinstance(tensor, LazyNumpyTensor):
+            return cls.__quantize_lazy(tensor)
+        else:
+            return cls.__quantize_array(tensor)
+
+    @classmethod
+    def dequantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
+        if isinstance(tensor, LazyNumpyTensor):
+            return cls.__dequantize_lazy(tensor)
+        else:
+            return cls.__dequantize_array(tensor)
+
+
+class BF16(__Quant, qtype=GGMLQuantizationType.BF16):
+    @classmethod
+    # same as ggml_compute_fp32_to_bf16 in ggml-impl.h
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n = blocks.view(np.uint32)
+        # force nan to quiet
+        n = np.where((n & 0x7fffffff) > 0x7f800000, (n & np.uint32(0xffff0000)) | np.uint32(64 << 16), n)
+        # round to nearest even
+        n = (np.uint64(n) + (0x7fff + ((n >> 16) & 1))) >> 16
+        return n.astype(np.uint16).view(np.uint8)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32)
+
+
+class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        imax = abs(blocks).argmax(axis=-1, keepdims=True)
+        max = np.take_along_axis(blocks, imax, axis=-1)
+
+        d = max / -8
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        qs = np.trunc((blocks * id) + np.float32(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
+
+        qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
+        qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
+
+        d = d.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([d, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, qs = np.hsplit(blocks, [2])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8)
+
+        return (d * qs.astype(np.float32))
+
+
+class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        max = blocks.max(axis=-1, keepdims=True)
+        min = blocks.min(axis=-1, keepdims=True)
+
+        d = (max - min) / 15
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
+
+        qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
+        qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
+
+        d = d.astype(np.float16).view(np.uint8)
+        m = min.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([d, m, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        m, qs = np.hsplit(rest, [2])
+
+        d = d.view(np.float16).astype(np.float32)
+        m = m.view(np.float16).astype(np.float32)
+
+        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32)
+
+        return (d * qs) + m
+
+
+class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        imax = abs(blocks).argmax(axis=-1, keepdims=True)
+        max = np.take_along_axis(blocks, imax, axis=-1)
+
+        d = max / -16
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        q = np.trunc((blocks * id) + np.float32(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
+
+        qs = q.reshape((n_blocks, 2, cls.block_size // 2))
+        qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
+
+        qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
+
+        d = d.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([d, qh, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qh, qs = np.hsplit(rest, [4])
+
+        d = d.view(np.float16).astype(np.float32)
+        qh = qh.view(np.uint32)
+
+        qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
+        ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qh = (qh & np.uint32(0x01)).astype(np.uint8)
+        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
+
+        qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16)
+
+        return (d * qs.astype(np.float32))
+
+
+class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        max = blocks.max(axis=-1, keepdims=True)
+        min = blocks.min(axis=-1, keepdims=True)
+
+        d = (max - min) / 31
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
+
+        qs = q.reshape((n_blocks, 2, cls.block_size // 2))
+        qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
+
+        qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
+
+        d = d.astype(np.float16).view(np.uint8)
+        m = min.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([d, m, qh, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        m, rest = np.hsplit(rest, [2])
+        qh, qs = np.hsplit(rest, [4])
+
+        d = d.view(np.float16).astype(np.float32)
+        m = m.view(np.float16).astype(np.float32)
+        qh = qh.view(np.uint32)
+
+        qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
+        ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qh = (qh & np.uint32(0x01)).astype(np.uint8)
+        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
+
+        qs = (ql | (qh << np.uint8(4))).astype(np.float32)
+
+        return (d * qs) + m
+
+
+class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0):
+    @classmethod
+    # Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+
+        d = abs(blocks).max(axis=1, keepdims=True) / 127
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        qs = np_roundf(blocks * id)
+
+        # (n_blocks, 2)
+        d = d.astype(np.float16).view(np.uint8)
+        # (n_blocks, block_size)
+        qs = qs.astype(np.int8).view(np.uint8)
+
+        return np.concatenate([d, qs], axis=1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        d, x = np.split(blocks, [2], axis=1)
+        d = d.view(np.float16).astype(np.float32)
+        x = x.view(np.int8).astype(np.float32)
+
+        return (x * d)
+
+
+class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K):
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        scales, rest = np.hsplit(blocks, [QK_K // 16])
+        qs, rest = np.hsplit(rest, [QK_K // 4])
+        d, dmin = np.hsplit(rest, [2])
+
+        d = d.view(np.float16).astype(np.float32)
+        dmin = dmin.view(np.float16).astype(np.float32)
+
+        # (n_blocks, 16, 1)
+        dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
+        ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
+
+        shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
+
+        qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3)
+
+        qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32)
+
+        qs = dl * qs - ml
+
+        return qs.reshape((n_blocks, -1))
+
+
+class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K):
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        hmask, rest = np.hsplit(blocks, [QK_K // 8])
+        qs, rest = np.hsplit(rest, [QK_K // 4])
+        scales, d = np.hsplit(rest, [12])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        # The scales are packed at 6-bit each in this pattern:
+        #  0: IIIIAAAA
+        #  1: JJJJBBBB
+        #  2: KKKKCCCC
+        #  3: LLLLDDDD
+        #  4: MMMMEEEE
+        #  5: NNNNFFFF
+        #  6: OOOOGGGG
+        #  7: PPPPHHHH
+        #  8: MMIIEEAA
+        #  9: NNJJFFBB
+        # 10: OOKKGGCC
+        # 11: PPLLHHDD
+        lscales, hscales = np.hsplit(scales, [8])
+        lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
+        lscales = lscales.reshape((n_blocks, 16))
+        hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1))
+        hscales = hscales.reshape((n_blocks, 16))
+        scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4))
+        scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32)
+
+        dl = (d * scales).reshape((n_blocks, 16, 1))
+
+        ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
+        ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3)
+        qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1))
+        qh = qh ^ np.uint8(1)  # strangely, the offset is zero when the bitmask is 1
+        q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32)
+
+        return (dl * q).reshape((n_blocks, QK_K))
+
+
+class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K):
+    K_SCALE_SIZE = 12
+
+    @staticmethod
+    def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        n_blocks = scales.shape[0]
+        scales = scales.view(np.uint8)
+        ### Unpacking the following: ###
+        #  0 EEAAAAAA
+        #  1 FFBBBBBB
+        #  2 GGCCCCCC
+        #  3 HHDDDDDD
+        #  4 eeaaaaaa
+        #  5 ffbbbbbb
+        #  6 ggcccccc
+        #  7 hhdddddd
+        #  8 eeeeEEEE
+        #  9 ffffFFFF
+        # 10 ggggGGGG
+        # 11 hhhhHHHH
+        scales = scales.reshape((n_blocks, 3, 4))
+        d, m, m_d = np.split(scales, 3, axis=-2)
+
+        sc = np.concatenate([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], axis=-1)
+        min = np.concatenate([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], axis=-1)
+
+        return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8)))
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        dmin, rest = np.hsplit(rest, [2])
+        scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE])
+
+        d = d.view(np.float16).astype(np.float32)
+        dmin = dmin.view(np.float16).astype(np.float32)
+
+        sc, m = Q4_K.get_scale_min(scales)
+
+        d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
+        dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
+
+        qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 32)).astype(np.float32)
+
+        return (d * qs - dm).reshape((n_blocks, QK_K))
+
+
+class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K):
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        dmin, rest = np.hsplit(rest, [2])
+        scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE])
+        qh, qs = np.hsplit(rest, [QK_K // 8])
+
+        d = d.view(np.float16).astype(np.float32)
+        dmin = dmin.view(np.float16).astype(np.float32)
+
+        sc, m = Q4_K.get_scale_min(scales)
+
+        d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
+        dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
+
+        ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
+        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
+        qh = (qh & np.uint8(0x01)).reshape((n_blocks, -1, 32))
+        q = (ql | (qh << np.uint8(4))).astype(np.float32)
+
+        return (d * q - dm).reshape((n_blocks, QK_K))
+
+
+class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K):
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        ql, rest = np.hsplit(blocks, [QK_K // 2])
+        qh, rest = np.hsplit(rest, [QK_K // 4])
+        scales, d = np.hsplit(rest, [QK_K // 16])
+
+        scales = scales.view(np.int8).astype(np.float32)
+        d = d.view(np.float16).astype(np.float32)
+        d = (d * scales).reshape((n_blocks, QK_K // 16, 1))
+
+        ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
+        qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32))
+        q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32)
+        q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32)
+
+        return (d * q).reshape((n_blocks, QK_K))
+
+
+class TQ1_0(__Quant, qtype=GGMLQuantizationType.TQ1_0):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d = abs(blocks).max(axis=-1, keepdims=True)
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        qs = np_roundf(blocks * id)
+        qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
+
+        qs0, qs1, qh = qs[..., :(32 * 5)], qs[..., (32 * 5):(48 * 5)], qs[..., (48 * 5):]
+        qs0 = qs0.reshape((n_blocks, -1, 5, 32)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
+        qs0 = np.sum(qs0, axis=-2).reshape((n_blocks, -1))
+        qs1 = qs1.reshape((n_blocks, -1, 5, 16)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
+        qs1 = np.sum(qs1, axis=-2).reshape((n_blocks, -1))
+        qh = qh.reshape((n_blocks, -1, 4, 4)) * np.array([81, 27, 9, 3], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qh = np.sum(qh, axis=-2).reshape((n_blocks, -1))
+        qs = np.concatenate([qs0, qs1, qh], axis=-1)
+        qs = (qs.astype(np.uint16) * 256 + (243 - 1)) // 243
+
+        qs = qs.astype(np.uint8)
+        d = d.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([qs, d], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        qs, rest = np.hsplit(blocks, [(QK_K - 4 * QK_K // 64) // 5])
+        qh, d = np.hsplit(rest, [QK_K // 64])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        qs0, qs1 = qs[..., :32], qs[..., 32:]
+        qs0 = qs0.reshape((n_blocks, -1, 1, 32)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
+        qs0 = qs0.reshape((n_blocks, -1))
+        qs1 = qs1.reshape((n_blocks, -1, 1, 16)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
+        qs1 = qs1.reshape((n_blocks, -1))
+        qh = qh.reshape((n_blocks, -1, 1, 4)) * np.array([1, 3, 9, 27], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qh = qh.reshape((n_blocks, -1))
+        qs = np.concatenate([qs0, qs1, qh], axis=-1)
+        qs = ((qs.astype(np.uint16) * 3) >> 8).astype(np.int8) - np.int8(1)
+
+        return (d * qs.astype(np.float32))
+
+
+class TQ2_0(__Quant, qtype=GGMLQuantizationType.TQ2_0):
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d = abs(blocks).max(axis=-1, keepdims=True)
+        with np.errstate(divide="ignore"):
+            id = np.where(d == 0, 0, 1 / d)
+        qs = np_roundf(blocks * id)
+        qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
+
+        qs = qs.reshape((n_blocks, -1, 4, 32)) << np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qs = qs[..., 0, :] | qs[..., 1, :] | qs[..., 2, :] | qs[..., 3, :]
+        qs = qs.reshape((n_blocks, -1))
+
+        d = d.astype(np.float16).view(np.uint8)
+
+        return np.concatenate([qs, d], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        qs, d = np.hsplit(blocks, [QK_K // 4])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
+        qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1)
+
+        return (d * qs.astype(np.float32))
+
+
+class MXFP4(__Quant, qtype=GGMLQuantizationType.MXFP4):
+    # e2m1 values (doubled)
+    # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+    kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12)
+
+    @staticmethod
+    # see ggml_e8m0_to_fp32_half in ggml-impl.h
+    def e8m0_to_fp32_half(x: np.ndarray) -> np.ndarray:
+        bits = np.where(x < 2, np.uint32(0x00200000) << np.uint32(x), np.uint32(x - 1) << np.uint32(23))
+        return bits.view(np.float32)
+
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d = abs(blocks).max(axis=-1, keepdims=True)
+
+        with np.errstate(divide="ignore"):
+            e = np.where(d > 0, np.floor(np.log2(d)) - 2 + 127, 0).astype(np.uint8)
+
+        d = cls.e8m0_to_fp32_half(e)
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape((1, 1, 16))
+
+        errs = np.abs(d.reshape((n_blocks, 1, 1)) * kvalues.astype(np.float32) - blocks.reshape((n_blocks, cls.block_size, 1)))
+        best = np.argmin(errs, axis=-1, keepdims=True)
+
+        qs = best.reshape(n_blocks, 2, cls.block_size // 2).astype(np.uint8)
+        qs = qs[:, 0] | (qs[:, 1] << np.uint8(4))
+
+        qs = qs.reshape((n_blocks, cls.block_size // 2))
+
+        return np.concatenate([e, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        e, qs = np.hsplit(blocks, [1])
+
+        d = cls.e8m0_to_fp32_half(e)
+
+        qs = qs.reshape((n_blocks, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
+        qs = (qs & np.uint8(0x0F)).view(np.int8)
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
+        qs = np.take_along_axis(kvalues, qs, axis=-1).reshape((n_blocks, cls.block_size))
+
+        return (d * qs.astype(np.float32))
+
+
+class NVFP4(__Quant, qtype=GGMLQuantizationType.NVFP4):
+    # E2M1 values doubled (kvalues_mxfp4 convention)
+    kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12)
+
+    @staticmethod
+    def ue4m3_to_fp32(x: np.ndarray) -> np.ndarray:
+        """Decode unsigned E4M3 (bias=7) to float, with 0.5 factor for kvalues convention."""
+        exp = (x >> 3).astype(np.int32) & 0xF
+        man = (x & 0x7).astype(np.float32)
+        raw = np.where(
+            exp == 0,
+            man * 2**-9,
+            (1.0 + man / 8.0) * (2.0 ** (exp.astype(np.float32) - 7)))
+        return np.where((x == 0) | (x == 0x7F), 0.0, raw * 0.5)
+
+    @staticmethod
+    def fp32_to_ue4m3(x: np.ndarray) -> np.ndarray:
+        """Vectorized float32 to unsigned E4M3, matching ggml_fp32_to_ue4m3 in C."""
+        x = np.clip(x, 0.0, 448.0).astype(np.float32)
+        bits = x.view(np.uint32)
+        fp32_exp = ((bits >> 23) & 0xFF).astype(np.int32) - 127
+        fp32_man = ((bits >> 20) & 0x7).astype(np.int32)
+        ue4m3_exp = fp32_exp + 7
+
+        # Subnormal
+        sub_man = np.clip((x * 512.0 + 0.5).astype(np.int32), 0, 7)
+        sub_result = np.where(sub_man >= 1, sub_man, 0).astype(np.uint8)
+
+        # Normal with rounding
+        round_bit = ((bits >> 19) & 1).astype(np.int32)
+        man = fp32_man + round_bit
+        exp = ue4m3_exp.copy()
+        overflow = man > 7
+        man = np.where(overflow, 0, man)
+        exp = np.where(overflow, exp + 1, exp)
+        normal_result = np.where(exp >= 15, np.uint8(0x7E), ((exp << 3) | man).astype(np.uint8))
+
+        return np.where(x <= 0.0, np.uint8(0),
+                        np.where(ue4m3_exp <= 0, sub_result,
+                        np.where(ue4m3_exp >= 15, np.uint8(0x7E), normal_result)))
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_super = blocks.shape[0]
+
+        d_bytes, qs = np.hsplit(blocks, [4])
+        d = cls.ue4m3_to_fp32(d_bytes).reshape(n_super, 4, 1)  # (n_super, 4, 1)
+
+        qs = qs.reshape(n_super, 4, 8)
+        lo = (qs & np.uint8(0x0F)).view(np.int8)
+        hi = (qs >> np.uint8(4)).view(np.int8)
+        vals = np.concatenate([lo, hi], axis=-1)  # (n_super, 4, 16)
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
+        vals = np.take_along_axis(kvalues, vals, axis=-1)
+
+        return (d * vals.astype(np.float32)).reshape(n_super, 64)
+
+
+class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS):
+    ksigns: bytes = (
+        b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f"
+        b"\x90\x11\x12\x93\x14\x95\x96\x17\x18\x99\x9a\x1b\x9c\x1d\x1e\x9f"
+        b"\xa0\x21\x22\xa3\x24\xa5\xa6\x27\x28\xa9\xaa\x2b\xac\x2d\x2e\xaf"
+        b"\x30\xb1\xb2\x33\xb4\x35\x36\xb7\xb8\x39\x3a\xbb\x3c\xbd\xbe\x3f"
+        b"\xc0\x41\x42\xc3\x44\xc5\xc6\x47\x48\xc9\xca\x4b\xcc\x4d\x4e\xcf"
+        b"\x50\xd1\xd2\x53\xd4\x55\x56\xd7\xd8\x59\x5a\xdb\x5c\xdd\xde\x5f"
+        b"\x60\xe1\xe2\x63\xe4\x65\x66\xe7\xe8\x69\x6a\xeb\x6c\xed\xee\x6f"
+        b"\xf0\x71\x72\xf3\x74\xf5\xf6\x77\x78\xf9\xfa\x7b\xfc\x7d\x7e\xff"
+    )
+
+    # iq2xxs_grid, but with each byte of the original packed in 2 bits,
+    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
+    grid_shape = (256, 8)
+    grid_map = (0x08, 0x19, 0x2b)
+    grid_hex = (
+        b"00000200050008000a00110014002000220028002a0041004400500058006100"
+        b"6400800082008a00a20001010401100115014001840198010002020222028202"
+        b"010404041004210424044004420448046004810484049004a404000502050805"
+        b"200546056905800591050906100640068406a406000805080808140828084108"
+        b"440850085208880804094009020a140a01100410101021104010601084109010"
+        b"951000110811201150115a118011241245120014081420142514491480141815"
+        b"6215001616160118041810184018811800190519a019511a002002200a204420"
+        b"6120802082202921482100220222012404241024402456240025412564259026"
+        b"082820289428442a014004401040184021402440404048405640604081408440"
+        b"9040004120416141804185410142104248425642684200440844204480449944"
+        b"124524450046014804481048404845480049584961498249454a904a00500850"
+        b"1150195020508050885004514251a4519152905492540a550156545600581158"
+        b"195864584059085a046010604060686000615561186260620064056410651265"
+        b"84654268008002800a8041808280048118814081118201840484108415844084"
+        b"608400854685948509864086608602880489118a0490109024904090a1901691"
+        b"8091459200942294449451958198209902a050a085a009a100a218a450a804a9"
+    )
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, qs = np.hsplit(blocks, [2])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        qs = qs.view(np.uint32).reshape(n_blocks, -1, 2)
+
+        db = d * (np.float32(0.5) + (qs[..., 1] >> 28).astype(np.float32)) * np.float32(0.25)
+        db = db.reshape((n_blocks, -1, 1, 1))
+
+        # get the sign indices and unpack the bits
+        signs = qs[..., 1].reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
+        ksigns = np.frombuffer(cls.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
+        signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
+        signs = np.take_along_axis(ksigns, signs, axis=-1)
+        signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8))
+        signs = signs & np.uint8(0x01)
+        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
+        signs = signs.reshape((n_blocks, -1, 4, 8))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs[..., 0].copy().view(np.uint8).reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 4, 8))
+
+        return (db * grid * signs).reshape((n_blocks, -1))
+
+
+class IQ2_XS(__Quant, qtype=GGMLQuantizationType.IQ2_XS):
+    # iq2xs_grid, but with each byte of the original packed in 2 bits,
+    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
+    grid_shape = (512, 8)
+    grid_map = (0x08, 0x19, 0x2b)
+    grid_hex = (
+        b"00000200050008000a0011001400160019002000220025002800410044004600"
+        b"49005000520055005800610064008000820085008800910094009900a0000101"
+        b"04010601090110011201150118011a0121012401400142014501480151015401"
+        b"6001680181018401900100020202050208021102140220024102440250025502"
+        b"80028a0201040404060409041004120415041804210424044004420445044804"
+        b"5104540456046004810484049004000502050505080511051405200541054405"
+        b"500561058005010604061006260640064206840600080208050808080a081108"
+        b"14082008250841084408500858088008a008aa08010904091009400981098909"
+        b"000a200a280a960aa00a01100410061009101010121015101810211024104010"
+        b"4210451048105110541060106a10811084109010001102110511081111111411"
+        b"2011411144115011801194119611011204120612101240126012001402140514"
+        b"0814111414142014411444144914501464148014011504151015401500161416"
+        b"49160118041810181218401854188618001905196619511aa91a002002200520"
+        b"08200a201120142020204120442050208020a020012104211021402148216521"
+        b"002222228022a82201240424102429244024002541255225992501261a26a626"
+        b"002808280a28202855288828a22868299029082a202a822a882a8a2a01400440"
+        b"0640094010401240154018402140244040404240454048404a40514054406040"
+        b"6540814084409040004102410541084111411441204141414441504180418541"
+        b"a241014204421042124229424042004402440544084411441444194420444144"
+        b"4444504480449444014504451045244540459a4500460a464446504601480448"
+        b"1048404845485448624800491149444950496949044a00500250055008501150"
+        b"145020502850415044505050805001510451105115514051425100524452aa52"
+        b"0154045410542154405460548154a154005508558055885521566856a1560058"
+        b"14584158505899581a5940594259855a0160046010604060546062608660a960"
+        b"006124624a62926200641664106540654565a46501686a682569066a546a626a"
+        b"00800280058008801180148020802a8041804480508080808280a880aa800181"
+        b"0481068110814081518159810082208280828282a082a8820184048410841284"
+        b"158440846084898400854485a58518866a860088088825885a8880888288a888"
+        b"0689228a808a888a968aa88a0190049010904090569084900091229164915692"
+        b"89920094059444945094589429959095929541965198a6984999159a609a00a0"
+        b"02a008a00aa020a02aa0a0a051a159a1a6a100a202a208a22aa280a2a0a240a4"
+        b"95a465a698a60aa820a822a828a8a0a8a8a804a984a986a928aa2aaa91aaaaaa"
+    )
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qs, scales = np.hsplit(rest, [2 * QK_K // 8])
+
+        d = d.view(np.float16).astype(np.float32)
+        qs = qs.view(np.uint16)
+
+        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
+        scales = (scales & 0x0F).reshape((n_blocks, -1))
+        db = d * (np.float32(0.5) + scales) * np.float32(0.25)
+        db = db.reshape((n_blocks, -1, 1, 1))
+
+        # get the sign indices and unpack the bits
+        signs = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape(1, 1, 128)
+        signs = np.take_along_axis(signs, (qs >> 9).reshape((n_blocks, -1, 1)), axis=-1)
+        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
+        signs = signs & np.uint8(0x01)
+        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
+        signs = signs.reshape((n_blocks, -1, 2, 8))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, (qs & np.uint16(511)).reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 2, 8))
+
+        return (db * grid * signs).reshape((n_blocks, -1))
+
+
+class IQ2_S(__Quant, qtype=GGMLQuantizationType.IQ2_S):
+    # iq2s_grid, but with each byte of the original packed in 2 bits,
+    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
+    grid_shape = (1024, 8)
+    grid_map = (0x08, 0x19, 0x2b)
+    grid_hex = (
+        b"00000200050008000a0011001400160019002000220025002800410044004600"
+        b"490050005200550058006100640066006900800082008500880091009400a000"
+        b"a500aa0001010401060109011001120115011801210124014001420145014801"
+        b"510154015601590160016501680181018401900192019501a101a40100020202"
+        b"050208021102140220022a02410244024602490250025502800285028a029402"
+        b"a202010404040604090410041204150418042104240426042904400442044504"
+        b"48044a0451045404560459046004620465048104840486048904900495049804"
+        b"a104a40400050205050508050a05110514051605190520052505280541054405"
+        b"46054905500552055505580561056405800582058505880591059405a0050106"
+        b"0406060609061006150640064506480651065406600681068406900600080208"
+        b"050808081108140816081908200825082a084108440846084908500852085508"
+        b"580861086408800885089408aa08010904091009120915091809210940094509"
+        b"480951095409600981099009000a110a140a220a280a2a0a500a990a01100410"
+        b"0610091010101210151018102110241026104010421045104810511054105610"
+        b"59106010621065106810811084108610901095109810a110a410001102110511"
+        b"08110a1111111411161119112011221125112811411144114611491150115211"
+        b"5511581161116411801182118511881191119411011204120912101215122112"
+        b"2412401245125112541281128412901200140214051408141114141416141914"
+        b"2014251428144114441446144914501452145514581461146414801482148514"
+        b"881491149414a014011504150615091510151215151518152115241540154215"
+        b"4515481551155415601581158415901500160516081611161416201641164416"
+        b"50168016aa160118041806180918101815181818211840184218451848185118"
+        b"541860188118841800190219051908191119141920194119441950196919a219"
+        b"041a101a401a561a00200220052008201120142016201920202025202a204120"
+        b"4420502052205520642080208a209420aa200121042110211221152121214021"
+        b"4221452151215421602181218421902100220a22222228222a22442250228822"
+        b"8a22a82201240424062409241024152418242124242440244224452448245124"
+        b"5424602481248424902400250525082511251425202541254425502566258025"
+        b"0126042610264026592600280528112814284128442850288a28aa2801290429"
+        b"102995290a2a222a642a882a8a2a014004400640094010401240154018401a40"
+        b"21402440264040404240454048404a4051405440564059406040624065408140"
+        b"8440904095409840a140a4400041024105410841114114411641194120412241"
+        b"2541414144414641494150415241554158416141644180418241854188419141"
+        b"9441a04101420442104212421542184224424042454248425142544260428142"
+        b"844200440244054408440a441144144416441944204422442544284441444444"
+        b"46444944504452445544584461446444804482448544884491449444a0440145"
+        b"0445064509451045124515451845214524454045424545454845514554456045"
+        b"6a4581458445904500460246054608461146144620464146444650468046a546"
+        b"0148044809481048124815481848214824484048424845484848514854486048"
+        b"84489048004902490549084911491449204941494449504980499649014a044a"
+        b"104a404a00500250055008501150145016501950205022502550285041504450"
+        b"4650495050505250555058506150645080508250855088509150945001510451"
+        b"0651095110511251155118512151245140514251455148515151545160518151"
+        b"8451905100520552085211521452205241524452505269528052015404540654"
+        b"0954105412541554185421542454405442544554485451545454605481548454"
+        b"9054005502550555085511551455205541554455505580550156045610562656"
+        b"405600580258055808581158145820584158445850585a588058015904591059"
+        b"4059005a195a855aa85a01600460066010601260156018602160246040604560"
+        b"4860516054606060846090600061026105610861116114612061416144615061"
+        b"806199610462106240625662a162006405640864116414642064416444645064"
+        b"806401650465106540654a656865926500669466016804681068656898680069"
+        b"2a69426aa16a0080028005800880118014801980208025804180448050805280"
+        b"5580588061808080858091809480018104810981108112811581188121812481"
+        b"408142814581488151815481818184819081a981008205820a82118214824182"
+        b"4482508201840484068409841084128415841884218440844284458448845184"
+        b"5484608481848484908400850285058508851185148520854185448550858085"
+        b"8a85018604861086298640860088058811881488418844885088a28801890489"
+        b"40896589228a588a5a8a828aa28a019004900990109012901590189024904090"
+        b"4290459048905190549060908190849090900091059111911491419144915091"
+        b"5a910192049210924092a6920094029405940894119414942094419444945094"
+        b"8094969401950495109540959895a19500964696649601980498109826984098"
+        b"a998009949995299909a00a005a00aa014a022a02aa041a044a050a0a2a0aaa0"
+        b"40a165a102a20aa222a228a22aa282a288a28aa2a8a201a404a410a440a489a4"
+        b"a4a400a519a551a60aa828a8a2a854a986a908aa0aaa20aa22aa28aa88aaaaaa"
+    )
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qs, rest = np.hsplit(rest, [QK_K // 8])
+        signs, rest = np.hsplit(rest, [QK_K // 8])
+        qh, scales = np.hsplit(rest, [QK_K // 32])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
+        scales = (scales & 0x0F).reshape((n_blocks, -1))
+        db = d * (np.float32(0.5) + scales) * np.float32(0.25)
+        db = db.reshape((n_blocks, -1, 1, 1))
+
+        # unpack the sign bits
+        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
+        signs = signs & np.uint8(0x01)
+        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
+        signs = signs.reshape((n_blocks, -1, 2, 8))
+
+        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4))
+        qs = qs.astype(np.uint16) | ((qh & 0x03).astype(np.uint16) << 8).reshape((n_blocks, -1))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 2, 8))
+
+        return (db * grid * signs).reshape((n_blocks, -1))
+
+
+class IQ3_XXS(__Quant, qtype=GGMLQuantizationType.IQ3_XXS):
+    grid_shape = (256, 4)
+    grid_map = (0x04, 0x0c, 0x14, 0x1c, 0x24, 0x2c, 0x34, 0x3e)
+    grid_hex = (
+        b"0000020004001100130017002000220031004200730075000101030110011201"
+        b"2101250130013201410154017001000202020402110220022202310233023702"
+        b"5102570275020103070310031203250370031304370444045704730475040105"
+        b"0705320552053506640610071407160743076107011003101010121021102310"
+        b"3010321034104710501000110211111120112211011203121012121221123012"
+        b"7212001302132013311346136613011405145014201524154615711505162217"
+        b"4017002002201120132020202220262031204220012103210521102112212121"
+        b"3021632167217021002202221122172220222222372240225522012310231423"
+        b"7023742335245324032527254125742501270327162745270130103012302130"
+        b"2330503065307230003102312031313144314631013203321032253252327232"
+        b"1133333330344734723400350635223555351436363663363337603704401740"
+        b"3540374053405740744120423742404260426642074345430444514464442545"
+        b"4345704505471047124730471250415070500051065126515551145232527252"
+        b"0253535310542354275472540255315550562457425724604460466064602161"
+        b"6161176264623063366344640565526533660367216703700570077010703270"
+        b"5270267140711272457252720073157333736073217441740075027524753076"
+    )
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qs, scales = np.hsplit(rest, [QK_K // 4])
+
+        d = d.view(np.float16).astype(np.float32)
+        scales = scales.view(np.uint32)
+
+        db = d * (np.float32(0.5) + (scales >> 28).astype(np.float32)) * np.float32(0.5)
+        db = db.reshape((n_blocks, -1, 1, 1))
+
+        # get the sign indices and unpack the bits
+        signs = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
+        ksigns = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
+        signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
+        signs = np.take_along_axis(ksigns, signs, axis=-1)
+        signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8))
+        signs = signs & np.uint8(0x01)
+        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
+        signs = signs.reshape((n_blocks, -1, 4, 8))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 4, 8))
+
+        return (db * grid * signs).reshape((n_blocks, -1))
+
+
+class IQ3_S(__Quant, qtype=GGMLQuantizationType.IQ3_S):
+    grid_shape = (512, 4)
+    grid_map = (0x01, 0x03, 0x05, 0x07, 0x09, 0x0b, 0x0d, 0x0f)
+    grid_hex = (
+        b"0000010002000500070010001100120014001600200021002500330040004200"
+        b"4500470051005300600062007100740077000001010102010401100111011501"
+        b"2001230127013101350144016101650172010002010205020702100213021602"
+        b"2102250230023402420245024702510253027002730203031103150320032203"
+        b"3103330336034403500352036703710375030004130417042104240432044004"
+        b"4304510470040205040520052205260533054105450547056605730506061106"
+        b"1306310652067106000702070407200722072607330750075407001001100210"
+        b"0410101011101310151017102010221031103410361054105610611072100011"
+        b"0111031106111011141121113011331141115011521170117611001212121512"
+        b"1712201224123212401243125512601272120113041307131013131321132713"
+        b"3013341341136213701303140514121414143114331442144614501454140115"
+        b"1015131521153015321551152016241627164416461601170317101712172117"
+        b"3517411762177017002001200320052007201020122014201620212023202720"
+        b"3020322041204320452050205220672070207320752000210221102113211721"
+        b"2221252131213421422151210122042207222122232230223722412253225722"
+        b"7122742200230223052311232223242331233323422350236623012407242024"
+        b"2324322435244124722475240425112522253725402553257025002602260726"
+        b"2126552661260527112726273027432750270230113013301530173022303130"
+        b"3330353042304430473051306330713001310331053114312131233140316031"
+        b"7231763100321232203232323432503201331033143321332333273330334133"
+        b"4333473355337333033411341634223431345234603464340135103512352535"
+        b"3235443556357335163641360137033720372237353700400440124020402440"
+        b"2740324041405040704002410741114113412241304135414341514155410142"
+        b"0342104215422142334240425742624270420443114313432043224331433543"
+        b"0044024424443744404471440545074521456245134634466046104715473047"
+        b"4347514702501050145022504050445047505250665074500151035105511251"
+        b"2151325172510052115223523052365253520253075310532753445351536553"
+        b"7353015404542054325446541255265551555355425602570457225711601360"
+        b"1560316033606060006120612761646112623462426255626262706200631463"
+        b"2163406325644364626400650365346560650566406611671367007004700770"
+        b"2070227036704070547062700271117124714371457101720472107216722172"
+        b"3072517202733273357353730174057413742074507422754275027631760077"
+    )
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qs, rest = np.hsplit(rest, [QK_K // 4])
+        qh, rest = np.hsplit(rest, [QK_K // 32])
+        signs, scales = np.hsplit(rest, [QK_K // 8])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
+        scales = (scales & 0x0F).reshape((n_blocks, -1))
+        db = d * (1 + 2 * scales)
+        db = db.reshape((n_blocks, -1, 1, 1))
+
+        # unpack the sign bits
+        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
+        signs = signs & np.uint8(0x01)
+        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
+        signs = signs.reshape((n_blocks, -1, 4, 8))
+
+        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8)
+        qh = (qh & 0x01).astype(np.uint16).reshape((n_blocks, -1))
+        qs = qs.astype(np.uint16) | (qh << 8)
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 4, 8))
+
+        return (db * grid * signs).reshape((n_blocks, -1))
+
+
+class IQ1_S(__Quant, qtype=GGMLQuantizationType.IQ1_S):
+    # iq1s_grid, with each byte packed into 2 bits
+    # -1, 0, 1 <=> 0, 1, 2
+    grid_shape = (2048, 8)
+    grid_map = (-1, 0, 1)
+    grid_hex = (
+        b"00000200050008000a00110015002000220028002a0045005100540056006500"
+        b"8000820088008a009500a000a200a800aa000401050111011401160119011a01"
+        b"2501410146014901520155015a0161016401660168018501910194019601a501"
+        b"0002020208020a0215022002220228022a024502510259026402690280028202"
+        b"88028a02910295029902a002a202a802aa021104140416042504410449045504"
+        b"5a046404650491049904a5040105040505050605150518051a05290540054505"
+        b"4a0550055105540555055605590560056205650568056a058105910595059805"
+        b"9a05a105a405a505a605a9051406190641064406500652065506580660066106"
+        b"6606690685069106940699060008020808080a0815082008220828082a084508"
+        b"5108560865088008820888088a089508a008a208a808aa080509110914091909"
+        b"2409250941095009510955096109640969099109940996099909a509000a020a"
+        b"080a0a0a150a200a220a280a2a0a450a510a590a610a650a800a820a850a880a"
+        b"8a0a950aa00aa20aa80aaa0a1010111014101910241025104110441050105510"
+        b"58106110641065106910911094109610a110a510011104110611091110111211"
+        b"1511181121112411291145114a11501151115211541155115611591160116511"
+        b"841192119511a111a41111121412161225124012461249125212551258125a12"
+        b"641266128512911294129612a512011406140914141415141814191421142614"
+        b"41144514461448144a1451145414551456145914621465146814841489149014"
+        b"94149514981499149a14a114a414a514a914021505150a151115141515151615"
+        b"191520152215251528152a154115441545154615511552155415551556155915"
+        b"5a1561156415651566156915801582158415851588158a159015911594159515"
+        b"961599159a15a015a215a51501160416051606161516161618161a1621162616"
+        b"401642164416451648164a165116551656165816591661166416651668166916"
+        b"6a1686168a1692169516a416a916111816182518411844184618491850185518"
+        b"58185a1860186118641866186918851891189418a5181019121915191a192119"
+        b"25194219441945194819511954195519561959195a19601965196a1989199119"
+        b"921995199819a119a619a919091a161a241a261a441a461a491a501a521a551a"
+        b"581a611a661a691a851a911a961a9a1a0020022008200a201520202022202520"
+        b"28202a20452051205920612065208020822088208a209520a020a220a520a820"
+        b"aa2005211121142119212521422144214921552158215a216121642165216621"
+        b"8521902196219921a521012208220a22112215222022222228222a2245225122"
+        b"562259226522812288228a2291229522a022a222a822aa220524142416241924"
+        b"252444244524462449245224552458245a2466248524912494249924a124a524"
+        b"0925152521252925402545254825512554255525592562256525682589259025"
+        b"9425952598259a25a125a425a625a92505261026122619262526412649265526"
+        b"6026612669268426862690269a260028022808280a2815282028222828282a28"
+        b"45285128542865288028822888288a28a028a228a828aa280929112914291929"
+        b"2529462949295229552961296429662969298529902996299929a429a529002a"
+        b"022a082a0a2a202a222a282a2a2a452a512a562a592a652a802a822a882a8a2a"
+        b"952aa02aa22aa82aaa2a054011401640254049405240554058405a4061406440"
+        b"664094409940a140a6400041014104410641094112411541164118411a412141"
+        b"26412941454148414a41514154415541564159415a41654168416a4181418441"
+        b"8641904192419541a041a141a241054211421442164225424142524255425a42"
+        b"6442694289429442a5420144154419442944454448444a445144544455445644"
+        b"61446244654468446a44814486448944904492449544a044a144a94401450245"
+        b"05450a4511451445154516451945204525452a45414544454545464549455045"
+        b"5145544555455645584559456145644565456645694582458445854588459145"
+        b"94459545964599459a45a545a845aa450146054609461446154618461a462146"
+        b"2446294640464246454648465046514652465546564659466246654668468146"
+        b"85468a4694469546a146a446a6460548114815481a4825484248494850485548"
+        b"5848614864486648694885489148944896489948a5480149054906490a491049"
+        b"144915491849214924492649404945494a495149524954495549564959496049"
+        b"6249654966496a49864989499249954996499849a149a449a649a949164a444a"
+        b"464a494a554a584a5a4a644a694a944aa54a0150045005500650095012501550"
+        b"1a50215024502950405045504850515054505550565059506550685086508950"
+        b"95509850a050a150a650a9500551085109510a51115114511551165118511951"
+        b"20512551265128512a5141514451455146514951505151515251545155515651"
+        b"585159515a51615164516551665169518251855191519451955196519951a051"
+        b"a551aa5101520652125215521a5221522452425245524a525152545255525652"
+        b"595262526552855290529252955299529a52a452045405541154145415541654"
+        b"185419542154255428542a54415444544554465449544a545054515454545554"
+        b"5654585459545a54615462546454655466546954805488548a54915494549554"
+        b"96549954a154a454a554aa540155025504550555065509551055115512551455"
+        b"1555165519551a55215524552555265529554055415542554455455546554855"
+        b"4955505551555255545555555655585559555a55605561556455655566556855"
+        b"69556a5581558455855589558a559055915594559555965598559955a155a455"
+        b"a555a655a9550056015602560456065608560956115614561556185619562056"
+        b"2156225624562556265628562956415645564656485649564a56505651565256"
+        b"545655565656585659565a566156645665566956825685568656885689568a56"
+        b"915695569a56a256a556a656a856a95604580558065809581058155818582158"
+        b"2a58455848584a58515854585558565858585958605862586458655882588958"
+        b"9058925895589858a158a9580159025905590a59115914591559165919592559"
+        b"41594459455946594959505951595259545955595659585959595a5961596459"
+        b"655966596959815985598959915994599559965998599959a559045a085a155a"
+        b"1a5a205a255a265a295a455a485a495a515a555a565a585a595a625a655a685a"
+        b"6a5a815a8a5a925a955a965a985a9a5aa15a0560146016601960256044605060"
+        b"5560566058605a60616064606660696081609660a56001610461066109611261"
+        b"15612161226126612961456149615161556156615961656166616a6184618a61"
+        b"92619561a161a661a96111621662196240624162466255625662586260628562"
+        b"91629662a56211641264156416641a6421642664296440644264456448644a64"
+        b"516454645564566459645a646064626465648464856489649064926494649564"
+        b"966498649a64a164a464a964056508650a651165156516651965446545654665"
+        b"496550655165546555655665596561656465656566656965866589658a659165"
+        b"9565966599659a65a265a565a665a86502660966156620662666286629664066"
+        b"456648664a66516654665566566658665a666066656668668066826685668a66"
+        b"9466966698669966a066a466a666aa661668196825684168526855685a686168"
+        b"6968856891689868a66801690469106915692169246926692969406941694569"
+        b"4669486951695469556956695969606965696a69826984698a699569a169a469"
+        b"a569a969116a166a186a416a446a496a506a556a586a5a6a646a656a696a866a"
+        b"946a986a9a6aa66a0080028008800a802080228028802a804580508051805480"
+        b"5680598065808080828088808a809580a080a280a880aa800581118114811681"
+        b"1981258141814481498150815281558156815881598164816681698185818981"
+        b"948196819981a5810082028208820a8215822082228228822a82518254825982"
+        b"65828082828288828a829582a082a282a882aa82148419844184448451845584"
+        b"5a846184648469849484998401850985128515851a8526852985408541854585"
+        b"4885518554855585568559855a856585668568856a8581858485868589859085"
+        b"928595859885a68511861686198625864186448649864a865086558659865a86"
+        b"618666866a86858691869a86a4860088028808880a8815882088228828882a88"
+        b"41884588518854885988658869888088828888888a889588a088a288a888aa88"
+        b"05890689118914891689258941894489468949895089528955895a8961896489"
+        b"858996899989a589008a028a088a0a8a158a208a228a288a2a8a458a518a548a"
+        b"568a808a828a888a8a8a958aa08aa28aa88aaa8a059011901690189019902590"
+        b"419046904990559058905a9069906a9085909190949096909990a59001910491"
+        b"069109911091159118911a912191249126912991409145915091519154915591"
+        b"569159916291659184918691929195919891a191a491a691a991059211921492"
+        b"19922592449246924992509252925592589266926992859294929692a9920194"
+        b"04940694109415941894269440944a9451945494559456945894599460946194"
+        b"62946594849486949294949495949894a194a9940095059508950a9510951195"
+        b"14951595169519952195259529952a9541954495459546954995509551955295"
+        b"549555955695589559955a956195649565956695699581958595889591959295"
+        b"94959595969599959a95a095a295a595a895aa95019604961096159619962096"
+        b"2696299645964896499651965296559656965996659668968296849689968a96"
+        b"929694969596a496a696a9960598169819982598419846985098529855985698"
+        b"5a98649865988598919896989998a59804990699099910991299159918991a99"
+        b"209921992499269940994299459948994a995199549955995699599962996599"
+        b"66996a99819984999099929995999a99a199a699059a159a259a449a469a499a"
+        b"509a559a589a619a859a919a949a959a969a00a002a008a00aa015a020a022a0"
+        b"28a02aa045a051a054a056a059a080a082a088a08aa095a0a0a0a2a0a8a0aaa0"
+        b"05a109a111a114a116a119a11aa146a149a151a155a158a15aa161a164a185a1"
+        b"90a192a196a199a102a208a20aa210a219a222a228a22aa245a251a256a259a2"
+        b"65a280a282a288a28aa295a2a0a2a2a2a8a2aaa219a425a441a444a450a454a4"
+        b"55a458a45aa461a465a466a468a469a485a406a509a510a512a515a518a526a5"
+        b"29a542a545a551a554a555a556a559a565a56aa581a584a585a586a589a592a5"
+        b"95a598a505a611a616a61aa621a625a644a646a64aa652a655a656a658a660a6"
+        b"62a686a690a695a696a699a6a1a6a4a6a6a600a802a808a80aa820a822a828a8"
+        b"2aa851a854a856a859a880a882a888a88aa895a8a0a8a2a8a8a8aaa805a914a9"
+        b"19a921a925a941a950a955a95aa961a966a969a990a996a900aa02aa08aa0aaa"
+        b"20aa22aa28aa2aaa51aa54aa56aa80aa82aa88aa8aaa95aaa0aaa2aaa8aaaaaa"
+    )
+
+    delta = np.float32(0.125)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        qs, qh = np.hsplit(rest, [QK_K // 8])
+
+        d = d.view(np.float16).astype(np.float32)
+        qh = qh.view(np.uint16)
+
+        dl = d * (2 * ((qh >> 12) & 7) + 1)
+        dl = dl.reshape((n_blocks, -1, 1, 1))
+        delta = np.where((qh & np.uint16(0x8000)) == 0, cls.delta, -cls.delta)
+        delta = delta.reshape((n_blocks, -1, 1, 1))
+
+        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
+        qs = qs.astype(np.uint16) | ((qh & 7) << 8).reshape((n_blocks, -1))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 4, 8))
+
+        return (dl * (grid + delta)).reshape((n_blocks, -1))
+
+
+class IQ1_M(__Quant, qtype=GGMLQuantizationType.IQ1_M):
+    grid_shape = IQ1_S.grid_shape
+    grid_map = IQ1_S.grid_map
+    grid_hex = IQ1_S.grid_hex
+
+    delta = IQ1_S.delta
+
+    # Okay *this* type is weird. It's the only one which stores the f16 scales in multiple parts.
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        qs, rest = np.hsplit(blocks, [QK_K // 8])
+        qh, scales = np.hsplit(rest, [QK_K // 16])
+
+        # The f16 scale is packed across multiple bytes
+        scales = scales.view(np.uint16)
+        d = (scales.reshape((n_blocks, 4)) & np.uint16(0xF000)) >> np.array([12, 8, 4, 0], dtype=np.uint16).reshape((1, 4))
+        d = d[..., 0] | d[..., 1] | d[..., 2] | d[..., 3]
+        d = d.view(np.float16).astype(np.float32).reshape((n_blocks, 1))
+
+        scales = scales.reshape(n_blocks, -1, 1) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
+        scales = (scales & 0x07).reshape((n_blocks, -1))
+        dl = d * (2 * scales + 1)
+        dl = dl.reshape((n_blocks, -1, 2, 1, 1))
+
+        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
+        qs = qs.astype(np.uint16) | ((qh & 0x07).astype(np.uint16) << 8).reshape((n_blocks, -1))
+
+        delta = np.where(qh & 0x08 == 0, cls.delta, -cls.delta)
+        delta = delta.reshape((n_blocks, -1, 2, 2, 1))
+
+        assert cls.grid is not None
+        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
+        grid = grid.reshape((n_blocks, -1, 2, 2, 8))
+
+        return (dl * (grid + delta)).reshape((n_blocks, -1))
+
+
+class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL):
+    kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, qs = np.hsplit(blocks, [2])
+
+        d = d.view(np.float16).astype(np.float32)
+
+        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+
+        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1))
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
+        qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1))
+
+        return (d * qs)
+
+
+class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS):
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d, rest = np.hsplit(blocks, [2])
+        scales_h, rest = np.hsplit(rest, [2])
+        scales_l, qs = np.hsplit(rest, [QK_K // 64])
+
+        d = d.view(np.float16).astype(np.float32)
+        scales_h = scales_h.view(np.uint16)
+
+        scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
+        scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array([2 * i for i in range(QK_K // 32)], dtype=np.uint16).reshape((1, -1, 1))
+        scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F)
+        scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03)
+
+        scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32)
+        dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1))
+
+        qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
+        qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F)
+
+        kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1))
+        qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32))
+
+        return (dl * qs).reshape((n_blocks, -1))
diff --git a/tests/kernels/test_ggml.py b/tests/kernels/test_ggml.py
new file mode 100644
index 0000000000..ea6a5d6830
--- /dev/null
+++ b/tests/kernels/test_ggml.py
@@ -0,0 +1,84 @@
+"""Packed GGUF ops against gguf-py's CPU dequantization and PyTorch math."""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn.functional as F
+
+from freetoken.kernel import ggml
+from freetoken.third_party.gguf import GGMLQuantizationType as QType
+from freetoken.third_party.gguf import quants
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU")
+DTYPES = [torch.float32, torch.bfloat16, torch.float16]
+
+
+def _q4_weights(shape, seed=690):
+    values = np.random.default_rng(seed).normal(0, 0.1, shape).astype(np.float32)
+    packed = quants.quantize(values, QType.Q4_0)
+    reference = torch.from_numpy(quants.dequantize(packed, QType.Q4_0))
+    # ggml's last MMQ tile may read an extra partial row.
+    storage = torch.zeros(*packed.shape[:-2], packed.shape[-2] + 1, packed.shape[-1], dtype=torch.uint8, device="cuda")
+    storage[..., :-1, :].copy_(torch.from_numpy(packed))
+    return storage[..., :-1, :], reference
+
+
+def _activation(shape, dtype, seed=19):
+    rng = np.random.default_rng(seed)
+    half = rng.integers(-127, 128, (*shape[:-1], shape[-1] // 32, 16))
+    half[..., 0] = 127
+    # Exactly representable q8_1 inputs with zero block sums isolate weight math.
+    values = np.concatenate([half, -half], axis=-1).reshape(shape).astype(np.float32) / 64
+    return torch.from_numpy(values).to("cuda", dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_q4_dequantize_and_rows(dtype):
+    w, ref = _q4_weights((64, 256))
+    actual = ggml.dequantize(w, int(QType.Q4_0), 256, dtype)
+    torch.testing.assert_close(actual.cpu(), ref.to(dtype), atol=0, rtol=0)
+    ids = torch.tensor([63, 0, 17, 17], dtype=torch.int32, device="cuda")
+    rows = ggml.get_rows(w, int(QType.Q4_0), 256, ids, dtype)
+    torch.testing.assert_close(rows.cpu(), ref[ids.cpu().long()].to(dtype), atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("batch", [1, 17])
+def test_q4_matrix_matches_cpu(dtype, batch):
+    w, ref = _q4_weights((64, 256))
+    x = _activation((batch, 256), dtype)
+    actual = ggml.mul_mat(w, int(QType.Q4_0), x)
+    expected = F.linear(x.cpu().float(), ref).to(dtype)
+    torch.testing.assert_close(actual.cpu(), expected, atol=2e-4, rtol=1e-3)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("batch", [1, 17])
+def test_q4_moe_routes_match_cpu(dtype, batch):
+    w, ref = _q4_weights((4, 64, 256))
+    x = _activation((batch, 256), dtype)
+    ids = torch.tensor([[i % 4, (i + 1) % 4] for i in range(batch)], dtype=torch.int32, device="cuda")
+    actual = ggml.mul_mat_id(w, int(QType.Q4_0), x, ids)
+    expected = torch.einsum("tk, tjok -> tjo", x.cpu().float(), ref[ids.cpu().long()]).to(dtype)
+    torch.testing.assert_close(actual.cpu(), expected, atol=2e-4, rtol=1e-3)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_q4_shared_q8_matches_internal_quantization(dtype):
+    w, _ = _q4_weights((64, 256))
+    x = _activation((1, 256), dtype)
+    expected = ggml.mul_mat(w, int(QType.Q4_0), x)
+    actual = ggml.mul_mat(w, int(QType.Q4_0), x, xq=ggml.quantize_q8_1(x))
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+def test_q4_nondefault_stream():
+    w, ref = _q4_weights((64, 256))
+    torch.cuda.synchronize()
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        x = _activation((1, 256), torch.float32)
+        actual = ggml.mul_mat(w, int(QType.Q4_0), x)
+        copied = actual.clone()
+    stream.synchronize()
+    torch.testing.assert_close(copied.cpu(), F.linear(x.cpu(), ref), atol=2e-4, rtol=1e-3)
diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py
index 98b19e2de4..83c0ec61ff 100644
--- a/tests/models/deepseek_v41/harness.py
+++ b/tests/models/deepseek_v41/harness.py
@@ -62,7 +62,7 @@ def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req:
             from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache
 
             self.banks = load_expert_banks(
-                checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False,
+                checkpoint, mc, method=method, device=self.device, parallel=False,
             )
             self.expert_cache = OffloadMoeCache(
                 num_layers=mc.num_moe_layers, num_experts=mc.num_experts,
@@ -76,7 +76,7 @@ def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req:
             from freetoken.moe.expert_banks import attach_resident_banks
 
             attach_resident_banks(list(iter_moe_layers(self.model)), load_expert_banks(
-                checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, resident=True,
+                checkpoint, mc, method=method, device=self.device, parallel=False, resident=True,
             ))
         for layer in self.model.engram_layers():
             layer.attach_table(engram_table if engram_table is not None else ZeroEngramTable(layer.width, self.device))
diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py
index 7babb5d4ce..9452f95c7e 100644
--- a/tests/moe/test_cpu_moe_q4_0.py
+++ b/tests/moe/test_cpu_moe_q4_0.py
@@ -3,7 +3,7 @@
 
 The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same*
 packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the
-K-loop. We check it against the reference dequant (models/gguf/dequant.py) + the
+K-loop. We check it against the reference dequant (gguf-py's numpy dequantize) + the
 production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the
 only spread is weight bf16-rounding + reduction order -> tight relative tolerance.
 
@@ -39,11 +39,11 @@ def _pack_q4_0(nibbles: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
 
 def _dequant_bank(packed: torch.Tensor, K: int, dev) -> torch.Tensor:
     """[S, OUT, K//32*18] packed Q4_0 -> [S, OUT, K] bf16 (storage order == elem order)."""
-    from freetoken.models.gguf.dequant import GGML_Q4_0, dequantize
+    from freetoken.third_party.gguf import GGMLQuantizationType, quants
 
     S, OUT, _ = packed.shape
-    flat = dequantize(packed.reshape(-1), GGML_Q4_0, torch.bfloat16)
-    return flat.reshape(S, OUT, K).to(dev)
+    values = quants.dequantize(packed.numpy(), GGMLQuantizationType.Q4_0)
+    return torch.from_numpy(values).reshape(S, OUT, K).to(dev, torch.bfloat16)
 
 
 def _make_q4_0_cache(L, E, H, I, seed=0):
@@ -115,8 +115,10 @@ def test_cpu_decode_q4_0_matches_dequant_then_gpu(bs):
 def test_cpu_decode_q4_0_matches_ggml_mmvq():
     """Sanity: the CPU W4A16 GEMV lands close to the GPU ggml MMVQ (W4A8) kernel the
     offload path uses -- a looser tol since MMVQ quantizes activations to int8."""
+    from freetoken.kernel import ggml
+    from freetoken.layers.activation import gated_act_and_mul
     from freetoken.moe.cpu_executor import CpuMoeExecutor
-    from freetoken.moe.fused_q4_0 import fused_experts_gguf_q4_0
+    from freetoken.third_party.gguf import GGMLQuantizationType
 
     torch.manual_seed(77)
     L, E, H, I, top_k = 2, 16, 2816, 704, 8
@@ -137,11 +139,12 @@ def test_cpu_decode_q4_0_matches_ggml_mmvq():
     torch.cuda.synchronize()
 
     b = cache.bank_sources
-    gate_up_layer = b["gate_up"][layer].to(dev)
-    down_layer = b["down"][layer].to(dev)
-    gpu_out = fused_experts_gguf_q4_0(
-        hidden, gate_up_layer, down_layer, w, ids.clone(), "gelu_tanh"
-    ).float()
+    q4_0 = int(GGMLQuantizationType.Q4_0)
+    h = ggml.mul_mat_id(b["gate_up"][layer].to(dev), q4_0, hidden.float(), ids)
+    inter = torch.empty(bs, top_k, I, device=dev)
+    gated_act_and_mul("gelu_tanh", h.view(-1, 2 * I), inter.view(-1, I))
+    y = ggml.mul_mat_id(b["down"][layer].to(dev), q4_0, inter, ids)
+    gpu_out = ggml.moe_weighted_sum(y, w)
 
     rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)
     assert rel < 6e-2, f"q4_0 vs ggml-mmvq rel err {rel.item()}"
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index d3c8d9d843..a5a7adb9f0 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -681,7 +681,6 @@ def decode(steps: int) -> None:
         check_resident()
 
     decode(20)
-    cache.zero_inactive_route_slots = True
     cache.begin_prefill()
     for layer in range(L):
         cache.prefetch_prefill_layer(layer)
@@ -735,7 +734,6 @@ def make(pool_of_layer, pool_sizes):
     cache.set_bank_sources(sources, layer_tags=tags)
     assert cache.bank_caches["down"].shape == (20, 32, 144)
     # the staging buffers cover pool 0 and slots [0, 6) of pool 1, so both slot 0 rows end up with another quant's bytes
-    cache.zero_inactive_route_slots = True
     cache.begin_prefill()
     for layer in range(L):
         cache.prefetch_prefill_layer(layer)
@@ -748,6 +746,56 @@ def make(pool_of_layer, pool_sizes):
             assert not view[0].any()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("wide", [(), (1, 3)], ids=["cpu-mode", "two-layouts"])
+def test_slot_pools_and_staging_buffers_end_in_zero_mmq_tails(wide):
+    from freetoken.layers.quantization.ggml_types import mmq_tail_bytes, row_bytes
+    from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.third_party.gguf import GGMLQuantizationType
+
+    _init_tp()
+    torch.manual_seed(0)
+    L, E, top_k = 4, 8, 4
+
+    def layout(ggml_type):
+        """(rows, row bytes, tail bytes) of gate_up and down for Gemma-4 26B-A4B's K = 2816 and 704, with fewer rows."""
+        return tuple((rows, row_bytes(ggml_type, k), mmq_tail_bytes(ggml_type, k)) for rows, k in ((32, 2816), (64, 704)))
+
+    q4_0, q8_0 = layout(GGMLQuantizationType.Q4_0), layout(GGMLQuantizationType.Q8_0)
+    shapes = [q8_0 if layer in wide else q4_0 for layer in range(L)]
+    roles = ("gate_up", "down")
+    sources = {role: [torch.randint(1, 255, (E, *shapes[l][k][:2]), dtype=torch.uint8).pin_memory() for l in range(L)] for k, role in enumerate(roles)}
+    tails = {role: [shapes[l][k][2] for l in range(L)] for k, role in enumerate(roles)}
+    # cpu mode: the slots are exactly the two staging buffers
+    pools = dict(pool_of_layer=[2, 0, 3, 1], pool_sizes=[6, 5, 12, 10], prefill_hit_d2d=True) if wide else dict(pool_of_layer=[0] * L, pool_sizes=[2 * E])
+    cache = OffloadMoeCache(num_layers=L, num_experts=E, cache_size=sum(pools["pool_sizes"]), device=torch.device("cuda"), prefill_overlap=True, **pools)
+    cache.set_bank_sources(sources, row_tail_bytes=tails)
+    assert cache.prefill_overlap
+
+    def assert_zero_tails(layer: int, views) -> None:
+        torch.cuda.synchronize()
+        for k, ((_, bank), view) in enumerate(zip(cache.banks, views)):
+            end, tail = view.storage_offset() + view.numel(), shapes[layer][k][2]
+            assert end + tail <= bank.numel() and not bank[end : end + tail].any(), (layer, k)
+
+    for _ in range(20):
+        for layer in range(L):
+            cache.ensure_experts(layer, torch.randperm(E)[:top_k].to(torch.int32).view(1, top_k).cuda())
+            cache.copy_missing()
+    cache.begin_prefill()
+    for layer in range(L):
+        cache.prefetch_prefill_layer(layer)
+        cache.prefetch_prefill_layer(layer + 1)
+        assert_zero_tails(layer, cache.wait_prefill_layer(layer))
+        cache.release_prefill_layer(layer)
+    for layer in range(L):
+        assert_zero_tails(layer, cache.pool_bank_views(layer))
+    if wide:
+        cache.materialize_layer(0)
+        cache.copy_missing()
+        assert_zero_tails(0, cache.stage_views(0, 0))
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 def test_offload_cache_rebuild_frees_the_old_slots_first():
     from freetoken.moe.offload_cache import OffloadMoeCache

From 7d6490c5d6bae3cbaee3d0f0394886d301d4d2fe Mon Sep 17 00:00:00 2001
From: Zhang Jason <ning.zhang2@amd.com>
Date: Fri, 16 Oct 2026 11:34:20 +0800
Subject: [PATCH 100/249] fix(rocm): clamp QSA KV tile to the device LDS budget
 on RDNA (#587)

The qsa_sparse_paged_attention tiles are tuned for the GB300's 228KB shared
memory; RDNA gfx11/gfx12 have ~64KB/CU, so k+v staging overflows LDS. Query the
device's actual shared-memory limit and shrink the KV tile only when it doesn't
fit, ROCm-gated so NVIDIA tiles stay byte-identical.

The gguf ROCm-compile fixes this branch also carried are now covered upstream by
the ggml-cuda vendoring in #690 (separate DEFAULT_HIP_CFLAGS + vendors/hip.h), so
this branch reduces to the QSA clamp only.
---
 python/freetoken/kernel/triton/qsa/attend.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py
index 541e27c680..4c55f8201a 100644
--- a/python/freetoken/kernel/triton/qsa/attend.py
+++ b/python/freetoken/kernel/triton/qsa/attend.py
@@ -277,6 +277,18 @@ def qsa_sparse_paged_attention(
     else:
         block_n, target_splits, partial_warps = 64, 1, 2
 
+    # These tiles were tuned for the GB300's 228KB LDS; AMD parts have less and it varies by arch
+    # (e.g. 64KB/CU on RDNA gfx11/gfx12), so k+v staging can overflow. Query the device's actual
+    # shared-memory limit rather than assuming a fixed size, and shrink the KV tile only when it
+    # doesn't fit -- so a larger-LDS arch keeps the wide tile. ROCm-gated so the NVIDIA tiles stay
+    # byte-identical (a small NV part could report a 48KB default and would otherwise clamp).
+    if torch.version.hip is not None:
+        _props = torch.cuda.get_device_properties(q.device)
+        _lds_budget = getattr(_props, "shared_memory_per_block", 64 * 1024) or 64 * 1024
+        _kv_bytes = k_cache.element_size()
+        while block_n > 16 and 2 * block_n * q.shape[2] * _kv_bytes + 2048 > _lds_budget:
+            block_n //= 2
+
     num_tiles = triton.cdiv(logical_indices.shape[1], block_n)
     # Avoid empty splits when the selection width is smaller than the profile.
     max_useful_splits = 1 << (num_tiles.bit_length() - 1)

From b6193e86f7fdef2ed553f9d8f83e716387a57d45 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 20:35:02 -0700
Subject: [PATCH 101/249] build(deps): upgrade transformers to 5.18 (#698)

---
 pyproject.toml                              |  2 +-
 python/freetoken/models/gguf/tokenizer.py   | 79 +++++++--------------
 python/freetoken/models/glm5_next/args.py   | 10 +--
 python/freetoken/models/qwen4_exp/config.py |  4 +-
 python/freetoken/utils/hf.py                |  2 +-
 tests/attention/test_dsa_kpool.py           |  2 +-
 tests/models/test_glm_dsa.py                |  4 +-
 7 files changed, 41 insertions(+), 62 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ab1dfd0731..384b905f2c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,7 +57,7 @@ dependencies = [
     "torch>=2.14,<2.15",
     "torchvision>=0.29,<0.30",
     "tqdm>=4.66,<5",
-    "transformers>=5.16,<5.17",
+    "transformers>=5.18,<5.19",
     # CUDA torch 2.14 resolves Triton 3.8; AMD's ROCm 7.14 image supplies its
     # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
     # runtime-specific wheel selected by the PyTorch distribution.
diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index e6a5d4aca0..af715c6598 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -1,74 +1,49 @@
 """Build a HF fast tokenizer from a GGUF file's embedded tokenizer metadata.
 
-transformers' ``AutoTokenizer.from_pretrained(gguf_file=...)`` first builds the HF
-config, which the gemma4 strict dataclass rejects (per-layer ``num_key_value_heads``
-array). So we call the GGUF->fast tokenizer converter directly on the
-``tokenizer.ggml.*`` metadata, bypassing config entirely.
+transformers' ``AutoTokenizer.from_pretrained(gguf_file=...)`` first builds the HF model config from the GGUF, which needs the PyPI ``gguf`` package and an architecture transformers knows (``qwen4exp`` is not one).
+So we run only its tokenizer half: ``get_gguf_tokenizer`` and ``convert_gguf_tokenizer``.
 """
 
 from __future__ import annotations
 
-from typing import Any
+from .reader import load_gguf_metadata
 
-from .reader import gguf_architecture, load_gguf_metadata
-
-# GGUF architecture -> transformers GGUF tokenizer-converter key.
-_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen4exp": "qwen3"}
+# end-of-generation token names from llama.cpp (llama_vocab::impl::load in src/llama-vocab.cpp): a GGUF stores one eos id and no stop list
+_EOG_NAMES = frozenset((
+    "<|eot_id|>", "<|im_end|>", "<|end|>", "<|return|>", "<|call|>", "<|flush|>", "<|calls|>", "<end_of_turn>",
+    "<|endoftext|>", "</s>", "<|eom_id|>", "<EOT>", "_<EOT>", "[EOT]", "[EOS]", "<|end_of_text|>", "<end_of_utterance>",
+    "<eos>", "<turn|>", "<|tool_response>", "<\uff5cend\u2581of\u2581sentence\uff5c>", "[e~[", "<|ifm|im_end|>",
+))
 
 
 def load_gguf_tokenizer(model_path: str):
     from transformers import PreTrainedTokenizerFast
-    from transformers.integrations.ggml import convert_gguf_tokenizer
+    from transformers.integrations.gguf import convert_gguf_tokenizer, get_gguf_tokenizer
 
+    arch, tok_dict, special = get_gguf_tokenizer(model_path)
+    fast, _ = convert_gguf_tokenizer(arch, tok_dict)
     meta = load_gguf_metadata(model_path)
-    arch = gguf_architecture(model_path)
-    conv_arch = _TOKENIZER_ARCH.get(arch, arch)
-    tok_dict: dict[str, Any] = {
-        k[len("tokenizer.ggml.") :]: v
-        for k, v in meta.items()
-        if k.startswith("tokenizer.ggml.")
-    }
-    fast, _extra = convert_gguf_tokenizer(conv_arch, tok_dict)
-
-    tokens = tok_dict["tokens"]
-
-    def tok_for(id_key: str, default: str) -> str:
-        tid = meta.get(f"tokenizer.ggml.{id_key}")
-        return tokens[int(tid)] if tid is not None and int(tid) < len(tokens) else default
-
-    # gemma4 chat turns end with <turn|>; prefer it as eos so chat generation halts
-    # (the formal <eos> is also a stop id, see gguf_eos_token_ids).
-    turn_end = "<turn|>" if "<turn|>" in tokens else None
-    tokenizer = PreTrainedTokenizerFast(
+    # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template
+    return PreTrainedTokenizerFast(
         tokenizer_object=fast,
-        bos_token=tok_for("bos_token_id", "<bos>"),
-        eos_token=turn_end or tok_for("eos_token_id", "<eos>"),
-        unk_token=tok_for("unknown_token_id", "<unk>"),
-        pad_token=tok_for("padding_token_id", "<pad>"),
+        **special,
+        add_bos_token=bool(meta.get("tokenizer.ggml.add_bos_token")),
+        add_eos_token=bool(meta.get("tokenizer.ggml.add_eos_token")),
     )
-    chat_template = meta.get("tokenizer.chat_template")
-    if chat_template:
-        tokenizer.chat_template = chat_template
-    return tokenizer
 
 
-def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]:
-    """Stop ids for GGUF generation: the formal <eos> plus the chat turn end <turn|>."""
+def gguf_eos_token_ids(model_path: str) -> set[int]:
+    """Stop ids for GGUF generation, picked the way llama.cpp picks its end-of-generation tokens."""
     meta = load_gguf_metadata(model_path)
     tokens = meta["tokenizer.ggml.tokens"]
-    ids: set[int] = set()
-    if tokenizer.eos_token_id is not None:
-        ids.add(int(tokenizer.eos_token_id))
-    eid = meta.get("tokenizer.ggml.eos_token_id")
-    if eid is not None:
-        ids.add(int(eid))
-    # Look the stop tokens up in the vocab directly (convert_tokens_to_ids would map an
-    # absent name to <unk>, wrongly adding it as a stop id).
-    for name in ("<eos>", "<turn|>"):
-        try:
-            ids.add(tokens.index(name))
-        except ValueError:
-            pass
+    ids = {i for i, t in enumerate(tokens) if t in _EOG_NAMES}
+    ids |= {int(meta[f"tokenizer.ggml.{k}_token_id"]) for k in ("eos", "eot", "eom") if f"tokenizer.ggml.{k}_token_id" in meta}
+    names = {tokens[i]: i for i in ids if i < len(tokens)}
+    # llama.cpp's exceptions: <|end|> only ends a message when <|call|> and <|return|> end the turn (gpt-oss), and gemma4's </s> is a normal token
+    if "<|end|>" in names and names.keys() & {"<|call|>", "<|calls|>"} and names.keys() & {"<|return|>", "<|flush|>"}:
+        ids.discard(names["<|end|>"])
+    if "</s>" in names and names.keys() & {"<|tool_response>", "<|plamo:eos|>"}:
+        ids.discard(names["</s>"])
     return ids
 
 
diff --git a/python/freetoken/models/glm5_next/args.py b/python/freetoken/models/glm5_next/args.py
index c25fef4856..c771467b89 100644
--- a/python/freetoken/models/glm5_next/args.py
+++ b/python/freetoken/models/glm5_next/args.py
@@ -24,16 +24,16 @@
 from dataclasses import dataclass
 from typing import Any, Tuple
 
-# Layer-type strings used by the checkpoint's ``layer_types`` field.
+# Layer-type strings after transformers' legacy remap, which renames the checkpoint's deepseek_sparse_attention.
 KDA_LAYER = "linear_attention"
-DSA_LAYER = "deepseek_sparse_attention"
+DSA_LAYER = "indexed_attention"
 
 
 @dataclass(frozen=True)
 class Glm5NextArgs:
     hidden_size: int
     num_heads: int
-    # ---- MLA (the 11 "deepseek_sparse_attention" layers) ----
+    # ---- MLA (the 11 "indexed_attention" layers) ----
     q_lora_rank: int
     kv_lora_rank: int
     qk_nope_head_dim: int
@@ -119,8 +119,10 @@ def load_args(hf_config: Any) -> Glm5NextArgs:
     are accepted as fallbacks."""
     text = _get(hf_config, "text_config", hf_config)
 
+    from transformers.configuration_utils import remap_legacy_layer_types
+
     num_layers = int(text.num_hidden_layers)
-    layer_types = tuple(_get(text, "layer_types", ()) or ())
+    layer_types = tuple(remap_legacy_layer_types(list(_get(text, "layer_types", ()) or ())))
     if not layer_types:
         raise ValueError("glm5_next config is missing layer_types")
     if len(layer_types) != num_layers:
diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index b1e5f9aa1d..2bc96f67a1 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -101,9 +101,9 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]:
 def _layer_types(text: Any) -> list[str]:
     layer_types = getattr(text, "layer_types", None)
     if layer_types is not None:
-        # HF Qwen4ExpTextConfig rewrites full_attention to qwen_sparse_attention in __post_init__.
+        # HF Qwen4ExpTextConfig rewrites full_attention to indexed_attention in __post_init__.
         return [
-            "full_attention" if t == "qwen_sparse_attention" else t for t in layer_types
+            "full_attention" if t == "indexed_attention" else t for t in layer_types
         ]
     # Fall back to full_attention_interval: every Nth layer (1-indexed) is full.
     interval = int(getattr(text, "full_attention_interval", 4))
diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py
index 489585e8c2..5a7b7e6935 100644
--- a/python/freetoken/utils/hf.py
+++ b/python/freetoken/utils/hf.py
@@ -62,7 +62,7 @@ def load_eos_token_ids(
     if (gguf_src := gguf_config_source(model_path)) is not None:
         from freetoken.models.gguf.tokenizer import gguf_eos_token_ids
 
-        return frozenset(gguf_eos_token_ids(gguf_src, tokenizer))
+        return frozenset(gguf_eos_token_ids(gguf_src))
 
     ids: set[int] = set()
     if tokenizer.eos_token_id is not None:
diff --git a/tests/attention/test_dsa_kpool.py b/tests/attention/test_dsa_kpool.py
index fb856e1757..ec2bbb1fa2 100644
--- a/tests/attention/test_dsa_kpool.py
+++ b/tests/attention/test_dsa_kpool.py
@@ -41,7 +41,7 @@ def _args(num_layers=1):
         index_kpool_always_select_tail=True,
         linear_num_heads=0, linear_head_dim=0, linear_conv_kernel_dim=4,
         linear_lower_bound=-5.0,
-        layer_types=("deepseek_sparse_attention",) * num_layers,
+        layer_types=("indexed_attention",) * num_layers,
         mlp_layer_types=("dense",) * num_layers,
         mhc=False, mhc_num_residual_streams=1, hc_eps=1e-6,
         mhc_sinkhorn_iterations=0, mhc_tau=0.05, mhc_post_mult_value=2.0,
diff --git a/tests/models/test_glm_dsa.py b/tests/models/test_glm_dsa.py
index 0dceabadd0..7a82894dac 100644
--- a/tests/models/test_glm_dsa.py
+++ b/tests/models/test_glm_dsa.py
@@ -41,7 +41,9 @@ class Cfg:  # minimal duck-typed config for GlmMoeDsaIndexer
     freqs = torch.outer(pos.float(), inv)
     emb = torch.cat((freqs, freqs), dim=-1)
     cos, sin = emb.cos()[None].to(torch.bfloat16), emb.sin()[None].to(torch.bfloat16)
-    ref_topk = idx(x, q_resid, (cos, sin), None, pos[None])  # [1, S, topk]
+    # bool causal mask, True = visible: since transformers 5.18 the indexer no longer builds one when given None
+    causal = (pos[None, :] <= pos[:, None])[None]
+    ref_topk = idx(x, q_resid, (cos, sin), causal, pos[None])  # [1, S, topk]
     return idx, x, q_resid, pos, ref_topk
 
 

From df8d450c28affecf3023931b9f84f52faf38f1ea Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 21:25:32 -0700
Subject: [PATCH 102/249] fix(gguf): reject ggml weights past the 32-bit block
 offsets of MMQ / MMVQ (#699)

---
 .../kernel/csrc/ggml/freetoken/ggml_ops.cu       | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
index cc735c3b0e..69d9322736 100644
--- a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
@@ -123,6 +123,19 @@ int64_t mmq_tail_bytes(int64_t type_id, int64_t k) {
     return rem == 0 ? 0 : (int64_t) ggml_row_size((ggml_type) type_id, MATRIX_ROW_PADDING - rem);
 }
 
+// MMQ / MMVQ index weight blocks from the data pointer with 32-bit ints, MMQ tail reads included
+void check_block_offsets(TensorView w, ggml_type type, int64_t k) {
+    if (w.numel() == 0) {
+        return;
+    }
+    int64_t bytes = 1 + mmq_tail_bytes(type, k);
+    for (int d = 0; d < w.ndim(); ++d) {
+        bytes += (w.size(d) - 1) * w.stride(d);
+    }
+    const int64_t blocks = bytes / (int64_t) ggml_type_size(type);
+    RuntimeCheck(blocks <= INT32_MAX, "ggml weight spans ", blocks, " ", ggml_type_name(type), " blocks; MMQ / MMVQ address at most ", INT32_MAX);
+}
+
 struct Activation {
     const void * data;
     ggml_type type;
@@ -186,6 +199,7 @@ void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Opt
     const ggml_type type = (ggml_type) type_id;
     const int64_t m = x.size(0), k = x.size(1), n = w.size(0);
     check_weight(w, type, k);
+    check_block_offsets(w, type, k);
     RuntimeCheck(out.ndim() == 2 && out.size(0) == m && out.size(1) == n && out.stride(1) == 1, "mul_mat out must be a [M, N] view with contiguous rows");
     const ggml_type otype = io_type(x, out);
     if (m == 0) {
@@ -229,6 +243,7 @@ void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids,
     RuntimeCheck(x.ndim() == 3 && x.size(0) == t && (x.size(1) == 1 || x.size(1) == top_k), "x must be [T, 1, K] or [T, k, K]");
     const int64_t ky = x.size(1), k = x.size(2), e = w.size(0), n = w.size(1);
     check_weight(w, type, k);
+    check_block_offsets(w, type, k);
     RuntimeCheck(out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n && out.IsContiguous(),
                  "out must be a contiguous [T, k, N] tensor");
     const ggml_type otype = io_type(x, out);
@@ -285,6 +300,7 @@ void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView i
     const ggml_type type = (ggml_type) type_id;
     const int64_t t = ids.size(0), top_k = ids.size(1), k = x.size(1), e = w.size(0), n = w.size(1) / 2;
     check_weight(w, type, k);
+    check_block_offsets(w, type, k);
     RuntimeCheck(out.IsContiguous() && out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n,
                  "out must be a contiguous [T, k, I] tensor");
     const ggml_type otype = io_type(x, out);

From 18448d29a764d13360a4e7be3b66e9177007b426 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 22:15:28 -0700
Subject: [PATCH 103/249] feat(gguf): serve Qwen3.5 dense and MoE GGUF
 checkpoints (#701)

---
 python/freetoken/models/gguf/config.py        |   2 +
 .../freetoken/models/qwen3_5_moe/__init__.py  |   4 +
 python/freetoken/models/qwen3_5_moe/gdn.py    |   7 +-
 python/freetoken/models/qwen3_5_moe/gguf.py   | 220 ++++++++++++++++++
 python/freetoken/models/qwen3_5_moe/model.py  |   4 +
 python/freetoken/models/register.py           |  19 ++
 6 files changed, 255 insertions(+), 1 deletion(-)
 create mode 100644 python/freetoken/models/qwen3_5_moe/gguf.py

diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 4eb2a999db..f9210938fb 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -18,6 +18,8 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "qwen35": "Qwen3_5GGUFForCausalLM",
+    "qwen35moe": "Qwen3_5MoeGGUFForCausalLM",
     "qwen4exp": "Qwen4ExpGGUFForCausalLM",
 }
 
diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py
index 018b1c1e97..abe8fd1c2f 100644
--- a/python/freetoken/models/qwen3_5_moe/__init__.py
+++ b/python/freetoken/models/qwen3_5_moe/__init__.py
@@ -1,4 +1,5 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import (
     Qwen3_5ForCausalLM,
     Qwen3_5ForConditionalGeneration,
@@ -18,4 +19,7 @@
     "iter_weights_parallel",
     "iter_expert_pieces",
     "nvfp4_expert_spec",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py
index d5fe2aa8fe..0785ad4968 100644
--- a/python/freetoken/models/qwen3_5_moe/gdn.py
+++ b/python/freetoken/models/qwen3_5_moe/gdn.py
@@ -30,9 +30,10 @@ class Qwen3_5GatedDeltaNet(BaseOP):
     def __init__(
         self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim,
         conv_kernel_size, rms_norm_eps, layer_id, *, quant_config: QuantConfig | None = None,
-        prefix: str = "",
+        prefix: str = "", out_heads_tiled: bool = False,
     ):
         self.layer_id = layer_id
+        self._out_heads_tiled = out_heads_tiled
         # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as
         # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the
         # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head
@@ -195,6 +196,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         core_out = core_out.reshape(-1, self.head_v_dim)
         z = z.reshape(-1, self.head_v_dim)
         out = self.norm.forward(core_out, z).reshape(total, -1)
+        if self._out_heads_tiled:
+            # grouped value heads (k-head major) -> tiled (k-head minor), the column order of a GGUF ssm_out
+            r = self.num_v_heads // self.num_k_heads
+            out = out.view(total, self.num_k_heads, r, self.head_v_dim).transpose(1, 2).reshape(total, -1)
         return self.out_proj.forward(out)
 
 
diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py
new file mode 100644
index 0000000000..85af8ab85f
--- /dev/null
+++ b/python/freetoken/models/qwen3_5_moe/gguf.py
@@ -0,0 +1,220 @@
+"""Qwen3.5 / 3.6 / 3.8 dense and MoE from a llama.cpp ``qwen35`` / ``qwen35moe`` GGUF: config from the GGUF metadata, tensor names, and the converter's transforms undone.
+
+The llama.cpp converter (conversion/qwen.py) stores ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, and the GDN value heads in tiled order (k-head minor). Its zero-centred norms are stored as ``w + 1``, which is what this family loads. Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.layers.quantization.ggml_types import is_quantized
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_LAYER_LEAVES: dict[str, tuple[str, ...]] = {
+    "linear_attn.in_proj_qkv": ("attn_qkv",),
+    "linear_attn.in_proj_z": ("attn_gate",),
+    "linear_attn.in_proj_b": ("ssm_beta",),
+    "linear_attn.in_proj_a": ("ssm_alpha",),
+    "linear_attn.out_proj": ("ssm_out",),
+    "self_attn.q_proj": ("attn_q",),
+    "self_attn.k_proj": ("attn_k",),
+    "self_attn.v_proj": ("attn_v",),
+    "self_attn.o_proj": ("attn_output",),
+    "mlp.gate_proj": ("ffn_gate",),
+    "mlp.up_proj": ("ffn_up",),
+    "mlp.down_proj": ("ffn_down",),
+    "mlp.gate": ("ffn_gate_inp",),
+    "mlp.shared_expert_gate": ("ffn_gate_inp_shexp",),
+    "mlp.shared_expert.gate_proj": ("ffn_gate_shexp",),
+    "mlp.shared_expert.up_proj": ("ffn_up_shexp",),
+    "mlp.shared_expert.down_proj": ("ffn_down_shexp",),
+    "mlp.experts.0.gate_proj": ("ffn_gate_exps",),
+    "mlp.experts.0.up_proj": ("ffn_up_exps",),
+    "mlp.experts.0.down_proj": ("ffn_down_exps",),
+}
+_TOP_LEAVES: dict[str, tuple[str, ...]] = {
+    "model.language_model.embed_tokens": ("token_embd",),
+    "lm_head": ("output",),
+}
+_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return _TOP_LEAVES[name]
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors:
+        return (f"{blk}.ffn_gate_up_exps",)
+    return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ()))
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    md = shim.metadata
+    arch = shim.model_type
+
+    def kv(key: str):
+        return md[f"{arch}.{key}"]
+
+    _check_supported(shim)
+    head_dim = int(kv("attention.key_length"))
+    n_v = int(kv("ssm.time_step_rank"))
+    # block_count also counts the MTP layers a GGUF may carry after the decoder
+    num_layers = int(kv("block_count")) - int(md.get(f"{arch}.nextn_predict_layers", 0))
+    recurrent = md.get(f"{arch}.attention.recurrent_layers")
+    if recurrent is not None:
+        layer_types = ["linear_attention" if r else "full_attention" for r in list(recurrent)[:num_layers]]
+    else:
+        interval = int(kv("full_attention_interval"))
+        layer_types = ["full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(num_layers)]
+    moe = arch == "qwen35moe"
+    text = SimpleNamespace(
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=num_layers,
+        layer_types=layer_types,
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=head_dim,
+        rope_parameters={
+            "rope_type": "default",
+            "rope_theta": float(kv("rope.freq_base")),
+            "partial_rotary_factor": int(kv("rope.dimension_count")) / head_dim,
+            "mrope_section": [int(s) for s in kv("rope.dimension_sections")][:3],
+            "mrope_interleaved": True,
+        },
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=0 if moe else int(kv("feed_forward_length")),
+        num_experts=int(kv("expert_count")) if moe else 0,
+        num_experts_per_tok=int(kv("expert_used_count")) if moe else 0,
+        moe_intermediate_size=int(kv("expert_feed_forward_length")) if moe else 0,
+        shared_expert_intermediate_size=int(kv("expert_shared_feed_forward_length")) if moe else 0,
+        linear_num_key_heads=int(kv("ssm.group_count")),
+        linear_num_value_heads=n_v,
+        linear_key_head_dim=int(kv("ssm.state_size")),
+        linear_value_head_dim=int(kv("ssm.inner_size")) // n_v,
+        linear_conv_kernel_dim=int(kv("ssm.conv_kernel")),
+    )
+    hf_like = SimpleNamespace(
+        text_config=text,
+        model_type="qwen3_5_moe" if moe else "qwen3_5",
+        architectures=list(shim.architectures),
+        vision_config=None,
+    )
+    config = parse_config(hf_like)
+    return replace(config, expert_quant="gguf") if moe else config
+
+
+# ----------------------------------------------------------------------------------------------
+# Weights
+# ----------------------------------------------------------------------------------------------
+
+def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> torch.Tensor:
+    """Tiled value heads (k-head minor) back to grouped (k-head major) along ``dim``."""
+    n = t.shape[dim]
+    r = n // (num_k_heads * head_dim)
+    idx = torch.arange(n).view(r, num_k_heads, head_dim).transpose(0, 1).reshape(n)
+    return t.index_select(dim, idx)
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+    include_vision: bool = True,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    config = parse_gguf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    g = config.linear_attention_group()
+    hidden = config.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not config.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(config.num_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.post_attention_norm"))
+        if config.is_linear_layer(layer):
+            la = f"{pre}.linear_attn"
+            nk, dk = g.num_key_heads, g.key_head_dim
+            qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate")
+            qkv_rows = matrix(qkv)
+            v_start = 2 * nk * dk
+            qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)])
+            z_rows = _untile(matrix(gate), nk, g.value_head_dim)
+            beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha")
+            beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info))
+            if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type):
+                yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden)
+                yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden)
+            else:
+                yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden)
+            conv = float_tensor(t(f"{blk}.ssm_conv1d")).float()
+            conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)])
+            yield f"{la}.conv1d.weight", conv.unsqueeze(1)
+            yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1)
+            yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1)
+            yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm"))
+            yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim)
+        else:
+            sa = f"{pre}.self_attn"
+            parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+            yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden)
+            out = t(f"{blk}.attn_output")
+            yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+            yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+            yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+
+        mlp = f"{pre}.mlp"
+        if config.moe_enabled:
+            yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+            sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down"))
+            yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden)
+            yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size)
+            yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1)
+        else:
+            fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+            yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden)
+            yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size)
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py
index 8cdd41ae35..3fe373ca3d 100644
--- a/python/freetoken/models/qwen3_5_moe/model.py
+++ b/python/freetoken/models/qwen3_5_moe/model.py
@@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
                 layer_id=layer_id,
                 quant_config=config.quant,
                 prefix=f"{prefix}.linear_attn",
+                # llama.cpp stores every qwen35 ssm_out with tiled value-head columns
+                out_heads_tiled=config.quant is not None and config.quant.dialect == "gguf",
             )
         else:
             self.self_attn = Qwen3_5Attention(config, layer_id, prefix=f"{prefix}.self_attn")
@@ -79,6 +81,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 4b0cea25a9..86e3a734c8 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -201,6 +201,25 @@ class ModelSpec:
         packed_modules_mapping=_QWEN3_5_PACKED,
         unquantized_modules=_QWEN3_5_UNQUANTIZED,
     ),
+    # llama.cpp qwen35 / qwen35moe GGUF (text only): the same model classes with every tensor in its packed ggml type
+    "Qwen3_5GGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen3_5_moe",
+        "Qwen3_5ForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        checkpoint_roots=_LANGUAGE_MODEL_ROOT,
+        packed_modules_mapping=_QWEN3_5_PACKED,
+        unquantized_modules=_QWEN3_5_UNQUANTIZED,
+    ),
+    "Qwen3_5MoeGGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen3_5_moe",
+        "Qwen3_5MoeForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        checkpoint_roots=_LANGUAGE_MODEL_ROOT,
+        packed_modules_mapping=_QWEN3_5_PACKED,
+        unquantized_modules=_QWEN3_5_UNQUANTIZED,
+    ),
     # Qwen3.8-Flash-Next (model_type qwen4_exp): multimodal wrapper config (text tower in
     # text_config, weights under model.language_model.). 36 GDN + 12 QSA
     # compressed-sparse attention layers on 4 hyper-connection residual streams, a PLE

From 3381b6eb0a9e6bd4736e5c415953d6c881c734b5 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 23:20:12 -0700
Subject: [PATCH 104/249] feat(gguf): run experts whose gate_up and down mix
 float and quantized types (#702)

* feat(gguf): run experts whose gate_up and down mix float and quantized types

* fix(kernel): map an fp32 compute type to tl.float32 in fused_moe_kernel_triton

* fix(gguf): write float expert GEMM output in the activation dtype
---
 python/freetoken/kernel/moe_impl.py           |  7 +-
 .../freetoken/layers/quantization/moe/gguf.py | 66 ++++++++++++++-----
 2 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/python/freetoken/kernel/moe_impl.py b/python/freetoken/kernel/moe_impl.py
index 17fd77c4ed..488f8c2bdf 100644
--- a/python/freetoken/kernel/moe_impl.py
+++ b/python/freetoken/kernel/moe_impl.py
@@ -34,7 +34,12 @@ def fused_moe_kernel_triton(
         even_Ks = True
     else:
         even_Ks = False
-    dtype = tl.bfloat16 if compute_type == torch.bfloat16 else tl.float16
+    if compute_type == torch.float32:
+        dtype = tl.float32
+    elif compute_type == torch.bfloat16:
+        dtype = tl.bfloat16
+    else:
+        dtype = tl.float16
     fused_moe_kernel[grid](
         A,
         B,
diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py
index 78e051da67..dfd545bee6 100644
--- a/python/freetoken/layers/quantization/moe/gguf.py
+++ b/python/freetoken/layers/quantization/moe/gguf.py
@@ -32,11 +32,39 @@ def expert_types(cfg: MoEConfig) -> tuple[int, int]:
     gate, up, down = gguf_types(cfg.scheme)
     if gate != up:
         raise NotImplementedError(f"GGUF experts with gate {type_name(gate)} and up {type_name(up)} need one gate_up bank type")
-    if is_quantized(gate) != is_quantized(down):
-        raise NotImplementedError(f"GGUF experts with gate_up {type_name(gate)} and down {type_name(down)} mix float and quantized rows")
     return gate, down
 
 
+def _bank_spec(ggml_type: int, rows: int, k: int) -> BankSpec:
+    dtype = float_dtype(ggml_type)
+    return BankSpec((rows, k), dtype) if dtype is not None else BankSpec((rows, row_bytes(ggml_type, k)), torch.uint8)
+
+
+_DECODE_CONFIG = {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "num_warps": 8}
+
+
+def float_mul_mat_id(w: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, topk_weights: torch.Tensor, config: dict | None) -> torch.Tensor:
+    """``ggml.mul_mat_id`` for float experts ``w[E, N, K]``: same ``x`` and ``ids``, same ``[T, k, N]`` output in ``x``'s dtype, no router weights.
+
+    ``config`` runs the sorted-route kernel (prefill, resident layers); None runs the per-route decode kernel."""
+    from freetoken.kernel import fused_moe_decode_kernel_triton, fused_moe_kernel_triton
+    from freetoken.moe.fused import moe_align_block_size
+
+    t, top_k = ids.shape
+    per_route = x.dim() == 3
+    # f16 experts take f16 activations, as llama.cpp's f16 path does
+    a = x.reshape(-1, x.shape[-1]).to(w.dtype).contiguous()
+    out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device)
+    if config is None:
+        fused_moe_decode_kernel_triton(a, w, out, topk_weights, ids, False, per_route, _DECODE_CONFIG, compute_type=x.dtype)
+    else:
+        sorted_ids, expert_ids, num_post_padded = moe_align_block_size(ids, config["BLOCK_SIZE_M"], w.shape[0])
+        fused_moe_kernel_triton(
+            a, w, out, topk_weights, ids, sorted_ids, expert_ids, num_post_padded, False, 1 if per_route else top_k, config, compute_type=x.dtype
+        )
+    return out
+
+
 
 class GgmlMoEKernel(MoEKernel):
     name = "ggml"
@@ -52,13 +80,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         gate_up, down = expert_types(cfg)
         i, h = cfg.local_intermediate, cfg.hidden
-        dtype = float_dtype(gate_up)
-        if dtype is not None:
-            return {"gate_up": BankSpec((2 * i, h), dtype), "down": BankSpec((h, i), dtype)}
-        return {
-            "gate_up": BankSpec((2 * i, row_bytes(gate_up, h)), torch.uint8),
-            "down": BankSpec((h, row_bytes(down, i)), torch.uint8),
-        }
+        return {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)}
 
     def pack(self, pieces, cfg: MoEConfig, out):
         i = cfg.local_intermediate
@@ -71,23 +93,37 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         from freetoken.kernel import ggml
         from freetoken.layers.activation import gated_act_and_mul
 
-        gate_up_type, down_type = expert_types(layer.quant_method.cfg)
-        if not is_quantized(gate_up_type):
+        cfg = layer.quant_method.cfg
+        gate_up_type, down_type = expert_types(cfg)
+        if not is_quantized(gate_up_type) and not is_quantized(down_type):
             return self._apply_float(layer, x, topk_weights, topk_ids, view, gate_up_type, is_prefill=is_prefill)
         assert not layer.apply_router_weight_on_input, "GGUF experts apply the router weights on the output"
         ids = topk_ids.to(torch.int32)
         gate_up = view.tensors["gate_up"]
+        # a float half runs on the bf16 fused MoE's grouped GEMM, with its config for this layer's shape
+        config = None
+        if not (is_quantized(gate_up_type) and is_quantized(down_type)) and (is_prefill or is_resident(layer)):
+            from freetoken.moe.fused import try_get_optimal_moe_config
+
+            e, inter_size, hidden = gate_up.shape[0], cfg.local_intermediate, cfg.hidden
+            config = try_get_optimal_moe_config((e, 2 * inter_size, hidden), (e, hidden, inter_size), ids.shape[1], x.shape[0])
+
+        def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor:
+            if is_quantized(ggml_type):
+                return ggml.mul_mat_id(w, ggml_type, a, ids)
+            return float_mul_mat_id(w, a, ids, topk_weights, config)
+
         glu = fused_glu(layer)
-        small = x.shape[0] <= min(ggml.mmvq_max_tokens(gate_up_type), ggml.mmvq_max_tokens(down_type))
-        if glu is not None and small:
+        small = x.shape[0] <= min(ggml.mmvq_max_tokens(t) for t in (gate_up_type, down_type) if is_quantized(t))
+        if glu is not None and small and is_quantized(gate_up_type):
             inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu)
         else:
-            h = ggml.mul_mat_id(gate_up, gate_up_type, x, ids)
+            h = mul_mat_id(gate_up, gate_up_type, x)
             t, k, two_i = h.shape
             inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device)
             gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer))
             del h  # free the gate_up output before the down GEMM allocates its own
-        y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids)
+        y = mul_mat_id(view.tensors["down"], down_type, inter)
         return ggml.moe_weighted_sum(y, topk_weights)
 
     @staticmethod

From 9b68e0076769650bb51b92f035ec05488b92a4f0 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 15 Oct 2026 23:23:37 -0700
Subject: [PATCH 105/249] fix(engine): upload cuda graphs at capture so a
 stream wait cannot block the first launch (#697)

---
 python/freetoken/engine/graph.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index eb9684f609..a995316ddc 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import ctypes
+import functools
 import gc
+import sys
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List
 
@@ -102,6 +105,30 @@ def get_free_memory(device: torch.device) -> int:
     return torch.cuda.mem_get_info(device)[0]
 
 
+@functools.cache
+def _cu_graph_upload():
+    if sys.platform != "linux" or torch.version.hip is not None:
+        return None
+    fn = ctypes.CDLL("libcuda.so.1").cuGraphUpload
+    fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+    fn.restype = ctypes.c_int
+    return fn
+
+
+def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> None:
+    """Upload the graph at capture time.
+
+    Otherwise the first replay pushes every node behind the graph's own stream WAIT, and a large graph then blocks cudaGraphLaunch until the WAIT is released. PLE wait-sync releases it only after the launch returns, so that deadlocks.
+    """
+    upload = _cu_graph_upload()
+    if upload is None:
+        return
+    rc = upload(graph.raw_cuda_graph_exec(), stream.cuda_stream)
+    if rc != 0:
+        raise RuntimeError(f"cuGraphUpload failed with CUresult {rc}")
+    stream.synchronize()
+
+
 class GraphRunner:
     def __init__(
         self,
@@ -192,6 +219,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel
                 # CUDA graph capture to replay cold-cache expert copies.
                 with torch.cuda.graph(graph, pool=pool, stream=self.stream):
                     self.buffer.logits[:bs] = model.forward()
+                _upload_graph(graph, self.stream)
                 self._reset_moe_offload_cache()
             if pool is None:
                 pool = graph.pool()  # reuse cuda graph handle to reduce memory

From 6c21dd56157f43172ce8e75f51353f71b719abd0 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 13:04:29 -0700
Subject: [PATCH 106/249] feat(gguf): serve Qwen2, Qwen3 and Qwen3 MoE GGUF
 checkpoints (#707)

---
 python/freetoken/kernel/aot_models.py         |   9 +-
 python/freetoken/models/gguf/config.py        |   3 +
 python/freetoken/models/qwen2/__init__.py     |   3 +-
 python/freetoken/models/qwen2/gguf.py         | 144 +++++++++++++++++
 python/freetoken/models/qwen2/model.py        |   2 +
 python/freetoken/models/qwen3/__init__.py     |   3 +-
 python/freetoken/models/qwen3/gguf.py         | 143 +++++++++++++++++
 python/freetoken/models/qwen3/model.py        |   2 +
 python/freetoken/models/qwen3_moe/__init__.py |   4 +
 python/freetoken/models/qwen3_moe/gguf.py     | 149 ++++++++++++++++++
 python/freetoken/models/qwen3_moe/model.py    |   2 +
 python/freetoken/models/register.py           |  22 +++
 12 files changed, 480 insertions(+), 6 deletions(-)
 create mode 100644 python/freetoken/models/qwen2/gguf.py
 create mode 100644 python/freetoken/models/qwen3/gguf.py
 create mode 100644 python/freetoken/models/qwen3_moe/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index 308c2d6810..8c5d32ba9d 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -142,7 +142,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         moe_intermediate_size=768,
         expert_formats=("bf16",),
         aliases=("Qwen/Qwen3-30B-A3B-Thinking-2507",),
-        arch_aliases=("Qwen3VLMoeForConditionalGeneration",),  # Qwen3-VL-30B-A3B: same text tower
+        arch_aliases=("Qwen3VLMoeForConditionalGeneration", "Qwen3MoeGGUFForCausalLM"),  # Qwen3-VL-30B-A3B: same text tower
     ),
     AotModel(
         name="Qwen/Qwen3.5-35B-A3B",
@@ -152,7 +152,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         top_k=8,
         moe_intermediate_size=512,
         expert_formats=("bf16",),
-        arch_aliases=("Qwen3_5MoeForCausalLM",),  # text-only release of the same tower
+        arch_aliases=("Qwen3_5MoeForCausalLM", "Qwen3_5MoeGGUFForCausalLM"),  # text-only release of the same tower
     ),
     AotModel(
         name="Qwen/Qwen3.5-35B-A3B-FP8",
@@ -348,7 +348,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         hidden_size=5120,
         kv_groups=((4, 256),),
         aliases=("Qwen/Qwen3.6-27B-FP8", "nvidia/Qwen3.6-27B-NVFP4"),
-        arch_aliases=("Qwen3_5ForCausalLM",),  # text-only release of the same tower
+        arch_aliases=("Qwen3_5ForCausalLM", "Qwen3_5GGUFForCausalLM"),  # text-only release of the same tower
     ),
     AotModel(
         name="google/gemma-4-12B-it",
@@ -381,13 +381,14 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         architecture="Qwen2ForCausalLM",
         hidden_size=3584,
         kv_groups=((4, 128),),
+        arch_aliases=("Qwen2GGUFForCausalLM",),
     ),
     AotModel(
         name="Qwen/Qwen3-8B",
         architecture="Qwen3ForCausalLM",
         hidden_size=4096,
         kv_groups=((8, 128),),
-        arch_aliases=("Qwen3VLForConditionalGeneration",),  # Qwen3-VL-8B: same text tower
+        arch_aliases=("Qwen3VLForConditionalGeneration", "Qwen3GGUFForCausalLM"),  # Qwen3-VL-8B: same text tower
     ),
     AotModel(
         name="mistralai/Mistral-7B-Instruct-v0.3",
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index f9210938fb..518acfe4a6 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -18,6 +18,9 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "qwen2": "Qwen2GGUFForCausalLM",
+    "qwen3": "Qwen3GGUFForCausalLM",
+    "qwen3moe": "Qwen3MoeGGUFForCausalLM",
     "qwen35": "Qwen3_5GGUFForCausalLM",
     "qwen35moe": "Qwen3_5MoeGGUFForCausalLM",
     "qwen4exp": "Qwen4ExpGGUFForCausalLM",
diff --git a/python/freetoken/models/qwen2/__init__.py b/python/freetoken/models/qwen2/__init__.py
index 43ca84f82a..3917c188fd 100644
--- a/python/freetoken/models/qwen2/__init__.py
+++ b/python/freetoken/models/qwen2/__init__.py
@@ -1,5 +1,6 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen2ForCausalLM
 from .weight import iter_weights
 
-__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights"]
+__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen2/gguf.py b/python/freetoken/models/qwen2/gguf.py
new file mode 100644
index 0000000000..501c113a3a
--- /dev/null
+++ b/python/freetoken/models/qwen2/gguf.py
@@ -0,0 +1,144 @@
+"""Qwen2 / Qwen2.5 from a llama.cpp ``qwen2`` GGUF: config from the GGUF metadata, tensors by name.
+
+The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen2"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"}
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling == "yarn":
+        rope.update(
+            rope_type="yarn",
+            factor=float(kv("rope.scaling.factor")),
+            original_max_position_embeddings=int(kv("rope.scaling.original_context_length")),
+        )
+    elif scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    # Qwen2.5 files carry no key_length; parse_config then derives head_dim from the hidden size
+    key_length = md.get(f"{_ARCH}.attention.key_length")
+    return SimpleNamespace(
+        model_type="qwen2",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(key_length) if key_length is not None else None,
+        rope_scaling=rope,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        bias = index.get(f"{blk}.attn_qkv.bias")
+        yield f"{sa}.qkv_proj.bias", float_tensor(bias) if bias is not None else torch.cat([float_tensor(t(f"{blk}.attn_{p}.bias")) for p in ("q", "k", "v")])
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen2/model.py b/python/freetoken/models/qwen2/model.py
index ed56fde25a..94bc2beb6b 100644
--- a/python/freetoken/models/qwen2/model.py
+++ b/python/freetoken/models/qwen2/model.py
@@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/qwen3/__init__.py b/python/freetoken/models/qwen3/__init__.py
index 71bda9781f..fef7c8fb99 100644
--- a/python/freetoken/models/qwen3/__init__.py
+++ b/python/freetoken/models/qwen3/__init__.py
@@ -1,5 +1,6 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen3ForCausalLM
 from .weight import iter_weights
 
-__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights"]
+__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen3/gguf.py b/python/freetoken/models/qwen3/gguf.py
new file mode 100644
index 0000000000..06245ce4c2
--- /dev/null
+++ b/python/freetoken/models/qwen3/gguf.py
@@ -0,0 +1,143 @@
+"""Qwen3 dense from a llama.cpp ``qwen3`` GGUF: config from the GGUF metadata, tensors by name.
+
+The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen3"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"}
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling == "yarn":
+        rope.update(
+            rope_type="yarn",
+            factor=float(kv("rope.scaling.factor")),
+            original_max_position_embeddings=int(kv("rope.scaling.original_context_length")),
+        )
+    elif scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    key_length = md.get(f"{_ARCH}.attention.key_length")
+    return SimpleNamespace(
+        model_type="qwen3",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(key_length) if key_length is not None else None,
+        rope_scaling=rope,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+        yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen3/model.py b/python/freetoken/models/qwen3/model.py
index 6ac7d2fa8f..ae0d4394f5 100644
--- a/python/freetoken/models/qwen3/model.py
+++ b/python/freetoken/models/qwen3/model.py
@@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/qwen3_moe/__init__.py b/python/freetoken/models/qwen3_moe/__init__.py
index 59c74fb0cb..cb96a1921b 100644
--- a/python/freetoken/models/qwen3_moe/__init__.py
+++ b/python/freetoken/models/qwen3_moe/__init__.py
@@ -1,4 +1,5 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen3MoeForCausalLM
 from .weight import iter_weights, iter_weights_parallel
 
@@ -7,4 +8,7 @@
     "parse_config",
     "iter_weights",
     "iter_weights_parallel",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py
new file mode 100644
index 0000000000..7bf016292e
--- /dev/null
+++ b/python/freetoken/models/qwen3_moe/gguf.py
@@ -0,0 +1,149 @@
+"""Qwen3 MoE from a llama.cpp ``qwen3moe`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
+
+The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen3moe"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate": "ffn_gate_inp",
+    "mlp.experts.0.gate_proj": "ffn_gate_exps",
+    "mlp.experts.0.up_proj": "ffn_up_exps",
+    "mlp.experts.0.down_proj": "ffn_down_exps",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors:
+        return (f"{blk}.ffn_gate_up_exps",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"}
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling == "yarn":
+        rope.update(
+            rope_type="yarn",
+            factor=float(kv("rope.scaling.factor")),
+            original_max_position_embeddings=int(kv("rope.scaling.original_context_length")),
+        )
+    elif scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    key_length = md.get(f"{_ARCH}.attention.key_length")
+    return SimpleNamespace(
+        model_type="qwen3_moe",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(key_length) if key_length is not None else None,
+        rope_scaling=rope,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+        num_experts=int(kv("expert_count")),
+        num_experts_per_tok=int(kv("expert_used_count")),
+        moe_intermediate_size=int(kv("expert_feed_forward_length")),
+        # llama.cpp always renormalizes the top-k weights of a qwen3moe router
+        norm_topk_prob=True,
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return replace(parse_config(_hf_config(shim)), expert_quant="gguf")
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+        yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        yield f"{pre}.mlp.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/qwen3_moe/model.py b/python/freetoken/models/qwen3_moe/model.py
index 5d89656bda..7d81dd2673 100644
--- a/python/freetoken/models/qwen3_moe/model.py
+++ b/python/freetoken/models/qwen3_moe/model.py
@@ -47,6 +47,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 86e3a734c8..263b5f932c 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -110,6 +110,28 @@ class ModelSpec:
         "Qwen3MoeForCausalLM",
         packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED,
     ),
+    # llama.cpp qwen2 / qwen3 / qwen3moe GGUF: the same model classes with every tensor in its packed ggml type
+    "Qwen2GGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen2",
+        "Qwen2ForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_DENSE_PACKED,
+    ),
+    "Qwen3GGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen3",
+        "Qwen3ForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_DENSE_PACKED,
+    ),
+    "Qwen3MoeGGUFForCausalLM": ModelSpec(
+        "freetoken.models.qwen3_moe",
+        "Qwen3MoeForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED,
+    ),
     # Qwen3-VL: the Qwen3 text tower under model.language_model. plus the shared Qwen VL vision tower with DeepStack; the MoE variant ships its experts pre-stacked.
     "Qwen3VLForConditionalGeneration": ModelSpec(
         "freetoken.models.qwen3_vl",

From 3886ceeebf25a0075ea83c7629a19f950a03b3be Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 13:05:08 -0700
Subject: [PATCH 107/249] feat(gguf): serve Gemma 4 dense GGUF checkpoints
 (#706)

---
 python/freetoken/models/gemma4/gguf.py | 33 ++++++++++++++++++--------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py
index 95f89c3ca4..2b86cca127 100644
--- a/python/freetoken/models/gemma4/gguf.py
+++ b/python/freetoken/models/gemma4/gguf.py
@@ -56,8 +56,11 @@ def g(key: str):
             raise KeyError(f"missing GGUF metadata key gemma4.{key}")
         return val
 
+    if int(m.get("gemma4.embedding_length_per_layer_input", 0)) or int(m.get("gemma4.attention.shared_kv_layers", 0)):
+        raise NotImplementedError("gemma4 E2B / E4B GGUF (per-layer input embeddings, shared KV layers) is not supported")
     num_layers = int(g("block_count"))
     hidden = int(g("embedding_length"))
+    num_experts = int(m.get("gemma4.expert_count", 0))
     num_qo_heads = int(g("attention.head_count"))
     kv_per_layer = g("attention.head_count_kv")  # per-layer list
     # True -> sliding-window (SWA) layer, False -> full attention.
@@ -104,14 +107,14 @@ def g(key: str):
         rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")),
         tie_word_embeddings=bool(shim.tie_word_embeddings),
         rotary_config=full_rotary,
-        num_experts=int(g("expert_count")),
-        num_experts_per_tok=int(g("expert_used_count")),
-        moe_intermediate_size=int(g("expert_feed_forward_length")),
+        num_experts=num_experts,
+        num_experts_per_tok=int(g("expert_used_count")) if num_experts else 0,
+        moe_intermediate_size=int(g("expert_feed_forward_length")) if num_experts else 0,
         norm_topk_prob=True,
         model_type="gemma4",
         architectures=list(shim.architectures),
-        moe_enabled=True,
-        expert_quant="gguf",
+        moe_enabled=num_experts > 0,
+        expert_quant="gguf" if num_experts else "none",
         use_qk_norm=True,
         attn_sm_scale=1.0,
         final_logit_softcapping=float(g("final_logit_softcapping")),
@@ -149,11 +152,14 @@ def g(key: str):
     "attn_k_norm.weight": "self_attn.k_norm.weight",
     "post_attention_norm.weight": "post_attention_layernorm.weight",
     "ffn_norm.weight": "pre_feedforward_layernorm.weight",
-    "pre_ffw_norm_2.weight": "feed_forward.pre_feedforward_layernorm_2.weight",
     "post_ffw_norm.weight": "feed_forward.post_feedforward_layernorm.weight",
+    "layer_output_scale.weight": "feed_forward.layer_scalar",
+}
+# the MoE layers' second feed-forward branch and router
+_MOE_LAYER_SCALAR_MAP = {
+    "pre_ffw_norm_2.weight": "feed_forward.pre_feedforward_layernorm_2.weight",
     "post_ffw_norm_1.weight": "feed_forward.post_feedforward_layernorm_1.weight",
     "post_ffw_norm_2.weight": "feed_forward.post_feedforward_layernorm_2.weight",
-    "layer_output_scale.weight": "feed_forward.layer_scalar",
     "ffn_gate_inp.weight": "feed_forward.router.proj.weight",
     "ffn_gate_inp.scale": "feed_forward.router.scale",
     "ffn_down_exps.scale": "feed_forward.router.per_expert_scale",
@@ -184,6 +190,8 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
     if m is None:
         return ()
     blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
     if leaf == "self_attn.v_proj":
         # the full-attention layers store no attn_v: k is reused as v
         return (f"{blk}.attn_v",) if f"{blk}.attn_v.weight" in tensors else (f"{blk}.attn_k",)
@@ -233,13 +241,18 @@ def t(name: str):
         yield from fused("lm_head", [out], [matrix(out)], hidden)
     yield "model.norm.weight", float_tensor(t("output_norm"))
 
+    scalar_map = {**_LAYER_SCALAR_MAP, **_MOE_LAYER_SCALAR_MAP} if config.moe_enabled else _LAYER_SCALAR_MAP
     for layer in range(config.num_layers):
         blk, base = f"blk.{layer}", f"model.layers.{layer}"
-        for suffix, rel in _LAYER_SCALAR_MAP.items():
+        for suffix, rel in scalar_map.items():
             tensor = float_tensor(index[f"{blk}.{suffix}"])
             yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor
-        q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k")
-        qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)]
+        if f"{blk}.attn_qkv.weight" in index:
+            # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+            qkv = [t(f"{blk}.attn_qkv")]
+        else:
+            q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k")
+            qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)]
         yield from fused(f"{base}.self_attn.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
         o = t(f"{blk}.attn_output")
         yield from fused(f"{base}.self_attn.o_proj", [o], [matrix(o)], int(o.ne[0]))

From d5eb18febdcfe3f83b43c912cccc82e16cc76e8f Mon Sep 17 00:00:00 2001
From: LZ_QAQ <32613092+LZ-QWQ@users.noreply.github.com>
Date: Sat, 17 Oct 2026 04:49:17 +0800
Subject: [PATCH 108/249] fix(kernel): correct fp4 cache scale rounding (#720)

* fix(kernel): preserve fp4 scale rounding at subnormal ties

* fix(kernel): preserve native e4m3 scale calculation
---
 python/freetoken/kernel/triton/dsv41/pack.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/kernel/triton/dsv41/pack.py b/python/freetoken/kernel/triton/dsv41/pack.py
index 47a9164e6e..c0fca83162 100644
--- a/python/freetoken/kernel/triton/dsv41/pack.py
+++ b/python/freetoken/kernel/triton/dsv41/pack.py
@@ -87,7 +87,12 @@ def _pack_rows_kernel(
         amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1))
         if FMT == 2:
             amax = tl.maximum(amax, 6.0 * (2.0 ** -9))
-            s = _e4m3_round(amax * (1.0 / 6.0))
+            if e4m3_native_cx():
+                s = amax * (1.0 / 6.0)
+            else:
+                # Round before the emulated magic-add; fusing it changes subnormal ties.
+                s = tl.math.div_rn(amax, 6.0)
+            s = _e4m3_round(s)
             tl.store(dst + D // 2 + offs_g, _e4m3_codes(s))
         else:
             amax = tl.maximum(amax, 6.0 * (2.0 ** -126))

From e8c6bf11235d6f813b4aa4e15968e001f881aaec Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 15:13:59 -0700
Subject: [PATCH 109/249] feat(moe): serve experts of several formats through
 per-format slot pools (#700)

* feat(moe): serve experts of several formats through per-format slot pools

* refactor(cpu-moe): pass each layer's cpu format with its options

* fix(moe): read only each format's layers when loading experts of several formats

---------

Co-authored-by: andy-yang-1 <andy_yang@berkeley.edu>
---
 python/freetoken/checkpoint/convert.py        |  14 +-
 python/freetoken/checkpoint/ftw.py            |  31 +-
 python/freetoken/engine/engine.py             | 113 +++++--
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       |  65 ++--
 .../freetoken/layers/quantization/moe/base.py |   8 +
 .../layers/quantization/moe/mxfp4.py          |   3 +
 python/freetoken/models/deepseek_v4/weight.py |   7 +-
 .../freetoken/models/deepseek_v41/weight.py   |   7 +-
 python/freetoken/models/glm5_next/weight.py   |  12 +-
 python/freetoken/models/gpt_oss/weight.py     |   7 +-
 python/freetoken/models/nvfp4_banks.py        |   9 +-
 python/freetoken/models/qwen3_5_moe/weight.py |   9 +-
 python/freetoken/moe/benchbw.py               |   2 +-
 python/freetoken/moe/cpu_executor.py          | 186 ++++++-----
 python/freetoken/moe/expert_banks.py          | 125 ++++++--
 python/freetoken/moe/expert_pieces.py         |  45 ++-
 python/freetoken/moe/host_banks.py            |  21 +-
 python/freetoken/moe/offload_cache.py         |  42 ++-
 tests/models/test_models_loader.py            |  19 ++
 tests/models/test_qwen3_5_moe_weight.py       |   4 +
 tests/moe/test_cpu_moe.py                     |   4 +-
 tests/moe/test_dsfp4_prequant.py              |   2 +-
 tests/moe/test_dsfp4_reference_semantics.py   |   4 +-
 tests/moe/test_mixed_expert_formats.py        | 288 ++++++++++++++++++
 24 files changed, 808 insertions(+), 219 deletions(-)
 create mode 100644 tests/moe/test_mixed_expert_formats.py

diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py
index 054fade4eb..5551759f85 100644
--- a/python/freetoken/checkpoint/convert.py
+++ b/python/freetoken/checkpoint/convert.py
@@ -203,14 +203,14 @@ def convert_checkpoint(
     # every MoE strategy reads the routed experts from the banks, fused included
     offload = getattr(mc, "is_moe", False)
     include_moe_experts = not offload
-    method = None
+    method = layer_methods = None
     if offload:
-        from freetoken.engine.engine import offload_expert_method
+        from freetoken.engine.engine import offload_expert_methods
         from freetoken.layers import set_rope_device
 
         set_rope_device(dev)
         object.__setattr__(cfg, "moe_strategy", "offload")
-        method = offload_expert_method(cfg)
+        method, layer_methods = offload_expert_methods(cfg)
 
     from freetoken.utils.progress import byte_bar, count_bar
 
@@ -239,6 +239,7 @@ def convert_checkpoint(
 
     # 2) offload expert banks (post-repack) + alpha scales (slow path auto-picks parallel/serial)
     quant_format = None
+    layer_quant_formats = None
     num_layers = None
     if offload:
         # every method-packed format streams each layer to its own FTW entry as it completes (via the sink); the GGUF provider reports through ExpertBanks.streamed whether it engaged the sink or materialized the whole bank set first
@@ -247,9 +248,11 @@ def convert_checkpoint(
         names = legacy_bank_names(legacy_format_for(method.kind, method.kernel.name)) if method is not None else {}
         sink = _ConvertSink(writer, names=names)
         banks = load_expert_banks(
-            model_path, mc, method=method, device=dev, layer_sink=sink
+            model_path, mc, method=method, device=dev, layer_sink=sink, layer_methods=layer_methods
         )
         quant_format = banks.quant_format
+        layer_quant_formats = banks.layer_formats
+        assert layer_quant_formats is None or banks.streamed, "experts of several formats always stream to the sink"
         if banks.streamed:
             sink.close()
             num_layers = sink.num_layers  # however many distinct layers the sink actually saw
@@ -313,6 +316,9 @@ def convert_checkpoint(
         # read back at load (ftw.load_ftw_banks). dtype/moe_backend were dropped: each
         # tensor already carries its own dtype, and nothing reads a model-level backend.
         "quant_format": quant_format,
+        # Experts of several formats: each MoE layer's quant_format (quant_format above is the most common). A layer only
+        # stores the banks its kernel reads; the reader gives it zero-byte placeholders for the others.
+        "layer_quant_formats": layer_quant_formats,
         # The reader takes num_layers from the model config (copied into this
         # checkpoint); recording it here too gives load_ftw_banks a cross-check that
         # the banks match the config they ship with. None for non-offload checkpoints.
diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index e31e9d8278..ae97a98d5c 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -461,11 +461,18 @@ def _split_bank_entries(reader: FTWReader, path: str, num_layers: int):
 
     mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys()
     assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+    # experts of several formats: a layer stores only the banks its kernel reads
+    several = reader.meta("layer_quant_formats") is not None
+    assert not (several and flat_entries), "experts of several formats are stored per layer"
     for base, by_layer in per_layer_groups.items():
-        assert sorted(by_layer) == list(range(num_layers)), (
-            f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
+        layers = sorted(by_layer)
+        assert (set(layers) <= set(range(num_layers))) if several else layers == list(range(num_layers)), (
+            f"FTW bank {base!r} has per-layer entries for layers {layers}, "
             f"expected exactly range({num_layers})"
         )
+    if several:
+        covered = set().union(*(by_layer.keys() for by_layer in per_layer_groups.values()))
+        assert covered == set(range(num_layers)), f"FTW expert banks have no entries for layers {sorted(set(range(num_layers)) - covered)}"
     return bank_entries, alpha_entries, flat_entries, per_layer_groups
 
 
@@ -562,11 +569,16 @@ def _backing(layer_id: int) -> str:
             row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype))
             row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
 
+    num_experts = next((e["shape"][0] for by_layer in per_layer_groups.values() for e in by_layer.values()), None)
     for base, by_layer in per_layer_groups.items():
         row_hb[base] = []
         row_view_args[base] = []
         for layer_id in range(num_layers):
-            e = by_layer[layer_id]
+            e = by_layer.get(layer_id)
+            if e is None:  # experts of several formats: this layer's kernel reads no such bank
+                row_hb[base].append(None)
+                row_view_args[base].append("placeholder")
+                continue
             assert e["global_off"] % ALIGN == 0, (base, layer_id, e["global_off"])  # writer invariant
             bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id))
             row_hb[base].append(bank)
@@ -611,10 +623,15 @@ def _read_layer(job):
         bar.close()
         reader.close()
 
+    from freetoken.moe.offload_cache import placeholder_bank
+
     sources: dict[str, list] = {}
     for name, banks in row_hb.items():
         views = []
         for bank, view_args in zip(banks, row_view_args[name]):
+            if view_args == "placeholder":
+                views.append(placeholder_bank(num_experts))
+                continue
             if view_args is None:  # per-layer entry: already shaped [num_experts, ...]
                 views.append(bank.tensor)
                 continue
@@ -635,6 +652,8 @@ def _read_layer(job):
     applied = list(residency)
     for banks in row_hb.values():
         for layer_id, bank in enumerate(banks):
+            if bank is None:
+                continue
             if (applied[layer_id] == HostResidency.LOCKED.value
                     and bank.residency is not HostResidency.LOCKED):
                 applied[layer_id] = HostResidency.PAGEABLE.value
@@ -643,7 +662,7 @@ def _read_layer(job):
         by_layer = [0] * num_layers
         for name, banks in row_hb.items():
             for layer_id, bank in enumerate(banks):
-                by_layer[layer_id] += bank.nbytes
+                by_layer[layer_id] += bank.nbytes if bank is not None else 0
         locked = [i for i in unpinned if applied[i] == HostResidency.LOCKED.value]
         pageable = [i for i in unpinned if i not in set(locked)]
         pinned_b = sum(b for i, b in enumerate(by_layer) if i not in set(unpinned))
@@ -666,9 +685,11 @@ def _read_layer(job):
     # alphas are the small per-expert scale vectors, distinguished by their reserved names
     # (not a separate kind); everything else under experts_bank is a weight source.
     alpha_kw = {n: alpha_hb[n].tensor for n in alpha_hb}
+    layer_formats = reader.meta("layer_quant_formats")
     return ExpertBanks(
         quant_format, sources, **alpha_kw,
         layer_residency=applied, kind=kind, kernel=kernel,
+        layer_tags=layer_formats, layer_formats=layer_formats,
     )
 
 
@@ -683,6 +704,8 @@ def load_ftw_banks_to_device(path: str, *, num_layers: int, device: torch.device
         split = _split_bank_entries(reader, path, num_layers)
         if split is None:
             return None
+        if reader.meta("layer_quant_formats") is not None:
+            raise ValueError(f"{path!r} holds experts of several formats, which are served from the offload slot cache only")
         bank_entries, alpha_entries, flat_entries, per_layer_groups = split
 
         # (name, layer_id or None for an alpha, read offset, read length, head pad, shape, dtype)
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 81f42dbadd..980bb043bd 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -606,6 +606,7 @@ def _load_resident_experts(self, config: EngineConfig, layers) -> None:
                 dummy=config.use_dummy_weight,
                 parallel={"serial": False, "parallel": True}.get(config.expert_load, None),
                 resident=True,
+                layer_methods=[layer.quant_method for layer in sorted(layers, key=lambda layer: layer.layer_id)],
             )
         attach_resident_banks(layers, banks)
 
@@ -670,10 +671,21 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None,
         )
 
     def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
+        from freetoken.moe.expert_banks import is_mixed
+
         method = shared_offload_method(self.model)
         num_moe_layers = config.model_config.num_moe_layers
+        # experts of several formats: one method per MoE layer, one slot pool per format
+        layer_methods = offload_layer_methods(self.model)
+        if not is_mixed(layer_methods):
+            layer_methods = None
+        elif config.moe_strategy == "cpu" or config.moe_cpu_layers:
+            raise ValueError(
+                "this model's experts come in several formats, which are served from per-format GPU slot pools; "
+                "they do not combine with --moe-strategy cpu or --moe-cpu-layers (use offload or hybrid)"
+            )
         cpu_layer_ids = _resolve_cpu_layers(config, num_moe_layers, reserved=self._host_tables_bytes, method=method)
-        _check_pin_budget(config, reserved=self._host_tables_bytes, method=method)
+        _check_pin_budget(config, reserved=self._host_tables_bytes, method=method, layer_methods=layer_methods)
         # the kernels were picked for model_config.decode_target; --moe-cpu-layers auto may still find that every bank fits the pin budget
         decode_target = config.model_config.decode_target
         if decode_target == "cpu" and not cpu_layer_ids:
@@ -737,6 +749,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                     dummy=config.use_dummy_weight,
                     parallel=expert_parallel,
                     layer_residency=requested_residency,
+                    layer_methods=layer_methods,
                 )
         except PinFailed as exc:
             raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc
@@ -755,7 +768,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 f"--moe-cache-split {config.moe_cache_split} is ignored: all MoE layers share one slot pool "
                 f"(--moe-cache-scope per-layer gives each layer its own)"
             )
-        self._moe_slot_limit = method.slot_limit() if method is not None else None
+        self._moe_slot_limit = _slot_limit(layer_methods) if layer_methods is not None else (method.slot_limit() if method is not None else None)
         # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step
         pools = self._moe_pools(1)
         size_cap = sum(p.max_num_slots for p in pools)
@@ -798,8 +811,17 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             )
             if decode_target == "gpu":
                 _clamp_decode_batch_to_pools(config, pool_sizes)
-        layout = max_slots = None
-        if method is not None:
+        layout = layer_layouts = max_slots = None
+        if layer_methods is not None:
+            _check_layer_formats(banks, layer_methods)
+            layer_layouts = [m.layout() for m in layer_methods]
+            max_slots = self._moe_slot_limit
+        elif method is not None:
+            if banks.layer_formats is not None:
+                raise ValueError(
+                    f"expert banks hold experts of several formats ({sorted(set(banks.layer_formats))}) but the model "
+                    f"binds {method.kind} / {method.kernel.name} on every MoE layer; reconvert the checkpoint"
+                )
             if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name):
                 raise ValueError(
                     f"expert banks were packed for {banks.kind} / {banks.kernel} but the model "
@@ -822,6 +844,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             decode_target=decode_target,
             hybrid_max_fetch=config.moe_hybrid_max_fetch,
             layout=layout,
+            layer_layouts=layer_layouts,
             max_slots=max_slots,
             pool_of_layer=pool_of_layer(pools, config.model_config.num_moe_layers),
             pool_sizes=pool_sizes,
@@ -903,10 +926,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         # Decode batches never exceed max_running_req, but CUDA-graph padding can
         # round a batch up to the largest captured size; cover both.
         max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1)
-        method = sample.quant_method
-        formats = {layer.quant_method.cpu_format for layer in layers}
-        if len(formats) > 1:
-            raise NotImplementedError(f"CPU MoE backend needs one expert format across the MoE layers, got {sorted(map(str, formats))}")
+        by_id = sorted(layers, key=lambda layer: layer.layer_id)
+        layer_formats = [layer.quant_method.cpu_spec() for layer in by_id]
         executor = CpuMoeExecutor(
             cache,
             top_k=sample.top_k,
@@ -917,8 +938,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
             device=self.device,
             swiglu_alpha=float(sample.alpha),
             swiglu_limit=sample.limit,
-            fmt=method.cpu_format,
-            act_block=method.cfg.act_block,
+            layer_formats=layer_formats,
         )
         cache.set_cpu_executor(executor)
         self.cpu_moe_executor = executor
@@ -1475,10 +1495,10 @@ def _pin_budget_bytes(reserved: int = 0) -> int | None:
     return max(0, cap - reserved)
 
 
-def _bank_bytes(config: EngineConfig, method=None) -> int | None:
+def _bank_bytes(config: EngineConfig, method=None, layer_methods=None) -> int | None:
     from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes
 
-    return ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config, method)
+    return ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config, method, layer_methods)
 
 
 def _pin_hint(reserved: int) -> str:
@@ -1490,12 +1510,12 @@ def _pin_hint(reserved: int) -> str:
     )
 
 
-def _check_pin_budget(config: EngineConfig, *, reserved: int, method=None) -> None:
+def _check_pin_budget(config: EngineConfig, *, reserved: int, method=None, layer_methods=None) -> None:
     """Stop a plain offload boot whose banks exceed a known pin budget before any bank is read."""
     if config.moe_cpu_layers or config.moe_strategy not in ("offload", "hybrid"):
         return
     budget = _pin_budget_bytes(reserved)
-    bank_bytes = _bank_bytes(config, method) if budget is not None else None
+    bank_bytes = _bank_bytes(config, method, layer_methods) if budget is not None else None
     if bank_bytes and bank_bytes > budget:
         raise ValueError(
             f"expert banks need {bank_bytes / 2**30:.1f} GiB of pinned host RAM but the pin budget is "
@@ -1574,27 +1594,58 @@ def _adjust_ftw_quant_backend(model_path: str, quant_backend: QuantBackend) -> Q
     return QuantBackend(quant_backend.items + (((LayerKind.MOE, kind), kernel),))
 
 
-def shared_offload_method(model):
-    """The expert method every offload MoE layer of ``model`` uses, or None for models whose MoE layers carry none (GGUF).
-
-    The offload cache holds one bank layout, so the layers must agree on (kind, kernel)."""
+def offload_layer_methods(model) -> list | None:
+    """The expert method of each offload MoE layer of ``model``, by MoE-layer id; None for models whose MoE layers carry none (GGUF)."""
     layers = [l for l in iter_offload_moe_layers(model) if getattr(l, "quant_method", None) is not None]
     if not layers:
         return None
-    keys = {(layer.quant_method.kind, layer.quant_method.kernel.name) for layer in layers}
-    if len(keys) != 1:
-        found = sorted(f"{kind} / {kernel}" for kind, kernel in keys)
-        raise ValueError(f"expert layers disagree on format / kernel: {found}; per-layer mixed expert formats are not supported")
-    kind, kernel = keys.pop()
-    logger.info_rank0(f"MoE experts: {kind} via {kernel}")
-    return layers[0].quant_method
+    layers.sort(key=lambda layer: layer.layer_id)
+    assert [layer.layer_id for layer in layers] == list(range(len(layers))), [layer.layer_id for layer in layers]
+    return [layer.quant_method for layer in layers]
 
 
-def offload_expert_method(config: EngineConfig):
-    """The offload expert method of ``config``'s model, from a meta-device build.
+def shared_offload_method(model):
+    """The expert method the offload MoE layers of ``model`` use, or None for models whose MoE layers carry none (GGUF).
+
+    With experts of several formats (a checkpoint that keeps some layers unquantized) this is the format of the most
+    layers; the cache then takes every layer's method (``offload_layer_methods``) and keeps one slot pool per format."""
+    from freetoken.moe.expert_banks import method_groups
+
+    methods = offload_layer_methods(model)
+    if methods is None:
+        return None
+    groups = method_groups(methods)
+    found = ", ".join(f"{m.kind} via {m.kernel.name}" + (f" on {len(ids)} layers" if len(groups) > 1 else "") for m, ids in groups)
+    logger.info_rank0(f"MoE experts: {found}")
+    return groups[0][0]
+
+
+def _slot_limit(layer_methods) -> int | None:
+    """The tightest slot cap of the layers' expert kernels (a pool never mixes kernels, but the limit is checked per pool against the largest)."""
+    limits = [limit for limit in (m.slot_limit() for m in layer_methods) if limit is not None]
+    return min(limits) if limits else None
+
+
+def _check_layer_formats(banks, layer_methods) -> None:
+    from freetoken.moe.legacy_format import legacy_format_for
+
+    want = [legacy_format_for(m.kind, m.kernel.name) for m in layer_methods]
+    have = banks.layer_formats or [banks.quant_format] * len(want)
+    bad = [layer for layer, (w, h) in enumerate(zip(want, have)) if w != h]
+    if bad:
+        raise ValueError(
+            f"expert banks of MoE layers {bad[:8]} were packed as {[have[l] for l in bad[:8]]} but the model binds "
+            f"{[want[l] for l in bad[:8]]}; reconvert the checkpoint or select those kernels"
+        )
+
+
+def offload_expert_methods(config: EngineConfig) -> tuple:
+    """``(method, layer_methods)`` of ``config``'s model, from a meta-device build: the most common expert method and,
+    for experts of several formats, one per MoE layer (else None).
 
     For tools that load expert banks without an engine (the FTW converter); they pack for GPU decode."""
     from freetoken.models import create_model
+    from freetoken.moe.expert_banks import is_mixed
     from freetoken.utils.torch_utils import torch_dtype
 
     set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend)))
@@ -1602,7 +1653,13 @@ def offload_expert_method(config: EngineConfig):
     object.__setattr__(config.model_config, "decode_target", "gpu")
     with torch.device("meta"), torch_dtype(config.dtype):
         model = create_model(config.model_config)
-    return shared_offload_method(model)
+    layer_methods = offload_layer_methods(model)
+    return shared_offload_method(model), (layer_methods if is_mixed(layer_methods) else None)
+
+
+def offload_expert_method(config: EngineConfig):
+    """The (most common) offload expert method of ``config``'s model; see ``offload_expert_methods``."""
+    return offload_expert_methods(config)[0]
 
 
 def _adjust_config(config: EngineConfig):
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index c8b4d1b465..6236816eec 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -21,6 +21,7 @@
 #include <stdexcept>
 #include <string>
 #include <thread>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -103,6 +104,9 @@ std::unique_ptr<MoeFormat> make_format(const std::string& name, const MoeParams&
   throw std::invalid_argument("unknown CPU MoE weight format " + name);
 }
 
+// One format instance as cpu_executor.py passes it: name, bank tables by role, integer options.
+using FormatSpec = std::tuple<std::string, std::map<std::string, uint64_t>, std::map<std::string, int64_t>>;
+
 struct CpuMoeExecutor {
   // A host-function callback gets one pointer, so it carries the executor with the task.
   struct HostTask {
@@ -112,10 +116,9 @@ struct CpuMoeExecutor {
 
   const int num_layers, num_experts;
   const bool apply_on_input;
-  MoeParams params;  // the format keeps a reference; declared before fmt
-  std::unique_ptr<MoeFormat> fmt;
-  bool preps_inter = false;
-  bool rounds_routes = false;
+  MoeParams params;  // the formats keep a reference; declared before fmts
+  std::vector<std::unique_ptr<MoeFormat>> fmts;
+  std::vector<int> layer_fmt;  // [num_layers], index into fmts
   int n_iblk = 0, n_hblk = 0;  // kRowBlock blocks per intermediate row and per output row
 
   // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w.
@@ -167,9 +170,8 @@ struct CpuMoeExecutor {
 
   CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_,
                  int hidden_size, int inter_size, int max_tokens_, int activation_id,
-                 int apply_router_weight_on_input, const std::string& weight_format,
-                 std::map<std::string, uint64_t> banks, std::map<std::string, int64_t> options,
-                 double swiglu_alpha_, double swiglu_limit_,
+                 int apply_router_weight_on_input, std::vector<FormatSpec> formats,
+                 std::vector<int> layer_formats, double swiglu_alpha_, double swiglu_limit_,
                  std::vector<int> core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr,
                  uintptr_t y_ptr)
       : num_layers(num_layers_),
@@ -184,9 +186,13 @@ struct CpuMoeExecutor {
         num_threads(num_threads_ > 0 ? num_threads_ : 1),
         core_ids(std::move(core_ids_)) {
     if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK));
-    fmt = make_format(weight_format, params, FormatArgs{weight_format, std::move(banks), std::move(options)});
-    preps_inter = fmt->preps_inter();
-    rounds_routes = fmt->rounds_routes();
+    for (auto& [name, banks, options] : formats)
+      fmts.push_back(make_format(name, params, FormatArgs{name, std::move(banks), std::move(options)}));
+    if (static_cast<int>(layer_formats.size()) != num_layers)
+      throw std::invalid_argument("CPU MoE needs the format of every layer");
+    for (const int f : layer_formats)
+      if (f < 0 || f >= static_cast<int>(fmts.size())) throw std::out_of_range("CPU MoE layer format index out of range");
+    layer_fmt = std::move(layer_formats);
     n_iblk = (params.I + kRowBlock - 1) / kRowBlock;
     n_hblk = (params.H + kRowBlock - 1) / kRowBlock;
     g_scratch.assign(static_cast<size_t>(params.max_tokens) * params.top_k * params.I, 0);
@@ -217,7 +223,16 @@ struct CpuMoeExecutor {
     return MoeTask{static_cast<uint64_t>(layer_id), static_cast<uint64_t>(num_tokens)};
   }
 
-  std::string isa_name() const { return fmt->isa_name(); }
+  std::string isa_name() const {
+    std::vector<std::string> names;
+    for (const auto& f : fmts)
+      if (const std::string n = f->isa_name(); std::find(names.begin(), names.end(), n) == names.end()) names.push_back(n);
+    std::string out;
+    for (const auto& n : names) out += (out.empty() ? "" : ",") + n;
+    return out;
+  }
+
+  const MoeFormat& format_of(const MoeTask& t) const { return *fmts[layer_fmt[t.layer]]; }
 
   // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others.
   bool has_routes(const MoeTask& t) const {
@@ -249,8 +264,8 @@ struct CpuMoeExecutor {
     const float w_in = apply_on_input ? io_w[static_cast<size_t>(tok) * top_k + k] : 1.0f;
     const int i0 = static_cast<int>(ib) * kRowBlock;
     const int i1 = std::min(I, i0 + kRowBlock);
-    fmt->gate_up(static_cast<int>(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in,
-                 g_scratch.data() + ((size_t)tok * top_k + k) * I);
+    format_of(t).gate_up(static_cast<int>(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in,
+                         g_scratch.data() + ((size_t)tok * top_k + k) * I);
   }
 
   // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes.
@@ -272,7 +287,9 @@ struct CpuMoeExecutor {
       ++n;
     }
     float part[kMaxTopK * kRowBlock];
-    if (n > 0) fmt->down(static_cast<int>(t.layer), h0, h1, routes, n, part);
+    const MoeFormat& f = format_of(t);
+    if (n > 0) f.down(static_cast<int>(t.layer), h0, h1, routes, n, part);
+    const bool rounds_routes = f.rounds_routes();
     bf16_t* y_row = io_y + (size_t)tok * H;
     for (int h = h0; h < h1; ++h) {
       float acc = 0.0f;
@@ -294,11 +311,12 @@ struct CpuMoeExecutor {
     barrier(local_sense);
     // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter.
     // Needs all of pass1 done (a full row spans every iblk).
-    if (preps_inter) {
+    if (prt_total > 0) {
+      MoeFormat& f = *fmts[layer_fmt[t.layer]];
       for (;;) {
         int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed);
         if (r >= prt_total) break;
-        fmt->prep_inter(r, g_scratch.data() + (size_t)r * params.I);
+        f.prep_inter(r, g_scratch.data() + (size_t)r * params.I);
       }
       barrier(local_sense);
     }
@@ -348,9 +366,10 @@ struct CpuMoeExecutor {
 
   void submit(MoeTask t) {
     const int num_tokens = static_cast<int>(t.num_tokens);
+    MoeFormat& fmt = *fmts[layer_fmt[t.layer]];
     p1_total = static_cast<int64_t>(num_tokens) * params.top_k * n_iblk;
     p2_total = static_cast<int64_t>(num_tokens) * n_hblk;
-    prt_total = preps_inter ? static_cast<int64_t>(num_tokens) * params.top_k : 0;
+    prt_total = fmt.preps_inter() ? static_cast<int64_t>(num_tokens) * params.top_k : 0;
     p1_next.store(0, std::memory_order_relaxed);
     p2_next.store(0, std::memory_order_relaxed);
     prt_next.store(0, std::memory_order_relaxed);
@@ -359,7 +378,7 @@ struct CpuMoeExecutor {
     bar_sense.store(0, std::memory_order_relaxed);
     // Prepare the per-token input once, up front.
     // Single-threaded: tiny for decode, and done before the workers are woken below.
-    for (int tok = 0; tok < num_tokens; ++tok) fmt->prep_input(tok, io_x + (size_t)tok * params.H);
+    for (int tok = 0; tok < num_tokens; ++tok) fmt.prep_input(tok, io_x + (size_t)tok * params.H);
     {
       std::lock_guard<std::mutex> lk(task_mtx);
       cur_task = t;
@@ -510,14 +529,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   using freetoken::cpu_moe::CpuMoeExecutor;
   using freetoken::cpu_moe::GpuSync;
   py::class_<CpuMoeExecutor>(m, "CpuMoeExecutor")
-      .def(py::init<int, int, int, int, int, int, int, int, int, std::string,
-                    std::map<std::string, uint64_t>, std::map<std::string, int64_t>,
-                    double, double, std::vector<int>, uintptr_t, uintptr_t, uintptr_t, uintptr_t>(),
+      .def(py::init<int, int, int, int, int, int, int, int, int, std::vector<freetoken::cpu_moe::FormatSpec>,
+                    std::vector<int>, double, double, std::vector<int>, uintptr_t, uintptr_t, uintptr_t,
+                    uintptr_t>(),
            py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"),
            py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"),
            py::arg("max_tokens"), py::arg("activation_id"),
-           py::arg("apply_router_weight_on_input"), py::arg("weight_format"),
-           py::arg("banks"), py::arg("options"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"),
+           py::arg("apply_router_weight_on_input"), py::arg("formats"), py::arg("layer_formats"),
+           py::arg("swiglu_alpha"), py::arg("swiglu_limit"),
            py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"),
            py::arg("y_ptr"))
       .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream,
diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py
index de1ce004fe..5ae89c97fd 100644
--- a/python/freetoken/layers/quantization/moe/base.py
+++ b/python/freetoken/layers/quantization/moe/base.py
@@ -151,6 +151,10 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
     def worth_it(self, cfg: MoEConfig) -> bool:
         return True
 
+    def cpu_options(self, cfg: MoEConfig) -> dict[str, int]:
+        """Integer options of ``cpu_format`` for these experts."""
+        return {}
+
     def slot_limit(self, cfg: MoEConfig) -> int | None:
         """Most GPU cache slots the kernel can address for ``cfg``; None for no limit."""
         return self.max_slots
@@ -196,6 +200,10 @@ def apply(self, x: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Ten
     def cpu_format(self) -> str | None:
         return self.kernel.cpu_format
 
+    def cpu_spec(self) -> tuple[str | None, dict[str, int]]:
+        """The CPU executor format and options of this layer's experts."""
+        return self.cpu_format, self.kernel.cpu_options(self.cfg)
+
     def create_weights(self, layer: Any) -> None:
         # resident experts stay out of the state dict; attach_resident_banks fills these after the dense weights load
         layer._expert_banks = None
diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py
index 1af6d9f43e..a17b5bbe09 100644
--- a/python/freetoken/layers/quantization/moe/mxfp4.py
+++ b/python/freetoken/layers/quantization/moe/mxfp4.py
@@ -28,6 +28,9 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
             return "W4A8 ds_fp4 experts need the fp8 activation quant block"
         return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False)
 
+    def cpu_options(self, cfg: MoEConfig) -> dict[str, int]:
+        return {"act_block": int(cfg.act_block)}
+
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         i, h = cfg.intermediate, cfg.hidden
         return {
diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py
index fbcddb1c06..8bc4c1f098 100644
--- a/python/freetoken/models/deepseek_v4/weight.py
+++ b/python/freetoken/models/deepseek_v4/weight.py
@@ -237,7 +237,7 @@ def iter_vision_weights(model_path: str, device):
 _KIND_SUFFIX = {"weight": "", "scale": "_scale"}
 
 
-def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
+def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
     """Routed experts, one piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0
     ``_scale`` companions (``w1`` / ``w3`` / ``w2``). The MTP layer's experts are skipped."""
     if kind is not QuantKind.MXFP4:
@@ -249,10 +249,11 @@ def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bo
 
     args = load_args(model_path, max_batch_size=1)
     L, E = args.n_layers, args.n_routed_experts
+    keep = None if layers is None else set(layers)
 
     def locate(raw_name: str):
         m = _EXPERT_RE.match(raw_name)
-        if m is None or int(m["layer"]) >= L:
+        if m is None or int(m["layer"]) >= L or (keep is not None and int(m["layer"]) not in keep):
             return None
         return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]]
 
@@ -263,7 +264,7 @@ def locate(raw_name: str):
     def _serial():
         reader = _ShardReader(model_path, _weight_map(model_path), torch.device("cpu"))
         try:
-            for li in tqdm(range(L), desc="Loading DSV4 experts (serial)", disable=not get_tp_info().is_primary()):
+            for li in tqdm(range(L) if keep is None else sorted(keep), desc="Loading DSV4 experts (serial)", disable=not get_tp_info().is_primary()):
                 for e in range(E):
                     base = f"layers.{li}.ffn.experts.{e}"
                     for proj in ("w1", "w3", "w2"):
diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py
index 1ccd23eaef..962bd19589 100644
--- a/python/freetoken/models/deepseek_v41/weight.py
+++ b/python/freetoken/models/deepseek_v41/weight.py
@@ -200,7 +200,7 @@ def ftw_side_files(model_path: str, out_dir: str) -> list[str]:
 _KIND_SUFFIX = {"weight": "", "scale": "_scale"}
 
 
-def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
+def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
     """One piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0 ``_scale`` companions.
     The DSpark stage's experts (``mtp.*``) never match."""
     if kind is not QuantKind.MXFP4:
@@ -212,10 +212,11 @@ def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bo
 
     args = load_args(model_path)
     L, E = args.n_layers, args.n_routed_experts
+    keep = None if layers is None else set(layers)
 
     def locate(raw_name: str):
         m = _EXPERT_RE.match(raw_name)
-        if m is None or int(m["layer"]) >= L:
+        if m is None or int(m["layer"]) >= L or (keep is not None and int(m["layer"]) not in keep):
             return None
         return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]]
 
@@ -226,7 +227,7 @@ def locate(raw_name: str):
     def _serial():
         reader = _ShardReader(model_path, weight_map(model_path), torch.device("cpu"))
         try:
-            for li in tqdm(range(L), desc="Loading DeepSeek-V4.1 experts (serial)", disable=not get_tp_info().is_primary()):
+            for li in tqdm(range(L) if keep is None else sorted(keep), desc="Loading DeepSeek-V4.1 experts (serial)", disable=not get_tp_info().is_primary()):
                 for e in range(E):
                     base = f"layers.{li}.ffn.experts.{e}"
                     for proj in ("w1", "w3", "w2"):
diff --git a/python/freetoken/models/glm5_next/weight.py b/python/freetoken/models/glm5_next/weight.py
index 6f967dd8c7..09f9169419 100644
--- a/python/freetoken/models/glm5_next/weight.py
+++ b/python/freetoken/models/glm5_next/weight.py
@@ -250,7 +250,7 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple
         reader.close()
 
 
-def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
+def iter_expert_pieces(model_path, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
     """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their ``_scale`` companions; other kinds use the generic readers."""
     if kind is not QuantKind.FP8_BLOCK:
         return None
@@ -260,13 +260,14 @@ def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool |
     from freetoken.moe.expert_pieces import per_expert_pieces
 
     suffix = {"weight": "", "weight_scale_inv": "_scale"}
+    keep = None if layers is None else set(layers)
 
     def locate(raw_name: str):
         m = _FP8_EXPERT_RE.match(raw_name)
         if m is None:
             return None
         bank = _layer_to_bank(int(m["layer"]), config)
-        if bank is None:
+        if bank is None or (keep is not None and bank not in keep):
             return None
         return bank, int(m["expert"]), m["proj"] + suffix[m["kind"]]
 
@@ -284,8 +285,11 @@ def _serial():
             weight_map = json.load(f)["weight_map"]
         reader = _ShardReader(folder, weight_map, torch.device("cpu"))
         try:
-            layers = range(config.first_k_dense_replace, config.num_layers)
-            for layer in tqdm(layers, desc="Loading GLM-5.3 fp8 experts (serial)", disable=not get_tp_info().is_primary()):
+            ckpt_layers = [
+                layer for layer in range(config.first_k_dense_replace, config.num_layers)
+                if keep is None or _layer_to_bank(layer, config) in keep
+            ]
+            for layer in tqdm(ckpt_layers, desc="Loading GLM-5.3 fp8 experts (serial)", disable=not get_tp_info().is_primary()):
                 for e in range(config.num_experts):
                     base = f"{_CKPT}.layers.{layer}.mlp.experts.{e}"
                     for proj in ("gate", "up", "down"):
diff --git a/python/freetoken/models/gpt_oss/weight.py b/python/freetoken/models/gpt_oss/weight.py
index 686f02466f..cb5830c3db 100644
--- a/python/freetoken/models/gpt_oss/weight.py
+++ b/python/freetoken/models/gpt_oss/weight.py
@@ -288,7 +288,7 @@ def _source_slices(config, tp_info) -> dict[str, tuple[str, tuple[slice, ...]]]:
     }
 
 
-def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20):
+def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, layers=None, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20):
     """gpt-oss experts as HF ships them, one piece per layer: stacked ``[E, ...]`` mxfp4
     ``_blocks`` / ``_scales`` / ``_bias`` tensors sliced to this rank's intermediate range."""
     if kind is not QuantKind.MXFP4:
@@ -298,10 +298,11 @@ def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bo
     tp_info = get_tp_info()
     slices = _source_slices(config, tp_info)
     num_layers, num_experts = config.num_layers, config.num_experts
+    keep = None if layers is None else set(layers)
 
     def _is_expert(name: str) -> bool:
         info = _expert_layer_and_name(name)
-        return info is not None and 0 <= info[0] < num_layers
+        return info is not None and 0 <= info[0] < num_layers and (keep is None or info[0] in keep)
 
     def _tensors():
         if parallel:
@@ -320,6 +321,8 @@ def _tensors():
                     layer_id, source = info
                     if not 0 <= layer_id < num_layers:
                         raise ValueError(f"Unexpected GPT-OSS expert layer in checkpoint: {name}")
+                    if keep is not None and layer_id not in keep:
+                        continue
                     if source not in slices:
                         raise ValueError(f"Unexpected GPT-OSS expert source: {name}")
                     yield name, _read_safetensor_slice(f, name, slices[source][1])
diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py
index 2dfdaaa35f..04da050919 100644
--- a/python/freetoken/models/nvfp4_banks.py
+++ b/python/freetoken/models/nvfp4_banks.py
@@ -4,7 +4,7 @@
 import os
 import re
 from dataclasses import dataclass
-from typing import Callable
+from typing import Callable, Collection
 
 import safetensors
 import torch
@@ -68,6 +68,7 @@ def iter_nvfp4_expert_pieces(
     config,
     spec: Nvfp4ExpertSourceSpec,
     *,
+    layers: Collection[int] | None = None,
     parallel: bool = False,
     workers: int = 8,
     chunk: int = 8 << 20,
@@ -80,6 +81,7 @@ def iter_nvfp4_expert_pieces(
 
     Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either
     way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land.
+    ``layers`` (bank layer ids, default all) reads only those layers' experts.
     """
     from freetoken.models.loader import drop_page_cache as _drop
     from freetoken.models.loader import safetensors_weight_map
@@ -89,13 +91,14 @@ def iter_nvfp4_expert_pieces(
     folder = download_hf_weight(model_path)
     weight_map = safetensors_weight_map(folder)
 
+    keep = None if layers is None else set(layers)
     wanted: dict[str, tuple[int, int, str]] = {}
     for name in weight_map:
         match = spec.key_pattern.match(name)
         if match is None:
             continue
         bank_layer = _bank_layer(spec, int(match.group("layer")), config)
-        if bank_layer is None:
+        if bank_layer is None or (keep is not None and bank_layer not in keep):
             continue
         proj = match.group("proj")
         if proj not in spec.proj_to_role:
@@ -104,7 +107,7 @@ def iter_nvfp4_expert_pieces(
         if kind not in ("weight", "weight_scale", "weight_scale_2"):
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
         wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind))
-    expected = _num_moe_layers(config) * config.num_experts * 9
+    expected = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts * 9
     if len(wanted) != expected:
         raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}")
 
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index de5b9e8983..724ad15312 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -344,7 +344,7 @@ def _moe_dims(model_config):
     )
 
 
-def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
+def iter_expert_pieces(model_path, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20):
     """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their
     ``_scale`` (block scale) companions, named as the checkpoint's dialect stores them. Other expert kinds use the generic readers."""
     if kind is not QuantKind.FP8_BLOCK:
@@ -355,6 +355,7 @@ def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool |
     from freetoken.moe.expert_pieces import per_expert_pieces
 
     L, E, H, I, dense = _moe_dims(config)
+    keep = None if layers is None else set(layers)
     scale = get_quant_config().stored_tensors(QuantKind.FP8_BLOCK)["weight_scale_inv"].name
     key_re = re.compile(_FP8_EXPERT_KEY_RE.format(scale=re.escape(scale)))
     suffix = {"weight": "", scale: "_scale"}
@@ -366,20 +367,22 @@ def locate(raw_name: str):
         li = int(m["layer"]) - dense
         if not 0 <= li < L:
             raise ValueError(f"unexpected routed-expert layer in {raw_name}")
+        if keep is not None and li not in keep:
+            return None
         return li, int(m["expert"]), m["proj"] + suffix[m["kind"]]
 
     if parallel is None:
         parallel = experts_scattered(model_path)
     if parallel:
         tensors = iter_expert_tensors_parallel(
-            model_path, lambda n: key_re.match(n) is not None, workers=workers, chunk=chunk
+            model_path, lambda n: locate(n) is not None, workers=workers, chunk=chunk
         )
         return per_expert_pieces(tensors, locate, tensors_per_expert=6)
 
     def _serial():
         reader = ShardReader(model_path, torch.device("cpu"))
         try:
-            for li in tqdm(range(L), desc="Loading fp8 experts (serial)", disable=not get_tp_info().is_primary()):
+            for li in tqdm(range(L) if keep is None else sorted(keep), desc="Loading fp8 experts (serial)", disable=not get_tp_info().is_primary()):
                 for e in range(E):
                     base = f"model.language_model.layers.{dense + li}.mlp.experts.{e}"
                     for proj in ("gate", "up", "down"):
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index cd260cb2f6..86a03f4175 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -476,7 +476,7 @@ def _build_cpu_moe_executor(fmt: str, wl: Workload, banks: dict, num_threads: in
         cache, top_k=wl.top_k, activation=wl.activation,
         apply_router_weight_on_input=False, num_threads=num_threads, max_tokens=1,
         device=torch.device("cuda"), swiglu_alpha=wl.swiglu_alpha,
-        swiglu_limit=wl.swiglu_limit, act_block=wl.act_block,
+        swiglu_limit=wl.swiglu_limit, layer_formats=[(fmt, {"act_block": wl.act_block} if wl.act_block is not None else {})] * cache.num_layers,
     )
 
 
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index bb9ee56385..5d9bd3c9e4 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -163,28 +163,34 @@ def __init__(
         device: torch.device,
         swiglu_alpha: float = 1.702,
         swiglu_limit: float | None = None,
-        fmt: str | None = None,
-        act_block: int | None = None,
+        layer_formats: list[tuple[str, dict[str, int]]] | None = None,
     ) -> None:
         from freetoken.kernel import _cpu_moe
         from freetoken.moe.legacy_format import canonical_role
 
-        fmt = fmt or cache.quant_format
-        if fmt not in _cpu_moe.weight_formats():
-            raise NotImplementedError(
-                f"--moe-strategy cpu/hybrid computes experts on the CPU and supports "
-                f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are "
-                f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead."
-            )
+        # layer_formats[l] = (format, options) of layer l; None runs every layer in cache.quant_format
+        if layer_formats is None:
+            layer_formats = [(cache.quant_format, {})] * int(cache.num_layers)
+        assert len(layer_formats) == cache.num_layers, (len(layer_formats), cache.num_layers)
+        keys = [(name, tuple(sorted(opts.items()))) for name, opts in layer_formats]
+        specs = list(dict.fromkeys(keys))
+        names = sorted({name for name, _ in specs})
+        for f in names:
+            if f not in _cpu_moe.weight_formats():
+                raise NotImplementedError(
+                    f"--moe-strategy cpu/hybrid computes experts on the CPU and supports "
+                    f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are "
+                    f"{f!r}; use --moe-strategy offload (GPU-side dequant) instead."
+                )
         if activation not in _ACT_IDS:
             raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}")
-        if apply_router_weight_on_input and fmt in ("mxfp4_triton", "ds_fp4"):
-            raise NotImplementedError(f"CPU MoE backend: {fmt} experts apply the router weight on the output only")
+        if apply_router_weight_on_input and (output_only := {"mxfp4_triton", "ds_fp4"} & set(names)):
+            raise NotImplementedError(f"CPU MoE backend: {sorted(output_only)} experts apply the router weight on the output only")
         # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without
         # error and silently computes the wrong activation in the generic
         # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4
         # handles its act inside the kernel and predates the marker.)
-        if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton":
+        if _ACT_IDS[activation] >= 3 and names != ["mxfp4_triton"]:
             supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)()
             if _ACT_IDS[activation] > supported:
                 raise RuntimeError(
@@ -197,16 +203,24 @@ def __init__(
         self.num_layers = int(cache.num_layers)
         self.num_experts = int(cache.num_experts)
         self.top_k = int(top_k)
-        self.quant_format = fmt
+        self.quant_format = ",".join(names)
         self.device = device
         self.max_tokens = int(max_tokens)
         self.apply_router_weight_on_input = bool(apply_router_weight_on_input)
         # The per-layer tensors and their pointer tables must outlive the executor
         # (C++ holds raw addresses into both).
         self._banks: list[torch.Tensor] = []
-        banks, options, (self.H, self.I) = self._resolve_banks(
-            {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt
-        )
+        by_role = {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}
+        # every format's tables address every layer, but a format reads only its own layers' rows
+        formats, shapes = [], set()
+        for name, opts in specs:
+            banks, options, shape = self._resolve_banks(by_role, name, keys.index((name, opts)), dict(opts))
+            formats.append((name, banks, {**dict(opts), **options}))
+            shapes.add(shape)
+        if len(shapes) != 1:
+            raise NotImplementedError(f"CPU MoE formats disagree on the expert shape (hidden, intermediate): {sorted(shapes)}")
+        (self.H, self.I), = shapes
+        layer_counts = {spec: keys.count(spec) for spec in specs}
 
         # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain
         # that the C++ side runs single-threaded on the CUDA host-callback thread --
@@ -215,12 +229,15 @@ def __init__(
         # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell
         # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode:
         # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py).
-        self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" and _DSFP4_GPU_PREQUANT
         # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1)
-        self._act_block = act_block
-        if fmt == "ds_fp4":
-            assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block"
-            options.update(act_block=int(act_block), input_prequant=int(self._gpu_prequant))
+        blocks = {options.get("act_block") for name, _, options in formats if name == "ds_fp4"}
+        assert None not in blocks, "ds_fp4 experts need the act_block option, the checkpoint's fp8 activation quant block"
+        # the GPU round-trip rewrites the whole input, so only when every layer is ds_fp4 with one block
+        self._gpu_prequant = names == ["ds_fp4"] and len(blocks) == 1 and device.type == "cuda" and _DSFP4_GPU_PREQUANT
+        self._act_block = next(iter(blocks)) if self._gpu_prequant else None
+        for name, _, options in formats:
+            if name == "ds_fp4":
+                options["input_prequant"] = int(self._gpu_prequant)
         if self._gpu_prequant:
             logger.info_rank0(
                 "ds_fp4: input FP8 round-trip moved to the GPU "
@@ -271,9 +288,8 @@ def __init__(
             max_tokens=self.max_tokens,
             activation_id=_ACT_IDS[activation],
             apply_router_weight_on_input=1 if apply_router_weight_on_input else 0,
-            weight_format=fmt,
-            banks=banks,
-            options=options,
+            formats=formats,
+            layer_formats=[specs.index(k) for k in keys],
             swiglu_alpha=float(swiglu_alpha),
             swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"),
             core_ids=core_ids,
@@ -332,9 +348,11 @@ def __init__(
             self._watchdog.start()
 
 
+        fmt_label = "; ".join(f"{name}{dict(opts) if opts else ''}: {n} layers" for (name, opts), n in layer_counts.items())
         logger.info_rank0(
             f"CPU MoE executor ready: threads={nthreads} (pinned to cores "
-            f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} "
+            f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} "
+            f"fmt={fmt_label} "
             f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} "
             f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} "
             f"spin_wait={self.spin_wait} "
@@ -360,23 +378,23 @@ def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]:
         """Per-layer table addresses keyed by canonical bank role."""
         return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()}
 
-    def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
+    def _resolve_fp8_block_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride."""
         gu, gus = banks["gate_up"], banks["gate_up_scale"]
         dn, dns = banks["down"], banks["down_scale"]
-        assert gu[0].element_size() == 1 and dn[0].element_size() == 1, (gu[0].dtype, dn[0].dtype)
-        assert gus[0].dtype == torch.bfloat16 and dns[0].dtype == torch.bfloat16, (gus[0].dtype, dns[0].dtype)
-        I = int(gu[0].shape[1] // 2)
-        H = int(gu[0].shape[2])
-        assert gu[0].shape[1] == 2 * I
-        assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I)
-        assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I)
-        assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I)
+        assert gu[rep].element_size() == 1 and dn[rep].element_size() == 1, (gu[rep].dtype, dn[rep].dtype)
+        assert gus[rep].dtype == torch.bfloat16 and dns[rep].dtype == torch.bfloat16, (gus[rep].dtype, dns[rep].dtype)
+        I = int(gu[rep].shape[1] // 2)
+        H = int(gu[rep].shape[2])
+        assert gu[rep].shape[1] == 2 * I
+        assert tuple(dn[rep].shape[1:]) == (H, I), (dn[rep].shape, H, I)
+        assert gus[rep].shape[1] == (2 * I + 127) // 128 and gus[rep].shape[2] >= (H + 127) // 128, (gus[rep].shape, H, I)
+        assert dns[rep].shape[1] == (H + 127) // 128 and dns[rep].shape[2] >= (I + 127) // 128, (dns[rep].shape, H, I)
         banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns)
-        options = dict(gu_scale_stride=int(gus[0].stride(1)), dn_scale_stride=int(dns[0].stride(1)))
+        options = dict(gu_scale_stride=int(gus[rep].stride(1)), dn_scale_stride=int(dns[rep].stride(1)))
         return banks, options, (H, I)
 
-    def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int, int]]:
+    def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None = None) -> tuple[dict, dict, tuple[int, int]]:
         """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role
         the format reads to its per-layer table's address, options holds its integer settings.
 
@@ -389,102 +407,102 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int,
         if fmt == "bf16":
             gate_up = banks["gate_up"]
             down = banks["down"]
-            if gate_up[0].dtype != torch.bfloat16 or down[0].dtype != torch.bfloat16:
+            if gate_up[rep].dtype != torch.bfloat16 or down[rep].dtype != torch.bfloat16:
                 raise NotImplementedError(
-                    f"bf16 CPU MoE requires bf16 banks, got {gate_up[0].dtype}/{down[0].dtype}"
+                    f"bf16 CPU MoE requires bf16 banks, got {gate_up[rep].dtype}/{down[rep].dtype}"
                 )
-            H = int(gate_up[0].shape[2])
-            I = int(gate_up[0].shape[1] // 2)
-            assert gate_up[0].shape[1] == 2 * I
-            assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I)
+            H = int(gate_up[rep].shape[2])
+            I = int(gate_up[rep].shape[1] // 2)
+            assert gate_up[rep].shape[1] == 2 * I
+            assert tuple(down[rep].shape[1:]) == (H, I), (down[rep].shape, H, I)
             banks = self._tables(gate_up=gate_up, down=down)
             return banks, {}, (H, I)
 
         if fmt == "q4_0":
-            return self._resolve_q4_0_banks(banks)
+            return self._resolve_q4_0_banks(banks, rep)
 
         if fmt == "mxfp4_triton":
-            return self._resolve_mxfp4_banks(banks)
+            return self._resolve_mxfp4_banks(banks, rep)
 
         if fmt == "ds_fp4":
-            return self._resolve_dsfp4_banks(banks)
+            return self._resolve_dsfp4_banks(banks, rep)
 
         if fmt == "fp8_block":
-            return self._resolve_fp8_block_banks(banks)
+            return self._resolve_fp8_block_banks(banks, rep)
 
         # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals.
         gup, gus, gug = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_global"]
         dnp, dns, dng = banks["down"], banks["down_scale"], banks["down_global"]
-        assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype)
-        assert gus[0].element_size() == 1 and dns[0].element_size() == 1, "block scales must be 1 byte"
-        assert gug[0].dtype == torch.float16 and dng[0].dtype == torch.float16, (gug[0].dtype, dng[0].dtype)
-        I = int(gup[0].shape[1] // 2)
-        H = int(gup[0].shape[2] * 2)
-        assert gup[0].shape[1] == 2 * I
+        assert gup[rep].dtype == torch.uint8 and dnp[rep].dtype == torch.uint8, (gup[rep].dtype, dnp[rep].dtype)
+        assert gus[rep].element_size() == 1 and dns[rep].element_size() == 1, "block scales must be 1 byte"
+        assert gug[rep].dtype == torch.float16 and dng[rep].dtype == torch.float16, (gug[rep].dtype, dng[rep].dtype)
+        I = int(gup[rep].shape[1] // 2)
+        H = int(gup[rep].shape[2] * 2)
+        assert gup[rep].shape[1] == 2 * I
         assert H % 16 == 0 and I % 16 == 0, (H, I)
-        assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I)
-        assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H)
-        assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I)
-        assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,)
+        assert tuple(dnp[rep].shape[1:]) == (H, I // 2), (dnp[rep].shape, H, I)
+        assert tuple(gus[rep].shape[1:]) == (2 * I, H // 16), (gus[rep].shape, I, H)
+        assert tuple(dns[rep].shape[1:]) == (H, I // 16), (dns[rep].shape, H, I)
+        assert tuple(gug[rep].shape[1:]) == (2 * I,) and tuple(dng[rep].shape[1:]) == (H,)
         banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng)
         return banks, {}, (H, I)
 
-    def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
+    def _resolve_q4_0_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble
         bytes), row-major over K -- the *same* packed banks the GPU offload path streams.
         gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a
         row in place (18 bytes / 32 K) and dequantizes weights inside the K-loop."""
         gate_up, down = banks["gate_up"], banks["down"]
-        assert gate_up[0].dtype == torch.uint8 and down[0].dtype == torch.uint8, (
-            gate_up[0].dtype, down[0].dtype,
+        assert gate_up[rep].dtype == torch.uint8 and down[rep].dtype == torch.uint8, (
+            gate_up[rep].dtype, down[rep].dtype,
         )
-        I = int(gate_up[0].shape[1] // 2)
-        H = int(down[0].shape[1])
-        assert gate_up[0].shape[1] == 2 * I
+        I = int(gate_up[rep].shape[1] // 2)
+        H = int(down[rep].shape[1])
+        assert gate_up[rep].shape[1] == 2 * I
         assert H % 32 == 0 and I % 32 == 0, (H, I)
-        assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H)
-        assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I)
+        assert int(gate_up[rep].shape[2]) == (H // 32) * 18, (gate_up[rep].shape, H)
+        assert int(down[rep].shape[2]) == (I // 32) * 18, (down[rep].shape, I)
         banks = self._tables(gate_up=gate_up, down=down)
         return banks, {}, (H, I)
 
-    def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
+    def _resolve_mxfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales
         (N innermost) + per-output-row biases. The C++ kernel streams K and
         accumulates a contiguous N-block, so the GPU-tiled layout is read in place
         (no repack, no extra host memory). Block scales are e8m0 (1 byte / 32 K)."""
         gub, gus, gob = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_bias"]
         dnb, dns, dob = banks["down"], banks["down_scale"], banks["down_bias"]
-        assert gub[0].dtype == torch.uint8 and dnb[0].dtype == torch.uint8, (gub[0].dtype, dnb[0].dtype)
-        assert gus[0].dtype == torch.uint8 and dns[0].dtype == torch.uint8, (gus[0].dtype, dns[0].dtype)
-        assert gob[0].dtype == torch.bfloat16 and dob[0].dtype == torch.bfloat16, (gob[0].dtype, dob[0].dtype)
+        assert gub[rep].dtype == torch.uint8 and dnb[rep].dtype == torch.uint8, (gub[rep].dtype, dnb[rep].dtype)
+        assert gus[rep].dtype == torch.uint8 and dns[rep].dtype == torch.uint8, (gus[rep].dtype, dns[rep].dtype)
+        assert gob[rep].dtype == torch.bfloat16 and dob[rep].dtype == torch.bfloat16, (gob[rep].dtype, dob[rep].dtype)
         # gate_up_blocks [E, H//2, 2I]; down_blocks [E, I//2, H]
-        H = int(gub[0].shape[1] * 2)
-        I = int(gub[0].shape[2] // 2)
-        assert gub[0].shape[2] == 2 * I
+        H = int(gub[rep].shape[1] * 2)
+        I = int(gub[rep].shape[2] // 2)
+        assert gub[rep].shape[2] == 2 * I
         assert H % 32 == 0 and I % 32 == 0, (H, I)
-        assert tuple(dnb[0].shape[1:]) == (I // 2, H), (dnb[0].shape, H, I)
-        assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I)
-        assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I)
-        assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,)
+        assert tuple(dnb[rep].shape[1:]) == (I // 2, H), (dnb[rep].shape, H, I)
+        assert tuple(gus[rep].shape[1:]) == (H // 32, 2 * I), (gus[rep].shape, H, I)
+        assert tuple(dns[rep].shape[1:]) == (I // 32, H), (dns[rep].shape, H, I)
+        assert tuple(gob[rep].shape[1:]) == (2 * I,) and tuple(dob[rep].shape[1:]) == (H,)
         banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob)
         return banks, {}, (H, I)
 
-    def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]:
+    def _resolve_dsfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block
         scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row),
         so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the
         activations (the checkpoint's fp8 block) to match the DeepSeek W4A8 reference."""
         gup, gus = banks["gate_up"], banks["gate_up_scale"]
         dnp, dns = banks["down"], banks["down_scale"]
-        assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype)
-        assert gus[0].element_size() == 1 and dns[0].element_size() == 1, "block scales must be 1 byte"
-        I = int(gup[0].shape[1] // 2)
-        H = int(gup[0].shape[2] * 2)
-        assert gup[0].shape[1] == 2 * I
+        assert gup[rep].dtype == torch.uint8 and dnp[rep].dtype == torch.uint8, (gup[rep].dtype, dnp[rep].dtype)
+        assert gus[rep].element_size() == 1 and dns[rep].element_size() == 1, "block scales must be 1 byte"
+        I = int(gup[rep].shape[1] // 2)
+        H = int(gup[rep].shape[2] * 2)
+        assert gup[rep].shape[1] == 2 * I
         assert H % 32 == 0 and I % 32 == 0, (H, I)  # e8m0 scales per 32 along K
-        assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I)
-        assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H)
-        assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I)
+        assert tuple(dnp[rep].shape[1:]) == (H, I // 2), (dnp[rep].shape, H, I)
+        assert tuple(gus[rep].shape[1:]) == (2 * I, H // 32), (gus[rep].shape, I, H)
+        assert tuple(dns[rep].shape[1:]) == (H, I // 32), (dns[rep].shape, H, I)
         banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns)
         return banks, {}, (H, I)
 
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 03ce4b750b..7e25e58f73 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -18,8 +18,8 @@
 from freetoken.layers.quantization import QuantKind
 from freetoken.utils import init_logger
 
-from .host_banks import alloc_layer_banks
-from .offload_cache import _BANK_BYTES_PER_EXPERT
+from .host_banks import HostBank
+from .offload_cache import _BANK_BYTES_PER_EXPERT, placeholder_bank
 
 logger = init_logger(__name__)
 
@@ -52,6 +52,8 @@ class ExpertBanks:
     layer_tags: list | None = None
     # per bank, per layer: bytes past an expert matrix's last row that the kernel reads and that must read as zeros (ggml MMQ); None = no such reads
     row_tail_bytes: dict[str, list[int]] | None = None
+    # per-layer quant_format of experts in several formats (``quant_format``/``kind``/``kernel`` then name the most common one); None = every layer is ``quant_format``
+    layer_formats: list[str] | None = None
 
 
 def _dummy_fill(role: str, tensor: torch.Tensor) -> None:
@@ -80,6 +82,7 @@ def build_expert_banks(
     layer_sink=None,
     dummy: bool = False,
     resident: bool = False,
+    layers: list[int] | None = None,
 ) -> ExpertBanks:
     """Fill host banks in the kernel's layout from a stream of expert pieces.
 
@@ -89,6 +92,8 @@ def build_expert_banks(
     are pinned in the background, otherwise the sink receives them (converter). ``dummy``
     skips the pieces and fills the banks with finite random contents.
     ``resident`` packs into banks on ``device`` instead of host banks (the resident 'fused' experts).
+    ``layers`` fills only those bank layers (the ones ``method`` serves, for experts of several
+    formats): pieces of other layers are skipped and the other layers' sources are None.
     """
     from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, pin_banks
     from freetoken.moe.legacy_format import legacy_format_for
@@ -96,29 +101,34 @@ def build_expert_banks(
     kernel = method.kernel
     layout = method.layout()
     E = method.cfg.num_experts
+    ids = list(range(num_layers)) if layers is None else sorted(layers)
+    filled = set(ids)
     specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident}
     if resident:
         assert layer_sink is None
         hb = None
-        banks = {role: [torch.empty(shape, dtype=dtype, device=device) for _ in range(num_layers)] for role, (shape, dtype) in specs.items()}
+        banks = {role: {l: torch.empty(shape, dtype=dtype, device=device) for l in ids} for role, (shape, dtype) in specs.items()}
     else:
-        hb = alloc_layer_banks(specs, num_layers)
-        banks = {role: [b.tensor for b in hb[role]] for role in specs}
+        hb = {role: {l: HostBank(shape, dtype) for l in ids} for role, (shape, dtype) in specs.items()}
+        banks = {role: {l: b.tensor for l, b in per_layer.items()} for role, per_layer in hb.items()}
     alphas = {
         role: torch.empty(num_layers * E, dtype=spec.dtype, device=device)
         for role, spec in layout.items() if spec.resident
     }
 
+    def _sources() -> dict[str, list]:
+        return {role: [per_layer.get(l) for l in range(num_layers)] for role, per_layer in banks.items()}
+
     if dummy:
         for role, per_layer in banks.items():
-            for tensor in per_layer:
+            for tensor in per_layer.values():
                 _dummy_fill(role, tensor)
         for alpha in alphas.values():
             alpha.fill_(1.0)
         if hb is not None and torch.cuda.is_available():
             pin_banks(hb)
         return ExpertBanks(
-            legacy_format_for(method.kind, kernel.name), banks,
+            legacy_format_for(method.kind, kernel.name), _sources(),
             gate_up_alpha=alphas.get("gate_up_alpha"), down_alpha=alphas.get("down_alpha"),
             kind=method.kind, kernel=kernel.name, layout=layout,
         )
@@ -130,6 +140,8 @@ def _fill(sink) -> None:
         for layer_id, e0, e1, piece in pieces:
             if not (0 <= layer_id < num_layers and 0 <= e0 < e1 <= E):
                 raise ValueError(f"expert piece out of range: layer {layer_id}, experts {e0}:{e1} of {num_layers} x {E}")
+            if layer_id not in filled:
+                continue
             # refuse before writing: a duplicate row would also complete the layer early and hand the sink a half-filled bank
             if written[layer_id, e0:e1].any():
                 raise ValueError(f"expert rows written more than once: layer {layer_id}, experts {e0}:{e1}")
@@ -141,8 +153,9 @@ def _fill(sink) -> None:
             if tracker is not None:
                 for _ in range(e1 - e0):
                     tracker.note(layer_id)
-        missing = (written == 0).nonzero().tolist()
+        missing = (written[ids] == 0).nonzero().tolist()
         if missing:
+            missing = [[ids[l], e] for l, e in missing]
             raise ValueError(f"expert banks were not filled: {len(missing)} (layer, expert) rows missing (first {missing[:4]})")
 
     if layer_sink is not None:
@@ -154,12 +167,74 @@ def _fill(sink) -> None:
         _fill(None)
 
     return ExpertBanks(
-        legacy_format_for(method.kind, kernel.name), banks,
+        legacy_format_for(method.kind, kernel.name), _sources(),
         gate_up_alpha=alphas.get("gate_up_alpha"), down_alpha=alphas.get("down_alpha"),
         streamed=layer_sink is not None, kind=method.kind, kernel=kernel.name, layout=layout,
     )
 
 
+def method_groups(layer_methods) -> list[tuple]:
+    """``[(method, [bank layers])]`` of the layers' expert methods, one entry per (kind, kernel), most layers first."""
+    groups: dict = {}
+    for layer, method in enumerate(layer_methods):
+        groups.setdefault((method.kind, method.kernel.name), (method, []))[1].append(layer)
+    return sorted(groups.values(), key=lambda g: (-len(g[1]), g[1][0]))
+
+
+def is_mixed(layer_methods) -> bool:
+    return layer_methods is not None and len(method_groups(layer_methods)) > 1
+
+
+def merge_expert_banks(parts: list[ExpertBanks], layer_methods, num_experts: int) -> ExpertBanks:
+    """One ``ExpertBanks`` over the per-format ``parts`` (each filled only on its layers): the banks are the
+    union of the roles, a layer without a role gets a zero-byte placeholder (see ``OffloadMoeCache.layer_layouts``)."""
+    from freetoken.moe.legacy_format import legacy_format_for
+
+    num_layers = len(layer_methods)
+    sources: dict[str, list] = {}
+    for part in parts:
+        if part.gate_up_alpha is not None or part.down_alpha is not None:
+            raise ValueError(f"experts of several formats cannot use the {part.kernel} kernel, which keeps GPU-resident alphas")
+        for role, per_layer in part.sources.items():
+            merged = sources.setdefault(role, [None] * num_layers)
+            for layer, tensor in enumerate(per_layer):
+                if tensor is not None:
+                    assert merged[layer] is None, (role, layer)
+                    merged[layer] = tensor
+    for per_layer in sources.values():
+        for layer, tensor in enumerate(per_layer):
+            if tensor is None:
+                per_layer[layer] = placeholder_bank(num_experts)
+    formats = [legacy_format_for(m.kind, m.kernel.name) for m in layer_methods]
+    primary = parts[0]
+    return ExpertBanks(
+        primary.quant_format, sources, streamed=all(p.streamed for p in parts),
+        kind=primary.kind, kernel=primary.kernel, layer_tags=formats, layer_formats=formats,
+    )
+
+
+def _mixed_expert_banks(model_path, model_config, layer_methods, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks:
+    from freetoken.moe.expert_pieces import iter_expert_pieces
+
+    num_layers = model_config.num_moe_layers
+    assert len(layer_methods) == num_layers, (len(layer_methods), num_layers)
+    groups = method_groups(layer_methods)
+    # resolve every format's reader before any bank exists, so a missing parallel reader falls back to serial with nothing allocated
+    readers = [
+        None if dummy else iter_expert_pieces(
+            model_path, model_config, method.kind, layers=layers, parallel=parallel, workers=workers, chunk=chunk
+        )
+        for method, layers in groups
+    ]
+    parts = []
+    for (method, layers), pieces in zip(groups, readers):
+        parts.append(build_expert_banks(
+            method, num_layers, pieces, device=device, layer_sink=layer_sink, dummy=dummy, layers=layers,
+        ))
+        logger.info_rank0(f"expert banks: {method.kind} / {method.kernel.name} on {len(layers)} of {num_layers} MoE layers")
+    return merge_expert_banks(parts, layer_methods, model_config.num_experts)
+
+
 _PARALLEL_CHUNK = 8 << 20  # default O_DIRECT chunk for the parallel reader
 
 
@@ -216,19 +291,24 @@ def ftw_bank_bytes(model_path: str) -> int | None:
     return sum(t["nbytes"] for t in tensors if t.get("kind") == "experts_bank")
 
 
-def bank_bytes_estimate(model_config, method=None) -> int | None:
+def _method_expert_bytes(method) -> int:
+    return method.cfg.num_experts * sum(
+        math.prod(spec.shape) * torch.empty((), dtype=spec.dtype).element_size()
+        for spec in method.layout().values() if not spec.resident
+    )
+
+
+def bank_bytes_estimate(model_config, method=None, layer_methods=None) -> int | None:
     """Estimated total expert-bank bytes of a raw checkpoint before loading it.
 
-    With a bound expert ``method`` the kernel's layout gives the exact host bytes; otherwise the
-    format-tag table sizes the format. ``None`` for unknown formats or missing dims
-    (callers then skip the pre-load sizing)."""
+    With a bound expert ``method`` the kernel's layout gives the exact host bytes (``layer_methods``,
+    one per MoE layer, for experts of several formats); otherwise the format-tag table sizes the
+    format. ``None`` for unknown formats or missing dims (callers then skip the pre-load sizing)."""
     layers = getattr(model_config, "num_moe_layers", None)
+    if layer_methods is not None:
+        return sum(_method_expert_bytes(m) for m in layer_methods)
     if method is not None and layers:
-        per_expert = sum(
-            math.prod(spec.shape) * torch.empty((), dtype=spec.dtype).element_size()
-            for spec in method.layout().values() if not spec.resident
-        )
-        return layers * method.cfg.num_experts * per_expert
+        return layers * _method_expert_bytes(method)
     expert_quant = getattr(model_config, "expert_quant", "none")
     fmt = expert_quant if expert_quant != "none" else (
         getattr(model_config, "moe_weight_format", None) or "bf16"
@@ -256,6 +336,7 @@ def load_expert_banks(
     layer_sink=None,
     layer_residency: list[str] | None = None,
     resident: bool = False,
+    layer_methods: list | None = None,
 ) -> ExpertBanks:
     """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning
     the same normalized ``ExpertBanks`` and both pinning after fill:
@@ -282,9 +363,15 @@ def load_expert_banks(
     Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota).
 
     ``resident`` (with ``method``): the banks land on ``device`` for the resident 'fused' experts; no host bank outlives the load.
+
+    ``layer_methods`` (one bound method per MoE layer) serves experts of several formats: each
+    format fills its own layers and the banks merge over the union of roles (``merge_expert_banks``).
     """
     from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks, load_ftw_banks_to_device
 
+    mixed = is_mixed(layer_methods)
+    if mixed and (resident or layer_residency is not None):
+        raise ValueError("experts of several formats are served from the offload slot cache only (no resident experts, no --moe-cpu-layers)")
     if resident:
         assert method is not None and layer_sink is None and layer_residency is None
     if model_path and is_ftw_checkpoint(model_path) and not dummy:
@@ -330,6 +417,8 @@ def load_expert_banks(
     from freetoken.moe.host_banks import requested_residency
 
     def _build(par: bool) -> ExpertBanks:
+        if mixed:
+            return _mixed_expert_banks(model_path, model_config, layer_methods, device, dummy, par, workers, chunk, layer_sink)
         if method is not None and method.kind is QuantKind.GGUF:
             from freetoken.models.gguf.experts import load_gguf_expert_banks
 
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index 79f655a951..917d84bbcb 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Iterable, Iterator
+from typing import Callable, Collection, Iterable, Iterator
 
 import torch
 from freetoken.layers.quantization import QuantKind
@@ -38,24 +38,34 @@ def _model_hook(spec, name: str):
 
 
 def iter_expert_pieces(
-    model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20
+    model_path: str,
+    config,
+    kind: QuantKind,
+    *,
+    layers: Collection[int] | None = None,
+    parallel: bool = False,
+    workers: int = 8,
+    chunk: int = 8 << 20,
 ) -> Iterator[Piece]:
     """The pieces of ``model_path``'s routed experts, stored as ``kind``.
 
-    A family may ship its own ``iter_expert_pieces(model_path, config, kind, *, parallel, workers,
-    chunk)`` for the storage forms only it uses (it returns None for the rest); otherwise bf16
+    A family may ship its own ``iter_expert_pieces(model_path, config, kind, *, layers, parallel,
+    workers, chunk)`` for the storage forms only it uses (it returns None for the rest); otherwise bf16
     experts come from the family's stacked ``iter_weights`` and NVFP4 experts from its
     ``nvfp4_expert_spec``. The reader is resolved here, before any bank is allocated, so a
     missing parallel reader raises ``NotImplementedError`` while a serial fallback is still cheap.
+
+    ``layers`` (bank layer ids, default all) limits the pieces to those layers, for experts of several
+    formats: a reader skips the other layers' tensors and checks only these layers for completeness.
     """
     spec = get_model_spec(config.architectures[0])
     hook = _model_hook(spec, "iter_expert_pieces")
     if hook is not None:
-        pieces = hook(model_path, config, kind, parallel=parallel, workers=workers, chunk=chunk)
+        pieces = hook(model_path, config, kind, layers=layers, parallel=parallel, workers=workers, chunk=chunk)
         if pieces is not None:
             return pieces
     if kind is QuantKind.NONE:
-        return _bf16_pieces(model_path, config, spec, parallel=parallel, workers=workers, chunk=chunk)
+        return _bf16_pieces(model_path, config, spec, layers=layers, parallel=parallel, workers=workers, chunk=chunk)
     if kind is QuantKind.NVFP4:
         spec_hook = _model_hook(spec, "nvfp4_expert_spec")
         if spec_hook is None:
@@ -63,7 +73,7 @@ def iter_expert_pieces(
         from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces
 
         return iter_nvfp4_expert_pieces(
-            model_path, config, spec_hook(model_path, config), parallel=parallel, workers=workers, chunk=chunk
+            model_path, config, spec_hook(model_path, config), layers=layers, parallel=parallel, workers=workers, chunk=chunk
         )
     raise NotImplementedError(f"{spec.module} provides no expert reader for {kind!r} experts")
 
@@ -81,18 +91,31 @@ def packed_expert_source_info(key: str) -> tuple[int, str] | None:
         return None
 
 
-def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config) -> Iterator[Piece]:
-    """Whole-layer ``experts.gate_up_proj`` / ``experts.down_proj`` tensors -> one piece per layer."""
+def _source_layer(key: str) -> int | None:
+    """Checkpoint layer of a ``model.layers.N....`` key."""
+    parts = key.split(".")
+    return int(parts[2]) if len(parts) > 2 and parts[:2] == ["model", "layers"] and parts[2].isdigit() else None
+
+
+def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config, *, layers: Collection[int] | None = None) -> Iterator[Piece]:
+    """Whole-layer ``experts.gate_up_proj`` / ``experts.down_proj`` tensors -> one piece per layer, for ``layers`` (bank layer ids, default all)."""
+    keep = None if layers is None else set(layers)
     pending: dict[int, dict[str, torch.Tensor]] = {}
     num_experts = config.num_experts
     for name, tensor in tensors:
         info = packed_expert_source_info(name)
         if info is None:
+            # another format's layer stores its experts in another form
+            layer = _source_layer(name)
+            if keep is not None and layer is not None and bank_layer_of(config, layer) not in keep:
+                continue
             raise ValueError(f"Unexpected expert weight key: {name}")
         layer, packed_name = info
         bank_layer = bank_layer_of(config, layer)
         if bank_layer is None:
             raise ValueError(f"Unexpected MoE expert layer {layer}; expected a routed-expert layer")
+        if keep is not None and bank_layer not in keep:
+            continue
         if tensor.size(0) != num_experts:
             raise ValueError(f"Unexpected {packed_name} expert count {tensor.size(0)}; expected {num_experts}")
         piece = pending.setdefault(bank_layer, {})
@@ -104,7 +127,7 @@ def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config) -
         raise ValueError(f"Missing MoE expert source layers: {sorted(pending)}")
 
 
-def _bf16_pieces(model_path: str, config, spec, *, parallel: bool, workers: int, chunk: int) -> Iterator[Piece]:
+def _bf16_pieces(model_path: str, config, spec, *, layers: Collection[int] | None, parallel: bool, workers: int, chunk: int) -> Iterator[Piece]:
     device = torch.device("cpu")
     if parallel:
         iter_weights = _model_hook(spec, "iter_weights_parallel")
@@ -116,7 +139,7 @@ def _bf16_pieces(model_path: str, config, spec, *, parallel: bool, workers: int,
     else:
         iter_weights = _load_attr(spec.module, spec.iter_weights)
         tensors = iter_weights(model_path, device, include_moe_experts=True, include_non_moe=False)
-    return stacked_expert_pieces(tensors, config)
+    return stacked_expert_pieces(tensors, config, layers=layers)
 
 
 def per_expert_pieces(
diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py
index 91e2648773..95562a8855 100644
--- a/python/freetoken/moe/host_banks.py
+++ b/python/freetoken/moe/host_banks.py
@@ -204,18 +204,6 @@ def alloc_banks(specs: dict[str, tuple[tuple[int, ...], torch.dtype]]) -> dict[s
     return {name: HostBank(shape, dtype) for name, (shape, dtype) in specs.items()}
 
 
-def alloc_layer_banks(
-    specs: dict[str, tuple[tuple[int, ...], torch.dtype]], num_layers: int
-) -> dict[str, list[HostBank]]:
-    """Allocate per-layer host banks: ``{name: ([num_experts, ...] row shape, dtype)}``
-    -> one independently allocated (page-aligned, independently pin/lock-able)
-    ``HostBank`` per layer per name."""
-    return {
-        name: [HostBank(shape, dtype) for _ in range(num_layers)]
-        for name, (shape, dtype) in specs.items()
-    }
-
-
 class _ResidencyPlan:
     """Per-layer ``HostResidency`` labels, ambiently visible to the bank settle points.
 
@@ -265,13 +253,13 @@ def _settle(bank: HostBank, residency: str) -> None:
         bank.lock()
 
 
-def pin_banks(banks: dict[str, HostBank | list[HostBank]]) -> None:
+def pin_banks(banks: dict[str, HostBank | list[HostBank] | dict[int, HostBank]]) -> None:
     """Settle every bank after it has been filled -- pin-after-fill by default.
-    List-valued entries are per-layer and honor the ambient :func:`requested_residency` plan; scalar banks always pin."""
+    List-valued entries (indexed by layer id) and dict-valued ones (layer id -> bank) are per-layer and honor the ambient :func:`requested_residency` plan; scalar banks always pin."""
     plan = _requested_residency
     for bank in banks.values():
-        if isinstance(bank, list):
-            for layer_id, layer_bank in enumerate(bank):
+        if isinstance(bank, (list, dict)):
+            for layer_id, layer_bank in bank.items() if isinstance(bank, dict) else enumerate(bank):
                 residency = (
                     HostResidency.PINNED.value if plan is None
                     else plan.residency_for(layer_id)
@@ -479,7 +467,6 @@ def rd(i: int) -> None:
     "LayerCompletionTracker",
     "PinPipeline",
     "alloc_banks",
-    "alloc_layer_banks",
     "born_pinned_default",
     "pin_banks",
     "read_file_into",
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index a209f9a924..23ba325106 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -102,6 +102,11 @@
 MARLIN_MAX_CACHE_SIZE = 992
 
 
+def placeholder_bank(num_experts: int) -> torch.Tensor:
+    """The source of a bank whose role a layer's expert kernel does not read (experts of several formats): zero bytes per row."""
+    return torch.empty((num_experts, 0), dtype=torch.uint8)
+
+
 def bank_layouts(sources: dict[str, list[torch.Tensor]], layer_tags: list | None = None) -> list[tuple]:
     """Each layer's layout key: its quant tag and every bank's (dtype, row shape). A slot pool holds layers of one key."""
     num_layers = len(next(iter(sources.values())))
@@ -154,6 +159,9 @@ class OffloadMoeCache:
     # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots
     layout: dict | None = None
     max_slots: int | None = None
+    # Per-layer kernel layouts for experts of several formats (replaces ``layout``). The banks are the union of their roles;
+    # a layer whose kernel reads no such role carries a zero-byte [num_experts, 0] uint8 placeholder (see placeholder_bank).
+    layer_layouts: list[dict] | None = None
     # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers.
     pool_of_layer: list[int] | None = None
     pool_sizes: list[int] | None = None
@@ -162,7 +170,12 @@ def __post_init__(self) -> None:
         policy_ids = {"lru": 0}
         assert self.cache_policy in policy_ids
         assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target
-        if self.layout is None:
+        if self.layer_layouts is not None:
+            assert self.layout is None, "pass layout or layer_layouts, not both"
+            assert len(self.layer_layouts) == self.num_layers, (len(self.layer_layouts), self.num_layers)
+            if any(spec.resident for layout in self.layer_layouts for spec in layout.values()):
+                raise ValueError("experts of several formats cannot use a kernel with GPU-resident alphas (marlin, b12x)")
+        elif self.layout is None:
             assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}"
         if self.pool_sizes is None:
             self.pool_of_layer = [0] * self.num_layers
@@ -231,7 +244,11 @@ def __post_init__(self) -> None:
         # carry independent host attributes -- see layer_residency) and their GPU
         # slot caches, keyed by the format's bank schema (attached by
         # set_bank_sources). The GPU slot cache stays one unified pool per bank.
-        if self.layout is not None:
+        if self.layer_layouts is not None:
+            self.bank_schema = tuple(dict.fromkeys(
+                role for layout in self.layer_layouts for role, spec in layout.items() if not spec.resident
+            ))
+        elif self.layout is not None:
             self.bank_schema = tuple(role for role, spec in self.layout.items() if not spec.resident)
         else:
             self.bank_schema = _BANK_SCHEMAS[self.quant_format]
@@ -371,9 +388,9 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
             for s, r, d in zip(self.pool_sizes[:-1], row, pad):
                 offsets[-1].append(offsets[-1][-1] + s * r + d)
         self._pool_offsets = offsets
-        # local slots below this index overlap the prefill staging area in some bank
+        # local slots below this index overlap the prefill staging area in some bank; a pool's zero-byte rows (a role its kernel does not read) overlap nothing
         self._pool_stable_from = [
-            max(min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows)))
+            max((min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows)) if rows[k][p]), default=0)
             for p, s in enumerate(self.pool_sizes)
         ]
         self._inval = []
@@ -382,6 +399,8 @@ def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor:
             for k in range(len(rows)):
                 lo, hi = b * stage[k], (b + 1) * stage[k]
                 for p, (base, s) in enumerate(spans):
+                    if not rows[k][p]:
+                        continue
                     first = max((lo - offsets[k][p]) // rows[k][p], 0)
                     end = min(-(-(hi - offsets[k][p]) // rows[k][p]), s)
                     mask[base + first : base + max(first, end)] = True
@@ -488,8 +507,12 @@ def set_bank_sources(
             for layer_id, source in enumerate(per_layer):
                 assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous"
                 assert source.size(0) == self.num_experts, (name, layer_id, source.shape)
-                if self.layout is not None:
-                    spec = self.layout[name]
+                spec = self.layout[name] if self.layout is not None else None
+                if self.layer_layouts is not None:
+                    spec = self.layer_layouts[layer_id].get(name)
+                    if spec is None and source.numel() != 0:
+                        raise ValueError(f"layer {layer_id}'s expert kernel reads no {name!r} bank, so it takes a zero-byte placeholder")
+                if spec is not None:
                     if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype:
                         raise ValueError(
                             f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout "
@@ -586,6 +609,9 @@ def _build_fused_copy_plan(self) -> None:
                     # a 0 placeholder keeps the descriptor shape
                     layer_src_ptrs[layer_id].append(0)
                     continue
+                if source.numel() == 0:
+                    layer_src_ptrs[layer_id].append(0)  # placeholder bank: zero bytes per row, never dereferenced
+                    continue
                 # The kernel dereferences these on the GPU, so store each host bank's
                 # device alias (== data_ptr() under UVA identity; differs on
                 # Windows/WDDM).
@@ -980,6 +1006,8 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
             dst, src, nbytes = [], [], []
             for b, feat in enumerate(self._row_bytes[layout]):
                 stage = self._copy_dst_ptrs_host[b] + buffer_id * self._stage_bytes[b]
+                if feat == 0:
+                    continue
                 if feat < _SMALL_BANK_FEAT_BYTES:
                     # Whole layer as one entry, EVEN with zero misses: it keeps every
                     # batch entry above the driver's async floor and covers the hit
@@ -1229,6 +1257,8 @@ def copy_missing(self) -> None:
 
         views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id)
         for (per_layer, _), cache in zip(self.banks, views):
+            if cache.numel() == 0:
+                continue
             fast_index_copy_jit(
                 cache,
                 self.evict_slots,
diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py
index b0ac2b13f3..4f634c4c94 100644
--- a/tests/models/test_models_loader.py
+++ b/tests/models/test_models_loader.py
@@ -201,3 +201,22 @@ def test_stacked_expert_pieces_pair_each_layer_in_arrival_order():
     assert torch.equal(pieces[1][3]["gate_up"], torch.full((2, 3, 4), 2.0))
     with pytest.raises(ValueError, match="Missing MoE expert source layers"):
         list(stacked_expert_pieces(tensors[:3], config))
+
+
+def test_stacked_expert_pieces_skip_the_layers_of_other_formats():
+    from freetoken.moe.expert_pieces import stacked_expert_pieces
+
+    config = SimpleNamespace(num_layers=3, num_experts=2)
+    tensors = [
+        ("model.layers.0.mlp.experts.gate_up_proj", torch.full((2, 3, 4), 2.0)),
+        ("model.layers.1.mlp.experts.0.gate_proj.weight_scale", torch.zeros(3, 1)),
+        ("model.layers.0.mlp.experts.down_proj", torch.full((2, 4, 3), 10.0)),
+        ("model.layers.2.mlp.experts.gate_up_proj", torch.full((2, 3, 4), 3.0)),
+        ("model.layers.2.mlp.experts.down_proj", torch.full((2, 4, 3), 11.0)),
+    ]
+
+    pieces = list(stacked_expert_pieces(tensors, config, layers=[0]))
+
+    assert [(layer, e0, e1) for layer, e0, e1, _ in pieces] == [(0, 0, 2)]
+    with pytest.raises(ValueError, match="Unexpected expert weight key"):
+        list(stacked_expert_pieces(tensors, config))
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index 158f1f4598..16f439a9ec 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -507,6 +507,8 @@ def test_nvfp4_expert_pieces_read_either_dialect_from_a_single_file(checkpoint):
     spec = nvfp4_expert_spec(folder, config)
     pieces = list(iter_nvfp4_expert_pieces(folder, config, spec))
     assert len(pieces) == 2 * E
+    only = list(iter_nvfp4_expert_pieces(folder, config, spec, layers=[1]))
+    assert len(only) == E and {p[0] for p in only} == {1}
     layer, e0, e1, piece = next(p for p in pieces if p[0] == 1 and p[1] == 2)
     assert (e0, e1) == (2, 3)
     base = f"{LM}.layers.1.mlp.experts.2.gate_proj"
@@ -529,6 +531,8 @@ def test_block_fp8_expert_pieces_read_either_dialect(checkpoint, parallel):
     config = parse_config(cached_load_hf_config(folder))
     pieces = list(iter_expert_pieces(folder, config, QuantKind.FP8_BLOCK, parallel=parallel))
     assert len(pieces) == 2 * E
+    only = list(iter_expert_pieces(folder, config, QuantKind.FP8_BLOCK, layers=[1], parallel=parallel))
+    assert len(only) == E and {p[0] for p in only} == {1}
     layer, e0, e1, piece = next(p for p in pieces if p[0] == 1 and p[1] == 2)
     base = f"{LM}.layers.1.mlp.experts.2"
     scale = "weight_scale_inv" if name == "fp8_block" else "weight_scale"
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index 38e79b2743..56c6c18aa4 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -391,7 +391,7 @@ def test_cpu_decode_dsfp4_matches_gpu(bs):
         max_tokens=bs,
         device=dev,
         swiglu_limit=swiglu_limit,
-        act_block=128,
+        layer_formats=[("ds_fp4", {"act_block": 128})] * cache.num_layers,
     )
 
     hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16)
@@ -689,7 +689,7 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4():
     ex = CpuMoeExecutor(
         cache, top_k=top_k, activation="silu",
         apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev,
-        swiglu_limit=limit, act_block=128,
+        swiglu_limit=limit, layer_formats=[("ds_fp4", {"act_block": 128})] * cache.num_layers,
     )
     b = cache.bank_sources
     layer_banks = [b[n][layer].to(dev) for n in (
diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py
index 8187402d3f..1212c166bd 100644
--- a/tests/moe/test_dsfp4_prequant.py
+++ b/tests/moe/test_dsfp4_prequant.py
@@ -49,7 +49,7 @@ def run(cpu_roundtrip: bool) -> torch.Tensor:
         monkeypatch.setattr(cpu_executor, "_DSFP4_GPU_PREQUANT", not cpu_roundtrip)
         ex = CpuMoeExecutor(
             cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False,
-            num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128,
+            num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, layer_formats=[("ds_fp4", {"act_block": 128})] * cache.num_layers,
         )
         assert ex._gpu_prequant != cpu_roundtrip, "ds_fp4 on CUDA should enable GPU prequant"
         out = ex.decode(0, x, wts, ids.clone()).clone()
diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py
index 2f89cf334d..f537b24877 100644
--- a/tests/moe/test_dsfp4_reference_semantics.py
+++ b/tests/moe/test_dsfp4_reference_semantics.py
@@ -105,7 +105,7 @@ def test_cpu_executor_matches_the_transcription(block):
     cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E)
     dev = torch.device("cuda", 0)
     x, slots, w = _inputs(2, dev)
-    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=LIMIT, act_block=block)
+    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=LIMIT, layer_formats=[("ds_fp4", {"act_block": block})] * cache.num_layers)
     got = ex.decode(0, x, w, slots.clone()).clone().float()
     torch.cuda.synchronize()
     del ex
@@ -162,7 +162,7 @@ def test_inactive_routes_contribute_zero_without_reading_their_slots(block):
         buf.copy_(t.cpu())
         pinned[name] = buf
     cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E)
-    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=torch.device("cuda", 0), swiglu_limit=LIMIT, act_block=block)
+    ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=torch.device("cuda", 0), swiglu_limit=LIMIT, layer_formats=[("ds_fp4", {"act_block": block})] * cache.num_layers)
     cpu_ids = torch.where(on_gpu, -1, slots)
     cpu_out = ex.decode(0, x, w, cpu_ids.clone()).clone()
     torch.cuda.synchronize()
diff --git a/tests/moe/test_mixed_expert_formats.py b/tests/moe/test_mixed_expert_formats.py
new file mode 100644
index 0000000000..c9301dc132
--- /dev/null
+++ b/tests/moe/test_mixed_expert_formats.py
@@ -0,0 +1,288 @@
+"""Experts of several formats in one model: bf16 layers next to NVFP4 layers.
+
+A checkpoint may leave some MoE layers unquantized (Muse Spark keeps its first and last layer's
+experts bf16). Each format fills its own layers' banks, the banks merge over the union of roles with
+zero-byte placeholders, the slot cache keeps one pool per format, the FTW stores each layer's own
+banks, and the CPU executor computes each layer in its format.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.distributed import set_tp_info, try_get_tp_info
+
+cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+E, H, I, TOPK = 8, 256, 128, 2
+FORMATS = ["bf16", "nvfp4", "nvfp4", "nvfp4", "bf16"]
+L = len(FORMATS)
+
+_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
+
+
+def _init_tp():
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+
+
+def _layer(layer_id: int, fmt: str):
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig, QuantBackend, QuantConfig, set_quant_backend
+
+    _init_tp()
+    if fmt == "bf16":
+        quant = NoQuantConfig()
+    else:
+        set_quant_backend(QuantBackend.parse("moe.nvfp4=triton"))
+        quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}})
+    return OffloadMoELayer(layer_id, E, TOPK, H, I, quant_config=quant, prefix=f"model.layers.{layer_id}.mlp.experts")
+
+
+def _dequant(packed, scale, row_global):
+    n, k2 = packed.shape
+    codes = torch.stack([packed & 0xF, packed >> 4], dim=-1).view(n, 2 * k2).long()
+    return _E2M1.to(packed.device)[codes] * scale.float().repeat_interleave(16, dim=1) * row_global.float().unsqueeze(1)
+
+
+def _pieces(seed: int = 0):
+    """The checkpoint side: per layer, bf16 or ModelOpt NVFP4 pieces in the layer's format."""
+    g = torch.Generator().manual_seed(seed)
+    pieces = {}
+    for layer, fmt in enumerate(FORMATS):
+        if fmt == "bf16":
+            pieces[layer] = {
+                "gate_up": (torch.randn(E, 2 * I, H, generator=g) / 16).to(torch.bfloat16),
+                "down": (torch.randn(E, H, I, generator=g) / 16).to(torch.bfloat16),
+            }
+        else:
+            pieces[layer] = {
+                "gate_up": torch.randint(0, 256, (E, 2 * I, H // 2), dtype=torch.uint8, generator=g),
+                "gate_up_scale": (torch.rand(E, 2 * I, H // 16, generator=g) + 0.25).to(torch.float8_e4m3fn),
+                "gate_up_global": torch.full((E, 2 * I), 0.02, dtype=torch.float16),
+                "down": torch.randint(0, 256, (E, H, I // 2), dtype=torch.uint8, generator=g),
+                "down_scale": (torch.rand(E, H, I // 16, generator=g) + 0.25).to(torch.float8_e4m3fn),
+                "down_global": torch.full((E, H), 0.02, dtype=torch.float16),
+            }
+    return pieces
+
+
+def _reader(pieces, kind):
+    """A family reader for one storage form: it yields only the layers stored that way."""
+    from freetoken.layers.quantization import QuantKind
+
+    want = "bf16" if kind is QuantKind.NONE else "nvfp4"
+    for layer, fmt in enumerate(FORMATS):
+        if fmt == want:
+            yield layer, 0, E, pieces[layer]
+
+
+def _weights(pieces, layer):
+    """Dense fp32 (gate_up [E, 2I, H], down [E, H, I]) of one layer, the reference."""
+    p = pieces[layer]
+    if FORMATS[layer] == "bf16":
+        return p["gate_up"].float(), p["down"].float()
+    gu = torch.stack([_dequant(p["gate_up"][e], p["gate_up_scale"][e], p["gate_up_global"][e]) for e in range(E)])
+    dn = torch.stack([_dequant(p["down"][e], p["down_scale"][e], p["down_global"][e]) for e in range(E)])
+    return gu, dn
+
+
+def _reference(pieces, layer, hidden, weights, ids):
+    gu, dn = _weights(pieces, layer)
+    out = torch.zeros(hidden.shape, dtype=torch.float32)
+    x = hidden.float().cpu()
+    for t in range(hidden.shape[0]):
+        for k in range(ids.shape[1]):
+            e = int(ids[t, k])
+            h = gu[e] @ x[t]
+            out[t] += float(weights[t, k]) * (dn[e] @ (torch.nn.functional.silu(h[:I]) * h[I:]))
+    return out
+
+
+def _assert_close(out, ref):
+    torch.testing.assert_close(out.float().cpu(), ref, rtol=3e-2, atol=0.03 * float(ref.abs().max()))
+
+
+def _mixed_banks(pieces, device, layer_sink=None):
+    from freetoken.moe.expert_banks import build_expert_banks, merge_expert_banks, method_groups
+
+    methods = [_layer(layer, fmt).quant_method for layer, fmt in enumerate(FORMATS)]
+    parts = [
+        build_expert_banks(method, L, _reader(pieces, method.kind), device=device, layer_sink=layer_sink, layers=layers)
+        for method, layers in method_groups(methods)
+    ]
+    return merge_expert_banks(parts, methods, E), methods
+
+
+def _mixed_cache(device, pieces, *, decode_target="gpu", prefill_overlap=False):
+    from freetoken.moe.offload_cache import OffloadMoeCache, bank_layouts
+    from freetoken.moe.pool_policy import group_layers, make_pools, pool_of_layer, split_slots
+    from freetoken.engine.cache_budget import layer_slot_bytes
+
+    banks, methods = _mixed_banks(pieces, device)
+    slot_bytes = layer_slot_bytes(banks.sources)
+    groups = group_layers("unified", bank_layouts(banks.sources, banks.layer_tags), slot_bytes)
+    pools = make_pools(groups, E, TOPK, None)
+    sizes = split_slots(4 * E, pools)
+    cache = OffloadMoeCache(
+        num_layers=L, num_experts=E, cache_size=sum(sizes), device=device, quant_format=banks.quant_format,
+        decode_target=decode_target, hybrid_max_fetch=1, prefill_overlap=prefill_overlap,
+        layer_layouts=[m.layout() for m in methods], pool_of_layer=pool_of_layer(pools, L), pool_sizes=sizes,
+    )
+    cache.set_bank_sources(banks.sources, layer_tags=banks.layer_tags)
+    cache.reset()
+    layers = [_layer(layer, fmt) for layer, fmt in enumerate(FORMATS)]
+    for layer in layers:
+        layer.offload_cache = cache
+    return cache, layers, banks
+
+
+def test_engine_binds_the_most_common_format_and_checks_every_layer():
+    from freetoken.engine.engine import _check_layer_formats, offload_layer_methods, shared_offload_method
+    from freetoken.layers.base import BaseOP
+    from freetoken.layers.quantization import QuantKind
+    from freetoken.moe.expert_banks import ExpertBanks
+
+    class Model(BaseOP):
+        def __init__(self):
+            self.layers = [_layer(layer, fmt) for layer, fmt in reversed(list(enumerate(FORMATS)))]
+
+    model = Model()
+    methods = offload_layer_methods(model)
+    assert [m.kind for m in methods] == [QuantKind.NONE if f == "bf16" else QuantKind.NVFP4 for f in FORMATS]
+    assert shared_offload_method(model).kind is QuantKind.NVFP4
+    _check_layer_formats(ExpertBanks("nvfp4", {}, layer_formats=list(FORMATS)), methods)
+    with pytest.raises(ValueError, match=r"MoE layers \[0, 4\]"):
+        _check_layer_formats(ExpertBanks("nvfp4", {}), methods)
+
+
+def test_mixed_dummy_banks_merge_the_roles_with_placeholders():
+    from freetoken.moe.expert_banks import load_expert_banks
+
+    methods = [_layer(layer, fmt).quant_method for layer, fmt in enumerate(FORMATS)]
+    config = SimpleNamespace(num_moe_layers=L, num_experts=E)
+    banks = load_expert_banks("", config, method=methods[1], device=torch.device("cpu"), dummy=True, layer_methods=methods)
+    assert banks.layer_formats == FORMATS and banks.quant_format == "nvfp4"
+    assert set(banks.sources) == {"gate_up", "gate_up_scale", "gate_up_global", "down", "down_scale", "down_global"}
+    for layer, fmt in enumerate(FORMATS):
+        if fmt == "bf16":
+            assert banks.sources["gate_up"][layer].shape == (E, 2 * I, H) and banks.sources["gate_up"][layer].dtype == torch.bfloat16
+            assert banks.sources["down_scale"][layer].shape == (E, 0)
+        else:
+            assert banks.sources["gate_up"][layer].shape == (E, 2 * I, H // 2) and banks.sources["gate_up"][layer].dtype == torch.uint8
+    with pytest.raises(ValueError, match="several formats"):
+        load_expert_banks("", config, method=methods[1], device=torch.device("cpu"), dummy=True,
+                          layer_methods=methods, layer_residency=["pinned"] * L)
+
+
+@cuda
+def test_mixed_banks_give_each_reader_its_layers_and_resolve_every_reader_first(monkeypatch):
+    from freetoken.moe import expert_pieces
+    from freetoken.moe.expert_banks import _mixed_expert_banks
+
+    pieces = _pieces(seed=2)
+    methods = [_layer(layer, fmt).quant_method for layer, fmt in enumerate(FORMATS)]
+    events = []
+
+    def fake_reader(model_path, config, kind, *, layers=None, parallel=False, workers=8, chunk=0):
+        events.append(("resolve", sorted(layers)))
+
+        def read():
+            events.append(("read", sorted(layers)))
+            for layer in layers:
+                yield layer, 0, E, pieces[layer]
+
+        return read()
+
+    monkeypatch.setattr(expert_pieces, "iter_expert_pieces", fake_reader)
+    config = SimpleNamespace(num_moe_layers=L, num_experts=E)
+    banks = _mixed_expert_banks("", config, methods, torch.device("cpu"), False, False, 8, 0)
+    assert sorted(layers for step, layers in events if step == "resolve") == [[0, 4], [1, 2, 3]]
+    assert [step for step, _ in events[:2]] == ["resolve", "resolve"]
+    assert banks.layer_formats == FORMATS
+
+
+@cuda
+def test_mixed_layers_decode_and_prefill_through_their_pools():
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    pieces = _pieces()
+    cache, layers, banks = _mixed_cache(device, pieces, prefill_overlap=True)
+    assert len(cache.pool_sizes) == 2 and cache._raw_banks
+    # a staging buffer holds one bf16 layer (the widest rows); two of them still fit the 13 + 19 slots
+    assert cache.prefill_overlap
+    hidden = torch.randn(1, H, dtype=torch.bfloat16, device=device) / 4
+    weights = torch.rand(1, TOPK, device=device)
+    for step in range(6):
+        for layer in layers:
+            ids = torch.randperm(E)[:TOPK].to(torch.int32).view(1, TOPK)
+            out = layer._decode_routed(hidden, weights, ids.clone().to(device))
+            _assert_close(out, _reference(pieces, layer.layer_id, hidden, weights, ids))
+
+    M = 16
+    hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4
+    ids = torch.stack([torch.randperm(E)[:TOPK] for _ in range(M)]).to(torch.int32)
+    weights = torch.rand(M, TOPK, device=device)
+    for layer in layers:
+        # the bf16 prefill kernel overwrites its input
+        out = layer._prefill_routed(hidden.clone(), weights, ids.to(device))
+        _assert_close(out, _reference(pieces, layer.layer_id, hidden, weights, ids))
+    # decode after the prefill stomped the staging slots must reload, not read another format's bytes
+    for layer in layers:
+        ids = torch.randperm(E)[:TOPK].to(torch.int32).view(1, TOPK)
+        out = layer._decode_routed(hidden[:1], weights[:1], ids.clone().to(device))
+        _assert_close(out, _reference(pieces, layer.layer_id, hidden[:1], weights[:1], ids))
+
+
+@cuda
+def test_cpu_executor_computes_each_layer_in_its_format():
+    pytest.importorskip("freetoken.kernel._cpu_moe")
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    pieces = _pieces(seed=1)
+    cache, layers, _ = _mixed_cache(device, pieces, decode_target="hybrid")
+    ex = CpuMoeExecutor(
+        cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False, num_threads=2,
+        max_tokens=4, device=device, layer_formats=[(f, {}) for f in FORMATS],
+    )
+    cache.set_cpu_executor(ex)
+    bs = 3
+    hidden = torch.randn(bs, H, dtype=torch.bfloat16, device=device) / 4
+    weights = torch.rand(bs, TOPK, device=device)
+    for layer in range(L):
+        ids = torch.stack([torch.randperm(E)[:TOPK] for _ in range(bs)]).to(torch.int32)
+        _assert_close(ex.decode(layer, hidden, weights, ids.to(device)), _reference(pieces, layer, hidden, weights, ids))
+    # hybrid: capped PCIe fetch on the GPU, the overflow misses on the CPU, per layer format
+    for _ in range(3):
+        for layer in layers:
+            ids = torch.stack([torch.randperm(E)[:TOPK] for _ in range(bs)]).to(torch.int32)
+            out = layer._decode_routed(hidden, weights, ids.clone().to(device))
+            _assert_close(out, _reference(pieces, layer.layer_id, hidden, weights, ids))
+
+
+@cuda
+def test_ftw_stores_each_layers_own_banks(tmp_path):
+    from freetoken.checkpoint.convert import _ConvertSink
+    from freetoken.checkpoint.ftw import FTWWriter, load_ftw_banks
+
+    pieces = _pieces(seed=2)
+    reference, _ = _mixed_banks(pieces, torch.device("cuda"))
+    writer = FTWWriter(str(tmp_path))
+    sink = _ConvertSink(writer, names={"gate_up": "gate_up_packed", "down": "down_packed"})
+    banks, _ = _mixed_banks(pieces, torch.device("cuda"), layer_sink=sink)
+    assert banks.streamed and sink.num_layers == L
+    writer.finalize({"quant_format": banks.quant_format, "layer_quant_formats": banks.layer_formats, "expert_bank_num_layers": L})
+
+    loaded = load_ftw_banks(str(tmp_path), num_layers=L)
+    assert loaded.layer_formats == FORMATS and loaded.layer_tags == FORMATS and loaded.quant_format == "nvfp4"
+    assert set(loaded.sources) == set(reference.sources)
+    for role, per_layer in reference.sources.items():
+        for layer, tensor in enumerate(per_layer):
+            got = loaded.sources[role][layer]
+            assert got.shape == tensor.shape and got.dtype == tensor.dtype, (role, layer)
+            assert torch.equal(got.view(torch.uint8), tensor.view(torch.uint8)), (role, layer)

From 21d9578bde22840798b6063824416264fbdcdc25 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 15:50:29 -0700
Subject: [PATCH 110/249] feat(gguf): serve Llama and Mistral llama-arch GGUF
 checkpoints (#722)

---
 python/freetoken/kernel/aot_models.py     |   1 +
 python/freetoken/models/gguf/config.py    |   1 +
 python/freetoken/models/gguf/tokenizer.py |  58 ++++++-
 python/freetoken/models/llama/__init__.py |   3 +-
 python/freetoken/models/llama/gguf.py     | 191 ++++++++++++++++++++++
 python/freetoken/models/llama/model.py    |   2 +
 python/freetoken/models/register.py       |   8 +
 7 files changed, 261 insertions(+), 3 deletions(-)
 create mode 100644 python/freetoken/models/llama/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index 8c5d32ba9d..040236cb76 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -375,6 +375,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         architecture="LlamaForCausalLM",
         hidden_size=4096,
         kv_groups=((8, 128),),
+        arch_aliases=("LlamaGGUFForCausalLM",),
     ),
     AotModel(
         name="Qwen/Qwen2-7B",
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 518acfe4a6..9f860e66ed 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -18,6 +18,7 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "llama": "LlamaGGUFForCausalLM",
     "qwen2": "Qwen2GGUFForCausalLM",
     "qwen3": "Qwen3GGUFForCausalLM",
     "qwen3moe": "Qwen3MoeGGUFForCausalLM",
diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index af715c6598..d6c53fa407 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import json
+
 from .reader import load_gguf_metadata
 
 # end-of-generation token names from llama.cpp (llama_vocab::impl::load in src/llama-vocab.cpp): a GGUF stores one eos id and no stop list
@@ -15,6 +17,57 @@
     "<eos>", "<turn|>", "<|tool_response>", "<\uff5cend\u2581of\u2581sentence\uff5c>", "[e~[", "<|ifm|im_end|>",
 ))
 
+# BPE pre-tokenizers that llama.cpp adds a bos for when the file has no add_bos_token key (llama_vocab::impl::load)
+_BOS_PRE = frozenset((
+    "llama3", "llama-v3", "llama-bpe", "falcon3", "falcon-h1", "pixtral", "midm-2.0", "lfm2", "jina-v5-nano", "tekken", "chameleon",
+))
+
+
+def _add_bos(meta: dict) -> bool:
+    if "tokenizer.ggml.add_bos_token" in meta:
+        return bool(meta["tokenizer.ggml.add_bos_token"])
+    model = meta.get("tokenizer.ggml.model")
+    return model == "llama" or (model == "gpt2" and meta.get("tokenizer.ggml.pre") in _BOS_PRE)
+
+
+def _fix_added_tokens(fast, meta: dict):
+    """Register the control and user-defined tokens that transformers skipped at their file ids, and drop added tokens past the file's vocabulary."""
+    from tokenizers import Tokenizer
+    from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+    tokens = meta["tokenizer.ggml.tokens"]
+    spelling = bytes_to_unicode()
+    alphabet = set(spelling.values())
+    # transformers' byte-level vocabulary spells a token with characters outside the byte alphabet in bytes and then skips it (DeepSeek's full-width-bar tokens, MiniMax's tokens with spaces)
+    skipped = [
+        (i, text, kind == 3)
+        for i, (text, kind) in enumerate(zip(tokens, meta.get("tokenizer.ggml.token_type") or ()))
+        if kind in (3, 4) and meta.get("tokenizer.ggml.model") == "gpt2" and not set(text) <= alphabet and fast.token_to_id(text) != i
+    ]
+    # its GPT2Tokenizer also adds a default <|endoftext|> when the vocabulary lacks one
+    if not skipped and fast.get_vocab_size(with_added_tokens=True) <= len(tokens):
+        return fast
+    # a piece the merges use or build stays a plain BPE token, since renaming it breaks the merges (unsloth's MiniMax-M3 marks U+FFFD as control)
+    pieces = {p for m in meta.get("tokenizer.ggml.merges") or () for p in (*m.split(" "), m.replace(" ", ""))}
+    state = json.loads(fast.to_str())
+    vocab = state["model"]["vocab"]
+    added = [t for t in state["added_tokens"] if t["id"] < len(tokens)]
+    for i, text, special in skipped:
+        spelled = "".join(spelling[b] for b in text.encode("utf-8"))
+        if spelled in pieces or vocab.get(spelled) != i:
+            continue
+        # renamed so the added token keeps the file id instead of taking a new one
+        vocab[text] = vocab.pop(spelled)
+        added.append({"id": i, "content": text, "single_word": False, "lstrip": False, "rstrip": False, "normalized": False, "special": special})
+    state["added_tokens"] = sorted(added, key=lambda t: t["id"])
+    return Tokenizer.from_str(json.dumps(state))
+
+
+def _known_specials(special: dict, fast) -> dict:
+    """Drop a bos/eos/unk/pad the vocabulary lacks, which PreTrainedTokenizerFast would add past it (unsloth's MiniMax-M3 unk is U+FFFD stored decoded)."""
+    vocab = fast.get_vocab(with_added_tokens=True)
+    return {k: v for k, v in special.items() if k == "chat_template" or v is None or v in vocab}
+
 
 def load_gguf_tokenizer(model_path: str):
     from transformers import PreTrainedTokenizerFast
@@ -23,11 +76,12 @@ def load_gguf_tokenizer(model_path: str):
     arch, tok_dict, special = get_gguf_tokenizer(model_path)
     fast, _ = convert_gguf_tokenizer(arch, tok_dict)
     meta = load_gguf_metadata(model_path)
+    fast = _fix_added_tokens(fast, meta)
     # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template
     return PreTrainedTokenizerFast(
         tokenizer_object=fast,
-        **special,
-        add_bos_token=bool(meta.get("tokenizer.ggml.add_bos_token")),
+        **_known_specials(special, fast),
+        add_bos_token=_add_bos(meta),
         add_eos_token=bool(meta.get("tokenizer.ggml.add_eos_token")),
     )
 
diff --git a/python/freetoken/models/llama/__init__.py b/python/freetoken/models/llama/__init__.py
index 9ffad0e95f..31ffa447cd 100644
--- a/python/freetoken/models/llama/__init__.py
+++ b/python/freetoken/models/llama/__init__.py
@@ -1,5 +1,6 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import LlamaForCausalLM
 from .weight import iter_weights
 
-__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights"]
+__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py
new file mode 100644
index 0000000000..09d02627ae
--- /dev/null
+++ b/python/freetoken/models/llama/gguf.py
@@ -0,0 +1,191 @@
+"""Llama and Mistral from a llama.cpp ``llama`` GGUF: config from the GGUF metadata, tensors by name.
+
+llama.cpp runs interleaved rope on this arch, so its converter (conversion/llama.py) permutes the rows of each q and k head; the loader undoes that, and the model keeps HF's NeoX rope.
+The RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "llama"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    names = sorted(gguf_tensor_index(shim.model_path))
+    scaled = [n for n in names if n.endswith((".scale", ".input_scale"))]
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    # the llama arch also carries Mixtral experts and optional biases, which the llama model does not have
+    extra = [n for n in names if n.endswith(".bias") or ".ffn_gate_inp." in n]
+    if extra:
+        raise NotImplementedError(f"{shim.model_path}: llama GGUF tensors with biases or MoE experts are not supported (e.g. {extra[0]})")
+
+
+def _rope_scaling(shim: GgufConfigShim, base: float, head_dim: int) -> dict | None:
+    """The HF rope scaling of the file: none, or Llama 3's, which the converter stores only as the ``rope_freqs`` divisors."""
+    md = shim.metadata
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    # llama.cpp scales linearly by rope.scaling.factor (or the older rope.scale_linear) when the type key is absent
+    factor = float(md.get(f"{_ARCH}.rope.scaling.factor", md.get(f"{_ARCH}.rope.scale_linear", 0.0)))
+    attn_factor = float(md.get(f"{_ARCH}.rope.scaling.attn_factor", 1.0))
+    if scaling not in ("none", "linear") or factor not in (0.0, 1.0) or attn_factor != 1.0:
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} (factor {factor}, attn_factor {attn_factor}) is not supported")
+    index = gguf_tensor_index(shim.model_path)
+    if not index:
+        raise NotImplementedError(f"{shim.model_path}: a metadata-only llama GGUF has no rope_freqs tensor to tell Llama 3 rope from plain rope")
+    info = index.get("rope_freqs.weight")
+    if info is None:
+        return None
+    from freetoken.layers.rotary import _rope_frequencies
+
+    divisors = float_tensor(info).float()
+    # Llama 3.1-3.3 differ only in factor
+    llama3 = {
+        "rope_type": "llama3",
+        "factor": float(divisors.max()),
+        "low_freq_factor": 1.0,
+        "high_freq_factor": 4.0,
+        "original_max_position_embeddings": 8192,
+    }
+    plain, _ = _rope_frequencies(head_dim, head_dim, base, None)
+    scaled, _ = _rope_frequencies(head_dim, head_dim, base, llama3)
+    if not torch.allclose(plain / scaled, divisors, rtol=1e-5):
+        raise NotImplementedError(f"{shim.model_path}: rope_freqs does not match Llama 3 rope scaling")
+    return llama3
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    hidden, heads = int(kv("embedding_length")), int(kv("attention.head_count"))
+    # Mistral-7B files carry no key_length
+    head_dim = int(md.get(f"{_ARCH}.attention.key_length", hidden // heads))
+    q = gguf_tensor_index(shim.model_path).get("blk.0.attn_q.weight")
+    # early Mistral-Nemo conversions lack key_length, so hidden / heads gives the wrong head_dim; llama.cpp rejects them too
+    if q is not None and q.ne[1] != heads * head_dim:
+        raise NotImplementedError(f"{shim.model_path}: attn_q has {q.ne[1]} rows, not head_count x head_dim = {heads * head_dim}")
+    if int(md.get(f"{_ARCH}.rope.dimension_count", head_dim)) != head_dim:
+        raise NotImplementedError(f"{shim.model_path}: partial rope is not supported")
+    base = float(md.get(f"{_ARCH}.rope.freq_base", 10000.0))
+    return SimpleNamespace(
+        model_type="llama",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=hidden,
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=heads,
+        num_key_value_heads=int(md.get(f"{_ARCH}.attention.head_count_kv", heads)),
+        head_dim=head_dim,
+        rope_theta=base,
+        rope_scaling=_rope_scaling(shim, base, head_dim),
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def _unpermute(rows: torch.Tensor, heads: int) -> torch.Tensor:
+    """Undo the converter's rope permute of one q or k matrix (``LlamaModel.permute``); it moves whole rows, so packed ggml rows work too."""
+    return rows.reshape(heads, rows.shape[0] // heads // 2, 2, *rows.shape[1:]).swapaxes(1, 2).reshape(rows.shape)
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden, heads, kv_heads = hf.hidden_size, hf.num_attention_heads, hf.num_key_value_heads
+    q_rows, kv_rows = heads * hf.head_dim, kv_heads * hf.head_dim
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        if f"{blk}.attn_qkv.weight" in index:
+            qkv = [t(f"{blk}.attn_qkv")] * 3
+            q, k, v = matrix(qkv[0]).split([q_rows, kv_rows, kv_rows])
+        else:
+            qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+            q, k, v = (matrix(p) for p in qkv)
+        yield from fused(f"{sa}.qkv_proj", qkv, [_unpermute(q, heads), _unpermute(k, kv_heads), v], hidden)
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py
index eb7ff28fac..107d9e0a2c 100644
--- a/python/freetoken/models/llama/model.py
+++ b/python/freetoken/models/llama/model.py
@@ -48,6 +48,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 263b5f932c..057e2ae6ed 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -95,6 +95,14 @@ class ModelSpec:
         "LlamaForCausalLM",
         packed_modules_mapping=_DENSE_PACKED,
     ),
+    # llama.cpp llama GGUF (Llama, Mistral): the llama model classes with every tensor in its packed ggml type
+    "LlamaGGUFForCausalLM": ModelSpec(
+        "freetoken.models.llama",
+        "LlamaForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_DENSE_PACKED,
+    ),
     "Qwen2ForCausalLM": ModelSpec(
         "freetoken.models.qwen2",
         "Qwen2ForCausalLM",

From bc0e2feb7cf01fdd7b082381eda142fc84c029d2 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 16:48:24 -0700
Subject: [PATCH 111/249] perf(attention): copy triton metadata without
 blocking the host (#726)

---
 python/freetoken/attention/triton.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py
index 8bf9792d55..30bedf4ffa 100644
--- a/python/freetoken/attention/triton.py
+++ b/python/freetoken/attention/triton.py
@@ -230,17 +230,16 @@ def prepare_metadata(self, batch: Batch) -> None:
         cached_lens = [req.cached_len for req in reqs]
         num_query_tokens = sum(seqlens_q)
         is_decode = max(seqlens_q) == 1
-        prefix_lens = torch.tensor(cached_lens, dtype=torch.int32, device=device)
+        CPU_KWARGS = {"device": "cpu", "dtype": torch.int32, "pin_memory": device.type == "cuda"}
+        prefix_lens = torch.tensor(cached_lens, **CPU_KWARGS).to(device, non_blocking=True)
 
-        indptr = torch.tensor([0] + seqlens_k, dtype=torch.int32, device=device).cumsum_(0)
+        indptr = torch.tensor([0] + seqlens_k, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True)
         if is_decode:
             cu_seqlens_q_gpu = torch.arange(0, padded_size + 1, device=device, dtype=torch.int32)
         elif all(l == 0 for l in cached_lens):
             cu_seqlens_q_gpu = indptr
         else:
-            cu_seqlens_q_gpu = torch.tensor(
-                [0] + seqlens_q, dtype=torch.int32, device=device
-            ).cumsum_(0)
+            cu_seqlens_q_gpu = torch.tensor([0] + seqlens_q, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True)
         indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs])
         swa_indices = None
         if getattr(self.kvcache, "swa_paged", False):

From 312f324d65b3f49596bae424cfdf2cc5c0274f1b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 17:40:36 -0700
Subject: [PATCH 112/249] test(server): list the llama and qwen2 gguf archs in
 the parser auto-selection tables (#727)

---
 tests/server/test_parser_auto_selection.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py
index e1eb87982b..6bdc6579fd 100644
--- a/tests/server/test_parser_auto_selection.py
+++ b/tests/server/test_parser_auto_selection.py
@@ -31,14 +31,16 @@
 # parser, and a new architecture landing here is the bug this file exists to catch.
 NO_REASONING_FORMAT = {
     "LlamaForCausalLM",
+    "LlamaGGUFForCausalLM",
     "MistralForCausalLM",
     "Mistral3ForConditionalGeneration",
     "Qwen2ForCausalLM",
+    "Qwen2GGUFForCausalLM",
 }
 
 # `llama3` is the end of the cascade -- the answer when nothing matched.
 GENERIC_TOOL_CALL_FALLBACK = "llama3"
-NO_DEDICATED_TOOL_FORMAT = {"LlamaForCausalLM"}
+NO_DEDICATED_TOOL_FORMAT = {"LlamaForCausalLM", "LlamaGGUFForCausalLM"}
 
 
 class _Config:

From e0dbb7dee18fac3f752e040f1f2f5476b110c1a3 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 17:41:06 -0700
Subject: [PATCH 113/249] fix(gguf): set ignore_merges for llama 3 gguf
 tokenizers (#728)

---
 python/freetoken/models/gguf/tokenizer.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index d6c53fa407..7f0bf29a48 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -22,6 +22,9 @@
     "llama3", "llama-v3", "llama-bpe", "falcon3", "falcon-h1", "pixtral", "midm-2.0", "lfm2", "jina-v5-nano", "tekken", "chameleon",
 ))
 
+# Llama 3 pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it
+_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe"))
+
 
 def _add_bos(meta: dict) -> bool:
     if "tokenizer.ggml.add_bos_token" in meta:
@@ -77,6 +80,8 @@ def load_gguf_tokenizer(model_path: str):
     fast, _ = convert_gguf_tokenizer(arch, tok_dict)
     meta = load_gguf_metadata(model_path)
     fast = _fix_added_tokens(fast, meta)
+    if meta.get("tokenizer.ggml.model") == "gpt2" and meta.get("tokenizer.ggml.pre") in _IGNORE_MERGES_PRE:
+        fast.model.ignore_merges = True
     # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template
     return PreTrainedTokenizerFast(
         tokenizer_object=fast,

From 5d39d6408a23b0c978fbff47789b1126c58c3bcd Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 18:38:43 -0700
Subject: [PATCH 114/249] feat(gguf): serve Muse-Glimmer GGUF checkpoints
 (#724)

---
 python/freetoken/kernel/aot_models.py         |   1 +
 python/freetoken/models/gguf/config.py        |   1 +
 .../freetoken/models/muse_glimmer/__init__.py |   4 +
 python/freetoken/models/muse_glimmer/gguf.py  | 184 ++++++++++++++++++
 python/freetoken/models/muse_glimmer/model.py |   2 +
 python/freetoken/models/register.py           |  17 +-
 6 files changed, 205 insertions(+), 4 deletions(-)
 create mode 100644 python/freetoken/models/muse_glimmer/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index 040236cb76..eca18d0025 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -369,6 +369,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         hidden_size=6656,
         kv_groups=((2, 128),),  # sliding and full (NoPE) layers share the same geometry
         aliases=("RedHatAI/Muse-Glimmer-30B-NVFP4",),
+        arch_aliases=("MuseGlimmerGGUFForCausalLM",),
     ),
     AotModel(
         name="meta-llama/Llama-3.1-8B-Instruct",
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 9f860e66ed..6e31b8ff01 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -19,6 +19,7 @@
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
     "llama": "LlamaGGUFForCausalLM",
+    "muse-glimmer": "MuseGlimmerGGUFForCausalLM",
     "qwen2": "Qwen2GGUFForCausalLM",
     "qwen3": "Qwen3GGUFForCausalLM",
     "qwen3moe": "Qwen3MoeGGUFForCausalLM",
diff --git a/python/freetoken/models/muse_glimmer/__init__.py b/python/freetoken/models/muse_glimmer/__init__.py
index 4ee14808c4..ff8d3d1710 100644
--- a/python/freetoken/models/muse_glimmer/__init__.py
+++ b/python/freetoken/models/muse_glimmer/__init__.py
@@ -1,5 +1,6 @@
 from .attention import MuseGlimmerAttention
 from .config import VisionConfig, parse_config, parse_vision_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration
 from .vision import MuseGlimmerVisionModel
 from .weight import iter_vision_weights, iter_weights
@@ -14,4 +15,7 @@
     "parse_vision_config",
     "iter_vision_weights",
     "iter_weights",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/muse_glimmer/gguf.py b/python/freetoken/models/muse_glimmer/gguf.py
new file mode 100644
index 0000000000..7d5cee3b33
--- /dev/null
+++ b/python/freetoken/models/muse_glimmer/gguf.py
@@ -0,0 +1,184 @@
+"""Muse Glimmer text from a llama.cpp ``muse-glimmer`` GGUF: config from the GGUF metadata, tensors by name.
+
+The converter (conversion/muse_glimmer.py) adds 1 to the four norms of each block, permutes the rows of each q and k head for llama.cpp's interleaved rope, and writes constant q/k norms that carry ``qk_scale_factor``.
+The loader subtracts the 1, undoes the permute, and reads ``qk_scale_factor`` back from the q norm.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "muse-glimmer"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.gate_proj": "attn_gate",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+# checkpoint norm -> GGUF norm; the GGUF stores each as weight + 1
+_SHIFTED_NORMS = {
+    "input_layernorm": "attn_norm",
+    "post_attention_layernorm": "post_attention_norm",
+    "pre_feedforward_layernorm": "ffn_norm",
+    "post_feedforward_layernorm": "post_ffw_norm",
+}
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    index = gguf_tensor_index(shim.model_path)
+    if not index:
+        raise NotImplementedError(f"{shim.model_path}: a metadata-only muse-glimmer GGUF has no q norm to read qk_scale_factor from")
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in index if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _qk_scale_factor(model_path: str, num_layers: int) -> float:
+    """The constant the converter wrote into every q norm; the k norms are all ones."""
+    index = gguf_tensor_index(model_path)
+    scale = float(float_tensor(index["blk.0.attn_q_norm.weight"])[0])
+    for layer in range(num_layers):
+        q, k = (float_tensor(index[f"blk.{layer}.attn_{p}_norm.weight"]).float() for p in ("q", "k"))
+        if not (torch.all(q == scale) and torch.all(k == 1.0)):
+            raise NotImplementedError(f"{model_path}: the q/k norms of layer {layer} are not the converter's constant qk_scale_factor / ones")
+    return scale
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF text config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    num_layers = int(kv("block_count"))
+    # a bool per layer, or llama.cpp's period: every n-th layer is full attention, and 0 means none is
+    pattern = md.get(f"{_ARCH}.attention.sliding_window_pattern", 4)
+    if isinstance(pattern, list):
+        swa = [bool(p) for p in pattern]
+        if len(swa) != num_layers:
+            raise ValueError(f"{shim.model_path}: sliding_window_pattern has {len(swa)} entries for {num_layers} layers")
+    else:
+        swa = [int(pattern) == 0 or i % int(pattern) < int(pattern) - 1 for i in range(num_layers)]
+    base = float(kv("rope.freq_base"))
+    swa_base = float(md.get(f"{_ARCH}.rope.freq_base_swa", base))
+    return SimpleNamespace(
+        model_type="muse_glimmer",
+        architectures=list(shim.architectures),
+        hidden_activation="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=num_layers,
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(kv("attention.key_length")),
+        intermediate_size=int(kv("feed_forward_length")),
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        # llama.cpp hard-codes the post-norm eps and NoPE on the full layers (src/models/muse-glimmer.cpp)
+        post_norm_eps=1e-8,
+        layer_types=["sliding_attention" if s else "full_attention" for s in swa],
+        layer_rope_theta=[swa_base if s else 0.0 for s in swa],
+        sliding_window=int(kv("attention.sliding_window")),
+        qk_scale_factor=_qk_scale_factor(shim.model_path, num_layers),
+        # llama.cpp's default when the key is absent; 0 turns the cap off
+        final_logit_softcapping=float(md.get(f"{_ARCH}.final_logit_softcapping", 30.0)) or None,
+        output_multiplier=float(kv("logit_scale")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def _unpermute(rows: torch.Tensor, heads: int) -> torch.Tensor:
+    """Undo the converter's rope permute of one q or k matrix (``_unpermute_for_rope``); it moves whole rows, so packed ggml rows work too."""
+    return rows.reshape(heads, rows.shape[0] // heads // 2, 2, *rows.shape[1:]).swapaxes(1, 2).reshape(rows.shape)
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The text state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden, heads, kv_heads = hf.hidden_size, hf.num_attention_heads, hf.num_key_value_heads
+    q_rows, kv_rows = heads * hf.head_dim, kv_heads * hf.head_dim
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        for norm, tensor in _SHIFTED_NORMS.items():
+            yield f"{pre}.{norm}.weight", float_tensor(t(f"{blk}.{tensor}")).float() - 1.0
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        if f"{blk}.attn_qkv.weight" in index:
+            qkv = [t(f"{blk}.attn_qkv")] * 3
+            q, k, v = matrix(qkv[0]).split([q_rows, kv_rows, kv_rows])
+        else:
+            qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+            q, k, v = (matrix(p) for p in qkv)
+        gate = t(f"{blk}.attn_gate")
+        yield from fused(f"{sa}.qkvg_proj", [*qkv, gate], [_unpermute(q, heads), _unpermute(k, kv_heads), v, matrix(gate)], hidden)
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/muse_glimmer/model.py b/python/freetoken/models/muse_glimmer/model.py
index 50bd450670..d51c211743 100644
--- a/python/freetoken/models/muse_glimmer/model.py
+++ b/python/freetoken/models/muse_glimmer/model.py
@@ -105,6 +105,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         # NormedEmbedding: a weightless RMSNorm on top of the embeddings. Kept separate
         # from the embedding matrix (the reference cannot fold it either -- the DFlash
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 057e2ae6ed..08b70e832b 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -78,6 +78,10 @@ class ModelSpec:
 _GEMMA4_UNIFIED_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4UnifiedMMProcessor"
 _GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MUSE_GLIMMER_PROCESSOR = "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor"
+_MUSE_GLIMMER_PACKED = (
+    ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")),
+    ("gate_up_proj", ("gate_proj", "up_proj")),
+)
 _MUSE_GLIMMER_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MINIMAX_M3_PACKED = _DENSE_PACKED + (
     ("index_qk_proj", ("index_q_proj", "index_k_proj")),
@@ -293,13 +297,18 @@ class ModelSpec:
         "freetoken.models.muse_glimmer",
         "MuseGlimmerForConditionalGeneration",
         checkpoint_roots=_LANGUAGE_MODEL_ROOT,
-        packed_modules_mapping=(
-            ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")),
-            ("gate_up_proj", ("gate_proj", "up_proj")),
-        ),
+        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
         mm_processor=_MUSE_GLIMMER_PROCESSOR,
         encoders=_MUSE_GLIMMER_ENCODERS,
     ),
+    # llama.cpp muse-glimmer GGUF (text only): the same model classes with every tensor in its packed ggml type
+    "MuseGlimmerGGUFForCausalLM": ModelSpec(
+        "freetoken.models.muse_glimmer",
+        "MuseGlimmerForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
+    ),
     "MistralForCausalLM": ModelSpec(
         "freetoken.models.mistral",
         "MistralForCausalLM",

From 4663b2b89e60bd0fd7eacc1f7596794ac7984e3c Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 19:01:30 -0700
Subject: [PATCH 115/249] feat(gguf): serve MiniMax-M2 and MiniMax-M3 GGUF
 checkpoints (#725)

---
 python/freetoken/kernel/aot_models.py         |   3 +-
 python/freetoken/models/gguf/config.py        |   2 +
 .../freetoken/models/minimax_m2/__init__.py   |   4 +
 python/freetoken/models/minimax_m2/gguf.py    | 161 ++++++++++++++
 python/freetoken/models/minimax_m2/model.py   |   2 +
 .../freetoken/models/minimax_m3/__init__.py   |   4 +
 .../freetoken/models/minimax_m3/attention.py  |   1 +
 python/freetoken/models/minimax_m3/gguf.py    | 198 ++++++++++++++++++
 python/freetoken/models/minimax_m3/mlp.py     |   1 +
 python/freetoken/models/minimax_m3/model.py   |   2 +
 python/freetoken/models/register.py           |  19 ++
 python/freetoken/moe/expert_banks.py          |   2 +-
 12 files changed, 397 insertions(+), 2 deletions(-)
 create mode 100644 python/freetoken/models/minimax_m2/gguf.py
 create mode 100644 python/freetoken/models/minimax_m3/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index eca18d0025..f0b70e2024 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -300,6 +300,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         moe_intermediate_size=1536,  # experts reuse the dense intermediate_size
         expert_formats=_NVFP4_FORMATS,
         aliases=("MiniMaxAI/MiniMax-M2.5",),
+        arch_aliases=("MiniMaxM2GGUFForCausalLM",),
     ),
     AotModel(
         # GQA block-sparse (BSA): the BSAKVCache stores K/V through store_cache
@@ -318,7 +319,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         moe_intermediate_size=3072,
         expert_formats=("nvfp4",),
         aliases=("MiniMaxAI/MiniMax-M3",),
-        arch_aliases=("MiniMaxM3SparseForCausalLM",),
+        arch_aliases=("MiniMaxM3SparseForCausalLM", "MiniMaxM3GGUFForCausalLM"),
     ),
     AotModel(
         name="deepseek-ai/DeepSeek-V4-Flash",
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 6e31b8ff01..677adff621 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -19,6 +19,8 @@
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
     "llama": "LlamaGGUFForCausalLM",
+    "minimax-m2": "MiniMaxM2GGUFForCausalLM",
+    "minimax-m3": "MiniMaxM3GGUFForCausalLM",
     "muse-glimmer": "MuseGlimmerGGUFForCausalLM",
     "qwen2": "Qwen2GGUFForCausalLM",
     "qwen3": "Qwen3GGUFForCausalLM",
diff --git a/python/freetoken/models/minimax_m2/__init__.py b/python/freetoken/models/minimax_m2/__init__.py
index ad276f8457..edb9896d24 100644
--- a/python/freetoken/models/minimax_m2/__init__.py
+++ b/python/freetoken/models/minimax_m2/__init__.py
@@ -1,4 +1,5 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import MiniMaxM2ForCausalLM
 from .weight import iter_weights, nvfp4_expert_spec
 
@@ -7,4 +8,7 @@
     "MiniMaxM2ForCausalLM",
     "parse_config",
     "iter_weights",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/minimax_m2/gguf.py b/python/freetoken/models/minimax_m2/gguf.py
new file mode 100644
index 0000000000..99047b404c
--- /dev/null
+++ b/python/freetoken/models/minimax_m2/gguf.py
@@ -0,0 +1,161 @@
+"""MiniMax-M2 from a llama.cpp ``minimax-m2`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
+
+The llama.cpp converter (conversion/minimax.py) only stacks the w1 / w2 / w3 experts: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "minimax-m2"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "block_sparse_moe.gate": "ffn_gate_inp",
+    "block_sparse_moe.experts.0.w1": "ffn_gate_exps",
+    "block_sparse_moe.experts.0.w3": "ffn_up_exps",
+    "block_sparse_moe.experts.0.w2": "ffn_down_exps",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+_SIGMOID_GATING = 2
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    index = gguf_tensor_index(shim.model_path)
+    scaled = sorted(n for n in index if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    # a metadata-only GGUF (an FTW dir) has no tensor table to check
+    if index and "blk.0.exp_probs_b.bias" not in index:
+        raise NotImplementedError(f"{shim.model_path}: converted by an older llama.cpp that names the router bias exp_probs_b, which current llama.cpp refuses too")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    gating = int(md.get(f"{_ARCH}.expert_gating_func", 0))
+    if gating != _SIGMOID_GATING:
+        raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported; MiniMax-M2 routes with a sigmoid (2)")
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    # a list holds one size per layer
+    expert_ff = kv("expert_feed_forward_length")
+    if isinstance(expert_ff, list):
+        if len(set(expert_ff)) != 1:
+            raise NotImplementedError(f"{shim.model_path}: per-layer expert sizes {sorted(set(expert_ff))} are not supported")
+        expert_ff = expert_ff[0]
+    return SimpleNamespace(
+        model_type="minimax_m2",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(kv("attention.key_length")),
+        rotary_dim=int(kv("rope.dimension_count")),
+        rope_theta=float(kv("rope.freq_base")),
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(expert_ff),
+        num_local_experts=int(kv("expert_count")),
+        num_experts_per_tok=int(kv("expert_used_count")),
+        # llama.cpp always builds the per-layer q/k norms of this arch
+        use_qk_norm=True,
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return replace(parse_config(_hf_config(shim)), expert_quant="gguf")
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+    q_rows, kv_rows = hf.num_attention_heads * hf.head_dim, hf.num_key_value_heads * hf.head_dim
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        if f"{blk}.attn_qkv.weight" in index:
+            qkv = [t(f"{blk}.attn_qkv")] * 3
+            parts = list(matrix(qkv[0]).split([q_rows, kv_rows, kv_rows]))
+        else:
+            qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+            parts = [matrix(p) for p in qkv]
+        yield from fused(f"{sa}.qkv_proj", qkv, parts, hidden)
+        yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+        yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        moe = f"{pre}.block_sparse_moe"
+        yield f"{moe}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+        yield f"{moe}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias"))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/minimax_m2/model.py b/python/freetoken/models/minimax_m2/model.py
index 457ab56500..2f29a87353 100644
--- a/python/freetoken/models/minimax_m2/model.py
+++ b/python/freetoken/models/minimax_m2/model.py
@@ -41,6 +41,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/minimax_m3/__init__.py b/python/freetoken/models/minimax_m3/__init__.py
index 4abf9c8793..8d0aea79bf 100644
--- a/python/freetoken/models/minimax_m3/__init__.py
+++ b/python/freetoken/models/minimax_m3/__init__.py
@@ -1,4 +1,5 @@
 from .config import VisionConfig, parse_config, parse_vision_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import MiniMaxM3ForCausalLM, MiniMaxM3ForConditionalGeneration
 from .vision import MiniMaxM3VisionModel
 from .weight import (
@@ -17,4 +18,7 @@
     "parse_vision_config",
     "iter_vision_weights",
     "iter_weights",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/minimax_m3/attention.py b/python/freetoken/models/minimax_m3/attention.py
index 0d55c07294..4fcf57e103 100644
--- a/python/freetoken/models/minimax_m3/attention.py
+++ b/python/freetoken/models/minimax_m3/attention.py
@@ -45,6 +45,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
 
         self.qkv_proj = LinearReplicated(
             args.hidden_size, self.qo_attn_dim + 2 * self.kv_attn_dim, has_bias=False,
+            output_sizes=[self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim],
             quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
         )
         self.o_proj = LinearReplicated(
diff --git a/python/freetoken/models/minimax_m3/gguf.py b/python/freetoken/models/minimax_m3/gguf.py
new file mode 100644
index 0000000000..92742c3aae
--- /dev/null
+++ b/python/freetoken/models/minimax_m3/gguf.py
@@ -0,0 +1,198 @@
+"""MiniMax-M3 text from a llama.cpp ``minimax-m3`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
+
+The llama.cpp converter (conversion/minimax.py) stacks the w1 / w2 / w3 experts, stores the indexer projections as F32, and adds 1 to every RMSNorm weight so llama.cpp can run plain RMSNorm; the loader subtracts the 1 for this family's (1 + w) norms.
+Both sides run NeoX rope, so q/k carry no rope permute.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "minimax-m3"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "self_attn.index_q_proj": "indexer.q_proj",
+    "self_attn.index_k_proj": "indexer.k_proj",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+    "block_sparse_moe.gate": "ffn_gate_inp",
+    "block_sparse_moe.shared_experts.gate_proj": "ffn_gate_shexp",
+    "block_sparse_moe.shared_experts.up_proj": "ffn_up_shexp",
+    "block_sparse_moe.shared_experts.down_proj": "ffn_down_shexp",
+    "block_sparse_moe.experts.0.w1": "ffn_gate_exps",
+    "block_sparse_moe.experts.0.w3": "ffn_up_exps",
+    "block_sparse_moe.experts.0.w2": "ffn_down_exps",
+}
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+_SIGMOID_GATING = 2
+# llama.cpp hard-codes these for the arch (src/llama-graph.cpp LLM_FFN_SWIGLU_OAI_MOE)
+_SWIGLU_ALPHA, _SWIGLU_LIMIT = 1.702, 7.0
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    tensor = _LAYER_LEAVES.get(m.group(2))
+    return (f"blk.{m.group(1)}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF text config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    if f"{_ARCH}.attention.indexer.head_count" not in md:
+        raise NotImplementedError(f"{shim.model_path}: converted by an older llama.cpp without the sparse-attention indexer, which current llama.cpp refuses too")
+    gating = int(kv("expert_gating_func"))
+    if gating != _SIGMOID_GATING:
+        raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported; MiniMax-M3 routes with a sigmoid (2)")
+    if not md.get(f"{_ARCH}.expert_weights_norm", False):
+        raise NotImplementedError(f"{shim.model_path}: routing without renormalized expert weights is not supported")
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    expert_ff = kv("expert_feed_forward_length")
+    if isinstance(expert_ff, list):
+        if len(set(expert_ff)) != 1:
+            raise NotImplementedError(f"{shim.model_path}: per-layer expert sizes {sorted(set(expert_ff))} are not supported")
+        expert_ff = expert_ff[0]
+    num_layers = int(kv("block_count"))
+    dense = int(md.get(f"{_ARCH}.leading_dense_block_count", 0))
+    # llama.cpp runs sparse attention and the MoE FFN on every layer after the leading dense ones
+    sparse = [0] * dense + [1] * (num_layers - dense)
+    return SimpleNamespace(
+        model_type="minimax_m3",
+        architectures=list(shim.architectures),
+        hidden_act="swigluoai",
+        swiglu_alpha=_SWIGLU_ALPHA,
+        swiglu_limit=_SWIGLU_LIMIT,
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=num_layers,
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(kv("attention.key_length")),
+        rotary_dim=int(kv("rope.dimension_count")),
+        rope_theta=float(kv("rope.freq_base")),
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        dense_intermediate_size=int(kv("feed_forward_length")),
+        intermediate_size=int(expert_ff),
+        shared_intermediate_size=int(expert_ff),
+        n_shared_experts=int(kv("expert_shared_count")),
+        num_local_experts=int(kv("expert_count")),
+        num_experts_per_tok=int(kv("expert_used_count")),
+        # llama.cpp skips the scale when the key is absent or 0
+        routed_scaling_factor=float(md.get(f"{_ARCH}.expert_weights_scale", 0.0)) or 1.0,
+        use_routing_bias=True,
+        moe_layer_freq=sparse,
+        sparse_attention_config={
+            "use_sparse_attention": True,
+            "sparse_attention_freq": sparse,
+            "sparse_num_index_heads": int(kv("attention.indexer.head_count")),
+            "sparse_index_dim": int(kv("attention.indexer.key_length")),
+            "sparse_topk_blocks": int(kv("attention.indexer.top_k")),
+            "sparse_block_size": int(kv("attention.indexer.block_size")),
+            "sparse_local_block": int(kv("attention.indexer.local_blocks")),
+            "sparse_init_block": 0,
+            "sparse_score_type": "max",
+        },
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return replace(parse_config(_hf_config(shim)), expert_quant="gguf")
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense text state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    config = parse_gguf_config(cached_load_hf_config(model_path))
+    args = config.m3_args
+    index = gguf_tensor_index(model_path)
+    hidden = config.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    def norm(name: str) -> torch.Tensor:
+        return float_tensor(t(name)).float() - 1.0
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not config.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", norm("output_norm")
+
+    for layer in range(config.num_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm")
+        yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.ffn_norm")
+        sa = f"{pre}.self_attn"
+        qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm")
+        yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm")
+        if args.is_sparse_layer(layer):
+            iq, ik = t(f"{blk}.indexer.q_proj"), t(f"{blk}.indexer.k_proj")
+            yield from fused(f"{sa}.index_qk_proj", [iq, ik], [matrix(iq), matrix(ik)], hidden)
+            yield f"{sa}.index_q_norm.weight", norm(f"{blk}.indexer.q_norm")
+            yield f"{sa}.index_k_norm.weight", norm(f"{blk}.indexer.k_norm")
+        if layer in args.moe_layer_ids:
+            moe = f"{pre}.block_sparse_moe"
+            yield f"{moe}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")).float()
+            yield f"{moe}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias")).float()
+            mlp, parts = f"{moe}.shared_experts", [t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")]
+        else:
+            mlp, parts = f"{pre}.mlp", [t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")]
+        gate, up, down = parts
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/minimax_m3/mlp.py b/python/freetoken/models/minimax_m3/mlp.py
index d5332eed0c..62b02e8040 100644
--- a/python/freetoken/models/minimax_m3/mlp.py
+++ b/python/freetoken/models/minimax_m3/mlp.py
@@ -36,6 +36,7 @@ def __init__(
     ):
         self.gate_up_proj = LinearReplicated(
             hidden_size, 2 * intermediate_size, has_bias=False,
+            output_sizes=[intermediate_size, intermediate_size],
             quant_config=quant_config, prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = LinearReplicated(
diff --git a/python/freetoken/models/minimax_m3/model.py b/python/freetoken/models/minimax_m3/model.py
index d78fc10afd..b6f271b35c 100644
--- a/python/freetoken/models/minimax_m3/model.py
+++ b/python/freetoken/models/minimax_m3/model.py
@@ -71,6 +71,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 08b70e832b..bfa0a1d6ce 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -48,6 +48,8 @@ class ModelSpec:
 # per-expert checkpoints: probe the first expert's projections for the experts container
 _EXPERTS_PACKED = (("experts", ("experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj")),)
 _EXPERTS_W123_PACKED = (("experts", ("experts.0.w1", "experts.0.w2", "experts.0.w3")),)
+# the GGUF scheme reads expert types in (gate, up, down) order
+_EXPERTS_W132_PACKED = (("experts", ("experts.0.w1", "experts.0.w3", "experts.0.w2")),)
 # pre-stacked expert checkpoints (Qwen3-VL-MoE): the container's own two tensors are the probe
 _STACKED_EXPERTS_PACKED = (("experts", ("experts.gate_up_proj", "experts.down_proj")),)
 _QWEN3_5_PACKED = _DENSE_PACKED + (
@@ -86,6 +88,7 @@ class ModelSpec:
 _MINIMAX_M3_PACKED = _DENSE_PACKED + (
     ("index_qk_proj", ("index_q_proj", "index_k_proj")),
 ) + _EXPERTS_W123_PACKED
+_MINIMAX_M3_GGUF_PACKED = _DENSE_PACKED + (("index_qk_proj", ("index_q_proj", "index_k_proj")),) + _EXPERTS_W132_PACKED
 _MINIMAX_M3_PROCESSOR = "freetoken.mm.processors.minimax_m3:MiniMaxM3MMProcessor"
 _MINIMAX_M3_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 # DeepSeek-V4's vision dims sit at the top level of the checkpoint config, next to the
@@ -166,6 +169,14 @@ class ModelSpec:
         "MiniMaxM2ForCausalLM",
         packed_modules_mapping=_DENSE_PACKED + _EXPERTS_W123_PACKED,
     ),
+    # llama.cpp minimax-m2 GGUF: the same model classes with every tensor in its packed ggml type
+    "MiniMaxM2GGUFForCausalLM": ModelSpec(
+        "freetoken.models.minimax_m2",
+        "MiniMaxM2ForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_DENSE_PACKED + _EXPERTS_W132_PACKED,
+    ),
     # MiniMax-M3 (model_type minimax_m3_vl): multimodal wrapper config (text tower in
     # text_config, weights under language_model.); the CLIP-style tower under vision_tower.
     # serves image input. GQA + block-sparse attention (lightning indexer, top-k 128-token
@@ -185,6 +196,14 @@ class ModelSpec:
         "MiniMaxM3ForCausalLM",
         packed_modules_mapping=_MINIMAX_M3_PACKED,
     ),
+    # llama.cpp minimax-m3 GGUF (text only): the same model classes with every tensor in its packed ggml type
+    "MiniMaxM3GGUFForCausalLM": ModelSpec(
+        "freetoken.models.minimax_m3",
+        "MiniMaxM3ForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_MINIMAX_M3_GGUF_PACKED,
+    ),
     "DeepseekV4ForCausalLM": ModelSpec(
         "freetoken.models.deepseek_v4",
         "DeepseekV4ForCausalLM",
diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py
index 7e25e58f73..b79d28609e 100644
--- a/python/freetoken/moe/expert_banks.py
+++ b/python/freetoken/moe/expert_banks.py
@@ -423,7 +423,7 @@ def _build(par: bool) -> ExpertBanks:
             from freetoken.models.gguf.experts import load_gguf_expert_banks
 
             # each layer keeps its own ggml types, so the banks come from the file, not from one method's layout
-            return load_gguf_expert_banks(model_path, model_config.num_moe_layers, dummy=dummy, device=device if resident else None)
+            return load_gguf_expert_banks(model_path, model_config.num_moe_layers, first_layer=model_config.first_k_dense_replace, dummy=dummy, device=device if resident else None)
         if method is None:
             raise ValueError("the experts of a raw checkpoint are packed by the offload layers' MoE quant method; pass method=")
         return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident)

From 03c2c9c35eae76601fae48efd0eb352cdab1f170 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 22:03:59 -0700
Subject: [PATCH 116/249] perf(kernel): rewrite the triton nvfp4 moe prefill
 gemm (#723)

* perf(kernel): rewrite the triton nvfp4 moe prefill gemm

* perf(moe): tune nvfp4 prefill tiles per gpu class

GeForce runs fp32-accumulate tensor MMA at half rate, so RTX parts want smaller BLOCK_SIZE_M than H100 and their own BLOCK_N / BLOCK_KB / stages.
Tables come from tile sweeps on RTX PRO 6000, RTX 5090, RTX 4090 and RTX 3080 Ti; H100-class GPUs and AMD keep the previous tuning.
---
 .../kernel/triton/fp8_blockscale_moe.py       |  3 +-
 .../kernel/triton/nvfp4_fused_moe.py          | 75 +++++++++--------
 python/freetoken/moe/fused_nvfp4.py           | 80 ++++++++++++++++---
 3 files changed, 110 insertions(+), 48 deletions(-)

diff --git a/python/freetoken/kernel/triton/fp8_blockscale_moe.py b/python/freetoken/kernel/triton/fp8_blockscale_moe.py
index 1db54a8072..429134d23f 100644
--- a/python/freetoken/kernel/triton/fp8_blockscale_moe.py
+++ b/python/freetoken/kernel/triton/fp8_blockscale_moe.py
@@ -128,8 +128,7 @@ def fused_experts_decode_fp8_blockscale(
 
 
 # ======================================================================================
-# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size), mirrors
-# nvfp4_fused_moe._prefill_nvfp4_moe_kernel with fp8*bf16-block-scale dequant.
+# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size) with fp8*bf16-block-scale dequant.
 # BLOCK_SIZE_N == BLOCK_SIZE_K == 128 so each tile maps to one weight scale block.
 # ======================================================================================
 @triton.jit
diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
index b1de443787..ca9b9dbe9d 100644
--- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py
+++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
@@ -292,18 +292,24 @@ def _decode_nvfp4_marlin_kernel(
     tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask)
 
 
+@triton.jit
+def _e2m1_f16_x2p14(code):
+    """e2m1 code (low 4 bits) -> fp16 value * 2^-14: magnitude bits go to fp16 [11:9], the sign to [15]; fp16 subnormals line up the e == 0 codes."""
+    bits = ((code & 7) << 9) | ((code & 8) << 12)
+    return bits.to(tl.uint16).to(tl.float16, bitcast=True)
+
+
 @triton.jit
 def _prefill_nvfp4_moe_kernel(
     a_ptr,             # [M, K] activations
     packed_ptr,        # [S, N, K // 2] uint8
-    scale_ptr,         # [S, N, K // 16] fp8-e4m3
+    scale_ptr,         # [S, N, K // 16] fp8-e4m3 (uint8 view before sm_89)
     global_ptr,        # [S, N] fp16
     c_ptr,             # [num_valid_tokens, N] output (flat over M*top_k)
     topk_weights_ptr,
     sorted_token_ids_ptr,
     expert_ids_ptr,    # cache slot per M-block
     num_tokens_post_padded_ptr,
-    lut_ptr,
     N,
     K,
     EM,
@@ -314,14 +320,16 @@ def _prefill_nvfp4_moe_kernel(
     stride_ge, stride_gn,
     stride_cm, stride_cn,
     stride_tw,
-    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,   # tokens per block, the moe_align padding
     BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_KB: tl.constexpr,
+    BLOCK_SIZE_KB: tl.constexpr,  # packed bytes per K step (2 values each)
     GROUP_SIZE_M: tl.constexpr,
     MUL_ROUTED_WEIGHT: tl.constexpr,
     top_k: tl.constexpr,
     compute_type: tl.constexpr,
 ):
+    # out^T = W @ X^T: wgmma reads its A operand (the decoded weights) from registers, while B must be in shared memory.
+    # The token block is the MMA N dimension, so a 16-token block still fills the MMA.
     pid = tl.program_id(axis=0)
     num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
@@ -342,51 +350,50 @@ def _prefill_nvfp4_moe_kernel(
 
     offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
     offs_kb = tl.arange(0, BLOCK_SIZE_KB)
-    a_ptrs_lo = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb)[None, :] * stride_ak)
-    a_ptrs_hi = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb + 1)[None, :] * stride_ak)
+    offs_k = tl.arange(0, 2 * BLOCK_SIZE_KB)
+    offs_sb = tl.arange(0, BLOCK_SIZE_KB // 8)
+    xt_ptrs = a_ptr + (offs_token[None, :] // top_k * stride_am + offs_k[:, None] * stride_ak)
 
     slot = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
-    packed_base = packed_ptr + slot * stride_pe + offs_bn[None, :] * stride_pn
-    scale_base = scale_ptr + slot * stride_se + offs_bn[None, :] * stride_sn
+    p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb
+    s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk
 
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
     K_BYTES = K // 2
     for kb in range(0, tl.cdiv(K_BYTES, BLOCK_SIZE_KB)):
-        byte_idx = kb * BLOCK_SIZE_KB + offs_kb
-        byte_mask = byte_idx < K_BYTES
-
-        p_ptrs = packed_base + byte_idx[:, None] * stride_pkb
-        bytes_ = tl.load(p_ptrs, mask=byte_mask[:, None], other=0).to(tl.int32)
-        lo = bytes_ & 0xF
-        hi = (bytes_ >> 4) & 0xF
-        sblk = byte_idx // 8
-        s_ptrs = scale_base + sblk[:, None] * stride_sblk
+        byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES
+        sblk_mask = kb * (BLOCK_SIZE_KB // 8) + offs_sb < K_BYTES // 8
+        k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K
+        v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32)
         if e4m3_native_cx():
-            scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0).to(tl.float32)
+            s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32)
         else:
-            scale = e4m3_u8_to_f32(tl.load(s_ptrs, mask=byte_mask[:, None], other=0))
-        b_lo = tl.load(lut_ptr + lo) * scale  # [BLOCK_KB, BLOCK_N]
-        b_hi = tl.load(lut_ptr + hi) * scale
-
-        a_lo = tl.load(a_ptrs_lo, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)
-        a_hi = tl.load(a_ptrs_hi, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)
-        accumulator += tl.dot(a_lo, b_lo.to(a_lo.dtype))
-        accumulator += tl.dot(a_hi, b_hi.to(a_hi.dtype))
-
-        a_ptrs_lo += BLOCK_SIZE_KB * 2 * stride_ak
-        a_ptrs_hi += BLOCK_SIZE_KB * 2 * stride_ak
+            s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0))
+        # one scale per 16 values; the 2^14 undoes the decode's 2^-14
+        s = tl.reshape(tl.broadcast_to((s * 16384.0)[:, :, None], (BLOCK_SIZE_N, BLOCK_SIZE_KB // 8, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB))
+        s = s.to(tl.bfloat16)
+        # exact in bf16: e2m1 has <= 2 significant bits and e4m3 <= 4
+        w_lo = _e2m1_f16_x2p14(v & 0xF).to(tl.bfloat16) * s
+        w_hi = _e2m1_f16_x2p14(v >> 4).to(tl.bfloat16) * s
+        # interleave back to k = 2 * byte + {0, 1} so one dot covers the K step
+        w = tl.reshape(tl.join(w_lo, w_hi), (BLOCK_SIZE_N, 2 * BLOCK_SIZE_KB))
+        xt = tl.load(xt_ptrs, mask=k_mask[:, None] & token_mask[None, :], other=0.0)
+        accumulator = tl.dot(w.to(xt.dtype), xt, accumulator)
+        xt_ptrs += 2 * BLOCK_SIZE_KB * stride_ak
+        p_ptrs += BLOCK_SIZE_KB * stride_pkb
+        s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk
 
     g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32)
-    accumulator = accumulator * g[None, :]
+    accumulator = accumulator * g[:, None]
 
     if MUL_ROUTED_WEIGHT:
         moe_weight = tl.load(topk_weights_ptr + offs_token * stride_tw, mask=token_mask, other=0)
-        accumulator = accumulator * moe_weight[:, None]
+        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)
+    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)
 
 
diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py
index 375d2a547a..9202f5f38a 100644
--- a/python/freetoken/moe/fused_nvfp4.py
+++ b/python/freetoken/moe/fused_nvfp4.py
@@ -7,13 +7,16 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict
+import functools
+import re
+from typing import Any, Dict, Tuple
 
 import torch
 import triton
 import triton.language as tl
 
 from freetoken.kernel import moe_sum_reduce_triton
+from freetoken.kernel.backend import is_rocm
 from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view
 from freetoken.kernel.triton.nvfp4_fused_moe import (
     _decode_nvfp4_marlin_kernel,
@@ -23,6 +26,9 @@
 )
 from freetoken.layers import gated_act_and_mul
 from freetoken.moe.fused import moe_align_block_size
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
 
 # Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune,
 # which benchmarks at run time). Tuned offline against the NVFP4 decode kernels.
@@ -234,15 +240,66 @@ def fused_experts_decode_nvfp4_serial(
     )
 
 
-def _prefill_config(M: int) -> Dict[str, int]:
-    # ``BLOCK_SIZE_M`` is coupled to host-side ``moe_align_block_size`` (token padding),
-    # so it cannot be picked by triton.autotune; these were chosen by an offline sweep
-    # over (BLOCK_M, BLOCK_N, BLOCK_KB, num_warps, num_stages) for the MiniMax-M2 shapes.
-    if M <= 64:
-        return dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32,
-                    GROUP_SIZE_M=1, num_warps=8, num_stages=4)
-    return dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32,
-                GROUP_SIZE_M=8, num_warps=8, num_stages=4)
+# rho = bf16 tensor-core FLOP / fp32 ALU FLOP per SM per clock, with fp32 accumulate (the only accumulator for bf16 MMA).
+# Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M.
+# High rho makes decode the cost, so big blocks win; low rho makes padded MMA the cost, so small blocks win.
+# Spec, FLOP/clk/SM: GeForce RTX 30/40/50 512 / 256 = 2, as GeForce runs fp32 accumulate at half rate.
+# RTX A, RTX Ada Generation, RTX PRO Blackwell, A10, A40, L4, L40S 1024 / 256 = 4; A100 2048 / 128 = 16; H100 4096 / 256 = 16; H20 148 / 44 TFLOP/s = 3.4.
+# Measured (cuBLAS bf16 / Triton fp32 FMA): RTX PRO 6000 4.5, RTX 5090 2.3, RTX 4090 1.9, RTX 3080 Ti 1.6.
+_PRO_GPU_RE = re.compile(
+    r"RTX\s*A\d{3,4}"                        # RTX A400 .. A6000
+    r"|RTX\s*\d{3,4}E?\s*(SFF\s*)?-?\s*Ada"  # RTX 500 .. 6000 Ada Generation
+    r"|RTX\s*PRO\s*\d{3,4}"                  # RTX PRO 500 .. 6000 Blackwell
+    r"|RTX\s*\d{4}D\b"                       # RTX 6000D
+    r"|Quadro|TITAN"
+    r"|^(NVIDIA\s+)?(GRID\s+)?(A2|A10[GM]?|A16|A40|L2|L20|L4|L40S?)\b",
+    re.I,
+)
+# also matches names without "GeForce", e.g. "NVIDIA RTX 4090"
+_GEFORCE_GPU_RE = re.compile(r"GeForce|RTX\s*[2-6]0[5-9]0(?!\d)", re.I)
+
+
+def _tensor_rho(name: str, major: int, minor: int) -> float:
+    if major not in (8, 12) or (major, minor) == (8, 0):
+        return 3.4 if re.search(r"\bH20\b", name) else 16.0
+    if _PRO_GPU_RE.search(name):
+        return 4.0
+    if _GEFORCE_GPU_RE.search(name):
+        return 2.0
+    logger.warning("unknown GPU %r: NVFP4 MoE prefill assumes full-rate fp32-accumulate tensor cores", name)
+    return 4.0
+
+
+# (min rho, tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128), from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti
+_PREFILL_BLOCK_M_STEPS = ((8.0, (12, 24, 64)), (3.0, (16, 36, 150)), (0.0, (16, 150, float("inf"))))
+# BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages)
+_PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)}
+_PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)}
+_PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)}
+
+
+@functools.lru_cache(maxsize=None)
+def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]:
+    if is_rocm():
+        # TODO: tune for AMD GPUs; they keep the H100 tuning
+        return _PREFILL_BLOCK_M_STEPS[0][1], _PREFILL_TILES_H100
+    props = torch.cuda.get_device_properties(device_index)
+    rho = _tensor_rho(props.name, props.major, props.minor)
+    steps = next(s for min_rho, s in _PREFILL_BLOCK_M_STEPS if rho >= min_rho)
+    # the RTX tables come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 one
+    if (props.major, props.minor) not in ((8, 6), (8, 9), (12, 0)):
+        return steps, _PREFILL_TILES_H100
+    return steps, _PREFILL_TILES_RTX_PRO if rho >= 3.0 else _PREFILL_TILES_GEFORCE
+
+
+def _prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]:
+    # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it.
+    (to32, to64, to128), tiles = _prefill_tuning(device_index)
+    per_expert = M * top_k / num_experts
+    block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16
+    block_n, block_kb, stages = tiles[block_m]
+    return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb,
+                GROUP_SIZE_M=8, num_warps=4, num_stages=stages)
 
 
 def _prefill_gemm(
@@ -270,7 +327,6 @@ def _prefill_gemm(
     _prefill_nvfp4_moe_kernel[grid](
         a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids,
         num_tokens_post_padded,
-        _e2m1_lut(a.device.index),
         N, K, EM, num_valid_tokens,
         a.stride(0), a.stride(1),
         packed.stride(0), packed.stride(1), packed.stride(2),
@@ -309,7 +365,7 @@ def fused_experts_nvfp4(
     two_i = gate_up_packed.shape[1]
     inter = two_i // 2
     dev, dt = hidden_states.device, hidden_states.dtype
-    cfg = _prefill_config(M)
+    cfg = _prefill_config(M, top_k, num_experts, hidden_states.device.index)
 
     sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
     tw = topk_weights.reshape(-1).contiguous()

From c8559fdff68c517c5cab60d491a9870099b73d66 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 22:28:34 -0700
Subject: [PATCH 117/249] feat(cpu-moe): serve GGUF experts on the cpu and
 hybrid backends through ggml-cpu (#704)

* feat(cpu-moe): serve GGUF experts on the cpu and hybrid backends through ggml-cpu

* feat(cpu-moe): run float GGUF expert rows on the cpu through ggml-cpu's vec_dot

* feat(cpu-moe): build ggml-cpu for each llama.cpp x86 variant and run the best one the CPU supports
---
 docs/cli.md                                   |    2 +-
 python/freetoken/engine/engine.py             |    4 +
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       |    9 +-
 .../kernel/csrc/cpu_moe/kernels/gguf.cpp      |  107 +
 .../kernel/csrc/cpu_moe/kernels/q4_0.cpp      |  216 -
 python/freetoken/kernel/csrc/ggml/VERSION     |    5 +-
 .../kernel/csrc/ggml/freetoken/cpu_kernels.c  |  172 +
 .../kernel/csrc/ggml/freetoken/cpu_kernels.h  |   43 +
 .../csrc/ggml/freetoken/variant_module.cpp    |   32 +
 .../kernel/csrc/ggml/ggml-cpu/arch-fallback.h |  375 ++
 .../csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp |  327 +
 .../csrc/ggml/ggml-cpu/arch/x86/quants.c      | 4172 ++++++++++++
 .../kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h |  538 ++
 .../kernel/csrc/ggml/ggml-cpu/quants.c        | 1339 ++++
 .../kernel/csrc/ggml/ggml-cpu/quants.h        |  106 +
 .../kernel/csrc/ggml/ggml-cpu/simd-mappings.h | 1366 ++++
 .../kernel/csrc/ggml/ggml-cpu/vec.cpp         |  613 ++
 .../freetoken/kernel/csrc/ggml/ggml-cpu/vec.h | 1630 +++++
 .../kernel/csrc/ggml/include/ggml-cpu.h       |  152 +
 .../kernel/csrc/ggml/src/ggml-backend-impl.h  |  292 +
 .../kernel/csrc/ggml/src/ggml-quants.c        | 5639 +++++++++++++++++
 .../kernel/csrc/ggml/src/ggml-quants.h        |  115 +
 python/freetoken/kernel/ggml_cpu/__init__.py  |   37 +
 .../freetoken/layers/quantization/moe/gguf.py |   14 +-
 python/freetoken/moe/benchbw.py               |    2 +-
 python/freetoken/moe/cpu_executor.py          |   51 +-
 setup.py                                      |  155 +-
 tests/README.md                               |    2 +-
 tests/moe/test_cpu_moe_q4_0.py                |  213 -
 29 files changed, 17260 insertions(+), 468 deletions(-)
 create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
 delete mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp
 create mode 100644 python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h
 create mode 100644 python/freetoken/kernel/csrc/ggml/src/ggml-quants.c
 create mode 100644 python/freetoken/kernel/csrc/ggml/src/ggml-quants.h
 create mode 100644 python/freetoken/kernel/ggml_cpu/__init__.py
 delete mode 100644 tests/moe/test_cpu_moe_q4_0.py

diff --git a/docs/cli.md b/docs/cli.md
index 4918a3177a..69848f2e44 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -89,7 +89,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:<path>`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k |
 | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts |
 | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor |
-| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) |
+| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4, GGUF on x86 CPUs) |
 | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU |
 | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) |
 | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap |
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 980bb043bd..326c467d6d 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1479,6 +1479,10 @@ def _cpu_moe_executor_viable(model_config) -> bool:
         return False
     expert_quant = getattr(model_config, "expert_quant", "none")
     fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
+    if fmt == "gguf":
+        from freetoken.kernel import ggml_cpu
+
+        return fmt in _cpu_moe.weight_formats() and ggml_cpu.select() is not None
     return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats()
 
 
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index 6236816eec..c4887efbff 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -46,8 +46,9 @@ std::unique_ptr<MoeFormat> make_bf16_format(const MoeParams& p, const FormatArgs
 std::unique_ptr<MoeFormat> make_nvfp4_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_mxfp4_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_ds_fp4_format(const MoeParams& p, const FormatArgs& a);
-std::unique_ptr<MoeFormat> make_q4_0_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_fp8_block_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a);
+void set_ggml_cpu(uintptr_t addr);
 
 namespace {
 
@@ -94,8 +95,8 @@ inline MoeTask from_word(uint64_t w) {
 
 // Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names.
 const std::pair<const char*, std::unique_ptr<MoeFormat> (*)(const MoeParams&, const FormatArgs&)> kFormats[] = {
-    {"bf16", make_bf16_format},     {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format},
-    {"ds_fp4", make_ds_fp4_format}, {"q4_0", make_q4_0_format},   {"fp8_block", make_fp8_block_format},
+    {"bf16", make_bf16_format},     {"nvfp4", make_nvfp4_format},         {"mxfp4_triton", make_mxfp4_format},
+    {"ds_fp4", make_ds_fp4_format}, {"fp8_block", make_fp8_block_format}, {"gguf", make_gguf_format},
 };
 
 std::unique_ptr<MoeFormat> make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) {
@@ -573,6 +574,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   // (act_apply falls through to gelu_tanh); the probe turns a stale extension
   // into a loud rebuild instruction instead of wrong model outputs.
   m.def("max_generic_act_id", []() { return static_cast<int>(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); });
+  // The ggml-cpu variant the gguf format calls: the kernel table address that freetoken.kernel.ggml_cpu.select() returns.
+  m.def("set_ggml_cpu", &freetoken::cpu_moe::set_ggml_cpu, py::arg("api"));
   m.def("weight_formats", []() {
     py::list names;
     for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first);
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
new file mode 100644
index 0000000000..c1641af51d
--- /dev/null
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
@@ -0,0 +1,107 @@
+#include <atomic>
+#include <mutex>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "../format.h"
+#include "cpu_kernels.h"
+
+namespace freetoken::cpu_moe {
+namespace {
+
+std::atomic<const ft_ggml_cpu_api*> g_ggml{nullptr};
+
+const ft_ggml_cpu_api& ggml_api() {
+  const ft_ggml_cpu_api* api = g_ggml.load();
+  if (api == nullptr) throw std::logic_error("CPU MoE gguf experts need a ggml-cpu variant; call _cpu_moe.set_ggml_cpu first");
+  return *api;
+}
+
+}  // namespace
+
+void set_ggml_cpu(uintptr_t addr) {
+  const auto* api = reinterpret_cast<const ft_ggml_cpu_api*>(addr);
+  if (api == nullptr || api->abi != FT_GGML_CPU_ABI) throw std::invalid_argument("set_ggml_cpu: not a ggml-cpu kernel table of this build");
+  const ft_ggml_cpu_api* prev = nullptr;
+  if (!g_ggml.compare_exchange_strong(prev, api) && prev != api)
+    throw std::invalid_argument(std::string("set_ggml_cpu: this process already uses the ") + prev->name + " ggml-cpu variant");
+}
+
+namespace {
+
+// Packed ggml rows of the layers whose experts have this gate_up type and this down type, dotted by ggml-cpu against the activation converted to each type's vec_dot input (q8_0, q8_1 or q8_K, or the float type itself).
+class GgufFormat final : public MoeFormat {
+ public:
+  GgufFormat(const MoeParams& p, const FormatArgs& a)
+      : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {
+    static std::once_flag tables;
+    std::call_once(tables, ggml_.init);
+    gu_ = kernel(static_cast<int>(a.option("gate_up_type")), p.H, "gate_up");
+    dn_ = kernel(static_cast<int>(a.option("down_type")), p.I, "down");
+    xf_.resize(p.H);
+    xq_.resize(static_cast<size_t>(p.max_tokens) * gu_.act_row_bytes);
+    gq_.resize(static_cast<size_t>(p.max_tokens) * p.top_k * dn_.act_row_bytes);
+  }
+
+  std::string isa_name() const override { return std::string(ggml_.name) + "(ggml-cpu)"; }
+
+  void prep_input(int tok, const bf16_t* x) override {
+    for (int i = 0; i < p.H; ++i) xf_[i] = bf16_to_f32(x[i]);
+    gu_.from_float(xf_.data(), xq_.data() + static_cast<size_t>(tok) * gu_.act_row_bytes, p.H);
+  }
+
+  bool preps_inter() const override { return true; }
+
+  void prep_inter(int64_t r, bf16_t* g) override {
+    thread_local std::vector<float> gf;
+    gf.resize(p.I);
+    for (int i = 0; i < p.I; ++i) gf[i] = bf16_to_f32(g[i]);
+    dn_.from_float(gf.data(), gq_.data() + static_cast<size_t>(r) * dn_.act_row_bytes, p.I);
+  }
+
+  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, bf16_t* g) const override {
+    const uint8_t* w = static_cast<const uint8_t*>(tbl_at(gate_up_, layer)) + static_cast<size_t>(e) * 2 * p.I * gu_.row_bytes;
+    const uint8_t* xq = xq_.data() + static_cast<size_t>(tok) * gu_.act_row_bytes;
+    gated_rows(p, i0, i1, w_in, g, [&](int row) {
+      float s;
+      gu_.vec_dot(p.H, &s, 0, w + static_cast<size_t>(row) * gu_.row_bytes, 0, xq, 0, 1);
+      return s;
+    });
+  }
+
+  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
+    const uint8_t* w = static_cast<const uint8_t*>(tbl_at(down_, layer));
+    for (int h = h0; h < h1; ++h)
+      for (int j = 0; j < n; ++j)
+        dn_.vec_dot(p.I, &out[j * kRowBlock + h - h0], 0, w + (static_cast<size_t>(routes[j].e) * p.H + h) * dn_.row_bytes, 0,
+                    gq_.data() + static_cast<size_t>(routes[j].r) * dn_.act_row_bytes, 0, 1);
+  }
+
+ private:
+  ft_ggml_row_kernel kernel(int type, int k, const char* role) const {
+    ft_ggml_row_kernel out{};
+    const int rc = ggml_.row_kernel(type, k, &out);
+    const std::string name = ggml_.type_name(type);
+    if (rc == FT_GGML_NO_KERNEL) throw std::invalid_argument("ggml-cpu has no dot product for " + name + " experts");
+    if (rc != FT_GGML_OK)
+      throw std::invalid_argument("CPU MoE " + name + " " + role + " rows of " + std::to_string(k) + " values are not whole blocks");
+    return out;
+  }
+
+  const ft_ggml_cpu_api& ggml_;
+  const uint64_t* gate_up_;
+  const uint64_t* down_;
+  ft_ggml_row_kernel gu_{}, dn_{};
+  std::vector<float> xf_;
+  // the activation in the vec_dot input type: x per token, g per route
+  AlignedVector<uint8_t> xq_, gq_;
+};
+
+}  // namespace
+
+std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<GgufFormat>(p, a);
+}
+
+}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
deleted file mode 100644
index 600a194ebc..0000000000
--- a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp
+++ /dev/null
@@ -1,216 +0,0 @@
-#include <stdexcept>
-#include <vector>
-
-#include "../format.h"
-
-namespace freetoken::cpu_moe {
-namespace {
-
-// -------------------------------- Q4_0 (W4A8) --------------------------------
-// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed
-// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a
-// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches
-// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload
-// path streams.
-//
-// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0
-// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to
-// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high
-// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD
-// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then
-// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU
-// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close.
-using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int);
-
-float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  float acc = 0.0f;
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    const uint8_t* q = blk + 2;  // 16 nibble bytes
-    const int8_t* a = aq + (size_t)b * 32;
-    int isum = 0;
-    for (int j = 0; j < 16; ++j) {
-      isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j];       // elem j
-      isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j];    // elem 16+j
-    }
-    acc += fp16_to_f32(dh) * asb[b] * (float)isum;
-  }
-  return acc;
-}
-
-#if CPU_MOE_X86
-// fp16 block scale -> fp32 via HW F16C (single value in lane 0).
-CPU_MOE_TARGET("f16c")
-static inline float q4_scale(uint16_t h) {
-  return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h)));
-}
-
-// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low
-// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8).
-CPU_MOE_TARGET("avx2")
-static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) {
-  const __m128i qb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(blk + 2));
-  const __m128i lo = _mm_and_si128(qb, mask);
-  const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask);
-  return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight);
-}
-
-// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW +
-// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no
-// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the
-// avx512-tier fallback, since the block dot is 256-bit either way).
-CPU_MOE_TARGET("avx2,fma,f16c")
-float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  const __m128i mask = _mm_set1_epi8(0x0F);
-  const __m256i eight = _mm256_set1_epi8(8);
-  const __m256i ones16 = _mm256_set1_epi16(1);
-  __m256 accF = _mm256_setzero_ps();
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    __m256i wq = q4_unpack32(blk, mask, eight);
-    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
-    __m256i aw = _mm256_sign_epi8(wq, wq);              // |wq|          (unsigned operand)
-    __m256i sa = _mm256_sign_epi8(a, wq);               // sign(wq) * a  (signed operand)
-    __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16);  // 8 int32
-    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
-  }
-  return hsum256(accF);
-}
-
-// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs).
-CPU_MOE_TARGET("avx2,avxvnni,fma,f16c")
-float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) {
-  const __m128i mask = _mm_set1_epi8(0x0F);
-  const __m256i eight = _mm256_set1_epi8(8);
-  __m256 accF = _mm256_setzero_ps();
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const uint8_t* blk = w + (size_t)b * 18;
-    _mm_prefetch(reinterpret_cast<const char*>(blk) + 512, _MM_HINT_T0);
-    uint16_t dh;
-    std::memcpy(&dh, blk, sizeof(dh));
-    __m256i wq = q4_unpack32(blk, mask, eight);
-    __m256i a = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(aq + (size_t)b * 32));
-    __m256i aw = _mm256_sign_epi8(wq, wq);   // |wq|            (unsigned operand)
-    __m256i sa = _mm256_sign_epi8(a, wq);    // sign(wq) * a    (signed operand)
-    __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa);
-    // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and
-    // reduce at the end (equivalent to scale * block_total).
-    accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF);
-  }
-  return hsum256(accF);
-}
-#endif  // CPU_MOE_X86
-
-// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to
-// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when
-// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers.
-q4dot_fn select_q4dot() {
-  const IsaTier t = pick_isa();
-#if CPU_MOE_X86
-  if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni;
-  if (t >= ISA_AVX2) return q4_0_dot_i8_avx2;
-#endif
-  (void)t;
-  return q4_0_dot_i8_scalar;
-}
-
-// Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in
-// natural order + one fp32 scale (amax/127) per block. Done once per token/route,
-// amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0.
-void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) {
-  const int nb = K / 32;
-  for (int b = 0; b < nb; ++b) {
-    const bf16_t* xb = x + (size_t)b * 32;
-    float xf[32], amax = 0.0f;
-    for (int j = 0; j < 32; ++j) {
-      xf[j] = bf16_to_f32(xb[j]);
-      amax = std::max(amax, std::fabs(xf[j]));
-    }
-    const float d = amax > 0.0f ? amax / 127.0f : 1.0f;
-    asb[b] = d;
-    const float inv = amax > 0.0f ? 1.0f / d : 0.0f;
-    int8_t* o = aq + (size_t)b * 32;
-    for (int j = 0; j < 32; ++j)
-      o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv)));
-  }
-}
-
-class Q4_0Format final : public MoeFormat {
- public:
-  Q4_0Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_q4dot()) {
-    gate_up_ = a.bank("gate_up");
-    down_ = a.bank("down");
-    const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k;
-    xi8_.resize(tokens * p.H);
-    xas_.resize(tokens * (p.H / 32));
-    gi8_.resize(routes * p.I);
-    gas_.resize(routes * (p.I / 32));
-    if (p.H % 32 != 0 || p.I % 32 != 0)
-      throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32");
-    gu_row_bytes_ = (p.H / 32) * 18;  // K = H (gate_up rows)
-    dn_row_bytes_ = (p.I / 32) * 18;  // K = I (down rows)
-    // select_q4dot picks VPDPBUSD / VPMADDUBSW / scalar for the tier, so the tag reflects which of those it resolved to.
-    isa_ = std::string(tier_name(pick_isa())) + (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8");
-  }
-
-  std::string isa_name() const override { return isa_; }
-
-
-  void prep_input(int tok, const bf16_t* x) override {
-    quant_q8_0(x, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 32));
-  }
-
-  bool preps_inter() const override { return true; }
-
-  void prep_inter(int64_t r, bf16_t* g) override {
-    quant_q8_0(g, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 32));
-  }
-
-  void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in,
-               bf16_t* g) const override {
-    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
-    const int8_t* xi8 = xi8_.data() + (size_t)tok * p.H;
-    const float* xas = xas_.data() + (size_t)tok * (p.H / 32);
-    gated_rows(p, i0, i1, w_in, g, [&](int row) {
-      return dot_(packed + ((size_t)e * (2 * p.I) + row) * (size_t)gu_row_bytes_, xi8, xas, p.H);
-    });
-  }
-
-  void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override {
-    const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(down_, layer));
-    for (int h = h0; h < h1; ++h) {
-      for (int j = 0; j < n; ++j) {
-        const size_t r = (size_t)routes[j].r;
-        out[j * kRowBlock + h - h0] = dot_(packed + ((size_t)routes[j].e * p.H + h) * (size_t)dn_row_bytes_,
-                                          gi8_.data() + r * p.I, gas_.data() + r * (p.I / 32), p.I);
-      }
-    }
-  }
-
- private:
-  q4dot_fn dot_;
-  // GGUF Q4_0 rows of 32-element blocks (fp16 scale + 16 nibble bytes): gate_up [E,2I,H/32*18], down [E,H,I/32*18].
-  const uint64_t* gate_up_ = nullptr;
-  const uint64_t* down_ = nullptr;
-  int gu_row_bytes_ = 0, dn_row_bytes_ = 0;
-  std::string isa_;
-  // Q8_0 activations: per-32-block int8 + fp32 scale, for the input and the intermediate.
-  AlignedVector<int8_t> xi8_, gi8_;
-  AlignedVector<float> xas_, gas_;
-};
-
-}  // namespace
-
-std::unique_ptr<MoeFormat> make_q4_0_format(const MoeParams& p, const FormatArgs& a) {
-  return std::make_unique<Q4_0Format>(p, a);
-}
-
-}  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION
index 6eabb0c19f..3c66f9b7d0 100644
--- a/python/freetoken/kernel/csrc/ggml/VERSION
+++ b/python/freetoken/kernel/csrc/ggml/VERSION
@@ -1,7 +1,8 @@
 llama.cpp 8c2134a8755d8a9792af513dab20282affe903e4 (2026-10-14)
 
-Only the ggml files that the GGUF ops (freetoken/ggml_ops.cu, built by kernel/ggml.py) compile or include:
-ggml/include -> include, ggml/src -> src, ggml/src/ggml-cuda -> ggml-cuda.
+Only the ggml files that FreeToken compiles or includes:
+- the GGUF ops (freetoken/ggml_ops.cu, built by kernel/ggml.py): ggml/include -> include, ggml/src -> src, ggml/src/ggml-cuda -> ggml-cuda.
+- the CPU MoE row kernels (freetoken/cpu_kernels.c and freetoken/variant_module.cpp, one freetoken.kernel.ggml_cpu module per x86 variant, built by setup.py): ggml/src/{ggml-quants.c,ggml-quants.h,ggml-backend-impl.h} -> src, ggml/include/ggml-cpu.h -> include, ggml/src/ggml-cpu/{quants.c,quants.h,vec.cpp,vec.h,ggml-cpu-impl.h,simd-mappings.h,arch-fallback.h,arch/x86/quants.c,arch/x86/cpu-feats.cpp} -> ggml-cpu. These have no local changes.
 
 Local changes, kept in freetoken.patch (already applied):
 - quantize.cu/.cuh: the q8_1 quantize kernels read fp32, bf16 or fp16 activations.
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c
new file mode 100644
index 0000000000..79a036383a
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c
@@ -0,0 +1,172 @@
+// The ggml-base symbols and lookup tables that the vendored ggml-cpu quant kernels reference, and the row kernel of each weight type.
+
+#include "cpu_kernels.h"
+
+#include "ggml-impl.h"
+#include "ggml-cpu/ggml-cpu-impl.h"
+#include "ggml-cpu/quants.h"
+#include "ggml-cpu/simd-mappings.h"
+#include "ggml-cpu/vec.h"
+
+#include <stdarg.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+float ggml_table_f32_f16[1 << 16];
+float ggml_table_f32_e8m0_half[1 << 8];
+float ggml_table_f32_ue4m3[1 << 8];
+
+void ggml_abort(const char * file, int line, const char * fmt, ...) {
+    va_list args;
+    va_start(args, fmt);
+    fprintf(stderr, "ggml: %s:%d: ", file, line);
+    vfprintf(stderr, fmt, args);
+    va_end(args);
+    fputc('\n', stderr);
+    fflush(stderr);
+    abort();
+}
+
+struct type_info {
+    const char * name;
+    int64_t blck_size;
+    size_t type_size;
+};
+
+static const struct type_info k_types[GGML_TYPE_COUNT] = {
+    [GGML_TYPE_F32] = {"f32", 1, sizeof(float)},
+    [GGML_TYPE_F16] = {"f16", 1, sizeof(ggml_fp16_t)},
+    [GGML_TYPE_BF16] = {"bf16", 1, sizeof(ggml_bf16_t)},
+    [GGML_TYPE_Q1_0] = {"q1_0", QK1_0, sizeof(block_q1_0)},
+    [GGML_TYPE_Q2_0] = {"q2_0", QK2_0, sizeof(block_q2_0)},
+    [GGML_TYPE_Q4_0] = {"q4_0", QK4_0, sizeof(block_q4_0)},
+    [GGML_TYPE_Q4_1] = {"q4_1", QK4_1, sizeof(block_q4_1)},
+    [GGML_TYPE_Q5_0] = {"q5_0", QK5_0, sizeof(block_q5_0)},
+    [GGML_TYPE_Q5_1] = {"q5_1", QK5_1, sizeof(block_q5_1)},
+    [GGML_TYPE_Q8_0] = {"q8_0", QK8_0, sizeof(block_q8_0)},
+    [GGML_TYPE_Q8_1] = {"q8_1", QK8_1, sizeof(block_q8_1)},
+    [GGML_TYPE_MXFP4] = {"mxfp4", QK_MXFP4, sizeof(block_mxfp4)},
+    [GGML_TYPE_NVFP4] = {"nvfp4", QK_NVFP4, sizeof(block_nvfp4)},
+    [GGML_TYPE_Q2_K] = {"q2_K", QK_K, sizeof(block_q2_K)},
+    [GGML_TYPE_Q3_K] = {"q3_K", QK_K, sizeof(block_q3_K)},
+    [GGML_TYPE_Q4_K] = {"q4_K", QK_K, sizeof(block_q4_K)},
+    [GGML_TYPE_Q5_K] = {"q5_K", QK_K, sizeof(block_q5_K)},
+    [GGML_TYPE_Q6_K] = {"q6_K", QK_K, sizeof(block_q6_K)},
+    [GGML_TYPE_Q8_K] = {"q8_K", QK_K, sizeof(block_q8_K)},
+    [GGML_TYPE_TQ1_0] = {"tq1_0", QK_K, sizeof(block_tq1_0)},
+    [GGML_TYPE_TQ2_0] = {"tq2_0", QK_K, sizeof(block_tq2_0)},
+    [GGML_TYPE_IQ2_XXS] = {"iq2_xxs", QK_K, sizeof(block_iq2_xxs)},
+    [GGML_TYPE_IQ2_XS] = {"iq2_xs", QK_K, sizeof(block_iq2_xs)},
+    [GGML_TYPE_IQ2_S] = {"iq2_s", QK_K, sizeof(block_iq2_s)},
+    [GGML_TYPE_IQ3_XXS] = {"iq3_xxs", QK_K, sizeof(block_iq3_xxs)},
+    [GGML_TYPE_IQ3_S] = {"iq3_s", QK_K, sizeof(block_iq3_s)},
+    [GGML_TYPE_IQ1_S] = {"iq1_s", QK_K, sizeof(block_iq1_s)},
+    [GGML_TYPE_IQ1_M] = {"iq1_m", QK_K, sizeof(block_iq1_m)},
+    [GGML_TYPE_IQ4_NL] = {"iq4_nl", QK4_NL, sizeof(block_iq4_nl)},
+    [GGML_TYPE_IQ4_XS] = {"iq4_xs", QK_K, sizeof(block_iq4_xs)},
+};
+
+static int known_type(int type) { return type >= 0 && type < GGML_TYPE_COUNT && k_types[type].name != NULL; }
+
+int64_t ggml_blck_size(enum ggml_type type) { return k_types[type].blck_size; }
+
+size_t ggml_type_size(enum ggml_type type) { return k_types[type].type_size; }
+
+size_t ggml_row_size(enum ggml_type type, int64_t ne) { return k_types[type].type_size * ne / k_types[type].blck_size; }
+
+const char * ggml_type_name(enum ggml_type type) { return known_type(type) ? k_types[type].name : "unknown"; }
+
+const char * ft_ggml_type_name(int type) { return ggml_type_name((enum ggml_type) type); }
+
+struct dot_info {
+    ft_ggml_vec_dot_t vec_dot;
+    int act_type;
+};
+
+// The weight types of type_traits_cpu in ggml-cpu.c, with their vec_dot and vec_dot_type.
+static const struct dot_info k_dots[GGML_TYPE_COUNT] = {
+    [GGML_TYPE_F32] = {(ft_ggml_vec_dot_t) ggml_vec_dot_f32, GGML_TYPE_F32},
+    [GGML_TYPE_F16] = {(ft_ggml_vec_dot_t) ggml_vec_dot_f16, GGML_TYPE_F16},
+    [GGML_TYPE_BF16] = {(ft_ggml_vec_dot_t) ggml_vec_dot_bf16, GGML_TYPE_BF16},
+    [GGML_TYPE_Q1_0] = {ggml_vec_dot_q1_0_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_Q2_0] = {ggml_vec_dot_q2_0_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_Q4_0] = {ggml_vec_dot_q4_0_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_Q4_1] = {ggml_vec_dot_q4_1_q8_1, GGML_TYPE_Q8_1},
+    [GGML_TYPE_Q5_0] = {ggml_vec_dot_q5_0_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_Q5_1] = {ggml_vec_dot_q5_1_q8_1, GGML_TYPE_Q8_1},
+    [GGML_TYPE_Q8_0] = {ggml_vec_dot_q8_0_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_MXFP4] = {ggml_vec_dot_mxfp4_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_NVFP4] = {ggml_vec_dot_nvfp4_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_Q2_K] = {ggml_vec_dot_q2_K_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_Q3_K] = {ggml_vec_dot_q3_K_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_Q4_K] = {ggml_vec_dot_q4_K_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_Q5_K] = {ggml_vec_dot_q5_K_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_Q6_K] = {ggml_vec_dot_q6_K_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_TQ1_0] = {ggml_vec_dot_tq1_0_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_TQ2_0] = {ggml_vec_dot_tq2_0_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ2_XXS] = {ggml_vec_dot_iq2_xxs_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ2_XS] = {ggml_vec_dot_iq2_xs_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ2_S] = {ggml_vec_dot_iq2_s_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ3_XXS] = {ggml_vec_dot_iq3_xxs_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ3_S] = {ggml_vec_dot_iq3_s_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ1_S] = {ggml_vec_dot_iq1_s_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ1_M] = {ggml_vec_dot_iq1_m_q8_K, GGML_TYPE_Q8_K},
+    [GGML_TYPE_IQ4_NL] = {ggml_vec_dot_iq4_nl_q8_0, GGML_TYPE_Q8_0},
+    [GGML_TYPE_IQ4_XS] = {ggml_vec_dot_iq4_xs_q8_K, GGML_TYPE_Q8_K},
+};
+
+static void fp32_to_fp32_row(const float * x, void * y, int64_t k) { memcpy(y, x, (size_t) k * sizeof(float)); }
+
+static void fp32_to_fp16_row(const float * x, void * y, int64_t k) {
+    ggml_fp16_t * out = y;
+    for (int64_t i = 0; i < k; ++i) {
+        out[i] = GGML_CPU_FP32_TO_FP16(x[i]);
+    }
+}
+
+static void fp32_to_bf16_row(const float * x, void * y, int64_t k) {
+    ggml_bf16_t * out = y;
+    for (int64_t i = 0; i < k; ++i) {
+        out[i] = GGML_FP32_TO_BF16(x[i]);
+    }
+}
+
+static ft_ggml_from_float_t act_quantizer(int act_type) {
+    switch (act_type) {
+        case GGML_TYPE_F32: return fp32_to_fp32_row;
+        case GGML_TYPE_F16: return fp32_to_fp16_row;
+        case GGML_TYPE_BF16: return fp32_to_bf16_row;
+        case GGML_TYPE_Q8_0: return quantize_row_q8_0;
+        case GGML_TYPE_Q8_1: return quantize_row_q8_1;
+        case GGML_TYPE_Q8_K: return quantize_row_q8_K;
+        default: return NULL;
+    }
+}
+
+// The tables ggml_cpu_init fills for the quant kernels.
+void ft_ggml_cpu_init(void) {
+    for (int i = 0; i < (1 << 16); ++i) {
+        ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32((ggml_fp16_t) i);
+    }
+    for (int i = 0; i < (1 << 8); ++i) {
+        ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i);
+        ggml_table_f32_ue4m3[i] = ggml_ue4m3_to_fp32((uint8_t) i);
+    }
+}
+
+int ft_ggml_cpu_row_kernel(int type, int64_t k, struct ft_ggml_row_kernel * out) {
+    if (!known_type(type) || k_dots[type].vec_dot == NULL) {
+        return FT_GGML_NO_KERNEL;
+    }
+    const int act = k_dots[type].act_type;
+    if (k % k_types[type].blck_size != 0 || k % k_types[act].blck_size != 0) {
+        return FT_GGML_PARTIAL_BLOCK;
+    }
+    out->vec_dot = k_dots[type].vec_dot;
+    out->from_float = act_quantizer(act);
+    out->act_type = act;
+    out->row_bytes = ggml_row_size((enum ggml_type) type, k);
+    out->act_row_bytes = ggml_row_size((enum ggml_type) act, k);
+    return FT_GGML_OK;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h
new file mode 100644
index 0000000000..d58c174854
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h
@@ -0,0 +1,43 @@
+// ggml-cpu row kernels for the CPU MoE executor. setup.py builds them once per x86 variant into freetoken.kernel.ggml_cpu._<variant>; call them only through the table of a variant whose score() is not 0.
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef void (*ft_ggml_vec_dot_t)(int n, float * s, size_t bs, const void * x, size_t bx, const void * y, size_t by, int nrc);
+typedef void (*ft_ggml_from_float_t)(const float * x, void * y, int64_t k);
+
+struct ft_ggml_row_kernel {
+    ft_ggml_vec_dot_t vec_dot;        // one weight row . one activation row in act_type
+    ft_ggml_from_float_t from_float;  // fp32 activation row -> act_type
+    int act_type;
+    size_t row_bytes;                 // one weight row of k values
+    size_t act_row_bytes;             // one activation row of k values in act_type
+};
+
+enum { FT_GGML_OK = 0, FT_GGML_NO_KERNEL = 1, FT_GGML_PARTIAL_BLOCK = 2 };
+
+void ft_ggml_cpu_init(void);
+
+// The kernels for weight rows of ggml type `type` over k values.
+int ft_ggml_cpu_row_kernel(int type, int64_t k, struct ft_ggml_row_kernel * out);
+
+const char * ft_ggml_type_name(int type);
+
+// One variant's functions, as its module's api() hands them to _cpu_moe.set_ggml_cpu.
+#define FT_GGML_CPU_ABI 1
+struct ft_ggml_cpu_api {
+    int abi;           // FT_GGML_CPU_ABI of the module that built it
+    const char * name; // the variant, e.g. "haswell"
+    void (*init)(void);
+    int (*row_kernel)(int type, int64_t k, struct ft_ggml_row_kernel * out);
+    const char * (*type_name)(int type);
+};
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp b/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp
new file mode 100644
index 0000000000..8f3d3e5246
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp
@@ -0,0 +1,32 @@
+// One ggml-cpu variant as a Python module: score() is 0 unless this CPU has every ISA extension the variant is built for, and api() is the address of its kernel table.
+// setup.py builds this file without the variant's ISA flags, like upstream's <variant>-feats target, so importing the module and calling score() is safe on any x86 CPU.
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#include "cpu_kernels.h"
+#include "ggml-cpu/arch/x86/cpu-feats.cpp"
+
+#define FT_CAT_(a, b) a##b
+#define FT_CAT(a, b) FT_CAT_(a, b)
+#define FT_STR_(a) #a
+#define FT_STR(a) FT_STR_(a)
+
+namespace {
+
+const ft_ggml_cpu_api kApi = {FT_GGML_CPU_ABI, FT_STR(FT_GGML_VARIANT), ft_ggml_cpu_init, ft_ggml_cpu_row_kernel, ft_ggml_type_name};
+
+PyObject* score(PyObject*, PyObject*) { return PyLong_FromLong(ggml_backend_cpu_x86_score()); }
+
+PyObject* api(PyObject*, PyObject*) { return PyLong_FromVoidPtr(const_cast<ft_ggml_cpu_api*>(&kApi)); }
+
+PyMethodDef kMethods[] = {
+    {"score", score, METH_NOARGS, "0 if this CPU lacks an ISA extension the variant is built for, else upstream's variant score."},
+    {"api", api, METH_NOARGS, "The address of the variant's kernel table, for _cpu_moe.set_ggml_cpu."},
+    {nullptr, nullptr, 0, nullptr},
+};
+
+PyModuleDef kModule = {PyModuleDef_HEAD_INIT, "_" FT_STR(FT_GGML_VARIANT), nullptr, -1, kMethods};
+
+}  // namespace
+
+PyMODINIT_FUNC FT_CAT(PyInit__, FT_GGML_VARIANT)(void) { return PyModule_Create(&kModule); }
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h
new file mode 100644
index 0000000000..4e59d0ae0a
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h
@@ -0,0 +1,375 @@
+
+#pragma once
+
+// Rename `_generic` functions if no native implementation is available.
+// This effectively selects the generic implementation.
+
+#if defined(GGML_CPU_GENERIC)
+// quants.c
+#define quantize_row_q8_0_generic quantize_row_q8_0
+#define quantize_row_q8_1_generic quantize_row_q8_1
+#define quantize_row_q8_K_generic quantize_row_q8_K
+#define ggml_vec_dot_q4_0_q8_0_generic ggml_vec_dot_q4_0_q8_0
+#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1
+#define ggml_vec_dot_q5_0_q8_0_generic ggml_vec_dot_q5_0_q8_0
+#define ggml_vec_dot_q5_1_q8_1_generic ggml_vec_dot_q5_1_q8_1
+#define ggml_vec_dot_q8_0_q8_0_generic ggml_vec_dot_q8_0_q8_0
+#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
+#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
+#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
+#define ggml_vec_dot_q3_K_q8_K_generic ggml_vec_dot_q3_K_q8_K
+#define ggml_vec_dot_q4_K_q8_K_generic ggml_vec_dot_q4_K_q8_K
+#define ggml_vec_dot_q5_K_q8_K_generic ggml_vec_dot_q5_K_q8_K
+#define ggml_vec_dot_q6_K_q8_K_generic ggml_vec_dot_q6_K_q8_K
+#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K
+#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K
+#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K
+#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K
+#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
+#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
+#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0
+#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64) || defined(_M_ARM64EC)
+// repack.cpp
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__POWERPC__) || defined(__powerpc__)
+// ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679
+// quants.c
+#define quantize_row_q8_K_generic quantize_row_q8_K
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
+#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
+#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__loongarch64)
+// quants.c
+#define quantize_row_q8_K_generic quantize_row_q8_K
+#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
+#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
+#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__riscv)
+// quants.c
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__s390x__)
+// quants.c
+#define quantize_row_q8_K_generic quantize_row_q8_K
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
+#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
+#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
+#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K
+#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K
+#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K
+#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K
+#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
+#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
+#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#elif defined(__wasm__)
+// quants.c
+#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
+#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
+#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K
+#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K
+#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K
+#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K
+#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
+#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
+#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0
+#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K
+#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
+#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+// repack.cpp
+#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
+#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
+#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
+#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
+#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
+#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
+#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
+#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
+#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K
+#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
+#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
+#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
+#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K
+#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
+#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0
+#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
+#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
+#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
+#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
+#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
+#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
+#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
+#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
+#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
+#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K
+#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
+#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
+#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K
+#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K
+#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
+#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0
+#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
+#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
+#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
+#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp
new file mode 100644
index 0000000000..d775a03638
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp
@@ -0,0 +1,327 @@
+#include "ggml-backend-impl.h"
+
+#if defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64))
+
+#ifdef _MSC_VER
+#include <intrin.h>
+#endif
+
+#include <cstring>
+#include <vector>
+#include <bitset>
+#include <array>
+#include <string>
+
+// ref: https://cdrdv2-public.intel.com/782156/325383-sdm-vol-2abcd.pdf
+struct cpuid_x86 {
+    bool SSE3(void) { return f_1_ecx[0]; }
+    bool PCLMULQDQ(void) { return f_1_ecx[1]; }
+    bool MONITOR(void) { return f_1_ecx[3]; }
+    bool SSSE3(void) { return f_1_ecx[9]; }
+    bool FMA(void) { return f_1_ecx[12]; }
+    bool CMPXCHG16B(void) { return f_1_ecx[13]; }
+    bool SSE41(void) { return f_1_ecx[19]; }
+    bool SSE42(void) { return f_1_ecx[20]; }
+    bool MOVBE(void) { return f_1_ecx[22]; }
+    bool POPCNT(void) { return f_1_ecx[23]; }
+    bool AES(void) { return f_1_ecx[25]; }
+    bool XSAVE(void) { return f_1_ecx[26]; }
+    bool OSXSAVE(void) { return f_1_ecx[27]; }
+    bool AVX(void) { return f_1_ecx[28]; }
+    bool F16C(void) { return f_1_ecx[29]; }
+    bool RDRAND(void) { return f_1_ecx[30]; }
+
+    bool MSR(void) { return f_1_edx[5]; }
+    bool CX8(void) { return f_1_edx[8]; }
+    bool SEP(void) { return f_1_edx[11]; }
+    bool CMOV(void) { return f_1_edx[15]; }
+    bool CLFSH(void) { return f_1_edx[19]; }
+    bool MMX(void) { return f_1_edx[23]; }
+    bool FXSR(void) { return f_1_edx[24]; }
+    bool SSE(void) { return f_1_edx[25]; }
+    bool SSE2(void) { return f_1_edx[26]; }
+
+    bool FSGSBASE(void) { return f_7_ebx[0]; }
+    bool BMI1(void) { return f_7_ebx[3]; }
+    bool HLE(void) { return is_intel && f_7_ebx[4]; }
+    bool AVX2(void) { return f_7_ebx[5]; }
+    bool BMI2(void) { return f_7_ebx[8]; }
+    bool ERMS(void) { return f_7_ebx[9]; }
+    bool INVPCID(void) { return f_7_ebx[10]; }
+    bool RTM(void) { return is_intel && f_7_ebx[11]; }
+    bool AVX512F(void) { return f_7_ebx[16]; }
+    bool AVX512DQ(void) { return f_7_ebx[17]; }
+    bool RDSEED(void) { return f_7_ebx[18]; }
+    bool ADX(void) { return f_7_ebx[19]; }
+    bool AVX512PF(void) { return f_7_ebx[26]; }
+    bool AVX512ER(void) { return f_7_ebx[27]; }
+    bool AVX512CD(void) { return f_7_ebx[28]; }
+    bool AVX512BW(void) { return f_7_ebx[30]; }
+    bool AVX512VL(void) { return f_7_ebx[31]; }
+
+    bool SHA(void) { return f_7_ebx[29]; }
+
+    bool PREFETCHWT1(void) { return f_7_ecx[0]; }
+
+    bool LAHF(void) { return f_81_ecx[0]; }
+    bool LZCNT(void) { return is_intel && f_81_ecx[5]; }
+    bool ABM(void) { return is_amd && f_81_ecx[5]; }
+    bool SSE4a(void) { return is_amd && f_81_ecx[6]; }
+    bool XOP(void) { return is_amd && f_81_ecx[11]; }
+    bool TBM(void) { return is_amd && f_81_ecx[21]; }
+
+    bool SYSCALL(void) { return is_intel && f_81_edx[11]; }
+    bool MMXEXT(void) { return is_amd && f_81_edx[22]; }
+    bool RDTSCP(void) { return is_intel && f_81_edx[27]; }
+    bool _3DNOWEXT(void) { return is_amd && f_81_edx[30]; }
+    bool _3DNOW(void) { return is_amd && f_81_edx[31]; }
+
+    bool AVX512_VBMI(void) { return f_7_ecx[1]; }
+    bool AVX512_VNNI(void) { return f_7_ecx[11]; }
+    bool AVX512_FP16(void) { return f_7_edx[23]; }
+    bool AVX512_BF16(void) { return f_7_1_eax[5]; }
+    bool AVX_VNNI(void) { return f_7_1_eax[4]; }
+
+    bool AMX_TILE(void) { return f_7_edx[24]; }
+    bool AMX_INT8(void) { return f_7_edx[25]; }
+    bool AMX_FP16(void) { return f_7_1_eax[21]; }
+    bool AMX_BF16(void) { return f_7_edx[22]; }
+
+#ifdef _MSC_VER
+    static void cpuid(int cpu_info[4], int eax) {
+        __cpuid(cpu_info, eax);
+    }
+    static void cpuidex(int cpu_info[4], int eax, int ecx) {
+        __cpuidex(cpu_info, eax, ecx);
+    }
+#else
+    static void cpuid(int cpu_info[4], int eax) {
+        __asm__ __volatile__(
+            "cpuid"
+            : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3])
+            : "a"(eax), "c"(0));
+    }
+    static void cpuidex(int cpu_info[4], int eax, int ecx) {
+        __asm__ __volatile__(
+            "cpuid"
+            : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3])
+            : "a"(eax), "c"(ecx));
+    }
+#endif
+
+    cpuid_x86() {
+        std::array<int, 4> cpui;
+        std::vector<std::array<int, 4>> data;
+
+        // calling __cpuid with 0x0 as the function_id argument
+        // gets the number of the highest valid function ID.
+        cpuid(cpui.data(), 0);
+        int n_ids = cpui[0];
+
+        for (int i = 0; i <= n_ids; ++i) {
+            cpuidex(cpui.data(), i, 0);
+            data.push_back(cpui);
+        }
+
+        // capture vendor string
+        char vendor[0x20] = {};
+        *reinterpret_cast<int *>(vendor)     = data[0][1];
+        *reinterpret_cast<int *>(vendor + 4) = data[0][3];
+        *reinterpret_cast<int *>(vendor + 8) = data[0][2];
+        this->vendor = vendor;
+        if (this->vendor == "GenuineIntel") {
+            is_intel = true;
+        } else if (this->vendor == "AuthenticAMD") {
+            is_amd = true;
+        }
+
+        // load bitset with flags for function 0x00000001
+        if (n_ids >= 1) {
+            f_1_ecx = data[1][2];
+            f_1_edx = data[1][3];
+        }
+
+        // load bitset with flags for function 0x00000007
+        if (n_ids >= 7) {
+            f_7_ebx = data[7][1];
+            f_7_ecx = data[7][2];
+            f_7_edx = data[7][3];
+            cpuidex(cpui.data(), 7, 1);
+            f_7_1_eax = cpui[0];
+        }
+
+        // calling __cpuid with 0x80000000 as the function_id argument
+        // gets the number of the highest valid extended ID.
+        cpuid(cpui.data(), 0x80000000);
+        unsigned int n_ex_ids = cpui[0];
+
+        std::vector<std::array<int, 4>> ext_data;
+        for (unsigned int i = 0x80000000; i <= n_ex_ids; ++i) {
+            cpuidex(cpui.data(), i, 0);
+            ext_data.push_back(cpui);
+        }
+
+        // load bitset with flags for function 0x80000001
+        if (n_ex_ids >= 0x80000001) {
+            f_81_ecx = ext_data[1][2];
+            f_81_edx = ext_data[1][3];
+        }
+
+        // interpret CPU brand string if reported
+        char brand[0x40] = {};
+        if (n_ex_ids >= 0x80000004) {
+            std::memcpy(brand, ext_data[2].data(), sizeof(cpui));
+            std::memcpy(brand + 16, ext_data[3].data(), sizeof(cpui));
+            std::memcpy(brand + 32, ext_data[4].data(), sizeof(cpui));
+            this->brand = brand;
+        }
+    }
+
+    bool is_intel = false;
+    bool is_amd = false;
+    std::string vendor;
+    std::string brand;
+    std::bitset<32> f_1_ecx;
+    std::bitset<32> f_1_edx;
+    std::bitset<32> f_7_ebx;
+    std::bitset<32> f_7_ecx;
+    std::bitset<32> f_7_edx;
+    std::bitset<32> f_7_1_eax;
+    std::bitset<32> f_81_ecx;
+    std::bitset<32> f_81_edx;
+};
+
+#if 0
+void test_x86_is() {
+    cpuid_x86 is;
+    printf("CPU Vendor: %s\n", is.vendor.c_str());
+    printf("Brand: %s\n", is.brand.c_str());
+    printf("is_intel: %d\n", is.is_intel);
+    printf("is_amd: %d\n", is.is_amd);
+    printf("sse3: %d\n", is.SSE3());
+    printf("pclmulqdq: %d\n", is.PCLMULQDQ());
+    printf("ssse3: %d\n", is.SSSE3());
+    printf("fma: %d\n", is.FMA());
+    printf("cmpxchg16b: %d\n", is.CMPXCHG16B());
+    printf("sse41: %d\n", is.SSE41());
+    printf("sse42: %d\n", is.SSE42());
+    printf("movbe: %d\n", is.MOVBE());
+    printf("popcnt: %d\n", is.POPCNT());
+    printf("aes: %d\n", is.AES());
+    printf("xsave: %d\n", is.XSAVE());
+    printf("osxsave: %d\n", is.OSXSAVE());
+    printf("avx: %d\n", is.AVX());
+    printf("f16c: %d\n", is.F16C());
+    printf("rdrand: %d\n", is.RDRAND());
+    printf("msr: %d\n", is.MSR());
+    printf("cx8: %d\n", is.CX8());
+    printf("sep: %d\n", is.SEP());
+    printf("cmov: %d\n", is.CMOV());
+    printf("clflush: %d\n", is.CLFSH());
+    printf("mmx: %d\n", is.MMX());
+    printf("fxsr: %d\n", is.FXSR());
+    printf("sse: %d\n", is.SSE());
+    printf("sse2: %d\n", is.SSE2());
+    printf("fsgsbase: %d\n", is.FSGSBASE());
+    printf("bmi1: %d\n", is.BMI1());
+    printf("hle: %d\n", is.HLE());
+    printf("avx2: %d\n", is.AVX2());
+    printf("bmi2: %d\n", is.BMI2());
+    printf("erms: %d\n", is.ERMS());
+    printf("invpcid: %d\n", is.INVPCID());
+    printf("rtm: %d\n", is.RTM());
+    printf("avx512f: %d\n", is.AVX512F());
+    printf("rdseed: %d\n", is.RDSEED());
+    printf("adx: %d\n", is.ADX());
+    printf("avx512pf: %d\n", is.AVX512PF());
+    printf("avx512er: %d\n", is.AVX512ER());
+    printf("avx512cd: %d\n", is.AVX512CD());
+    printf("sha: %d\n", is.SHA());
+    printf("prefetchwt1: %d\n", is.PREFETCHWT1());
+    printf("lahf: %d\n", is.LAHF());
+    printf("lzcnt: %d\n", is.LZCNT());
+    printf("abm: %d\n", is.ABM());
+    printf("sse4a: %d\n", is.SSE4a());
+    printf("xop: %d\n", is.XOP());
+    printf("tbm: %d\n", is.TBM());
+    printf("syscall: %d\n", is.SYSCALL());
+    printf("mmxext: %d\n", is.MMXEXT());
+    printf("rdtscp: %d\n", is.RDTSCP());
+    printf("3dnowext: %d\n", is._3DNOWEXT());
+    printf("3dnow: %d\n", is._3DNOW());
+    printf("avx512_vbmi: %d\n", is.AVX512_VBMI());
+    printf("avx512_vnni: %d\n", is.AVX512_VNNI());
+    printf("avx512_fp16: %d\n", is.AVX512_FP16());
+    printf("avx512_bf16: %d\n", is.AVX512_BF16());
+    printf("amx_tile: %d\n", is.AMX_TILE());
+    printf("amx_int8: %d\n", is.AMX_INT8());
+    printf("amx_fp16: %d\n", is.AMX_FP16());
+    printf("amx_bf16: %d\n", is.AMX_BF16());
+}
+#endif
+
+static int ggml_backend_cpu_x86_score() {
+    // FIXME: this does not check for OS support
+
+    int score = 1;
+    cpuid_x86 is;
+
+#ifdef GGML_FMA
+    if (!is.FMA()) { return 0; }
+    score += 1;
+#endif
+#ifdef GGML_F16C
+    if (!is.F16C()) { return 0; }
+    score += 1<<1;
+#endif
+#ifdef GGML_SSE42
+    if (!is.SSE42()) { return 0; }
+    score += 1<<2;
+#endif
+#ifdef GGML_BMI2
+    if (!is.BMI2()) { return 0; }
+    score += 1<<3;
+#endif
+#ifdef GGML_AVX
+    if (!is.AVX()) { return 0; }
+    score += 1<<4;
+#endif
+#ifdef GGML_AVX2
+    if (!is.AVX2()) { return 0; }
+    score += 1<<5;
+#endif
+#ifdef GGML_AVX_VNNI
+    if (!is.AVX_VNNI()) { return 0; }
+    score += 1<<6;
+#endif
+#ifdef GGML_AVX512
+    if (!is.AVX512F()) { return 0; }
+    if (!is.AVX512CD()) { return 0; }
+    if (!is.AVX512VL()) { return 0; }
+    if (!is.AVX512DQ()) { return 0; }
+    if (!is.AVX512BW()) { return 0; }
+    score += 1<<7;
+#endif
+#ifdef GGML_AVX512_VBMI
+    if (!is.AVX512_VBMI()) { return 0; }
+    score += 1<<8;
+#endif
+#ifdef GGML_AVX512_BF16
+    if (!is.AVX512_BF16()) { return 0; }
+    score += 1<<9;
+#endif
+#ifdef GGML_AVX512_VNNI
+    if (!is.AVX512_VNNI()) { return 0; }
+    score += 1<<10;
+#endif
+#ifdef GGML_AMX_INT8
+    if (!is.AMX_INT8()) { return 0; }
+    score += 1<<11;
+#endif
+
+    return score;
+}
+
+GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_x86_score)
+
+#endif // defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64))
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c
new file mode 100644
index 0000000000..21cd6312a1
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c
@@ -0,0 +1,4172 @@
+#define GGML_COMMON_IMPL_C
+#include "ggml-common.h"
+#include "ggml-quants.h"
+#include "ggml-impl.h"
+#include "ggml-cpu.h"
+#include "simd-mappings.h"
+
+#include "../../quants.h"
+#include "../../ggml-cpu-impl.h"
+
+#include <math.h>
+#include <string.h>
+#include <assert.h>
+#include <stdlib.h> // for qsort
+#include <stdio.h>  // for GGML_ASSERT
+
+#define GROUP_MAX_EPS 1e-15f
+#define GROUP_MAX_EPS_IQ3_XXS 1e-8f
+#define GROUP_MAX_EPS_IQ2_S 1e-8f
+#define GROUP_MAX_EPS_IQ1_M 1e-7f
+#define GROUP_MAX_EPS_IQ1_S 1e-12f
+
+#define UNUSED GGML_UNUSED
+
+// some compilers don't provide _mm256_set_m128i, e.g. gcc 7
+#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1)
+
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__)
+// multiply int8_t, add results pairwise twice
+static inline __m128i mul_sum_i8_pairs(const __m128i x, const __m128i y) {
+    // Get absolute values of x vectors
+    const __m128i ax = _mm_sign_epi8(x, x);
+    // Sign the values of the y vectors
+    const __m128i sy = _mm_sign_epi8(y, x);
+    // Perform multiplication and create 16-bit values
+    const __m128i dot = _mm_maddubs_epi16(ax, sy);
+    const __m128i ones = _mm_set1_epi16(1);
+    return _mm_madd_epi16(ones, dot);
+}
+
+#if __AVX__ || __AVX2__ || __AVX512F__
+// horizontally add 8 floats
+static inline float hsum_float_8(const __m256 x) {
+    __m128 res = _mm256_extractf128_ps(x, 1);
+    res = _mm_add_ps(res, _mm256_castps256_ps128(x));
+    res = _mm_add_ps(res, _mm_movehl_ps(res, res));
+    res = _mm_add_ss(res, _mm_movehdup_ps(res));
+    return _mm_cvtss_f32(res);
+}
+
+// horizontally add 8 int32_t
+static inline int hsum_i32_8(const __m256i a) {
+    const __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(a), _mm256_extractf128_si256(a, 1));
+    const __m128i hi64 = _mm_unpackhi_epi64(sum128, sum128);
+    const __m128i sum64 = _mm_add_epi32(hi64, sum128);
+    const __m128i hi32  = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1));
+    return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32));
+}
+
+// horizontally add 4 int32_t
+static inline int hsum_i32_4(const __m128i a) {
+    const __m128i hi64 = _mm_unpackhi_epi64(a, a);
+    const __m128i sum64 = _mm_add_epi32(hi64, a);
+    const __m128i hi32  = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1));
+    return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32));
+}
+
+#if defined(__AVX2__) || defined(__AVX512F__)
+static inline __m256i mul_add_epi8(const __m256i x, const __m256i y) {
+    const __m256i ax = _mm256_sign_epi8(x, x);
+    const __m256i sy = _mm256_sign_epi8(y, x);
+    return _mm256_maddubs_epi16(ax, sy);
+}
+
+// spread 32 bits to 32 bytes { 0x00, 0xFF }
+static inline __m256i bytes_from_bits_32(const uint8_t * x) {
+    uint32_t x32;
+    memcpy(&x32, x, sizeof(uint32_t));
+    const __m256i shuf_mask = _mm256_set_epi64x(
+            0x0303030303030303, 0x0202020202020202,
+            0x0101010101010101, 0x0000000000000000);
+    __m256i bytes = _mm256_shuffle_epi8(_mm256_set1_epi32(x32), shuf_mask);
+    const __m256i bit_mask = _mm256_set1_epi64x(0x7fbfdfeff7fbfdfe);
+    bytes = _mm256_or_si256(bytes, bit_mask);
+    return _mm256_cmpeq_epi8(bytes, _mm256_set1_epi64x(-1));
+}
+
+// Unpack 32 4-bit fields into 32 bytes
+// The output vector contains 32 bytes, each one in [ 0 .. 15 ] interval
+static inline __m256i bytes_from_nibbles_32(const uint8_t * rsi)
+{
+    const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi);
+    const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp);
+    const __m256i lowMask = _mm256_set1_epi8( 0xF );
+    return _mm256_and_si256(lowMask, bytes);
+}
+
+// add int16_t pairwise and return as float vector
+static inline __m256 sum_i16_pairs_float(const __m256i x) {
+    const __m256i ones = _mm256_set1_epi16(1);
+    const __m256i summed_pairs = _mm256_madd_epi16(ones, x);
+    return _mm256_cvtepi32_ps(summed_pairs);
+}
+
+static inline __m256 mul_sum_us8_pairs_float(const __m256i ax, const __m256i sy) {
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
+    const __m256i zero = _mm256_setzero_si256();
+    const __m256i summed_pairs = _mm256_dpbusd_epi32(zero, ax, sy);
+    return _mm256_cvtepi32_ps(summed_pairs);
+#elif defined(__AVXVNNI__)
+    const __m256i zero = _mm256_setzero_si256();
+    const __m256i summed_pairs = _mm256_dpbusd_avx_epi32(zero, ax, sy);
+    return _mm256_cvtepi32_ps(summed_pairs);
+#else
+    // Perform multiplication and create 16-bit values
+    const __m256i dot = _mm256_maddubs_epi16(ax, sy);
+    return sum_i16_pairs_float(dot);
+#endif
+}
+
+// multiply int8_t, add results pairwise twice and return as float vector
+static inline __m256 mul_sum_i8_pairs_float(const __m256i x, const __m256i y) {
+#if __AVXVNNIINT8__
+    const __m256i zero = _mm256_setzero_si256();
+    const __m256i summed_pairs = _mm256_dpbssd_epi32(zero, x, y);
+    return _mm256_cvtepi32_ps(summed_pairs);
+#else
+    // Get absolute values of x vectors
+    const __m256i ax = _mm256_sign_epi8(x, x);
+    // Sign the values of the y vectors
+    const __m256i sy = _mm256_sign_epi8(y, x);
+    return mul_sum_us8_pairs_float(ax, sy);
+#endif
+}
+
+static inline __m128i packNibbles( __m256i bytes )
+{
+    // Move bits within 16-bit lanes from 0000_abcd_0000_efgh into 0000_0000_abcd_efgh
+#if __AVX512F__
+    const __m256i bytes_srli_4 = _mm256_srli_epi16(bytes, 4);   // 0000_0000_abcd_0000
+    bytes = _mm256_or_si256(bytes, bytes_srli_4);               // 0000_abcd_abcd_efgh
+    return _mm256_cvtepi16_epi8(bytes);                         // abcd_efgh
+#else
+    const __m256i lowByte = _mm256_set1_epi16( 0xFF );
+    __m256i high = _mm256_andnot_si256( lowByte, bytes );
+    __m256i low = _mm256_and_si256( lowByte, bytes );
+    high = _mm256_srli_epi16( high, 4 );
+    bytes = _mm256_or_si256( low, high );
+
+    // Compress uint16_t lanes into bytes
+    __m128i r0 = _mm256_castsi256_si128( bytes );
+    __m128i r1 = _mm256_extracti128_si256( bytes, 1 );
+    return _mm_packus_epi16( r0, r1 );
+#endif
+}
+#elif defined(__AVX__)
+static inline __m128i packNibbles( __m128i bytes1, __m128i bytes2 )
+{
+    // Move bits within 16-bit lanes from 0000_abcd_0000_efgh into 0000_0000_abcd_efgh
+    const __m128i lowByte = _mm_set1_epi16( 0xFF );
+    __m128i high = _mm_andnot_si128( lowByte, bytes1 );
+    __m128i low = _mm_and_si128( lowByte, bytes1 );
+    high = _mm_srli_epi16( high, 4 );
+    bytes1 = _mm_or_si128( low, high );
+    high = _mm_andnot_si128( lowByte, bytes2 );
+    low = _mm_and_si128( lowByte, bytes2 );
+    high = _mm_srli_epi16( high, 4 );
+    bytes2 = _mm_or_si128( low, high );
+
+    return _mm_packus_epi16( bytes1, bytes2);
+}
+
+static inline __m128i mul_add_epi8_sse(const __m128i x, const __m128i y) {
+    const __m128i ax = _mm_sign_epi8(x, x);
+    const __m128i sy = _mm_sign_epi8(y, x);
+    return _mm_maddubs_epi16(ax, sy);
+}
+
+// spread 32 bits to 32 bytes { 0x00, 0xFF }
+static inline __m256i bytes_from_bits_32(const uint8_t * x) {
+    uint32_t x32;
+    memcpy(&x32, x, sizeof(uint32_t));
+    const __m128i shuf_maskl = _mm_set_epi64x(0x0101010101010101, 0x0000000000000000);
+    const __m128i shuf_maskh = _mm_set_epi64x(0x0303030303030303, 0x0202020202020202);
+    __m128i bytesl = _mm_shuffle_epi8(_mm_set1_epi32(x32), shuf_maskl);
+    __m128i bytesh = _mm_shuffle_epi8(_mm_set1_epi32(x32), shuf_maskh);
+    const __m128i bit_mask = _mm_set1_epi64x(0x7fbfdfeff7fbfdfe);
+    bytesl = _mm_or_si128(bytesl, bit_mask);
+    bytesh = _mm_or_si128(bytesh, bit_mask);
+    bytesl = _mm_cmpeq_epi8(bytesl, _mm_set1_epi64x(-1));
+    bytesh = _mm_cmpeq_epi8(bytesh, _mm_set1_epi64x(-1));
+    return MM256_SET_M128I(bytesh, bytesl);
+}
+
+// Unpack 32 4-bit fields into 32 bytes
+// The output vector contains 32 bytes, each one in [ 0 .. 15 ] interval
+static inline __m256i bytes_from_nibbles_32(const uint8_t * rsi)
+{
+    // Load 16 bytes from memory
+    __m128i tmpl = _mm_loadu_si128((const __m128i *)rsi);
+    __m128i tmph = _mm_srli_epi16(tmpl, 4);
+    const __m128i lowMask = _mm_set1_epi8(0xF);
+    tmpl = _mm_and_si128(lowMask, tmpl);
+    tmph = _mm_and_si128(lowMask, tmph);
+    return MM256_SET_M128I(tmph, tmpl);
+}
+
+// add int16_t pairwise and return as float vector
+static inline __m256 sum_i16_pairs_float(const __m128i xh, const __m128i xl) {
+    const __m128i ones = _mm_set1_epi16(1);
+    const __m128i summed_pairsl = _mm_madd_epi16(ones, xl);
+    const __m128i summed_pairsh = _mm_madd_epi16(ones, xh);
+    const __m256i summed_pairs = MM256_SET_M128I(summed_pairsh, summed_pairsl);
+    return _mm256_cvtepi32_ps(summed_pairs);
+}
+
+static inline __m256 mul_sum_us8_pairs_float(const __m256i ax, const __m256i sy) {
+    const __m128i axl = _mm256_castsi256_si128(ax);
+    const __m128i axh = _mm256_extractf128_si256(ax, 1);
+    const __m128i syl = _mm256_castsi256_si128(sy);
+    const __m128i syh = _mm256_extractf128_si256(sy, 1);
+    // Perform multiplication and create 16-bit values
+    const __m128i dotl = _mm_maddubs_epi16(axl, syl);
+    const __m128i doth = _mm_maddubs_epi16(axh, syh);
+    return sum_i16_pairs_float(doth, dotl);
+}
+
+// multiply int8_t, add results pairwise twice and return as float vector
+static inline __m256 mul_sum_i8_pairs_float(const __m256i x, const __m256i y) {
+    const __m128i xl = _mm256_castsi256_si128(x);
+    const __m128i xh = _mm256_extractf128_si256(x, 1);
+    const __m128i yl = _mm256_castsi256_si128(y);
+    const __m128i yh = _mm256_extractf128_si256(y, 1);
+    // Get absolute values of x vectors
+    const __m128i axl = _mm_sign_epi8(xl, xl);
+    const __m128i axh = _mm_sign_epi8(xh, xh);
+    // Sign the values of the y vectors
+    const __m128i syl = _mm_sign_epi8(yl, xl);
+    const __m128i syh = _mm_sign_epi8(yh, xh);
+    // Perform multiplication and create 16-bit values
+    const __m128i dotl = _mm_maddubs_epi16(axl, syl);
+    const __m128i doth = _mm_maddubs_epi16(axh, syh);
+    return sum_i16_pairs_float(doth, dotl);
+}
+
+// larger version of mul_sum_i8_pairs_float where x and y are each represented by four 128-bit vectors
+static inline __m256 mul_sum_i8_quad_float(const __m128i x_1_0, const __m128i x_1_1, const __m128i x_2_0, const __m128i x_2_1,
+                                           const __m128i y_1_0, const __m128i y_1_1, const __m128i y_2_0, const __m128i y_2_1) {
+    const __m128i mone = _mm_set1_epi16(1);
+
+    const __m128i p16_1_0 = mul_add_epi8_sse(x_1_0, y_1_0);
+    const __m128i p16_1_1 = mul_add_epi8_sse(x_1_1, y_1_1);
+    const __m128i p16_2_0 = mul_add_epi8_sse(x_2_0, y_2_0);
+    const __m128i p16_2_1 = mul_add_epi8_sse(x_2_1, y_2_1);
+    const __m128i p_1_0 = _mm_madd_epi16(p16_1_0, mone);
+    const __m128i p_1_1 = _mm_madd_epi16(p16_1_1, mone);
+    const __m128i p_2_0 = _mm_madd_epi16(p16_2_0, mone);
+    const __m128i p_2_1 = _mm_madd_epi16(p16_2_1, mone);
+    const __m128i p_1 = _mm_add_epi32(p_1_0, p_1_1);
+    const __m128i p_2 = _mm_add_epi32(p_2_0, p_2_1);
+    return _mm256_cvtepi32_ps(MM256_SET_M128I(p_2, p_1));
+}
+
+// quad fp16 delta calculation
+static inline __m256 quad_fp16_delta_float(const float x0, const float y0, const float x1, const float y1) {
+    // GGML_CPU_FP16_TO_FP32 is faster than Intel F16C
+    return _mm256_set_m128(_mm_set1_ps(GGML_CPU_FP16_TO_FP32(x1) * GGML_CPU_FP16_TO_FP32(y1)),
+                           _mm_set1_ps(GGML_CPU_FP16_TO_FP32(x0) * GGML_CPU_FP16_TO_FP32(y0)));
+}
+
+static inline __m256 quad_mx_delta_float(const uint8_t x0, const float y0, const uint8_t x1, const float y1) {
+    return _mm256_set_m128(_mm_set1_ps(GGML_CPU_E8M0_TO_FP32_HALF(x1) * GGML_CPU_FP16_TO_FP32(y1)),
+                           _mm_set1_ps(GGML_CPU_E8M0_TO_FP32_HALF(x0) * GGML_CPU_FP16_TO_FP32(y0)));
+}
+#endif
+#elif defined(__SSSE3__)
+static inline __m128i bytes_from_bits_16(const uint8_t * x) {
+    uint16_t x16;
+    memcpy(&x16, x, sizeof(uint16_t));
+
+    const __m128i shuf_mask = _mm_set_epi64x(0x0101010101010101, 0x0000000000000000);
+    __m128i bytes = _mm_shuffle_epi8(_mm_set1_epi16((short) x16), shuf_mask);
+    const __m128i bit_mask = _mm_set_epi64x(0x7fbfdfeff7fbfdfe, 0x7fbfdfeff7fbfdfe);
+    bytes = _mm_or_si128(bytes, bit_mask);
+
+    return _mm_cmpeq_epi8(bytes, _mm_set1_epi64x(-1));
+}
+
+// horizontally add 4x4 floats
+static inline float hsum_float_4x4(const __m128 a, const __m128 b, const __m128 c, const __m128 d) {
+    __m128 res_0 =_mm_hadd_ps(a, b);
+    __m128 res_1 =_mm_hadd_ps(c, d);
+    __m128 res =_mm_hadd_ps(res_0, res_1);
+    res =_mm_hadd_ps(res, res);
+    res =_mm_hadd_ps(res, res);
+
+    return _mm_cvtss_f32(res);
+}
+#endif // __AVX__ || __AVX2__ || __AVX512F__
+#endif // defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__)
+
+void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(QK8_0 == 32);
+    assert(k % QK8_0 == 0);
+    const int nb = k / QK8_0;
+
+    block_q8_0 * GGML_RESTRICT y = vy;
+
+#if defined(__AVX2__) || defined(__AVX__)
+    for (int i = 0; i < nb; i++) {
+        // Load elements into 4 AVX vectors
+        __m256 v0 = _mm256_loadu_ps( x );
+        __m256 v1 = _mm256_loadu_ps( x + 8 );
+        __m256 v2 = _mm256_loadu_ps( x + 16 );
+        __m256 v3 = _mm256_loadu_ps( x + 24 );
+        x += 32;
+
+        // Compute max(abs(e)) for the block
+        const __m256 signBit = _mm256_set1_ps( -0.0f );
+        __m256 maxAbs = _mm256_andnot_ps( signBit, v0 );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) );
+
+        __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) );
+        max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) );
+        max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) );
+        const float maxScalar = _mm_cvtss_f32( max4 );
+
+        // Quantize these floats
+        const float d = maxScalar / 127.f;
+        y[i].d = GGML_CPU_FP32_TO_FP16(d);
+        const float id = ( maxScalar != 0.0f ) ? 127.f / maxScalar : 0.0f;
+        const __m256 mul = _mm256_set1_ps( id );
+
+        // Apply the multiplier
+        v0 = _mm256_mul_ps( v0, mul );
+        v1 = _mm256_mul_ps( v1, mul );
+        v2 = _mm256_mul_ps( v2, mul );
+        v3 = _mm256_mul_ps( v3, mul );
+
+        // Round to nearest integer
+        v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST );
+        v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST );
+        v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST );
+        v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST );
+
+        // Convert floats to integers
+        __m256i i0 = _mm256_cvtps_epi32( v0 );
+        __m256i i1 = _mm256_cvtps_epi32( v1 );
+        __m256i i2 = _mm256_cvtps_epi32( v2 );
+        __m256i i3 = _mm256_cvtps_epi32( v3 );
+
+#if defined(__AVX2__)
+        // Convert int32 to int16
+        i0 = _mm256_packs_epi32( i0, i1 );	// 0, 1, 2, 3,  8, 9, 10, 11,  4, 5, 6, 7, 12, 13, 14, 15
+        i2 = _mm256_packs_epi32( i2, i3 );	// 16, 17, 18, 19,  24, 25, 26, 27,  20, 21, 22, 23, 28, 29, 30, 31
+                                            // Convert int16 to int8
+        i0 = _mm256_packs_epi16( i0, i2 );	// 0, 1, 2, 3,  8, 9, 10, 11,  16, 17, 18, 19,  24, 25, 26, 27,  4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
+
+        // We got our precious signed bytes, but the order is now wrong
+        // These AVX2 pack instructions process 16-byte pieces independently
+        // The following instruction is fixing the order
+        const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 );
+        i0 = _mm256_permutevar8x32_epi32( i0, perm );
+
+        _mm256_storeu_si256((__m256i *)y[i].qs, i0);
+#else
+        // Since we don't have in AVX some necessary functions,
+        // we split the registers in half and call AVX2 analogs from SSE
+        __m128i ni0 = _mm256_castsi256_si128( i0 );
+        __m128i ni1 = _mm256_extractf128_si256( i0, 1);
+        __m128i ni2 = _mm256_castsi256_si128( i1 );
+        __m128i ni3 = _mm256_extractf128_si256( i1, 1);
+        __m128i ni4 = _mm256_castsi256_si128( i2 );
+        __m128i ni5 = _mm256_extractf128_si256( i2, 1);
+        __m128i ni6 = _mm256_castsi256_si128( i3 );
+        __m128i ni7 = _mm256_extractf128_si256( i3, 1);
+
+        // Convert int32 to int16
+        ni0 = _mm_packs_epi32( ni0, ni1 );
+        ni2 = _mm_packs_epi32( ni2, ni3 );
+        ni4 = _mm_packs_epi32( ni4, ni5 );
+        ni6 = _mm_packs_epi32( ni6, ni7 );
+        // Convert int16 to int8
+        ni0 = _mm_packs_epi16( ni0, ni2 );
+        ni4 = _mm_packs_epi16( ni4, ni6 );
+
+        _mm_storeu_si128((__m128i *)(y[i].qs +  0), ni0);
+        _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4);
+#endif
+    }
+#else
+    GGML_UNUSED(nb);
+    // scalar
+    quantize_row_q8_0_ref(x, y, k);
+#endif
+}
+
+void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK8_1 == 0);
+    const int nb = k / QK8_1;
+
+    block_q8_1 * GGML_RESTRICT y = vy;
+#if defined(__AVX2__) || defined(__AVX__)
+    for (int i = 0; i < nb; i++) {
+        // Load elements into 4 AVX vectors
+        __m256 v0 = _mm256_loadu_ps( x );
+        __m256 v1 = _mm256_loadu_ps( x + 8 );
+        __m256 v2 = _mm256_loadu_ps( x + 16 );
+        __m256 v3 = _mm256_loadu_ps( x + 24 );
+        x += 32;
+
+        // Compute max(abs(e)) for the block
+        const __m256 signBit = _mm256_set1_ps( -0.0f );
+        __m256 maxAbs = _mm256_andnot_ps( signBit, v0 );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) );
+        maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) );
+
+        __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) );
+        max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) );
+        max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) );
+        const float max_scalar = _mm_cvtss_f32( max4 );
+
+        // Quantize these floats
+        const float d = max_scalar / 127.f;
+        y[i].d = GGML_CPU_FP32_TO_FP16(d);
+        const float id = ( max_scalar != 0.0f ) ? 127.f / max_scalar : 0.0f;
+        const __m256 mul = _mm256_set1_ps( id );
+
+        // Apply the multiplier
+        v0 = _mm256_mul_ps( v0, mul );
+        v1 = _mm256_mul_ps( v1, mul );
+        v2 = _mm256_mul_ps( v2, mul );
+        v3 = _mm256_mul_ps( v3, mul );
+
+        // Round to nearest integer
+        v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST );
+        v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST );
+        v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST );
+        v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST );
+
+        // Convert floats to integers
+        __m256i i0 = _mm256_cvtps_epi32( v0 );
+        __m256i i1 = _mm256_cvtps_epi32( v1 );
+        __m256i i2 = _mm256_cvtps_epi32( v2 );
+        __m256i i3 = _mm256_cvtps_epi32( v3 );
+
+#if defined(__AVX2__)
+        // Compute the sum of the quants and set y[i].s
+        y[i].s = GGML_CPU_FP32_TO_FP16(d * hsum_i32_8(_mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3))));
+
+        // Convert int32 to int16
+        i0 = _mm256_packs_epi32( i0, i1 );	// 0, 1, 2, 3,  8, 9, 10, 11,  4, 5, 6, 7, 12, 13, 14, 15
+        i2 = _mm256_packs_epi32( i2, i3 );	// 16, 17, 18, 19,  24, 25, 26, 27,  20, 21, 22, 23, 28, 29, 30, 31
+                                            // Convert int16 to int8
+        i0 = _mm256_packs_epi16( i0, i2 );	// 0, 1, 2, 3,  8, 9, 10, 11,  16, 17, 18, 19,  24, 25, 26, 27,  4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
+
+        // We got our precious signed bytes, but the order is now wrong
+        // These AVX2 pack instructions process 16-byte pieces independently
+        // The following instruction is fixing the order
+        const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 );
+        i0 = _mm256_permutevar8x32_epi32( i0, perm );
+
+        _mm256_storeu_si256((__m256i *)y[i].qs, i0);
+#else
+        // Since we don't have in AVX some necessary functions,
+        // we split the registers in half and call AVX2 analogs from SSE
+        __m128i ni0 = _mm256_castsi256_si128( i0 );
+        __m128i ni1 = _mm256_extractf128_si256( i0, 1);
+        __m128i ni2 = _mm256_castsi256_si128( i1 );
+        __m128i ni3 = _mm256_extractf128_si256( i1, 1);
+        __m128i ni4 = _mm256_castsi256_si128( i2 );
+        __m128i ni5 = _mm256_extractf128_si256( i2, 1);
+        __m128i ni6 = _mm256_castsi256_si128( i3 );
+        __m128i ni7 = _mm256_extractf128_si256( i3, 1);
+
+        // Compute the sum of the quants and set y[i].s
+        const __m128i s0 = _mm_add_epi32(_mm_add_epi32(ni0, ni1), _mm_add_epi32(ni2, ni3));
+        const __m128i s1 = _mm_add_epi32(_mm_add_epi32(ni4, ni5), _mm_add_epi32(ni6, ni7));
+        y[i].s = GGML_CPU_FP32_TO_FP16(d * hsum_i32_4(_mm_add_epi32(s0, s1)));
+
+        // Convert int32 to int16
+        ni0 = _mm_packs_epi32( ni0, ni1 );
+        ni2 = _mm_packs_epi32( ni2, ni3 );
+        ni4 = _mm_packs_epi32( ni4, ni5 );
+        ni6 = _mm_packs_epi32( ni6, ni7 );
+        // Convert int16 to int8
+        ni0 = _mm_packs_epi16( ni0, ni2 );
+        ni4 = _mm_packs_epi16( ni4, ni6 );
+
+        _mm_storeu_si128((__m128i *)(y[i].qs +  0), ni0);
+        _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4);
+#endif
+    }
+#else
+    GGML_UNUSED(nb);
+    // scalar
+    quantize_row_q8_1_ref(x, y, k);
+#endif
+}
+
+// placeholder implementation for Apple targets
+void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q8_K_ref(x, y, k);
+}
+
+//===================================== Dot products =================================
+
+//
+// Helper functions
+//
+
+#if __AVX__ || __AVX2__ || __AVX512F__
+
+// shuffles to pick the required scales in dot products
+static inline __m256i get_scale_shuffle_q3k(int i) {
+    static const uint8_t k_shuffle[128] = {
+         0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1,     2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3,
+         4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5,     6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7,
+         8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9,    10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,
+        12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,    14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,
+    };
+    return _mm256_loadu_si256((const __m256i*)k_shuffle + i);
+}
+static inline __m256i get_scale_shuffle_k4(int i) {
+    static const uint8_t k_shuffle[256] = {
+         0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1,
+         2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3,
+         4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5,
+         6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7,
+         8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9,
+        10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,
+        12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,
+        14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15
+    };
+    return _mm256_loadu_si256((const __m256i*)k_shuffle + i);
+}
+static inline __m128i get_scale_shuffle(int i) {
+    static const uint8_t k_shuffle[128] = {
+         0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
+         2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
+         4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5,
+         6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7,
+         8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9,
+        10,10,10,10,10,10,10,10, 11,11,11,11,11,11,11,11,
+        12,12,12,12,12,12,12,12, 13,13,13,13,13,13,13,13,
+        14,14,14,14,14,14,14,14, 15,15,15,15,15,15,15,15
+    };
+    return _mm_loadu_si128((const __m128i*)k_shuffle + i);
+}
+#endif
+
+void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK1_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q1_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+#if defined(__AVX2__)
+    const __m256i ones_8 = _mm256_set1_epi8(1);
+    const __m256i ones_16 = _mm256_set1_epi16(1);
+    const __m256i byte_shuf = _mm256_setr_epi8(
+            0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
+            2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3);
+    const __m256i bit_masks = _mm256_setr_epi8(
+            1, 2, 4, 8, 16, 32, 64, (char) -128, 1, 2, 4, 8, 16, 32, 64, (char) -128,
+            1, 2, 4, 8, 16, 32, 64, (char) -128, 1, 2, 4, 8, 16, 32, 64, (char) -128);
+    const __m256i zero = _mm256_setzero_si256();
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int ib = 0; ib < nb; ++ib) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d);
+        const uint32_t * GGML_RESTRICT qs32 = (const uint32_t *) x[ib].qs;
+        const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4];
+
+        __m256 acc_block;
+        {
+            const __m256i qy = _mm256_loadu_si256((const __m256i *) y_ptr[0].qs);
+            const __m256i sm = _mm256_cmpeq_epi8(
+                    _mm256_and_si256(_mm256_shuffle_epi8(_mm256_set1_epi32((int) qs32[0]), byte_shuf), bit_masks), zero);
+            const __m256i sy = _mm256_sub_epi8(_mm256_xor_si256(qy, sm), sm);
+            const __m256i s32 = _mm256_madd_epi16(_mm256_maddubs_epi16(ones_8, sy), ones_16);
+            acc_block = _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[0].d)), _mm256_cvtepi32_ps(s32));
+        }
+        for (int K = 1; K < 4; ++K) {
+            const __m256i qy = _mm256_loadu_si256((const __m256i *) y_ptr[K].qs);
+            const __m256i sm = _mm256_cmpeq_epi8(
+                    _mm256_and_si256(_mm256_shuffle_epi8(_mm256_set1_epi32((int) qs32[K]), byte_shuf), bit_masks), zero);
+            const __m256i sy = _mm256_sub_epi8(_mm256_xor_si256(qy, sm), sm);
+            const __m256i s32 = _mm256_madd_epi16(_mm256_maddubs_epi16(ones_8, sy), ones_16);
+            acc_block = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[K].d)), _mm256_cvtepi32_ps(s32), acc_block);
+        }
+        acc = _mm256_fmadd_ps(_mm256_set1_ps(d0), acc_block, acc);
+    }
+
+    *s = hsum_float_8(acc);
+#elif defined(__AVX__)
+    const __m128i ones_8 = _mm_set1_epi8(1);
+    const __m128i ones_16 = _mm_set1_epi16(1);
+    const __m128i zero = _mm_setzero_si128();
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int ib = 0; ib < nb; ++ib) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d);
+        const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4];
+        __m256 acc_block;
+        {
+            const __m256i bit_mask = bytes_from_bits_32(&x[ib].qs[0]);
+            const __m128i bit_mask_0 = _mm256_castsi256_si128(bit_mask);
+            const __m128i bit_mask_1 = _mm256_extractf128_si256(bit_mask, 1);
+            const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[0].qs[0]);
+            const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[0].qs[16]);
+            const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero);
+            const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero);
+            const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0);
+            const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1);
+            const __m128i sum16_0 = _mm_maddubs_epi16(ones_8, sy_0);
+            const __m128i sum16_1 = _mm_maddubs_epi16(ones_8, sy_1);
+            const __m128i sum32_0 = _mm_madd_epi16(sum16_0, ones_16);
+            const __m128i sum32_1 = _mm_madd_epi16(sum16_1, ones_16);
+            const __m256 q = _mm256_cvtepi32_ps(MM256_SET_M128I(sum32_1, sum32_0));
+            acc_block = _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[0].d)), q);
+        }
+        for(int K = 1; K < 4; ++K) {
+            const __m256i bit_mask = bytes_from_bits_32(&x[ib].qs[(K) * 4]);
+            const __m128i bit_mask_0 = _mm256_castsi256_si128(bit_mask);
+            const __m128i bit_mask_1 = _mm256_extractf128_si256(bit_mask, 1);
+            const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[(K)].qs[0]);
+            const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[(K)].qs[16]);
+            const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero);
+            const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero);
+            const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0);
+            const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1);
+            const __m128i sum16_0 = _mm_maddubs_epi16(ones_8, sy_0);
+            const __m128i sum16_1 = _mm_maddubs_epi16(ones_8, sy_1);
+            const __m128i sum32_0 = _mm_madd_epi16(sum16_0, ones_16);
+            const __m128i sum32_1 = _mm_madd_epi16(sum16_1, ones_16);
+            const __m256 q = _mm256_cvtepi32_ps(MM256_SET_M128I(sum32_1, sum32_0));
+            acc_block = _mm256_add_ps(acc_block, _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[(K)].d)), q));
+        }
+#undef Q1_AVX_BLOCK
+
+        acc = _mm256_add_ps(acc, _mm256_mul_ps(_mm256_set1_ps(d0), acc_block));
+    }
+
+    *s = hsum_float_8(acc);
+#elif defined(__SSSE3__)
+    const __m128i ones_8 = _mm_set1_epi8(1);
+    const __m128i ones_16 = _mm_set1_epi16(1);
+    const __m128i zero = _mm_setzero_si128();
+    __m128 acc_0 = _mm_setzero_ps();
+    __m128 acc_1 = _mm_setzero_ps();
+    __m128 acc_2 = _mm_setzero_ps();
+    __m128 acc_3 = _mm_setzero_ps();
+
+    for (int ib = 0; ib < nb; ++ib) {
+        const __m128 d0 = _mm_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d));
+        const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4];
+
+#define Q1_SSSE3_BLOCK(QS_OFF, Y_IDX, ACC) \
+        { \
+            const __m128i bit_mask_0 = bytes_from_bits_16(&x[ib].qs[(QS_OFF) + 0]); \
+            const __m128i bit_mask_1 = bytes_from_bits_16(&x[ib].qs[(QS_OFF) + 2]); \
+            const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[(Y_IDX)].qs[0]); \
+            const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[(Y_IDX)].qs[16]); \
+            const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero); \
+            const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero); \
+            const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0); \
+            const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1); \
+            const __m128i sum_0 = _mm_madd_epi16(_mm_maddubs_epi16(ones_8, sy_0), ones_16); \
+            const __m128i sum_1 = _mm_madd_epi16(_mm_maddubs_epi16(ones_8, sy_1), ones_16); \
+            const __m128 q = _mm_cvtepi32_ps(_mm_add_epi32(sum_0, sum_1)); \
+            (ACC) = _mm_add_ps((ACC), _mm_mul_ps(_mm_mul_ps(d0, _mm_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[(Y_IDX)].d))), q)); \
+        }
+        Q1_SSSE3_BLOCK(0,  0, acc_0)
+        Q1_SSSE3_BLOCK(4,  1, acc_1)
+        Q1_SSSE3_BLOCK(8,  2, acc_2)
+        Q1_SSSE3_BLOCK(12, 3, acc_3)
+#undef Q1_SSSE3_BLOCK
+    }
+
+    *s = hsum_float_4x4(acc_0, acc_1, acc_2, acc_3);
+#else
+    UNUSED(nb);
+    UNUSED(x);
+    UNUSED(y);
+    ggml_vec_dot_q1_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+#if defined(__AVX2__)
+// four consecutive Q2_0 weights per byte: byte b becomes b & 3, (b >> 2) & 3, (b >> 4) & 3, (b >> 6) & 3 (OR, never add: the overlapping copies only touch masked-off bits)
+static inline __m256i q2_0_unpack_32_avx2(const uint8_t * qs) {
+    __m256i v = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *) qs));
+    v = _mm256_or_si256(v, _mm256_slli_epi32(v, 12));
+    v = _mm256_or_si256(v, _mm256_slli_epi32(v,  6));
+    return _mm256_and_si256(v, _mm256_set1_epi32(0x03030303));
+}
+
+// 32 weights against one Q8_0 block: sum((v - 1) * y) = sum(v * y) - sum(y), exact for every int8 y
+static inline __m256 q2_0_dot_32_avx2(const uint8_t * qs, const int8_t * qy) {
+    const __m256i v = q2_0_unpack_32_avx2(qs);
+    const __m256i y = _mm256_loadu_si256((const __m256i *) qy);
+    const __m256i p = _mm256_sub_epi16(_mm256_maddubs_epi16(v, y), _mm256_maddubs_epi16(_mm256_set1_epi8(1), y));
+    return _mm256_cvtepi32_ps(_mm256_madd_epi16(p, _mm256_set1_epi16(1)));
+}
+
+static inline __m256 q2_0_acc_block_avx2(const block_q2_0 * x, const block_q8_0 * y, __m256 acc) {
+    const float  dx = GGML_CPU_FP16_TO_FP32(x->d);
+    const __m256 d0 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[0].d));
+    const __m256 d1 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[1].d));
+    const __m256 p0 = q2_0_dot_32_avx2(x->qs + 0, y[0].qs);
+    const __m256 p1 = q2_0_dot_32_avx2(x->qs + 8, y[1].qs);
+#if defined(__FMA__)
+    acc = _mm256_fmadd_ps(d0, p0, acc);
+    return _mm256_fmadd_ps(d1, p1, acc);
+#else
+    acc = _mm256_add_ps(acc, _mm256_mul_ps(d0, p0));
+    return _mm256_add_ps(acc, _mm256_mul_ps(d1, p1));
+#endif
+}
+#endif
+
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+#if defined(__AVX2__)
+    assert(n % QK2_0 == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+    const int nb = n / QK2_0;
+
+    // two Q2_0 blocks per step, each into its own accumulator (two FMA chains)
+    __m256 acc0 = _mm256_setzero_ps();
+    __m256 acc1 = _mm256_setzero_ps();
+    int i = 0;
+    for (; i + 1 < nb; i += 2) {
+        acc0 = q2_0_acc_block_avx2(x + i + 0, y + 2*i + 0, acc0);
+        acc1 = q2_0_acc_block_avx2(x + i + 1, y + 2*i + 2, acc1);
+    }
+    if (i < nb) {
+        acc0 = q2_0_acc_block_avx2(x + i, y + 2*i, acc0);
+    }
+
+    *s = hsum_float_8(_mm256_add_ps(acc0, acc1));
+#else
+    ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined(__AVX2__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        /* Compute combined scale for the block */
+        const __m256 d = _mm256_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) );
+
+        __m256i qx = bytes_from_nibbles_32(x[ib].qs);
+
+        // Now we have a vector with bytes in [ 0 .. 15 ] interval. Offset them into [ -8 .. +7 ] interval.
+        const __m256i off = _mm256_set1_epi8( 8 );
+        qx = _mm256_sub_epi8( qx, off );
+
+        __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_i8_pairs_float(qx, qy);
+
+        /* Multiply q with scale and accumulate */
+        acc = _mm256_fmadd_ps( d, q, acc );
+    }
+
+    sumf = hsum_float_8(acc);
+#elif defined(__AVX__)
+    __m256 accum = _mm256_setzero_ps();
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs);
+        const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs);
+        const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs);
+        const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1);
+        const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs);
+        const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1);
+
+        const __m128i q4b_1_0 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), q4bits_1), _mm_set1_epi8(8));
+        const __m128i q4b_1_1 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(q4bits_1, 4)), _mm_set1_epi8(8));
+        const __m128i q4b_2_0 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), q4bits_2), _mm_set1_epi8(8));
+        const __m128i q4b_2_1 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(q4bits_2, 4)), _mm_set1_epi8(8));
+
+        const __m128i p16_1_0 = mul_add_epi8_sse(q4b_1_0, q8b_1_0);
+        const __m128i p16_1_1 = mul_add_epi8_sse(q4b_1_1, q8b_1_1);
+        const __m128i p16_2_0 = mul_add_epi8_sse(q4b_2_0, q8b_2_0);
+        const __m128i p16_2_1 = mul_add_epi8_sse(q4b_2_1, q8b_2_1);
+        const __m128i p_1 = _mm_add_epi16(p16_1_0, p16_1_1);
+        const __m128i p_2 = _mm_add_epi16(p16_2_0, p16_2_1);
+        const __m256 p =  sum_i16_pairs_float(p_2, p_1);
+
+        const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d);
+        accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+#elif defined(__SSSE3__)
+    // set constants
+    const __m128i lowMask = _mm_set1_epi8(0xF);
+    const __m128i off = _mm_set1_epi8(8);
+
+    // Initialize accumulator with zeros
+    __m128 acc_0 = _mm_setzero_ps();
+    __m128 acc_1 = _mm_setzero_ps();
+    __m128 acc_2 = _mm_setzero_ps();
+    __m128 acc_3 = _mm_setzero_ps();
+
+    for (; ib + 1 < nb; ib += 2) {
+        _mm_prefetch(&x[ib] + sizeof(block_q4_0), _MM_HINT_T0);
+        _mm_prefetch(&y[ib] + sizeof(block_q8_0), _MM_HINT_T0);
+
+        // Compute combined scale for the block 0 and 1
+        const __m128 d_0_1 = _mm_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) );
+
+        const __m128i tmp_0_1 = _mm_loadu_si128((const __m128i *)x[ib].qs);
+
+        __m128i bx_0 = _mm_and_si128(lowMask, tmp_0_1);
+        __m128i by_0 = _mm_loadu_si128((const __m128i *)y[ib].qs);
+        bx_0 = _mm_sub_epi8(bx_0, off);
+        const __m128i i32_0 = mul_sum_i8_pairs(bx_0, by_0);
+
+        __m128i bx_1 = _mm_and_si128(lowMask, _mm_srli_epi64(tmp_0_1, 4));
+        __m128i by_1 = _mm_loadu_si128((const __m128i *)(y[ib].qs + 16));
+        bx_1 = _mm_sub_epi8(bx_1, off);
+        const __m128i i32_1 = mul_sum_i8_pairs(bx_1, by_1);
+
+        _mm_prefetch(&x[ib] + 2 * sizeof(block_q4_0), _MM_HINT_T0);
+        _mm_prefetch(&y[ib] + 2 * sizeof(block_q8_0), _MM_HINT_T0);
+
+        // Compute combined scale for the block 2 and 3
+        const __m128 d_2_3 = _mm_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib + 1].d) * GGML_CPU_FP16_TO_FP32(y[ib + 1].d) );
+
+        const __m128i tmp_2_3 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs);
+
+        __m128i bx_2 = _mm_and_si128(lowMask, tmp_2_3);
+        __m128i by_2 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs);
+        bx_2 = _mm_sub_epi8(bx_2, off);
+        const __m128i i32_2 = mul_sum_i8_pairs(bx_2, by_2);
+
+        __m128i bx_3 = _mm_and_si128(lowMask, _mm_srli_epi64(tmp_2_3, 4));
+        __m128i by_3 = _mm_loadu_si128((const __m128i *)(y[ib + 1].qs + 16));
+        bx_3 = _mm_sub_epi8(bx_3, off);
+        const __m128i i32_3 = mul_sum_i8_pairs(bx_3, by_3);
+
+        // Convert int32_t to float
+        __m128 p0 = _mm_cvtepi32_ps(i32_0);
+        __m128 p1 = _mm_cvtepi32_ps(i32_1);
+        __m128 p2 = _mm_cvtepi32_ps(i32_2);
+        __m128 p3 = _mm_cvtepi32_ps(i32_3);
+
+        // Apply the scale
+        __m128 p0_d = _mm_mul_ps( d_0_1, p0 );
+        __m128 p1_d = _mm_mul_ps( d_0_1, p1 );
+        __m128 p2_d = _mm_mul_ps( d_2_3, p2 );
+        __m128 p3_d = _mm_mul_ps( d_2_3, p3 );
+
+        // Acummulate
+        acc_0 = _mm_add_ps(p0_d, acc_0);
+        acc_1 = _mm_add_ps(p1_d, acc_1);
+        acc_2 = _mm_add_ps(p2_d, acc_2);
+        acc_3 = _mm_add_ps(p3_d, acc_3);
+    }
+
+    sumf = hsum_float_4x4(acc_0, acc_1, acc_2, acc_3);
+
+#endif
+    for (; ib < nb; ++ib) {
+        int sumi0 = 0;
+        int sumi1 = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int v0 = (x[ib].qs[j] & 0x0F) - 8;
+            const int v1 = (x[ib].qs[j] >>   4) - 8;
+
+            sumi0 += (v0 * y[ib].qs[j]);
+            sumi1 += (v1 * y[ib].qs[j + qk/2]);
+        }
+
+        int sumi = sumi0 + sumi1;
+        sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_1;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_1 * GGML_RESTRICT x = vx;
+    const block_q8_1 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+
+#if defined(__AVX2__) || defined(__AVX__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+
+    float summs = 0;
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d);
+        const float d1 = GGML_CPU_FP16_TO_FP32(y[ib].d);
+
+        summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s);
+
+        const __m256 d0v = _mm256_set1_ps( d0 );
+        const __m256 d1v = _mm256_set1_ps( d1 );
+
+        // Compute combined scales
+        const __m256 d0d1 = _mm256_mul_ps( d0v, d1v );
+
+        // Load 16 bytes, and unpack 4 bit fields into bytes, making 32 bytes
+        const __m256i qx = bytes_from_nibbles_32(x[ib].qs);
+        const __m256i qy = _mm256_loadu_si256( (const __m256i *)y[ib].qs );
+
+        const __m256 xy = mul_sum_us8_pairs_float(qx, qy);
+
+        // Accumulate d0*d1*x*y
+#if defined(__AVX2__)
+        acc = _mm256_fmadd_ps( d0d1, xy, acc );
+#else
+        acc = _mm256_add_ps( _mm256_mul_ps( d0d1, xy ), acc );
+#endif
+    }
+
+    *s = hsum_float_8(acc) + summs;
+#else
+    UNUSED(nb);
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(ib);
+    ggml_vec_dot_q4_1_q8_1_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_MXFP4 == 0);
+    static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same");
+
+    const block_mxfp4 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_MXFP4;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __AVX2__
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+    const __m256i mone = _mm256_set1_epi16(1);
+
+    __m256 accum1 = _mm256_setzero_ps();
+    __m256 accum2 = _mm256_setzero_ps();
+
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)x[ib + 0].qs);
+        const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)x[ib + 1].qs);
+        const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)y[ib + 0].qs);
+        const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)y[ib + 1].qs);
+        const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)),
+                                              _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)));
+        const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)),
+                                              _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)));
+        const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1);
+        const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2);
+        const __m256i p_1 = _mm256_madd_epi16(p16_1, mone);
+        const __m256i p_2 = _mm256_madd_epi16(p16_2, mone);
+        const __m256 scale0 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 0].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 0].e));
+        const __m256 scale1 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 1].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 1].e));
+        accum1 = _mm256_fmadd_ps(scale0, _mm256_cvtepi32_ps(p_1), accum1);
+        accum2 = _mm256_fmadd_ps(scale1, _mm256_cvtepi32_ps(p_2), accum2);
+    }
+
+    sumf = hsum_float_8(_mm256_add_ps(accum1, accum2));
+
+#elif defined __AVX__
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs);
+        const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs);
+        const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs);
+        const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1);
+        const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs);
+        const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1);
+
+        const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b));
+        const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b));
+        const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b));
+        const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b));
+
+        const __m256 p = mul_sum_i8_quad_float(q4b_1_0, q4b_1_1, q4b_2_0, q4b_2_1, q8b_1_0, q8b_1_1, q8b_2_0, q8b_2_1);
+        const __m256 deltas = quad_mx_delta_float(x[ib].e, y[ib].d, x[ib + 1].e, y[ib + 1].d);
+        accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib].e);
+        int sumi1 = 0;
+        int sumi2 = 0;
+        for (int j = 0; j < QK_MXFP4/2; ++j) {
+            sumi1 += y[ib].qs[j +          0] * kvalues_fp4[x[ib].qs[j] & 0xf];
+            sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_fp4[x[ib].qs[j] >>  4];
+        }
+        sumf += d * (sumi1 + sumi2);
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_NVFP4 == 0);
+
+    const block_nvfp4 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_NVFP4;
+    int ib = 0;
+    float sumf = 0;
+
+#if defined(__AVX2__)
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+    const __m256i mone = _mm256_set1_epi16(1);
+
+    __m256 accum = _mm256_setzero_ps();
+    for(; ib < nb; ib++){
+
+        const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0));
+        const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16));
+
+        const __m256i q8_01 = _mm256_loadu_si256((const __m256i *)y[2*ib + 0].qs);
+        const __m256i q8_23 = _mm256_loadu_si256((const __m256i *)y[2*ib + 1].qs);
+
+        const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b));
+        const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b));
+        const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b));
+        const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b));
+
+        //reordering
+        const __m256i q4_01 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_01_lo,q4_01_hi), _mm_unpacklo_epi64(q4_01_lo,q4_01_hi));
+        const __m256i q4_23 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_23_lo,q4_23_hi),_mm_unpacklo_epi64(q4_23_lo,q4_23_hi));
+
+        const __m256i p01 = mul_add_epi8(q4_01,q8_01);
+        const __m256i p_1 = _mm256_madd_epi16(p01, mone);
+
+        const __m256i p23 = mul_add_epi8(q4_23,q8_23);
+        const __m256i p_2 = _mm256_madd_epi16(p23, mone);
+
+        const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
+        const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
+
+        const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0;
+        const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0;
+        const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1;
+        const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1;
+
+        const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0));
+        const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2));
+
+        accum = _mm256_fmadd_ps(scales01, _mm256_cvtepi32_ps(p_1), accum);
+        accum = _mm256_fmadd_ps(scales23, _mm256_cvtepi32_ps(p_2), accum);
+    }
+    sumf = hsum_float_8(accum);
+
+#elif defined(__AVX__)
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for(; ib < nb; ib++){
+
+        const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0));
+        const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16));
+
+        const __m128i q8_0 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 0));
+        const __m128i q8_1 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 16));
+        const __m128i q8_2 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 0));
+        const __m128i q8_3 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 16));
+
+        const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b));
+        const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b));
+        const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b));
+        const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b));
+
+        const __m128i q4_0 = _mm_unpacklo_epi64(q4_01_lo, q4_01_hi);
+        const __m128i q4_1 = _mm_unpackhi_epi64(q4_01_lo, q4_01_hi);
+        const __m128i q4_2 = _mm_unpacklo_epi64(q4_23_lo, q4_23_hi);
+        const __m128i q4_3 = _mm_unpackhi_epi64(q4_23_lo, q4_23_hi);
+
+        const __m128i p0_i32 = mul_sum_i8_pairs(q4_0, q8_0);
+        const __m128i p1_i32 = mul_sum_i8_pairs(q4_1, q8_1);
+        const __m128i p2_i32 = mul_sum_i8_pairs(q4_2, q8_2);
+        const __m128i p3_i32 = mul_sum_i8_pairs(q4_3, q8_3);
+
+        const __m128 p0 = _mm_cvtepi32_ps(p0_i32);
+        const __m128 p1 = _mm_cvtepi32_ps(p1_i32);
+        const __m128 p2 = _mm_cvtepi32_ps(p2_i32);
+        const __m128 p3 = _mm_cvtepi32_ps(p3_i32);
+
+        const __m256 p01 = _mm256_set_m128(p1, p0);
+        const __m256 p23 = _mm256_set_m128(p3, p2);
+
+        const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
+        const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
+
+        const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0;
+        const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0;
+        const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1;
+        const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1;
+
+        const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0));
+        const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2));
+
+        accum = _mm256_add_ps(accum, _mm256_mul_ps(p01, scales01));
+        accum = _mm256_add_ps(accum, _mm256_mul_ps(p23, scales23));
+    }
+    sumf = hsum_float_8(accum);
+
+#endif
+
+    for (;ib < nb; ++ib) {
+        for (int s_idx = 0; s_idx < 4; ++s_idx) {
+            const float d = GGML_CPU_UE4M3_TO_FP32(x[ib].d[s_idx]);
+            const int q8_block = s_idx / 2;
+            const int q8_off   = (s_idx % 2) * QK_NVFP4_SUB;
+            const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d);
+
+            int sumi_lo = 0, sumi_hi = 0;
+            for (int j = 0; j < QK_NVFP4_SUB/2; ++j) {
+                const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j];
+                sumi_lo += y[2*ib + q8_block].qs[q8_off + j +               0] * kvalues_fp4[qv & 0xf];
+                sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_fp4[qv >>  4];
+            }
+
+            sumf += dy * d * (sumi_lo + sumi_hi);
+        }
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    int ib = 0;
+
+    assert(n % qk == 0);
+    assert(qk == QK5_0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+#if defined(__AVX2__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        /* Compute combined scale for the block */
+        const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d));
+
+        __m256i qx = bytes_from_nibbles_32(x[ib].qs);
+        __m256i bxhi = bytes_from_bits_32(x[ib].qh);
+        bxhi = _mm256_andnot_si256(bxhi, _mm256_set1_epi8((char)0xF0));
+        qx = _mm256_or_si256(qx, bxhi);
+
+        __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_i8_pairs_float(qx, qy);
+
+        /* Multiply q with scale and accumulate */
+        acc = _mm256_fmadd_ps(d, q, acc);
+    }
+
+    *s = hsum_float_8(acc);
+#elif defined(__AVX__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+    __m128i mask = _mm_set1_epi8((char)0xF0);
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        /* Compute combined scale for the block */
+        const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d));
+
+        __m256i bx_0 = bytes_from_nibbles_32(x[ib].qs);
+        const __m256i bxhi = bytes_from_bits_32(x[ib].qh);
+        __m128i bxhil = _mm256_castsi256_si128(bxhi);
+        __m128i bxhih = _mm256_extractf128_si256(bxhi, 1);
+        bxhil = _mm_andnot_si128(bxhil, mask);
+        bxhih = _mm_andnot_si128(bxhih, mask);
+        __m128i bxl = _mm256_castsi256_si128(bx_0);
+        __m128i bxh = _mm256_extractf128_si256(bx_0, 1);
+        bxl = _mm_or_si128(bxl, bxhil);
+        bxh = _mm_or_si128(bxh, bxhih);
+        bx_0 = MM256_SET_M128I(bxh, bxl);
+
+        const __m256i by_0 = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_i8_pairs_float(bx_0, by_0);
+
+        /* Multiply q with scale and accumulate */
+        acc = _mm256_add_ps(_mm256_mul_ps(d, q), acc);
+    }
+
+    *s = hsum_float_8(acc);
+#else
+    UNUSED(nb);
+    UNUSED(ib);
+    UNUSED(x);
+    UNUSED(y);
+    ggml_vec_dot_q5_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_1;
+    const int nb = n / qk;
+
+    int ib = 0;
+
+    assert(n % qk == 0);
+    assert(qk == QK5_1);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_1 * GGML_RESTRICT x = vx;
+    const block_q8_1 * GGML_RESTRICT y = vy;
+
+#if defined(__AVX2__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+
+    float summs = 0.0f;
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        const __m256 dx = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d));
+
+        summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s);
+
+        __m256i qx = bytes_from_nibbles_32(x[ib].qs);
+        __m256i bxhi = bytes_from_bits_32(x[ib].qh);
+        bxhi = _mm256_and_si256(bxhi, _mm256_set1_epi8(0x10));
+        qx = _mm256_or_si256(qx, bxhi);
+
+        const __m256 dy = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d));
+        const __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_us8_pairs_float(qx, qy);
+
+        acc = _mm256_fmadd_ps(q, _mm256_mul_ps(dx, dy), acc);
+    }
+
+    *s = hsum_float_8(acc) + summs;
+#elif defined(__AVX__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+    __m128i mask = _mm_set1_epi8(0x10);
+
+    float summs = 0.0f;
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        const __m256 dx = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d));
+
+        summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s);
+
+        __m256i bx_0 = bytes_from_nibbles_32(x[ib].qs);
+        const __m256i bxhi = bytes_from_bits_32(x[ib].qh);
+        __m128i bxhil = _mm256_castsi256_si128(bxhi);
+        __m128i bxhih = _mm256_extractf128_si256(bxhi, 1);
+        bxhil = _mm_and_si128(bxhil, mask);
+        bxhih = _mm_and_si128(bxhih, mask);
+        __m128i bxl = _mm256_castsi256_si128(bx_0);
+        __m128i bxh = _mm256_extractf128_si256(bx_0, 1);
+        bxl = _mm_or_si128(bxl, bxhil);
+        bxh = _mm_or_si128(bxh, bxhih);
+        bx_0 = MM256_SET_M128I(bxh, bxl);
+
+        const __m256 dy = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d));
+        const __m256i by_0 = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_us8_pairs_float(bx_0, by_0);
+
+        acc = _mm256_add_ps(_mm256_mul_ps(q, _mm256_mul_ps(dx, dy)), acc);
+    }
+
+    *s = hsum_float_8(acc) + summs;
+#else
+    UNUSED(nb);
+    UNUSED(ib);
+    UNUSED(x);
+    UNUSED(y);
+    ggml_vec_dot_q5_1_q8_1_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q8_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined(__AVX2__)
+    // Initialize accumulator with zeros
+    __m256 acc = _mm256_setzero_ps();
+
+    // Main loop
+    for (; ib < nb; ++ib) {
+        // Compute combined scale for the block
+        const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d));
+        __m256i qx = _mm256_loadu_si256((const __m256i *)x[ib].qs);
+        __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+
+        const __m256 q = mul_sum_i8_pairs_float(qx, qy);
+
+        // Multiply q with scale and accumulate
+        acc = _mm256_fmadd_ps( d, q, acc );
+    }
+
+    sumf = hsum_float_8(acc);
+#elif defined(__AVX__)
+    __m256 accum = _mm256_setzero_ps();
+
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i qx_1_0 = _mm_loadu_si128((const __m128i *)x[ib].qs);
+        const __m128i qx_1_1 = _mm_loadu_si128((const __m128i *)x[ib].qs + 1);
+        const __m128i qx_2_0 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs);
+        const __m128i qx_2_1 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs + 1);
+        const __m128i qy_1_0 = _mm_loadu_si128((const __m128i *)y[ib].qs);
+        const __m128i qy_1_1 = _mm_loadu_si128((const __m128i *)y[ib].qs + 1);
+        const __m128i qy_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs);
+        const __m128i qy_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1);
+
+        const __m256 p = mul_sum_i8_quad_float(qx_1_0, qx_1_1, qx_2_0, qx_2_1, qy_1_0, qy_1_1, qy_2_0, qy_2_1);
+        const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d);
+        accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+#endif
+    for (; ib < nb; ++ib) {
+        int sumi = 0;
+
+        for (int j = 0; j < qk; j++) {
+            sumi += x[ib].qs[j]*y[ib].qs[j];
+        }
+
+        sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d));
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_tq1_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_tq1_0 * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+    __m256 sumf = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+        // 16-bit sums
+        __m256i sumi0 = _mm256_setzero_si256();
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+
+        // first 32 bytes of 5 elements
+        {
+            __m256i qx0 = _mm256_loadu_si256((const __m256i *) (x[i].qs));
+            // 8-bit multiplies with shifts, masks and adds
+            __m256i qx1 = _mm256_add_epi8(qx0, _mm256_add_epi8(qx0, qx0)); // 1 * 3
+            __m256i qx2 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx0, 3), _mm256_set1_epi8(-8)), qx0); // 1 * 9
+            __m256i qx3 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx1, 3), _mm256_set1_epi8(-8)), qx1); // 3 * 9
+            __m256i qx4 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx2, 3), _mm256_set1_epi8(-8)), qx2); // 9 * 9
+
+            // TODO: can _mm256_mulhi_epu16 be faster even if 16-bits?
+
+            // Cancel the +1 from avg so that it behaves like a halving add
+            qx0 = _mm256_subs_epu8(qx0, _mm256_set1_epi8(1));
+            qx1 = _mm256_subs_epu8(qx1, _mm256_set1_epi8(1));
+            qx2 = _mm256_subs_epu8(qx2, _mm256_set1_epi8(1));
+            qx3 = _mm256_subs_epu8(qx3, _mm256_set1_epi8(1));
+            qx4 = _mm256_subs_epu8(qx4, _mm256_set1_epi8(1));
+            // Multiply by 3 and get the top 2 bits
+            qx0 = _mm256_avg_epu8(qx0, _mm256_avg_epu8(qx0, _mm256_setzero_si256()));
+            qx1 = _mm256_avg_epu8(qx1, _mm256_avg_epu8(qx1, _mm256_setzero_si256()));
+            qx2 = _mm256_avg_epu8(qx2, _mm256_avg_epu8(qx2, _mm256_setzero_si256()));
+            qx3 = _mm256_avg_epu8(qx3, _mm256_avg_epu8(qx3, _mm256_setzero_si256()));
+            qx4 = _mm256_avg_epu8(qx4, _mm256_avg_epu8(qx4, _mm256_setzero_si256()));
+            qx0 = _mm256_and_si256(_mm256_srli_epi16(qx0, 6), _mm256_set1_epi8(3));
+            qx1 = _mm256_and_si256(_mm256_srli_epi16(qx1, 6), _mm256_set1_epi8(3));
+            qx2 = _mm256_and_si256(_mm256_srli_epi16(qx2, 6), _mm256_set1_epi8(3));
+            qx3 = _mm256_and_si256(_mm256_srli_epi16(qx3, 6), _mm256_set1_epi8(3));
+            qx4 = _mm256_and_si256(_mm256_srli_epi16(qx4, 6), _mm256_set1_epi8(3));
+
+            const __m256i qy0 = _mm256_loadu_si256((const __m256i *) (y[i].qs +   0));
+            const __m256i qy1 = _mm256_loadu_si256((const __m256i *) (y[i].qs +  32));
+            const __m256i qy2 = _mm256_loadu_si256((const __m256i *) (y[i].qs +  64));
+            const __m256i qy3 = _mm256_loadu_si256((const __m256i *) (y[i].qs +  96));
+            const __m256i qy4 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 128));
+
+            qx0 = _mm256_maddubs_epi16(qx0, qy0);
+            qx1 = _mm256_maddubs_epi16(qx1, qy1);
+            qx2 = _mm256_maddubs_epi16(qx2, qy2);
+            qx3 = _mm256_maddubs_epi16(qx3, qy3);
+            qx4 = _mm256_maddubs_epi16(qx4, qy4);
+
+            sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(qx0, qx1));
+            sumi1 = _mm256_add_epi16(sumi1, _mm256_add_epi16(qx2, qx3));
+            sumi2 = _mm256_add_epi16(sumi2, qx4);
+        }
+
+        // last 16 bytes of 5-element, along with the 4 bytes of 4 elements
+        {
+            __m128i qx0 = _mm_loadu_si128((const __m128i *) (x[i].qs + 32));
+            uint32_t qh;
+            memcpy(&qh, x[i].qh, sizeof(qh)); // potentially unaligned
+            __m256i qx5_l = _mm256_cvtepu8_epi16(_mm_set1_epi32(qh));
+            __m128i qx1 = _mm_add_epi8(qx0, _mm_add_epi8(qx0, qx0)); // 1 * 3
+            __m128i qx2 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx0, 3), _mm_set1_epi8(-8)), qx0); // 1 * 9
+            __m128i qx3 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx1, 3), _mm_set1_epi8(-8)), qx1); // 3 * 9
+            __m128i qx4 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx2, 3), _mm_set1_epi8(-8)), qx2); // 9 * 9
+            __m256i qx01 = MM256_SET_M128I(qx1, qx0);
+            __m256i qx23 = MM256_SET_M128I(qx3, qx2);
+
+            // avx2 does not have 8-bit multiplies, so 16-bit it is.
+            qx5_l = _mm256_mullo_epi16(qx5_l, _mm256_set_epi16(27, 27, 27, 27, 9, 9, 9, 9, 3, 3, 3, 3, 1, 1, 1, 1));
+            qx5_l = _mm256_and_si256(qx5_l, _mm256_set1_epi16(0xFF));
+            __m128i qx5 = _mm_packus_epi16(_mm256_castsi256_si128(qx5_l), _mm256_extracti128_si256(qx5_l, 1));
+
+            __m256i qx45 = MM256_SET_M128I(qx5, qx4);
+
+            // Cancel the +1 from avg so that it behaves like a halving add
+            qx01 = _mm256_subs_epu8(qx01, _mm256_set1_epi8(1));
+            qx23 = _mm256_subs_epu8(qx23, _mm256_set1_epi8(1));
+            qx45 = _mm256_subs_epu8(qx45, _mm256_set1_epi8(1));
+            // Multiply by 3 and get the top 2 bits
+            qx01 = _mm256_avg_epu8(qx01, _mm256_avg_epu8(qx01, _mm256_setzero_si256()));
+            qx23 = _mm256_avg_epu8(qx23, _mm256_avg_epu8(qx23, _mm256_setzero_si256()));
+            qx45 = _mm256_avg_epu8(qx45, _mm256_avg_epu8(qx45, _mm256_setzero_si256()));
+            qx01 = _mm256_and_si256(_mm256_srli_epi16(qx01, 6), _mm256_set1_epi8(3));
+            qx23 = _mm256_and_si256(_mm256_srli_epi16(qx23, 6), _mm256_set1_epi8(3));
+            qx45 = _mm256_and_si256(_mm256_srli_epi16(qx45, 6), _mm256_set1_epi8(3));
+
+            const __m256i qy01 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 160));
+            const __m256i qy23 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 192));
+            const __m256i qy45 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 224));
+
+            qx01 = _mm256_maddubs_epi16(qx01, qy01);
+            qx23 = _mm256_maddubs_epi16(qx23, qy23);
+            qx45 = _mm256_maddubs_epi16(qx45, qy45);
+
+            sumi0 = _mm256_add_epi16(sumi0, qx01);
+            sumi1 = _mm256_add_epi16(sumi1, qx23);
+            sumi2 = _mm256_add_epi16(sumi2, qx45);
+        }
+
+        const __m256i ysum = _mm256_loadu_si256((const __m256i *) y[i].bsums);
+        const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d));
+
+        sumi0 = _mm256_sub_epi16(sumi0, ysum);
+        sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(sumi1, sumi2));
+        sumi0 = _mm256_madd_epi16(sumi0, _mm256_set1_epi16(1));
+
+        sumf = _mm256_add_ps(_mm256_mul_ps(_mm256_cvtepi32_ps(sumi0), d), sumf);
+    }
+
+    *s = hsum_float_8(sumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_tq1_0_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_tq2_0 * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+    __m256 sumf = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+        // 16-bit sums, because 256*127 still fits
+        __m256i sumi0 = _mm256_setzero_si256();
+        __m256i sumi1 = _mm256_setzero_si256();
+
+        for (size_t j = 0; j < sizeof(x->qs); j += 32) {
+            __m256i qx0 = _mm256_loadu_si256((const __m256i *) (x[i].qs + j));
+            __m256i qx1 = _mm256_srli_epi16(qx0, 2);
+            __m256i qx2 = _mm256_srli_epi16(qx0, 4);
+            __m256i qx3 = _mm256_srli_epi16(qx0, 6);
+
+            // 0, 1, 2 (should not be 3)
+            qx0 = _mm256_and_si256(qx0, _mm256_set1_epi8(3));
+            qx1 = _mm256_and_si256(qx1, _mm256_set1_epi8(3));
+            qx2 = _mm256_and_si256(qx2, _mm256_set1_epi8(3));
+            qx3 = _mm256_and_si256(qx3, _mm256_set1_epi8(3));
+
+            const __m256i qy0 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 +  0));
+            const __m256i qy1 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 32));
+            const __m256i qy2 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 64));
+            const __m256i qy3 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 96));
+
+            qx0 = _mm256_maddubs_epi16(qx0, qy0);
+            qx1 = _mm256_maddubs_epi16(qx1, qy1);
+            qx2 = _mm256_maddubs_epi16(qx2, qy2);
+            qx3 = _mm256_maddubs_epi16(qx3, qy3);
+
+            sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(qx0, qx1));
+            sumi1 = _mm256_add_epi16(sumi1, _mm256_add_epi16(qx2, qx3));
+        }
+
+        const __m256i ysum = _mm256_loadu_si256((const __m256i *) y[i].bsums);
+        const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d));
+
+        sumi0 = _mm256_add_epi16(sumi0, sumi1);
+        sumi0 = _mm256_sub_epi16(sumi0, ysum);
+        sumi0 = _mm256_madd_epi16(sumi0, _mm256_set1_epi16(1));
+
+        sumf = _mm256_add_ps(_mm256_mul_ps(_mm256_cvtepi32_ps(sumi0), d), sumf);
+    }
+
+    *s = hsum_float_8(sumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_tq2_0_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined __AVX2__
+
+    const __m256i m3 = _mm256_set1_epi8(3);
+    const __m128i m4 = _mm_set1_epi8(0xF);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales);
+        const __m128i scales8 = _mm_and_si128(mins_and_scales, m4);
+        const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4);
+        const __m256i mins = _mm256_cvtepi8_epi16(mins8);
+        const __m256i prod = _mm256_madd_epi16(mins, _mm256_loadu_si256((const __m256i*)y[i].bsums));
+
+        acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&dmin), _mm256_cvtepi32_ps(prod), acc);
+
+        const __m256i all_scales = _mm256_cvtepi8_epi16(scales8);
+        const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0);
+        const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1);
+        const __m256i scales[2] = {MM256_SET_M128I(l_scales, l_scales), MM256_SET_M128I(h_scales, h_scales)};
+
+        __m256i sumi = _mm256_setzero_si256();
+
+        for (int j = 0; j < QK_K/128; ++j) {
+
+            const __m256i q2bits = _mm256_loadu_si256((const __m256i*)q2); q2 += 32;
+
+            const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            const __m256i q2_0 = _mm256_and_si256(q2bits, m3);
+            const __m256i q2_1 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), m3);
+            const __m256i q2_2 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), m3);
+            const __m256i q2_3 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), m3);
+
+            __m256i p0 = _mm256_maddubs_epi16(q2_0, q8_0);
+            __m256i p1 = _mm256_maddubs_epi16(q2_1, q8_1);
+            __m256i p2 = _mm256_maddubs_epi16(q2_2, q8_2);
+            __m256i p3 = _mm256_maddubs_epi16(q2_3, q8_3);
+
+            p0 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(0)), p0);
+            p1 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(1)), p1);
+            p2 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(2)), p2);
+            p3 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(3)), p3);
+
+            p0 = _mm256_add_epi32(p0, p1);
+            p2 = _mm256_add_epi32(p2, p3);
+
+            sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p0, p2));
+        }
+
+        acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc);
+
+    }
+
+    *s = hsum_float_8(acc);
+
+#elif defined __AVX__
+
+    const __m128i m3 = _mm_set1_epi8(0x3);
+    const __m128i m4 = _mm_set1_epi8(0xF);
+    const __m128i m2 = _mm_set1_epi8(0x2);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        // load mins and scales from block_q2_K.scales[QK_K/16]
+        const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales);
+        const __m128i scales16 = _mm_and_si128(mins_and_scales, m4);
+        const __m128i mins16 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4);
+        const __m128i mins_0 = _mm_cvtepi8_epi16(mins16);
+        const __m128i mins_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(mins16, mins16));
+
+        // summs = y[i].bsums * (x[i].scales >> 4) in 16bits*8*2 to 32bits*4*2
+        const __m128i summs_0 = _mm_madd_epi16(mins_0, _mm_loadu_si128((const __m128i*)&y[i].bsums[0]));
+        const __m128i summs_1 = _mm_madd_epi16(mins_1, _mm_loadu_si128((const __m128i*)&y[i].bsums[8]));
+
+        // sumf += -dmin * summs in 32bits*8
+        acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&dmin), _mm256_cvtepi32_ps(MM256_SET_M128I(summs_1, summs_0))), acc);
+
+        const __m128i scales_0 = _mm_cvtepi8_epi16(scales16);
+        const __m128i scales_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(scales16, scales16));
+        const __m128i scales[2] = { scales_0, scales_1 };
+
+        __m128i sumi_0 = _mm_setzero_si128();
+        __m128i sumi_1 = _mm_setzero_si128();
+
+        for (int j = 0; j < QK_K/128; ++j) {
+
+            // load Q8 quants int8*16*8 from block_q8_K.qs[QK_K]
+            const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+
+            // load 2bits*16*8 from block_q2_K.qs[QK_K/4]
+            __m128i q2bits = _mm_loadu_si128((const __m128i*)q2); q2 += 16;
+            const __m128i q2_0 = _mm_and_si128(q2bits, m3);
+            const __m128i q2_2 = _mm_and_si128(_mm_srli_epi16(q2bits, 2), m3);
+            const __m128i q2_4 = _mm_and_si128(_mm_srli_epi16(q2bits, 4), m3);
+            const __m128i q2_6 = _mm_and_si128(_mm_srli_epi16(q2bits, 6), m3);
+            q2bits = _mm_loadu_si128((const __m128i*)q2); q2 += 16;
+            const __m128i q2_1 = _mm_and_si128(q2bits, m3);
+            const __m128i q2_3 = _mm_and_si128(_mm_srli_epi16(q2bits, 2), m3);
+            const __m128i q2_5 = _mm_and_si128(_mm_srli_epi16(q2bits, 4), m3);
+            const __m128i q2_7 = _mm_and_si128(_mm_srli_epi16(q2bits, 6), m3);
+
+            // isuml = q8[l] * ((q2[l] >> shift) & 3) in 8bits*16*8 to 16bits*8*8
+            __m128i p0 = _mm_maddubs_epi16(q2_0, q8_0);
+            __m128i p1 = _mm_maddubs_epi16(q2_1, q8_1);
+            __m128i p2 = _mm_maddubs_epi16(q2_2, q8_2);
+            __m128i p3 = _mm_maddubs_epi16(q2_3, q8_3);
+            __m128i p4 = _mm_maddubs_epi16(q2_4, q8_4);
+            __m128i p5 = _mm_maddubs_epi16(q2_5, q8_5);
+            __m128i p6 = _mm_maddubs_epi16(q2_6, q8_6);
+            __m128i p7 = _mm_maddubs_epi16(q2_7, q8_7);
+
+            // isum += (x[i].scales[is++] & 0xF) * isuml in 16bits*8*8 to 32bits*4*8
+            __m128i shuffle = _mm_set1_epi16(0x0100);
+            p0 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p0);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p1 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p1);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p2 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p2);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p3 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p3);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p4 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p4);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p5 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p5);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p6 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p6);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p7 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p7);
+
+            p0 = _mm_add_epi32(p0, p1);
+            p2 = _mm_add_epi32(p2, p3);
+            p4 = _mm_add_epi32(p4, p5);
+            p6 = _mm_add_epi32(p6, p7);
+
+            // isum in 32bits*4*2
+            sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p0, p2));
+            sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p4, p6));
+        }
+
+        // sumf += dall * isum - dmin * summs in 32bits
+        __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0);
+        acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&dall), _mm256_cvtepi32_ps(sumi)), acc);
+    }
+
+    *s = hsum_float_8(acc);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_q2_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const uint32_t kmask1 = 0x03030303;
+    const uint32_t kmask2 = 0x0f0f0f0f;
+
+    const block_q3_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined __AVX2__
+
+    const __m256i m3 = _mm256_set1_epi8(3);
+    const __m256i mone = _mm256_set1_epi8(1);
+    const __m128i m32 = _mm_set1_epi8(32);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    uint32_t aux[3];
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        // Set up scales
+        memcpy(aux, x[i].scales, 12);
+        __m128i scales128 = _mm_set_epi32(
+                ((aux[1] >> 4) & kmask2) | (((aux[2] >> 6) & kmask1) << 4),
+                ((aux[0] >> 4) & kmask2) | (((aux[2] >> 4) & kmask1) << 4),
+                (aux[1] & kmask2) | (((aux[2] >> 2) & kmask1) << 4),
+                (aux[0] & kmask2) | (((aux[2] >> 0) & kmask1) << 4));
+        scales128 = _mm_sub_epi8(scales128, m32);
+        const __m256i all_scales = _mm256_cvtepi8_epi16(scales128);
+        const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0);
+        const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1);
+        const __m256i scales[2] = {MM256_SET_M128I(l_scales, l_scales), MM256_SET_M128I(h_scales, h_scales)};
+
+        // high bit
+        const __m256i hbits = _mm256_loadu_si256((const __m256i*)x[i].hmask);
+
+        // integer accumulator
+        __m256i sumi = _mm256_setzero_si256();
+
+        int bit = 0;
+        int is  = 0;
+
+        for (int j = 0; j < QK_K/128; ++j) {
+            // load low 2 bits
+            const __m256i q3bits = _mm256_loadu_si256((const __m256i*)q3); q3 += 32;
+
+            // prepare low and high bits
+            const __m256i q3l_0 = _mm256_and_si256(q3bits, m3);
+            const __m256i q3h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2);
+            ++bit;
+
+            const __m256i q3l_1 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 2), m3);
+            const __m256i q3h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2);
+            ++bit;
+
+            const __m256i q3l_2 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 4), m3);
+            const __m256i q3h_2 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2);
+            ++bit;
+
+            const __m256i q3l_3 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 6), m3);
+            const __m256i q3h_3 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2);
+            ++bit;
+
+            // load Q8 quants
+            const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            // Dot product: we multiply the 2 low bits and 1 high bit part separately, so we can use _mm256_maddubs_epi16,
+            // and then subtract. The high bit part has the 2 already subtracted (and so, it is zero if the high bit was not set,
+            // and 2 if the high bit was set)
+            __m256i q8s_0 = _mm256_maddubs_epi16(q3h_0, q8_0);
+            __m256i q8s_1 = _mm256_maddubs_epi16(q3h_1, q8_1);
+            __m256i q8s_2 = _mm256_maddubs_epi16(q3h_2, q8_2);
+            __m256i q8s_3 = _mm256_maddubs_epi16(q3h_3, q8_3);
+
+            __m256i p16_0 = _mm256_maddubs_epi16(q3l_0, q8_0);
+            __m256i p16_1 = _mm256_maddubs_epi16(q3l_1, q8_1);
+            __m256i p16_2 = _mm256_maddubs_epi16(q3l_2, q8_2);
+            __m256i p16_3 = _mm256_maddubs_epi16(q3l_3, q8_3);
+
+            p16_0 = _mm256_sub_epi16(p16_0, q8s_0);
+            p16_1 = _mm256_sub_epi16(p16_1, q8s_1);
+            p16_2 = _mm256_sub_epi16(p16_2, q8s_2);
+            p16_3 = _mm256_sub_epi16(p16_3, q8s_3);
+
+            // multiply with scales
+            p16_0 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 0)), p16_0);
+            p16_1 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 1)), p16_1);
+            p16_2 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 2)), p16_2);
+            p16_3 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 3)), p16_3);
+
+            // accumulate
+            p16_0 = _mm256_add_epi32(p16_0, p16_1);
+            p16_2 = _mm256_add_epi32(p16_2, p16_3);
+            sumi  = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_2));
+
+        }
+
+        // multiply with block scale and accumulate
+        acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc);
+
+    }
+
+    *s = hsum_float_8(acc);
+
+#elif defined __AVX__
+
+    const __m128i m3 = _mm_set1_epi8(3);
+    const __m128i mone = _mm_set1_epi8(1);
+    const __m128i m32 = _mm_set1_epi8(32);
+    const __m128i m2 = _mm_set1_epi8(2);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    const uint32_t *aux;
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        // Set up scales
+        aux = (const uint32_t *)x[i].scales;
+        __m128i scales128 = _mm_set_epi32(
+                ((aux[1] >> 4) & kmask2) | (((aux[2] >> 6) & kmask1) << 4),
+                ((aux[0] >> 4) & kmask2) | (((aux[2] >> 4) & kmask1) << 4),
+                (aux[1] & kmask2) | (((aux[2] >> 2) & kmask1) << 4),
+                (aux[0] & kmask2) | (((aux[2] >> 0) & kmask1) << 4));
+        scales128 = _mm_sub_epi8(scales128, m32);
+        const __m128i scales_0 = _mm_cvtepi8_epi16(scales128);
+        const __m128i scales_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(scales128, scales128));
+        const __m128i scales[2] = { scales_0, scales_1 };
+
+        // high bit *128*2 from block_q3_K.hmask[QK_K/8]
+        const __m128i hbits_0 = _mm_loadu_si128((const __m128i*)&x[i].hmask[0]);
+        const __m128i hbits_1 = _mm_loadu_si128((const __m128i*)&x[i].hmask[16]);
+
+        // integer accumulator
+        __m128i sumi_0 = _mm_setzero_si128();
+        __m128i sumi_1 = _mm_setzero_si128();
+
+        for (int j = 0; j < QK_K/128; ++j) {
+            // load low 2 bits *64*2 from block_q3_K.qs[QK_K/4]
+            const __m128i q3bits_0 = _mm_loadu_si128((const __m128i*)q3); q3 += 16;
+            const __m128i q3bits_1 = _mm_loadu_si128((const __m128i*)q3); q3 += 16;
+
+            // prepare low and high bits
+            const int bit = j << 2;
+
+            const __m128i q3l_0 = _mm_and_si128(q3bits_0, m3);
+            const __m128i q3l_1 = _mm_and_si128(q3bits_1, m3);
+            const __m128i q3h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit)), bit), 2);
+            const __m128i q3h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit)), bit), 2);
+
+            const __m128i q3l_2 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 2), m3);
+            const __m128i q3l_3 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 2), m3);
+            const __m128i q3h_2 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+1)), bit+1), 2);
+            const __m128i q3h_3 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+1)), bit+1), 2);
+
+            const __m128i q3l_4 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 4), m3);
+            const __m128i q3l_5 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 4), m3);
+            const __m128i q3h_4 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+2)), bit+2), 2);
+            const __m128i q3h_5 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+2)), bit+2), 2);
+
+            const __m128i q3l_6 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 6), m3);
+            const __m128i q3l_7 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 6), m3);
+            const __m128i q3h_6 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+3)), bit+3), 2);
+            const __m128i q3h_7 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+3)), bit+3), 2);
+
+            // load Q8 quants from block_q8_K.qs[QK_K]
+            const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+
+            // Dot product: we multiply the 2 low bits and 1 high bit part separately, so we can use _mm256_maddubs_epi16,
+            // and then subtract. The high bit part has the 2 already subtracted (and so, it is zero if the high bit was not set,
+            // and 2 if the high bit was set)
+            __m128i q8s_0 = _mm_maddubs_epi16(q3h_0, q8_0);
+            __m128i q8s_1 = _mm_maddubs_epi16(q3h_1, q8_1);
+            __m128i q8s_2 = _mm_maddubs_epi16(q3h_2, q8_2);
+            __m128i q8s_3 = _mm_maddubs_epi16(q3h_3, q8_3);
+            __m128i q8s_4 = _mm_maddubs_epi16(q3h_4, q8_4);
+            __m128i q8s_5 = _mm_maddubs_epi16(q3h_5, q8_5);
+            __m128i q8s_6 = _mm_maddubs_epi16(q3h_6, q8_6);
+            __m128i q8s_7 = _mm_maddubs_epi16(q3h_7, q8_7);
+
+            __m128i p16_0 = _mm_maddubs_epi16(q3l_0, q8_0);
+            __m128i p16_1 = _mm_maddubs_epi16(q3l_1, q8_1);
+            __m128i p16_2 = _mm_maddubs_epi16(q3l_2, q8_2);
+            __m128i p16_3 = _mm_maddubs_epi16(q3l_3, q8_3);
+            __m128i p16_4 = _mm_maddubs_epi16(q3l_4, q8_4);
+            __m128i p16_5 = _mm_maddubs_epi16(q3l_5, q8_5);
+            __m128i p16_6 = _mm_maddubs_epi16(q3l_6, q8_6);
+            __m128i p16_7 = _mm_maddubs_epi16(q3l_7, q8_7);
+
+            p16_0 = _mm_sub_epi16(p16_0, q8s_0);
+            p16_1 = _mm_sub_epi16(p16_1, q8s_1);
+            p16_2 = _mm_sub_epi16(p16_2, q8s_2);
+            p16_3 = _mm_sub_epi16(p16_3, q8s_3);
+            p16_4 = _mm_sub_epi16(p16_4, q8s_4);
+            p16_5 = _mm_sub_epi16(p16_5, q8s_5);
+            p16_6 = _mm_sub_epi16(p16_6, q8s_6);
+            p16_7 = _mm_sub_epi16(p16_7, q8s_7);
+
+            // multiply with scales
+            __m128i shuffle = _mm_set1_epi16(0x0100);
+            p16_0 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_0);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_1 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_1);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_2 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_2);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_3 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_3);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_4 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_4);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_5 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_5);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_6 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_6);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            p16_7 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_7);
+
+            // accumulate
+            p16_0 = _mm_add_epi32(p16_0, p16_1);
+            p16_2 = _mm_add_epi32(p16_2, p16_3);
+            p16_4 = _mm_add_epi32(p16_4, p16_5);
+            p16_6 = _mm_add_epi32(p16_6, p16_7);
+            sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2));
+            sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_4, p16_6));
+
+        }
+
+        // multiply with block scale and accumulate
+        __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0);
+        acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi)), acc);
+
+    }
+
+    *s = hsum_float_8(acc);
+
+#else
+    UNUSED(kmask1);
+    UNUSED(kmask2);
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_q3_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    static const uint32_t kmask1 = 0x3f3f3f3f;
+    static const uint32_t kmask2 = 0x0f0f0f0f;
+    static const uint32_t kmask3 = 0x03030303;
+
+    uint32_t utmp[4];
+
+#if defined __AVX2__
+
+    const __m256i m4 = _mm256_set1_epi8(0xF);
+
+    __m256 acc = _mm256_setzero_ps();
+    __m128 acc_m = _mm_setzero_ps();
+
+   for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        const uint8_t * GGML_RESTRICT q4 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]));
+
+        const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums);
+        const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
+        const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s);
+        acc_m = _mm_fmadd_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod), acc_m);
+
+        const __m128i sc128  = _mm256_extracti128_si256(mins_and_scales, 0);
+        const __m256i scales = MM256_SET_M128I(sc128, sc128);
+
+        __m256i sumi = _mm256_setzero_si256();
+
+        for (int j = 0; j < QK_K/64; ++j) {
+
+            const __m256i scale_l = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+0));
+            const __m256i scale_h = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+1));
+
+            const __m256i q4bits = _mm256_loadu_si256((const __m256i*)q4); q4 += 32;
+            const __m256i q4l = _mm256_and_si256(q4bits, m4);
+            const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), m4);
+
+            const __m256i q8l = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            __m256i p16l = _mm256_maddubs_epi16(q4l, q8l);
+            p16l = _mm256_madd_epi16(scale_l, p16l);
+
+            const __m256i q8h = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            __m256i p16h = _mm256_maddubs_epi16(q4h, q8h);
+            p16h = _mm256_madd_epi16(scale_h, p16h);
+            const __m256i sumj = _mm256_add_epi32(p16l, p16h);
+
+            sumi = _mm256_add_epi32(sumi, sumj);
+        }
+
+        __m256 vd = _mm256_set1_ps(d);
+        acc = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi), acc);
+
+    }
+
+    acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m));
+    acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m));
+
+    *s = hsum_float_8(acc) + _mm_cvtss_f32(acc_m);
+
+#elif defined __AVX__
+
+    const __m128i m4 = _mm_set1_epi8(0xF);
+    const __m128i m2 = _mm_set1_epi8(0x2);
+
+    __m256 acc = _mm256_setzero_ps();
+    __m128 acc_m = _mm_setzero_ps();
+
+   for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * GGML_RESTRICT q4 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        const __m128i utmps = _mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]);
+        const __m128i scales = _mm_cvtepu8_epi16(utmps);
+        const __m128i mins = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(utmps, utmps));
+
+        const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)&y[i].bsums[0]);
+        const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)&y[i].bsums[8]);
+        const __m128i q8s = _mm_hadd_epi16(q8sums_0, q8sums_1);
+        const __m128i prod = _mm_madd_epi16(mins, q8s);
+        acc_m = _mm_add_ps(_mm_mul_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod)), acc_m);
+
+        __m128i sumi_0 = _mm_setzero_si128();
+        __m128i sumi_1 = _mm_setzero_si128();
+
+        __m128i shuffle = _mm_set1_epi16(0x0100);
+        for (int j = 0; j < QK_K/64; ++j) {
+
+            const __m128i scale_l = _mm_shuffle_epi8(scales, shuffle);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            const __m128i scale_h = _mm_shuffle_epi8(scales, shuffle);
+            shuffle = _mm_add_epi16(shuffle, m2);
+
+            __m128i q4bits = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+            const __m128i q4l_0 = _mm_and_si128(q4bits, m4);
+            const __m128i q4h_0 = _mm_and_si128(_mm_srli_epi16(q4bits, 4), m4);
+            q4bits = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+            const __m128i q4l_1 = _mm_and_si128(q4bits, m4);
+            const __m128i q4h_1 = _mm_and_si128(_mm_srli_epi16(q4bits, 4), m4);
+
+            const __m128i q8l_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            __m128i p16l = _mm_maddubs_epi16(q4l_0, q8l_0);
+            p16l = _mm_madd_epi16(scale_l, p16l);
+            sumi_0 = _mm_add_epi32(sumi_0, p16l);
+            const __m128i q8l_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            p16l = _mm_maddubs_epi16(q4l_1, q8l_1);
+            p16l = _mm_madd_epi16(scale_l, p16l);
+            sumi_1 = _mm_add_epi32(sumi_1, p16l);
+
+            const __m128i q8h_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            __m128i p16h = _mm_maddubs_epi16(q4h_0, q8h_0);
+            p16h = _mm_madd_epi16(scale_h, p16h);
+            sumi_0 = _mm_add_epi32(sumi_0, p16h);
+            const __m128i q8h_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            p16h = _mm_maddubs_epi16(q4h_1, q8h_1);
+            p16h = _mm_madd_epi16(scale_h, p16h);
+            sumi_1 = _mm_add_epi32(sumi_1, p16h);
+
+        }
+
+        __m256 vd = _mm256_set1_ps(d);
+        __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0);
+        acc = _mm256_add_ps(_mm256_mul_ps(vd, _mm256_cvtepi32_ps(sumi)), acc);
+
+    }
+
+    acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m));
+    acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m));
+
+    *s = hsum_float_8(acc) + _mm_cvtss_f32(acc_m);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    UNUSED(kmask1);
+    UNUSED(kmask2);
+    UNUSED(kmask3);
+    UNUSED(utmp);
+    ggml_vec_dot_q4_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy,  size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    static const uint32_t kmask1 = 0x3f3f3f3f;
+    static const uint32_t kmask2 = 0x0f0f0f0f;
+    static const uint32_t kmask3 = 0x03030303;
+
+    uint32_t utmp[4];
+
+#if defined __AVX2__
+
+    const __m256i m4 = _mm256_set1_epi8(0xF);
+    const __m128i mzero = _mm_setzero_si128();
+    const __m256i mone  = _mm256_set1_epi8(1);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    float summs = 0.f;
+
+    for (int i = 0; i < nb; ++i) {
+        const uint8_t * GGML_RESTRICT q5 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]));
+
+        const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums);
+        const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
+        const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s);
+        const __m128i hsum = _mm_hadd_epi32(_mm_hadd_epi32(prod, mzero), mzero);
+        summs += dmin * _mm_extract_epi32(hsum, 0);
+
+        const __m128i sc128  = _mm256_extracti128_si256(mins_and_scales, 0);
+        const __m256i scales = MM256_SET_M128I(sc128, sc128);
+
+        const __m256i hbits = _mm256_loadu_si256((const __m256i*)x[i].qh);
+        __m256i hmask = mone;
+
+        __m256i sumi = _mm256_setzero_si256();
+
+        int bit = 0;
+
+        for (int j = 0; j < QK_K/64; ++j) {
+
+            const __m256i scale_0 = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+0));
+            const __m256i scale_1 = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+1));
+
+            const __m256i q5bits = _mm256_loadu_si256((const __m256i*)q5); q5 += 32;
+
+            const __m256i q5l_0 = _mm256_and_si256(q5bits, m4);
+            const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), bit++), 4);
+            const __m256i q5_0  = _mm256_add_epi8(q5l_0, q5h_0);
+            hmask = _mm256_slli_epi16(hmask, 1);
+
+            const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), m4);
+            const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), bit++), 4);
+            const __m256i q5_1  = _mm256_add_epi8(q5l_1, q5h_1);
+            hmask = _mm256_slli_epi16(hmask, 1);
+
+            const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            __m256i p16_0 = _mm256_maddubs_epi16(q5_0, q8_0);
+            __m256i p16_1 = _mm256_maddubs_epi16(q5_1, q8_1);
+
+            p16_0 = _mm256_madd_epi16(scale_0, p16_0);
+            p16_1 = _mm256_madd_epi16(scale_1, p16_1);
+
+            sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1));
+
+        }
+
+        __m256 vd = _mm256_set1_ps(d);
+        acc = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi), acc);
+
+    }
+
+    *s = hsum_float_8(acc) + summs;
+
+#elif defined __AVX__
+
+    const __m128i m4 = _mm_set1_epi8(0xF);
+    const __m128i mzero = _mm_setzero_si128();
+    const __m128i mone  = _mm_set1_epi8(1);
+    const __m128i m2 = _mm_set1_epi8(2);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    float summs = 0.f;
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * GGML_RESTRICT q5 = x[i].qs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        const __m128i utmps = _mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]);
+        const __m128i scales = _mm_cvtepu8_epi16(utmps);
+        const __m128i mins = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(utmps, utmps));
+
+        const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)&y[i].bsums[0]);
+        const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)&y[i].bsums[8]);
+        const __m128i q8s = _mm_hadd_epi16(q8sums_0, q8sums_1);
+        const __m128i prod = _mm_madd_epi16(mins, q8s);
+        const __m128i hsum = _mm_hadd_epi32(_mm_hadd_epi32(prod, mzero), mzero);
+        summs += dmin * _mm_extract_epi32(hsum, 0);
+
+        const __m128i hbits_0 = _mm_loadu_si128((const __m128i*)&x[i].qh[0]);
+        const __m128i hbits_1 = _mm_loadu_si128((const __m128i*)&x[i].qh[16]);
+        __m128i hmask = mone;
+
+        __m128i sumi_0 = _mm_setzero_si128();
+        __m128i sumi_1 = _mm_setzero_si128();
+
+        int bit = 0;
+
+        __m128i shuffle = _mm_set1_epi16(0x0100);
+        for (int j = 0; j < QK_K/64; ++j) {
+
+            const __m128i scale_0 = _mm_shuffle_epi8(scales, shuffle);
+            shuffle = _mm_add_epi16(shuffle, m2);
+            const __m128i scale_1 = _mm_shuffle_epi8(scales, shuffle);
+            shuffle = _mm_add_epi16(shuffle, m2);
+
+            const __m128i q5bits_0 = _mm_loadu_si128((const __m128i*)q5); q5 += 16;
+            const __m128i q5bits_1 = _mm_loadu_si128((const __m128i*)q5); q5 += 16;
+
+            __m128i q5l_0 = _mm_and_si128(q5bits_0, m4);
+            __m128i q5l_1 = _mm_and_si128(q5bits_1, m4);
+            __m128i q5h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_0, hmask), bit), 4);
+            __m128i q5h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_1, hmask), bit++), 4);
+            __m128i q5_0  = _mm_add_epi8(q5l_0, q5h_0);
+            __m128i q5_1  = _mm_add_epi8(q5l_1, q5h_1);
+            hmask = _mm_slli_epi16(hmask, 1);
+
+            __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            __m128i p16_0 = _mm_maddubs_epi16(q5_0, q8_0);
+            __m128i p16_1 = _mm_maddubs_epi16(q5_1, q8_1);
+            p16_0 = _mm_madd_epi16(scale_0, p16_0);
+            p16_1 = _mm_madd_epi16(scale_0, p16_1);
+
+            q5l_0 = _mm_and_si128(_mm_srli_epi16(q5bits_0, 4), m4);
+            q5l_1 = _mm_and_si128(_mm_srli_epi16(q5bits_1, 4), m4);
+            q5h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_0, hmask), bit), 4);
+            q5h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_1, hmask), bit++), 4);
+            q5_0  = _mm_add_epi8(q5l_0, q5h_0);
+            q5_1  = _mm_add_epi8(q5l_1, q5h_1);
+            hmask = _mm_slli_epi16(hmask, 1);
+
+            q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            __m128i p16_2 = _mm_maddubs_epi16(q5_0, q8_0);
+            __m128i p16_3 = _mm_maddubs_epi16(q5_1, q8_1);
+            p16_2 = _mm_madd_epi16(scale_1, p16_2);
+            p16_3 = _mm_madd_epi16(scale_1, p16_3);
+
+            sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2));
+            sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_1, p16_3));
+
+        }
+
+        __m256 vd = _mm256_set1_ps(d);
+        __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0);
+        acc = _mm256_add_ps(_mm256_mul_ps(vd, _mm256_cvtepi32_ps(sumi)), acc);
+
+    }
+
+    *s = hsum_float_8(acc) + summs;
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    UNUSED(kmask1);
+    UNUSED(kmask2);
+    UNUSED(kmask3);
+    UNUSED(utmp);
+    ggml_vec_dot_q5_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q6_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined __AVX2__
+
+    const __m256i m3 = _mm256_set1_epi8(3);
+    const __m256i m15 = _mm256_set1_epi8(15);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q4 = x[i].ql;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums);
+        const __m128i scales = _mm_loadu_si128((const __m128i*)x[i].scales);
+        const __m256i scales_16 = _mm256_cvtepi8_epi16(scales);
+        const __m256i q8sclsub = _mm256_slli_epi32(_mm256_madd_epi16(q8sums, scales_16), 5);
+
+        __m256i sumi = _mm256_setzero_si256();
+
+        int is = 0;
+
+        for (int j = 0; j < QK_K/128; ++j) {
+            const __m256i q4bits1 = _mm256_loadu_si256((const __m256i*)q4); q4 += 32;
+            const __m256i q4bits2 = _mm256_loadu_si256((const __m256i*)q4); q4 += 32;
+            const __m256i q4bitsH = _mm256_loadu_si256((const __m256i*)qh); qh += 32;
+
+            const __m256i q4h_0 = _mm256_slli_epi16(_mm256_and_si256(q4bitsH, m3), 4);
+            const __m256i q4h_1 = _mm256_slli_epi16(_mm256_and_si256(q4bitsH, _mm256_set1_epi8(12)), 2);
+            const __m256i q4h_2 = _mm256_and_si256(q4bitsH, _mm256_set1_epi8(48));
+            const __m256i q4h_3 = _mm256_srli_epi16(_mm256_and_si256(q4bitsH, _mm256_set1_epi8(-64)), 2);
+
+            const __m256i q4_0 = _mm256_or_si256(_mm256_and_si256(q4bits1, m15), q4h_0);
+            const __m256i q4_1 = _mm256_or_si256(_mm256_and_si256(q4bits2, m15), q4h_1);
+            const __m256i q4_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q4bits1, 4), m15), q4h_2);
+            const __m256i q4_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q4bits2, 4), m15), q4h_3);
+
+            const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            __m256i p16_0 = _mm256_maddubs_epi16(q4_0, q8_0);
+            __m256i p16_1 = _mm256_maddubs_epi16(q4_1, q8_1);
+            __m256i p16_2 = _mm256_maddubs_epi16(q4_2, q8_2);
+            __m256i p16_3 = _mm256_maddubs_epi16(q4_3, q8_3);
+
+            const __m128i scale_0 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 0));
+            const __m128i scale_1 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 1));
+            const __m128i scale_2 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 2));
+            const __m128i scale_3 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 3));
+            is += 4;
+
+            p16_0 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_0), p16_0);
+            p16_1 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_1), p16_1);
+            p16_2 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_2), p16_2);
+            p16_3 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_3), p16_3);
+
+            sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1));
+            sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_2, p16_3));
+
+        }
+
+        sumi = _mm256_sub_epi32(sumi, q8sclsub);
+        acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc);
+    }
+
+    *s = hsum_float_8(acc);
+
+#elif defined __AVX__
+
+    const __m128i m3 = _mm_set1_epi8(3);
+    const __m128i m15 = _mm_set1_epi8(15);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; ++i) {
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q4 = x[i].ql;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        // handle the q6_k -32 offset separately using bsums
+        const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)y[i].bsums);
+        const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)y[i].bsums + 1);
+        const __m128i scales = _mm_loadu_si128((const __m128i*)x[i].scales);
+        const __m128i scales_16_0 = _mm_cvtepi8_epi16(scales);
+        const __m128i scales_16_1 = _mm_cvtepi8_epi16(_mm_bsrli_si128(scales, 8));
+        const __m128i q8sclsub_0 = _mm_slli_epi32(_mm_madd_epi16(q8sums_0, scales_16_0), 5);
+        const __m128i q8sclsub_1 = _mm_slli_epi32(_mm_madd_epi16(q8sums_1, scales_16_1), 5);
+
+        __m128i sumi_0 = _mm_setzero_si128();
+        __m128i sumi_1 = _mm_setzero_si128();
+
+        int is = 0;
+
+        for (int j = 0; j < QK_K/128; ++j) {
+
+            const __m128i q4bitsH_0 = _mm_loadu_si128((const __m128i*)qh); qh += 16;
+            const __m128i q4bitsH_1 = _mm_loadu_si128((const __m128i*)qh); qh += 16;
+
+            const __m128i q4h_0 = _mm_slli_epi16(_mm_and_si128(q4bitsH_0, m3), 4);
+            const __m128i q4h_1 = _mm_slli_epi16(_mm_and_si128(q4bitsH_1, m3), 4);
+            const __m128i q4h_2 = _mm_slli_epi16(_mm_and_si128(q4bitsH_0, _mm_set1_epi8(12)), 2);
+            const __m128i q4h_3 = _mm_slli_epi16(_mm_and_si128(q4bitsH_1, _mm_set1_epi8(12)), 2);
+            const __m128i q4h_4 = _mm_and_si128(q4bitsH_0, _mm_set1_epi8(48));
+            const __m128i q4h_5 = _mm_and_si128(q4bitsH_1, _mm_set1_epi8(48));
+            const __m128i q4h_6 = _mm_srli_epi16(_mm_and_si128(q4bitsH_0, _mm_set1_epi8(-64)), 2);
+            const __m128i q4h_7 = _mm_srli_epi16(_mm_and_si128(q4bitsH_1, _mm_set1_epi8(-64)), 2);
+
+            const __m128i q4bits1_0 = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+            const __m128i q4bits1_1 = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+            const __m128i q4bits2_0 = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+            const __m128i q4bits2_1 = _mm_loadu_si128((const __m128i*)q4); q4 += 16;
+
+            const __m128i q4_0 = _mm_or_si128(_mm_and_si128(q4bits1_0, m15), q4h_0);
+            const __m128i q4_1 = _mm_or_si128(_mm_and_si128(q4bits1_1, m15), q4h_1);
+            const __m128i q4_2 = _mm_or_si128(_mm_and_si128(q4bits2_0, m15), q4h_2);
+            const __m128i q4_3 = _mm_or_si128(_mm_and_si128(q4bits2_1, m15), q4h_3);
+            const __m128i q4_4 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits1_0, 4), m15), q4h_4);
+            const __m128i q4_5 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits1_1, 4), m15), q4h_5);
+            const __m128i q4_6 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits2_0, 4), m15), q4h_6);
+            const __m128i q4_7 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits2_1, 4), m15), q4h_7);
+
+            const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+            const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16;
+
+            __m128i p16_0 = _mm_maddubs_epi16(q4_0, q8_0);
+            __m128i p16_1 = _mm_maddubs_epi16(q4_1, q8_1);
+            __m128i p16_2 = _mm_maddubs_epi16(q4_2, q8_2);
+            __m128i p16_3 = _mm_maddubs_epi16(q4_3, q8_3);
+            __m128i p16_4 = _mm_maddubs_epi16(q4_4, q8_4);
+            __m128i p16_5 = _mm_maddubs_epi16(q4_5, q8_5);
+            __m128i p16_6 = _mm_maddubs_epi16(q4_6, q8_6);
+            __m128i p16_7 = _mm_maddubs_epi16(q4_7, q8_7);
+
+            const __m128i scale_0 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 0));
+            const __m128i scale_1 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 1));
+            const __m128i scale_2 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 2));
+            const __m128i scale_3 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 3));
+            is += 4;
+
+            p16_0 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_0), p16_0);
+            p16_1 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_0, 8)), p16_1);
+            p16_2 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_1), p16_2);
+            p16_3 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_1, 8)), p16_3);
+            p16_4 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_2), p16_4);
+            p16_5 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_2, 8)), p16_5);
+            p16_6 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_3), p16_6);
+            p16_7 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_3, 8)), p16_7);
+
+            sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2));
+            sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_1, p16_3));
+            sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_4, p16_6));
+            sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_5, p16_7));
+
+        }
+
+        sumi_0 = _mm_sub_epi32(sumi_0, q8sclsub_0);
+        sumi_1 = _mm_sub_epi32(sumi_1, q8sclsub_1);
+        const __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0);
+        acc = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(sumi)), acc);
+    }
+
+    *s = hsum_float_8(acc);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_q6_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+#if defined (__AVX__) || defined (__AVX2__)
+static const int8_t keven_signs_q2xs[1024] = {
+     1,  1,  1,  1,  1,  1,  1,  1, -1,  1,  1,  1,  1,  1,  1, -1,  1, -1,  1,  1,  1,  1,  1, -1, -1, -1,  1,  1,  1,  1,  1,  1,
+     1,  1, -1,  1,  1,  1,  1, -1, -1,  1, -1,  1,  1,  1,  1,  1,  1, -1, -1,  1,  1,  1,  1,  1, -1, -1, -1,  1,  1,  1,  1, -1,
+     1,  1,  1, -1,  1,  1,  1, -1, -1,  1,  1, -1,  1,  1,  1,  1,  1, -1,  1, -1,  1,  1,  1,  1, -1, -1,  1, -1,  1,  1,  1, -1,
+     1,  1, -1, -1,  1,  1,  1,  1, -1,  1, -1, -1,  1,  1,  1, -1,  1, -1, -1, -1,  1,  1,  1, -1, -1, -1, -1, -1,  1,  1,  1,  1,
+     1,  1,  1,  1, -1,  1,  1, -1, -1,  1,  1,  1, -1,  1,  1,  1,  1, -1,  1,  1, -1,  1,  1,  1, -1, -1,  1,  1, -1,  1,  1, -1,
+     1,  1, -1,  1, -1,  1,  1,  1, -1,  1, -1,  1, -1,  1,  1, -1,  1, -1, -1,  1, -1,  1,  1, -1, -1, -1, -1,  1, -1,  1,  1,  1,
+     1,  1,  1, -1, -1,  1,  1,  1, -1,  1,  1, -1, -1,  1,  1, -1,  1, -1,  1, -1, -1,  1,  1, -1, -1, -1,  1, -1, -1,  1,  1,  1,
+     1,  1, -1, -1, -1,  1,  1, -1, -1,  1, -1, -1, -1,  1,  1,  1,  1, -1, -1, -1, -1,  1,  1,  1, -1, -1, -1, -1, -1,  1,  1, -1,
+     1,  1,  1,  1,  1, -1,  1, -1, -1,  1,  1,  1,  1, -1,  1,  1,  1, -1,  1,  1,  1, -1,  1,  1, -1, -1,  1,  1,  1, -1,  1, -1,
+     1,  1, -1,  1,  1, -1,  1,  1, -1,  1, -1,  1,  1, -1,  1, -1,  1, -1, -1,  1,  1, -1,  1, -1, -1, -1, -1,  1,  1, -1,  1,  1,
+     1,  1,  1, -1,  1, -1,  1,  1, -1,  1,  1, -1,  1, -1,  1, -1,  1, -1,  1, -1,  1, -1,  1, -1, -1, -1,  1, -1,  1, -1,  1,  1,
+     1,  1, -1, -1,  1, -1,  1, -1, -1,  1, -1, -1,  1, -1,  1,  1,  1, -1, -1, -1,  1, -1,  1,  1, -1, -1, -1, -1,  1, -1,  1, -1,
+     1,  1,  1,  1, -1, -1,  1,  1, -1,  1,  1,  1, -1, -1,  1, -1,  1, -1,  1,  1, -1, -1,  1, -1, -1, -1,  1,  1, -1, -1,  1,  1,
+     1,  1, -1,  1, -1, -1,  1, -1, -1,  1, -1,  1, -1, -1,  1,  1,  1, -1, -1,  1, -1, -1,  1,  1, -1, -1, -1,  1, -1, -1,  1, -1,
+     1,  1,  1, -1, -1, -1,  1, -1, -1,  1,  1, -1, -1, -1,  1,  1,  1, -1,  1, -1, -1, -1,  1,  1, -1, -1,  1, -1, -1, -1,  1, -1,
+     1,  1, -1, -1, -1, -1,  1,  1, -1,  1, -1, -1, -1, -1,  1, -1,  1, -1, -1, -1, -1, -1,  1, -1, -1, -1, -1, -1, -1, -1,  1,  1,
+     1,  1,  1,  1,  1,  1, -1, -1, -1,  1,  1,  1,  1,  1, -1,  1,  1, -1,  1,  1,  1,  1, -1,  1, -1, -1,  1,  1,  1,  1, -1, -1,
+     1,  1, -1,  1,  1,  1, -1,  1, -1,  1, -1,  1,  1,  1, -1, -1,  1, -1, -1,  1,  1,  1, -1, -1, -1, -1, -1,  1,  1,  1, -1,  1,
+     1,  1,  1, -1,  1,  1, -1,  1, -1,  1,  1, -1,  1,  1, -1, -1,  1, -1,  1, -1,  1,  1, -1, -1, -1, -1,  1, -1,  1,  1, -1,  1,
+     1,  1, -1, -1,  1,  1, -1, -1, -1,  1, -1, -1,  1,  1, -1,  1,  1, -1, -1, -1,  1,  1, -1,  1, -1, -1, -1, -1,  1,  1, -1, -1,
+     1,  1,  1,  1, -1,  1, -1,  1, -1,  1,  1,  1, -1,  1, -1, -1,  1, -1,  1,  1, -1,  1, -1, -1, -1, -1,  1,  1, -1,  1, -1,  1,
+     1,  1, -1,  1, -1,  1, -1, -1, -1,  1, -1,  1, -1,  1, -1,  1,  1, -1, -1,  1, -1,  1, -1,  1, -1, -1, -1,  1, -1,  1, -1, -1,
+     1,  1,  1, -1, -1,  1, -1, -1, -1,  1,  1, -1, -1,  1, -1,  1,  1, -1,  1, -1, -1,  1, -1,  1, -1, -1,  1, -1, -1,  1, -1, -1,
+     1,  1, -1, -1, -1,  1, -1,  1, -1,  1, -1, -1, -1,  1, -1, -1,  1, -1, -1, -1, -1,  1, -1, -1, -1, -1, -1, -1, -1,  1, -1,  1,
+     1,  1,  1,  1,  1, -1, -1,  1, -1,  1,  1,  1,  1, -1, -1, -1,  1, -1,  1,  1,  1, -1, -1, -1, -1, -1,  1,  1,  1, -1, -1,  1,
+     1,  1, -1,  1,  1, -1, -1, -1, -1,  1, -1,  1,  1, -1, -1,  1,  1, -1, -1,  1,  1, -1, -1,  1, -1, -1, -1,  1,  1, -1, -1, -1,
+     1,  1,  1, -1,  1, -1, -1, -1, -1,  1,  1, -1,  1, -1, -1,  1,  1, -1,  1, -1,  1, -1, -1,  1, -1, -1,  1, -1,  1, -1, -1, -1,
+     1,  1, -1, -1,  1, -1, -1,  1, -1,  1, -1, -1,  1, -1, -1, -1,  1, -1, -1, -1,  1, -1, -1, -1, -1, -1, -1, -1,  1, -1, -1,  1,
+     1,  1,  1,  1, -1, -1, -1, -1, -1,  1,  1,  1, -1, -1, -1,  1,  1, -1,  1,  1, -1, -1, -1,  1, -1, -1,  1,  1, -1, -1, -1, -1,
+     1,  1, -1,  1, -1, -1, -1,  1, -1,  1, -1,  1, -1, -1, -1, -1,  1, -1, -1,  1, -1, -1, -1, -1, -1, -1, -1,  1, -1, -1, -1,  1,
+     1,  1,  1, -1, -1, -1, -1,  1, -1,  1,  1, -1, -1, -1, -1, -1,  1, -1,  1, -1, -1, -1, -1, -1, -1, -1,  1, -1, -1, -1, -1,  1,
+     1,  1, -1, -1, -1, -1, -1, -1, -1,  1, -1, -1, -1, -1, -1,  1,  1, -1, -1, -1, -1, -1, -1,  1, -1, -1, -1, -1, -1, -1, -1, -1,
+};
+#endif
+
+void ggml_vec_dot_iq2_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_xxs * GGML_RESTRICT x = vx;
+    const block_q8_K    * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+
+    const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs;
+
+    uint32_t aux32[4];
+    const uint8_t * aux8 = (const uint8_t *)aux32;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            memcpy(aux32, q2, 4*sizeof(uint32_t)); q2 += 8;
+            const __m256i q2_1 = _mm256_set_epi64x(iq2xxs_grid[aux8[ 3]], iq2xxs_grid[aux8[ 2]], iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[0]]);
+            const __m256i q2_2 = _mm256_set_epi64x(iq2xxs_grid[aux8[11]], iq2xxs_grid[aux8[10]], iq2xxs_grid[aux8[9]], iq2xxs_grid[aux8[8]]);
+            const __m256i s2_1 = _mm256_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127],
+                                                   signs64[(aux32[1] >>  7) & 127], signs64[(aux32[1] >>  0) & 127]);
+            const __m256i s2_2 = _mm256_set_epi64x(signs64[(aux32[3] >> 21) & 127], signs64[(aux32[3] >> 14) & 127],
+                                                   signs64[(aux32[3] >>  7) & 127], signs64[(aux32[3] >>  0) & 127]);
+            const __m256i q8s_1 = _mm256_sign_epi8(q8_1, s2_1);
+            const __m256i q8s_2 = _mm256_sign_epi8(q8_2, s2_2);
+            const __m256i dot1  = _mm256_maddubs_epi16(q2_1, q8s_1);
+            const __m256i dot2  = _mm256_maddubs_epi16(q2_2, q8s_2);
+            const uint16_t ls1 = aux32[1] >> 28;
+            const uint16_t ls2 = aux32[3] >> 28;
+            const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1));
+            const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1));
+            sumi1 = _mm256_add_epi32(sumi1, p1);
+            sumi2 = _mm256_add_epi32(sumi2, p2);
+        }
+
+        accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#elif defined(__AVX__)
+    const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs;
+
+    uint32_t aux32[4];
+    const uint8_t * aux8 = (const uint8_t *)aux32;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            memcpy(aux32, q2, 4*sizeof(uint32_t)); q2 += 8;
+            const __m128i q2_1_0 = _mm_set_epi64x(iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[0]]);
+            const __m128i q2_1_1 = _mm_set_epi64x(iq2xxs_grid[aux8[3]], iq2xxs_grid[aux8[2]]);
+            const __m128i q2_2_0 = _mm_set_epi64x(iq2xxs_grid[aux8[9]], iq2xxs_grid[aux8[8]]);
+            const __m128i q2_2_1 = _mm_set_epi64x(iq2xxs_grid[aux8[11]], iq2xxs_grid[aux8[10]]);
+            const __m128i s2_1_0 = _mm_set_epi64x(signs64[(aux32[1] >>  7) & 127], signs64[(aux32[1] >>  0) & 127]);
+            const __m128i s2_1_1 = _mm_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127]);
+            const __m128i s2_2_0 = _mm_set_epi64x(signs64[(aux32[3] >>  7) & 127], signs64[(aux32[3] >>  0) & 127]);
+            const __m128i s2_2_1 = _mm_set_epi64x(signs64[(aux32[3] >> 21) & 127], signs64[(aux32[3] >> 14) & 127]);
+            const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, s2_1_0);
+            const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, s2_1_1);
+            const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, s2_2_0);
+            const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, s2_2_1);
+            const __m128i dot1_0  = _mm_maddubs_epi16(q2_1_0, q8s_1_0);
+            const __m128i dot1_1  = _mm_maddubs_epi16(q2_1_1, q8s_1_1);
+            const __m128i dot2_0  = _mm_maddubs_epi16(q2_2_0, q8s_2_0);
+            const __m128i dot2_1  = _mm_maddubs_epi16(q2_2_1, q8s_2_1);
+            const uint16_t ls1 = aux32[1] >> 28;
+            const uint16_t ls2 = aux32[3] >> 28;
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1));
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1));
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1));
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1));
+            sumi1_0 = _mm_add_epi32(sumi1_0, p1_0);
+            sumi1_1 = _mm_add_epi32(sumi1_1, p1_1);
+            sumi2_0 = _mm_add_epi32(sumi2_0, p2_0);
+            sumi2_1 = _mm_add_epi32(sumi2_1, p2_1);
+        }
+
+        accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq2_xxs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq2_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_xs * GGML_RESTRICT x = vx;
+    const block_q8_K   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+
+    const __m256i mone = _mm256_set1_epi8(1);
+    static const char block_sign_shuffle_mask_1[32] = {
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
+        0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06,
+    };
+    static const char block_sign_shuffle_mask_2[32] = {
+        0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a,
+        0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e,
+    };
+    static const uint8_t bit_selector_mask_bytes[32] = {
+        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m256i bit_selector_mask = _mm256_loadu_si256((const __m256i*)bit_selector_mask_bytes);
+    const __m256i block_sign_shuffle_1 = _mm256_loadu_si256((const __m256i*)block_sign_shuffle_mask_1);
+    const __m256i block_sign_shuffle_2 = _mm256_loadu_si256((const __m256i*)block_sign_shuffle_mask_2);
+
+    static const uint8_t k_bit_helper[32] = {
+        0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00,
+        0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00,
+    };
+    const __m256i bit_helper = _mm256_loadu_si256((const __m256i*)k_bit_helper);
+    const __m256i m511 = _mm256_set1_epi16(511);
+    const __m128i m4 = _mm_set1_epi8(0xf);
+    const __m128i m1 = _mm_set1_epi8(1);
+
+    uint64_t aux64;
+
+    // somewhat hacky, but gives a significant boost in performance
+    __m256i aux_gindex;
+    const uint16_t * gindex = (const uint16_t *)&aux_gindex;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(&aux64, x[i].scales, 8);
+        __m128i stmp = _mm_set1_epi64x(aux64);
+        stmp = _mm_unpacklo_epi8(_mm_and_si128(stmp, m4), _mm_and_si128(_mm_srli_epi16(stmp, 4), m4));
+        const __m128i scales = _mm_add_epi8(_mm_slli_epi16(stmp, 1), m1);
+
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 4) {
+
+            const __m256i q2_data = _mm256_loadu_si256((const __m256i*)q2);  q2 += 16;
+            aux_gindex = _mm256_and_si256(q2_data, m511);
+
+            const __m256i partial_sign_bits = _mm256_srli_epi16(q2_data, 9);
+            const __m256i partial_sign_bits_upper = _mm256_srli_epi16(q2_data, 13);
+            const __m256i partial_sign_bits_for_counting = _mm256_xor_si256(partial_sign_bits, partial_sign_bits_upper);
+
+            const __m256i odd_bits = _mm256_shuffle_epi8(bit_helper, partial_sign_bits_for_counting);
+            const __m256i full_sign_bits = _mm256_or_si256(partial_sign_bits, odd_bits);
+
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_3 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_4 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+
+            const __m256i q2_1 = _mm256_set_epi64x(iq2xs_grid[gindex[ 3]], iq2xs_grid[gindex[ 2]],
+                                                   iq2xs_grid[gindex[ 1]], iq2xs_grid[gindex[ 0]]);
+            const __m256i q2_2 = _mm256_set_epi64x(iq2xs_grid[gindex[ 7]], iq2xs_grid[gindex[ 6]],
+                                                   iq2xs_grid[gindex[ 5]], iq2xs_grid[gindex[ 4]]);
+            const __m256i q2_3 = _mm256_set_epi64x(iq2xs_grid[gindex[11]], iq2xs_grid[gindex[10]],
+                                                   iq2xs_grid[gindex[ 9]], iq2xs_grid[gindex[ 8]]);
+            const __m256i q2_4 = _mm256_set_epi64x(iq2xs_grid[gindex[15]], iq2xs_grid[gindex[14]],
+                                                   iq2xs_grid[gindex[13]], iq2xs_grid[gindex[12]]);
+
+            const __m128i full_signs_l = _mm256_castsi256_si128(full_sign_bits);
+            const __m128i full_signs_h = _mm256_extractf128_si256(full_sign_bits, 1);
+            const __m256i full_signs_1 = MM256_SET_M128I(full_signs_l, full_signs_l);
+            const __m256i full_signs_2 = MM256_SET_M128I(full_signs_h, full_signs_h);
+
+            __m256i signs;
+            signs = _mm256_shuffle_epi8(full_signs_1, block_sign_shuffle_1);
+            signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask);
+            const __m256i q8s_1 = _mm256_sign_epi8(q8_1, _mm256_or_si256(signs, mone));
+
+            signs = _mm256_shuffle_epi8(full_signs_1, block_sign_shuffle_2);
+            signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask);
+            const __m256i q8s_2 = _mm256_sign_epi8(q8_2, _mm256_or_si256(signs, mone));
+
+            signs = _mm256_shuffle_epi8(full_signs_2, block_sign_shuffle_1);
+            signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask);
+            const __m256i q8s_3 = _mm256_sign_epi8(q8_3, _mm256_or_si256(signs, mone));
+
+            signs = _mm256_shuffle_epi8(full_signs_2, block_sign_shuffle_2);
+            signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask);
+            const __m256i q8s_4 = _mm256_sign_epi8(q8_4, _mm256_or_si256(signs, mone));
+
+            const __m256i dot1  = _mm256_maddubs_epi16(q2_1, q8s_1);
+            const __m256i dot2  = _mm256_maddubs_epi16(q2_2, q8s_2);
+            const __m256i dot3  = _mm256_maddubs_epi16(q2_3, q8s_3);
+            const __m256i dot4  = _mm256_maddubs_epi16(q2_4, q8s_4);
+
+            const __m256i sc1 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+0)));
+            const __m256i sc2 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+1)));
+            const __m256i sc3 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+2)));
+            const __m256i sc4 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+3)));
+
+            sumi1 = _mm256_add_epi32(sumi1, _mm256_madd_epi16(dot1, sc1));
+            sumi2 = _mm256_add_epi32(sumi2, _mm256_madd_epi16(dot2, sc2));
+            sumi1 = _mm256_add_epi32(sumi1, _mm256_madd_epi16(dot3, sc3));
+            sumi2 = _mm256_add_epi32(sumi2, _mm256_madd_epi16(dot4, sc4));
+        }
+
+        accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#elif defined(__AVX__)
+    const __m128i mone = _mm_set1_epi8(1);
+    static const char block_sign_shuffle_mask_1[32] = {
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
+        0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06,
+    };
+    static const char block_sign_shuffle_mask_2[32] = {
+        0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a,
+        0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e,
+    };
+    static const uint8_t bit_selector_mask_bytes[32] = {
+        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m128i bit_selector_mask_0 = _mm_loadu_si128((const __m128i*)bit_selector_mask_bytes);
+    const __m128i bit_selector_mask_1 = _mm_loadu_si128((const __m128i*)bit_selector_mask_bytes + 1);
+    const __m128i block_sign_shuffle_1_0 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_1);
+    const __m128i block_sign_shuffle_1_1 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_1 + 1);
+    const __m128i block_sign_shuffle_2_0 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_2);
+    const __m128i block_sign_shuffle_2_1 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_2 + 1);
+
+    static const uint8_t k_bit_helper[32] = {
+        0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00,
+        0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00,
+    };
+    const __m128i bit_helper_0 = _mm_loadu_si128((const __m128i*)k_bit_helper);
+    const __m128i bit_helper_1 = _mm_loadu_si128((const __m128i*)k_bit_helper + 1);
+    const __m128i m511 = _mm_set1_epi16(511);
+    const __m128i m4 = _mm_set1_epi8(0xf);
+    const __m128i m1 = _mm_set1_epi8(1);
+
+    uint64_t aux64;
+
+    // somewhat hacky, but gives a significant boost in performance
+    __m256i aux_gindex;
+    const uint16_t * gindex = (const uint16_t *)&aux_gindex;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(&aux64, x[i].scales, 8);
+        __m128i stmp = _mm_set1_epi64x(aux64);
+        stmp = _mm_unpacklo_epi8(_mm_and_si128(stmp, m4), _mm_and_si128(_mm_srli_epi16(stmp, 4), m4));
+        const __m128i scales = _mm_add_epi8(_mm_slli_epi16(stmp, 1), m1);
+
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 4) {
+
+            const __m128i q2_data_0 = _mm_loadu_si128((const __m128i*)q2);
+            const __m128i q2_data_1 = _mm_loadu_si128((const __m128i*)q2 + 1);  q2 += 16;
+            aux_gindex = MM256_SET_M128I(_mm_and_si128(q2_data_1, m511), _mm_and_si128(q2_data_0, m511));
+
+            const __m128i partial_sign_bits_0 = _mm_srli_epi16(q2_data_0, 9);
+            const __m128i partial_sign_bits_1 = _mm_srli_epi16(q2_data_1, 9);
+            const __m128i partial_sign_bits_upper_0 = _mm_srli_epi16(q2_data_0, 13);
+            const __m128i partial_sign_bits_upper_1 = _mm_srli_epi16(q2_data_1, 13);
+            const __m128i partial_sign_bits_for_counting_0 = _mm_xor_si128(partial_sign_bits_0, partial_sign_bits_upper_0);
+            const __m128i partial_sign_bits_for_counting_1 = _mm_xor_si128(partial_sign_bits_1, partial_sign_bits_upper_1);
+
+            const __m128i odd_bits_0 = _mm_shuffle_epi8(bit_helper_0, partial_sign_bits_for_counting_0);
+            const __m128i odd_bits_1 = _mm_shuffle_epi8(bit_helper_1, partial_sign_bits_for_counting_1);
+            const __m128i full_sign_bits_0 = _mm_or_si128(partial_sign_bits_0, odd_bits_0);
+            const __m128i full_sign_bits_1 = _mm_or_si128(partial_sign_bits_1, odd_bits_1);
+
+            const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_3_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_3_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_4_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_4_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+
+            const __m128i q2_1_0 = _mm_set_epi64x(iq2xs_grid[gindex[1]], iq2xs_grid[gindex[0]]);
+            const __m128i q2_1_1 = _mm_set_epi64x(iq2xs_grid[gindex[3]], iq2xs_grid[gindex[2]]);
+            const __m128i q2_2_0 = _mm_set_epi64x(iq2xs_grid[gindex[5]], iq2xs_grid[gindex[4]]);
+            const __m128i q2_2_1 = _mm_set_epi64x(iq2xs_grid[gindex[7]], iq2xs_grid[gindex[6]]);
+            const __m128i q2_3_0 = _mm_set_epi64x(iq2xs_grid[gindex[9]], iq2xs_grid[gindex[8]]);
+            const __m128i q2_3_1 = _mm_set_epi64x(iq2xs_grid[gindex[11]], iq2xs_grid[gindex[10]]);
+            const __m128i q2_4_0 = _mm_set_epi64x(iq2xs_grid[gindex[13]], iq2xs_grid[gindex[12]]);
+            const __m128i q2_4_1 = _mm_set_epi64x(iq2xs_grid[gindex[15]], iq2xs_grid[gindex[14]]);
+
+            // AVX2 full_signs_1 is full_sign_bits_0 here
+            // AVX2 full_signs_2 is full_sign_bits_1 here
+            __m128i signs_0, signs_1;
+            signs_0 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_1_0);
+            signs_1 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_1_1);
+            signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0);
+            signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1);
+            const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, _mm_or_si128(signs_0, mone));
+            const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, _mm_or_si128(signs_1, mone));
+
+            signs_0 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_2_0);
+            signs_1 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_2_1);
+            signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0);
+            signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1);
+            const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, _mm_or_si128(signs_0, mone));
+            const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, _mm_or_si128(signs_1, mone));
+
+            signs_0 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_1_0);
+            signs_1 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_1_1);
+            signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0);
+            signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1);
+            const __m128i q8s_3_0 = _mm_sign_epi8(q8_3_0, _mm_or_si128(signs_0, mone));
+            const __m128i q8s_3_1 = _mm_sign_epi8(q8_3_1, _mm_or_si128(signs_1, mone));
+
+            signs_0 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_2_0);
+            signs_1 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_2_1);
+            signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0);
+            signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1);
+            const __m128i q8s_4_0 = _mm_sign_epi8(q8_4_0, _mm_or_si128(signs_0, mone));
+            const __m128i q8s_4_1 = _mm_sign_epi8(q8_4_1, _mm_or_si128(signs_1, mone));
+
+            const __m128i dot1_0  = _mm_maddubs_epi16(q2_1_0, q8s_1_0);
+            const __m128i dot1_1  = _mm_maddubs_epi16(q2_1_1, q8s_1_1);
+            const __m128i dot2_0  = _mm_maddubs_epi16(q2_2_0, q8s_2_0);
+            const __m128i dot2_1  = _mm_maddubs_epi16(q2_2_1, q8s_2_1);
+            const __m128i dot3_0  = _mm_maddubs_epi16(q2_3_0, q8s_3_0);
+            const __m128i dot3_1  = _mm_maddubs_epi16(q2_3_1, q8s_3_1);
+            const __m128i dot4_0  = _mm_maddubs_epi16(q2_4_0, q8s_4_0);
+            const __m128i dot4_1  = _mm_maddubs_epi16(q2_4_1, q8s_4_1);
+
+            __m128i sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+0));
+            const __m128i sc1_0 = _mm_cvtepi8_epi16(sc_tmp);
+            const __m128i sc1_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8));
+            sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+1));
+            const __m128i sc2_0 = _mm_cvtepi8_epi16(sc_tmp);
+            const __m128i sc2_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8));
+            sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+2));
+            const __m128i sc3_0 = _mm_cvtepi8_epi16(sc_tmp);
+            const __m128i sc3_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8));
+            sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+3));
+            const __m128i sc4_0 = _mm_cvtepi8_epi16(sc_tmp);
+            const __m128i sc4_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8));
+
+            sumi1_0 = _mm_add_epi32(sumi1_0, _mm_madd_epi16(dot1_0, sc1_0));
+            sumi1_1 = _mm_add_epi32(sumi1_1, _mm_madd_epi16(dot1_1, sc1_1));
+            sumi2_0 = _mm_add_epi32(sumi2_0, _mm_madd_epi16(dot2_0, sc2_0));
+            sumi2_1 = _mm_add_epi32(sumi2_1, _mm_madd_epi16(dot2_1, sc2_1));
+            sumi1_0 = _mm_add_epi32(sumi1_0, _mm_madd_epi16(dot3_0, sc3_0));
+            sumi1_1 = _mm_add_epi32(sumi1_1, _mm_madd_epi16(dot3_1, sc3_1));
+            sumi2_0 = _mm_add_epi32(sumi2_0, _mm_madd_epi16(dot4_0, sc4_0));
+            sumi2_1 = _mm_add_epi32(sumi2_1, _mm_madd_epi16(dot4_1, sc4_1));
+        }
+
+        accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq2_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq2_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+
+   static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+                                       0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03
+   };
+
+    static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+                                        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m128i m4 = _mm_set1_epi8(0xf);
+    const __m128i m1 = _mm_set1_epi8(1);
+
+    const __m256i mask1 = _mm256_loadu_si256((const __m256i*)k_mask1);
+    const __m256i mask2 = _mm256_loadu_si256((const __m256i*)k_mask2);
+
+    uint64_t aux64;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT qs = x[i].qs;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const uint16_t * GGML_RESTRICT signs = (const uint16_t *)(x[i].qs + QK_K/8);
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(&aux64, x[i].scales, 8);
+        const __m128i scales8 = _mm_add_epi8(_mm_slli_epi16(_mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), m4), 1), m1);
+        const __m256i scales16 = _mm256_cvtepi8_epi16(scales8); // 0 2 4 6 8 10 12 14 1 3 5 7 9 11 13 15
+
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q2_1 = _mm256_set_epi64x(iq2s_grid[qs[3] | ((qh[ib32+0] << 2) & 0x300)],
+                                                   iq2s_grid[qs[2] | ((qh[ib32+0] << 4) & 0x300)],
+                                                   iq2s_grid[qs[1] | ((qh[ib32+0] << 6) & 0x300)],
+                                                   iq2s_grid[qs[0] | ((qh[ib32+0] << 8) & 0x300)]);
+            const __m256i q2_2 = _mm256_set_epi64x(iq2s_grid[qs[7] | ((qh[ib32+1] << 2) & 0x300)],
+                                                   iq2s_grid[qs[6] | ((qh[ib32+1] << 4) & 0x300)],
+                                                   iq2s_grid[qs[5] | ((qh[ib32+1] << 6) & 0x300)],
+                                                   iq2s_grid[qs[4] | ((qh[ib32+1] << 8) & 0x300)]);
+            qs += 8;
+
+            __m256i aux256 = _mm256_set1_epi32(signs[0] | ((uint32_t) signs[1] << 16));
+            aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2);
+            const __m256i s2_1 = _mm256_cmpeq_epi8(aux256, mask2);
+            const __m256i q8s_1 = _mm256_sub_epi8(_mm256_xor_si256(s2_1, q8_1), s2_1);
+
+            aux256 = _mm256_set1_epi32(signs[2] | ((uint32_t) signs[3] << 16));
+            aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2);
+            const __m256i s2_2 = _mm256_cmpeq_epi8(aux256, mask2);
+            const __m256i q8s_2 = _mm256_sub_epi8(_mm256_xor_si256(s2_2, q8_2), s2_2);
+
+            signs += 4;
+
+            const __m256i dot1  = _mm256_maddubs_epi16(q2_1, q8s_1); // blocks 2*ib32+0, 2*ib32+1
+            const __m256i dot2  = _mm256_maddubs_epi16(q2_2, q8s_2); // blocks 2*ib32+2, 2*ib32+3
+
+            const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_shuffle_epi8(scales16, get_scale_shuffle_k4(ib32+0)));
+            const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_shuffle_epi8(scales16, get_scale_shuffle_k4(ib32+1)));
+            sumi1 = _mm256_add_epi32(sumi1, p1);
+            sumi2 = _mm256_add_epi32(sumi2, p2);
+        }
+
+        accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#elif defined(__AVX__)
+   static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+                                       0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03
+   };
+
+    static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+                                        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m128i m4 = _mm_set1_epi8(0xf);
+    const __m128i m1 = _mm_set1_epi8(1);
+
+    const __m128i mask1_0 = _mm_loadu_si128((const __m128i*)k_mask1);
+    const __m128i mask1_1 = _mm_loadu_si128((const __m128i*)k_mask1 + 1);
+    const __m128i mask2_0 = _mm_loadu_si128((const __m128i*)k_mask2);
+    const __m128i mask2_1 = _mm_loadu_si128((const __m128i*)k_mask2 + 1);
+
+    uint64_t aux64;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT qs = x[i].qs;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const uint16_t * GGML_RESTRICT signs = (const uint16_t *)(x[i].qs + QK_K/8);
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+
+        memcpy(&aux64, x[i].scales, 8);
+        const __m128i scales8 = _mm_add_epi8(_mm_slli_epi16(_mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), m4), 1), m1);
+        const __m128i scales16_0 = _mm_cvtepi8_epi16(scales8);
+        const __m128i scales16_1 = _mm_cvtepi8_epi16(_mm_srli_si128(scales8, 8));
+
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q2_1_0 = _mm_set_epi64x(iq2s_grid[qs[1] | ((qh[ib32+0] << 6) & 0x300)],
+                                                  iq2s_grid[qs[0] | ((qh[ib32+0] << 8) & 0x300)]);
+            const __m128i q2_1_1 = _mm_set_epi64x(iq2s_grid[qs[3] | ((qh[ib32+0] << 2) & 0x300)],
+                                                  iq2s_grid[qs[2] | ((qh[ib32+0] << 4) & 0x300)]);
+            const __m128i q2_2_0 = _mm_set_epi64x(iq2s_grid[qs[5] | ((qh[ib32+1] << 6) & 0x300)],
+                                                  iq2s_grid[qs[4] | ((qh[ib32+1] << 8) & 0x300)]);
+            const __m128i q2_2_1 = _mm_set_epi64x(iq2s_grid[qs[7] | ((qh[ib32+1] << 2) & 0x300)],
+                                                  iq2s_grid[qs[6] | ((qh[ib32+1] << 4) & 0x300)]);
+            qs += 8;
+
+            __m128i aux128_0 = _mm_set1_epi32(signs[0] | ((uint32_t) signs[1] << 16));
+            __m128i aux128_1 = aux128_0;
+            aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0);
+            aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1);
+            const __m128i s2_1_0 = _mm_cmpeq_epi8(aux128_0, mask2_0);
+            const __m128i s2_1_1 = _mm_cmpeq_epi8(aux128_1, mask2_1);
+            const __m128i q8s_1_0 = _mm_sub_epi8(_mm_xor_si128(s2_1_0, q8_1_0), s2_1_0);
+            const __m128i q8s_1_1 = _mm_sub_epi8(_mm_xor_si128(s2_1_1, q8_1_1), s2_1_1);
+
+            aux128_0 = _mm_set1_epi32(signs[2] | ((uint32_t) signs[3] << 16));
+            aux128_1 = aux128_0;
+            aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0);
+            aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1);
+            const __m128i s2_2_0 = _mm_cmpeq_epi8(aux128_0, mask2_0);
+            const __m128i s2_2_1 = _mm_cmpeq_epi8(aux128_1, mask2_1);
+            const __m128i q8s_2_0 = _mm_sub_epi8(_mm_xor_si128(s2_2_0, q8_2_0), s2_2_0);
+            const __m128i q8s_2_1 = _mm_sub_epi8(_mm_xor_si128(s2_2_1, q8_2_1), s2_2_1);
+
+            signs += 4;
+
+            const __m128i dot1_0  = _mm_maddubs_epi16(q2_1_0, q8s_1_0);
+            const __m128i dot1_1  = _mm_maddubs_epi16(q2_1_1, q8s_1_1);
+            const __m128i dot2_0  = _mm_maddubs_epi16(q2_2_0, q8s_2_0);
+            const __m128i dot2_1  = _mm_maddubs_epi16(q2_2_1, q8s_2_1);
+
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_shuffle_epi8(scales16_0, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+0), 0)));
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_shuffle_epi8(scales16_1, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+0), 1)));
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_shuffle_epi8(scales16_0, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+1), 0)));
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_shuffle_epi8(scales16_1, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+1), 1)));
+            sumi1_0 = _mm_add_epi32(sumi1_0, p1_0);
+            sumi1_1 = _mm_add_epi32(sumi1_1, p1_1);
+            sumi2_0 = _mm_add_epi32(sumi2_0, p2_0);
+            sumi2_1 = _mm_add_epi32(sumi2_1, p2_1);
+        }
+
+        accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf);
+
+    }
+
+    *s = 0.125f * hsum_float_8(accumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq2_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq3_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq3_xxs * GGML_RESTRICT x = vx;
+    const block_q8_K    * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+
+    const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs;
+
+    uint32_t aux32[2];
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q2_1 = _mm256_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]],
+                                                  iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]);
+            q3 += 8;
+            const __m256i q2_2 = _mm256_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]],
+                                                  iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]);
+            q3 += 8;
+            memcpy(aux32, gas, 8); gas += 8;
+            const __m256i s2_1 = _mm256_set_epi64x(signs64[(aux32[0] >> 21) & 127], signs64[(aux32[0] >> 14) & 127],
+                                                   signs64[(aux32[0] >>  7) & 127], signs64[(aux32[0] >>  0) & 127]);
+            const __m256i s2_2 = _mm256_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127],
+                                                   signs64[(aux32[1] >>  7) & 127], signs64[(aux32[1] >>  0) & 127]);
+            const __m256i q8s_1 = _mm256_sign_epi8(q8_1, s2_1);
+            const __m256i q8s_2 = _mm256_sign_epi8(q8_2, s2_2);
+            const __m256i dot1  = _mm256_maddubs_epi16(q2_1, q8s_1);
+            const __m256i dot2  = _mm256_maddubs_epi16(q2_2, q8s_2);
+            const uint16_t ls1 = aux32[0] >> 28;
+            const uint16_t ls2 = aux32[1] >> 28;
+            const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1));
+            const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1));
+            sumi1 = _mm256_add_epi32(sumi1, p1);
+            sumi2 = _mm256_add_epi32(sumi2, p2);
+        }
+
+        accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf);
+
+    }
+
+    *s = 0.25f * hsum_float_8(accumf);
+
+#elif defined(__AVX__)
+    const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs;
+
+    uint32_t aux32[2];
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q2_1_0 = _mm_set_epi32(iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]);
+            const __m128i q2_1_1 = _mm_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]]);
+            q3 += 8;
+            const __m128i q2_2_0 = _mm_set_epi32(iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]);
+            const __m128i q2_2_1 = _mm_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]]);
+            q3 += 8;
+            memcpy(aux32, gas, 8); gas += 8;
+            const __m128i s2_1_0 = _mm_set_epi64x(signs64[(aux32[0] >>  7) & 127], signs64[(aux32[0] >>  0) & 127]);
+            const __m128i s2_1_1 = _mm_set_epi64x(signs64[(aux32[0] >> 21) & 127], signs64[(aux32[0] >> 14) & 127]);
+            const __m128i s2_2_0 = _mm_set_epi64x(signs64[(aux32[1] >>  7) & 127], signs64[(aux32[1] >>  0) & 127]);
+            const __m128i s2_2_1 = _mm_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127]);
+            const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, s2_1_0);
+            const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, s2_1_1);
+            const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, s2_2_0);
+            const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, s2_2_1);
+            const __m128i dot1_0  = _mm_maddubs_epi16(q2_1_0, q8s_1_0);
+            const __m128i dot1_1  = _mm_maddubs_epi16(q2_1_1, q8s_1_1);
+            const __m128i dot2_0  = _mm_maddubs_epi16(q2_2_0, q8s_2_0);
+            const __m128i dot2_1  = _mm_maddubs_epi16(q2_2_1, q8s_2_1);
+            const uint16_t ls1 = aux32[0] >> 28;
+            const uint16_t ls2 = aux32[1] >> 28;
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1));
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1));
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1));
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1));
+            sumi1_0 = _mm_add_epi32(sumi1_0, p1_0);
+            sumi1_1 = _mm_add_epi32(sumi1_1, p1_1);
+            sumi2_0 = _mm_add_epi32(sumi2_0, p2_0);
+            sumi2_1 = _mm_add_epi32(sumi2_1, p2_1);
+        }
+
+        accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf);
+
+    }
+
+    *s = 0.25f * hsum_float_8(accumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq3_xxs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq3_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq3_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined(__AVX2__)
+
+   static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+                                       0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03
+   };
+
+    static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+                                        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m256i mask1 = _mm256_loadu_si256((const __m256i*)k_mask1);
+    const __m256i mask2 = _mm256_loadu_si256((const __m256i*)k_mask2);
+
+    const __m256i idx_shift = _mm256_set_epi32(1, 2, 3, 4, 5, 6, 7, 8);
+    const __m256i idx_mask  = _mm256_set1_epi32(256);
+
+    typedef union {
+        __m256i  vec[2];
+        uint32_t index[16];
+    } index_t;
+
+    index_t idx;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT qs = x[i].qs;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const uint16_t * GGML_RESTRICT signs = (const uint16_t *)x[i].signs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i idx_l = _mm256_cvtepu8_epi16(_mm_loadu_si128((const __m128i *)qs)); qs += 16;
+            idx.vec[0] = _mm256_set1_epi32(qh[ib32+0]);
+            idx.vec[1] = _mm256_set1_epi32(qh[ib32+1]);
+            idx.vec[0] = _mm256_and_si256(_mm256_sllv_epi32(idx.vec[0], idx_shift), idx_mask);
+            idx.vec[1] = _mm256_and_si256(_mm256_sllv_epi32(idx.vec[1], idx_shift), idx_mask);
+            idx.vec[0] = _mm256_or_si256(idx.vec[0], _mm256_cvtepi16_epi32(_mm256_castsi256_si128(idx_l)));
+            idx.vec[1] = _mm256_or_si256(idx.vec[1], _mm256_cvtepi16_epi32(_mm256_extractf128_si256(idx_l, 1)));
+
+            // At leat on my CPU (Ryzen 7950X), using _mm256_i32gather_epi32 is slower than _mm256_set_epi32. Strange.
+            //const __m256i q2_1 = _mm256_i32gather_epi32((const int *)iq3s_grid, idx.vec[0], 4);
+            //const __m256i q2_2 = _mm256_i32gather_epi32((const int *)iq3s_grid, idx.vec[1], 4);
+            const __m256i q2_1 = _mm256_set_epi32(
+                    iq3s_grid[idx.index[7]], iq3s_grid[idx.index[6]], iq3s_grid[idx.index[5]], iq3s_grid[idx.index[4]],
+                    iq3s_grid[idx.index[3]], iq3s_grid[idx.index[2]], iq3s_grid[idx.index[1]], iq3s_grid[idx.index[0]]
+            );
+            const __m256i q2_2 = _mm256_set_epi32(
+                    iq3s_grid[idx.index[15]], iq3s_grid[idx.index[14]], iq3s_grid[idx.index[13]], iq3s_grid[idx.index[12]],
+                    iq3s_grid[idx.index[11]], iq3s_grid[idx.index[10]], iq3s_grid[idx.index[ 9]], iq3s_grid[idx.index[ 8]]
+            );
+
+            __m256i aux256 = _mm256_set1_epi32(signs[0] | (signs[1] << 16));
+            aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2);
+            const __m256i s2_1 = _mm256_cmpeq_epi8(aux256, mask2);
+            const __m256i q8s_1 = _mm256_sub_epi8(_mm256_xor_si256(s2_1, q8_1), s2_1);
+
+            aux256 = _mm256_set1_epi32(signs[2] | (signs[3] << 16));
+            aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2);
+            const __m256i s2_2 = _mm256_cmpeq_epi8(aux256, mask2);
+            const __m256i q8s_2 = _mm256_sub_epi8(_mm256_xor_si256(s2_2, q8_2), s2_2);
+
+            signs += 4;
+
+            const __m256i dot1  = _mm256_maddubs_epi16(q2_1, q8s_1);
+            const __m256i dot2  = _mm256_maddubs_epi16(q2_2, q8s_2);
+            const uint16_t ls1 = x[i].scales[ib32/2] & 0xf;
+            const uint16_t ls2 = x[i].scales[ib32/2] >>  4;
+            const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1));
+            const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1));
+            sumi1 = _mm256_add_epi32(sumi1, p1);
+            sumi2 = _mm256_add_epi32(sumi2, p2);
+        }
+
+        accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf);
+
+    }
+
+    *s = hsum_float_8(accumf);
+
+#elif defined(__AVX__)
+   static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+                                       0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03
+   };
+
+    static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+                                        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
+    };
+
+    const __m128i mask1_0 = _mm_loadu_si128((const __m128i*)k_mask1);
+    const __m128i mask1_1 = _mm_loadu_si128((const __m128i*)k_mask1 + 1);
+    const __m128i mask2_0 = _mm_loadu_si128((const __m128i*)k_mask2);
+    const __m128i mask2_1 = _mm_loadu_si128((const __m128i*)k_mask2 + 1);
+
+    const __m128i idx_mul_0 = _mm_set_epi32(32, 64, 128, 256);
+    const __m128i idx_mul_1 = _mm_set_epi32(2, 4, 8, 16);
+    const __m128i idx_mask  = _mm_set1_epi32(256);
+
+    typedef union {
+        __m128i  vec[4];
+        uint32_t index[16];
+    } index_t;
+
+    index_t idx;
+
+    __m256 accumf = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT qs = x[i].qs;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const uint16_t * GGML_RESTRICT signs = (const uint16_t *)x[i].signs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i qs_tmp = _mm_loadu_si128((const __m128i *)qs);
+            const __m128i idx_l_0 = _mm_cvtepu8_epi16(qs_tmp);
+            const __m128i idx_l_1 = _mm_cvtepu8_epi16(_mm_srli_si128(qs_tmp, 8)); qs += 16;
+            idx.vec[0] = _mm_set1_epi32(qh[ib32+0]);
+            idx.vec[1] = idx.vec[0];
+            idx.vec[2] = _mm_set1_epi32(qh[ib32+1]);
+            idx.vec[3] = idx.vec[2];
+
+            idx.vec[0] = _mm_and_si128(_mm_mullo_epi32(idx.vec[0], idx_mul_0), idx_mask);
+            idx.vec[1] = _mm_and_si128(_mm_mullo_epi32(idx.vec[1], idx_mul_1), idx_mask);
+            idx.vec[2] = _mm_and_si128(_mm_mullo_epi32(idx.vec[2], idx_mul_0), idx_mask);
+            idx.vec[3] = _mm_and_si128(_mm_mullo_epi32(idx.vec[3], idx_mul_1), idx_mask);
+
+            idx.vec[0] = _mm_or_si128(idx.vec[0], _mm_cvtepi16_epi32(idx_l_0));
+            idx.vec[1] = _mm_or_si128(idx.vec[1], _mm_cvtepi16_epi32(_mm_srli_si128(idx_l_0, 8)));
+            idx.vec[2] = _mm_or_si128(idx.vec[2], _mm_cvtepi16_epi32(idx_l_1));
+            idx.vec[3] = _mm_or_si128(idx.vec[3], _mm_cvtepi16_epi32(_mm_srli_si128(idx_l_1, 8)));
+
+            const __m128i q2_1_0 = _mm_set_epi32(iq3s_grid[idx.index[3]], iq3s_grid[idx.index[2]], iq3s_grid[idx.index[1]], iq3s_grid[idx.index[0]]);
+            const __m128i q2_1_1 = _mm_set_epi32(iq3s_grid[idx.index[7]], iq3s_grid[idx.index[6]], iq3s_grid[idx.index[5]], iq3s_grid[idx.index[4]]);
+            const __m128i q2_2_0 = _mm_set_epi32(iq3s_grid[idx.index[11]], iq3s_grid[idx.index[10]], iq3s_grid[idx.index[9]], iq3s_grid[idx.index[8]]);
+            const __m128i q2_2_1 = _mm_set_epi32(iq3s_grid[idx.index[15]], iq3s_grid[idx.index[14]], iq3s_grid[idx.index[13]], iq3s_grid[idx.index[12]]);
+
+            __m128i aux128_0 = _mm_set1_epi32(signs[0] | (signs[1] << 16));
+            __m128i aux128_1 = aux128_0;
+            aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0);
+            aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1);
+            const __m128i s2_1_0 = _mm_cmpeq_epi8(aux128_0, mask2_0);
+            const __m128i s2_1_1 = _mm_cmpeq_epi8(aux128_1, mask2_1);
+            const __m128i q8s_1_0 = _mm_sub_epi8(_mm_xor_si128(s2_1_0, q8_1_0), s2_1_0);
+            const __m128i q8s_1_1 = _mm_sub_epi8(_mm_xor_si128(s2_1_1, q8_1_1), s2_1_1);
+
+            aux128_0 = _mm_set1_epi32(signs[2] | (signs[3] << 16));
+            aux128_1 = aux128_0;
+            aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0);
+            aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1);
+            const __m128i s2_2_0 = _mm_cmpeq_epi8(aux128_0, mask2_0);
+            const __m128i s2_2_1 = _mm_cmpeq_epi8(aux128_1, mask2_1);
+            const __m128i q8s_2_0 = _mm_sub_epi8(_mm_xor_si128(s2_2_0, q8_2_0), s2_2_0);
+            const __m128i q8s_2_1 = _mm_sub_epi8(_mm_xor_si128(s2_2_1, q8_2_1), s2_2_1);
+
+            signs += 4;
+
+            const __m128i dot1_0  = _mm_maddubs_epi16(q2_1_0, q8s_1_0);
+            const __m128i dot1_1  = _mm_maddubs_epi16(q2_1_1, q8s_1_1);
+            const __m128i dot2_0  = _mm_maddubs_epi16(q2_2_0, q8s_2_0);
+            const __m128i dot2_1  = _mm_maddubs_epi16(q2_2_1, q8s_2_1);
+            const uint16_t ls1 = x[i].scales[ib32/2] & 0xf;
+            const uint16_t ls2 = x[i].scales[ib32/2] >>  4;
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1));
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1));
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1));
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1));
+            sumi1_0 = _mm_add_epi32(sumi1_0, p1_0);
+            sumi1_1 = _mm_add_epi32(sumi1_1, p1_1);
+            sumi2_0 = _mm_add_epi32(sumi2_0, p2_0);
+            sumi2_1 = _mm_add_epi32(sumi2_1, p2_1);
+        }
+
+        accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf);
+
+    }
+
+    *s = hsum_float_8(accumf);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq3_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq1_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq1_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined __AVX2__
+
+    __m256 accum = _mm256_setzero_ps();
+    float accum1 = 0;
+    for (int i = 0; i < nb; ++i) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint16_t * qh = x[i].qh;
+
+        __m256i sumi = _mm256_setzero_si256();
+        int sumi1 = 0;
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+#ifdef __BMI2__
+            const uint64_t packed_idx1 = _pdep_u64(*(const uint32_t *)qs, 0x00ff00ff00ff00ffULL) | _pdep_u64(qh[ib], 0x700070007000700ULL);
+            const uint64_t packed_idx2 = _pdep_u64(*(const uint32_t *)(qs + 4), 0x00ff00ff00ff00ffULL) | _pdep_u64(qh[ib + 1], 0x700070007000700ULL);
+            const uint16_t *idx1 = (const uint16_t *)(&packed_idx1);
+            const uint16_t *idx2 = (const uint16_t *)(&packed_idx2);
+            const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[idx1[3]], iq1s_grid[idx1[2]], iq1s_grid[idx1[1]], iq1s_grid[idx1[0]]);
+            const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[idx2[3]], iq1s_grid[idx2[2]], iq1s_grid[idx2[1]], iq1s_grid[idx2[0]]);
+#else
+            const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[qs[3] | ((qh[ib+0] >> 1) & 0x700)], iq1s_grid[qs[2] | ((qh[ib+0] << 2) & 0x700)],
+                                                    iq1s_grid[qs[1] | ((qh[ib+0] << 5) & 0x700)], iq1s_grid[qs[0] | ((qh[ib+0] << 8) & 0x700)]);
+            const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[qs[7] | ((qh[ib+1] >> 1) & 0x700)], iq1s_grid[qs[6] | ((qh[ib+1] << 2) & 0x700)],
+                                                    iq1s_grid[qs[5] | ((qh[ib+1] << 5) & 0x700)], iq1s_grid[qs[4] | ((qh[ib+1] << 8) & 0x700)]);
+#endif
+            qs += 8;
+            const __m256i q8b_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8b_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            const __m256i dot1 = mul_add_epi8(q1b_1, q8b_1);
+            const __m256i dot2 = mul_add_epi8(q1b_2, q8b_2);
+            const int16_t ls1 = 2*((qh[ib+0] >> 12) & 7) + 1;
+            const int16_t ls2 = 2*((qh[ib+1] >> 12) & 7) + 1;
+            const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(ls1));
+            const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(ls2));
+
+            sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p1, p2));
+            sumi1 += (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]) * (qh[ib+0] & 0x8000 ? -1 : 1) * ls1
+                   + (y[i].bsums[2*ib+2] + y[i].bsums[2*ib+3]) * (qh[ib+1] & 0x8000 ? -1 : 1) * ls2;
+        }
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        accum = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(sumi), accum);
+        accum1 += d * sumi1;
+
+    }
+
+    *s = hsum_float_8(accum) + IQ1S_DELTA * accum1;
+
+#elif defined __AVX__
+    __m256 accum = _mm256_setzero_ps();
+    float accum1 = 0;
+    for (int i = 0; i < nb; ++i) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint16_t * qh = x[i].qh;
+
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        int sumi1 = 0;
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+            const __m128i q1b_1_0 = _mm_set_epi64x(iq1s_grid[qs[1] | ((qh[ib+0] << 5) & 0x700)], iq1s_grid[qs[0] | ((qh[ib+0] << 8) & 0x700)]);
+            const __m128i q1b_1_1 = _mm_set_epi64x(iq1s_grid[qs[3] | ((qh[ib+0] >> 1) & 0x700)], iq1s_grid[qs[2] | ((qh[ib+0] << 2) & 0x700)]);
+            const __m128i q1b_2_0 = _mm_set_epi64x(iq1s_grid[qs[5] | ((qh[ib+1] << 5) & 0x700)], iq1s_grid[qs[4] | ((qh[ib+1] << 8) & 0x700)]);
+            const __m128i q1b_2_1 = _mm_set_epi64x(iq1s_grid[qs[7] | ((qh[ib+1] >> 1) & 0x700)], iq1s_grid[qs[6] | ((qh[ib+1] << 2) & 0x700)]);
+            qs += 8;
+            const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+
+            const __m128i dot1_0 = mul_add_epi8_sse(q1b_1_0, q8b_1_0);
+            const __m128i dot1_1 = mul_add_epi8_sse(q1b_1_1, q8b_1_1);
+            const __m128i dot2_0 = mul_add_epi8_sse(q1b_2_0, q8b_2_0);
+            const __m128i dot2_1 = mul_add_epi8_sse(q1b_2_1, q8b_2_1);
+            const int16_t ls1 = 2*((qh[ib+0] >> 12) & 7) + 1;
+            const int16_t ls2 = 2*((qh[ib+1] >> 12) & 7) + 1;
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(ls1));
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(ls1));
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(ls2));
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(ls2));
+
+            sumi1_0 = _mm_add_epi32(sumi1_0, _mm_add_epi32(p1_0, p2_0));
+            sumi1_1 = _mm_add_epi32(sumi1_1, _mm_add_epi32(p1_1, p2_1));
+            sumi1 += (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]) * (qh[ib+0] & 0x8000 ? -1 : 1) * ls1
+                   + (y[i].bsums[2*ib+2] + y[i].bsums[2*ib+3]) * (qh[ib+1] & 0x8000 ? -1 : 1) * ls2;
+        }
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        accum = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(sumi1_1, sumi1_0))), accum);
+        accum1 += d * sumi1;
+
+    }
+
+    *s = hsum_float_8(accum) + IQ1S_DELTA * accum1;
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq1_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq1_m_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq1_m * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    iq1m_scale_t scale;
+
+#if defined __AVX2__
+
+    const __m256i mask = _mm256_set1_epi16(0x7);
+    const __m256i mone = _mm256_set1_epi16(1);
+    const __m256i mone8 = _mm256_set1_epi8(1);
+    const __m256i mtwo8 = _mm256_set1_epi8(2);
+    // VPSHUFB cannot cross 128-bit lanes so odd shifts go to upper half.
+    const __m256i scales_shift = _mm256_set_epi64x(9, 3, 6, 0);
+
+    __m256 accum1 = _mm256_setzero_ps();
+    __m256 accum2 = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint8_t  * qh = x[i].qh;
+        const uint16_t * sc = (const uint16_t *)x[i].scales;
+
+        scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+        // Extract 3-bit scales (16 values)
+        __m256i scales = _mm256_set1_epi64x(*(const uint64_t*)sc);
+        scales = _mm256_srlv_epi64(scales, scales_shift);
+        scales = _mm256_add_epi16(_mm256_slli_epi16(_mm256_and_si256(scales, mask), 1), mone);
+
+        // Indices to repeat each scale 8 times.
+        __m256i scales_idx1 = _mm256_set1_epi16(0x0100);
+        __m256i scales_idx2 = _mm256_add_epi8(scales_idx1, _mm256_set1_epi8(8));
+
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+#ifdef __BMI2__
+            const uint64_t packed_idx1 = _pdep_u64(*(const uint32_t *)qs, 0x00ff00ff00ff00ffULL)
+                                       | _pdep_u64(*(const uint16_t*)(qh) & 0x7777, 0xf000f000f000f00ULL);
+            const uint64_t packed_idx2 = _pdep_u64(*(const uint32_t *)(qs + 4), 0x00ff00ff00ff00ffULL)
+                                       | _pdep_u64(*(const uint16_t*)(qh + 2) & 0x7777, 0xf000f000f000f00ULL);
+            const uint16_t *idx1 = (const uint16_t *)(&packed_idx1);
+            const uint16_t *idx2 = (const uint16_t *)(&packed_idx2);
+            const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[idx1[3]], iq1s_grid[idx1[2]], iq1s_grid[idx1[1]], iq1s_grid[idx1[0]]);
+            const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[idx2[3]], iq1s_grid[idx2[2]], iq1s_grid[idx2[1]], iq1s_grid[idx2[0]]);
+
+            // Convert signs to bytes 0x81 (negative) or 0x01 (positive)
+            const uint64_t delta_sign = _pdep_u64(*(const uint32_t*)(qh) & 0x88888888, 0xf0f0f0f0f0f0f0f0ULL);
+            const __m256i delta1 = _mm256_or_si256(mone8, _mm256_cvtepi8_epi64(_mm_set1_epi32(delta_sign)));
+            const __m256i delta2 = _mm256_or_si256(mone8, _mm256_cvtepi8_epi64(_mm_set1_epi32(delta_sign >> 32)));
+#else
+            const __m256i q1b_1 = _mm256_set_epi64x(
+                    iq1s_grid[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)], iq1s_grid[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)],
+                    iq1s_grid[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)], iq1s_grid[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)]
+            );
+            const __m256i q1b_2 = _mm256_set_epi64x(
+                    iq1s_grid[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)], iq1s_grid[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)],
+                    iq1s_grid[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)], iq1s_grid[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)]
+            );
+
+            const __m256i delta1 = _mm256_set_epi64x(qh[1] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[1] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[0] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[0] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+            const __m256i delta2 = _mm256_set_epi64x(qh[3] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[3] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[2] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[2] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+#endif
+            const __m256i q8b_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+            const __m256i q8b_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
+
+            const __m256i dot1 = mul_add_epi8(q1b_1, q8b_1);
+            const __m256i dot2 = mul_add_epi8(q1b_2, q8b_2);
+            const __m256i dot3 = _mm256_maddubs_epi16(mone8, _mm256_sign_epi8(q8b_1, delta1));
+            const __m256i dot4 = _mm256_maddubs_epi16(mone8, _mm256_sign_epi8(q8b_2, delta2));
+
+            __m256i scale1 = _mm256_shuffle_epi8(scales, scales_idx1);
+            __m256i scale2 = _mm256_shuffle_epi8(scales, scales_idx2);
+
+            scales_idx1 = _mm256_add_epi8(scales_idx1, mtwo8);
+            scales_idx2 = _mm256_add_epi8(scales_idx2, mtwo8);
+
+            const __m256i p1 = _mm256_madd_epi16(dot1, scale1);
+            const __m256i p2 = _mm256_madd_epi16(dot2, scale2);
+            const __m256i p3 = _mm256_madd_epi16(dot3, scale1);
+            const __m256i p4 = _mm256_madd_epi16(dot4, scale2);
+
+            sumi1 = _mm256_add_epi32(sumi1, _mm256_add_epi32(p1, p2));
+            sumi2 = _mm256_add_epi32(sumi2, _mm256_add_epi32(p3, p4));
+
+            qs += 8; qh += 4;
+        }
+
+        const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(scale.f16));
+
+        accum1 = _mm256_fmadd_ps(d, _mm256_cvtepi32_ps(sumi1), accum1);
+        accum2 = _mm256_fmadd_ps(d, _mm256_cvtepi32_ps(sumi2), accum2);
+    }
+
+    *s = hsum_float_8(accum1) + IQ1M_DELTA * hsum_float_8(accum2);
+
+#elif defined __AVX__
+    const __m128i mask = _mm_set1_epi16(0x7);
+    const __m128i mone = _mm_set1_epi16(1);
+
+    __m256 accum1 = _mm256_setzero_ps();
+    __m256 accum2 = _mm256_setzero_ps();
+    for (int i = 0; i < nb; ++i) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint8_t  * qh = x[i].qh;
+        const uint16_t * sc = (const uint16_t *)x[i].scales;
+
+        scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+            const __m128i q1b_1_0 = _mm_set_epi64x(
+                    iq1s_grid[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)], iq1s_grid[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)]);
+            const __m128i q1b_1_1 = _mm_set_epi64x(
+                    iq1s_grid[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)], iq1s_grid[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)]);
+            const __m128i q1b_2_0 = _mm_set_epi64x(
+                    iq1s_grid[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)], iq1s_grid[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)]);
+            const __m128i q1b_2_1 = _mm_set_epi64x(
+                    iq1s_grid[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)], iq1s_grid[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)]);
+            const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+
+            const __m128i dot1_0 = mul_add_epi8_sse(q1b_1_0, q8b_1_0);
+            const __m128i dot1_1 = mul_add_epi8_sse(q1b_1_1, q8b_1_1);
+            const __m128i dot2_0 = mul_add_epi8_sse(q1b_2_0, q8b_2_0);
+            const __m128i dot2_1 = mul_add_epi8_sse(q1b_2_1, q8b_2_1);
+
+            const __m128i delta1_0 = _mm_set_epi64x(qh[0] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[0] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+            const __m128i delta1_1 = _mm_set_epi64x(qh[1] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[1] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+            const __m128i delta2_0 = _mm_set_epi64x(qh[2] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[2] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+            const __m128i delta2_1 = _mm_set_epi64x(qh[3] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101,
+                                                     qh[3] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101);
+
+            const __m128i dot3_0 = mul_add_epi8_sse(delta1_0, q8b_1_0);
+            const __m128i dot3_1 = mul_add_epi8_sse(delta1_1, q8b_1_1);
+            const __m128i dot4_0 = mul_add_epi8_sse(delta2_0, q8b_2_0);
+            const __m128i dot4_1 = mul_add_epi8_sse(delta2_1, q8b_2_1);
+
+            __m128i scale1_0 = _mm_set1_epi16(sc[ib/2] >> 0);
+            __m128i scale1_1 = _mm_set1_epi16(sc[ib/2] >> 3);
+            __m128i scale2_0 = _mm_set1_epi16(sc[ib/2] >> 6);
+            __m128i scale2_1 = _mm_set1_epi16(sc[ib/2] >> 9);
+
+            scale1_0 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale1_0, mask), 1), mone);
+            scale1_1 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale1_1, mask), 1), mone);
+            scale2_0 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale2_0, mask), 1), mone);
+            scale2_1 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale2_1, mask), 1), mone);
+            const __m128i p1_0 = _mm_madd_epi16(dot1_0, scale1_0);
+            const __m128i p1_1 = _mm_madd_epi16(dot1_1, scale1_1);
+            const __m128i p2_0 = _mm_madd_epi16(dot2_0, scale2_0);
+            const __m128i p2_1 = _mm_madd_epi16(dot2_1, scale2_1);
+            const __m128i p3_0 = _mm_madd_epi16(dot3_0, scale1_0);
+            const __m128i p3_1 = _mm_madd_epi16(dot3_1, scale1_1);
+            const __m128i p4_0 = _mm_madd_epi16(dot4_0, scale2_0);
+            const __m128i p4_1 = _mm_madd_epi16(dot4_1, scale2_1);
+
+            sumi1_0 = _mm_add_epi32(sumi1_0, _mm_add_epi32(p1_0, p2_0));
+            sumi1_1 = _mm_add_epi32(sumi1_1, _mm_add_epi32(p1_1, p2_1));
+            sumi2_0 = _mm_add_epi32(sumi2_0, _mm_add_epi32(p3_0, p4_0));
+            sumi2_1 = _mm_add_epi32(sumi2_1, _mm_add_epi32(p3_1, p4_1));
+
+            qs += 8; qh += 4;
+        }
+
+        const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(scale.f16));
+
+        accum1 = _mm256_add_ps(_mm256_mul_ps(d, _mm256_cvtepi32_ps(MM256_SET_M128I(sumi1_1, sumi1_0))), accum1);
+        accum2 = _mm256_add_ps(_mm256_mul_ps(d, _mm256_cvtepi32_ps(MM256_SET_M128I(sumi2_1, sumi2_0))), accum2);
+    }
+
+    *s = hsum_float_8(accum1) + IQ1M_DELTA * hsum_float_8(accum2);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    UNUSED(scale);
+    ggml_vec_dot_iq1_m_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
+void ggml_vec_dot_iq4_nl_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK4_NL == 0);
+    static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same");
+
+    const block_iq4_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK4_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __AVX2__
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+    const __m256i mone = _mm256_set1_epi16(1);
+
+    __m256 accum1 = _mm256_setzero_ps();
+    __m256 accum2 = _mm256_setzero_ps();
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)x[ib + 0].qs);
+        const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)x[ib + 1].qs);
+        const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)y[ib + 0].qs);
+        const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)y[ib + 1].qs);
+        const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)),
+                                              _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)));
+        const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)),
+                                              _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)));
+        const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1);
+        const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2);
+        const __m256i p_1 = _mm256_madd_epi16(p16_1, mone);
+        const __m256i p_2 = _mm256_madd_epi16(p16_2, mone);
+        accum1 = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 0].d)*GGML_CPU_FP16_TO_FP32(x[ib + 0].d)),
+                _mm256_cvtepi32_ps(p_1), accum1);
+        accum2 = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 1].d)*GGML_CPU_FP16_TO_FP32(x[ib + 1].d)),
+                _mm256_cvtepi32_ps(p_2), accum2);
+    }
+
+    sumf = hsum_float_8(_mm256_add_ps(accum1, accum2));
+
+#elif defined __AVX__
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (; ib + 1 < nb; ib += 2) {
+        const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs);
+        const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs);
+        const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs);
+        const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1);
+        const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs);
+        const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1);
+
+        const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b));
+        const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b));
+        const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b));
+        const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b));
+
+        const __m256 p = mul_sum_i8_quad_float(q4b_1_0, q4b_1_1, q4b_2_0, q4b_2_1, q8b_1_0, q8b_1_1, q8b_2_0, q8b_2_1);
+        const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d);
+        accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi1 = 0, sumi2 = 0;
+        for (int j = 0; j < QK4_NL/2; ++j) {
+            sumi1 += y[ib].qs[j+       0] * kvalues_iq4nl[x[ib].qs[j] & 0xf];
+            sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >>  4];
+        }
+        sumf += d * (sumi1 + sumi2);
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_K == 0);
+
+    const block_iq4_xs * GGML_RESTRICT x = vx;
+    const block_q8_K   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+#if defined __AVX2__
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (int ibl = 0; ibl < nb; ++ibl) {
+        const uint8_t * qs = x[ibl].qs;
+        const int8_t  * q8 = y[ibl].qs;
+        uint16_t sh = x[ibl].scales_h;
+        __m256i sumi1 = _mm256_setzero_si256();
+        __m256i sumi2 = _mm256_setzero_si256();
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+            const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)qs);  qs += 16;
+            const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)qs);  qs += 16;
+            const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32;
+            const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)),
+                                                  _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)));
+            const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)),
+                                                  _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)));
+            const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1);
+            const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2);
+            const int16_t ls1 = ((x[ibl].scales_l[ib/2] & 0xf) | ((sh << 4) & 0x30)) - 32;
+            const int16_t ls2 = ((x[ibl].scales_l[ib/2] >>  4) | ((sh << 2) & 0x30)) - 32;
+            sh >>= 4;
+            const __m256i p_1 = _mm256_madd_epi16(p16_1, _mm256_set1_epi16(ls1));
+            const __m256i p_2 = _mm256_madd_epi16(p16_2, _mm256_set1_epi16(ls2));
+            sumi1 = _mm256_add_epi32(p_1, sumi1);
+            sumi2 = _mm256_add_epi32(p_2, sumi2);
+        }
+        accum = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ibl].d)*y[ibl].d),
+                _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accum);
+    }
+
+    *s = hsum_float_8(accum);
+
+#elif defined __AVX__
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (int ibl = 0; ibl < nb; ++ibl) {
+        const uint8_t * qs = x[ibl].qs;
+        const int8_t  * q8 = y[ibl].qs;
+        uint16_t sh = x[ibl].scales_h;
+        __m128i sumi1_0 = _mm_setzero_si128();
+        __m128i sumi1_1 = _mm_setzero_si128();
+        __m128i sumi2_0 = _mm_setzero_si128();
+        __m128i sumi2_1 = _mm_setzero_si128();
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+            const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)qs); qs += 16;
+            const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)qs); qs += 16;
+            const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16;
+            const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b));
+            const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b));
+            const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b));
+            const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b));
+            const __m128i p16_1_0 = mul_add_epi8_sse(q4b_1_0, q8b_1_0);
+            const __m128i p16_1_1 = mul_add_epi8_sse(q4b_1_1, q8b_1_1);
+            const __m128i p16_2_0 = mul_add_epi8_sse(q4b_2_0, q8b_2_0);
+            const __m128i p16_2_1 = mul_add_epi8_sse(q4b_2_1, q8b_2_1);
+            const int16_t ls1 = ((x[ibl].scales_l[ib/2] & 0xf) | ((sh << 4) & 0x30)) - 32;
+            const int16_t ls2 = ((x[ibl].scales_l[ib/2] >>  4) | ((sh << 2) & 0x30)) - 32;
+            sh >>= 4;
+            const __m128i p_1_0 = _mm_madd_epi16(p16_1_0, _mm_set1_epi16(ls1));
+            const __m128i p_1_1 = _mm_madd_epi16(p16_1_1, _mm_set1_epi16(ls1));
+            const __m128i p_2_0 = _mm_madd_epi16(p16_2_0, _mm_set1_epi16(ls2));
+            const __m128i p_2_1 = _mm_madd_epi16(p16_2_1, _mm_set1_epi16(ls2));
+            sumi1_0 = _mm_add_epi32(p_1_0, sumi1_0);
+            sumi1_1 = _mm_add_epi32(p_1_1, sumi1_1);
+            sumi2_0 = _mm_add_epi32(p_2_0, sumi2_0);
+            sumi2_1 = _mm_add_epi32(p_2_1, sumi2_1);
+        }
+        __m128i sumi12_0 = _mm_add_epi32(sumi1_0, sumi2_0);
+        __m128i sumi12_1 = _mm_add_epi32(sumi1_1, sumi2_1);
+        accum = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ibl].d)*y[ibl].d),
+                _mm256_cvtepi32_ps(MM256_SET_M128I(sumi12_1, sumi12_0))), accum);
+    }
+
+    *s = hsum_float_8(accum);
+
+#else
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(nb);
+    ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h
new file mode 100644
index 0000000000..82345b2063
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h
@@ -0,0 +1,538 @@
+#pragma once
+
+// GGML CPU internal header
+
+#include "ggml.h"
+#include "ggml-impl.h"
+
+#include <stdlib.h> // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/
+//#include <stddef.h>
+#include <stdbool.h>
+#include <string.h> // memcpy
+#include <math.h>   // fabsf
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct ggml_compute_params {
+    // ith = thread index, nth = number of threads
+    int ith, nth;
+
+    // work buffer for all threads
+    size_t wsize;
+    void * wdata;
+
+    struct ggml_threadpool * threadpool;
+
+    // use reference implementation
+    bool use_ref;
+};
+
+
+#if defined(_MSC_VER)
+
+#define m512bh(p) p
+#define m512i(p) p
+
+#else
+
+#define m512bh(p) (__m512bh)(p)
+#define m512i(p) (__m512i)(p)
+
+#endif
+
+// __FMA__ and __F16C__ are not defined in MSVC, however they are implied with AVX2/AVX512
+#if defined(_MSC_VER) && (defined(__AVX2__) || defined(__AVX512F__))
+#ifndef __FMA__
+#define __FMA__
+#endif
+#ifndef __F16C__
+#define __F16C__
+#endif
+#endif
+
+// __SSE3__ and __SSSE3__ are not defined in MSVC, but SSE3/SSSE3 are present when AVX/AVX2/AVX512 are available
+#if defined(_MSC_VER) && (defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__))
+#ifndef __SSE3__
+#define __SSE3__
+#endif
+#ifndef __SSSE3__
+#define __SSSE3__
+#endif
+#endif
+
+#if defined(__s390x__) && defined(__VEC__)
+#ifndef __VXE__
+#define __VXE__
+#endif  // __VXE__
+#ifndef __VXE2__
+#define __VXE2__
+#endif  // __VXE2__
+#endif  // __s390x__ && __VEC__
+
+#if defined(__ARM_FEATURE_SVE) && defined(__linux__)
+#include <sys/prctl.h>
+#endif
+
+#if defined(__ARM_NEON)
+
+// ref: https://github.com/ggml-org/llama.cpp/pull/5404
+#if defined(_MSC_VER) && !defined(__clang__)
+#define ggml_vld1q_u32(w,x,y,z) { ((w) + ((uint64_t)(x) << 32)), ((y) + ((uint64_t)(z) << 32)) }
+#else
+#define ggml_vld1q_u32(w,x,y,z) { (w), (x), (y), (z) }
+#endif // _MSC_VER
+
+#if !defined(__aarch64__)
+
+// 32-bit ARM compatibility
+
+// vaddlvq_s16
+// vpaddq_s16
+// vpaddq_s32
+// vaddvq_s32
+// vaddvq_f32
+// vmaxvq_f32
+// vcvtnq_s32_f32
+// vzip1_u8
+// vzip2_u8
+
+inline static int32_t vaddlvq_s16(int16x8_t v) {
+    int32x4_t v0 = vreinterpretq_s32_s64(vpaddlq_s32(vpaddlq_s16(v)));
+    return vgetq_lane_s32(v0, 0) + vgetq_lane_s32(v0, 2);
+}
+
+inline static int16x8_t vpaddq_s16(int16x8_t a, int16x8_t b) {
+    int16x4_t a0 = vpadd_s16(vget_low_s16(a), vget_high_s16(a));
+    int16x4_t b0 = vpadd_s16(vget_low_s16(b), vget_high_s16(b));
+    return vcombine_s16(a0, b0);
+}
+
+inline static int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) {
+    int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a));
+    int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b));
+    return vcombine_s32(a0, b0);
+}
+
+inline static int32_t vaddvq_s32(int32x4_t v) {
+    return vgetq_lane_s32(v, 0) + vgetq_lane_s32(v, 1) + vgetq_lane_s32(v, 2) + vgetq_lane_s32(v, 3);
+}
+
+inline static float vaddvq_f32(float32x4_t v) {
+    return vgetq_lane_f32(v, 0) + vgetq_lane_f32(v, 1) + vgetq_lane_f32(v, 2) + vgetq_lane_f32(v, 3);
+}
+
+inline static float vmaxvq_f32(float32x4_t v) {
+    return
+        MAX(MAX(vgetq_lane_f32(v, 0), vgetq_lane_f32(v, 1)),
+            MAX(vgetq_lane_f32(v, 2), vgetq_lane_f32(v, 3)));
+}
+
+inline static int32x4_t vcvtnq_s32_f32(float32x4_t v) {
+    int32x4_t res;
+
+    res[0] = roundf(vgetq_lane_f32(v, 0));
+    res[1] = roundf(vgetq_lane_f32(v, 1));
+    res[2] = roundf(vgetq_lane_f32(v, 2));
+    res[3] = roundf(vgetq_lane_f32(v, 3));
+
+    return res;
+}
+
+inline static uint8x8_t vzip1_u8(uint8x8_t a, uint8x8_t b) {
+    uint8x8_t res;
+
+    res[0] = a[0]; res[1] = b[0];
+    res[2] = a[1]; res[3] = b[1];
+    res[4] = a[2]; res[5] = b[2];
+    res[6] = a[3]; res[7] = b[3];
+
+    return res;
+}
+
+inline static uint8x8_t vzip2_u8(uint8x8_t a, uint8x8_t b) {
+    uint8x8_t res;
+
+    res[0] = a[4]; res[1] = b[4];
+    res[2] = a[5]; res[3] = b[5];
+    res[4] = a[6]; res[5] = b[6];
+    res[6] = a[7]; res[7] = b[7];
+
+    return res;
+}
+
+// vld1q_s16_x2
+// vld1q_u8_x2
+// vld1q_u8_x4
+// vld1q_s8_x2
+// vld1q_s8_x4
+// TODO: double-check these work correctly
+
+typedef struct ggml_int16x8x2_t {
+    int16x8_t val[2];
+} ggml_int16x8x2_t;
+
+inline static ggml_int16x8x2_t ggml_vld1q_s16_x2(const int16_t * ptr) {
+    ggml_int16x8x2_t res;
+
+    res.val[0] = vld1q_s16(ptr + 0);
+    res.val[1] = vld1q_s16(ptr + 8);
+
+    return res;
+}
+
+typedef struct ggml_uint8x16x2_t {
+    uint8x16_t val[2];
+} ggml_uint8x16x2_t;
+
+inline static ggml_uint8x16x2_t ggml_vld1q_u8_x2(const uint8_t * ptr) {
+    ggml_uint8x16x2_t res;
+
+    res.val[0] = vld1q_u8(ptr + 0);
+    res.val[1] = vld1q_u8(ptr + 16);
+
+    return res;
+}
+
+typedef struct ggml_uint8x16x4_t {
+    uint8x16_t val[4];
+} ggml_uint8x16x4_t;
+
+inline static ggml_uint8x16x4_t ggml_vld1q_u8_x4(const uint8_t * ptr) {
+    ggml_uint8x16x4_t res;
+
+    res.val[0] = vld1q_u8(ptr + 0);
+    res.val[1] = vld1q_u8(ptr + 16);
+    res.val[2] = vld1q_u8(ptr + 32);
+    res.val[3] = vld1q_u8(ptr + 48);
+
+    return res;
+}
+
+typedef struct ggml_int8x16x2_t {
+    int8x16_t val[2];
+} ggml_int8x16x2_t;
+
+inline static ggml_int8x16x2_t ggml_vld1q_s8_x2(const int8_t * ptr) {
+    ggml_int8x16x2_t res;
+
+    res.val[0] = vld1q_s8(ptr + 0);
+    res.val[1] = vld1q_s8(ptr + 16);
+
+    return res;
+}
+
+typedef struct ggml_int8x16x4_t {
+    int8x16_t val[4];
+} ggml_int8x16x4_t;
+
+inline static ggml_int8x16x4_t ggml_vld1q_s8_x4(const int8_t * ptr) {
+    ggml_int8x16x4_t res;
+
+    res.val[0] = vld1q_s8(ptr + 0);
+    res.val[1] = vld1q_s8(ptr + 16);
+    res.val[2] = vld1q_s8(ptr + 32);
+    res.val[3] = vld1q_s8(ptr + 48);
+
+    return res;
+}
+
+// NOTE: not tested
+inline static int8x16_t ggml_vqtbl1q_s8(int8x16_t a, uint8x16_t b) {
+    int8x16_t res;
+
+    res[ 0] = a[b[ 0]];
+    res[ 1] = a[b[ 1]];
+    res[ 2] = a[b[ 2]];
+    res[ 3] = a[b[ 3]];
+    res[ 4] = a[b[ 4]];
+    res[ 5] = a[b[ 5]];
+    res[ 6] = a[b[ 6]];
+    res[ 7] = a[b[ 7]];
+    res[ 8] = a[b[ 8]];
+    res[ 9] = a[b[ 9]];
+    res[10] = a[b[10]];
+    res[11] = a[b[11]];
+    res[12] = a[b[12]];
+    res[13] = a[b[13]];
+    res[14] = a[b[14]];
+    res[15] = a[b[15]];
+
+    return res;
+}
+
+// NOTE: not tested
+inline static uint8x16_t ggml_vqtbl1q_u8(uint8x16_t a, uint8x16_t b) {
+    uint8x16_t res;
+
+    res[ 0] = a[b[ 0]];
+    res[ 1] = a[b[ 1]];
+    res[ 2] = a[b[ 2]];
+    res[ 3] = a[b[ 3]];
+    res[ 4] = a[b[ 4]];
+    res[ 5] = a[b[ 5]];
+    res[ 6] = a[b[ 6]];
+    res[ 7] = a[b[ 7]];
+    res[ 8] = a[b[ 8]];
+    res[ 9] = a[b[ 9]];
+    res[10] = a[b[10]];
+    res[11] = a[b[11]];
+    res[12] = a[b[12]];
+    res[13] = a[b[13]];
+    res[14] = a[b[14]];
+    res[15] = a[b[15]];
+
+    return res;
+}
+
+#else
+
+#define ggml_int16x8x2_t  int16x8x2_t
+#define ggml_uint8x16x2_t uint8x16x2_t
+#define ggml_uint8x16x4_t uint8x16x4_t
+#define ggml_int8x16x2_t  int8x16x2_t
+#define ggml_int8x16x4_t  int8x16x4_t
+
+#define ggml_vld1q_s16_x2 vld1q_s16_x2
+#define ggml_vld1q_u8_x2  vld1q_u8_x2
+#define ggml_vld1q_u8_x4  vld1q_u8_x4
+#define ggml_vld1q_s8_x2  vld1q_s8_x2
+#define ggml_vld1q_s8_x4  vld1q_s8_x4
+#define ggml_vqtbl1q_s8   vqtbl1q_s8
+#define ggml_vqtbl1q_u8   vqtbl1q_u8
+
+#endif // !defined(__aarch64__)
+
+#if !defined(__ARM_FEATURE_DOTPROD)
+
+// NOTE: this fallback produces the same total sum as native vdotq_s32 but with different per-lane grouping — do not use when individual lane values matter.
+inline static int32x4_t ggml_vdotq_s32(int32x4_t acc, int8x16_t a, int8x16_t b) {
+    const int16x8_t p0 = vmull_s8(vget_low_s8 (a), vget_low_s8 (b));
+    const int16x8_t p1 = vmull_s8(vget_high_s8(a), vget_high_s8(b));
+
+    return vaddq_s32(acc, vaddq_s32(vpaddlq_s16(p0), vpaddlq_s16(p1)));
+}
+
+#else
+
+#define ggml_vdotq_s32(a, b, c) vdotq_s32(a, b, c)
+
+#endif // !defined(__ARM_FEATURE_DOTPROD)
+
+static inline int32x4_t ggml_nvfp4_dot8(const int8x8_t q4_lo, const int8x8_t q8_lo,
+                                         const int8x8_t q4_hi, const int8x8_t q8_hi) {
+    const int16x8_t p_lo = vmull_s8(q4_lo, q8_lo);
+    const int16x8_t p_hi = vmull_s8(q4_hi, q8_hi);
+    const int32x4_t sum_lo = vpaddlq_s16(p_lo);
+    const int32x4_t sum_hi = vpaddlq_s16(p_hi);
+    return vaddq_s32(sum_lo, sum_hi);
+}
+
+#endif // defined(__ARM_NEON)
+
+#ifdef __wasm_simd128__
+#include <wasm_simd128.h>
+#endif
+
+#ifdef __POWER9_VECTOR__
+#include <altivec.h>
+#endif
+
+#if defined(_MSC_VER) || defined(__MINGW32__)
+#include <intrin.h>
+#elif defined(__SSE__) || defined(__SSE3__) || defined(__SSSE3__) || defined(__AVX__) || defined(__F16C__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX512BF16__)
+#include <immintrin.h>
+#endif
+
+#ifdef __riscv_v_intrinsic
+#include <riscv_vector.h>
+#endif
+
+#if defined(__loongarch64)
+#if defined(__loongarch_asx)
+#include <lasxintrin.h>
+#endif
+#if defined(__loongarch_sx)
+#include <lsxintrin.h>
+#endif
+#endif
+
+#if defined(__VXE__) || defined(__VXE2__)
+#include <vecintrin.h>
+
+#define vec_neg(a)    (-(a))                // Vector Negate
+#define vec_add(a, b) ((a) + (b))           // Vector Add
+#define vec_sub(a, b) ((a) - (b))           // Vector Subtract
+#define vec_mul(a, b) ((a) * (b))           // Vector Multiply
+#define vec_div(a, b) ((a) / (b))           // Vector Divide
+#define vec_sl(a, b)  ((a) << (b))          // Vector Shift Left
+#define vec_sra(a, b) ((a) >> (b))          // Vector Shift Right
+#define vec_sr(a, b)  ((a) >> (b))          // Vector Shift Right Algebraic
+#define vec_slo(a, b) vec_slb(a, (b) << 64) // Vector Shift Left by Octet
+#define vec_sro(a, b) vec_srb(a, (b) << 64) // Vector Shift Right by Octet
+
+#ifndef vec_and
+#define vec_and(a, b) ((a) & (b)) // Vector AND
+#endif
+
+#ifndef vec_or
+#define vec_or(a, b)  ((a) | (b)) // Vector OR
+#endif
+
+#ifndef vec_xor
+#define vec_xor(a, b) ((a) ^ (b)) // Vector XOR
+#endif
+
+typedef signed   char char8x16_t  __attribute__((vector_size(16)));
+typedef unsigned char uchar8x16_t __attribute__((vector_size(16)));
+
+typedef int8_t  int8x16_t __attribute__((vector_size(16)));
+typedef int16_t int16x8_t __attribute__((vector_size(16)));
+typedef int32_t int32x4_t __attribute__((vector_size(16)));
+typedef int64_t int64x2_t __attribute__((vector_size(16)));
+
+typedef uint8_t  uint8x16_t __attribute__((vector_size(16)));
+typedef uint16_t uint16x8_t __attribute__((vector_size(16)));
+typedef uint32_t uint32x4_t __attribute__((vector_size(16)));
+typedef uint64_t uint64x2_t __attribute__((vector_size(16)));
+
+typedef float  float32x4_t  __attribute__((vector_size(16)));
+typedef double double64x2_t __attribute__((vector_size(16)));
+
+typedef struct ggml_uint8x16x2_t {
+    uint8x16_t val[2];
+} ggml_uint8x16x2_t;
+
+inline static ggml_uint8x16x2_t ggml_vec_xl_u8x2(const uint8_t * ptr) {
+    ggml_uint8x16x2_t res;
+
+    res.val[0] = vec_xl( 0, ptr);
+    res.val[1] = vec_xl(16, ptr);
+
+    return res;
+}
+
+typedef struct ggml_uint8x16x4_t {
+    uint8x16_t val[4];
+} ggml_uint8x16x4_t;
+
+inline static ggml_uint8x16x4_t ggml_vec_xl_u8x4(const uint8_t * ptr) {
+    ggml_uint8x16x4_t res;
+
+    res.val[0] = vec_xl( 0, ptr);
+    res.val[1] = vec_xl(16, ptr);
+    res.val[2] = vec_xl(32, ptr);
+    res.val[3] = vec_xl(48, ptr);
+
+    return res;
+}
+
+typedef struct ggml_int8x16x4_t {
+    int8x16_t val[4];
+} ggml_int8x16x4_t;
+
+inline static ggml_int8x16x4_t ggml_vec_xl_s8x4(const int8_t * ptr) {
+    ggml_int8x16x4_t res;
+
+    res.val[0] = vec_xl( 0, ptr);
+    res.val[1] = vec_xl(16, ptr);
+    res.val[2] = vec_xl(32, ptr);
+    res.val[3] = vec_xl(48, ptr);
+
+    return res;
+}
+
+typedef struct ggml_int16x8x2_t {
+    int16x8_t val[2];
+} ggml_int16x8x2_t;
+
+inline static ggml_int16x8x2_t ggml_vec_xl_s16x2(const int16_t * ptr) {
+    ggml_int16x8x2_t res;
+
+    res.val[0] = vec_xl( 0, ptr);
+    res.val[1] = vec_xl(16, ptr);
+
+    return res;
+}
+
+/*
+    ! WARNING: Very slow. Use vec_perm if possible. Refer to iq4_xs
+    !          or iq4_nl for example implementation.
+*/
+inline static int8x16_t ggml_vec_tbl(int8x16_t a, uint8x16_t b) {
+    int8x16_t res;
+
+    res[ 0] = a[b[ 0]];
+    res[ 1] = a[b[ 1]];
+    res[ 2] = a[b[ 2]];
+    res[ 3] = a[b[ 3]];
+    res[ 4] = a[b[ 4]];
+    res[ 5] = a[b[ 5]];
+    res[ 6] = a[b[ 6]];
+    res[ 7] = a[b[ 7]];
+    res[ 8] = a[b[ 8]];
+    res[ 9] = a[b[ 9]];
+    res[10] = a[b[10]];
+    res[11] = a[b[11]];
+    res[12] = a[b[12]];
+    res[13] = a[b[13]];
+    res[14] = a[b[14]];
+    res[15] = a[b[15]];
+
+    return res;
+}
+
+inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) {
+    const uchar8x16_t v_maske = {  0,  1,  4,  5,  8,  9, 12, 13,
+                                  16, 17, 20, 21, 24, 25, 28, 29 };
+
+    const int16x8_t v_abo = vec_pack((int32x4_t)a, (int32x4_t)b);
+    const int16x8_t v_abe = vec_perm(a, b, v_maske);
+    return v_abo + v_abe;
+}
+
+/**
+ * @see https://github.com/ggml-org/llama.cpp/pull/14037
+ */
+inline static float vec_hsum_f32x4(float32x4_t v) {
+    float32x4_t v_temp = v + vec_reve(v);
+    return v_temp[0] + v_temp[1];
+}
+
+inline static int32_t vec_hsum_i32x4(int32x4_t v) {
+    int32x4_t v_temp = v + vec_reve(v);
+    return v_temp[0] + v_temp[1];
+}
+
+inline static int32x4_t ggml_vec_dot(int32x4_t acc, int8x16_t a, int8x16_t b) {
+    const int16x8_t p = vec_mule(a, b) + vec_mulo(a, b);
+    return acc + (vec_unpackh(p) + vec_unpackl(p));
+}
+
+#endif
+
+#if defined(__loongarch_sx)
+/* float type data load instructions */
+static __m128 __lsx_vreplfr2vr_s(const float val) {
+    v4f32 res = {val, val, val, val};
+    return (__m128)res;
+}
+#endif
+
+#if defined(__loongarch_asx)
+static __m256 __lasx_xvreplfr2vr_s(const float val) {
+    v8f32 res = {val, val, val, val, val, val, val, val};
+    return (__m256)res;
+}
+#endif
+
+// TODO: move to ggml-threading
+void ggml_barrier(struct ggml_threadpool * tp);
+
+void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value);
+int  ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c
new file mode 100644
index 0000000000..5e36459f8c
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c
@@ -0,0 +1,1339 @@
+#define GGML_COMMON_IMPL_C
+#include "ggml-common.h"
+
+#include "ggml-cpu-impl.h"
+#include "simd-mappings.h"
+#include "ggml-quants.h"
+#include "quants.h"
+
+#include "arch-fallback.h"
+
+#include <string.h>
+#include <assert.h>
+#include <float.h>
+#include <stdlib.h> // for qsort
+#include <stdio.h>  // for GGML_ASSERT
+
+#define GROUP_MAX_EPS 1e-15f
+#define GROUP_MAX_EPS_IQ3_XXS 1e-8f
+#define GROUP_MAX_EPS_IQ2_S 1e-8f
+#define GROUP_MAX_EPS_IQ1_M 1e-7f
+#define GROUP_MAX_EPS_IQ1_S 1e-12f
+
+#define UNUSED GGML_UNUSED
+
+void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q1_0_ref(x, y, k);
+}
+
+void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q2_0_ref(x, y, k);
+}
+
+void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q4_0_ref(x, y, k);
+}
+
+void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q4_1_ref(x, y, k);
+}
+
+void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q5_0_ref(x, y, k);
+}
+
+void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q5_1_ref(x, y, k);
+}
+
+void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q8_0_ref(x, y, k);
+}
+
+void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q8_1_ref(x, y, k);
+}
+
+void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_mxfp4_ref(x, y, k);
+}
+
+void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_nvfp4_ref(x, y, k);
+}
+
+//
+// 2-6 bit quantization in super-blocks
+//
+
+//========================- 2-bit (de)-quantization
+
+void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    quantize_row_q2_K_ref(x, vy, k);
+}
+
+//========================= 3-bit (de)-quantization
+
+void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    quantize_row_q3_K_ref(x, vy, k);
+}
+
+// ====================== 4-bit (de)-quantization
+
+void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK_K == 0);
+    block_q4_K * GGML_RESTRICT y = vy;
+    quantize_row_q4_K_ref(x, y, k);
+}
+
+// ====================== 5-bit (de)-quantization
+
+void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK_K == 0);
+    block_q5_K * GGML_RESTRICT y = vy;
+    quantize_row_q5_K_ref(x, y, k);
+}
+
+// ====================== 6-bit (de)-quantization
+
+void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK_K == 0);
+    block_q6_K * GGML_RESTRICT y = vy;
+    quantize_row_q6_K_ref(x, y, k);
+}
+
+// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs)
+
+void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK_K == 0);
+    block_tq1_0 * GGML_RESTRICT y = vy;
+    quantize_row_tq1_0_ref(x, y, k);
+}
+
+void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(k % QK_K == 0);
+    block_tq2_0 * GGML_RESTRICT y = vy;
+    quantize_row_tq2_0_ref(x, y, k);
+}
+
+//===================================== Q8_K ==============================================
+
+void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q8_K_ref(x, y, k);
+}
+
+//===================================== Dot products =================================
+
+void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK1_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q1_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    float sumf = 0.0;
+
+    for (int i = 0; i < nb; i++) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        float sumi = 0.0f;
+
+        for (int k = 0; k < 4; k++) {
+            const block_q8_0 * GGML_RESTRICT yb = &y[i * 4 + k];
+            const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
+            int sumi_block = 0;
+
+            const uint8_t * GGML_RESTRICT bits = &x[i].qs[k * 4];
+            const int8_t  * GGML_RESTRICT qy   = yb->qs;
+
+            for (int b = 0; b < 4; ++b, qy += 8) {
+                const unsigned mask = bits[b];
+                sumi_block += ((mask & 0x01) ? qy[0] : -qy[0])
+                           +  ((mask & 0x02) ? qy[1] : -qy[1])
+                           +  ((mask & 0x04) ? qy[2] : -qy[2])
+                           +  ((mask & 0x08) ? qy[3] : -qy[3])
+                           +  ((mask & 0x10) ? qy[4] : -qy[4])
+                           +  ((mask & 0x20) ? qy[5] : -qy[5])
+                           +  ((mask & 0x40) ? qy[6] : -qy[6])
+                           +  ((mask & 0x80) ? qy[7] : -qy[7]);
+            }
+
+            sumi += d1 * sumi_block;
+        }
+
+        sumf += d0 * sumi;
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK2_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; i++) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        float sumi = 0.0f;
+
+        // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64)
+        for (int k = 0; k < 2; k++) {
+            const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
+            const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
+            int sumi_block = 0;
+
+            const uint8_t * GGML_RESTRICT qs = &x[i].qs[k * 8];
+            const int8_t  * GGML_RESTRICT qy = yb->qs;
+
+            for (int b = 0; b < 8; ++b) {
+                const uint8_t byte = qs[b];
+                // Extract 4 two-bit values, map {0,1,2,3} -> {-1,0,1,2}
+                sumi_block += ((int)((byte >> 0) & 3) - 1) * qy[b*4 + 0];
+                sumi_block += ((int)((byte >> 2) & 3) - 1) * qy[b*4 + 1];
+                sumi_block += ((int)((byte >> 4) & 3) - 1) * qy[b*4 + 2];
+                sumi_block += ((int)((byte >> 6) & 3) - 1) * qy[b*4 + 3];
+            }
+
+            sumi += d1 * sumi_block;
+        }
+
+        sumf += d0 * sumi;
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+    float sumf = 0;
+
+    for (; ib < nb; ++ib) {
+        int sumi0 = 0;
+        int sumi1 = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int v0 = (x[ib].qs[j] & 0x0F) - 8;
+            const int v1 = (x[ib].qs[j] >>   4) - 8;
+
+            sumi0 += (v0 * y[ib].qs[j]);
+            sumi1 += (v1 * y[ib].qs[j + qk/2]);
+        }
+
+        int sumi = sumi0 + sumi1;
+        sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d);
+    }
+
+    *s = sumf;
+}
+
+// TODO: add WASM SIMD
+void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_1;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_1 * GGML_RESTRICT x = vx;
+    const block_q8_1 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+    float sumf = 0;
+
+    for (; ib < nb; ++ib) {
+        int sumi0 = 0;
+        int sumi1 = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int v0 = (x[ib].qs[j] & 0x0F);
+            const int v1 = (x[ib].qs[j] >>   4);
+
+            sumi0 += (v0 * y[ib].qs[j]);
+            sumi1 += (v1 * y[ib].qs[j + qk/2]);
+        }
+
+        int sumi = sumi0 + sumi1;
+        sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_mxfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_MXFP4 == 0);
+    static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same");
+
+    const block_mxfp4 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_MXFP4;
+
+    int ib = 0;
+    float sumf = 0;
+
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_E8M0_TO_FP32_HALF(x[ib].e);
+
+        int sumi1 = 0;
+        int sumi2 = 0;
+        for (int j = 0; j < QK_MXFP4/2; ++j) {
+            sumi1 += y[ib].qs[j +          0] * kvalues_mxfp4[x[ib].qs[j] & 0xf];
+            sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >>  4];
+        }
+        sumf += d * (sumi1 + sumi2);
+    }
+    *s = sumf;
+}
+
+// NVFP4: super-block of 64 elements = 4 sub-blocks of 16 = 2 q8_0 blocks
+void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_NVFP4 == 0);
+
+    const block_nvfp4 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_NVFP4;
+
+    float sumf = 0;
+
+    for (int ib = 0; ib < nb; ++ib) {
+        for (int s_idx = 0; s_idx < 4; ++s_idx) {
+            const float d = ggml_ue4m3_to_fp32(x[ib].d[s_idx]);
+            const int q8_block = s_idx / 2;
+            const int q8_off   = (s_idx % 2) * QK_NVFP4_SUB;
+            const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d);
+
+            int sumi_lo = 0, sumi_hi = 0;
+            for (int j = 0; j < QK_NVFP4_SUB/2; ++j) {
+                const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j];
+                sumi_lo += y[2*ib + q8_block].qs[q8_off + j +               0] * kvalues_mxfp4[qv & 0xf];
+                sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_mxfp4[qv >>  4];
+            }
+
+            sumf += dy * d * (sumi_lo + sumi_hi);
+        }
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    int ib = 0;
+    float sumf = 0;
+
+    assert(n % qk == 0);
+    assert(qk == QK5_0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    for (; ib < nb; ++ib) {
+        uint32_t qh;
+        memcpy(&qh, x[ib].qh, sizeof(qh));
+
+        int sumi0 = 0;
+        int sumi1 = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const uint8_t xh_0 = ((qh & (1u << (j + 0 ))) >> (j + 0 )) << 4;
+            const uint8_t xh_1 = ((qh & (1u << (j + 16))) >> (j + 12));
+
+            const int32_t x0 = (int8_t)(((x[ib].qs[j] & 0x0F) | xh_0) - 16);
+            const int32_t x1 = (int8_t)(((x[ib].qs[j] >>   4) | xh_1) - 16);
+
+            sumi0 += (x0 * y[ib].qs[j]);
+            sumi1 += (x1 * y[ib].qs[j + qk/2]);
+        }
+
+        int sumi = sumi0 + sumi1;
+        sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)) * sumi;
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q5_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_1;
+    const int nb = n / qk;
+
+    int ib = 0;
+    float sumf = 0;
+
+    assert(n % qk == 0);
+    assert(qk == QK5_1);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_1 * GGML_RESTRICT x = vx;
+    const block_q8_1 * GGML_RESTRICT y = vy;
+
+    for (; ib < nb; ++ib) {
+        uint32_t qh;
+        memcpy(&qh, x[ib].qh, sizeof(qh));
+
+        int sumi0 = 0;
+        int sumi1 = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const uint8_t xh_0 = ((qh >> (j +  0)) << 4) & 0x10;
+            const uint8_t xh_1 = ((qh >> (j + 12))     ) & 0x10;
+
+            const int32_t x0 = (x[ib].qs[j] & 0xF) | xh_0;
+            const int32_t x1 = (x[ib].qs[j] >>  4) | xh_1;
+
+            sumi0 += (x0 * y[ib].qs[j]);
+            sumi1 += (x1 * y[ib].qs[j + qk/2]);
+        }
+
+        int sumi = sumi0 + sumi1;
+        sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q8_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q8_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    int ib = 0;
+    float sumf = 0;
+
+    for (; ib < nb; ++ib) {
+        int sumi = 0;
+
+        for (int j = 0; j < qk; j++) {
+            sumi += x[ib].qs[j]*y[ib].qs[j];
+        }
+
+        sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d));
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_tq1_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_tq1_0 * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243};
+
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; ++i) {
+        int sum = 0;
+
+        for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) {
+            for (size_t l = 0; l < 5; ++l) {
+                for (size_t m = 0; m < 32; ++m) {
+                    uint8_t q = x[i].qs[j + m] * pow3[l];
+                    uint16_t xi = ((uint16_t) q * 3) >> 8;
+                    sum += (xi - 1) * y[i].qs[j*5 + l*32 + m];
+                }
+            }
+        }
+        for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) {
+            for (size_t l = 0; l < 5; ++l) {
+                for (size_t m = 0; m < 16; ++m) {
+                    uint8_t q = x[i].qs[j + m] * pow3[l];
+                    uint16_t xi = ((uint16_t) q * 3) >> 8;
+                    sum += (xi - 1) * y[i].qs[j*5 + l*16 + m];
+                }
+            }
+        }
+
+        for (size_t l = 0; l < 4; ++l) {
+            for (size_t j = 0; j < sizeof(x->qh); ++j) {
+                uint8_t q = x[i].qh[j] * pow3[l];
+                uint16_t xi = ((uint16_t) q * 3) >> 8;
+                sum += (xi - 1) * y[i].qs[sizeof(x->qs)*5 + l*sizeof(x->qh) + j];
+            }
+        }
+
+        sumf += (float) sum * (GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_tq2_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_tq2_0 * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; ++i) {
+        int32_t sumi = 0;
+
+        for (size_t j = 0; j < sizeof(x->qs); j += 32) {
+            for (size_t l = 0; l < 4; ++l) {
+                for (size_t k = 0; k < 32; ++k) {
+                    sumi += y[i].qs[j*4 + l*32 + k] * (((x[i].qs[j + k] >> (l*2)) & 3) - 1);
+                }
+            }
+        }
+
+        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        sumf += (float) sumi * d;
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_q2_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0;
+
+    for (int i = 0; i < nb; ++i) {
+
+        const uint8_t * q2 = x[i].qs;
+        const  int8_t * q8 = y[i].qs;
+        const uint8_t * sc = x[i].scales;
+
+        int summs = 0;
+        for (int j = 0; j < 16; ++j) {
+            summs += y[i].bsums[j] * (sc[j] >> 4);
+        }
+
+        const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
+        const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin);
+
+        int isum = 0;
+        int is = 0;
+        int d;
+        for (int k = 0; k < QK_K/128; ++k) {
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {
+                d = sc[is++] & 0xF;
+                int isuml = 0;
+                for (int l =  0; l < 16; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3);
+                isum += d * isuml;
+                d = sc[is++] & 0xF;
+                isuml = 0;
+                for (int l = 16; l < 32; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3);
+                isum += d * isuml;
+                shift += 2;
+                q8 += 32;
+            }
+            q2 += 32;
+        }
+        sumf += dall * isum - dmin * summs;
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_q3_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const uint32_t kmask1 = 0x03030303;
+    const uint32_t kmask2 = 0x0f0f0f0f;
+
+    const block_q3_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    // scalar version
+    // This function is written like this so the compiler can manage to vectorize most of it
+    // Using -Ofast, GCC and clang manage to produce code that is within a factor of 2 or so from the
+    // manually vectorized version above. Every other version I tried would run at least 4 times slower.
+    // The ideal situation would be if we could just write the code once, and the compiler would
+    // automatically produce the best possible set of machine instructions, instead of us having to manually
+    // write vectorized versions for AVX, ARM_NEON, etc.
+
+    int8_t  aux8[QK_K];
+    int16_t aux16[8];
+    float   sums [8];
+    int32_t aux32[8];
+    memset(sums, 0, 8*sizeof(float));
+
+    uint32_t auxs[4];
+    const int8_t * scales = (const int8_t*)auxs;
+
+    float sumf = 0;
+    for (int i = 0; i < nb; ++i) {
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const uint8_t * GGML_RESTRICT hm = x[i].hmask;
+        const  int8_t * GGML_RESTRICT q8 = y[i].qs;
+        memset(aux32, 0, 8*sizeof(int32_t));
+        int8_t * GGML_RESTRICT a = aux8;
+        uint8_t m = 1;
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) a[l] = q3[l] & 3;
+            for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4);
+            a += 32; m <<= 1;
+            for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 2) & 3;
+            for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4);
+            a += 32; m <<= 1;
+            for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 4) & 3;
+            for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4);
+            a += 32; m <<= 1;
+            for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 6) & 3;
+            for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4);
+            a += 32; m <<= 1;
+            q3 += 32;
+        }
+        a = aux8;
+
+        memcpy(auxs, x[i].scales, 12);
+        uint32_t tmp = auxs[2];
+        auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+        auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+        auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
+        auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+        for (int j = 0; j < QK_K/16; ++j) {
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l];
+            q8 += 8; a += 8;
+        }
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l];
+    }
+    for (int l = 0; l < 8; ++l) sumf += sums[l];
+    *s = sumf;
+}
+
+void ggml_vec_dot_q4_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    static const uint32_t kmask1 = 0x3f3f3f3f;
+    static const uint32_t kmask2 = 0x0f0f0f0f;
+    static const uint32_t kmask3 = 0x03030303;
+
+    uint32_t utmp[4];
+
+    const uint8_t * scales = (const uint8_t*)&utmp[0];
+    const uint8_t * mins   = (const uint8_t*)&utmp[2];
+
+    int8_t  aux8[QK_K];
+    int16_t aux16[8];
+    float   sums [8];
+    int32_t aux32[8];
+    memset(sums, 0, 8*sizeof(float));
+
+    float sumf = 0;
+    for (int i = 0; i < nb; ++i) {
+        const uint8_t * GGML_RESTRICT q4 = x[i].qs;
+        const  int8_t * GGML_RESTRICT q8 = y[i].qs;
+        memset(aux32, 0, 8*sizeof(int32_t));
+        int8_t * GGML_RESTRICT a = aux8;
+        for (int j = 0; j < QK_K/64; ++j) {
+            for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF);
+            a += 32;
+            for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l]  >> 4);
+            a += 32; q4 += 32;
+        }
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        int sumi = 0;
+        for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2];
+        a = aux8;
+        int is = 0;
+        for (int j = 0; j < QK_K/32; ++j) {
+            int32_t scale = scales[is++];
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+        }
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l];
+        const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d;
+        sumf -= dmin * sumi;
+    }
+    for (int l = 0; l < 8; ++l) sumf += sums[l];
+    *s = sumf;
+}
+
+void ggml_vec_dot_q5_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy,  size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q5_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    static const uint32_t kmask1 = 0x3f3f3f3f;
+    static const uint32_t kmask2 = 0x0f0f0f0f;
+    static const uint32_t kmask3 = 0x03030303;
+
+    uint32_t utmp[4];
+
+    const uint8_t * scales = (const uint8_t*)&utmp[0];
+    const uint8_t * mins   = (const uint8_t*)&utmp[2];
+
+    int8_t  aux8[QK_K];
+    int16_t aux16[8];
+    float   sums [8];
+    int32_t aux32[8];
+    memset(sums, 0, 8*sizeof(float));
+
+    float sumf = 0;
+    for (int i = 0; i < nb; ++i) {
+        const uint8_t * GGML_RESTRICT q4 = x[i].qs;
+        const uint8_t * GGML_RESTRICT hm = x[i].qh;
+        const  int8_t * GGML_RESTRICT q8 = y[i].qs;
+        memset(aux32, 0, 8*sizeof(int32_t));
+        int8_t * GGML_RESTRICT a = aux8;
+        uint8_t m = 1;
+        for (int j = 0; j < QK_K/64; ++j) {
+            for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF);
+            for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0);
+            a += 32; m <<= 1;
+            for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l]  >> 4);
+            for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0);
+            a += 32; m <<= 1;
+            q4 += 32;
+        }
+        memcpy(utmp, x[i].scales, 12);
+        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
+        const uint32_t uaux = utmp[1] & kmask1;
+        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
+        utmp[2] = uaux;
+        utmp[0] &= kmask1;
+
+        int sumi = 0;
+        for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2];
+        a = aux8;
+        int is = 0;
+        for (int j = 0; j < QK_K/32; ++j) {
+            int32_t scale = scales[is++];
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+        }
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l];
+        const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d;
+        sumf -= dmin * sumi;
+    }
+    for (int l = 0; l < 8; ++l) sumf += sums[l];
+    *s = sumf;
+}
+
+void ggml_vec_dot_q6_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q6_K * GGML_RESTRICT x = vx;
+    const block_q8_K * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    int8_t  aux8[QK_K];
+    int16_t aux16[8];
+    float   sums [8];
+    int32_t aux32[8];
+    memset(sums, 0, 8*sizeof(float));
+
+    float sumf = 0;
+    for (int i = 0; i < nb; ++i) {
+        const uint8_t * GGML_RESTRICT q4 = x[i].ql;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const  int8_t * GGML_RESTRICT q8 = y[i].qs;
+        memset(aux32, 0, 8*sizeof(int32_t));
+        int8_t * GGML_RESTRICT a = aux8;
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                a[l +  0] = (int8_t)((q4[l +  0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
+                a[l + 32] = (int8_t)((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
+                a[l + 64] = (int8_t)((q4[l +  0] >>  4) | (((qh[l] >> 4) & 3) << 4)) - 32;
+                a[l + 96] = (int8_t)((q4[l + 32] >>  4) | (((qh[l] >> 6) & 3) << 4)) - 32;
+            }
+            a  += 128;
+            q4 += 64;
+            qh += 32;
+        }
+        a = aux8;
+        int is = 0;
+        for (int j = 0; j < QK_K/16; ++j) {
+            int scale = x[i].scales[is++];
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+            for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l];
+            for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l];
+            q8 += 8; a += 8;
+        }
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l];
+    }
+    for (int l = 0; l < 8; ++l) sumf += sums[l];
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq2_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_xxs * GGML_RESTRICT x = vx;
+    const block_q8_K    * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    uint32_t aux32[2];
+    const uint8_t * aux8 = (const uint8_t *)aux32;
+
+    float sumf = 0.f;
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+        int32_t bsum = 0;
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            memcpy(aux32, q2, 2*sizeof(uint32_t));
+            q2 += 4;
+            const uint32_t ls = 2*(aux32[1] >> 28) + 1;
+            int32_t sumi = 0;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]);
+                const uint8_t  signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127];
+                for (int j = 0; j < 8; ++j) {
+                    sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            bsum += sumi * ls;
+        }
+        sumf += d * bsum;
+    }
+    *s = 0.125f * sumf;
+}
+
+void ggml_vec_dot_iq2_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_xs * GGML_RESTRICT x = vx;
+    const block_q8_K   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0.f;
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint16_t * GGML_RESTRICT q2 = x[i].qs;
+        const uint8_t  * GGML_RESTRICT sc = x[i].scales;
+        const int8_t   * GGML_RESTRICT q8 = y[i].qs;
+        int32_t bsum = 0;
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            const uint16_t ls1 = 2*(sc[ib32] & 0xf) + 1;
+            const uint16_t ls2 = 2*(sc[ib32] >>  4) + 1;
+            int32_t sumi = 0;
+            for (int l = 0; l < 2; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511));
+                const uint8_t  signs = ksigns_iq2xs[q2[l] >> 9];
+                for (int j = 0; j < 8; ++j) {
+                    sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            bsum += sumi * ls1;
+            sumi = 0;
+            for (int l = 2; l < 4; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511));
+                const uint8_t  signs = ksigns_iq2xs[q2[l] >> 9];
+                for (int j = 0; j < 8; ++j) {
+                    sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            bsum += sumi * ls2;
+            q2 += 4;
+        }
+        sumf += d * bsum;
+    }
+    *s = 0.125f * sumf;
+}
+
+void ggml_vec_dot_iq2_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq2_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0;
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const int8_t  * q8 = y[i].qs;
+        const uint8_t * qs = x[i].qs;
+        const uint8_t * qh = x[i].qh;
+        const uint8_t * signs = qs + QK_K/8;
+
+        int bsum = 0;
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            int ls1 = 1 + 2*(x[i].scales[ib32] & 0xf);
+            int ls2 = 1 + 2*(x[i].scales[ib32] >>  4);
+            int sumi1 = 0, sumi2 = 0;
+            for (int l = 0; l < 2; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300)));
+                for (int j = 0; j < 8; ++j) {
+                    sumi1 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            for (int l = 2; l < 4; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300)));
+                for (int j = 0; j < 8; ++j) {
+                    sumi2 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            bsum += ls1 * sumi1 + ls2 * sumi2;
+            qs += 4;
+            signs += 4;
+        }
+
+        sumf += d * bsum;
+    }
+
+    *s = 0.125f * sumf;
+}
+
+void ggml_vec_dot_iq3_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq3_xxs * GGML_RESTRICT x = vx;
+    const block_q8_K    * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    uint32_t aux32;
+
+    float sumf = 0.f;
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
+        const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        int32_t bsum = 0;
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            memcpy(&aux32, gas, sizeof(uint32_t)); gas += sizeof(uint32_t);
+            const uint32_t ls = 2*(aux32 >> 28) + 1;
+            int32_t sumi = 0;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*l+0]);
+                const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*l+1]);
+                const uint8_t  signs = ksigns_iq2xs[(aux32 >> 7*l) & 127];
+                for (int j = 0; j < 4; ++j) {
+                    sumi += grid1[j] * q8[j+0] * (signs & kmask_iq2xs[j+0] ? -1 : 1);
+                    sumi += grid2[j] * q8[j+4] * (signs & kmask_iq2xs[j+4] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            q3 += 8;
+            bsum += sumi * ls;
+        }
+        sumf += d * bsum;
+    }
+    *s = 0.25f * sumf;
+}
+
+void ggml_vec_dot_iq3_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq3_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0.f;
+    for (int i = 0; i < nb; ++i) {
+        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
+        const uint8_t * GGML_RESTRICT qs = x[i].qs;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const uint8_t * GGML_RESTRICT signs = x[i].signs;
+        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
+        int32_t bsum = 0;
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const uint32_t ls1 = 2*(x[i].scales[ib32/2] & 0xf) + 1;
+            const uint32_t ls2 = 2*(x[i].scales[ib32/2] >>  4) + 1;
+            int32_t sumi = 0;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+0] << (8-2*l)) & 256)));
+                const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+0] << (7-2*l)) & 256)));
+                for (int j = 0; j < 4; ++j) {
+                    sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1);
+                    sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            qs += 8;
+            signs += 4;
+            bsum += sumi * ls1;
+            sumi = 0;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+1] << (8-2*l)) & 256)));
+                const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+1] << (7-2*l)) & 256)));
+                for (int j = 0; j < 4; ++j) {
+                    sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1);
+                    sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1);
+                }
+                q8 += 8;
+            }
+            qs += 8;
+            signs += 4;
+            bsum += sumi * ls2;
+        }
+        sumf += d * bsum;
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq1_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq1_s * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0;
+    for (int i = 0; i < nb; i++) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint16_t * qh = x[i].qh;
+
+        int sumi = 0, sumi1 = 0;
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const int ls = 2*((qh[ib] >> 12) & 7) + 1;
+            const int delta = qh[ib] & 0x8000 ? -1 : 1;
+            int lsum = 0;
+            for (int l = 0; l < 4; ++l) {
+                const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8)));
+                for (int j = 0; j < 8; ++j) {
+                    lsum += q8[j] * grid[j];
+                }
+                q8 += 8;
+            }
+            sumi  += ls * lsum;
+            sumi1 += ls * delta * (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]);
+            qs += 4;
+        }
+
+        sumf += GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d * (sumi + IQ1S_DELTA * sumi1);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq1_m_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_iq1_m * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    iq1m_scale_t scale;
+
+    int sum1[2], sum2[2], delta[4];
+
+    float sumf = 0;
+    for (int i = 0; i < nb; i++) {
+
+        const int8_t   * q8 = y[i].qs;
+        const uint8_t  * qs = x[i].qs;
+        const uint8_t  * qh = x[i].qh;
+        const uint16_t * sc = (const uint16_t *)x[i].scales;
+
+        scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+
+        int sumi1 = 0, sumi2 = 0;
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            delta[0] = qh[0] & 0x08 ? -1 : 1;
+            delta[1] = qh[0] & 0x80 ? -1 : 1;
+            delta[2] = qh[1] & 0x08 ? -1 : 1;
+            delta[3] = qh[1] & 0x80 ? -1 : 1;
+            sum1[0] = sum1[1] = sum2[0] = sum2[1] = 0;
+            for (int l = 0; l < 4; ++l) {
+                const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((uint16_t)qh[l/2] << (8 - 4*(l%2))) & 0x700)));
+                int lsum1 = 0, lsum2 = 0;
+                for (int j = 0; j < 8; ++j) {
+                    lsum1 += q8[j] * grid[j];
+                    lsum2 += q8[j];
+                }
+                q8 += 8;
+                sum1[l/2] += lsum1;
+                sum2[l/2] += lsum2*delta[l];
+            }
+
+            const int ls1 = 2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1;
+            const int ls2 = 2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1;
+
+            sumi1 += sum1[0] * ls1 + sum1[1] * ls2;
+            sumi2 += sum2[0] * ls1 + sum2[1] * ls2;
+            qs += 4;
+            qh += 2;
+        }
+
+        sumf += GGML_CPU_FP16_TO_FP32(scale.f16) * y[i].d * (sumi1 + IQ1M_DELTA * sumi2);
+    }
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq4_nl_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK4_NL == 0);
+    static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same");
+
+    const block_iq4_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK4_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi1 = 0, sumi2 = 0;
+        for (int j = 0; j < QK4_NL/2; ++j) {
+            sumi1 += y[ib].qs[j+       0] * kvalues_iq4nl[x[ib].qs[j] & 0xf];
+            sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >>  4];
+        }
+        sumf += d * (sumi1 + sumi2);
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq4_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_K == 0);
+
+    const block_iq4_xs * GGML_RESTRICT x = vx;
+    const block_q8_K   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_K;
+
+    float sumf = 0;
+    for (int ibl = 0; ibl < nb; ++ibl) {
+        const float d4d8 = GGML_CPU_FP16_TO_FP32(x[ibl].d) * y[ibl].d;
+        uint16_t h = x[ibl].scales_h;
+        const uint8_t * qs = x[ibl].qs;
+        const int8_t  * q8 = y[ibl].qs;
+        for (int ib = 0; ib < QK_K/32; ib += 2) {
+            const uint8_t ls1 = (x[ibl].scales_l[ib/2] & 0xf) | ((h << 4) & 0x30);
+            const uint8_t ls2 = (x[ibl].scales_l[ib/2] >>  4) | ((h << 2) & 0x30);
+            h >>= 4;
+            const float d1 = d4d8*(ls1 - 32);
+            const float d2 = d4d8*(ls2 - 32);
+            int sumi1 = 0, sumi2 = 0;
+            for (int j = 0; j < 16; ++j) {
+                sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf];
+                sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >>  4];
+            }
+            sumf += d1 * (sumi1 + sumi2);
+            qs += 16;
+            q8 += 32;
+            sumi1 = sumi2 = 0;
+            for (int j = 0; j < 16; ++j) {
+                sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf];
+                sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >>  4];
+            }
+            sumf += d2 * (sumi1 + sumi2);
+            qs += 16;
+            q8 += 32;
+        }
+    }
+    *s = sumf;
+}
+
+// ============================ 4-bit non-linear quants
+
+void quantize_row_iq4_nl(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK4_NL == 0);
+    quantize_row_iq4_nl_ref(x, y, k);
+}
+
+void quantize_row_iq4_xs(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    quantize_iq4_xs(x, y, 1, k, NULL);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h
new file mode 100644
index 0000000000..93ea7eeffe
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h
@@ -0,0 +1,106 @@
+#pragma once
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+
+#include "ggml.h"
+
+// GGML CPU internal header
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Quantization
+void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+
+void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+
+void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+
+void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+
+void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+
+// Dot product
+void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_tq1_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_iq2_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq2_xs_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq2_s_q8_K  (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq3_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq1_s_q8_K  (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq1_m_q8_K  (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq4_nl_q8_0 (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq4_xs_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq3_s_q8_K  (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+// Generic implementation
+void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
+void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
+void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q8_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_mxfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_tq1_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_tq2_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+void ggml_vec_dot_q2_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q3_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q4_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q5_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy,  size_t by, int nrc);
+void ggml_vec_dot_q6_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq2_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq2_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq2_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq3_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq3_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq1_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq1_m_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq4_nl_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_iq4_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h
new file mode 100644
index 0000000000..fb53ecdf42
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h
@@ -0,0 +1,1366 @@
+#pragma once
+
+#include "ggml-cpu-impl.h"
+
+#ifdef __ARM_FEATURE_SVE
+#include <arm_sve.h>
+#endif // __ARM_FEATURE_SVE
+
+#if defined(__ARM_NEON) && !defined(__CUDACC__) && !defined(__MUSACC__)
+// if YCM cannot find <arm_neon.h>, make a symbolic link to it, for example:
+//
+//   $ ln -sfn /Library/Developer/CommandLineTools/usr/lib/clang/13.1.6/include/arm_neon.h ./src/
+//
+#include <arm_neon.h>
+#endif
+
+#if defined(__riscv_v_intrinsic)
+#include <riscv_vector.h>
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//
+// simd mappings
+//
+
+// FP16 to FP32 conversion
+
+// 16-bit float
+// on Arm, we use __fp16, which requires the IEEE fp16 format: implied on
+// AArch64, selected by -mfp16-format=ieee on 32 bit Arm, where the compiler
+// may otherwise reject the type
+// on x86, we use uint16_t
+//
+// for old CUDA compilers (<= 11), we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/10616
+// for     MUSA compilers        , we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/11843
+//
+#if defined(__ARM_NEON) && defined(__ARM_FP16_FORMAT_IEEE) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__)
+    #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) neon_compute_fp16_to_fp32(x)
+    #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x)
+
+    #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+
+    static inline float neon_compute_fp16_to_fp32(ggml_fp16_t h) {
+        __fp16 tmp;
+        memcpy(&tmp, &h, sizeof(ggml_fp16_t));
+        return (float)tmp;
+    }
+
+    static inline ggml_fp16_t neon_compute_fp32_to_fp16(float f) {
+        ggml_fp16_t res;
+        __fp16 tmp = f;
+        memcpy(&res, &tmp, sizeof(ggml_fp16_t));
+        return res;
+    }
+#elif defined(__F16C__)
+    #ifdef _MSC_VER
+        #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128(x)))
+        #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) _mm_extract_epi16(_mm_cvtps_ph(_mm_set_ss(x), 0), 0)
+    #else
+        #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) _cvtsh_ss(x)
+        #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) _cvtss_sh(x, 0)
+    #endif
+#elif defined(__POWER9_VECTOR__)
+    #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) power_compute_fp16_to_fp32(x)
+    #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) power_compute_fp32_to_fp16(x)
+    /* the inline asm below is about 12% faster than the lookup method */
+    #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+    #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
+
+    static inline float power_compute_fp16_to_fp32(ggml_fp16_t h) {
+        float f;
+        double d;
+        __asm__(
+            "mtfprd %0,%2\n"
+            "xscvhpdp %0,%0\n"
+            "frsp %1,%0\n" :
+            /* temp */ "=d"(d),
+            /* out */  "=f"(f):
+            /* in */   "r"(h));
+        return f;
+    }
+
+    static inline ggml_fp16_t power_compute_fp32_to_fp16(float f) {
+        double d;
+        ggml_fp16_t r;
+        __asm__( /* xscvdphp can work on double or single precision */
+            "xscvdphp %0,%2\n"
+            "mffprd %1,%0\n" :
+            /* temp */ "=d"(d),
+            /* out */  "=r"(r):
+            /* in */   "f"(f));
+        return r;
+    }
+#elif defined(__riscv) && defined(__riscv_zfhmin)
+    static inline float riscv_compute_fp16_to_fp32(ggml_fp16_t h) {
+        _Float16 hf;
+        memcpy(&hf, &h, sizeof(ggml_fp16_t));
+        return hf;
+    }
+
+    static inline ggml_fp16_t riscv_compute_fp32_to_fp16(float f) {
+        ggml_fp16_t res;
+        _Float16 hf = (_Float16)f;
+        memcpy(&res, &hf, sizeof(ggml_fp16_t));
+        return res;
+    }
+
+    #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) riscv_compute_fp16_to_fp32(x)
+    #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) riscv_compute_fp32_to_fp16(x)
+    #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+    #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
+#endif
+
+// precomputed f32 table for f16 (256 KB)
+// defined in ggml-cpu.c, initialized in ggml_cpu_init()
+extern float ggml_table_f32_f16[1 << 16];
+
+// precomputed f32 table for e8m0 half (1 KB)
+// defined in ggml-cpu.c, initialized in ggml_cpu_init()
+extern float ggml_table_f32_e8m0_half[1 << 8];
+
+// precomputed f32 table for ue4m3 (1 KB)
+// defined in ggml-cpu.c, initialized in ggml_cpu_init()
+extern float ggml_table_f32_ue4m3[1 << 8];
+
+// Use lookup table for E8M0 on x86 (faster than bit manipulation)
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+#define GGML_CPU_E8M0_TO_FP32_HALF(x) ggml_table_f32_e8m0_half[(uint8_t)(x)]
+#else
+#define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x)
+#endif
+
+// Use lookup table for UE4M3 on x86 and ARM (faster than bit manipulation)
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__ARM_NEON)
+#define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)]
+#else
+#define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x)
+#endif
+
+// Pre AVX512 CPUs (AVX2 with F16C, such as i5-8th gen or i7-4790)
+// do not produce codegen as checks below don't test for __F16C__
+// but this check does.
+#if defined(__F16C__)
+#define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+#define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
+#endif
+
+// On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32,
+// so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON.
+// This is also true for POWER9.
+#if !defined(GGML_CPU_FP16_TO_FP32)
+inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) {
+    uint16_t s;
+    memcpy(&s, &f, sizeof(uint16_t));
+    return ggml_table_f32_f16[s];
+}
+
+#define GGML_CPU_FP16_TO_FP32(x) ggml_lookup_fp16_to_fp32(x)
+#endif
+
+#if !defined(GGML_CPU_FP32_TO_FP16)
+#define GGML_CPU_FP32_TO_FP16(x) GGML_COMPUTE_FP32_TO_FP16(x)
+#endif
+
+
+// we define a common set of C macros which map to specific intrinsics based on the current architecture
+// we then implement the fundamental computation operations below using only these macros
+// adding support for new architectures requires to define the corresponding SIMD macros
+//
+// GGML_F32_STEP / GGML_F16_STEP
+//   number of elements to process in a single step
+//
+// GGML_F32_EPR / GGML_F16_EPR
+//   number of elements to fit in a single register
+//
+
+#if defined(__ARM_FEATURE_SVE) && defined(__ARM_FEATURE_FMA)
+
+#define GGML_SIMD
+
+// F32 SVE
+#define GGML_F32_EPR 8
+#define DEFAULT_PG svptrue_b32()
+
+#define GGML_F32xt                        svfloat32_t
+#define GGML_F32xt_ZERO                   svdup_n_f32(0.0f)
+#define GGML_F32xt_SET1(x)                svdup_n_f32(x)
+#define GGML_F32xt_LOAD_IMPL(pg, a)       svld1_f32(pg, a)
+#define GGML_F32xt_LOAD(a)                GGML_F32xt_LOAD_IMPL(DEFAULT_PG, a)
+#define GGML_F32xt_STORE_IMPL(pg, a, b)   svst1_f32(pg, a, b)
+#define GGML_F32xt_STORE(a, b)            GGML_F32xt_STORE_IMPL(DEFAULT_PG, a, b)
+#define GGML_F32xt_FMA_IMPL(pg, a, b, c)  svmad_f32_m(pg, b, c, a)
+#define GGML_F32xt_FMA(a, b, c)           GGML_F32xt_FMA_IMPL(DEFAULT_PG, a, b, c)
+#define GGML_F32xt_ADD_IMPL(pg, a, b)     svadd_f32_m(pg, a, b)
+#define GGML_F32xt_ADD(a, b)              GGML_F32xt_ADD_IMPL(DEFAULT_PG, a, b)
+#define GGML_F32xt_MUL_IMPL(pg, a, b)     svmul_f32_m(pg, a, b)
+#define GGML_F32xt_MUL(a, b)              GGML_F32xt_MUL_IMPL(DEFAULT_PG, a, b)
+#define GGML_F32xt_REDUCE_ONE_IMPL(pg, a) svaddv(pg, a)
+#define GGML_F32xt_REDUCE_ONE(a)          GGML_F32xt_REDUCE_ONE_IMPL(DEFAULT_PG, a)
+#define GGML_F32xt_REDUCE_IMPL(pg, res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8)  \
+{                                                      \
+    sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum2);        \
+    sum3 = svadd_f32_m(DEFAULT_PG, sum3, sum4);        \
+    sum5 = svadd_f32_m(DEFAULT_PG, sum5, sum6);        \
+    sum7 = svadd_f32_m(DEFAULT_PG, sum7, sum8);        \
+    sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum3);        \
+    sum5 = svadd_f32_m(DEFAULT_PG, sum5, sum7);        \
+    sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum5);        \
+    (res) = (ggml_float) GGML_F32xt_REDUCE_ONE(sum1);  \
+}
+#define GGML_F32xt_REDUCE(res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8)  \
+        GGML_F32xt_REDUCE_IMPL(DEFAULT_PG, res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8)
+
+#define GGML_F32_VEC        GGML_F32xt
+#define GGML_F32_VEC_ZERO   GGML_F32xt_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32xt_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32xt_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32xt_STORE
+#define GGML_F32_VEC_FMA    GGML_F32xt_FMA
+#define GGML_F32_VEC_ADD    GGML_F32xt_ADD
+#define GGML_F32_VEC_MUL    GGML_F32xt_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32xt_REDUCE
+
+// F16 SVE
+#define DEFAULT_PG32    svptrue_b32()
+#define DEFAULT_PG16    svptrue_b16()
+
+#define GGML_F32Cxt                         svfloat16_t
+#define GGML_F32Cxt_ZERO                    svdup_n_f16(0.0f)
+#define GGML_F32Cxt_SET1(x)                 svdup_n_f16(x)
+#define GGML_F32Cxt_LOAD(p)                 svld1_f16(DEFAULT_PG16, (const __fp16 *)(p))
+#define GGML_F32Cxt_STORE(dst_ptr, src_vec) svst1_f16(DEFAULT_PG16, (__fp16 *)(dst_ptr), (src_vec))
+
+#define GGML_F32Cxt_FMA_IMPL(pg, a, b, c)   svmad_f16_x(pg, b, c, a)
+#define GGML_F32Cxt_FMA(a, b, c)            GGML_F32Cxt_FMA_IMPL(DEFAULT_PG16, a, b, c)
+#define GGML_F32Cxt_ADD_IMPL(pg, a, b)      svadd_f16_x(pg, a, b)
+#define GGML_F32Cxt_ADD(a, b)               GGML_F32Cxt_ADD_IMPL(DEFAULT_PG16, a, b)
+#define GGML_F32Cxt_MUL_IMPL(pg, a, b)      svmul_f16_x(pg, a, b)
+#define GGML_F32Cxt_MUL(a, b)               GGML_F32Cxt_MUL_IMPL(DEFAULT_PG16, a, b)
+#define GGML_F32Cxt_REDUCE                  GGML_F16xt_REDUCE_MIXED
+
+#define GGML_F16x_VEC                GGML_F32Cxt
+#define GGML_F16x_VEC_ZERO           GGML_F32Cxt_ZERO
+#define GGML_F16x_VEC_SET1           GGML_F32Cxt_SET1
+#define GGML_F16x_VEC_LOAD(p, i)     GGML_F32Cxt_LOAD(p)
+#define GGML_F16x_VEC_STORE(p, r, i) GGML_F32Cxt_STORE((__fp16 *)(p), r)
+#define GGML_F16x_VEC_FMA            GGML_F32Cxt_FMA
+#define GGML_F16x_VEC_ADD            GGML_F32Cxt_ADD
+#define GGML_F16x_VEC_MUL            GGML_F32Cxt_MUL
+#define GGML_F16x_VEC_REDUCE         GGML_F32Cxt_REDUCE
+
+#define GGML_F16xt_REDUCE_ONE_IMPL(pg, a) svaddv_f16(pg, a)
+#define GGML_F16xt_REDUCE_ONE(a)          GGML_F16xt_REDUCE_ONE_IMPL(DEFAULT_PG16, a)
+
+#define GGML_F16xt_REDUCE_MIXED_IMPL(pg16, res, sum1, sum2, sum3, sum4)  \
+{                                                      \
+    sum1 = svadd_f16_x(pg16, sum1, sum2);              \
+    sum3 = svadd_f16_x(pg16, sum3, sum4);              \
+    sum1 = svadd_f16_x(pg16, sum1, sum3);              \
+    __fp16 sum_f16 = svaddv_f16(pg16, sum1);           \
+    (res) = (ggml_float) sum_f16;                      \
+}
+#define GGML_F16xt_REDUCE_MIXED(res, sum1, sum2, sum3, sum4)  \
+        GGML_F16xt_REDUCE_MIXED_IMPL(DEFAULT_PG16, res, sum1, sum2, sum3, sum4)
+
+// F16 NEON
+
+#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)
+    #define GGML_F16_STEP 32
+    #define GGML_F16_EPR  8
+
+    #define GGML_F16x8              float16x8_t
+    #define GGML_F16x8_ZERO         vdupq_n_f16(0.0f)
+    #define GGML_F16x8_SET1(x)      vdupq_n_f16(x)
+    #define GGML_F16x8_LOAD(x)      vld1q_f16((const __fp16 *)(x))
+    #define GGML_F16x8_STORE        vst1q_f16
+    #define GGML_F16x8_FMA(a, b, c) vfmaq_f16(a, b, c)
+    #define GGML_F16x8_ADD          vaddq_f16
+    #define GGML_F16x8_MUL          vmulq_f16
+    #define GGML_F16x8_REDUCE(res, x)                               \
+    do {                                                            \
+        int offset = GGML_F16_ARR >> 1;                             \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        offset >>= 1;                                               \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        offset >>= 1;                                               \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        const float32x4_t t0 = vcvt_f32_f16(vget_low_f16 ((x)[0])); \
+        const float32x4_t t1 = vcvt_f32_f16(vget_high_f16((x)[0])); \
+        (res) = (ggml_float) vaddvq_f32(vaddq_f32(t0, t1));         \
+    } while (0)
+
+    #define GGML_F16_VEC                GGML_F16x8
+    #define GGML_F16_VEC_ZERO           GGML_F16x8_ZERO
+    #define GGML_F16_VEC_SET1           GGML_F16x8_SET1
+    #define GGML_F16_VEC_LOAD(p, i)     GGML_F16x8_LOAD(p)
+    #define GGML_F16_VEC_STORE(p, r, i) GGML_F16x8_STORE((__fp16 *)(p), (r)[i])
+    #define GGML_F16_VEC_FMA            GGML_F16x8_FMA
+    #define GGML_F16_VEC_ADD            GGML_F16x8_ADD
+    #define GGML_F16_VEC_MUL            GGML_F16x8_MUL
+    #define GGML_F16_VEC_REDUCE         GGML_F16x8_REDUCE
+#else
+    // if FP16 vector arithmetic is not supported, we use FP32 instead
+    // and take advantage of the vcvt_ functions to convert to/from FP16
+
+    #define GGML_F16_STEP 16
+    #define GGML_F16_EPR  4
+
+    #define GGML_F32Cx4              float32x4_t
+    #define GGML_F32Cx4_ZERO         vdupq_n_f32(0.0f)
+    #define GGML_F32Cx4_SET1(x)      vdupq_n_f32(x)
+    #define GGML_F32Cx4_LOAD(x)      vcvt_f32_f16(vld1_f16((const __fp16 *)(x)))
+    #define GGML_F32Cx4_STORE(x, y)  vst1_f16(x, vcvt_f16_f32(y))
+    #define GGML_F32Cx4_FMA(a, b, c) vfmaq_f32(a, b, c)
+    #define GGML_F32Cx4_ADD          vaddq_f32
+    #define GGML_F32Cx4_MUL          vmulq_f32
+    #define GGML_F32Cx4_REDUCE       GGML_F32x4_REDUCE
+
+    #define GGML_F16_VEC                GGML_F32Cx4
+    #define GGML_F16_VEC_ZERO           GGML_F32Cx4_ZERO
+    #define GGML_F16_VEC_SET1           GGML_F32Cx4_SET1
+    #define GGML_F16_VEC_LOAD(p, i)     GGML_F32Cx4_LOAD(p)
+    #define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE((__fp16 *)(p), r[i])
+    #define GGML_F16_VEC_FMA            GGML_F32Cx4_FMA
+    #define GGML_F16_VEC_ADD            GGML_F32Cx4_ADD
+    #define GGML_F16_VEC_MUL            GGML_F32Cx4_MUL
+    #define GGML_F16_VEC_REDUCE         GGML_F32Cx4_REDUCE
+#endif
+
+#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) && defined(__ARM_FP16_FORMAT_IEEE)
+
+#define GGML_SIMD
+
+// F32 NEON
+
+#define GGML_F32_STEP 16
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4              float32x4_t
+#define GGML_F32x4_ZERO         vdupq_n_f32(0.0f)
+#define GGML_F32x4_SET1(x)      vdupq_n_f32(x)
+#define GGML_F32x4_LOAD         vld1q_f32
+#define GGML_F32x4_STORE        vst1q_f32
+#define GGML_F32x4_FMA(a, b, c) vfmaq_f32(a, b, c)
+#define GGML_F32x4_ADD          vaddq_f32
+#define GGML_F32x4_MUL          vmulq_f32
+#define GGML_F32x4_REDUCE_ONE(x) vaddvq_f32(x)
+#define GGML_F32x4_REDUCE(res, x)                       \
+{                                                       \
+    int offset = GGML_F32_ARR >> 1;                     \
+    for (int i = 0; i < offset; ++i) {                  \
+        (x)[i] = vaddq_f32((x)[i], (x)[offset+i]);      \
+    }                                                   \
+    offset >>= 1;                                       \
+    for (int i = 0; i < offset; ++i) {                  \
+        (x)[i] = vaddq_f32((x)[i], (x)[offset+i]);      \
+    }                                                   \
+    offset >>= 1;                                       \
+    for (int i = 0; i < offset; ++i) {                  \
+        (x)[i] = vaddq_f32((x)[i], (x)[offset+i]);      \
+    }                                                   \
+    (res) = (ggml_float) GGML_F32x4_REDUCE_ONE((x)[0]); \
+}
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 NEON
+
+#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)
+    #define GGML_F16_STEP 32
+    #define GGML_F16_EPR  8
+
+    #define GGML_F16x8              float16x8_t
+    #define GGML_F16x8_ZERO         vdupq_n_f16(0.0f)
+    #define GGML_F16x8_SET1(x)      vdupq_n_f16(x)
+    #define GGML_F16x8_LOAD(x)      vld1q_f16((const __fp16 *)(x))
+    #define GGML_F16x8_STORE        vst1q_f16
+    #define GGML_F16x8_FMA(a, b, c) vfmaq_f16(a, b, c)
+    #define GGML_F16x8_ADD          vaddq_f16
+    #define GGML_F16x8_MUL          vmulq_f16
+    #define GGML_F16x8_REDUCE(res, x)                               \
+    do {                                                            \
+        int offset = GGML_F16_ARR >> 1;                             \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        offset >>= 1;                                               \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        offset >>= 1;                                               \
+        for (int i = 0; i < offset; ++i) {                          \
+            (x)[i] = vaddq_f16((x)[i], (x)[offset+i]);              \
+        }                                                           \
+        const float32x4_t t0 = vcvt_f32_f16(vget_low_f16 ((x)[0])); \
+        const float32x4_t t1 = vcvt_f32_f16(vget_high_f16((x)[0])); \
+        (res) = (ggml_float) vaddvq_f32(vaddq_f32(t0, t1));         \
+    } while (0)
+
+    #define GGML_F16_VEC                GGML_F16x8
+    #define GGML_F16_VEC_ZERO           GGML_F16x8_ZERO
+    #define GGML_F16_VEC_SET1           GGML_F16x8_SET1
+    #define GGML_F16_VEC_LOAD(p, i)     GGML_F16x8_LOAD(p)
+    #define GGML_F16_VEC_STORE(p, r, i) GGML_F16x8_STORE((__fp16 *)(p), (r)[i])
+    #define GGML_F16_VEC_FMA            GGML_F16x8_FMA
+    #define GGML_F16_VEC_ADD            GGML_F16x8_ADD
+    #define GGML_F16_VEC_MUL            GGML_F16x8_MUL
+    #define GGML_F16_VEC_REDUCE         GGML_F16x8_REDUCE
+#else
+    // if FP16 vector arithmetic is not supported, we use FP32 instead
+    // and take advantage of the vcvt_ functions to convert to/from FP16
+
+    #define GGML_F16_STEP 16
+    #define GGML_F16_EPR  4
+
+    #define GGML_F32Cx4              float32x4_t
+    #define GGML_F32Cx4_ZERO         vdupq_n_f32(0.0f)
+    #define GGML_F32Cx4_SET1(x)      vdupq_n_f32(x)
+    #define GGML_F32Cx4_LOAD(x)      vcvt_f32_f16(vld1_f16((const __fp16 *)(x)))
+    #define GGML_F32Cx4_STORE(x, y)  vst1_f16(x, vcvt_f16_f32(y))
+    #define GGML_F32Cx4_FMA(a, b, c) vfmaq_f32(a, b, c)
+    #define GGML_F32Cx4_ADD          vaddq_f32
+    #define GGML_F32Cx4_MUL          vmulq_f32
+    #define GGML_F32Cx4_REDUCE       GGML_F32x4_REDUCE
+
+    #define GGML_F16_VEC                GGML_F32Cx4
+    #define GGML_F16_VEC_ZERO           GGML_F32Cx4_ZERO
+    #define GGML_F16_VEC_SET1           GGML_F32Cx4_SET1
+    #define GGML_F16_VEC_LOAD(p, i)     GGML_F32Cx4_LOAD(p)
+    #define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE((__fp16 *)(p), r[i])
+    #define GGML_F16_VEC_FMA            GGML_F32Cx4_FMA
+    #define GGML_F16_VEC_ADD            GGML_F32Cx4_ADD
+    #define GGML_F16_VEC_MUL            GGML_F32Cx4_MUL
+    #define GGML_F16_VEC_REDUCE         GGML_F32Cx4_REDUCE
+#endif
+
+#elif defined(__AVX512F__)
+
+#define GGML_SIMD
+
+// F32 AVX512
+
+#define GGML_F32_STEP 64
+#define GGML_F32_EPR  16
+
+#define GGML_F32x16         __m512
+#define GGML_F32x16_ZERO    _mm512_setzero_ps()
+#define GGML_F32x16_SET1(x) _mm512_set1_ps(x)
+#define GGML_F32x16_LOAD    _mm512_loadu_ps
+#define GGML_F32x16_STORE   _mm512_storeu_ps
+// _mm512_fmadd_ps is defined in AVX512F so no guard is required
+#define GGML_F32x16_FMA(a, b, c) _mm512_fmadd_ps(b, c, a)
+#define GGML_F32x16_ADD     _mm512_add_ps
+#define GGML_F32x16_MUL     _mm512_mul_ps
+#define GGML_F32x16_REDUCE(res, x)                                    \
+do {                                                                  \
+    int offset = GGML_F32_ARR >> 1;                                   \
+    for (int i = 0; i < offset; ++i) {                                \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                      \
+    }                                                                 \
+    offset >>= 1;                                                     \
+    for (int i = 0; i < offset; ++i) {                                \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                      \
+    }                                                                 \
+    offset >>= 1;                                                     \
+    for (int i = 0; i < offset; ++i) {                                \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                      \
+    }                                                                 \
+    res = (ggml_float) _mm512_reduce_add_ps(x[0]);                    \
+} while (0)
+
+// TODO: is this optimal ?
+
+#define GGML_F32_VEC        GGML_F32x16
+#define GGML_F32_VEC_ZERO   GGML_F32x16_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x16_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x16_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x16_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x16_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x16_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x16_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x16_REDUCE
+
+// F16 AVX512
+
+#if defined(__AVX512FP16__)
+
+#define GGML_F16_STEP 128
+#define GGML_F16_EPR  32
+
+#define GGML_F16x32              __m512h
+#define GGML_F16x32_ZERO         _mm512_setzero_ph()
+#define GGML_F16x32_SET1(x)      _mm512_set1_ph(__extension__(_Float16)(x))
+#define GGML_F16x32_LOAD(x)      _mm512_loadu_ph(x)
+#define GGML_F16x32_STORE(x, y)  _mm512_storeu_ph(x, y)
+#define GGML_F16x32_FMA(a, b, c) _mm512_fmadd_ph(b, c, a)
+#define GGML_F16x32_ADD          _mm512_add_ph
+#define GGML_F16x32_MUL          _mm512_mul_ph
+#define GGML_F16x32_REDUCE(res, x)                                     \
+do {                                                                   \
+    int offset = GGML_F16_ARR >> 1;                                    \
+    for (int i = 0; i < offset; ++i) {                                 \
+        x[i] = _mm512_add_ph(x[i], x[offset+i]);                       \
+    }                                                                  \
+    offset >>= 1;                                                      \
+    for (int i = 0; i < offset; ++i) {                                 \
+        x[i] = _mm512_add_ph(x[i], x[offset+i]);                       \
+    }                                                                  \
+    offset >>= 1;                                                      \
+    for (int i = 0; i < offset; ++i) {                                 \
+        x[i] = _mm512_add_ph(x[i], x[offset+i]);                       \
+    }                                                                  \
+    res = (ggml_float) _mm512_reduce_add_ph(x[0]);                     \
+} while (0)
+
+#define GGML_F16_VEC                GGML_F16x32
+#define GGML_F16_VEC_ZERO           GGML_F16x32_ZERO
+#define GGML_F16_VEC_SET1           GGML_F16x32_SET1
+#define GGML_F16_VEC_LOAD(p, i)     GGML_F16x32_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i) GGML_F16x32_STORE(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F16x32_FMA
+#define GGML_F16_VEC_ADD            GGML_F16x32_ADD
+#define GGML_F16_VEC_MUL            GGML_F16x32_MUL
+#define GGML_F16_VEC_REDUCE         GGML_F16x32_REDUCE
+
+#else // Fallback FP16 <-> FP32
+
+#define GGML_F16_STEP 64
+#define GGML_F16_EPR  16
+
+#define GGML_F32Cx16             __m512
+#define GGML_F32Cx16_ZERO        _mm512_setzero_ps()
+#define GGML_F32Cx16_SET1(x)     _mm512_set1_ps(x)
+
+// unlike  _mm256_cvt intrinsics that require F16C, _mm512_cvt is defined in AVX512F
+// so F16C guard isn't required
+#define GGML_F32Cx16_LOAD(x)     _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x)))
+#define GGML_F32Cx16_STORE(x, y) _mm256_storeu_si256((__m256i *)(x), _mm512_cvtps_ph(y, 0))
+
+#define GGML_F32Cx16_FMA(a, b, c) _mm512_fmadd_ps(b, c, a)
+#define GGML_F32Cx16_ADD         _mm512_add_ps
+#define GGML_F32Cx16_MUL         _mm512_mul_ps
+#define GGML_F32Cx16_REDUCE(res, x)                               \
+do {                                                              \
+    int offset = GGML_F32_ARR >> 1;                               \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm512_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    res = (ggml_float) _mm512_reduce_add_ps(x[0]);                \
+} while (0)
+
+#define GGML_F16_VEC                GGML_F32Cx16
+#define GGML_F16_VEC_ZERO           GGML_F32Cx16_ZERO
+#define GGML_F16_VEC_SET1           GGML_F32Cx16_SET1
+#define GGML_F16_VEC_LOAD(p, i)     GGML_F32Cx16_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx16_STORE(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F32Cx16_FMA
+#define GGML_F16_VEC_ADD            GGML_F32Cx16_ADD
+#define GGML_F16_VEC_MUL            GGML_F32Cx16_MUL
+
+#define GGML_F16_VEC_REDUCE         GGML_F32Cx16_REDUCE
+
+#endif // __AVX512FP16__
+#elif defined(__AVX__)
+
+#define GGML_SIMD
+
+// F32 AVX
+
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  8
+
+#define GGML_F32x8         __m256
+#define GGML_F32x8_ZERO    _mm256_setzero_ps()
+#define GGML_F32x8_SET1(x) _mm256_set1_ps(x)
+#define GGML_F32x8_LOAD    _mm256_loadu_ps
+#define GGML_F32x8_STORE   _mm256_storeu_ps
+#if defined(__FMA__)
+    #define GGML_F32x8_FMA(a, b, c) _mm256_fmadd_ps(b, c, a)
+#else
+    #define GGML_F32x8_FMA(a, b, c) _mm256_add_ps(_mm256_mul_ps(b, c), a)
+#endif
+#define GGML_F32x8_ADD     _mm256_add_ps
+#define GGML_F32x8_MUL     _mm256_mul_ps
+#define GGML_F32x8_REDUCE(res, x)                                 \
+do {                                                              \
+    int offset = GGML_F32_ARR >> 1;                               \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm256_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm256_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm256_add_ps(x[i], x[offset+i]);                  \
+    }                                                             \
+    const __m128 t0 = _mm_add_ps(_mm256_castps256_ps128(x[0]),    \
+                                 _mm256_extractf128_ps(x[0], 1)); \
+    const __m128 t1 = _mm_hadd_ps(t0, t0);                        \
+    res = (ggml_float) _mm_cvtss_f32(_mm_hadd_ps(t1, t1));        \
+} while (0)
+// TODO: is this optimal ?
+
+#define GGML_F32_VEC        GGML_F32x8
+#define GGML_F32_VEC_ZERO   GGML_F32x8_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x8_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x8_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x8_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x8_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x8_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x8_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x8_REDUCE
+
+// F16 AVX
+
+#define GGML_F16_STEP 32
+#define GGML_F16_EPR  8
+
+// F16 arithmetic is not supported by AVX, so we use F32 instead
+
+#define GGML_F32Cx8             __m256
+#define GGML_F32Cx8_ZERO        _mm256_setzero_ps()
+#define GGML_F32Cx8_SET1(x)     _mm256_set1_ps(x)
+
+#if defined(__F16C__)
+// the  _mm256_cvt intrinsics require F16C
+#define GGML_F32Cx8_LOAD(x)     _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x)))
+#define GGML_F32Cx8_STORE(x, y) _mm_storeu_si128((__m128i *)(x), _mm256_cvtps_ph(y, 0))
+#else
+static inline __m256 __avx_f32cx8_load(const ggml_fp16_t * x) {
+    float tmp[8];
+
+    for (int i = 0; i < 8; i++) {
+        tmp[i] = GGML_CPU_FP16_TO_FP32(x[i]);
+    }
+
+    return _mm256_loadu_ps(tmp);
+}
+static inline void __avx_f32cx8_store(ggml_fp16_t *x, __m256 y) {
+    float arr[8];
+
+    _mm256_storeu_ps(arr, y);
+
+    for (int i = 0; i < 8; i++)
+        x[i] = GGML_CPU_FP32_TO_FP16(arr[i]);
+}
+#define GGML_F32Cx8_LOAD(x)     __avx_f32cx8_load(x)
+#define GGML_F32Cx8_STORE(x, y) __avx_f32cx8_store(x, y)
+#endif
+
+#define GGML_F32Cx8_FMA         GGML_F32x8_FMA
+#define GGML_F32Cx8_ADD         _mm256_add_ps
+#define GGML_F32Cx8_MUL         _mm256_mul_ps
+#define GGML_F32Cx8_REDUCE      GGML_F32x8_REDUCE
+
+#define GGML_F16_VEC                GGML_F32Cx8
+#define GGML_F16_VEC_ZERO           GGML_F32Cx8_ZERO
+#define GGML_F16_VEC_SET1           GGML_F32Cx8_SET1
+#define GGML_F16_VEC_LOAD(p, i)     GGML_F32Cx8_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx8_STORE(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F32Cx8_FMA
+#define GGML_F16_VEC_ADD            GGML_F32Cx8_ADD
+#define GGML_F16_VEC_MUL            GGML_F32Cx8_MUL
+#define GGML_F16_VEC_REDUCE         GGML_F32Cx8_REDUCE
+
+#elif defined(__POWER9_VECTOR__)
+
+#define GGML_SIMD
+
+// F32 POWER9
+
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4              vector float
+#define GGML_F32x4_ZERO         {0.0f}
+#define GGML_F32x4_SET1         vec_splats
+#define GGML_F32x4_LOAD(p)      vec_xl(0, p)
+#define GGML_F32x4_STORE(p, r)  vec_xst(r, 0, p)
+#define GGML_F32x4_FMA(a, b, c) vec_madd(b, c, a)
+#define GGML_F32x4_ADD          vec_add
+#define GGML_F32x4_MUL          vec_mul
+#define GGML_F32x4_REDUCE(res, x)              \
+{                                              \
+    int offset = GGML_F32_ARR >> 1;            \
+    for (int i = 0; i < offset; ++i) {         \
+        x[i] = vec_add(x[i], x[offset+i]);     \
+    }                                          \
+    offset >>= 1;                              \
+    for (int i = 0; i < offset; ++i) {         \
+        x[i] = vec_add(x[i], x[offset+i]);     \
+    }                                          \
+    offset >>= 1;                              \
+    for (int i = 0; i < offset; ++i) {         \
+        x[i] = vec_add(x[i], x[offset+i]);     \
+    }                                          \
+    res = vec_extract(x[0], 0) +               \
+          vec_extract(x[0], 1) +               \
+          vec_extract(x[0], 2) +               \
+          vec_extract(x[0], 3);                \
+}
+#define GGML_F32x4_REDUCE_4(res, s0, s1, s2, s3)        \
+{                                                       \
+    vector float v = vec_add(vec_add(s0, s1),           \
+                             vec_add(s2, s3));          \
+    v = vec_add(v, vec_sld(v, v, 8));                   \
+    v = vec_add(v, vec_sld(v, v, 4));                   \
+    res += (ggml_float) vec_extract(v, 0);              \
+}
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 POWER9
+#define GGML_F16_STEP       GGML_F32_STEP
+#define GGML_F16_EPR        GGML_F32_EPR
+#define GGML_F16_VEC        GGML_F32x4
+#define GGML_F16_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F16_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F16_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F16_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F16_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F16_VEC_REDUCE GGML_F32x4_REDUCE
+// Use vec_xl, not vec_ld, in case the load address is not aligned.
+#define GGML_F16_VEC_LOAD(p, i) (i & 0x1) ?                   \
+  vec_extract_fp32_from_shorth(vec_xl(0, p - GGML_F16_EPR)) : \
+  vec_extract_fp32_from_shortl(vec_xl(0, p))
+static inline unsigned char ggml_endian_byte(int i) {
+       uint16_t tmp_val = 1;
+       return ((unsigned char *)&tmp_val)[i];
+}
+#define GGML_ENDIAN_BYTE(i) ggml_endian_byte(i)
+#define GGML_F16_VEC_STORE(p, r, i)                             \
+  if (i & 0x1)                                                  \
+    vec_xst(vec_pack_to_short_fp32(r[i - GGML_ENDIAN_BYTE(1)],  \
+                                   r[i - GGML_ENDIAN_BYTE(0)]), \
+            0, p - GGML_F16_EPR)
+
+//BF16 POWER9
+#define GGML_BF16_STEP 16
+#define GGML_BF16_EPR  8
+
+#define GGML_BF16x8         vector unsigned short
+#define GGML_BF16x8_ZERO    vec_splats((unsigned short)0)
+#define GGML_BF16x8_LOAD(p) vec_xl(0, (const unsigned short *)(p))
+
+#define GGML_BF16_VEC          GGML_BF16x8
+#define GGML_BF16_VEC_ZERO     GGML_BF16x8_ZERO
+#define GGML_BF16_VEC_LOAD     GGML_BF16x8_LOAD
+#if defined(__LITTLE_ENDIAN__)
+#define GGML_BF16_TO_F32_LO(v) ((vector float) vec_mergel(GGML_BF16_VEC_ZERO, (v)))
+#define GGML_BF16_TO_F32_HI(v) ((vector float) vec_mergeh(GGML_BF16_VEC_ZERO, (v)))
+#else
+#define GGML_BF16_TO_F32_LO(v) ((vector float) vec_mergel((v), GGML_BF16_VEC_ZERO))
+#define GGML_BF16_TO_F32_HI(v) ((vector float) vec_mergeh((v), GGML_BF16_VEC_ZERO))
+#endif
+#define GGML_BF16_FMA_LO(acc, x, y) \
+    (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_LO(x), GGML_BF16_TO_F32_LO(y))
+#define GGML_BF16_FMA_HI(acc, x, y) \
+    (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_HI(x), GGML_BF16_TO_F32_HI(y))
+
+#elif defined(__wasm_simd128__)
+
+#define GGML_SIMD
+
+// F32 WASM
+
+#define GGML_F32_STEP 16
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4              v128_t
+#define GGML_F32x4_ZERO         wasm_f32x4_splat(0.0f)
+#define GGML_F32x4_SET1(x)      wasm_f32x4_splat(x)
+#define GGML_F32x4_LOAD         wasm_v128_load
+#define GGML_F32x4_STORE        wasm_v128_store
+#define GGML_F32x4_FMA(a, b, c) wasm_f32x4_add(wasm_f32x4_mul(b, c), a)
+#define GGML_F32x4_ADD          wasm_f32x4_add
+#define GGML_F32x4_MUL          wasm_f32x4_mul
+#define GGML_F32x4_REDUCE(res, x)                  \
+{                                                  \
+    int offset = GGML_F32_ARR >> 1;                \
+    for (int i = 0; i < offset; ++i) {             \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);  \
+    }                                              \
+    offset >>= 1;                                  \
+    for (int i = 0; i < offset; ++i) {             \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);  \
+    }                                              \
+    offset >>= 1;                                  \
+    for (int i = 0; i < offset; ++i) {             \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);  \
+    }                                              \
+    res = wasm_f32x4_extract_lane(x[0], 0) +       \
+          wasm_f32x4_extract_lane(x[0], 1) +       \
+          wasm_f32x4_extract_lane(x[0], 2) +       \
+          wasm_f32x4_extract_lane(x[0], 3);        \
+}
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 WASM
+
+#define GGML_F16_STEP 16
+#define GGML_F16_EPR  4
+
+inline static v128_t __wasm_f16x4_load(const ggml_fp16_t * p) {
+    float tmp[4];
+
+    tmp[0] = GGML_CPU_FP16_TO_FP32(p[0]);
+    tmp[1] = GGML_CPU_FP16_TO_FP32(p[1]);
+    tmp[2] = GGML_CPU_FP16_TO_FP32(p[2]);
+    tmp[3] = GGML_CPU_FP16_TO_FP32(p[3]);
+
+    return wasm_v128_load(tmp);
+}
+
+inline static void __wasm_f16x4_store(ggml_fp16_t * p, v128_t x) {
+    float tmp[4];
+
+    wasm_v128_store(tmp, x);
+
+    p[0] = GGML_CPU_FP32_TO_FP16(tmp[0]);
+    p[1] = GGML_CPU_FP32_TO_FP16(tmp[1]);
+    p[2] = GGML_CPU_FP32_TO_FP16(tmp[2]);
+    p[3] = GGML_CPU_FP32_TO_FP16(tmp[3]);
+}
+
+#define GGML_F16x4             v128_t
+#define GGML_F16x4_ZERO        wasm_f32x4_splat(0.0f)
+#define GGML_F16x4_SET1(x)     wasm_f32x4_splat(x)
+#define GGML_F16x4_LOAD(x)     __wasm_f16x4_load(x)
+#define GGML_F16x4_STORE(x, y) __wasm_f16x4_store(x, y)
+#define GGML_F16x4_FMA         GGML_F32x4_FMA
+#define GGML_F16x4_ADD         wasm_f32x4_add
+#define GGML_F16x4_MUL         wasm_f32x4_mul
+#define GGML_F16x4_REDUCE(res, x)                           \
+{                                                           \
+    int offset = GGML_F16_ARR >> 1;                         \
+    for (int i = 0; i < offset; ++i) {                      \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);           \
+    }                                                       \
+    offset >>= 1;                                           \
+    for (int i = 0; i < offset; ++i) {                      \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);           \
+    }                                                       \
+    offset >>= 1;                                           \
+    for (int i = 0; i < offset; ++i) {                      \
+        x[i] = wasm_f32x4_add(x[i], x[offset+i]);           \
+    }                                                       \
+    res = (ggml_float) (wasm_f32x4_extract_lane(x[0], 0) +  \
+          wasm_f32x4_extract_lane(x[0], 1) +                \
+          wasm_f32x4_extract_lane(x[0], 2) +                \
+          wasm_f32x4_extract_lane(x[0], 3));                \
+}
+
+#define GGML_F16_VEC                GGML_F16x4
+#define GGML_F16_VEC_ZERO           GGML_F16x4_ZERO
+#define GGML_F16_VEC_SET1           GGML_F16x4_SET1
+#define GGML_F16_VEC_LOAD(p, i)     GGML_F16x4_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i) GGML_F16x4_STORE(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F16x4_FMA
+#define GGML_F16_VEC_ADD            GGML_F16x4_ADD
+#define GGML_F16_VEC_MUL            GGML_F16x4_MUL
+#define GGML_F16_VEC_REDUCE         GGML_F16x4_REDUCE
+
+#elif defined(__SSE3__)
+
+#define GGML_SIMD
+
+// F32 SSE
+
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4         __m128
+#define GGML_F32x4_ZERO    _mm_setzero_ps()
+#define GGML_F32x4_SET1(x) _mm_set1_ps(x)
+#define GGML_F32x4_LOAD    _mm_loadu_ps
+#define GGML_F32x4_STORE   _mm_storeu_ps
+#if defined(__FMA__)
+    // TODO: Does this work?
+    #define GGML_F32x4_FMA(a, b, c) _mm_fmadd_ps(b, c, a)
+#else
+    #define GGML_F32x4_FMA(a, b, c) _mm_add_ps(_mm_mul_ps(b, c), a)
+#endif
+#define GGML_F32x4_ADD     _mm_add_ps
+#define GGML_F32x4_MUL     _mm_mul_ps
+#define GGML_F32x4_REDUCE(res, x)                                 \
+{                                                                 \
+    int offset = GGML_F32_ARR >> 1;                               \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm_add_ps(x[i], x[offset+i]);                     \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm_add_ps(x[i], x[offset+i]);                     \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = _mm_add_ps(x[i], x[offset+i]);                     \
+    }                                                             \
+    const __m128 t0 = _mm_hadd_ps(x[0], x[0]);                    \
+    res = (ggml_float) _mm_cvtss_f32(_mm_hadd_ps(t0, t0));        \
+}
+// TODO: is this optimal ?
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 SSE
+
+#define GGML_F16_STEP 32
+#define GGML_F16_EPR  4
+
+static inline __m128 __sse_f16x4_load(const ggml_fp16_t * x) {
+    float tmp[4];
+
+    tmp[0] = GGML_CPU_FP16_TO_FP32(x[0]);
+    tmp[1] = GGML_CPU_FP16_TO_FP32(x[1]);
+    tmp[2] = GGML_CPU_FP16_TO_FP32(x[2]);
+    tmp[3] = GGML_CPU_FP16_TO_FP32(x[3]);
+
+    return _mm_loadu_ps(tmp);
+}
+
+static inline void __sse_f16x4_store(ggml_fp16_t * x, __m128 y) {
+    float arr[4];
+
+    _mm_storeu_ps(arr, y);
+
+    x[0] = GGML_CPU_FP32_TO_FP16(arr[0]);
+    x[1] = GGML_CPU_FP32_TO_FP16(arr[1]);
+    x[2] = GGML_CPU_FP32_TO_FP16(arr[2]);
+    x[3] = GGML_CPU_FP32_TO_FP16(arr[3]);
+}
+
+#define GGML_F32Cx4             __m128
+#define GGML_F32Cx4_ZERO        _mm_setzero_ps()
+#define GGML_F32Cx4_SET1(x)     _mm_set1_ps(x)
+#define GGML_F32Cx4_LOAD(x)     __sse_f16x4_load(x)
+#define GGML_F32Cx4_STORE(x, y) __sse_f16x4_store(x, y)
+#define GGML_F32Cx4_FMA         GGML_F32x4_FMA
+#define GGML_F32Cx4_ADD         _mm_add_ps
+#define GGML_F32Cx4_MUL         _mm_mul_ps
+#define GGML_F32Cx4_REDUCE      GGML_F32x4_REDUCE
+
+#define GGML_F16_VEC                 GGML_F32Cx4
+#define GGML_F16_VEC_ZERO            GGML_F32Cx4_ZERO
+#define GGML_F16_VEC_SET1            GGML_F32Cx4_SET1
+#define GGML_F16_VEC_LOAD(p, i)      GGML_F32Cx4_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i)  GGML_F32Cx4_STORE(p, r[i])
+#define GGML_F16_VEC_FMA             GGML_F32Cx4_FMA
+#define GGML_F16_VEC_ADD             GGML_F32Cx4_ADD
+#define GGML_F16_VEC_MUL             GGML_F32Cx4_MUL
+#define GGML_F16_VEC_REDUCE          GGML_F32Cx4_REDUCE
+
+#elif defined(__loongarch_asx)
+
+#define GGML_SIMD
+
+// F32 LASX
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  8
+
+#define GGML_F32x8         __m256
+#define GGML_F32x8_ZERO    (__m256)__lasx_xvldi(0)
+#define GGML_F32x8_SET1(x) (__m256)__lasx_xvreplfr2vr_s((x))
+#define GGML_F32x8_LOAD(x) (__m256)__lasx_xvld((x), 0)
+#define GGML_F32x8_STORE(x,y)   __lasx_xvst((y), (x), 0)
+#define GGML_F32x8_FMA(a, b, c) __lasx_xvfmadd_s(b, c, a)
+#define GGML_F32x8_ADD     __lasx_xvfadd_s
+#define GGML_F32x8_MUL     __lasx_xvfmul_s
+#define GGML_F32x8_REDUCE(res, x)                                 \
+do {                                                              \
+    int offset = GGML_F32_ARR >> 1;                               \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = __lasx_xvfadd_s(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = __lasx_xvfadd_s(x[i], x[offset+i]);                  \
+    }                                                             \
+    offset >>= 1;                                                 \
+    for (int i = 0; i < offset; ++i) {                            \
+        x[i] = __lasx_xvfadd_s(x[i], x[offset+i]);                  \
+    }                                                             \
+    float *tmp_p = (float *)&x[0]; \
+    res = tmp_p[0] + tmp_p[1] + tmp_p[2] + tmp_p[3] + tmp_p[4] + tmp_p[5] + tmp_p[6] + tmp_p[7];  \
+} while (0)
+// TODO: is this optimal ?
+
+#define GGML_F32_VEC        GGML_F32x8
+#define GGML_F32_VEC_ZERO   GGML_F32x8_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x8_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x8_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x8_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x8_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x8_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x8_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x8_REDUCE
+
+// F16 LASX
+
+#define GGML_F16_STEP 32
+#define GGML_F16_EPR  8
+
+// F16 arithmetic is not supported by LASX, so we use F32 instead
+
+#define GGML_F32Cx8          __m256
+#define GGML_F32Cx8_ZERO    (__m256)__lasx_xvldi(0)
+#define GGML_F32Cx8_SET1(x) (__m256)__lasx_xvreplfr2vr_s((x))
+
+static inline __m256 __lasx_f32cx8_load(const ggml_fp16_t * x) {
+    __m256i a;
+    memcpy(&a, x, sizeof(ggml_fp16_t) * 8);
+    a = __lasx_xvpermi_d(a, 0 | (1 << 4));
+    return __lasx_xvfcvtl_s_h(a);
+}
+
+static inline void __lasx_f32cx8_store(ggml_fp16_t * x, __m256 y) {
+    __m256i a = __lasx_xvfcvt_h_s(y, y);
+    a = __lasx_xvpermi_d(a, 0 | (2 << 2));
+    memcpy(x, &a, sizeof(ggml_fp16_t) * 8);
+}
+#define GGML_F32Cx8_LOAD(x)     __lasx_f32cx8_load(x)
+#define GGML_F32Cx8_STORE(x, y) __lasx_f32cx8_store(x, y)
+
+#define GGML_F32Cx8_FMA         GGML_F32x8_FMA
+#define GGML_F32Cx8_ADD         __lasx_xvfadd_s
+#define GGML_F32Cx8_MUL         __lasx_xvfmul_s
+#define GGML_F32Cx8_REDUCE      GGML_F32x8_REDUCE
+
+#define GGML_F16_VEC                GGML_F32Cx8
+#define GGML_F16_VEC_ZERO           GGML_F32Cx8_ZERO
+#define GGML_F16_VEC_SET1           GGML_F32Cx8_SET1
+#define GGML_F16_VEC_LOAD(p, i)     GGML_F32Cx8_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx8_STORE(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F32Cx8_FMA
+#define GGML_F16_VEC_ADD            GGML_F32Cx8_ADD
+#define GGML_F16_VEC_MUL            GGML_F32Cx8_MUL
+#define GGML_F16_VEC_REDUCE         GGML_F32Cx8_REDUCE
+
+#elif defined(__loongarch_sx)
+
+#define GGML_SIMD
+
+// F32 LSX
+
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4         __m128
+#define GGML_F32x4_ZERO    (__m128)__lsx_vldi(0)
+#define GGML_F32x4_SET1(x) (__m128)__lsx_vreplfr2vr_s((x))
+#define GGML_F32x4_LOAD(x) (__m128)__lsx_vld((x), 0)
+#define GGML_F32x4_STORE(x, y)   __lsx_vst(y, x, 0)
+#define GGML_F32x4_FMA(a, b, c) __lsx_vfmadd_s(b, c, a)
+#define GGML_F32x4_ADD     __lsx_vfadd_s
+#define GGML_F32x4_MUL     __lsx_vfmul_s
+
+#define GGML_F32x4_REDUCE(res, x)                               \
+{                                                               \
+    int offset = GGML_F32_ARR >> 1;                             \
+    for (int i = 0; i < offset; ++i) {                          \
+        x[i] = __lsx_vfadd_s(x[i], x[offset+i]);                \
+    }                                                           \
+    offset >>= 1;                                               \
+    for (int i = 0; i < offset; ++i) {                          \
+        x[i] = __lsx_vfadd_s(x[i], x[offset+i]);                \
+    }                                                           \
+    offset >>= 1;                                               \
+    for (int i = 0; i < offset; ++i) {                          \
+        x[i] = __lsx_vfadd_s(x[i], x[offset+i]);                \
+    }                                                           \
+    __m128i t0 = __lsx_vpickev_w((__m128i)x[0], (__m128i)x[0]); \
+    __m128i t1 = __lsx_vpickod_w((__m128i)x[0], (__m128i)x[0]); \
+    __m128 t2 = __lsx_vfadd_s((__m128)t0, (__m128)t1);          \
+    __m128i t3 = __lsx_vpickev_w((__m128i)t2, (__m128i)t2);     \
+    __m128i t4 = __lsx_vpickod_w((__m128i)t2, (__m128i)t2);     \
+    __m128 t5 = __lsx_vfadd_s((__m128)t3, (__m128)t4);          \
+    res = (ggml_float) ((v4f32)t5)[0];                          \
+}
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 LSX
+
+#define GGML_F16_STEP 32
+#define GGML_F16_EPR  4
+
+static inline __m128 __lsx_f16x4_load(const ggml_fp16_t * x) {
+    return __lsx_vfcvtl_s_h(__lsx_vld((const void *)x, 0));
+}
+
+static inline void __lsx_f16x4_store(ggml_fp16_t * x, __m128 y) {
+    __m128i a = __lsx_vfcvt_h_s(y, y);
+    memcpy(x, &a, sizeof(ggml_fp16_t) * 4);
+}
+
+#define GGML_F32Cx4             __m128
+#define GGML_F32Cx4_ZERO        (__m128)__lsx_vldi(0)
+#define GGML_F32Cx4_SET1(x)     (__m128)__lsx_vreplfr2vr_s((x))
+#define GGML_F32Cx4_LOAD(x)     (__m128)__lsx_f16x4_load(x)
+#define GGML_F32Cx4_STORE(x, y) __lsx_f16x4_store(x, y)
+#define GGML_F32Cx4_FMA         GGML_F32x4_FMA
+#define GGML_F32Cx4_ADD         __lsx_vfadd_s
+#define GGML_F32Cx4_MUL         __lsx_vfmul_s
+#define GGML_F32Cx4_REDUCE      GGML_F32x4_REDUCE
+
+#define GGML_F16_VEC                 GGML_F32Cx4
+#define GGML_F16_VEC_ZERO            GGML_F32Cx4_ZERO
+#define GGML_F16_VEC_SET1            GGML_F32Cx4_SET1
+#define GGML_F16_VEC_LOAD(p, i)      GGML_F32Cx4_LOAD(p)
+#define GGML_F16_VEC_STORE(p, r, i)  GGML_F32Cx4_STORE(p, r[i])
+#define GGML_F16_VEC_FMA             GGML_F32Cx4_FMA
+#define GGML_F16_VEC_ADD             GGML_F32Cx4_ADD
+#define GGML_F16_VEC_MUL             GGML_F32Cx4_MUL
+#define GGML_F16_VEC_REDUCE          GGML_F32Cx4_REDUCE
+
+#elif defined(__VXE__) || defined(__VXE2__)
+
+#define GGML_SIMD
+
+// F32 s390x
+
+#define GGML_F32_STEP 32
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4              float32x4_t
+#define GGML_F32x4_ZERO         vec_splats(0.0f)
+#define GGML_F32x4_SET1         vec_splats
+#define GGML_F32x4_LOAD(p)      vec_xl(0, p)
+#define GGML_F32x4_STORE(p, r)  vec_xst(r, 0, p)
+#define GGML_F32x4_FMA(a, b, c) vec_madd(b, c, a)
+#define GGML_F32x4_ADD          vec_add
+#define GGML_F32x4_MUL          vec_mul
+#define GGML_F32x4_REDUCE(res, x)                   \
+{                                                   \
+    int offset = GGML_F32_ARR >> 1;                 \
+    for (int i = 0; i < offset; ++i) {              \
+        x[i] = vec_add(x[i], x[offset + i]);        \
+    }                                               \
+    offset >>= 1;                                   \
+    for (int i = 0; i < offset; ++i) {              \
+        x[i] = vec_add(x[i], x[offset + i]);        \
+    }                                               \
+    offset >>= 1;                                   \
+    for (int i = 0; i < offset; ++i) {              \
+        x[i] = vec_add(x[i], x[offset + i]);        \
+    }                                               \
+    float32x4_t tmp = x[0] + vec_reve(x[0]);        \
+    res = tmp[0] + tmp[1];                          \
+}
+#define GGML_F32x4_REDUCE_4(res, s0, s1, s2, s3) \
+{                                                \
+    float32x4_t v = vec_add(vec_add(s0, s1),     \
+                            vec_add(s2, s3));    \
+    v = vec_add(v, vec_sld(v, v, 8));            \
+    v = vec_add(v, vec_sld(v, v, 4));            \
+    res += (ggml_float)vec_extract(v, 0);        \
+}
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+// F16 s390x
+#define GGML_F16_STEP GGML_F32_STEP
+#define GGML_F16_EPR  GGML_F32_EPR
+
+static inline uint32x4_t __lzs_f32cx4_to_f16(float32x4_t v_f) {
+    float32x4_t v_base = vec_mul(vec_mul(vec_abs(v_f), vec_splats(0x1.0p+112f)), vec_splats(0x1.0p-110f));
+
+    const uint32x4_t v_w      = (uint32x4_t)v_f;
+    const uint32x4_t v_shl1_w = vec_add(v_w, v_w);
+    const uint32x4_t v_sign   = vec_and(v_w, vec_splats(UINT32_C(0x80000000)));
+    const uint32x4_t v_bias   = vec_max(vec_and(v_shl1_w, vec_splats(UINT32_C(0xFF000000))), vec_splats(UINT32_C(0x71000000)));
+
+    v_base = vec_add((float32x4_t)vec_add(vec_sr(v_bias, 1), vec_splats(UINT32_C(0x07800000))), v_base);
+
+    const uint32x4_t v_bits    = (uint32x4_t)v_base;
+    const uint32x4_t v_nonsign = vec_add(vec_and(vec_sr(v_bits, 13), vec_splats(UINT32_C(0x00007C00))),
+                                         vec_and(v_bits, vec_splats(UINT32_C(0x00000FFF))));
+    const uint32x4_t v_is_nan  = (uint32x4_t)vec_cmpgt(v_shl1_w, vec_splats(UINT32_C(0xFF000000)));
+
+    return vec_or(vec_sr(v_sign, 16), vec_sel(v_nonsign, vec_splats(UINT32_C(0x7E00)), v_is_nan));
+}
+
+static inline float32x4_t __lzs_f16cx4_load(const ggml_fp16_t * x) {
+    float tmp[4];
+
+    for (int i = 0; i < 4; i++) {
+        tmp[i] = GGML_CPU_FP16_TO_FP32(x[i]);
+    }
+
+    // note: keep type-cast here to prevent compiler bugs
+    // see: https://github.com/ggml-org/llama.cpp/issues/12846
+    return vec_xl(0, (const float *)(tmp));
+}
+
+static inline void __lzs_f16cx4_store(ggml_fp16_t * x, float32x4_t v_y) {
+    const uint32x4_t v_h = __lzs_f32cx4_to_f16(v_y);
+    const uint64_t   tmp = ((uint64x2_t)vec_pack(v_h, v_h))[0];
+    memcpy(x, &tmp, sizeof(tmp));
+}
+
+#define GGML_F16_VEC                GGML_F32x4
+#define GGML_F16_VEC_ZERO           GGML_F32x4_ZERO
+#define GGML_F16_VEC_SET1           GGML_F32x4_SET1
+#define GGML_F16_VEC_LOAD(p, i)     __lzs_f16cx4_load(p)
+#define GGML_F16_VEC_STORE(p, r, i) __lzs_f16cx4_store(p, r[i])
+#define GGML_F16_VEC_FMA            GGML_F32x4_FMA
+#define GGML_F16_VEC_ADD            GGML_F32x4_ADD
+#define GGML_F16_VEC_MUL            GGML_F32x4_MUL
+#define GGML_F16_VEC_REDUCE         GGML_F32x4_REDUCE
+
+// BF16 s390x
+#define GGML_BF16_STEP 16
+#define GGML_BF16_EPR  8
+
+#define GGML_BF16x8         __vector unsigned short
+#define GGML_BF16x8_ZERO    vec_splats((unsigned short)0)
+#define GGML_BF16x8_LOAD(p) vec_xl(0, (const unsigned short *)(p))
+
+#define GGML_BF16_VEC      GGML_BF16x8
+#define GGML_BF16_VEC_ZERO GGML_BF16x8_ZERO
+#define GGML_BF16_VEC_LOAD GGML_BF16x8_LOAD
+#define GGML_BF16_TO_F32_LO(v) ((float32x4_t) vec_mergel((v), GGML_BF16_VEC_ZERO))
+#define GGML_BF16_TO_F32_HI(v) ((float32x4_t) vec_mergeh((v), GGML_BF16_VEC_ZERO))
+#define GGML_BF16_FMA_LO(acc, x, y) \
+    (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_LO(x), GGML_BF16_TO_F32_LO(y))
+#define GGML_BF16_FMA_HI(acc, x, y) \
+    (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_HI(x), GGML_BF16_TO_F32_HI(y))
+
+#elif defined(__riscv_v_intrinsic)
+
+// compatible with vlen >= 128
+
+#define GGML_SIMD
+
+// F32
+
+#define GGML_F32_STEP 16
+#define GGML_F32_EPR  4
+
+#define GGML_F32x4              vfloat32m1_t
+#define GGML_F32x4_ZERO         __riscv_vfmv_v_f_f32m1(0.0f, GGML_F32_EPR)
+#define GGML_F32x4_SET1(x)      __riscv_vfmv_v_f_f32m1(x, GGML_F32_EPR)
+#define GGML_F32x4_LOAD(x)      __riscv_vle32_v_f32m1(x, GGML_F32_EPR)
+#define GGML_F32x4_STORE(b, v)  __riscv_vse32_v_f32m1(b, v, GGML_F32_EPR)
+#define GGML_F32x4_FMA(a, b, c) __riscv_vfmacc_vv_f32m1(a, b, c, GGML_F32_EPR)
+#define GGML_F32x4_ADD(a, b)    __riscv_vfadd_vv_f32m1(a, b, GGML_F32_EPR)
+#define GGML_F32x4_MUL(a, b)    __riscv_vfmul_vv_f32m1(a, b, GGML_F32_EPR)
+
+#define GGML_F32_VEC        GGML_F32x4
+#define GGML_F32_VEC_ZERO   GGML_F32x4_ZERO
+#define GGML_F32_VEC_SET1   GGML_F32x4_SET1
+#define GGML_F32_VEC_LOAD   GGML_F32x4_LOAD
+#define GGML_F32_VEC_STORE  GGML_F32x4_STORE
+#define GGML_F32_VEC_FMA    GGML_F32x4_FMA
+#define GGML_F32_VEC_ADD    GGML_F32x4_ADD
+#define GGML_F32_VEC_MUL    GGML_F32x4_MUL
+#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE
+
+#endif
+
+// GGML_F32_ARR / GGML_F16_ARR
+//   number of registers to use per step
+#ifdef GGML_SIMD
+#define GGML_F32_ARR (GGML_F32_STEP/GGML_F32_EPR)
+#define GGML_F16_ARR (GGML_F16_STEP/GGML_F16_EPR)
+#endif
+
+// GGML_F16_DOT_*
+// like GGML_F16_* but for dot products which need F32 accumulation on AVX512-FP16
+
+#if defined(__AVX512FP16__)
+
+#define GGML_F16_DOT_STEP           GGML_F32_STEP
+#define GGML_F16_DOT_EPR            GGML_F32_EPR
+#define GGML_F16_DOT_ARR            GGML_F32_ARR
+#define GGML_F16_DOT_VEC            GGML_F32x16
+#define GGML_F16_DOT_VEC_ZERO       GGML_F32x16_ZERO
+#define GGML_F16_DOT_VEC_LOAD(p, i) _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(p)))
+#define GGML_F16_DOT_VEC_FMA        GGML_F32x16_FMA
+#define GGML_F16_DOT_VEC_REDUCE     GGML_F32x16_REDUCE
+
+#else
+
+#define GGML_F16_DOT_STEP           GGML_F16_STEP
+#define GGML_F16_DOT_EPR            GGML_F16_EPR
+#define GGML_F16_DOT_ARR            GGML_F16_ARR
+#define GGML_F16_DOT_VEC            GGML_F16_VEC
+#define GGML_F16_DOT_VEC_ZERO       GGML_F16_VEC_ZERO
+#define GGML_F16_DOT_VEC_LOAD       GGML_F16_VEC_LOAD
+#define GGML_F16_DOT_VEC_FMA        GGML_F16_VEC_FMA
+#define GGML_F16_DOT_VEC_REDUCE     GGML_F16_VEC_REDUCE
+
+#endif // defined(__AVX512FP16__)
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp
new file mode 100644
index 0000000000..3918a6f4d9
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp
@@ -0,0 +1,613 @@
+#include "vec.h"
+
+#include <cassert>
+
+// precomputed gelu table for f16 (128 KB)
+ggml_fp16_t ggml_table_gelu_f16[1 << 16];
+
+// precomputed quick gelu table for f16 (128 KB)
+ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16];
+
+void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) {
+   assert(nrc == 1);
+   GGML_UNUSED(nrc);
+   GGML_UNUSED(bx);
+   GGML_UNUSED(by);
+   GGML_UNUSED(bs);
+
+#if defined(GGML_SIMD)
+    float sumf = 0.0f;
+
+    #if defined(__ARM_FEATURE_SVE)
+        const int sve_register_length = ggml_cpu_get_sve_cnt() * 8;
+        const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16
+        const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers
+
+        const int np = (n & ~(ggml_f32_step - 1));
+        svfloat32_t sum1 = svdup_n_f32(0.0f);
+        svfloat32_t sum2 = svdup_n_f32(0.0f);
+        svfloat32_t sum3 = svdup_n_f32(0.0f);
+        svfloat32_t sum4 = svdup_n_f32(0.0f);
+        svfloat32_t sum5 = svdup_n_f32(0.0f);
+        svfloat32_t sum6 = svdup_n_f32(0.0f);
+        svfloat32_t sum7 = svdup_n_f32(0.0f);
+        svfloat32_t sum8 = svdup_n_f32(0.0f);
+        svfloat32_t ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8;
+        svfloat32_t ay1,ay2,ay3,ay4,ay5,ay6,ay7,ay8;
+        for (int i = 0; i < np; i += ggml_f32_step) {
+            ax1 = GGML_F32_VEC_LOAD(x + i);
+            ay1 = GGML_F32_VEC_LOAD(y + i);
+            sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1);
+
+            ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr);
+            ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr);
+            sum2 = GGML_F32_VEC_FMA(sum2, ax2, ay2);
+
+            ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr);
+            ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr);
+            sum3 = GGML_F32_VEC_FMA(sum3, ax3, ay3);
+
+            ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr);
+            ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr);
+            sum4 = GGML_F32_VEC_FMA(sum4, ax4, ay4);
+
+            ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr);
+            ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr);
+            sum5 = GGML_F32_VEC_FMA(sum5, ax5, ay5);
+
+            ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr);
+            ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr);
+            sum6 = GGML_F32_VEC_FMA(sum6, ax6, ay6);
+
+            ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr);
+            ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr);
+            sum7 = GGML_F32_VEC_FMA(sum7, ax7, ay7);
+
+            ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr);
+            ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr);
+            sum8 = GGML_F32_VEC_FMA(sum8, ax8, ay8);
+        }
+        // leftovers
+        // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop
+        const int np2 = (n & ~(ggml_f32_epr - 1));
+        for (int i = np; i < np2; i += ggml_f32_epr) {
+            ax1 = GGML_F32_VEC_LOAD(x + i);
+            ay1 = GGML_F32_VEC_LOAD(y + i);
+            sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1);
+        }
+        // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmla on available elements only
+        if (np2 < n) {
+            svbool_t pg = svwhilelt_b32(np2, n);
+            ax1 = svld1_f32(pg, x + np2);
+            ay1 = svld1_f32(pg, y + np2);
+            sum1 = svmla_f32_m(pg, sum1, ax1, ay1);
+        }
+        // reduce sum1,sum2 to sum1
+        GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8);
+    #elif defined(__riscv_v_intrinsic)
+        int vl = __riscv_vsetvlmax_e32m8();
+        vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1);
+        vfloat32m8_t vsum;
+        vfloat32m8_t ax;
+        vfloat32m8_t ay;
+        vsum = __riscv_vfmv_v_f_f32m8_tu(vsum, 0.0f, vl);
+        for (int i = 0; i < n; i += vl) {
+            vl = __riscv_vsetvl_e32m8(n - i);
+            ax = __riscv_vle32_v_f32m8_tu(ax, &x[i], vl);
+            ay = __riscv_vle32_v_f32m8_tu(ay, &y[i], vl);
+            vsum = __riscv_vfmacc_vv_f32m8_tu(vsum, ax, ay, vl);
+        }
+        vl = __riscv_vsetvlmax_e32m8();
+        vs = __riscv_vfredusum_vs_f32m8_f32m1(vsum, vs, vl);
+        sumf += __riscv_vfmv_f_s_f32m1_f32(vs);
+    #else
+        const int np = (n & ~(GGML_F32_STEP - 1));
+
+        GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO };
+
+        GGML_F32_VEC ax[GGML_F32_ARR];
+        GGML_F32_VEC ay[GGML_F32_ARR];
+
+        for (int i = 0; i < np; i += GGML_F32_STEP) {
+            for (int j = 0; j < GGML_F32_ARR; j++) {
+                ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR);
+                ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR);
+
+                sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], ay[j]);
+            }
+        }
+
+        // reduce sum0..sum3 to sum0
+        GGML_F32_VEC_REDUCE(sumf, sum);
+
+        // leftovers
+        for (int i = np; i < n; ++i) {
+            sumf += x[i]*y[i];
+        }
+    #endif
+#else
+    // scalar
+    ggml_float sumf = 0.0;
+    for (int i = 0; i < n; ++i) {
+        sumf += (ggml_float)(x[i]*y[i]);
+    }
+#endif
+
+    *s = sumf;
+}
+
+void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc) {
+    assert(nrc == 1);
+    GGML_UNUSED(nrc);
+    GGML_UNUSED(bx);
+    GGML_UNUSED(by);
+    GGML_UNUSED(bs);
+    int i = 0;
+    ggml_float sumf = 0;
+
+#if defined(__AVX512BF16__)
+    __m512 c1 = _mm512_setzero_ps();
+    __m512 c2 = _mm512_setzero_ps();
+    for (; i + 64 <= n; i += 64) {
+        c1 = _mm512_dpbf16_ps(c1, m512bh(_mm512_loadu_si512((x + i))),
+                             m512bh(_mm512_loadu_si512((y + i))));
+        c2 = _mm512_dpbf16_ps(c2, m512bh(_mm512_loadu_si512((x + i + 32))),
+                             m512bh(_mm512_loadu_si512((y + i + 32))));
+    }
+    sumf += (ggml_float)_mm512_reduce_add_ps(c1);
+    sumf += (ggml_float)_mm512_reduce_add_ps(c2);
+
+#elif defined(__AVX512F__)
+#define LOAD(p) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(p))), 16))
+    __m512 c1 = _mm512_setzero_ps();
+    __m512 c2 = _mm512_setzero_ps();
+    for (; i + 32 <= n; i += 32) {
+        c1 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i), LOAD(y + i)), c1);
+        c2 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c2);
+    }
+    sumf += (ggml_float)_mm512_reduce_add_ps(c1);
+    sumf += (ggml_float)_mm512_reduce_add_ps(c2);
+
+#undef LOAD
+#elif defined(__AVX2__) || defined(__AVX__)
+#if defined(__AVX2__)
+#define LOAD(p) _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16))
+#else
+#define LOAD(p) _mm256_castsi256_ps(_mm256_insertf128_si256(_mm256_castsi128_si256(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)), (_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_bsrli_si128(_mm_loadu_si128((const __m128i *)(p)), 8)), 16)), 1))
+#endif
+    __m256 c1 = _mm256_setzero_ps();
+    __m256 c2 = _mm256_setzero_ps();
+    __m256 c3 = _mm256_setzero_ps();
+    __m256 c4 = _mm256_setzero_ps();
+    for (; i + 32 <= n; i += 32) {
+        c1 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i), LOAD(y + i)), c1);
+        c2 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 8), LOAD(y + i + 8)), c2);
+        c3 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c3);
+        c4 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 24), LOAD(y + i + 24)), c4);
+    }
+    __m128 g;
+    c1 = _mm256_add_ps(_mm256_add_ps(c1, c3),
+                       _mm256_add_ps(c2, c4));
+    g = _mm_add_ps(_mm256_extractf128_ps(c1, 1),
+                   _mm256_castps256_ps128(c1));
+    g = _mm_add_ps(g, _mm_movehl_ps(g, g));
+    g = _mm_add_ss(g, _mm_movehdup_ps(g));
+    sumf += (ggml_float)_mm_cvtss_f32(g);
+
+#undef LOAD
+#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfwma)
+    size_t vl = __riscv_vsetvlmax_e32m4();
+
+    // initialize accumulators to all zeroes
+    vfloat32m4_t vsum0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+    vfloat32m4_t vsum1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+
+    // calculate step size
+    const size_t epr = __riscv_vsetvlmax_e16m2();
+    const size_t step = epr * 2;
+    const int np = (n & ~(step - 1));
+
+    // unroll by 2
+    for (; i < np; i += step) {
+        vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], epr);
+        vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], epr);
+        vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, epr);
+        __asm__ __volatile__ ("" ::: "memory");
+
+        vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i + epr], epr);
+        vbfloat16m2_t ay1 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i + epr], epr);
+        vsum1 = __riscv_vfwmaccbf16_vv_f32m4(vsum1, ax1, ay1, epr);
+        __asm__ __volatile__ ("" ::: "memory");
+    }
+
+    // accumulate in 1 register
+    vsum0 = __riscv_vfadd_vv_f32m4(vsum0, vsum1, vl);
+
+    // leftovers
+    for (i = np; i < n; i += vl) {
+        vl = __riscv_vsetvl_e16m2(n - i);
+        vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], vl);
+        vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], vl);
+        vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, vl);
+    }
+
+    // reduce
+    vl = __riscv_vsetvlmax_e32m4();
+    vfloat32m1_t redsum = __riscv_vfredusum_vs_f32m4_f32m1(vsum0, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
+    sumf += __riscv_vfmv_f_s_f32m1_f32(redsum);
+
+#elif defined(__POWER9_VECTOR__) || defined(__VXE__) || defined(__VXE2__)
+    const int np = (n & ~(GGML_BF16_STEP - 1));
+    if (np > 0) {
+        GGML_F32_VEC sum[4] = {GGML_F32_VEC_ZERO};
+        for (; i < np; i += GGML_BF16_STEP) {
+            GGML_BF16_VEC vx0 = GGML_BF16_VEC_LOAD(x + i);
+            GGML_BF16_VEC vx1 = GGML_BF16_VEC_LOAD(x + i + 8);
+            GGML_BF16_VEC vy0 = GGML_BF16_VEC_LOAD(y + i);
+            GGML_BF16_VEC vy1 = GGML_BF16_VEC_LOAD(y + i + 8);
+            GGML_BF16_FMA_LO(sum[0], vx0, vy0);
+            GGML_BF16_FMA_HI(sum[1], vx0, vy0);
+            GGML_BF16_FMA_LO(sum[2], vx1, vy1);
+            GGML_BF16_FMA_HI(sum[3], vx1, vy1);
+        }
+        GGML_F32x4_REDUCE_4(sumf, sum[0], sum[1], sum[2], sum[3]);
+    }
+#endif
+
+    for (; i < n; ++i) {
+        sumf += (ggml_float)(GGML_BF16_TO_FP32(x[i]) *
+                             GGML_BF16_TO_FP32(y[i]));
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc) {
+    assert(nrc == 1);
+    GGML_UNUSED(nrc);
+    GGML_UNUSED(bx);
+    GGML_UNUSED(by);
+    GGML_UNUSED(bs);
+
+    ggml_float sumf = 0.0;
+
+
+#if defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+        const int ggml_f16_epr = svcnth();
+        const int ggml_f16_step = 8 * ggml_f16_epr;
+        const int np = n - (n % ggml_f16_step);
+        const int np2 = n - (n % ggml_f16_epr);
+
+        svfloat32_t sum1_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum1_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum2_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum2_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum3_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum3_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum4_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum4_hi = svdup_n_f32(0.0f);
+
+        for (int i = 0; i < np; i += ggml_f16_step) {
+            ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0), GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0));
+            ggml_sve_f16_fma_widened(&sum2_lo, &sum2_hi, GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1), GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1));
+            ggml_sve_f16_fma_widened(&sum3_lo, &sum3_hi, GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2), GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2));
+            ggml_sve_f16_fma_widened(&sum4_lo, &sum4_hi, GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3), GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3));
+            ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4), GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4));
+            ggml_sve_f16_fma_widened(&sum2_lo, &sum2_hi, GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5), GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5));
+            ggml_sve_f16_fma_widened(&sum3_lo, &sum3_hi, GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6), GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6));
+            ggml_sve_f16_fma_widened(&sum4_lo, &sum4_hi, GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7), GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7));
+        }
+
+        for (int i = np; i < np2; i += ggml_f16_epr) {
+            ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i, 0), GGML_F16x_VEC_LOAD(y + i, 0));
+        }
+
+        if (np2 < n) {
+            const svbool_t pg = svwhilelt_b16(np2, n);
+            const svfloat16_t rx = svld1_f16(pg, (const __fp16 *)(x + np2));
+            const svfloat16_t ry = svld1_f16(pg, (const __fp16 *)(y + np2));
+
+            ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, rx, ry);
+        }
+
+        sum1_lo = svadd_f32_m(DEFAULT_PG32, sum1_lo, sum2_lo);
+        sum1_hi = svadd_f32_m(DEFAULT_PG32, sum1_hi, sum2_hi);
+        sum3_lo = svadd_f32_m(DEFAULT_PG32, sum3_lo, sum4_lo);
+        sum3_hi = svadd_f32_m(DEFAULT_PG32, sum3_hi, sum4_hi);
+        sum1_lo = svadd_f32_m(DEFAULT_PG32, sum1_lo, sum3_lo);
+        sum1_hi = svadd_f32_m(DEFAULT_PG32, sum1_hi, sum3_hi);
+
+        sumf = ggml_sve_sum_f32x2(sum1_lo, sum1_hi);
+    #elif defined(__riscv_v_intrinsic)
+        #if defined(__riscv_zvfh)
+            int vl = __riscv_vsetvlmax_e32m2();
+            vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1);
+            vfloat32m2_t vsum;
+            vfloat16m1_t ax;
+            vfloat16m1_t ay;
+            vsum = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vmv_v_x_u32m2(0, vl));
+            for (int i = 0; i < n; i += vl) {
+                vl = __riscv_vsetvl_e16m1(n - i);
+                ax = __riscv_vle16_v_f16m1_tu(ax, (const _Float16 *)&x[i], vl);
+                ay = __riscv_vle16_v_f16m1_tu(ay, (const _Float16 *)&y[i], vl);
+                vsum = __riscv_vfwmacc_vv_f32m2_tu(vsum, ax, ay, vl);
+            }
+            vl = __riscv_vsetvlmax_e32m1();
+            vfloat32m1_t ac0 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(vsum, 0), __riscv_vget_v_f32m2_f32m1(vsum, 1), vl);
+            vs = __riscv_vfredusum_vs_f32m1_f32m1(ac0, vs, vl);
+            sumf += __riscv_vfmv_f_s_f32m1_f32(vs);
+        #else
+            for (int i = 0; i < n; ++i) {
+                sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i]));
+            }
+        #endif // __riscv_zvfh
+    #else
+        const int np = (n & ~(GGML_F16_DOT_STEP - 1));
+
+        GGML_F16_DOT_VEC sum[GGML_F16_DOT_ARR] = { GGML_F16_DOT_VEC_ZERO };
+
+        GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR];
+        GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR];
+
+        for (int i = 0; i < np; i += GGML_F16_DOT_STEP) {
+            for (int j = 0; j < GGML_F16_DOT_ARR; j++) {
+                ax[j] = GGML_F16_DOT_VEC_LOAD(x + i + j*GGML_F16_DOT_EPR, j);
+                ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j);
+
+                sum[j] = GGML_F16_DOT_VEC_FMA(sum[j], ax[j], ay[j]);
+            }
+        }
+
+        // reduce sum0..sum3 to sum0
+        GGML_F16_DOT_VEC_REDUCE(sumf, sum);
+
+        // leftovers
+        for (int i = np; i < n; ++i) {
+            sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i]));
+        }
+        // if you hit this, you are likely running outside the FP range
+        assert(!isnan(sumf) && !isinf(sumf));
+    #endif
+#else
+    for (int i = 0; i < n; ++i) {
+        sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i]));
+    }
+#endif // GGML_SIMD
+
+    *s = sumf;
+}
+
+void ggml_vec_silu_f32(const int n, float * y, const float * x) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, ggml_v_silu(_mm512_loadu_ps(x + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, ggml_v_silu(_mm256_loadu_ps(x + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, ggml_v_silu(_mm_loadu_ps(x + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, ggml_v_silu(pg, svld1_f32(pg, x + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, ggml_v_silu(vld1q_f32(x + i)));
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vy = ggml_v_silu_m2(vx, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_silu_f32(x[i]);
+    }
+}
+
+void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_silu(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_silu(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_silu(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_silu(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, vmulq_f32(ggml_v_silu(vld1q_f32(x + i)), vld1q_f32(g + i)));
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl);
+        vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_silu_m2(vx, vl), vg, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_silu_f32(x[i]) * g[i];
+    }
+}
+
+ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) {
+    int i = 0;
+    ggml_float sum = 0;
+// TODO: optimize to process the remaining elements in groups using the smaller vector sizes from AVX2 and SSE
+// ref: https://github.com/ggml-org/llama.cpp/pull/15953#pullrequestreview-3310928344
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        __m512 val = _mm512_sub_ps(_mm512_loadu_ps(x + i),
+                                   _mm512_set1_ps(mean));
+        _mm512_storeu_ps(y + i, val);
+        sum += (ggml_float)_mm512_reduce_add_ps(_mm512_mul_ps(val, val));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        __m256 val = _mm256_sub_ps(_mm256_loadu_ps(x + i),
+                                   _mm256_set1_ps(mean));
+        _mm256_storeu_ps(y + i, val);
+        val = _mm256_mul_ps(val,val);
+        __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1),
+                                 _mm256_castps256_ps128(val));
+        val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2));
+        val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2));
+        sum += (ggml_float)_mm_cvtss_f32(val2);
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        __m128 val = _mm_sub_ps(_mm_loadu_ps(x + i),
+                                _mm_set1_ps(mean));
+        _mm_storeu_ps(y + i, val);
+        val = _mm_mul_ps(val, val);
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+        val = _mm_add_ps(val, _mm_movehl_ps(val, val));
+        val = _mm_add_ss(val, _mm_movehdup_ps(val));
+#else
+        __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1));
+        val = _mm_add_ps(val, tmp);
+        tmp = _mm_movehl_ps(tmp, val);
+        val = _mm_add_ss(val, tmp);
+#endif  // __AVX__ || __AVX2__ || __AVX512F__
+        sum += (ggml_float)_mm_cvtss_f32(val);
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        float32x4_t val = vsubq_f32(vld1q_f32(x + i),
+                                    vdupq_n_f32(mean));
+        vst1q_f32(y + i, val);
+        val = vmulq_f32(val, val);
+        sum += (ggml_float)vaddvq_f32(val);
+    }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        float32x4_t val = vec_sub(vec_xl(0, x + i), vec_splats(mean));
+        vec_xst(val, 0, y + i);
+        val = vec_mul(val, val);
+        sum += (ggml_float)vec_hsum_f32x4(val);
+    }
+#elif defined(__riscv_v_intrinsic)
+    vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1);
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t val = __riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], vl), mean, vl);
+        __riscv_vse32_v_f32m2(&y[i], val, vl);
+        val = __riscv_vfmul_vv_f32m2(val, val, vl);
+        vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, vl);
+    }
+    sum = (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum);
+#endif
+    for (; i < n; ++i) {
+        float val = x[i] - mean;
+        y[i] = val;
+        val *= val;
+        sum += (ggml_float)val;
+    }
+    return sum/n;
+}
+
+ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max) {
+    int i = 0;
+    ggml_float sum = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        __m512 val = ggml_v_expf(_mm512_sub_ps(_mm512_loadu_ps(x + i),
+                                               _mm512_set1_ps(max)));
+        _mm512_storeu_ps(y + i, val);
+        sum += (ggml_float)_mm512_reduce_add_ps(val);
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        __m256 val = ggml_v_expf(_mm256_sub_ps(_mm256_loadu_ps(x + i),
+                                               _mm256_set1_ps(max)));
+        _mm256_storeu_ps(y + i, val);
+        __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1),
+                                 _mm256_castps256_ps128(val));
+        val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2));
+        val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2));
+        sum += (ggml_float)_mm_cvtss_f32(val2);
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        __m128 val = ggml_v_expf(_mm_sub_ps(_mm_loadu_ps(x + i),
+                                            _mm_set1_ps(max)));
+        _mm_storeu_ps(y + i, val);
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+        val = _mm_add_ps(val, _mm_movehl_ps(val, val));
+        val = _mm_add_ss(val, _mm_movehdup_ps(val));
+#else
+        __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1));
+        val = _mm_add_ps(val, tmp);
+        tmp = _mm_movehl_ps(tmp, val);
+        val = _mm_add_ss(val, tmp);
+#endif
+        sum += (ggml_float)_mm_cvtss_f32(val);
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svfloat32_t val = ggml_v_expf(pg, svsub_f32_x(pg, svld1_f32(pg, x + i),
+                                                svdup_n_f32_x(pg, max)));
+        svst1_f32(pg, y + i, val);
+        sum += (ggml_float)svaddv_f32(pg, val);
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        float32x4_t val = ggml_v_expf(vsubq_f32(vld1q_f32(x + i),
+                                                vdupq_n_f32(max)));
+        vst1q_f32(y + i, val);
+        sum += (ggml_float)vaddvq_f32(val);
+    }
+#elif defined(__riscv_v_intrinsic)
+    vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1);
+    for (int avl; i < n; i += avl) {
+        avl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t val = ggml_v_expf_m2(__riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], avl), max, avl), avl);
+        __riscv_vse32_v_f32m2(&y[i], val, avl);
+        vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, avl);
+    }
+    return (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum);
+#endif
+    for (; i < n; ++i) {
+        float val = expf(x[i] - max);
+        sum += (ggml_float)val;
+        y[i] = val;
+    }
+    return sum;
+}
+
+ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max) {
+    // log(soft_max) = log(soft_max_i / soft_max_sum) = log(soft_max_i) - log(soft_max_sum) = (logit_i - max) - log(soft_max_i)
+
+    int i = 0;
+    ggml_float sum = 0;
+    for (; i < n; ++i) {
+        float val = x[i] - max;
+        y[i] = val;
+        sum += (ggml_float)expf(val);
+    }
+    return sum = (ggml_float)logf(sum);
+}
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h
new file mode 100644
index 0000000000..59fa88f2c9
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h
@@ -0,0 +1,1630 @@
+// Vectorized functions for fundamental operations
+
+#pragma once
+
+#include "ggml-impl.h"
+#include "simd-mappings.h"
+#include "ggml.h"
+#include "ggml-cpu.h"
+
+#if defined(GGML_USE_ACCELERATE)
+#include <Accelerate/Accelerate.h>
+#endif
+
+// floating point type used to accumulate sums
+typedef double ggml_float;
+
+#if defined(__ARM_FEATURE_SVE)
+inline static void ggml_sve_f16_fma_widened(
+        svfloat32_t * acc_lo,
+        svfloat32_t * acc_hi,
+        svfloat16_t x,
+        svfloat16_t y) {
+#if defined(__ARM_FEATURE_SVE2)
+    *acc_lo = svmlalb_f32(*acc_lo, x, y);
+    *acc_hi = svmlalt_f32(*acc_hi, x, y);
+#else
+    // Plain SVE fallback path if SVE2 instructions not available
+    svfloat16_t x_even = svtrn1_f16(x, x);
+    svfloat16_t x_odd = svtrn2_f16(x, x);
+
+    svfloat16_t y_even = svtrn1_f16(y, y);
+    svfloat16_t y_odd = svtrn2_f16(y, y);
+
+    svbool_t pg = svptrue_b32();
+
+    *acc_lo = svmla_f32_x(pg, *acc_lo, svcvt_f32_f16_x(pg, x_even), svcvt_f32_f16_x(pg, y_even));
+    *acc_hi = svmla_f32_x(pg, *acc_hi, svcvt_f32_f16_x(pg, x_odd), svcvt_f32_f16_x(pg, y_odd));
+#endif
+}
+
+inline static ggml_float ggml_sve_sum_f32x2(svfloat32_t sum_lo, svfloat32_t sum_hi) {
+    return (ggml_float) (svaddv_f32(svptrue_b32(), sum_lo) + svaddv_f32(svptrue_b32(), sum_hi));
+}
+#endif
+
+#define GGML_GELU_FP16
+#define GGML_GELU_QUICK_FP16
+
+#define GGML_SOFT_MAX_UNROLL 4
+#define GGML_VEC_DOT_UNROLL  2
+#define GGML_VEC_MAD_UNROLL  32
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//
+// global data
+//
+
+// precomputed gelu table for f16 (128 KB)
+extern ggml_fp16_t ggml_table_gelu_f16[1 << 16];
+
+// precomputed quick gelu table for f16 (128 KB)
+extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16];
+
+//
+// fundamental operations
+//
+
+void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc);
+void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc);
+void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc);
+
+void ggml_vec_silu_f32(const int n, float * y, const float * x);
+ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean )
+ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max);
+ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max);
+
+inline static void ggml_vec_set_i8(const int n, int8_t * x, const int8_t v) { for (int i = 0; i < n; ++i) x[i] = v; }
+inline static void ggml_vec_set_i16(const int n, int16_t * x, const int16_t v) { for (int i = 0; i < n; ++i) x[i] = v; }
+
+inline static void ggml_vec_set_i32(const int n, int32_t * x, const int32_t   v) { for (int i = 0; i < n; ++i) x[i] = v;    }
+inline static void ggml_vec_cpy_i32(const int n, int32_t * y, const int32_t * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; }
+
+inline static void ggml_vec_set_f16(const int n, ggml_fp16_t * x, const ggml_fp16_t v) { for (int i = 0; i < n; ++i) x[i] = v; }
+inline static void ggml_vec_set_bf16(const int n, ggml_bf16_t * x, const ggml_bf16_t v) { for (int i = 0; i < n; ++i) x[i] = v; }
+
+inline static void ggml_vec_add_f32 (const int n, float * z, const float * x, const float * y) {
+    int i = 0;
+#if defined(__AVX2__)
+    for (; i + 7 < n; i += 8) {
+        __m256 vx = _mm256_loadu_ps(x + i);
+        __m256 vy = _mm256_loadu_ps(y + i);
+        __m256 vz = _mm256_add_ps(vx, vy);
+        _mm256_storeu_ps(z + i, vz);
+    }
+#endif
+    for (; i < n; ++i) {
+        z[i] = x[i] + y[i];
+    }
+}
+
+inline static void ggml_vec_add_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) {
+    for (int i = 0; i < n; ++i) {
+        z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) + GGML_CPU_FP16_TO_FP32(y[i]));
+    }
+}
+inline static void ggml_vec_add1_f32(const int n, float * z, const float * x, const float   v) { for (int i = 0; i < n; ++i) z[i]  = x[i] + v;    }
+inline static void ggml_vec_acc_f32 (const int n, float * y, const float * x)                  { for (int i = 0; i < n; ++i) y[i] += x[i];        }
+inline static void ggml_vec_acc1_f32(const int n, float * y, const float   v)                  { for (int i = 0; i < n; ++i) y[i] += v;           }
+inline static void ggml_vec_sub_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i]  = x[i] - y[i]; }
+inline static void ggml_vec_sub_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) {
+    for (int i = 0; i < n; ++i) {
+        z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) - GGML_CPU_FP16_TO_FP32(y[i]));
+    }
+}
+inline static void ggml_vec_set_f32 (const int n, float * x, const float   v)                  { for (int i = 0; i < n; ++i) x[i]  = v;           }
+inline static void ggml_vec_cpy_f32 (const int n, float * y, const float * x)                  { for (int i = 0; i < n; ++i) y[i]  = x[i];        }
+inline static void ggml_vec_neg_f32 (const int n, float * y, const float * x)                  { for (int i = 0; i < n; ++i) y[i]  = -x[i];       }
+inline static void ggml_vec_neg_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(-GGML_CPU_FP16_TO_FP32(x[i]));
+    }
+}
+
+inline static void ggml_vec_mul_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i]  = x[i]*y[i];   }
+inline static void ggml_vec_mul_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) {
+    for (int i = 0; i < n; ++i) {
+        z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) * GGML_CPU_FP16_TO_FP32(y[i]));
+    }
+}
+inline static void ggml_vec_div_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i]  = x[i]/y[i];   }
+inline static void ggml_vec_div_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) {
+    for (int i = 0; i < n; ++i) {
+        z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) / GGML_CPU_FP16_TO_FP32(y[i]));
+    }
+}
+
+// compute GGML_VEC_DOT_UNROLL dot products at once
+// xs - x row stride in bytes
+inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GGML_RESTRICT s, void * GGML_RESTRICT xv, ggml_fp16_t * GGML_RESTRICT y) {
+    ggml_float sumf[GGML_VEC_DOT_UNROLL] = { 0.0 };
+
+    ggml_fp16_t * GGML_RESTRICT x[GGML_VEC_DOT_UNROLL];
+
+    for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) {
+        x[i] = (ggml_fp16_t *) ((char *) xv + i*xs);
+    }
+
+#if defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+
+        const int ggml_f16_epr = svcnth();
+        const int ggml_f16_step = 2 * ggml_f16_epr;
+        int np = n - (n % ggml_f16_step);
+        int np2 = n - (n % ggml_f16_epr);
+
+        svfloat32_t sum_0_0_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum_0_0_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum_0_1_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum_0_1_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum_1_0_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum_1_0_hi = svdup_n_f32(0.0f);
+        svfloat32_t sum_1_1_lo = svdup_n_f32(0.0f);
+        svfloat32_t sum_1_1_hi = svdup_n_f32(0.0f);
+
+        for (int i = 0; i < np; i += ggml_f16_step) {
+            const svfloat16_t ay0 = GGML_F16x_VEC_LOAD(y + i, 0);
+            const svfloat16_t ax00 = GGML_F16x_VEC_LOAD(x[0] + i, 0);
+            const svfloat16_t ax01 = GGML_F16x_VEC_LOAD(x[1] + i, 0);
+
+            ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, ax00, ay0);
+            ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, ax01, ay0);
+
+            const svfloat16_t ay1 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 0);
+            const svfloat16_t ax10 = GGML_F16x_VEC_LOAD(x[0] + i + 1 * ggml_f16_epr, 0);
+            const svfloat16_t ax11 = GGML_F16x_VEC_LOAD(x[1] + i + 1 * ggml_f16_epr, 0);
+
+            ggml_sve_f16_fma_widened(&sum_0_1_lo, &sum_0_1_hi, ax10, ay1);
+            ggml_sve_f16_fma_widened(&sum_1_1_lo, &sum_1_1_hi, ax11, ay1);
+        }
+
+        for (int i = np; i < np2; i += ggml_f16_epr) {
+            const svfloat16_t ry = GGML_F16x_VEC_LOAD(y + i, 0);
+            const svfloat16_t rx0 = GGML_F16x_VEC_LOAD(x[0] + i, 0);
+            const svfloat16_t rx1 = GGML_F16x_VEC_LOAD(x[1] + i, 0);
+
+            ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, rx0, ry);
+            ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, rx1, ry);
+        }
+
+        if (np2 < n) {
+            const svbool_t pg = svwhilelt_b16(np2, n);
+            const svfloat16_t ay = svld1_f16(pg, (const __fp16 *)(y + np2));
+            const svfloat16_t ax0 = svld1_f16(pg, (const __fp16 *)(x[0] + np2));
+            const svfloat16_t ax1 = svld1_f16(pg, (const __fp16 *)(x[1] + np2));
+
+            ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, ax0, ay);
+            ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, ax1, ay);
+        }
+
+        svfloat32_t sum_0_lo = svadd_f32_x(DEFAULT_PG32, sum_0_0_lo, sum_0_1_lo);
+        svfloat32_t sum_0_hi = svadd_f32_x(DEFAULT_PG32, sum_0_0_hi, sum_0_1_hi);
+        svfloat32_t sum_1_lo = svadd_f32_x(DEFAULT_PG32, sum_1_0_lo, sum_1_1_lo);
+        svfloat32_t sum_1_hi = svadd_f32_x(DEFAULT_PG32, sum_1_0_hi, sum_1_1_hi);
+        sumf[0] = ggml_sve_sum_f32x2(sum_0_lo, sum_0_hi);
+        sumf[1] = ggml_sve_sum_f32x2(sum_1_lo, sum_1_hi);
+        np = n;
+    #elif defined(__riscv_v_intrinsic)
+        #if defined(__riscv_zvfh)
+            size_t vl = __riscv_vsetvlmax_e32m4();
+
+            // initialize accumulators to all zeroes
+            vfloat32m4_t vsum0_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+            vfloat32m4_t vsum0_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+            vfloat32m4_t vsum1_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+            vfloat32m4_t vsum1_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+
+            // calculate step size
+            const size_t epr = __riscv_vsetvlmax_e16m2();
+            const size_t step = epr * 2;
+            int np = (n & ~(step - 1));
+
+            // unroll by 2 along the row dimension
+            for (int i = 0; i < np; i += step) {
+                vfloat16m2_t ay0 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), epr);
+                vfloat16m2_t ax0_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), epr);
+                vfloat16m2_t ax1_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), epr);
+                vsum0_0 = __riscv_vfwmacc_vv_f32m4(vsum0_0, ax0_0, ay0, epr);
+                vsum1_0 = __riscv_vfwmacc_vv_f32m4(vsum1_0, ax1_0, ay0, epr);
+
+                vfloat16m2_t ay1 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i + epr), epr);
+                vfloat16m2_t ax0_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i + epr), epr);
+                vfloat16m2_t ax1_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i + epr), epr);
+                vsum0_1 = __riscv_vfwmacc_vv_f32m4(vsum0_1, ax0_1, ay1, epr);
+                vsum1_1 = __riscv_vfwmacc_vv_f32m4(vsum1_1, ax1_1, ay1, epr);
+            }
+
+            vfloat32m4_t vsum0 = __riscv_vfadd_vv_f32m4(vsum0_0, vsum0_1, vl);
+            vfloat32m4_t vsum1 = __riscv_vfadd_vv_f32m4(vsum1_0, vsum1_1, vl);
+
+            // leftovers
+            for (int i = np; i < n; i += vl) {
+                vl = __riscv_vsetvl_e16m2(n - i);
+                vfloat16m2_t ay = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), vl);
+                vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), vl);
+                vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), vl);
+
+                vsum0 = __riscv_vfwmacc_vv_f32m4(vsum0, ax0, ay, vl);
+                vsum1 = __riscv_vfwmacc_vv_f32m4(vsum1, ax1, ay, vl);
+            }
+
+            // reduce
+            vl = __riscv_vsetvlmax_e32m2();
+            vfloat32m2_t acc0_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum0, 0),
+                                        __riscv_vget_v_f32m4_f32m2(vsum0, 1), vl);
+            vl = __riscv_vsetvlmax_e32m1();
+            vfloat32m1_t acc0_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc0_0, 0),
+            __riscv_vget_v_f32m2_f32m1(acc0_0, 1), vl);
+            vfloat32m1_t redsum0 = __riscv_vfredusum_vs_f32m1_f32m1(
+                                        acc0_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
+
+            vl = __riscv_vsetvlmax_e32m2();
+            vfloat32m2_t acc1_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum1, 0),
+                                        __riscv_vget_v_f32m4_f32m2(vsum1, 1), vl);
+            vl = __riscv_vsetvlmax_e32m1();
+            vfloat32m1_t acc1_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc1_0, 0),
+                                        __riscv_vget_v_f32m2_f32m1(acc1_0, 1), vl);
+            vfloat32m1_t redsum1 = __riscv_vfredusum_vs_f32m1_f32m1(
+                                        acc1_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
+            sumf[0] = __riscv_vfmv_f_s_f32m1_f32(redsum0);
+            sumf[1] = __riscv_vfmv_f_s_f32m1_f32(redsum1);
+            np = n;
+        #else
+            const int np = 0;
+        #endif
+    #else
+        const int np = (n & ~(GGML_F16_DOT_STEP - 1));
+
+        GGML_F16_DOT_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_DOT_ARR] = { { GGML_F16_DOT_VEC_ZERO } };
+
+        GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR];
+        GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR];
+
+        for (int i = 0; i < np; i += GGML_F16_DOT_STEP) {
+            for (int j = 0; j < GGML_F16_DOT_ARR; j++) {
+                ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j);
+
+                for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) {
+                    ax[j] = GGML_F16_DOT_VEC_LOAD(x[k] + i + j*GGML_F16_DOT_EPR, j);
+
+                    sum[k][j] = GGML_F16_DOT_VEC_FMA(sum[k][j], ax[j], ay[j]);
+                }
+            }
+        }
+
+        // reduce sum0..sum3 to sum0
+        for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) {
+            GGML_F16_DOT_VEC_REDUCE(sumf[k], sum[k]);
+        }
+    #endif
+#else
+    // scalar path
+    const int np = 0;
+#endif
+    // scalar and leftovers
+    for (int i = np; i < n; ++i) {
+        for (int j = 0; j < GGML_VEC_DOT_UNROLL; ++j) {
+            sumf[j] += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[j][i])*GGML_CPU_FP16_TO_FP32(y[i]));
+        }
+    }
+
+    for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) {
+        s[i] = (float)sumf[i];
+    }
+}
+
+inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const float * GGML_RESTRICT x, const float v) {
+#if defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+
+        const int sve_register_length = ggml_cpu_get_sve_cnt() * 8;
+        const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16
+        const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers
+        GGML_F32_VEC vx = GGML_F32_VEC_SET1(v);
+
+        const int np = (n & ~(ggml_f32_step - 1));
+        svfloat32_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8;
+        svfloat32_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8;
+        for (int i = 0; i < np; i += ggml_f32_step) {
+
+            ax1 = GGML_F32_VEC_LOAD(x + i);
+            ay1 = GGML_F32_VEC_LOAD(y + i);
+            ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx);
+
+            GGML_F32_VEC_STORE(y + i, ay1);
+
+            ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr);
+            ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr);
+            ay2 = GGML_F32_VEC_FMA(ay2, ax2, vx);
+
+            GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2);
+
+            ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr);
+            ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr);
+            ay3 = GGML_F32_VEC_FMA(ay3, ax3, vx);
+
+            GGML_F32_VEC_STORE(y + i + 2*ggml_f32_epr, ay3);
+
+            ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr);
+            ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr);
+            ay4 = GGML_F32_VEC_FMA(ay4, ax4, vx);
+
+            GGML_F32_VEC_STORE(y + i + 3*ggml_f32_epr, ay4);
+
+            ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr);
+            ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr);
+            ay5 = GGML_F32_VEC_FMA(ay5, ax5, vx);
+
+            GGML_F32_VEC_STORE(y + i + 4*ggml_f32_epr, ay5);
+
+            ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr);
+            ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr);
+            ay6 = GGML_F32_VEC_FMA(ay6, ax6, vx);
+
+            GGML_F32_VEC_STORE(y + i + 5*ggml_f32_epr, ay6);
+
+            ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr);
+            ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr);
+            ay7 = GGML_F32_VEC_FMA(ay7, ax7, vx);
+
+            GGML_F32_VEC_STORE(y + i + 6*ggml_f32_epr, ay7);
+
+            ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr);
+            ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr);
+            ay8 = GGML_F32_VEC_FMA(ay8, ax8, vx);
+
+            GGML_F32_VEC_STORE(y + i + 7*ggml_f32_epr, ay8);
+        }
+        // leftovers
+        // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop
+        const int np2 = (n & ~(ggml_f32_epr - 1));
+        for (int i = np; i < np2; i += ggml_f32_epr) {
+            ax1 = GGML_F32_VEC_LOAD(x + i);
+            ay1 = GGML_F32_VEC_LOAD(y + i);
+            ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx);
+
+            GGML_F32_VEC_STORE(y + i, ay1);
+        }
+        // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only
+        if (np2 < n) {
+            svbool_t pg =svwhilelt_b32(np2, n);
+            ax1 = svld1_f32(pg, x + np2);
+            ay1 = svld1_f32(pg, y + np2);
+            ay1 = svmad_f32_m(pg, ax1, vx, ay1);
+
+            svst1_f32(pg, y + np2, ay1);
+        }
+    #elif defined(__riscv_v_intrinsic)
+        for (int i = 0, avl; i < n; i += avl) {
+            avl = __riscv_vsetvl_e32m8(n - i);
+            vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl);
+            vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl);
+            vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, v, ay, avl);
+            __riscv_vse32_v_f32m8(&y[i], ny, avl);
+        }
+    #else
+        const int np = (n & ~(GGML_F32_STEP - 1));
+
+        GGML_F32_VEC vx = GGML_F32_VEC_SET1(v);
+
+        GGML_F32_VEC ax[GGML_F32_ARR];
+        GGML_F32_VEC ay[GGML_F32_ARR];
+
+        for (int i = 0; i < np; i += GGML_F32_STEP) {
+            for (int j = 0; j < GGML_F32_ARR; j++) {
+                ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR);
+                ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR);
+                ay[j] = GGML_F32_VEC_FMA(ay[j], ax[j], vx);
+
+                GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]);
+            }
+        }
+
+        // leftovers
+        for (int i = np; i < n; ++i) {
+            y[i] += x[i]*v;
+        }
+    #endif
+#else
+    // scalar
+    for (int i = 0; i < n; ++i) {
+        y[i] += x[i]*v;
+    }
+#endif
+}
+
+inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) {
+#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE)
+    const int sve_register_length = svcntb() * 8;
+    const int ggml_f16_epr = sve_register_length / 16;
+    const int ggml_f16_step = 8 * ggml_f16_epr;
+
+    GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v);
+
+    int np = (n & ~(ggml_f16_step - 1));
+
+    svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8;
+    svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8;
+    for (int i = 0; i < np; i += ggml_f16_step) {
+        ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0);
+        ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0);
+        ay1 = GGML_F16x_VEC_FMA(ay1, ax1, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 0 * ggml_f16_epr, ay1, 0);
+
+        ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1);
+        ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1);
+        ay2 = GGML_F16x_VEC_FMA(ay2, ax2, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 1 * ggml_f16_epr, ay2, 1);
+
+        ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2);
+        ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2);
+        ay3 = GGML_F16x_VEC_FMA(ay3, ax3, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 2 * ggml_f16_epr, ay3, 2);
+
+        ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3);
+        ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3);
+        ay4 = GGML_F16x_VEC_FMA(ay4, ax4, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 3 * ggml_f16_epr, ay4, 3);
+
+        ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4);
+        ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4);
+        ay5 = GGML_F16x_VEC_FMA(ay5, ax5, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 4 * ggml_f16_epr, ay5, 4);
+
+        ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5);
+        ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5);
+        ay6 = GGML_F16x_VEC_FMA(ay6, ax6, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 5 * ggml_f16_epr, ay6, 5);
+
+        ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6);
+        ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6);
+        ay7 = GGML_F16x_VEC_FMA(ay7, ax7, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 6 * ggml_f16_epr, ay7, 6);
+
+        ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7);
+        ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7);
+        ay8 = GGML_F16x_VEC_FMA(ay8, ax8, vx);
+
+        GGML_F16x_VEC_STORE(y + i + 7 * ggml_f16_epr, ay8, 7);
+    }
+    const int np2 = (n & ~(ggml_f16_epr - 1));
+    for (int k = np; k < np2; k += ggml_f16_epr) {
+        svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0);
+        svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0);
+        ry = GGML_F16x_VEC_FMA(ry, rx, vx);
+
+        GGML_F16x_VEC_STORE(y + k, ry, 0);
+    }
+
+    if (np2 < n) {
+        svbool_t pg = svwhilelt_b16(np2, n);
+        svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2));
+        svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2));
+        hy = svmad_f16_x(pg, hx, vx, hy);
+        svst1_f16(pg, (__fp16 *)(y + np2), hy);
+    }
+    np = n;
+#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic
+    #if defined (__riscv_zvfh)
+        const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v);
+        const _Float16 scale = *(const _Float16*)(&s);
+
+        // calculate step size
+        const int epr = __riscv_vsetvlmax_e16m4();
+        const int step = epr * 2;
+        int np = (n & ~(step - 1));
+
+        // unroll by 2
+        for (int i = 0; i < np; i += step) {
+            vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, epr);
+            vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr);
+            ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, epr);
+            __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr);
+            __asm__ __volatile__ ("" ::: "memory");
+
+            vfloat16m4_t ax1 = __riscv_vle16_v_f16m4((const _Float16*)x + i + epr, epr);
+            vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr);
+            ay1 = __riscv_vfmacc_vf_f16m4(ay1, scale, ax1, epr);
+            __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr);
+            __asm__ __volatile__ ("" ::: "memory");
+        }
+
+        // leftovers
+        int vl;
+        for (int i = np; i < n; i += vl) {
+            vl = __riscv_vsetvl_e16m4(n - i);
+            vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, vl);
+            vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl);
+            ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, vl);
+            __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl);
+        }
+        np = n;
+    #else
+        // fall to scalar path
+        const int np = 0;
+    #endif
+#elif defined(GGML_SIMD)
+    const int np = (n & ~(GGML_F16_STEP - 1));
+
+    GGML_F16_VEC vx = GGML_F16_VEC_SET1(v);
+
+    GGML_F16_VEC ax[GGML_F16_ARR];
+    GGML_F16_VEC ay[GGML_F16_ARR];
+
+    for (int i = 0; i < np; i += GGML_F16_STEP) {
+        for (int j = 0; j < GGML_F16_ARR; j++) {
+            ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j);
+            ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j);
+            ay[j] = GGML_F16_VEC_FMA(ay[j], ax[j], vx);
+
+            GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j);
+        }
+    }
+#else
+    // scalar path
+    const int np = 0;
+#endif
+
+    // scalar and leftovers
+    for (int i = np; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i]) + GGML_CPU_FP16_TO_FP32(x[i])*v);
+    }
+}
+
+// xs and vs are byte strides of x and v
+inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int vs, float * GGML_RESTRICT y, const float * GGML_RESTRICT xv, const float * GGML_RESTRICT vv) {
+
+    const float * GGML_RESTRICT x[GGML_VEC_MAD_UNROLL];
+    const float * GGML_RESTRICT v[GGML_VEC_MAD_UNROLL];
+
+    for (int i = 0; i < GGML_VEC_MAD_UNROLL; ++i) {
+        x[i] = (const float *) ((const char *) xv + i*xs);
+        v[i] = (const float *) ((const char *) vv + i*vs);
+    }
+
+#if defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+        // scalar Route to scalar implementation       //TODO: Write SVE code
+        for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) {
+            for (int i = 0; i < n; ++i) {
+                y[i] += x[k][i]*v[k][0];
+            }
+        }
+    #elif defined(__riscv_v_intrinsic)
+        for (int i = 0, avl; i < n; i += avl) {
+            avl = __riscv_vsetvl_e32m8(n - i);
+            vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl);
+            for (int k = 0; k < GGML_VEC_MAD_UNROLL; k++) {
+                vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[k][i], avl);
+                ay = __riscv_vfmadd_vf_f32m8(ax, v[k][0], ay, avl);
+            }
+            __riscv_vse32_v_f32m8(&y[i], ay, avl);
+        }
+    #else
+        const int np = (n & ~(GGML_F32_STEP - 1));
+
+        GGML_F32_VEC vx[GGML_VEC_MAD_UNROLL];
+
+        for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) {
+            vx[k] = GGML_F32_VEC_SET1(v[k][0]);
+        }
+
+        GGML_F32_VEC ax[GGML_VEC_MAD_UNROLL][GGML_F32_ARR];
+        GGML_F32_VEC ay[GGML_F32_ARR];
+
+        for (int i = 0; i < np; i += GGML_F32_STEP) {
+            for (int j = 0; j < GGML_F32_ARR; j++) {
+                ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR);
+
+                for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) {
+                    ax[k][j] = GGML_F32_VEC_LOAD(x[k] + i + j*GGML_F32_EPR);
+                    ay[j] = GGML_F32_VEC_FMA(ay[j], ax[k][j], vx[k]);
+                }
+
+                GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]);
+            }
+        }
+
+        // leftovers
+        for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) {
+            for (int i = np; i < n; ++i) {
+                y[i] += x[k][i]*v[k][0];
+            }
+        }
+    #endif
+#else
+    // scalar
+    for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) {
+        for (int i = 0; i < n; ++i) {
+            y[i] += x[k][i]*v[k][0];
+        }
+    }
+#endif
+}
+
+inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, const float s, const float b) {
+#if defined(GGML_USE_ACCELERATE)
+    vDSP_vsmsa(x, 1, &s, &b, y, 1, n);
+#elif defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+        // scalar ; TODO: Write SVE code
+        for (int i = 0; i < n; ++i) {
+            y[i] = x[i]*s + b;
+        }
+    #elif defined(__riscv_v_intrinsic)
+        for (int i = 0, avl; i < n; i += avl) {
+            avl = __riscv_vsetvl_e32m8(n - i);
+            vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl);
+            vfloat32m8_t vb = __riscv_vfmv_v_f_f32m8(b, avl);
+            vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, s, vb, avl);
+            __riscv_vse32_v_f32m8(&y[i], ny, avl);
+        }
+    #else
+        const int np = (n & ~(GGML_F32_STEP - 1));
+
+        GGML_F32_VEC vs = GGML_F32_VEC_SET1(s);
+        GGML_F32_VEC vb = GGML_F32_VEC_SET1(b);
+
+        GGML_F32_VEC ay[GGML_F32_ARR];
+
+        for (int i = 0; i < np; i += GGML_F32_STEP) {
+            for (int j = 0; j < GGML_F32_ARR; j++) {
+                ay[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR);
+                ay[j] = GGML_F32_VEC_FMA(vb, ay[j], vs);
+
+                GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]);
+            }
+        }
+
+        // leftovers
+        for (int i = np; i < n; ++i) {
+            y[i] = x[i]*s + b;
+        }
+    #endif
+#else
+    // scalar
+    for (int i = 0; i < n; ++i) {
+        y[i] = x[i]*s + b;
+    }
+#endif
+}
+
+//inline static void ggml_vec_scale_f32(const int n, float * y, const float   v) { for (int i = 0; i < n; ++i) y[i] *= v;          }
+inline static void ggml_vec_scale_f32(const int n, float * y, const float   v) {
+#if defined(GGML_USE_ACCELERATE)
+    vDSP_vsmul(y, 1, &v, y, 1, n);
+#elif defined(GGML_SIMD)
+    #if defined(__ARM_FEATURE_SVE)
+        const int sve_register_length = ggml_cpu_get_sve_cnt() * 8;
+        const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16
+        const int ggml_f32_step = 2 * ggml_f32_epr;
+
+        GGML_F32_VEC vx = GGML_F32_VEC_SET1(v);
+        const int np = (n & ~(ggml_f32_step - 1));
+        svfloat32_t ay1;
+        svfloat32_t ay2;
+        for (int i = 0; i < np; i += ggml_f32_step) {
+            ay1 = GGML_F32_VEC_LOAD(y + i);
+            ay1 = GGML_F32_VEC_MUL(ay1, vx);
+            GGML_F32_VEC_STORE(y + i, ay1);
+
+            ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr);
+            ay2 = GGML_F32_VEC_MUL(ay2, vx);
+            GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2);
+        }
+        // leftovers
+        // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only
+        for (int i = np; i < n; i += ggml_f32_epr) {
+            svbool_t pg = svwhilelt_b32(i, n);
+            ay1 = svld1_f32(pg, y + i);
+            ay1 = svmul_f32_m(pg, ay1, vx);
+            svst1_f32(pg, y + i, ay1);
+        }
+    #elif defined(__riscv_v_intrinsic)
+        for (int i = 0, avl; i < n; i += avl) {
+            avl = __riscv_vsetvl_e32m8(n - i);
+            vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl);
+            vfloat32m8_t ny = __riscv_vfmul_vf_f32m8(ay, v, avl);
+            __riscv_vse32_v_f32m8(&y[i], ny, avl);
+        }
+    #else
+        const int np = (n & ~(GGML_F32_STEP - 1));
+
+        GGML_F32_VEC vx = GGML_F32_VEC_SET1(v);
+
+        GGML_F32_VEC ay[GGML_F32_ARR];
+
+        for (int i = 0; i < np; i += GGML_F32_STEP) {
+            for (int j = 0; j < GGML_F32_ARR; j++) {
+                ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR);
+                ay[j] = GGML_F32_VEC_MUL(ay[j], vx);
+
+                GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]);
+            }
+        }
+
+        // leftovers
+        for (int i = np; i < n; ++i) {
+            y[i] *= v;
+        }
+    #endif
+#else
+    // scalar
+    for (int i = 0; i < n; ++i) {
+        y[i] *= v;
+    }
+#endif
+}
+
+inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float v) {
+#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE)
+    const int sve_register_length = svcntb() * 8;
+    const int ggml_f16_epr = sve_register_length / 16;
+    const int ggml_f16_step = 2 * ggml_f16_epr;
+
+    GGML_F16x_VEC vx =  GGML_F16x_VEC_SET1(v);
+    int np = (n & ~(ggml_f16_step - 1));
+    svfloat16_t ay1, ay2;
+
+    for (int i = 0; i < np; i += ggml_f16_step) {
+        ay1 = GGML_F16x_VEC_LOAD(y + i + 0*ggml_f16_epr, 0);
+        ay1 = GGML_F16x_VEC_MUL(ay1, vx);
+        GGML_F16x_VEC_STORE(y + i + 0*ggml_f16_epr, ay1, 0);
+
+        ay2 = GGML_F16x_VEC_LOAD(y + i + 1*ggml_f16_epr, 1);
+        ay2 = GGML_F16x_VEC_MUL(ay2, vx);
+        GGML_F16x_VEC_STORE(y + i + 1*ggml_f16_epr, ay2, 1);
+    }
+    // leftovers
+    // maximum number of leftover elements will be less that ggmlF_16x_epr. Apply predicated svmad on available elements only
+    if (np < n) {
+        svbool_t pg = svwhilelt_b16(np, n);
+        svfloat16_t hy = svld1_f16(pg, (__fp16 *)(y + np));
+        svfloat16_t out = svmul_f16_m(pg, hy, vx);
+        svst1_f16(pg, (__fp16 *)(y + np), out);
+    }
+    np = n;
+#elif defined(__riscv_v_intrinsic)
+    #if defined(__riscv_zvfh)
+        const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v);
+        const _Float16 scale = *(const _Float16*)(&s);
+
+        // calculate step size
+        const int epr = __riscv_vsetvlmax_e16m4();
+        const int step = epr * 2;
+        int np = (n & ~(step - 1));
+
+        // unroll by 2
+        for (int i = 0; i < np; i += step) {
+            vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr);
+            ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, epr);
+            __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr);
+            __asm__ __volatile__ ("" ::: "memory");
+
+            vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr);
+            ay1 = __riscv_vfmul_vf_f16m4(ay1, scale, epr);
+            __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr);
+            __asm__ __volatile__ ("" ::: "memory");
+        }
+
+        // leftovers
+        int vl;
+        for (int i = np; i < n; i += vl) {
+            vl = __riscv_vsetvl_e16m4(n - i);
+            vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl);
+            ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, vl);
+            __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl);
+        }
+        np = n;
+    #else
+        // fall to scalar path
+        const int np = 0;
+    #endif
+#elif defined(GGML_SIMD)
+    const int np = (n & ~(GGML_F16_STEP - 1));
+
+    GGML_F16_VEC vx = GGML_F16_VEC_SET1(v);
+
+    GGML_F16_VEC ay[GGML_F16_ARR];
+
+    for (int i = 0; i < np; i += GGML_F16_STEP) {
+        for (int j = 0; j < GGML_F16_ARR; j++) {
+            ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j);
+            ay[j] = GGML_F16_VEC_MUL(ay[j], vx);
+
+            GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j);
+        }
+    }
+#else
+    // scalar path
+    const int np = 0;
+#endif
+    // scalar and leftovers
+    for (int i = np; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i])*v);
+    }
+}
+
+inline static void ggml_vec_norm_f32 (const int n, float * s, const float * x) { ggml_vec_dot_f32(n, s, 0, x, 0, x, 0, 1); *s = sqrtf(*s);   }
+inline static void ggml_vec_sqr_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]*x[i];   }
+inline static void ggml_vec_sqr_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(v*v);
+    }
+}
+inline static void ggml_vec_sqrt_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sqrtf(x[i]); }
+inline static void ggml_vec_sqrt_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(sqrtf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_log_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = logf(x[i]);  }
+inline static void ggml_vec_log_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(logf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_sin_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sinf(x[i]);  }
+inline static void ggml_vec_sin_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(sinf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_cos_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = cosf(x[i]);  }
+inline static void ggml_vec_cos_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(cosf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_abs_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fabsf(x[i]); }
+inline static void ggml_vec_abs_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(fabsf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_sgn_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : ((x[i] < 0.f) ? -1.f : 0.f); }
+inline static void ggml_vec_sgn_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? 1.f : ((v < 0.f) ? -1.f : 0.f));
+    }
+}
+inline static void ggml_vec_step_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : 0.f; }
+inline static void ggml_vec_step_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16((GGML_CPU_FP16_TO_FP32(x[i]) > 0.f) ? 1.f : 0.f);
+    }
+}
+inline static void ggml_vec_tanh_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = tanhf(x[i]);  }
+inline static void ggml_vec_tanh_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(tanhf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+inline static void ggml_vec_elu_f32  (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : expm1f(x[i]); }
+inline static void ggml_vec_elu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        const float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : expm1f(v));
+    }
+}
+inline static void ggml_vec_relu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : 0.f; }
+inline static void ggml_vec_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : 0.f);
+    }
+}
+inline static void ggml_vec_leaky_relu_f32 (const int n, float * y, const float * x, const float ns) { for (int i = 0; i < n; ++i) y[i] = ((x[i] > 0.f) ? x[i] : 0.f) + ns * ((x[i] < 0.0f) ? x[i] : 0.f); }
+inline static void ggml_vec_leaky_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const float ns) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(((v > 0.f) ? v : 0.f) + ns * ((v < 0.0f) ? v : 0.f));
+    }
+}
+inline static void ggml_vec_sigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = 1.f / (1.f + expf(-x[i])); }
+inline static void ggml_vec_sigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(1.f / (1.f + expf(-GGML_CPU_FP16_TO_FP32(x[i]))));
+    }
+}
+// TODO: optimize performance
+inline static void ggml_vec_hardswish_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i] * fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); }
+inline static void ggml_vec_hardswish_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(v * fminf(1.0f, fmaxf(0.0f, (v + 3.0f) / 6.0f)));
+    }
+}
+inline static void ggml_vec_hardsigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); }
+inline static void ggml_vec_hardsigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(fminf(1.0f, fmaxf(0.0f, (GGML_CPU_FP16_TO_FP32(x[i]) + 3.0f) / 6.0f)));
+    }
+}
+inline static void ggml_vec_exp_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = expf(x[i]); }
+inline static void ggml_vec_exp_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_CPU_FP32_TO_FP16(expf(GGML_CPU_FP16_TO_FP32(x[i])));
+    }
+}
+
+static const float GELU_COEF_A     = 0.044715f;
+static const float GELU_QUICK_COEF = -1.702f;
+static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
+static const float SQRT_2_INV      = 0.70710678118654752440084436210484f;
+
+inline static float ggml_gelu_f32(float x) {
+    return 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x)));
+}
+
+inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    const uint16_t * i16 = (const uint16_t *) x;
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_table_gelu_f16[i16[i]];
+    }
+}
+
+inline static void ggml_vec_gelu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i])));
+    }
+}
+
+inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_CPU_FP16_TO_FP32(x[i]);
+        float res = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV));
+        y[i] = GGML_CPU_FP32_TO_FP16(res);
+    }
+}
+
+inline static void ggml_vec_gelu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_BF16_TO_FP32(x[i]);
+        y[i] = GGML_FP32_TO_BF16(0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)));
+    }
+}
+
+#ifdef GGML_GELU_FP16
+inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) {
+    uint16_t t;
+    for (int i = 0; i < n; ++i) {
+        if (x[i] <= -10.0f) {
+            y[i] = 0.0f;
+        } else if (x[i] >= 10.0f) {
+            y[i] = x[i];
+        } else {
+            ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]);
+            memcpy(&t, &fp16, sizeof(uint16_t));
+            y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]);
+        }
+    }
+}
+#else
+inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_gelu_f32(x[i]);
+    }
+}
+#endif
+
+inline static void ggml_vec_gelu_erf_f32(const int n, float * y, const float * x) {
+    for (int i = 0; i < n; ++i) {
+        float xi = x[i];
+        y[i] = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV));
+    }
+}
+
+inline static float ggml_gelu_quick_f32(float x) {
+    return x*(1.0f/(1.0f+expf(GELU_QUICK_COEF*x)));
+}
+
+inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    const uint16_t * i16 = (const uint16_t *) x;
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_table_gelu_quick_f16[i16[i]];
+    }
+}
+
+inline static void ggml_vec_gelu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i])));
+    }
+}
+
+#ifdef GGML_GELU_QUICK_FP16
+inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
+    uint16_t t;
+    for (int i = 0; i < n; ++i) {
+        ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]);
+        memcpy(&t, &fp16, sizeof(uint16_t));
+        y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]);
+    }
+}
+#else
+inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]);
+    }
+}
+#endif
+
+// Sigmoid Linear Unit (SiLU) function
+inline static float ggml_silu_f32(float x) {
+    return x/(1.0f + expf(-x));
+}
+inline static ggml_fp16_t ggml_silu_f16(ggml_fp16_t x) {
+    float v = GGML_CPU_FP16_TO_FP32(x);
+    return GGML_CPU_FP32_TO_FP16(v/(1.0f + expf(-v)));
+}
+
+#if __FINITE_MATH_ONLY__
+#error "some routines in ggml.c require non-finite math arithmetics -- pass -fno-finite-math-only to the compiler to fix"
+#error "ref: https://github.com/ggml-org/llama.cpp/pull/7154#issuecomment-2143844461"
+#endif
+
+/* Below function was borrowed from the GitHub repository:
+https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.hpp */
+#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    inline static svfloat32_t exp_ps_sve(svbool_t pg, svfloat32_t src) {
+        // Constants
+        const svfloat32_t log2_e = svdup_n_f32(1.4426950409f);
+        const svfloat32_t ln2 = svdup_n_f32(0.6931473921f);
+        const svfloat32_t half_ln2_sq = svdup_n_f32(0.2413862043f);
+        const svuint32_t not_mask17 = svdup_n_u32(~((1u << 17) - 1));
+        const svfloat32_t one = svdup_n_f32(1.0f);
+        const svfloat32_t inactive1 = svdup_n_f32(0.0f);
+        const svint32_t inactive2 = svdup_n_s32(0);
+
+        // Algorithm starts here
+        svfloat32_t t0 = svmul_f32_m(pg, src, log2_e);  // y = x * log2(e)
+        svfloat32_t t1 = svrintm_f32_m(inactive1, pg, t0);         // rount to int (float)
+        svint32_t t2 = svcvt_s32_f32_m(inactive2, pg, t1);         // n
+
+        t1 = svsub_f32_m(pg, t0, t1);   // a = y - floor(y)
+        t1 = svadd_f32_m(pg, t1, one);  // b = a + 1
+
+        svuint32_t t3 = svlsr_n_u32_m(pg, svreinterpret_u32_f32(t1), 17);  // v = b >> 17 (u32)
+        svfloat32_t t4 = svexpa_f32(t3);                                   // c = fexpa(v)
+        t4 = svscale_f32_m(pg, t4, t2);                                    // fexpa(v) * 2^(n)
+
+        // and_(t2.d, t1.d, not_mask17.d)
+        svfloat32_t t5 = svreinterpret_f32_u32(svand_u32_m(pg, svreinterpret_u32_f32(t1), not_mask17));
+        t5 = svsub_f32_m(pg, t1, t5);                // z
+        t0 = svmla_f32_m(pg, ln2, t5, half_ln2_sq);  // ln2 + half_ln2_sq * z
+        t0 = svmla_f32_m(pg, one, t5, t0);           // 1 + (ln2 * z) + (half_ln2_sq * z * z)
+        t0 = svmul_f32_m(pg, t0, t4);                // Final result
+
+        return t0;
+    }
+#endif
+
+#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+
+inline static svfloat32_t ggml_v_expf(svbool_t pg, svfloat32_t x) {
+    const svfloat32_t r = svdup_n_f32_x(pg, 0x1.8p23f);
+    const svfloat32_t z = svmla_n_f32_x(pg, r, x, 0x1.715476p+0f);
+    const svfloat32_t n = svsub_f32_x(pg, z, r);
+    const svfloat32_t b = svmls_n_f32_x(pg, svmls_n_f32_x(pg, x, n, 0x1.62e4p-1f), n, 0x1.7f7d1cp-20f);
+    const svuint32_t e = svlsl_n_u32_x(pg, svreinterpret_u32_f32(z), 23);
+    const svfloat32_t k = svreinterpret_f32_u32(svadd_u32_x(pg, e, svreinterpret_u32_f32(svdup_n_f32_x(pg, 1))));
+    const svbool_t c = svacgt_n_f32(pg, n, 126);
+    const svfloat32_t u = svmul_f32_x(pg, b, b);
+    const svfloat32_t j = svmla_f32_x(pg,
+        svmul_n_f32_x(pg, b, 0x1.ffffecp-1f),
+        svmla_f32_x(pg, svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.fffdb6p-2f), svdup_n_f32_x(pg, 0x1.555e66p-3f), b),
+                        svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.573e2ep-5f), svdup_n_f32_x(pg, 0x1.0e4020p-7f), b), u), u);
+    const svuint32_t d = svdup_n_u32_z(svcmple_n_f32(pg, n, 0.0), 0x82000000);
+    const svfloat32_t s1 = svreinterpret_f32_u32(svadd_n_u32_x(pg, d, 0x7f000000));
+    const svfloat32_t s2 = svreinterpret_f32_u32(svsub_u32_x(pg, e, d));
+    return svsel_f32(svacgt_f32(pg, n, svdup_n_f32_x(pg, 192)), svmul_f32_x(pg, s1, s1),
+                     svsel_f32(c, svmul_f32_x(pg, svmla_f32_x(pg, s2, s2, j), s1), svmla_f32_x(pg, k, k, j)));
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) {
+    const svfloat32_t one = svdup_n_f32_x(pg, 1.0f);
+    const svfloat32_t zero = svdup_n_f32_x(pg, 0.0f);
+    const svfloat32_t neg_x = svsub_f32_x(pg, zero, x);
+    const svfloat32_t exp_neg_x = ggml_v_expf(pg, neg_x);
+    const svfloat32_t one_plus_exp_neg_x = svadd_f32_x(pg, one, exp_neg_x);
+    return svdiv_f32_x(pg, x, one_plus_exp_neg_x);
+}
+
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static float32x4_t ggml_v_expf(float32x4_t x) {
+    const float32x4_t r = vdupq_n_f32(0x1.8p23f);
+    const float32x4_t z = vfmaq_f32(r, x, vdupq_n_f32(0x1.715476p+0f));
+    const float32x4_t n = vsubq_f32(z, r);
+    const float32x4_t b = vfmsq_f32(vfmsq_f32(x, n, vdupq_n_f32(0x1.62e4p-1f)), n,
+                                    vdupq_n_f32(0x1.7f7d1cp-20f));
+    const uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_f32(z), 23);
+    const float32x4_t k = vreinterpretq_f32_u32(vaddq_u32(e, vreinterpretq_u32_f32(vdupq_n_f32(1))));
+    const uint32x4_t c = vcagtq_f32(n, vdupq_n_f32(126));
+    const float32x4_t u = vmulq_f32(b, b);
+    const float32x4_t j = vfmaq_f32(
+        vmulq_f32(vdupq_n_f32(0x1.ffffecp-1f), b),
+        vfmaq_f32(vfmaq_f32(vdupq_n_f32(0x1.fffdb6p-2f), vdupq_n_f32(0x1.555e66p-3f), b),
+                  vfmaq_f32(vdupq_n_f32(0x1.573e2ep-5f), vdupq_n_f32(0x1.0e4020p-7f), b), u), u);
+    if (!vpaddd_u64(vreinterpretq_u64_u32(c)))
+        return vfmaq_f32(k, j, k);
+    const uint32x4_t d = vandq_u32(vclezq_f32(n), vdupq_n_u32(0x82000000));
+    const float32x4_t s1 = vreinterpretq_f32_u32(vaddq_u32(d, vdupq_n_u32(0x7f000000)));
+    const float32x4_t s2 = vreinterpretq_f32_u32(vsubq_u32(e, d));
+    return vbslq_f32(vcagtq_f32(n, vdupq_n_f32(192)), vmulq_f32(s1, s1),
+                     vbslq_f32(c, vmulq_f32(vfmaq_f32(s2, s2, j), s1), vfmaq_f32(k, k, j)));
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static float32x4_t ggml_v_silu(float32x4_t x) {
+    const float32x4_t one = vdupq_n_f32(1.0f);
+    const float32x4_t zero = vdupq_n_f32(0.0f);
+    const float32x4_t neg_x = vsubq_f32(zero, x);
+    const float32x4_t exp_neg_x = ggml_v_expf(neg_x);
+    const float32x4_t one_plus_exp_neg_x = vaddq_f32(one, exp_neg_x);
+    return vdivq_f32(x, one_plus_exp_neg_x);
+}
+
+#elif defined(__AVX512F__) && defined(__AVX512DQ__)
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static __m512 ggml_v_expf(__m512 x) {
+  const __m512 r = _mm512_set1_ps(0x1.8p23f);
+  const __m512 z = _mm512_fmadd_ps(x, _mm512_set1_ps(0x1.715476p+0f), r);
+  const __m512 n = _mm512_sub_ps(z, r);
+  const __m512 b =
+      _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.7f7d1cp-20f),
+                       _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.62e4p-1f), x));
+  const __mmask16 d =
+      _mm512_cmp_ps_mask(_mm512_abs_ps(n), _mm512_set1_ps(192), _CMP_GT_OQ);
+  const __m512 u = _mm512_mul_ps(b, b);
+  const __m512 j = _mm512_fmadd_ps(
+      _mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_set1_ps(0x1.0e4020p-7f), b,
+                                      _mm512_set1_ps(0x1.573e2ep-5f)),
+                      u,
+                      _mm512_fmadd_ps(_mm512_set1_ps(0x1.555e66p-3f), b,
+                                      _mm512_set1_ps(0x1.fffdb6p-2f))),
+      u,
+      _mm512_fmadd_ps(_mm512_set1_ps(0x1.ffffecp-1f), b, _mm512_set1_ps(1.0F)));
+  const __m512 res = _mm512_scalef_ps(j, n);
+  if (_mm512_kortestz(d, d))
+    return res;
+  const __m512 zero = _mm512_setzero_ps();
+  const __m512 alt = _mm512_mask_blend_ps(
+      _mm512_cmp_ps_mask(n, zero, _CMP_LE_OQ), _mm512_set1_ps(INFINITY), zero);
+  return _mm512_mask_blend_ps(d, res, alt);
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static __m512 ggml_v_silu(__m512 x) {
+    const __m512 one = _mm512_set1_ps(1);
+    const __m512 zero = _mm512_setzero_ps();
+    const __m512 neg_x = _mm512_sub_ps(zero, x);
+    const __m512 exp_neg_x = ggml_v_expf(neg_x);
+    const __m512 one_plus_exp_neg_x = _mm512_add_ps(one, exp_neg_x);
+    return _mm512_div_ps(x, one_plus_exp_neg_x);
+}
+
+#elif defined(__AVX2__) && defined(__FMA__)
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static __m256 ggml_v_expf(__m256 x) {
+  const __m256 r = _mm256_set1_ps(0x1.8p23f);
+  const __m256 z = _mm256_fmadd_ps(x, _mm256_set1_ps(0x1.715476p+0f), r);
+  const __m256 n = _mm256_sub_ps(z, r);
+  const __m256 b = _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.7f7d1cp-20f),
+                                    _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.62e4p-1f), x));
+  const __m256i e = _mm256_slli_epi32(_mm256_castps_si256(z), 23);
+  const __m256 k = _mm256_castsi256_ps(
+      _mm256_add_epi32(e, _mm256_castps_si256(_mm256_set1_ps(1))));
+  const __m256i c = _mm256_castps_si256(
+      _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n),
+                    _mm256_set1_ps(126), _CMP_GT_OQ));
+  const __m256 u = _mm256_mul_ps(b, b);
+  const __m256 j = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_set1_ps(0x1.0e4020p-7f), b,
+                                                                   _mm256_set1_ps(0x1.573e2ep-5f)), u,
+                                                   _mm256_fmadd_ps(_mm256_set1_ps(0x1.555e66p-3f), b,
+                                                                   _mm256_set1_ps(0x1.fffdb6p-2f))),
+                                   u, _mm256_mul_ps(_mm256_set1_ps(0x1.ffffecp-1f), b));
+  if (!_mm256_movemask_ps(_mm256_castsi256_ps(c)))
+    return _mm256_fmadd_ps(j, k, k);
+  const __m256i g = _mm256_and_si256(
+      _mm256_castps_si256(_mm256_cmp_ps(n, _mm256_setzero_ps(), _CMP_LE_OQ)),
+      _mm256_set1_epi32(0x82000000u));
+  const __m256 s1 =
+      _mm256_castsi256_ps(_mm256_add_epi32(g, _mm256_set1_epi32(0x7f000000u)));
+  const __m256 s2 = _mm256_castsi256_ps(_mm256_sub_epi32(e, g));
+  const __m256i d = _mm256_castps_si256(
+      _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n),
+                    _mm256_set1_ps(192), _CMP_GT_OQ));
+  return _mm256_or_ps(
+      _mm256_and_ps(_mm256_castsi256_ps(d), _mm256_mul_ps(s1, s1)),
+      _mm256_andnot_ps(
+          _mm256_castsi256_ps(d),
+          _mm256_or_ps(
+              _mm256_and_ps(_mm256_castsi256_ps(c),
+                            _mm256_mul_ps(_mm256_fmadd_ps(s2, j, s2), s1)),
+              _mm256_andnot_ps(_mm256_castsi256_ps(c), _mm256_fmadd_ps(k, j, k)))));
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static __m256 ggml_v_silu(__m256 x) {
+    const __m256 one = _mm256_set1_ps(1);
+    const __m256 zero = _mm256_setzero_ps();
+    const __m256 neg_x = _mm256_sub_ps(zero, x);
+    const __m256 exp_neg_x = ggml_v_expf(neg_x);
+    const __m256 one_plus_exp_neg_x = _mm256_add_ps(one, exp_neg_x);
+    return _mm256_div_ps(x, one_plus_exp_neg_x);
+}
+
+#elif defined(__SSE2__) // __AVX2__ / __ARM_NEON
+
+#if defined(__FMA__)
+#define MADD128(x, y, z) _mm_fmadd_ps(x, y, z)
+#define NMADD128(x, y, z) _mm_fnmadd_ps(x, y, z)
+#else
+#define MADD128(x, y, z) _mm_add_ps(_mm_mul_ps(x, y), z)
+#define NMADD128(x, y, z) _mm_sub_ps(z, _mm_mul_ps(x, y))
+#endif
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static __m128 ggml_v_expf(__m128 x) {
+    const __m128 r = _mm_set1_ps(0x1.8p23f);
+    const __m128 z = MADD128(x, _mm_set1_ps(0x1.715476p+0f), r);
+    const __m128 n = _mm_sub_ps(z, r);
+    const __m128 b =
+        NMADD128(n, _mm_set1_ps(0x1.7f7d1cp-20f), NMADD128(n, _mm_set1_ps(0x1.62e4p-1f), x));
+    const __m128i e = _mm_slli_epi32(_mm_castps_si128(z), 23);
+    const __m128 k = _mm_castsi128_ps(_mm_add_epi32(e, _mm_castps_si128(_mm_set1_ps(1))));
+    const __m128i c =
+        _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(126)));
+    const __m128 u = _mm_mul_ps(b, b);
+    const __m128 j =
+        MADD128(MADD128(MADD128(_mm_set1_ps(0x1.0e4020p-7f), b, _mm_set1_ps(0x1.573e2ep-5f)), u,
+                        MADD128(_mm_set1_ps(0x1.555e66p-3f), b, _mm_set1_ps(0x1.fffdb6p-2f))),
+                u, _mm_mul_ps(_mm_set1_ps(0x1.ffffecp-1f), b));
+    if (!_mm_movemask_epi8(c))
+        return MADD128(j, k, k);
+    const __m128i g = _mm_and_si128(_mm_castps_si128(_mm_cmple_ps(n, _mm_setzero_ps())),
+                                    _mm_set1_epi32(0x82000000u));
+    const __m128 s1 = _mm_castsi128_ps(_mm_add_epi32(g, _mm_set1_epi32(0x7f000000u)));
+    const __m128 s2 = _mm_castsi128_ps(_mm_sub_epi32(e, g));
+    const __m128i d =
+        _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(192)));
+    return _mm_or_ps(
+        _mm_and_ps(_mm_castsi128_ps(d), _mm_mul_ps(s1, s1)),
+        _mm_andnot_ps(_mm_castsi128_ps(d),
+                      _mm_or_ps(_mm_and_ps(_mm_castsi128_ps(c), _mm_mul_ps(MADD128(s2, j, s2), s1)),
+                                _mm_andnot_ps(_mm_castsi128_ps(c), MADD128(k, j, k)))));
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static __m128 ggml_v_silu(__m128 x) {
+    const __m128 one = _mm_set1_ps(1);
+    const __m128 zero = _mm_setzero_ps();
+    const __m128 neg_x = _mm_sub_ps(zero, x);
+    const __m128 exp_neg_x = ggml_v_expf(neg_x);
+    const __m128 one_plus_exp_neg_x = _mm_add_ps(one, exp_neg_x);
+    return _mm_div_ps(x, one_plus_exp_neg_x);
+}
+
+#elif defined(__riscv_v_intrinsic)
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static vfloat32m2_t ggml_v_expf_m2(vfloat32m2_t x, int vl) {
+    const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl);
+#ifdef __riscv_xtheadvector
+    // workaround for compiler bug (gcc 14.3.0: Error: unrecognized opcode `th.vmv1r.v v2,v4')
+    vfloat32m2_t z = __riscv_vfadd_vf_f32m2(r, 0.0f, vl);
+    z = __riscv_vfmacc_vf_f32m2(z, 0x1.715476p+0f, x, vl);
+#else
+    const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl);
+#endif
+    const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl);
+    const vfloat32m2_t b = __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl),
+                                                    0x1.7f7d1cp-20f, n, vl);
+    const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl);
+    const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); // 1.0f
+    const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl);
+    const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl);
+    const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2(
+        __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl),
+        __riscv_vfmacc_vv_f32m2(
+            __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl),
+            __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl),
+            u, vl), u, vl);
+    if (!__riscv_vcpop_m_b16(c, vl))
+        return __riscv_vfmacc_vv_f32m2(k, j, k, vl);
+    const vbool16_t  dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl);
+    const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl);
+    const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl));
+    const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl));
+    const vfloat32m2_t r1 = __riscv_vmerge_vvm_f32m2(
+        __riscv_vfmacc_vv_f32m2(k, k, j, vl),
+        __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl),
+        c, vl);
+    return __riscv_vmerge_vvm_f32m2(
+        r1, __riscv_vfmul_vv_f32m2(s1, s1, vl),
+        __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl),
+        vl);
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) {
+    const vfloat32m2_t neg_x = __riscv_vfneg_v_f32m2(x, vl);
+    const vfloat32m2_t exp_neg_x = ggml_v_expf_m2(neg_x, vl);
+    const vfloat32m2_t one_plus_exp_neg_x = __riscv_vfadd_vf_f32m2(exp_neg_x, 1.0f, vl);
+    return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl);
+}
+
+#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic
+
+inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_silu_f16(x[i]);
+    }
+}
+
+inline static void ggml_vec_silu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_FP32_TO_BF16(ggml_silu_f32(GGML_BF16_TO_FP32(x[i])));
+    }
+}
+
+inline static float ggml_silu_backward_f32(float x, float dy) {
+    const float s = 1.0f/(1.0f + expf(-x));
+    return dy*s*(1.0f + x*(1.0f - s));
+}
+
+inline static ggml_fp16_t ggml_silu_backward_f16(ggml_fp16_t x, ggml_fp16_t dy) {
+    const float v = GGML_CPU_FP16_TO_FP32(x);
+    const float s = 1.0f/(1.0f + expf(-v));
+    return GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(dy)*s*(1.0f + v*(1.0f - s)));
+}
+
+inline static void ggml_vec_silu_backward_f32(const int n, float * dx, const float * x, const float * dy) {
+    for (int i = 0; i < n; ++i) {
+        dx[i] = ggml_silu_backward_f32(x[i], dy[i]);
+    }
+}
+
+inline static void ggml_vec_silu_backward_f16(const int n, ggml_fp16_t * dx, const ggml_fp16_t * x, const ggml_fp16_t * dy) {
+    for (int i = 0; i < n; ++i) {
+        dx[i] = ggml_silu_backward_f16(x[i], dy[i]);
+    }
+}
+
+inline static void ggml_vec_reglu_f32 (const int n, float * y, const float * x, const float * g) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = (x[i] > 0.f) ? x[i] * g[i] : 0.f;
+    }
+}
+
+inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(x[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v * GGML_CPU_FP16_TO_FP32(g[i]) : 0.f);
+    }
+}
+
+inline static void ggml_vec_reglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_BF16_TO_FP32(x[i]);
+        y[i] = GGML_FP32_TO_BF16((v > 0.f) ? v * GGML_BF16_TO_FP32(g[i]) : 0.f);
+    }
+}
+
+#ifdef GGML_GELU_FP16
+inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) {
+    uint16_t t;
+    for (int i = 0; i < n; ++i) {
+        if (x[i] <= -10.0f) {
+            y[i] = 0.0f;
+        } else if (x[i] >= 10.0f) {
+            y[i] = x[i] * g[i];
+        } else {
+            ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]);
+            memcpy(&t, &fp16, sizeof(uint16_t));
+            y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]) * g[i];
+        }
+    }
+}
+#else
+inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_gelu_f32(x[i]) * g[i];
+    }
+}
+#endif
+
+inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
+    const uint16_t * i16 = (const uint16_t *) x;
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(g[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[i16[i]]) * v);
+    }
+}
+
+inline static void ggml_vec_geglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i]));
+    }
+}
+
+void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g);
+
+inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_CPU_FP16_TO_FP32(x[i]);
+        float gi = GGML_CPU_FP16_TO_FP32(g[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16((xi/(1.0f + expf(-xi))) * gi);
+    }
+}
+
+inline static void ggml_vec_swiglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_BF16_TO_FP32(x[i]);
+        float gi = GGML_BF16_TO_FP32(g[i]);
+        y[i] = GGML_FP32_TO_BF16((xi/(1.0f + expf(-xi))) * gi);
+    }
+}
+
+inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) {
+    for (int i = 0; i < n; ++i) {
+        float xi = x[i];
+        y[i] = 0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * g[i];
+    }
+}
+
+inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_CPU_FP16_TO_FP32(x[i]);
+        float gi = GGML_CPU_FP16_TO_FP32(g[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi);
+    }
+}
+
+inline static void ggml_vec_geglu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        float xi = GGML_BF16_TO_FP32(x[i]);
+        float gi = GGML_BF16_TO_FP32(g[i]);
+        y[i] = GGML_FP32_TO_BF16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi);
+    }
+}
+
+#ifdef GGML_GELU_QUICK_FP16
+inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
+    uint16_t t;
+    for (int i = 0; i < n; ++i) {
+        ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]);
+        memcpy(&t, &fp16, sizeof(uint16_t));
+        y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]) * g[i];
+    }
+}
+#else
+inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
+    }
+}
+#endif
+
+inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
+    const uint16_t * i16 = (const uint16_t *) x;
+    for (int i = 0; i < n; ++i) {
+        float v = GGML_CPU_FP16_TO_FP32(g[i]);
+        y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[i16[i]]) * v);
+    }
+}
+
+inline static void ggml_vec_geglu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
+    for (int i = 0; i < n; ++i) {
+        y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i]));
+    }
+}
+
+inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) {
+#ifndef GGML_USE_ACCELERATE
+    ggml_float sum = 0.0;
+    for (int i = 0; i < n; ++i) {
+        sum += (ggml_float)x[i];
+    }
+    *s = (float)sum;
+#else
+    vDSP_sve(x, 1, s, n);
+#endif
+}
+
+inline static void ggml_vec_cumsum_f32(const int n, float * y, const float * x) {
+    for (int i = 0; i < n; ++i) {
+        if (i == 0) {
+            y[i] = x[i];
+        } else {
+            y[i] = y[i - 1] + x[i];
+        }
+    }
+}
+
+inline static void ggml_vec_sum_f32_ggf(const int n, ggml_float * s, const float * x) {
+    ggml_float sum = 0.0;
+    for (int i = 0; i < n; ++i) {
+        sum += (ggml_float)x[i];
+    }
+    *s = sum;
+}
+
+inline static void ggml_vec_sum_f16_ggf(const int n, float * s, const ggml_fp16_t * x) {
+    float sum = 0.0f;
+    for (int i = 0; i < n; ++i) {
+        sum += GGML_CPU_FP16_TO_FP32(x[i]);
+    }
+    *s = sum;
+}
+
+inline static void ggml_vec_sum_bf16_ggf(const int n, float * s, const ggml_bf16_t * x) {
+    float sum = 0.0f;
+    for (int i = 0; i < n; ++i) {
+        sum += GGML_BF16_TO_FP32(x[i]);
+    }
+    *s = sum;
+}
+
+inline static void ggml_vec_max_f32(const int n, float * s, const float * x) {
+#ifndef GGML_USE_ACCELERATE
+    float max = -INFINITY;
+    for (int i = 0; i < n; ++i) {
+        max = MAX(max, x[i]);
+    }
+    *s = max;
+#else
+    vDSP_maxv(x, 1, s, n);
+#endif
+}
+
+inline static void ggml_vec_norm_inv_f32(const int n, float * s, const float * x) {
+    ggml_vec_norm_f32(n, s, x);
+    *s = 1.f/(*s);
+}
+
+inline static void ggml_vec_argmax_f32(const int n, int * s, const float * x) {
+    float max = -INFINITY;
+    int idx = 0;
+    for (int i = 0; i < n; ++i) {
+        max = MAX(max, x[i]);
+        if (max == x[i]) { idx = i; }
+    }
+    *s = idx;
+}
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h b/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h
new file mode 100644
index 0000000000..dc6453c6ea
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h
@@ -0,0 +1,152 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    // the compute plan that needs to be prepared for ggml_graph_compute()
+    // since https://github.com/ggml-org/ggml/issues/287
+    struct ggml_cplan {
+        size_t    work_size; // size of work buffer, calculated by `ggml_graph_plan()`
+        uint8_t * work_data; // work buffer, to be allocated by caller before calling to `ggml_graph_compute()`
+
+        int n_threads;
+        struct ggml_threadpool * threadpool;
+
+        // abort ggml_graph_compute when true
+        ggml_abort_callback abort_callback;
+        void *              abort_callback_data;
+
+        // use only reference implementations
+        bool use_ref;
+    };
+
+    // numa strategies
+    enum ggml_numa_strategy {
+        GGML_NUMA_STRATEGY_DISABLED   = 0,
+        GGML_NUMA_STRATEGY_DISTRIBUTE = 1,
+        GGML_NUMA_STRATEGY_ISOLATE    = 2,
+        GGML_NUMA_STRATEGY_NUMACTL    = 3,
+        GGML_NUMA_STRATEGY_MIRROR     = 4,
+        GGML_NUMA_STRATEGY_COUNT
+    };
+
+    GGML_BACKEND_API void    ggml_numa_init(enum ggml_numa_strategy numa); // call once for better performance on NUMA systems
+    GGML_BACKEND_API bool    ggml_is_numa(void); // true if init detected that system has >1 NUMA node
+
+    GGML_BACKEND_API struct ggml_tensor * ggml_new_i32(struct ggml_context * ctx, int32_t value);
+    GGML_BACKEND_API struct ggml_tensor * ggml_new_f32(struct ggml_context * ctx, float value);
+
+    GGML_BACKEND_API struct ggml_tensor * ggml_set_i32 (struct ggml_tensor * tensor, int32_t value);
+    GGML_BACKEND_API struct ggml_tensor * ggml_set_f32 (struct ggml_tensor * tensor, float value);
+
+    GGML_BACKEND_API int32_t ggml_get_i32_1d(const struct ggml_tensor * tensor, int i);
+    GGML_BACKEND_API void    ggml_set_i32_1d(const struct ggml_tensor * tensor, int i, int32_t value);
+
+    GGML_BACKEND_API int32_t ggml_get_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3);
+    GGML_BACKEND_API void    ggml_set_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, int32_t value);
+
+    GGML_BACKEND_API float   ggml_get_f32_1d(const struct ggml_tensor * tensor, int i);
+    GGML_BACKEND_API void    ggml_set_f32_1d(const struct ggml_tensor * tensor, int i, float value);
+
+    GGML_BACKEND_API float   ggml_get_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3);
+    GGML_BACKEND_API void    ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, float value);
+
+    GGML_BACKEND_API struct ggml_threadpool *      ggml_threadpool_new           (struct ggml_threadpool_params  * params);
+    GGML_BACKEND_API void                          ggml_threadpool_free          (struct ggml_threadpool * threadpool);
+    GGML_BACKEND_API int                           ggml_threadpool_get_n_threads (struct ggml_threadpool * threadpool);
+    GGML_BACKEND_API void                          ggml_threadpool_pause         (struct ggml_threadpool * threadpool);
+    GGML_BACKEND_API void                          ggml_threadpool_resume        (struct ggml_threadpool * threadpool);
+
+    // ggml_graph_plan() has to be called before ggml_graph_compute()
+    // when plan.work_size > 0, caller must allocate memory for plan.work_data
+    GGML_BACKEND_API struct ggml_cplan ggml_graph_plan(
+                  const struct ggml_cgraph * cgraph,
+                                       int   n_threads, /* = GGML_DEFAULT_N_THREADS */
+                    struct ggml_threadpool * threadpool /* = NULL */ );
+    GGML_BACKEND_API enum ggml_status  ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan);
+
+    // same as ggml_graph_compute() but the work data is allocated as a part of the context
+    // note: the drawback of this API is that you must have ensured that the context has enough memory for the work data
+    GGML_BACKEND_API enum ggml_status  ggml_graph_compute_with_ctx(struct ggml_context * ctx, struct ggml_cgraph * cgraph, int n_threads);
+
+    //
+    // system info
+    //
+
+    // x86
+    GGML_BACKEND_API int ggml_cpu_has_sse3       (void);
+    GGML_BACKEND_API int ggml_cpu_has_ssse3      (void);
+    GGML_BACKEND_API int ggml_cpu_has_avx        (void);
+    GGML_BACKEND_API int ggml_cpu_has_avx_vnni   (void);
+    GGML_BACKEND_API int ggml_cpu_has_avx2       (void);
+    GGML_BACKEND_API int ggml_cpu_has_bmi2       (void);
+    GGML_BACKEND_API int ggml_cpu_has_f16c       (void);
+    GGML_BACKEND_API int ggml_cpu_has_fma        (void);
+    GGML_BACKEND_API int ggml_cpu_has_avx512     (void);
+    GGML_BACKEND_API int ggml_cpu_has_avx512_vbmi(void);
+    GGML_BACKEND_API int ggml_cpu_has_avx512_vnni(void);
+    GGML_BACKEND_API int ggml_cpu_has_avx512_bf16(void);
+    GGML_BACKEND_API int ggml_cpu_has_amx_int8   (void);
+    // ARM
+    GGML_BACKEND_API int ggml_cpu_has_neon       (void);
+    GGML_BACKEND_API int ggml_cpu_has_arm_fma    (void);
+    GGML_BACKEND_API int ggml_cpu_has_fp16_va    (void);
+    GGML_BACKEND_API int ggml_cpu_has_dotprod    (void);
+    GGML_BACKEND_API int ggml_cpu_has_matmul_int8(void);
+    GGML_BACKEND_API int ggml_cpu_has_sve        (void);
+    GGML_BACKEND_API int ggml_cpu_get_sve_cnt    (void);  // sve vector length in bytes
+    GGML_BACKEND_API int ggml_cpu_has_sme        (void);
+    GGML_BACKEND_API int ggml_cpu_has_sme2       (void);
+    // other
+    GGML_BACKEND_API int ggml_cpu_has_riscv_v    (void);
+    GGML_BACKEND_API int ggml_cpu_get_rvv_vlen   (void);  // risc-v vector length in bytes
+    GGML_BACKEND_API int ggml_cpu_has_vsx        (void);
+    GGML_BACKEND_API int ggml_cpu_has_vxe        (void);
+    GGML_BACKEND_API int ggml_cpu_has_wasm_simd  (void);
+    GGML_BACKEND_API int ggml_cpu_has_llamafile  (void);
+
+    // Internal types and functions exposed for tests and benchmarks
+
+    typedef void (*ggml_vec_dot_t)  (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT x, size_t bx,
+                                       const void * GGML_RESTRICT y, size_t by, int nrc);
+
+    struct ggml_type_traits_cpu {
+        ggml_from_float_t        from_float;
+        ggml_vec_dot_t           vec_dot;
+        enum ggml_type           vec_dot_type;
+        int64_t                  nrows; // number of rows to process simultaneously
+    };
+
+    GGML_BACKEND_API const struct ggml_type_traits_cpu * ggml_get_type_traits_cpu(enum ggml_type type);
+
+    GGML_BACKEND_API void ggml_cpu_init(void);
+
+    //
+    // CPU backend
+    //
+
+    GGML_BACKEND_API ggml_backend_t ggml_backend_cpu_init(void);
+
+    GGML_BACKEND_API bool ggml_backend_is_cpu                (ggml_backend_t backend);
+    GGML_BACKEND_API void ggml_backend_cpu_set_n_threads     (ggml_backend_t backend_cpu, int n_threads);
+    GGML_BACKEND_API void ggml_backend_cpu_set_threadpool    (ggml_backend_t backend_cpu, ggml_threadpool_t threadpool);
+    GGML_BACKEND_API void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data);
+
+    GGML_BACKEND_API void ggml_backend_cpu_set_use_ref(ggml_backend_t backend_cpu, bool use_ref);
+
+    GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cpu_reg(void);
+
+    GGML_BACKEND_API void ggml_cpu_fp32_to_fp32(const float *,       float *, int64_t);
+    GGML_BACKEND_API void ggml_cpu_fp32_to_i32 (const float *,     int32_t *, int64_t);
+    GGML_BACKEND_API void ggml_cpu_fp32_to_fp16(const float *, ggml_fp16_t *, int64_t);
+    GGML_BACKEND_API void ggml_cpu_fp16_to_fp32(const ggml_fp16_t *, float *, int64_t);
+    GGML_BACKEND_API void ggml_cpu_fp32_to_bf16(const float *, ggml_bf16_t *, int64_t);
+    GGML_BACKEND_API void ggml_cpu_bf16_to_fp32(const ggml_bf16_t *, float *, int64_t);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h b/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h
new file mode 100644
index 0000000000..dc7367dec4
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h
@@ -0,0 +1,292 @@
+#pragma once
+
+// ggml-backend internal header
+
+#include "ggml-backend.h"
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    #define GGML_BACKEND_API_VERSION 3
+
+    //
+    // Backend buffer type
+    //
+
+    struct ggml_backend_buffer_type_i {
+        const char *          (*get_name)        (ggml_backend_buffer_type_t buft);
+        // allocate a buffer of this type
+        ggml_backend_buffer_t (*alloc_buffer)    (ggml_backend_buffer_type_t buft, size_t size);
+        // (optional) allocate tensors from a list into a buffer of this type (defaults to alloc_buffer + linear allocator)
+        ggml_backend_buffer_t (*alloc_buffer_n)  (ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors);
+        // tensor alignment
+        size_t                (*get_alignment)   (ggml_backend_buffer_type_t buft);
+        // (optional) max buffer size that can be allocated (defaults to SIZE_MAX)
+        size_t                (*get_max_size)    (ggml_backend_buffer_type_t buft);
+        // (optional) data size needed to allocate the tensor, including padding (defaults to ggml_nbytes)
+        size_t                (*get_alloc_size)  (ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor);
+        // (optional) total data size needed to allocate the given tensors, including padding and splitting (defaults to per-tensor get_alloc_size)
+        size_t                (*get_alloc_size_n)(ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors);
+        // (optional) check if tensor data is in host memory and uses standard ggml tensor layout (defaults to false)
+        bool                  (*is_host)         (ggml_backend_buffer_type_t buft);
+    };
+
+    struct ggml_backend_buffer_type {
+        struct ggml_backend_buffer_type_i  iface;
+        ggml_backend_dev_t device;
+        void * context;
+    };
+
+    // [TAG_ALLOC_SIZE_EXPAND]
+    // returns true for ops that may require additional memory for fleeting data on some backends,
+    // i.e. the backend buffer type's get_alloc_size may return more than ggml_nbytes for the output tensor
+    GGML_API bool ggml_op_alloc_size_may_expand(enum ggml_op op);
+
+    //
+    // Backend buffer
+    //
+
+    struct ggml_backend_buffer_i {
+        // (optional) free the buffer
+        void         (*free_buffer)  (ggml_backend_buffer_t buffer);
+        // base address of the buffer
+        void *       (*get_base)     (ggml_backend_buffer_t buffer);
+        // (optional) initialize a tensor in the buffer (eg. add tensor extras)
+        enum ggml_status (*init_tensor)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor);
+        // tensor data access
+        void         (*memset_tensor)(ggml_backend_buffer_t buffer,       struct ggml_tensor * tensor,     uint8_t value, size_t offset, size_t size);
+        void         (*set_tensor)   (ggml_backend_buffer_t buffer,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+        void         (*get_tensor)   (ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+        // (optional) 2d data copies
+        void         (*set_tensor_2d)(ggml_backend_buffer_t buffer,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+        void         (*get_tensor_2d)(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+
+        // (optional) tensor copy: dst is in the buffer, src may be in any buffer, including buffers from a different backend (return false if not supported)
+        bool         (*cpy_tensor)   (ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst);
+        // clear the entire buffer
+        void         (*clear)        (ggml_backend_buffer_t buffer, uint8_t value);
+        // (optional) reset any internal state due to tensor initialization, such as tensor extras
+        void         (*reset)        (ggml_backend_buffer_t buffer);
+    };
+
+    struct ggml_backend_buffer {
+        struct ggml_backend_buffer_i  iface;
+        ggml_backend_buffer_type_t    buft;
+        void * context;
+        size_t size;
+        enum ggml_backend_buffer_usage usage;
+    };
+
+    GGML_API ggml_backend_buffer_t ggml_backend_buffer_init(
+                   ggml_backend_buffer_type_t buft,
+            struct ggml_backend_buffer_i      iface,
+                   void *                     context,
+                   size_t                     size);
+
+    // do not use directly, use ggml_backend_tensor_copy instead
+    GGML_API bool ggml_backend_buffer_copy_tensor(const struct ggml_tensor * src, struct ggml_tensor * dst);
+
+    // multi-buffer
+    // buffer that contains a collection of buffers
+    GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers);
+    GGML_API bool                  ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer);
+    GGML_API void                  ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+    GGML_API void                  ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+
+    //
+    // Backend (meta)
+    //
+
+    GGML_API bool ggml_backend_is_meta       (ggml_backend_t backend);
+    GGML_API bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf);
+    GGML_API bool ggml_backend_buft_is_meta  (ggml_backend_buffer_type_t buft);
+
+    GGML_API size_t         ggml_backend_meta_n_backends    (ggml_backend_t meta_backend);
+    GGML_API ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index);
+
+    //
+    // Backend (stream)
+    //
+
+    // passed to graph_optimize so the backend can add allocation dependencies:
+    // if the backend executes parts of the graph out of order (e.g. on concurrent streams),
+    // it must keep the affected tensors allocated until a node where execution is known to have joined
+    struct ggml_backend_graph_optimize_params {
+        // keep `tensor` allocated at least until `until` (a node of the same graph) has been computed
+        // can be called multiple times for the same tensor: the longest lifetime applies
+        void (*add_alloc_dep)(void * user_data, struct ggml_tensor * tensor, struct ggml_tensor * until);
+        void * user_data;
+    };
+
+    struct ggml_backend_i {
+        const char * (*get_name)(ggml_backend_t backend);
+
+        void (*free)(ggml_backend_t backend);
+
+        // (optional) asynchronous tensor data access
+        void (*set_tensor_async)   (ggml_backend_t backend,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+        void (*get_tensor_async)   (ggml_backend_t backend, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+        void (*set_tensor_2d_async)(ggml_backend_t backend,       struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+        void (*get_tensor_2d_async)(ggml_backend_t backend, const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+        bool (*cpy_tensor_async)(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);
+
+        // (optional) complete all pending operations (required if the backend supports async operations)
+        void (*synchronize)(ggml_backend_t backend);
+
+        // (optional) graph plans (not used currently)
+        // compute graph with a plan
+        ggml_backend_graph_plan_t (*graph_plan_create) (ggml_backend_t backend, const struct ggml_cgraph * cgraph);
+        void                      (*graph_plan_free)   (ggml_backend_t backend, ggml_backend_graph_plan_t plan);
+        // update the plan with a new graph - this should be faster than creating a new plan when the graph has the same topology
+        void                      (*graph_plan_update) (ggml_backend_t backend, ggml_backend_graph_plan_t plan, const struct ggml_cgraph * cgraph);
+        // compute the graph with the plan
+        enum ggml_status          (*graph_plan_compute)(ggml_backend_t backend, ggml_backend_graph_plan_t plan);
+
+        // compute graph (always async if supported by the backend)
+        enum ggml_status          (*graph_compute)     (ggml_backend_t backend, struct ggml_cgraph * cgraph);
+
+        // (optional) event synchronization
+        // record an event on this stream
+        void (*event_record)(ggml_backend_t backend, ggml_backend_event_t event);
+        // wait for an event on on a different stream
+        void (*event_wait)  (ggml_backend_t backend, ggml_backend_event_t event);
+
+        // (optional) sort/optimize the nodes in the graph
+        void                      (*graph_optimize)    (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params);
+    };
+
+    struct ggml_backend {
+        ggml_guid_t guid;
+        struct ggml_backend_i iface;
+        ggml_backend_dev_t device;
+        void * context;
+    };
+
+    struct ggml_backend_event {
+        struct ggml_backend_device * device;
+        void * context;
+    };
+
+    //
+    // Backend device
+    //
+
+    // Note: if additional properties are needed, we should add a struct with all of them
+    //       the current functions to obtain the properties can remain, since they are more convenient for often used properties
+    struct ggml_backend_device_i {
+        // device name: short identifier for this device, such as "CPU" or "CUDA0"
+        const char * (*get_name)(ggml_backend_dev_t dev);
+
+        // device description: short informative description of the device, could be the model name
+        const char * (*get_description)(ggml_backend_dev_t dev);
+
+        // device memory in bytes: 0 bytes to indicate no memory to report
+        void         (*get_memory)(ggml_backend_dev_t dev, size_t * free, size_t * total);
+
+        // device type
+        enum ggml_backend_dev_type (*get_type)(ggml_backend_dev_t dev);
+
+        // device properties
+        void (*get_props)(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props);
+
+        // backend (stream) initialization
+        ggml_backend_t (*init_backend)(ggml_backend_dev_t dev, const char * params);
+
+        // preferred buffer type
+        ggml_backend_buffer_type_t (*get_buffer_type)(ggml_backend_dev_t dev);
+
+        // (optional) host buffer type (in system memory, typically this is a pinned memory buffer for faster transfers between host and device)
+        ggml_backend_buffer_type_t (*get_host_buffer_type)(ggml_backend_dev_t dev);
+
+        // (optional) buffer from pointer: create a buffer from a host pointer (useful for memory mapped models and importing data from other libraries)
+        ggml_backend_buffer_t (*buffer_from_host_ptr)(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size);
+
+        // check if the backend can compute an operation
+        bool (*supports_op)(ggml_backend_dev_t dev, const struct ggml_tensor * op);
+
+        // check if the backend can use tensors allocated in a buffer type
+        bool (*supports_buft)(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft);
+
+        // (optional) check if the backend wants to run an operation, even if the weights are allocated in an incompatible buffer
+        // these should be expensive operations that may benefit from running on this backend instead of the CPU backend
+        bool (*offload_op)(ggml_backend_dev_t dev, const struct ggml_tensor * op);
+
+        // (optional) event synchronization
+        ggml_backend_event_t (*event_new)         (ggml_backend_dev_t dev);
+        void                 (*event_free)        (ggml_backend_dev_t dev, ggml_backend_event_t event);
+        void                 (*event_synchronize) (ggml_backend_dev_t dev, ggml_backend_event_t event);
+    };
+
+    struct ggml_backend_device {
+        struct ggml_backend_device_i iface;
+        ggml_backend_reg_t reg;
+        void * context;
+    };
+
+    //
+    // Backend (reg)
+    //
+
+    struct ggml_backend_reg_i {
+        const char * (*get_name)(ggml_backend_reg_t reg);
+
+        // enumerate available devices
+        size_t             (*get_device_count)(ggml_backend_reg_t reg);
+        ggml_backend_dev_t (*get_device)(ggml_backend_reg_t reg, size_t index);
+
+        // (optional) get a pointer to a function in the backend
+        // backends can add custom functions that are not part of the standard ggml-backend interface
+        void * (*get_proc_address)(ggml_backend_reg_t reg, const char * name);
+    };
+
+    struct ggml_backend_reg {
+        int api_version; // initialize to GGML_BACKEND_API_VERSION
+        struct ggml_backend_reg_i iface;
+        void * context;
+    };
+
+    // Add backend dynamic loading support to the backend
+
+    // Initialize the backend
+    typedef ggml_backend_reg_t (*ggml_backend_init_t)(void);
+    // Optional: obtain a score for the backend based on the system configuration
+    // Higher scores are preferred, 0 means the backend is not supported in the current system
+    typedef int                (*ggml_backend_score_t)(void);
+
+#ifdef GGML_BACKEND_DL
+#    ifdef __cplusplus
+#        define GGML_BACKEND_DL_IMPL(reg_fn)                             \
+            extern "C" {                                                 \
+            GGML_BACKEND_API ggml_backend_reg_t ggml_backend_init(void); \
+            }                                                            \
+            ggml_backend_reg_t ggml_backend_init(void) {                 \
+                return reg_fn();                                         \
+            }
+#        define GGML_BACKEND_DL_SCORE_IMPL(score_fn)       \
+            extern "C" {                                   \
+            GGML_BACKEND_API int ggml_backend_score(void); \
+            }                                              \
+            int ggml_backend_score(void) {                 \
+                return score_fn();                         \
+            }
+#    else
+#        define GGML_BACKEND_DL_IMPL(reg_fn)                              \
+            GGML_BACKEND_API ggml_backend_reg_t ggml_backend_init(void);  \
+            ggml_backend_reg_t                  ggml_backend_init(void) { \
+                return reg_fn();                                          \
+            }
+#        define GGML_BACKEND_DL_SCORE_IMPL(score_fn)        \
+            GGML_BACKEND_API int ggml_backend_score(void);  \
+            int                  ggml_backend_score(void) { \
+                return score_fn();                          \
+            }
+#    endif
+#else
+#    define GGML_BACKEND_DL_IMPL(reg_fn)
+#    define GGML_BACKEND_DL_SCORE_IMPL(score_fn)
+#endif
+
+#ifdef  __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c
new file mode 100644
index 0000000000..7750a72ce6
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c
@@ -0,0 +1,5639 @@
+#define GGML_COMMON_IMPL_C
+#include "ggml-common.h"
+
+#include "ggml-quants.h"
+#include "ggml-impl.h"
+#include "ggml-cpu/ggml-cpu-impl.h"
+#include "ggml-cpu.h"
+
+#include <math.h>
+#include <string.h>
+#include <assert.h>
+#include <float.h>
+#include <stdlib.h> // for qsort
+#include <stdio.h>  // for GGML_ASSERT
+
+#ifdef GGML_USE_OPENMP
+#include <omp.h>
+#endif
+
+#define GROUP_MAX_EPS 1e-15f
+#define GROUP_MAX_EPS_IQ3_XXS 1e-8f
+#define GROUP_MAX_EPS_IQ2_S 1e-8f
+#define GROUP_MAX_EPS_IQ1_M 1e-7f
+#define GROUP_MAX_EPS_IQ1_S 1e-12f
+
+#define UNUSED GGML_UNUSED
+
+static inline int best_index_int8(int n, const int8_t * val, float x) {
+    if (x <= val[0]) return 0;
+    if (x >= val[n-1]) return n-1;
+    int ml = 0, mu = n-1;
+    while (mu-ml > 1) {
+        int mav = (ml+mu)/2;
+        if (x < val[mav]) mu = mav; else ml = mav;
+    }
+    return x - val[mu-1] < val[mu] - x ? mu-1 : mu;
+}
+
+// reference implementation for deterministic creation of model files
+void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK1_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float sum_abs = 0.0f;
+        for (int j = 0; j < qk; j++) {
+            sum_abs += fabsf(x[i*qk + j]);
+        }
+        const float d = sum_abs / qk;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        // Clear all bits first
+        for (int j = 0; j < qk / 8; ++j) {
+            y[i].qs[j] = 0;
+        }
+
+        // Just store sign of each weight directly (no normalization)
+        for (int j = 0; j < qk; ++j) {
+            const int bit_index = j;
+            const int byte_index = bit_index / 8;
+            const int bit_offset = bit_index % 8;
+
+            if (x[i*qk + j] >= 0.0f) {
+                y[i].qs[byte_index] |= (1 << bit_offset);
+            }
+        }
+    }
+}
+
+void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK2_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        // Compute scale as max absolute value in the block
+        float amax = 0.0f;
+        for (int j = 0; j < qk; j++) {
+            const float a = fabsf(x[i*qk + j]);
+            if (a > amax) amax = a;
+        }
+        const float d = amax;
+        const float id = d > 0.0f ? 1.0f / d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        // Clear quant bytes
+        for (int j = 0; j < qk / 4; ++j) {
+            y[i].qs[j] = 0;
+        }
+
+        // Encode 2-bit values: round(w/d) clamped to [-1, 2], then add 1
+        // 00 (-1) = -scale, 01 (0) = 0, 10 (+1) = +scale, 11 (+2) = 2*scale
+        for (int j = 0; j < qk; ++j) {
+            const float w = x[i*qk + j];
+            int q = (int)roundf(w * id) + 1;
+            if (q < 0) q = 0;
+            if (q > 3) q = 3;
+            const int byte_index = j / 4;
+            const int bit_offset = (j % 4) * 2;
+            y[i].qs[byte_index] |= ((uint8_t)q << bit_offset);
+        }
+    }
+}
+
+// reference implementation for deterministic creation of model files
+void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK4_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+        float max  = 0.0f;
+
+        for (int j = 0; j < qk; j++) {
+            const float v = x[i*qk + j];
+            if (amax < fabsf(v)) {
+                amax = fabsf(v);
+                max  = v;
+            }
+        }
+
+        const float d  = max / -8;
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        for (int j = 0; j < qk/2; ++j) {
+            const float x0 = x[i*qk + 0    + j]*id;
+            const float x1 = x[i*qk + qk/2 + j]*id;
+
+            const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f));
+            const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f));
+
+            y[i].qs[j]  = xi0;
+            y[i].qs[j] |= xi1 << 4;
+        }
+    }
+}
+
+void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k) {
+    const int qk = QK4_1;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float min = FLT_MAX;
+        float max = -FLT_MAX;
+
+        for (int j = 0; j < qk; j++) {
+            const float v = x[i*qk + j];
+
+            if (v < min) min = v;
+            if (v > max) max = v;
+        }
+
+        const float d  = (max - min) / ((1 << 4) - 1);
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+        y[i].m = GGML_FP32_TO_FP16(min);
+
+        for (int j = 0; j < qk/2; ++j) {
+            const float x0 = (x[i*qk + 0    + j] - min)*id;
+            const float x1 = (x[i*qk + qk/2 + j] - min)*id;
+
+            const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f));
+            const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f));
+
+            y[i].qs[j]  = xi0;
+            y[i].qs[j] |= xi1 << 4;
+        }
+    }
+}
+
+void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK5_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+        float max  = 0.0f;
+
+        for (int j = 0; j < qk; j++) {
+            const float v = x[i*qk + j];
+            if (amax < fabsf(v)) {
+                amax = fabsf(v);
+                max  = v;
+            }
+        }
+
+        const float d  = max / -16;
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        uint32_t qh = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const float x0 = x[i*qk + 0    + j]*id;
+            const float x1 = x[i*qk + qk/2 + j]*id;
+
+            const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f));
+            const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f));
+
+            y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
+
+            // get the 5-th bit and store it in qh at the right position
+            qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+            qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2);
+        }
+
+        memcpy(&y[i].qh, &qh, sizeof(qh));
+    }
+}
+
+void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k) {
+    const int qk = QK5_1;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float min = FLT_MAX;
+        float max = -FLT_MAX;
+
+        for (int j = 0; j < qk; j++) {
+            const float v = x[i*qk + j];
+
+            if (v < min) min = v;
+            if (v > max) max = v;
+        }
+
+        const float d  = (max - min) / ((1 << 5) - 1);
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+        y[i].m = GGML_FP32_TO_FP16(min);
+
+        uint32_t qh = 0;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const float x0 = (x[i*qk + 0    + j] - min)*id;
+            const float x1 = (x[i*qk + qk/2 + j] - min)*id;
+
+            const uint8_t xi0 = (uint8_t)(x0 + 0.5f);
+            const uint8_t xi1 = (uint8_t)(x1 + 0.5f);
+
+            y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
+
+            // get the 5-th bit and store it in qh at the right position
+            qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+            qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2);
+        }
+
+        memcpy(&y[i].qh, &qh, sizeof(y[i].qh));
+    }
+}
+
+// reference implementation for deterministic creation of model files
+void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK8_0 == 0);
+    const int nb = k / QK8_0;
+
+    for (int i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+
+        for (int j = 0; j < QK8_0; j++) {
+            const float v = x[i*QK8_0 + j];
+            amax = MAX(amax, fabsf(v));
+        }
+
+        const float d = amax / ((1 << 7) - 1);
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        for (int j = 0; j < QK8_0; ++j) {
+            const float x0 = x[i*QK8_0 + j]*id;
+
+            y[i].qs[j] = roundf(x0);
+        }
+    }
+}
+
+// reference implementation for deterministic creation of model files
+void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k) {
+    assert(QK8_1 == 32);
+    assert(k % QK8_1 == 0);
+    const int nb = k / QK8_1;
+
+    for (int i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+
+        for (int j = 0; j < QK8_1; j++) {
+            const float v = x[i*QK8_1 + j];
+            amax = MAX(amax, fabsf(v));
+        }
+
+        const float d = amax / ((1 << 7) - 1);
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        int sum = 0;
+
+        for (int j = 0; j < QK8_1/2; ++j) {
+            const float v0 = x[i*QK8_1           + j]*id;
+            const float v1 = x[i*QK8_1 + QK8_1/2 + j]*id;
+
+            y[i].qs[          j] = roundf(v0);
+            y[i].qs[QK8_1/2 + j] = roundf(v1);
+
+            sum += y[i].qs[          j];
+            sum += y[i].qs[QK8_1/2 + j];
+        }
+
+        y[i].s = GGML_FP32_TO_FP16(sum*d);
+    }
+}
+
+static inline int best_index_mxfp4(float x, float e) {
+    int best_index = 0;
+    float best_err = fabsf(kvalues_mxfp4[0]*e - x);
+    for (int i = 1; i < 16; i++) {
+        float err = fabsf(kvalues_mxfp4[i]*e - x);
+        if (err < best_err) {
+            best_index = i;
+            best_err = err;
+        }
+    }
+    return best_index;
+}
+
+void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK_MXFP4;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+
+        for (int j = 0; j < qk; j++) {
+            const float v = x[i*qk + j];
+
+            if (amax < fabsf(v)) {
+                amax = fabsf(v);
+            }
+        }
+
+        const uint8_t e = amax > 0.0f ? (uint8_t) (floorf(log2f(amax)) - 2 + 127) : 0;
+
+        const float d = GGML_E8M0_TO_FP32_HALF(e);
+
+        y[i].e = e;
+
+        for (int j = 0; j < qk/2; ++j) {
+            const uint8_t x0 = best_index_mxfp4(x[i*qk + 0    + j], d);
+            const uint8_t x1 = best_index_mxfp4(x[i*qk + qk/2 + j], d);
+
+            y[i].qs[j]  = x0;
+            y[i].qs[j] |= x1 << 4;
+        }
+    }
+}
+
+void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK_NVFP4;
+    static const int qk_sub = QK_NVFP4_SUB;
+    static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        for (int s = 0; s < n_sub; s++) {
+            const float * xb = x + i*qk + s*qk_sub;
+
+            float amax = 0.0f;
+            for (int j = 0; j < qk_sub; j++) {
+                if (amax < fabsf(xb[j])) {
+                    amax = fabsf(xb[j]);
+                }
+            }
+
+            // UE4M3 scale: amax / 6.0 maps the max E2M1 value (6.0) to amax
+            const uint8_t ue = ggml_fp32_to_ue4m3(amax / 6.0f);
+            y[i].d[s] = ue;
+            const float d = ggml_ue4m3_to_fp32(ue);
+
+            for (int j = 0; j < qk_sub/2; ++j) {
+                const uint8_t x0 = best_index_mxfp4(xb[0        + j], d);
+                const uint8_t x1 = best_index_mxfp4(xb[qk_sub/2 + j], d);
+
+                y[i].qs[s*(qk_sub/2) + j] = x0 | (x1 << 4);
+            }
+        }
+    }
+}
+
+void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK1_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float neg_d = -d;
+
+        for (int j = 0; j < qk; ++j) {
+            const int byte_index = j / 8;
+            const int bit_offset = j % 8;
+            const uint8_t bit = (x[i].qs[byte_index] >> bit_offset) & 1;
+            y[i*qk + j] = bit ? d : neg_d;
+        }
+    }
+}
+
+void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK2_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int j = 0; j < qk; ++j) {
+            const int byte_index = j / 4;
+            const int bit_offset = (j % 4) * 2;
+            const uint8_t q = (x[i].qs[byte_index] >> bit_offset) & 0x03;
+            // 00=-1, 01=0, 10=+1, 11=+2
+            y[i*qk + j] = ((int)q - 1) * d;
+        }
+    }
+}
+
+void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK4_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int x0 = (x[i].qs[j] & 0x0F) - 8;
+            const int x1 = (x[i].qs[j] >>   4) - 8;
+
+            y[i*qk + j + 0   ] = x0*d;
+            y[i*qk + j + qk/2] = x1*d;
+        }
+    }
+}
+
+void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK4_1;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float m = GGML_FP16_TO_FP32(x[i].m);
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int x0 = (x[i].qs[j] & 0x0F);
+            const int x1 = (x[i].qs[j] >>   4);
+
+            y[i*qk + j + 0   ] = x0*d + m;
+            y[i*qk + j + qk/2] = x1*d + m;
+        }
+    }
+}
+
+void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK5_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        uint32_t qh;
+        memcpy(&qh, x[i].qh, sizeof(qh));
+
+        for (int j = 0; j < qk/2; ++j) {
+            const uint8_t xh_0 = ((qh >> (j +  0)) << 4) & 0x10;
+            const uint8_t xh_1 = ((qh >> (j + 12))     ) & 0x10;
+
+            const int32_t x0 = ((x[i].qs[j] & 0x0F) | xh_0) - 16;
+            const int32_t x1 = ((x[i].qs[j] >>   4) | xh_1) - 16;
+
+            y[i*qk + j + 0   ] = x0*d;
+            y[i*qk + j + qk/2] = x1*d;
+        }
+    }
+}
+
+void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK5_1;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float m = GGML_FP16_TO_FP32(x[i].m);
+
+        uint32_t qh;
+        memcpy(&qh, x[i].qh, sizeof(qh));
+
+        for (int j = 0; j < qk/2; ++j) {
+            const uint8_t xh_0 = ((qh >> (j +  0)) << 4) & 0x10;
+            const uint8_t xh_1 = ((qh >> (j + 12))     ) & 0x10;
+
+            const int x0 = (x[i].qs[j] & 0x0F) | xh_0;
+            const int x1 = (x[i].qs[j] >>   4) | xh_1;
+
+            y[i*qk + j + 0   ] = x0*d + m;
+            y[i*qk + j + qk/2] = x1*d + m;
+        }
+    }
+}
+
+void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK8_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int j = 0; j < qk; ++j) {
+            y[i*qk + j] = x[i].qs[j]*d;
+        }
+    }
+}
+
+void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK_MXFP4;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_E8M0_TO_FP32_HALF(x[i].e);
+
+        for (int j = 0; j < qk/2; ++j) {
+            const int8_t x0 = kvalues_mxfp4[x[i].qs[j] & 0x0F];
+            const int8_t x1 = kvalues_mxfp4[x[i].qs[j] >>   4];
+
+            y[i*qk + j + 0   ] = x0*d;
+            y[i*qk + j + qk/2] = x1*d;
+        }
+    }
+}
+
+void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK_NVFP4;
+    static const int qk_sub = QK_NVFP4_SUB;
+    static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        for (int s = 0; s < n_sub; s++) {
+            const float d = ggml_ue4m3_to_fp32(x[i].d[s]);
+            float * yb = y + i*qk + s*qk_sub;
+
+            for (int j = 0; j < qk_sub/2; ++j) {
+                const int8_t v0 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] & 0x0F];
+                const int8_t v1 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] >>   4];
+
+                yb[j + 0       ] = v0*d;
+                yb[j + qk_sub/2] = v1*d;
+            }
+        }
+    }
+}
+
+//
+// 2-6 bit quantization in super-blocks
+//
+
+//
+// ===================== Helper functions
+//
+static inline int nearest_int(float fval) {
+    assert(fabsf(fval) <= 4194303.f);
+    float val = fval + 12582912.f;
+    int i; memcpy(&i, &val, sizeof(int));
+    return (i & 0x007fffff) - 0x00400000;
+}
+
+static float make_qx_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, int rmse_type,
+        const float * GGML_RESTRICT qw) {
+    float max = 0;
+    float amax = 0;
+    for (int i = 0; i < n; ++i) {
+        float ax = fabsf(x[i]);
+        if (ax > amax) { amax = ax; max = x[i]; }
+    }
+    if (amax < GROUP_MAX_EPS) { // all zero
+        for (int i = 0; i < n; ++i) {
+            L[i] = 0;
+        }
+        return 0.f;
+    }
+    float iscale = -nmax / max;
+    if (rmse_type == 0) {
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale * x[i]);
+            L[i] = nmax + MAX(-nmax, MIN(nmax-1, l));
+        }
+        return 1/iscale;
+    }
+    bool return_early = false;
+    if (rmse_type < 0) {
+        rmse_type = -rmse_type;
+        return_early = true;
+    }
+    float sumlx = 0;
+    float suml2 = 0;
+#ifdef HAVE_BUGGY_APPLE_LINKER
+    // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7
+    for (volatile int i = 0; i < n; ++i) {
+#else
+    for (int i = 0; i < n; ++i) {
+#endif
+        int l = nearest_int(iscale * x[i]);
+        l = MAX(-nmax, MIN(nmax-1, l));
+        L[i] = l + nmax;
+        float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i]));
+        sumlx += w*x[i]*l;
+        suml2 += w*l*l;
+    }
+    float scale = suml2 ? sumlx/suml2 : 0.0f;
+    if (return_early) return suml2 > 0 ? 0.5f*(scale + 1/iscale) : 1/iscale;
+    float best = scale * sumlx;
+    for (int is = -9; is <= 9; ++is) {
+        if (is == 0) {
+            continue;
+        }
+        iscale = -(nmax + 0.1f*is) / max;
+        sumlx = suml2 = 0;
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale * x[i]);
+            l = MAX(-nmax, MIN(nmax-1, l));
+            float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i]));
+            sumlx += w*x[i]*l;
+            suml2 += w*l*l;
+        }
+        if (suml2 > 0 && sumlx*sumlx > best*suml2) {
+            for (int i = 0; i < n; ++i) {
+                int l = nearest_int(iscale * x[i]);
+                L[i] = nmax + MAX(-nmax, MIN(nmax-1, l));
+            }
+            scale = sumlx/suml2; best = scale*sumlx;
+        }
+    }
+    return scale;
+}
+
+static float make_q3_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, bool do_rmse) {
+    float max = 0;
+    float amax = 0;
+    for (int i = 0; i < n; ++i) {
+        float ax = fabsf(x[i]);
+        if (ax > amax) { amax = ax; max = x[i]; }
+    }
+    if (amax < GROUP_MAX_EPS) { // all zero
+        for (int i = 0; i < n; ++i) { L[i] = 0; }
+        return 0.f;
+    }
+    float iscale = -nmax / max;
+    if (do_rmse) {
+        float sumlx = 0;
+        float suml2 = 0;
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale * x[i]);
+            l = MAX(-nmax, MIN(nmax-1, l));
+            L[i] = l;
+            float w = x[i]*x[i];
+            sumlx += w*x[i]*l;
+            suml2 += w*l*l;
+        }
+        for (int itry = 0; itry < 5; ++itry) {
+            int n_changed = 0;
+            for (int i = 0; i < n; ++i) {
+                float w = x[i]*x[i];
+                float slx = sumlx - w*x[i]*L[i];
+                if (slx > 0) {
+                    float sl2 = suml2 - w*L[i]*L[i];
+                    int new_l = nearest_int(x[i] * sl2 / slx);
+                    new_l = MAX(-nmax, MIN(nmax-1, new_l));
+                    if (new_l != L[i]) {
+                        slx += w*x[i]*new_l;
+                        sl2 += w*new_l*new_l;
+                        if (sl2 > 0 && slx*slx*suml2 > sumlx*sumlx*sl2) {
+                            L[i] = new_l; sumlx = slx; suml2 = sl2;
+                            ++n_changed;
+                        }
+                    }
+                }
+            }
+            if (!n_changed) {
+                break;
+            }
+        }
+        for (int i = 0; i < n; ++i) {
+            L[i] += nmax;
+        }
+        return suml2 > 0.0f ? sumlx / suml2 : 0.0f;
+    }
+    for (int i = 0; i < n; ++i) {
+        int l = nearest_int(iscale * x[i]);
+        l = MAX(-nmax, MIN(nmax-1, l));
+        L[i] = l + nmax;
+    }
+    return 1/iscale;
+}
+
+static float make_qkx1_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min,
+        int ntry, float alpha) {
+    float min = x[0];
+    float max = x[0];
+    for (int i = 1; i < n; ++i) {
+        if (x[i] < min) min = x[i];
+        if (x[i] > max) max = x[i];
+    }
+    if (max == min) {
+        for (int i = 0; i < n; ++i) L[i] = 0;
+        *the_min = 0;
+        return 0.f;
+    }
+    if (min > 0) min = 0;
+    float iscale = nmax/(max - min);
+    float scale = 1/iscale;
+    for (int itry = 0; itry < ntry; ++itry) {
+        float sumlx = 0; int suml2 = 0;
+        bool did_change = false;
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale*(x[i] - min));
+            l = MAX(0, MIN(nmax, l));
+            if (l != L[i]) {
+                L[i] = l;
+                did_change = true;
+            }
+            sumlx += (x[i] - min)*l;
+            suml2 += l*l;
+        }
+        scale = sumlx/suml2;
+        float sum = 0;
+        for (int i = 0; i < n; ++i) {
+            sum += x[i] - scale*L[i];
+        }
+        min = alpha*min + (1 - alpha)*sum/n;
+        if (min > 0) min = 0;
+        iscale = 1/scale;
+        if (!did_change) break;
+    }
+    *the_min = -min;
+    return scale;
+}
+
+static float make_qkx2_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights,
+        uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux,
+        float rmin, float rdelta, int nstep, bool use_mad) {
+    float min = x[0];
+    float max = x[0];
+    float sum_w = weights[0];
+    float sum_x = sum_w * x[0];
+#ifdef HAVE_BUGGY_APPLE_LINKER
+    // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7
+    for (volatile int i = 1; i < n; ++i) {
+#else
+    for (int i = 1; i < n; ++i) {
+#endif
+        if (x[i] < min) min = x[i];
+        if (x[i] > max) max = x[i];
+        float w = weights[i];
+        sum_w += w;
+        sum_x += w * x[i];
+    }
+    if (min > 0) min = 0;
+    if (max == min) {
+        for (int i = 0; i < n; ++i) L[i] = 0;
+        *the_min = -min;
+        return 0.f;
+    }
+    float iscale = nmax/(max - min);
+    float scale = 1/iscale;
+    float best_error = 0;
+    for (int i = 0; i < n; ++i) {
+        int l = nearest_int(iscale*(x[i] - min));
+        L[i] = MAX(0, MIN(nmax, l));
+        float diff = scale * L[i] + min - x[i];
+        diff = use_mad ? fabsf(diff) : diff * diff;
+        float w = weights[i];
+        best_error += w * diff;
+    }
+    if (nstep < 1) {
+        *the_min = -min;
+        return scale;
+    }
+    for (int is = 0; is <= nstep; ++is) {
+        iscale = (rmin + rdelta*is + nmax)/(max - min);
+        float sum_l = 0, sum_l2 = 0, sum_xl = 0;
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale*(x[i] - min));
+            l = MAX(0, MIN(nmax, l));
+            Laux[i] = l;
+            float w = weights[i];
+            sum_l += w*l;
+            sum_l2 += w*l*l;
+            sum_xl += w*l*x[i];
+        }
+        float D = sum_w * sum_l2 - sum_l * sum_l;
+        if (D > 0) {
+            float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D;
+            float this_min   = (sum_l2 * sum_x - sum_l * sum_xl)/D;
+            if (this_min > 0) {
+                this_min = 0;
+                this_scale = sum_xl / sum_l2;
+            }
+            float cur_error = 0;
+            for (int i = 0; i < n; ++i) {
+                float diff = this_scale * Laux[i] + this_min - x[i];
+                diff = use_mad ? fabsf(diff) : diff * diff;
+                float w = weights[i];
+                cur_error += w * diff;
+            }
+            if (cur_error < best_error) {
+                for (int i = 0; i < n; ++i) {
+                    L[i] = Laux[i];
+                }
+                best_error = cur_error;
+                scale = this_scale;
+                min = this_min;
+            }
+        }
+    }
+    *the_min = -min;
+    return scale;
+}
+
+static inline void get_scale_min_k4(int j, const uint8_t * GGML_RESTRICT q, uint8_t * GGML_RESTRICT d, uint8_t * GGML_RESTRICT m) {
+    if (j < 4) {
+        *d = q[j] & 63; *m = q[j + 4] & 63;
+    } else {
+        *d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
+        *m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
+    }
+}
+
+//========================- 2-bit (de)-quantization
+
+void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    uint8_t L[QK_K];
+    uint8_t Laux[16];
+    float   weights[16];
+    float mins[QK_K/16];
+    float scales[QK_K/16];
+
+    const float q4scale = 15.f;
+
+    for (int i = 0; i < nb; i++) {
+        float max_scale = 0; // as we are deducting the min, scales are always positive
+        float max_min = 0;
+        for (int j = 0; j < QK_K/16; ++j) {
+            for (int l = 0; l < 16; ++l) weights[l] = fabsf(x[16*j + l]);
+            scales[j] = make_qkx2_quants(16, 3, x + 16*j, weights, L + 16*j, &mins[j], Laux, -0.5f, 0.1f, 15, true);
+            float scale = scales[j];
+            if (scale > max_scale) {
+                max_scale = scale;
+            }
+            float min = mins[j];
+            if (min > max_min) {
+                max_min = min;
+            }
+        }
+
+        if (max_scale > 0) {
+            float iscale = q4scale/max_scale;
+            for (int j = 0; j < QK_K/16; ++j) {
+                int l = nearest_int(iscale*scales[j]);
+                y[i].scales[j] = l;
+            }
+            y[i].d = GGML_FP32_TO_FP16(max_scale/q4scale);
+        } else {
+            for (int j = 0; j < QK_K/16; ++j) y[i].scales[j] = 0;
+            y[i].d = GGML_FP32_TO_FP16(0.f);
+        }
+        if (max_min > 0) {
+            float iscale = q4scale/max_min;
+            for (int j = 0; j < QK_K/16; ++j) {
+                int l = nearest_int(iscale*mins[j]);
+                y[i].scales[j] |= (l << 4);
+            }
+            y[i].dmin = GGML_FP32_TO_FP16(max_min/q4scale);
+        } else {
+            y[i].dmin = GGML_FP32_TO_FP16(0.f);
+        }
+        for (int j = 0; j < QK_K/16; ++j) {
+            const float d = GGML_FP16_TO_FP32(y[i].d) * (y[i].scales[j] & 0xF);
+            if (!d) continue;
+            const float dm = GGML_FP16_TO_FP32(y[i].dmin) * (y[i].scales[j] >> 4);
+            for (int ii = 0; ii < 16; ++ii) {
+                int l = nearest_int((x[16*j + ii] + dm)/d);
+                l = MAX(0, MIN(3, l));
+                L[16*j + ii] = l;
+            }
+        }
+
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6);
+            }
+        }
+
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float min = GGML_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * q = x[i].qs;
+
+        int is = 0;
+        float dl, ml;
+        for (int n = 0; n < QK_K; n += 128) {
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {
+
+                uint8_t sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml;
+
+                sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml;
+
+                shift += 2;
+            }
+            q += 32;
+        }
+    }
+}
+
+static float make_qkx3_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights,
+        uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux,
+        float rmin, float rdelta, int nstep, bool use_mad) {
+    float min = x[0];
+    float max = x[0];
+    float sum_w = weights ? weights[0] : x[0]*x[0];
+    float sum_x = sum_w * x[0];
+#ifdef HAVE_BUGGY_APPLE_LINKER
+    // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7
+    for (volatile int i = 1; i < n; ++i) {
+#else
+    for (int i = 1; i < n; ++i) {
+#endif
+        if (x[i] < min) min = x[i];
+        if (x[i] > max) max = x[i];
+        float w = weights ? weights[i] : x[i]*x[i];
+        sum_w += w;
+        sum_x += w * x[i];
+    }
+    if (min > 0) {
+        min = 0;
+    }
+    if (max <= min) {
+        memset(L, 0, n);
+        *the_min = -min;
+        return 0.f;
+    }
+    float iscale = nmax/(max - min);
+    float scale = 1/iscale;
+    float best_mad = 0;
+    for (int i = 0; i < n; ++i) {
+        int l = nearest_int(iscale*(x[i] - min));
+        L[i] = MAX(0, MIN(nmax, l));
+        float diff = scale * L[i] + min - x[i];
+        diff = use_mad ? fabsf(diff) : diff*diff;
+        float w = weights ? weights[i] : x[i]*x[i];
+        best_mad += w * diff;
+    }
+    if (nstep < 1) {
+        *the_min = -min;
+        return scale;
+    }
+    for (int is = 0; is <= nstep; ++is) {
+        iscale = (rmin + rdelta*is + nmax)/(max - min);
+        float sum_l = 0, sum_l2 = 0, sum_xl = 0;
+        for (int i = 0; i < n; ++i) {
+            // min is the best fit so far and can be at or near max, so v can be inf, nan or out of range for nearest_int
+            const float v = iscale*(x[i] - min);
+            const int l = v > 0 ? nearest_int(MIN(v, nmax)) : 0;
+            Laux[i] = l;
+            float w = weights ? weights[i] : x[i]*x[i];
+            sum_l  += w*l;
+            sum_l2 += w*l*l;
+            sum_xl += w*l*x[i];
+        }
+        float D = sum_w * sum_l2 - sum_l * sum_l;
+        if (D > 0) {
+            float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D;
+            float this_min   = (sum_l2 * sum_x - sum_l * sum_xl)/D;
+            if (this_min > 0) {
+                this_min = 0;
+                this_scale = sum_xl / sum_l2;
+            }
+            float mad = 0;
+            for (int i = 0; i < n; ++i) {
+                float diff = this_scale * Laux[i] + this_min - x[i];
+                diff = use_mad ? fabsf(diff) : diff*diff;
+                float w = weights ? weights[i] : x[i]*x[i];
+                mad += w * diff;
+            }
+            if (mad < best_mad) {
+                for (int i = 0; i < n; ++i) {
+                    L[i] = Laux[i];
+                }
+                best_mad = mad;
+                scale = this_scale;
+                min = this_min;
+            }
+        }
+    }
+    *the_min = -min;
+    return scale;
+}
+
+static float make_qp_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, const float * quant_weights) {
+    float max = 0;
+    for (int i = 0; i < n; ++i) {
+        max = MAX(max, x[i]);
+    }
+    if (max < GROUP_MAX_EPS) { // all zero
+        for (int i = 0; i < n; ++i) { L[i] = 0; }
+        return 0.f;
+    }
+    float iscale = nmax / max;
+    for (int i = 0; i < n; ++i) {
+        L[i] = nearest_int(iscale * x[i]);
+    }
+    float scale = 1/iscale;
+    float best_mse = 0;
+    for (int i = 0; i < n; ++i) {
+        float diff = x[i] - scale*L[i];
+        float w = quant_weights[i];
+        best_mse += w*diff*diff;
+    }
+    for (int is = -4; is <= 4; ++is) {
+        if (is == 0) continue;
+        float iscale_is = (0.1f*is + nmax)/max;
+        float scale_is = 1/iscale_is;
+        float mse = 0;
+        for (int i = 0; i < n; ++i) {
+            int l = nearest_int(iscale_is*x[i]);
+            l = MIN(nmax, l);
+            float diff = x[i] - scale_is*l;
+            float w = quant_weights[i];
+            mse += w*diff*diff;
+        }
+        if (mse < best_mse) {
+            best_mse = mse;
+            iscale = iscale_is;
+        }
+    }
+    float sumlx = 0;
+    float suml2 = 0;
+    for (int i = 0; i < n; ++i) {
+        int l = nearest_int(iscale * x[i]);
+        l = MIN(nmax, l);
+        L[i] = l;
+        float w = quant_weights[i];
+        sumlx += w*x[i]*l;
+        suml2 += w*l*l;
+    }
+    for (int itry = 0; itry < 5; ++itry) {
+        int n_changed = 0;
+        for (int i = 0; i < n; ++i) {
+            float w = quant_weights[i];
+            float slx = sumlx - w*x[i]*L[i];
+            float sl2 = suml2 - w*L[i]*L[i];
+            if (slx > 0 && sl2 > 0) {
+                int new_l = nearest_int(x[i] * sl2 / slx);
+                new_l = MIN(nmax, new_l);
+                if (new_l != L[i]) {
+                    slx += w*x[i]*new_l;
+                    sl2 += w*new_l*new_l;
+                    if (slx*slx*suml2 > sumlx*sumlx*sl2) {
+                        L[i] = new_l; sumlx = slx; suml2 = sl2;
+                        ++n_changed;
+                    }
+                }
+            }
+        }
+        if (!n_changed) {
+            break;
+        }
+    }
+    return suml2 > 0.0f ? sumlx / suml2 : 0.0f;
+}
+
+static void quantize_row_q2_K_impl(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int k, const float * GGML_RESTRICT quant_weights) {
+    GGML_ASSERT(quant_weights);
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+    const bool requantize = true;
+
+    uint8_t L[QK_K];
+    uint8_t Laux[16];
+    float mins[QK_K/16];
+    float scales[QK_K/16];
+    float sw[QK_K/16];
+    float weight[16];
+    uint8_t Ls[QK_K/16], Lm[QK_K/16];
+
+    for (int i = 0; i < nb; i++) {
+        memset(sw, 0, QK_K/16*sizeof(float));
+        float sumx2 = 0;
+        for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j];
+        float sigma2 = sumx2/QK_K;
+        for (int j = 0; j < QK_K/16; ++j) {
+            const float * GGML_RESTRICT qw = quant_weights + QK_K * i + 16*j;
+            for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j + l]*x[16*j + l]);
+            for (int l = 0; l < QK_K/16; ++l) sw[j] += weight[l];
+            scales[j] = make_qkx3_quants(16, 3, x + 16*j, weight, L + 16*j, &mins[j], Laux, -0.9f, 0.05f, 36, false);
+        }
+
+        float dm, mm;
+        dm  = make_qp_quants(QK_K/16, 15, scales, Ls, sw);
+        mm  = make_qp_quants(QK_K/16, 15, mins,   Lm, sw);
+
+        y[i].d    = GGML_FP32_TO_FP16(dm);
+        y[i].dmin = GGML_FP32_TO_FP16(mm);
+        dm        = GGML_FP16_TO_FP32(y[i].d);
+        mm        = GGML_FP16_TO_FP32(y[i].dmin);
+
+        for (int j = 0; j < QK_K/16; ++j) {
+            y[i].scales[j] = Ls[j] | (Lm[j] << 4);
+        }
+
+        if (requantize) {
+            for (int j = 0; j < QK_K/16; ++j) {
+                const float d = dm * (y[i].scales[j] & 0xF);
+                if (!d) continue;
+                const float m = mm * (y[i].scales[j] >> 4);
+                for (int ii = 0; ii < 16; ++ii) {
+                    int l = nearest_int((x[16*j + ii] + m)/d);
+                    l = MAX(0, MIN(3, l));
+                    L[16*j + ii] = l;
+                }
+            }
+        }
+
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6);
+            }
+        }
+
+        x += QK_K;
+    }
+}
+
+size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    size_t row_size = ggml_row_size(GGML_TYPE_Q2_K, n_per_row);
+    if (!quant_weights) {
+        quantize_row_q2_K_ref(src, dst, (int64_t)nrow*n_per_row);
+    }
+    else {
+        char * qrow = (char *)dst;
+        for (int64_t row = 0; row < nrow; ++row) {
+            quantize_row_q2_K_impl(src, (block_q2_K*)qrow, n_per_row, quant_weights);
+            src += n_per_row;
+            qrow += row_size;
+        }
+    }
+    return nrow * row_size;
+}
+
+//========================= 3-bit (de)-quantization
+
+void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    int8_t L[QK_K];
+    float scales[QK_K / 16];
+
+    for (int i = 0; i < nb; i++) {
+
+        float max_scale = 0;
+        float amax = 0;
+        for (int j = 0; j < QK_K/16; ++j) {
+            scales[j] = make_q3_quants(16, 4, x + 16*j, L + 16*j, true);
+            float scale = fabsf(scales[j]);
+            if (scale > amax) {
+                amax = scale; max_scale = scales[j];
+            }
+        }
+
+        memset(y[i].scales, 0, 12);
+        if (max_scale) {
+            float iscale = -32.f/max_scale;
+            for (int j = 0; j < QK_K/16; ++j) {
+                int8_t l = nearest_int(iscale*scales[j]);
+                l = MAX(-32, MIN(31, l)) + 32;
+                if (j < 8) {
+                    y[i].scales[j] = l & 0xF;
+                } else {
+                    y[i].scales[j-8] |= ((l & 0xF) << 4);
+                }
+                l >>= 4;
+                y[i].scales[j%4 + 8] |= (l << (2*(j/4)));
+            }
+            y[i].d = GGML_FP32_TO_FP16(1/iscale);
+        } else {
+            y[i].d = GGML_FP32_TO_FP16(0.f);
+        }
+
+        int8_t sc;
+        for (int j = 0; j < QK_K/16; ++j) {
+            sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4;
+            sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32;
+            float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) {
+                continue;
+            }
+            for (int ii = 0; ii < 16; ++ii) {
+                int l = nearest_int(x[16*j + ii]/d);
+                l = MAX(-4, MIN(3, l));
+                L[16*j + ii] = l + 4;
+            }
+        }
+
+        memset(y[i].hmask, 0, QK_K/8);
+        // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc.
+        int m = 0;
+        uint8_t hm = 1;
+        for (int j = 0; j < QK_K; ++j) {
+            if (L[j] > 3) {
+                y[i].hmask[m] |= hm;
+                L[j] -= 4;
+            }
+            if (++m == QK_K/8) {
+                m = 0; hm <<= 1;
+            }
+        }
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6);
+            }
+        }
+
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    const uint32_t kmask1 = 0x03030303;
+    const uint32_t kmask2 = 0x0f0f0f0f;
+
+    uint32_t aux[4];
+    const int8_t * scales = (const int8_t*)aux;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d_all = GGML_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q = x[i].qs;
+        const uint8_t * GGML_RESTRICT hm = x[i].hmask;
+        uint8_t m = 1;
+
+        memcpy(aux, x[i].scales, 12);
+        uint32_t tmp = aux[2];
+        aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+        aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+        aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
+        aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+
+        int is = 0;
+        float dl;
+        for (int n = 0; n < QK_K; n += 128) {
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {
+                    *y++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4));
+                }
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {
+                    *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4));
+                }
+
+                shift += 2;
+                m <<= 1;
+            }
+            q += 32;
+        }
+
+    }
+}
+
+static void quantize_row_q3_K_impl(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t n_per_row, const float * GGML_RESTRICT quant_weights) {
+    assert(n_per_row % QK_K == 0);
+    const int nb = n_per_row / QK_K;
+
+    int8_t L[QK_K];
+    float scales[QK_K / 16];
+    float weight[16];
+    float sw[QK_K / 16];
+    int8_t Ls[QK_K / 16];
+
+    for (int i = 0; i < nb; i++) {
+
+        float sumx2 = 0;
+        for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int j = 0; j < QK_K/16; ++j) {
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K * i + 16*j;
+                for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j+l]*x[16*j+l]);
+            } else {
+                for (int l = 0; l < 16; ++l) weight[l] = x[16*j+l]*x[16*j+l];
+            }
+            float sumw = 0;
+            for (int l = 0; l < 16; ++l) sumw += weight[l];
+            sw[j] = sumw;
+
+            scales[j] = make_qx_quants(16, 4, x + 16*j, L + 16*j, 1, weight);
+
+        }
+
+        memset(y[i].scales, 0, 12);
+
+        float d_block = make_qx_quants(QK_K/16, 32, scales, Ls, 1, sw);
+        for (int j = 0; j < QK_K/16; ++j) {
+            int l = Ls[j];
+            if (j < 8) {
+                y[i].scales[j] = l & 0xF;
+            } else {
+                y[i].scales[j-8] |= ((l & 0xF) << 4);
+            }
+            l >>= 4;
+            y[i].scales[j%4 + 8] |= (l << (2*(j/4)));
+        }
+        y[i].d = GGML_FP32_TO_FP16(d_block);
+
+        int8_t sc;
+        for (int j = 0; j < QK_K/16; ++j) {
+            sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4;
+            sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32;
+            float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) {
+                continue;
+            }
+            for (int ii = 0; ii < 16; ++ii) {
+                int l = nearest_int(x[16*j + ii]/d);
+                l = MAX(-4, MIN(3, l));
+                L[16*j + ii] = l + 4;
+            }
+        }
+
+        memset(y[i].hmask, 0, QK_K/8);
+        // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc.
+        int m = 0;
+        uint8_t hm = 1;
+        for (int j = 0; j < QK_K; ++j) {
+            if (L[j] > 3) {
+                y[i].hmask[m] |= hm;
+                L[j] -= 4;
+            }
+            if (++m == QK_K/8) {
+                m = 0; hm <<= 1;
+            }
+        }
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6);
+            }
+        }
+
+        x += QK_K;
+    }
+}
+
+size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    size_t row_size = ggml_row_size(GGML_TYPE_Q3_K, n_per_row);
+    if (!quant_weights) {
+        quantize_row_q3_K_ref(src, dst, (int64_t)nrow*n_per_row);
+    }
+    else {
+        char * qrow = (char *)dst;
+        for (int64_t row = 0; row < nrow; ++row) {
+            quantize_row_q3_K_impl(src, (block_q3_K*)qrow, n_per_row, quant_weights);
+            src += n_per_row;
+            qrow += row_size;
+        }
+    }
+    return nrow * row_size;
+}
+
+// ====================== 4-bit (de)-quantization
+
+void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    uint8_t L[QK_K];
+    uint8_t Laux[32];
+    float   weights[32];
+    float mins[QK_K/32];
+    float scales[QK_K/32];
+
+    for (int i = 0; i < nb; i++) {
+        float max_scale = 0; // as we are deducting the min, scales are always positive
+        float max_min = 0;
+        for (int j = 0; j < QK_K/32; ++j) {
+            //scales[j] = make_qkx1_quants(32, 15, x + 32*j, L + 32*j, &mins[j], 9, 0.5f);
+            float sum_x2 = 0;
+            for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l];
+            float av_x = sqrtf(sum_x2/32);
+            for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]);
+            scales[j] = make_qkx2_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -1.f, 0.1f, 20, false);
+            float scale = scales[j];
+            if (scale > max_scale) {
+                max_scale = scale;
+            }
+            float min = mins[j];
+            if (min > max_min) {
+                max_min = min;
+            }
+        }
+
+        float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f;
+        float inv_min   = max_min   > 0 ? 63.f/max_min   : 0.f;
+        for (int j = 0; j < QK_K/32; ++j) {
+            uint8_t ls = nearest_int(inv_scale*scales[j]);
+            uint8_t lm = nearest_int(inv_min*mins[j]);
+            ls = MIN(63, ls);
+            lm = MIN(63, lm);
+            if (j < 4) {
+                y[i].scales[j] = ls;
+                y[i].scales[j+4] = lm;
+            } else {
+                y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4);
+                y[i].scales[j-4] |= ((ls >> 4) << 6);
+                y[i].scales[j-0] |= ((lm >> 4) << 6);
+            }
+        }
+        y[i].d = GGML_FP32_TO_FP16(max_scale/63.f);
+        y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f);
+
+        uint8_t sc, m;
+        for (int j = 0; j < QK_K/32; ++j) {
+            get_scale_min_k4(j, y[i].scales, &sc, &m);
+            const float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) continue;
+            const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m;
+            for (int ii = 0; ii < 32; ++ii) {
+                int l = nearest_int((x[32*j + ii] + dm)/d);
+                l = MAX(0, MIN(15, l));
+                L[32*j + ii] = l;
+            }
+        }
+
+        uint8_t * q = y[i].qs;
+        for (int j = 0; j < QK_K; j += 64) {
+            for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4);
+            q += 32;
+        }
+
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+        const uint8_t * q = x[i].qs;
+
+        const float d   = GGML_FP16_TO_FP32(x[i].d);
+        const float min = GGML_FP16_TO_FP32(x[i].dmin);
+
+        int is = 0;
+        uint8_t sc, m;
+        for (int j = 0; j < QK_K; j += 64) {
+            get_scale_min_k4(is + 0, x[i].scales, &sc, &m);
+            const float d1 = d * sc; const float m1 = min * m;
+            get_scale_min_k4(is + 1, x[i].scales, &sc, &m);
+            const float d2 = d * sc; const float m2 = min * m;
+            for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 0xF) - m1;
+            for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l]  >> 4) - m2;
+            q += 32; is += 2;
+        }
+    }
+}
+
+static void quantize_row_q4_K_impl(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    assert(n_per_row % QK_K == 0);
+    const int64_t nb = n_per_row / QK_K;
+
+    uint8_t L[QK_K];
+    uint8_t Laux[32];
+    uint8_t Ls[QK_K/32];
+    uint8_t Lm[QK_K/32];
+    float   weights[32];
+    float   sw[QK_K/32];
+    float   mins[QK_K/32];
+    float   scales[QK_K/32];
+
+    for (int i = 0; i < nb; i++) {
+
+        float sum_x2 = 0;
+        for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l];
+        float sigma2 = 2*sum_x2/QK_K;
+        float av_x = sqrtf(sigma2);
+
+        for (int j = 0; j < QK_K/32; ++j) {
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*i + 32*j;
+                for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]);
+            } else {
+                for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]);
+            }
+            float sumw = 0;
+            for (int l = 0; l < 32; ++l) sumw += weights[l];
+            sw[j] = sumw;
+            scales[j] = make_qkx3_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false);
+        }
+
+        float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw);
+        float m_block = make_qp_quants(QK_K/32, 63, mins,   Lm, sw);
+        for (int j = 0; j < QK_K/32; ++j) {
+            uint8_t ls = Ls[j];
+            uint8_t lm = Lm[j];
+            if (j < 4) {
+                y[i].scales[j] = ls;
+                y[i].scales[j+4] = lm;
+            } else {
+                y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4);
+                y[i].scales[j-4] |= ((ls >> 4) << 6);
+                y[i].scales[j-0] |= ((lm >> 4) << 6);
+            }
+        }
+        y[i].d = GGML_FP32_TO_FP16(d_block);
+        y[i].dmin = GGML_FP32_TO_FP16(m_block);
+
+        uint8_t sc, m;
+        for (int j = 0; j < QK_K/32; ++j) {
+            get_scale_min_k4(j, y[i].scales, &sc, &m);
+            const float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) continue;
+            const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m;
+            for (int ii = 0; ii < 32; ++ii) {
+                int l = nearest_int((x[32*j + ii] + dm)/d);
+                l = MAX(0, MIN(15, l));
+                L[32*j + ii] = l;
+            }
+        }
+        uint8_t * q = y[i].qs;
+        for (int j = 0; j < QK_K; j += 64) {
+            for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4);
+            q += 32;
+        }
+
+        x += QK_K;
+
+    }
+}
+
+size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    size_t row_size = ggml_row_size(GGML_TYPE_Q4_K, n_per_row);
+    if (!quant_weights) {
+        quantize_row_q4_K_ref(src, dst, (int64_t)nrow*n_per_row);
+    }
+    else {
+        char * qrow = (char *)dst;
+        for (int64_t row = 0; row < nrow; ++row) {
+            quantize_row_q4_K_impl(src, (block_q4_K*)qrow, n_per_row, quant_weights);
+            src += n_per_row;
+            qrow += row_size;
+        }
+    }
+    return nrow * row_size;
+}
+
+// ====================== 5-bit (de)-quantization
+
+void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    uint8_t L[QK_K];
+    float mins[QK_K/32];
+    float scales[QK_K/32];
+    float weights[32];
+    uint8_t Laux[32];
+
+    for (int i = 0; i < nb; i++) {
+        float max_scale = 0; // as we are deducting the min, scales are always positive
+        float max_min = 0;
+        for (int j = 0; j < QK_K/32; ++j) {
+            //scales[j] = make_qkx1_quants(32, 31, x + 32*j, L + 32*j, &mins[j], 9, 0.5f);
+            float sum_x2 = 0;
+            for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l];
+            float av_x = sqrtf(sum_x2/32);
+            for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]);
+            scales[j] = make_qkx2_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.5f, 0.1f, 15, false);
+            float scale = scales[j];
+            if (scale > max_scale) {
+                max_scale = scale;
+            }
+            float min = mins[j];
+            if (min > max_min) {
+                max_min = min;
+            }
+        }
+
+        float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f;
+        float inv_min   = max_min   > 0 ? 63.f/max_min   : 0.f;
+        for (int j = 0; j < QK_K/32; ++j) {
+            uint8_t ls = nearest_int(inv_scale*scales[j]);
+            uint8_t lm = nearest_int(inv_min*mins[j]);
+            ls = MIN(63, ls);
+            lm = MIN(63, lm);
+            if (j < 4) {
+                y[i].scales[j] = ls;
+                y[i].scales[j+4] = lm;
+            } else {
+                y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4);
+                y[i].scales[j-4] |= ((ls >> 4) << 6);
+                y[i].scales[j-0] |= ((lm >> 4) << 6);
+            }
+        }
+        y[i].d = GGML_FP32_TO_FP16(max_scale/63.f);
+        y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f);
+
+        uint8_t sc, m;
+        for (int j = 0; j < QK_K/32; ++j) {
+            get_scale_min_k4(j, y[i].scales, &sc, &m);
+            const float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) continue;
+            const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m;
+            for (int ii = 0; ii < 32; ++ii) {
+                int l = nearest_int((x[32*j + ii] + dm)/d);
+                l = MAX(0, MIN(31, l));
+                L[32*j + ii] = l;
+            }
+        }
+
+        uint8_t * GGML_RESTRICT qh = y[i].qh;
+        uint8_t * GGML_RESTRICT ql = y[i].qs;
+        memset(qh, 0, QK_K/8);
+
+        uint8_t m1 = 1, m2 = 2;
+        for (int n = 0; n < QK_K; n += 64) {
+            for (int j = 0; j < 32; ++j) {
+                int l1 = L[n + j];
+                if (l1 > 15) {
+                    l1 -= 16; qh[j] |= m1;
+                }
+                int l2 = L[n + j + 32];
+                if (l2 > 15) {
+                    l2 -= 16; qh[j] |= m2;
+                }
+                ql[j] = l1 | (l2 << 4);
+            }
+            m1 <<= 2; m2 <<= 2;
+            ql += 32;
+        }
+
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+        const uint8_t * ql = x[i].qs;
+        const uint8_t * qh = x[i].qh;
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float min = GGML_FP16_TO_FP32(x[i].dmin);
+
+        int is = 0;
+        uint8_t sc, m;
+        uint8_t u1 = 1, u2 = 2;
+        for (int j = 0; j < QK_K; j += 64) {
+            get_scale_min_k4(is + 0, x[i].scales, &sc, &m);
+            const float d1 = d * sc; const float m1 = min * m;
+            get_scale_min_k4(is + 1, x[i].scales, &sc, &m);
+            const float d2 = d * sc; const float m2 = min * m;
+            for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1;
+            for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l]  >> 4) + (qh[l] & u2 ? 16 : 0)) - m2;
+            ql += 32; is += 2;
+            u1 <<= 2; u2 <<= 2;
+        }
+    }
+}
+
+static void quantize_row_q5_K_impl(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    assert(n_per_row % QK_K == 0);
+    const int64_t nb = n_per_row / QK_K;
+
+    uint8_t L[QK_K];
+    uint8_t Laux[32];
+    uint8_t Ls[QK_K/32];
+    uint8_t Lm[QK_K/32];
+    float   mins[QK_K/32];
+    float   scales[QK_K/32];
+    float   sw[QK_K/32];
+    float   weights[32];
+
+    for (int i = 0; i < nb; i++) {
+
+        float sum_x2 = 0;
+        for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l];
+        float sigma2 = 2*sum_x2/QK_K;
+        float av_x = sqrtf(sigma2);
+
+        for (int j = 0; j < QK_K/32; ++j) {
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*i + 32*j;
+                for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]);
+            } else {
+                for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]);
+            }
+            float sumw = 0;
+            for (int l = 0; l < 32; ++l) sumw += weights[l];
+            sw[j] = sumw;
+
+            scales[j] = make_qkx3_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false);
+        }
+
+        float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw);
+        float m_block = make_qp_quants(QK_K/32, 63, mins,   Lm, sw);
+
+        for (int j = 0; j < QK_K/32; ++j) {
+            uint8_t ls = Ls[j];
+            uint8_t lm = Lm[j];
+            ls = MIN(63, ls);
+            lm = MIN(63, lm);
+            if (j < 4) {
+                y[i].scales[j] = ls;
+                y[i].scales[j+4] = lm;
+            } else {
+                y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4);
+                y[i].scales[j-4] |= ((ls >> 4) << 6);
+                y[i].scales[j-0] |= ((lm >> 4) << 6);
+            }
+        }
+        y[i].d = GGML_FP32_TO_FP16(d_block);
+        y[i].dmin = GGML_FP32_TO_FP16(m_block);
+
+        uint8_t sc, m;
+        for (int j = 0; j < QK_K/32; ++j) {
+            get_scale_min_k4(j, y[i].scales, &sc, &m);
+            const float d = GGML_FP16_TO_FP32(y[i].d) * sc;
+            if (!d) continue;
+            const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m;
+            for (int ii = 0; ii < 32; ++ii) {
+                int l = nearest_int((x[32*j + ii] + dm)/d);
+                l = MAX(0, MIN(31, l));
+                L[32*j + ii] = l;
+            }
+        }
+
+        uint8_t * GGML_RESTRICT qh = y[i].qh;
+        uint8_t * GGML_RESTRICT ql = y[i].qs;
+        memset(qh, 0, QK_K/8);
+
+        uint8_t m1 = 1, m2 = 2;
+        for (int n = 0; n < QK_K; n += 64) {
+            for (int j = 0; j < 32; ++j) {
+                int l1 = L[n + j];
+                if (l1 > 15) {
+                    l1 -= 16; qh[j] |= m1;
+                }
+                int l2 = L[n + j + 32];
+                if (l2 > 15) {
+                    l2 -= 16; qh[j] |= m2;
+                }
+                ql[j] = l1 | (l2 << 4);
+            }
+            m1 <<= 2; m2 <<= 2;
+            ql += 32;
+        }
+
+        x += QK_K;
+
+    }
+}
+
+size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    size_t row_size = ggml_row_size(GGML_TYPE_Q5_K, n_per_row);
+    if (!quant_weights) {
+        quantize_row_q5_K_ref(src, dst, (int64_t)nrow*n_per_row);
+    }
+    else {
+        char * qrow = (char *)dst;
+        for (int64_t row = 0; row < nrow; ++row) {
+            quantize_row_q5_K_impl(src, (block_q5_K*)qrow, n_per_row, quant_weights);
+            src += n_per_row;
+            qrow += row_size;
+        }
+    }
+    return nrow * row_size;
+}
+
+// ====================== 6-bit (de)-quantization
+
+void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    int8_t L[QK_K];
+    float   scales[QK_K/16];
+
+    for (int i = 0; i < nb; i++) {
+
+        float max_scale = 0;
+        float max_abs_scale = 0;
+
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+
+            const float scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL);
+            scales[ib] = scale;
+
+            const float abs_scale = fabsf(scale);
+            if (abs_scale > max_abs_scale) {
+                max_abs_scale = abs_scale;
+                max_scale = scale;
+            }
+
+        }
+
+        if (max_abs_scale < GROUP_MAX_EPS) {
+            memset(&y[i], 0, sizeof(block_q6_K));
+            y[i].d = GGML_FP32_TO_FP16(0.f);
+            x += QK_K;
+            continue;
+        }
+
+        float iscale = -128.f/max_scale;
+        y[i].d = GGML_FP32_TO_FP16(1/iscale);
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib]));
+        }
+
+        for (int j = 0; j < QK_K/16; ++j) {
+            float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j];
+            if (!d) {
+                continue;
+            }
+            for (int ii = 0; ii < 16; ++ii) {
+                int l = nearest_int(x[16*j + ii]/d);
+                l = MAX(-32, MIN(31, l));
+                L[16*j + ii] = l + 32;
+            }
+        }
+
+        uint8_t * GGML_RESTRICT ql = y[i].ql;
+        uint8_t * GGML_RESTRICT qh = y[i].qh;
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                const uint8_t q1 = L[j + l +  0] & 0xF;
+                const uint8_t q2 = L[j + l + 32] & 0xF;
+                const uint8_t q3 = L[j + l + 64] & 0xF;
+                const uint8_t q4 = L[j + l + 96] & 0xF;
+                ql[l+ 0] = q1 | (q3 << 4);
+                ql[l+32] = q2 | (q4 << 4);
+                qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6);
+            }
+            ql += 64;
+            qh += 32;
+        }
+
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT ql = x[i].ql;
+        const uint8_t * GGML_RESTRICT qh = x[i].qh;
+        const int8_t  * GGML_RESTRICT sc = x[i].scales;
+
+        for (int n = 0; n < QK_K; n += 128) {
+            for (int l = 0; l < 32; ++l) {
+                int is = l/16;
+                const int8_t q1 = (int8_t)((ql[l +  0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
+                const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
+                const int8_t q3 = (int8_t)((ql[l +  0]  >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
+                const int8_t q4 = (int8_t)((ql[l + 32]  >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
+                y[l +  0] = d * sc[is + 0] * q1;
+                y[l + 32] = d * sc[is + 2] * q2;
+                y[l + 64] = d * sc[is + 4] * q3;
+                y[l + 96] = d * sc[is + 6] * q4;
+            }
+            y  += 128;
+            ql += 64;
+            qh += 32;
+            sc += 8;
+        }
+    }
+}
+
+static void quantize_row_q6_K_impl(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    assert(n_per_row % QK_K == 0);
+    const int64_t nb = n_per_row / QK_K;
+
+    int8_t L[QK_K];
+    float   scales[QK_K/16];
+    //float   weights[16];
+
+    for (int i = 0; i < nb; i++) {
+
+        //float sum_x2 = 0;
+        //for (int j = 0; j < QK_K; ++j) sum_x2 += x[j]*x[j];
+        //float sigma2 = sum_x2/QK_K;
+
+        float max_scale = 0;
+        float max_abs_scale = 0;
+
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+
+            float scale;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*i + 16*ib;
+                //for (int j = 0; j < 16; ++j) weights[j] = qw[j] * sqrtf(sigma2 + x[16*ib + j]*x[16*ib + j]);
+                //scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, weights);
+                scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, qw);
+            } else {
+                scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL);
+            }
+            scales[ib] = scale;
+
+            const float abs_scale = fabsf(scale);
+            if (abs_scale > max_abs_scale) {
+                max_abs_scale = abs_scale;
+                max_scale = scale;
+            }
+
+        }
+
+        if (max_abs_scale < GROUP_MAX_EPS) {
+            memset(&y[i], 0, sizeof(block_q6_K));
+            y[i].d = GGML_FP32_TO_FP16(0.f);
+            x += QK_K;
+            continue;
+        }
+
+        float iscale = -128.f/max_scale;
+        y[i].d = GGML_FP32_TO_FP16(1/iscale);
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib]));
+        }
+
+        for (int j = 0; j < QK_K/16; ++j) {
+            float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j];
+            if (!d) {
+                continue;
+            }
+            for (int ii = 0; ii < 16; ++ii) {
+                int l = nearest_int(x[16*j + ii]/d);
+                l = MAX(-32, MIN(31, l));
+                L[16*j + ii] = l + 32;
+            }
+        }
+
+        uint8_t * GGML_RESTRICT ql = y[i].ql;
+        uint8_t * GGML_RESTRICT qh = y[i].qh;
+        for (int j = 0; j < QK_K; j += 128) {
+            for (int l = 0; l < 32; ++l) {
+                const uint8_t q1 = L[j + l +  0] & 0xF;
+                const uint8_t q2 = L[j + l + 32] & 0xF;
+                const uint8_t q3 = L[j + l + 64] & 0xF;
+                const uint8_t q4 = L[j + l + 96] & 0xF;
+                ql[l+ 0] = q1 | (q3 << 4);
+                ql[l+32] = q2 | (q4 << 4);
+                qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6);
+            }
+            ql += 64;
+            qh += 32;
+        }
+
+        x += QK_K;
+
+    }
+}
+
+size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    size_t row_size = ggml_row_size(GGML_TYPE_Q6_K, n_per_row);
+    if (!quant_weights) {
+        quantize_row_q6_K_ref(src, dst, (int64_t)nrow*n_per_row);
+    }
+    else {
+        char * qrow = (char *)dst;
+        for (int64_t row = 0; row < nrow; ++row) {
+            quantize_row_q6_K_impl(src, (block_q6_K*)qrow, n_per_row, quant_weights);
+            src += n_per_row;
+            qrow += row_size;
+        }
+    }
+    return nrow * row_size;
+}
+
+static void quantize_row_q4_0_impl(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    static_assert(QK4_0 == 32, "QK4_0 must be 32");
+
+    if (!quant_weights) {
+        quantize_row_q4_0_ref(x, y, n_per_row);
+        return;
+    }
+
+    float weight[QK4_0];
+    int8_t L[QK4_0];
+
+    float sum_x2 = 0;
+    for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j];
+    float sigma2 = sum_x2/n_per_row;
+
+    const int64_t nb = n_per_row/QK4_0;
+    for (int ib = 0; ib < nb; ++ib) {
+        const float * xb = x + QK4_0 * ib;
+        const float * qw = quant_weights + QK4_0 * ib;
+        for (int j = 0; j < QK4_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]);
+        float d = make_qx_quants(QK4_0, 8, xb, L, 1, weight);
+        y[ib].d = GGML_FP32_TO_FP16(d);
+        for (int j = 0; j < 16; ++j) {
+            y[ib].qs[j] = L[j] | (L[j+16] << 4);
+        }
+    }
+}
+
+size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q1_0_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q1_0, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q1_0, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q1_0_ref(src, (block_q1_0*)qrow, n_per_row);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q2_0_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q2_0_ref(src, (block_q2_0*)qrow, n_per_row);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q4_0_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q4_0, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q4_0, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q4_0_impl(src, (block_q4_0*)qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+static void quantize_row_q4_1_impl(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    static_assert(QK4_1 == 32, "QK4_1 must be 32");
+
+    if (!quant_weights) {
+        quantize_row_q4_1_ref(x, y, n_per_row);
+        return;
+    }
+
+    float weight[QK4_1];
+    uint8_t L[QK4_1], Laux[QK4_1];
+
+    float sum_x2 = 0;
+    for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j];
+    float sigma2 = sum_x2/n_per_row;
+
+    const int64_t nb = n_per_row/QK4_1;
+    for (int ib = 0; ib < nb; ++ib) {
+        const float * xb = x + QK4_1 * ib;
+        const float * qw = quant_weights + QK4_1 * ib;
+        for (int j = 0; j < QK4_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]);
+        float min;
+        float d = make_qkx3_quants(QK4_1, 15, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false);
+        y[ib].d = GGML_FP32_TO_FP16(d);
+        y[ib].m = GGML_FP32_TO_FP16(-min);
+        for (int j = 0; j < 16; ++j) {
+            y[ib].qs[j] = L[j] | (L[j+16] << 4);
+        }
+    }
+}
+
+size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q4_1_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q4_1, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q4_1, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q4_1_impl(src, (block_q4_1*)qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+static void quantize_row_q5_0_impl(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    static_assert(QK5_0 == 32, "QK5_0 must be 32");
+
+    if (!quant_weights) {
+        quantize_row_q5_0_ref(x, y, n_per_row);
+        return;
+    }
+
+    float weight[QK5_0];
+    int8_t L[QK5_0];
+
+    float sum_x2 = 0;
+    for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j];
+    float sigma2 = sum_x2/n_per_row;
+
+    const int64_t nb = n_per_row/QK5_0;
+    for (int ib = 0; ib < nb; ++ib) {
+        const float * xb = x + QK5_0 * ib;
+        const float * qw = quant_weights + QK5_0 * ib;
+        for (int j = 0; j < QK5_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]);
+        float d = make_qx_quants(QK5_0, 16, xb, L, 1, weight);
+        y[ib].d = GGML_FP32_TO_FP16(d);
+
+        uint32_t qh = 0;
+
+        for (int j = 0; j < 16; ++j) {
+            const uint8_t xi0 = L[j];
+            const uint8_t xi1 = L[j+16];
+            y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
+
+            // get the 5-th bit and store it in qh at the right position
+            qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+            qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2);
+        }
+
+        memcpy(&y[ib].qh, &qh, sizeof(qh));
+    }
+}
+
+size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q5_0_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q5_0, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q5_0, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q5_0_impl(src, (block_q5_0*)qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+static void quantize_row_q5_1_impl(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) {
+    static_assert(QK5_1 == 32, "QK5_1 must be 32");
+
+    if (!quant_weights) {
+        quantize_row_q5_1_ref(x, y, n_per_row);
+        return;
+    }
+
+    float weight[QK5_1];
+    uint8_t L[QK5_1], Laux[QK5_1];
+
+    float sum_x2 = 0;
+    for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j];
+    float sigma2 = sum_x2/n_per_row;
+
+    const int64_t nb = n_per_row/QK5_1;
+    for (int ib = 0; ib < nb; ++ib) {
+        const float * xb = x + QK5_1 * ib;
+        const float * qw = quant_weights + QK5_1 * ib;
+        for (int j = 0; j < QK5_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]);
+        float min;
+        float d = make_qkx3_quants(QK5_1, 31, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false);
+        y[ib].d = GGML_FP32_TO_FP16(d);
+        y[ib].m = GGML_FP32_TO_FP16(-min);
+
+        uint32_t qh = 0;
+        for (int j = 0; j < 16; ++j) {
+            const uint8_t xi0 = L[j];
+            const uint8_t xi1 = L[j+16];
+            y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
+            // get the 5-th bit and store it in qh at the right position
+            qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+            qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2);
+        }
+        memcpy(&y[ib].qh, &qh, sizeof(qh));
+    }
+}
+
+size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q5_1_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q5_1, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q5_1, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q5_1_impl(src, (block_q5_1*)qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
+
+size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    (void)quant_weights; // not used
+    const size_t row_size = ggml_row_size(GGML_TYPE_Q8_0, n_per_row);
+    quantize_row_q8_0_ref(src, dst, (int64_t)nrow*n_per_row);
+    return nrow * row_size;
+}
+
+size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_UNUSED(quant_weights);
+    quantize_row_mxfp4_ref(src, dst, (int64_t)nrow*n_per_row);
+    return nrow * ggml_row_size(GGML_TYPE_MXFP4, n_per_row);
+}
+
+size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_UNUSED(quant_weights);
+    quantize_row_nvfp4_ref(src, dst, (int64_t)nrow*n_per_row);
+    return nrow * ggml_row_size(GGML_TYPE_NVFP4, n_per_row);
+}
+
+// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs)
+
+void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int64_t i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+
+        for (int j = 0; j < QK_K; j++) {
+            const float v = x[j];
+            amax = MAX(amax, fabsf(v));
+        }
+
+        const float d = amax;
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        // 5 elements per byte, along 32 bytes
+        for (size_t j = 0; j < sizeof(y->qs) - sizeof(y->qs) % 32; j += 32) {
+            for (size_t m = 0; m < 32; ++m) {
+                uint8_t q = 0;
+                for (size_t n = 0; n < 5; ++n) {
+                    int xi = lroundf(x[m + n*32] * id) + 1; // -1, 0, 1 -> 0, 1, 2
+                    q *= 3;
+                    q += xi;
+                }
+                // ceiling division (243 == pow(3, 5))
+                q = ((uint16_t)q * 256 + (243 - 1)) / 243;
+                y[i].qs[j + m] = q;
+            }
+            x += 5*32;
+        }
+        // along 16 bytes
+        for (size_t j = sizeof(y->qs) - sizeof(y->qs) % 32; j < sizeof(y->qs); j += 16) {
+            for (size_t m = 0; m < 16; ++m) {
+                uint8_t q = 0;
+                for (size_t n = 0; n < 5; ++n) {
+                    int xi = lroundf(x[m + n*16] * id) + 1; // -1, 0, 1 -> 0, 1, 2
+                    q *= 3;
+                    q += xi;
+                }
+                // ceiling division (243 == pow(3, 5))
+                q = ((uint16_t)q * 256 + (243 - 1)) / 243;
+                y[i].qs[j + m] = q;
+            }
+            x += 5*16;
+        }
+        // 4 elements per byte
+        for (size_t j = 0; j < sizeof(y->qh); ++j) {
+            uint8_t q = 0;
+            for (size_t m = 0; m < 4; ++m) {
+                // -1, 0, 1 -> 0, 1, 2
+                int xi = lroundf(x[j + m*sizeof(y->qh)] * id) + 1;
+                q *= 3;
+                q += xi;
+            }
+            // shift the first value to the most significant trit
+            q *= 3;
+            // ceiling division (243 == pow(3, 5))
+            q = ((uint16_t)q * 256 + (243 - 1)) / 243;
+            y[i].qh[j] = q;
+        }
+        x += 4*sizeof(y->qh);
+    }
+}
+
+void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int64_t i = 0; i < nb; i++) {
+        float amax = 0.0f; // absolute max
+
+        for (int j = 0; j < QK_K; j++) {
+            const float v = x[j];
+            amax = MAX(amax, fabsf(v));
+        }
+
+        const float d = amax;
+        const float id = d ? 1.0f/d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        for (size_t j = 0; j < sizeof(y->qs); j += 32) {
+            for (size_t m = 0; m < 32; ++m) {
+                uint8_t q = 0;
+                for (size_t n = 0; n < 4; ++n) {
+                    // -1, 0, 1 -> 0, 1, 2
+                    int xi = lroundf(x[m + n*32] * id) + 1;
+                    q += (xi & 3) << (2*n);
+                }
+                y[i].qs[j + m] = q;
+            }
+            x += 4*32;
+        }
+    }
+}
+
+size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    (void)quant_weights; // not used
+    const size_t row_size = ggml_row_size(GGML_TYPE_TQ1_0, n_per_row);
+    quantize_row_tq1_0_ref(src, dst, (int64_t)nrow*n_per_row);
+    return nrow * row_size;
+}
+
+size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    (void)quant_weights; // not used
+    const size_t row_size = ggml_row_size(GGML_TYPE_TQ2_0, n_per_row);
+    quantize_row_tq2_0_ref(src, dst, (int64_t)nrow*n_per_row);
+    return nrow * row_size;
+}
+
+void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243};
+
+    for (int64_t i = 0; i < nb; ++i) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) {
+            for (size_t n = 0; n < 5; ++n) {
+                for (size_t m = 0; m < 32; ++m) {
+                    uint8_t q = x[i].qs[j + m] * pow3[n];
+                    int16_t xi = ((uint16_t) q * 3) >> 8;
+                    *y++ = (float) (xi - 1) * d;
+                }
+            }
+        }
+        for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) {
+            for (size_t n = 0; n < 5; ++n) {
+                for (size_t m = 0; m < 16; ++m) {
+                    uint8_t q = x[i].qs[j + m] * pow3[n];
+                    int16_t xi = ((uint16_t) q * 3) >> 8;
+                    *y++ = (float) (xi - 1) * d;
+                }
+            }
+        }
+
+        for (size_t n = 0; n < 4; ++n) {
+            for (size_t j = 0; j < sizeof(x->qh); ++j) {
+                uint8_t q = x[i].qh[j] * pow3[n];
+                int16_t xi = ((uint16_t) q * 3) >> 8;
+                *y++ = (float) (xi - 1) * d;
+            }
+        }
+    }
+}
+
+void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int64_t i = 0; i < nb; ++i) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (size_t j = 0; j < sizeof(x->qs); j += 32) {
+            for (size_t l = 0; l < 4; ++l) {
+                for (size_t m = 0; m < 32; ++m) {
+                    int8_t q = (x[i].qs[j + m] >> (l*2)) & 3;
+                    *y++ = (float) (q - 1) * d;
+                }
+            }
+        }
+    }
+}
+
+// ====================== "True" 2-bit (de)-quantization
+
+void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    uint32_t aux32[2];
+    const uint8_t * aux8 = (const uint8_t *)aux32;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            memcpy(aux32, x[i].qs + 4*ib32, 2*sizeof(uint32_t));
+            const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]);
+                const uint8_t  signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127];
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+        }
+    }
+}
+
+// ====================== 2.3125 bpw (de)-quantization
+
+void dequantize_row_iq2_xs(const block_iq2_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    float db[2];
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f;
+            db[1] = d * (0.5f + (x[i].scales[ib32] >>  4)) * 0.25f;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (x[i].qs[4*ib32 + l] & 511));
+                const uint8_t  signs = ksigns_iq2xs[x[i].qs[4*ib32 + l] >> 9];
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = db[l/2] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+        }
+    }
+}
+
+// ====================== 2.5625 bpw (de)-quantization
+
+void dequantize_row_iq2_s(const block_iq2_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    float db[2];
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const uint8_t * qs = x[i].qs;
+        const uint8_t * qh = x[i].qh;
+        const uint8_t * signs = qs + QK_K/8;
+
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f;
+            db[1] = d * (0.5f + (x[i].scales[ib32] >>  4)) * 0.25f;
+            for (int l = 0; l < 4; ++l) {
+                const float dl = db[l/2];
+                const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300)));
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+            qs += 4;
+            signs += 4;
+        }
+    }
+}
+
+// ====================== 3.0625 bpw (de)-quantization
+
+void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    uint32_t aux32;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const uint8_t * qs = x[i].qs;
+        const uint8_t * scales_and_signs = qs + QK_K/4;
+
+        for (int ib32 = 0; ib32 < QK_K/32; ++ib32) {
+            memcpy(&aux32, scales_and_signs + 4*ib32, sizeof(uint32_t));
+            const float db = d * (0.5f + (aux32 >> 28)) * 0.5f;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t  signs = ksigns_iq2xs[(aux32 >> 7*l) & 127];
+                const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + qs[2*l+0]);
+                const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + qs[2*l+1]);
+                for (int j = 0; j < 4; ++j) {
+                    y[j+0] = db * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
+                    y[j+4] = db * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+            qs += 8;
+        }
+    }
+}
+
+// ====================== 3.3125 bpw (de)-quantization
+
+void dequantize_row_iq3_s(const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const uint8_t * qs = x[i].qs;
+        const uint8_t * qh = x[i].qh;
+        const uint8_t * signs = x[i].signs;
+
+        for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) {
+            const float db1 = d * (1 + 2*(x[i].scales[ib32/2] & 0xf));
+            const float db2 = d * (1 + 2*(x[i].scales[ib32/2] >>  4));
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[0] << (8-2*l)) & 256)));
+                const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[0] << (7-2*l)) & 256)));
+                for (int j = 0; j < 4; ++j) {
+                    y[j+0] = db1 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f);
+                    y[j+4] = db1 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+            qs += 8;
+            signs += 4;
+            for (int l = 0; l < 4; ++l) {
+                const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[1] << (8-2*l)) & 256)));
+                const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[1] << (7-2*l)) & 256)));
+                for (int j = 0; j < 4; ++j) {
+                    y[j+0] = db2 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f);
+                    y[j+4] = db2 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f);
+                }
+                y += 8;
+            }
+            qh += 2;
+            qs += 8;
+            signs += 4;
+        }
+    }
+}
+
+// ====================== 1.5625 bpw (de)-quantization
+
+void dequantize_row_iq1_s(const block_iq1_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const uint8_t  * qs = x[i].qs;
+        const uint16_t * qh = x[i].qh;
+
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const float dl = d * (2*((qh[ib] >> 12) & 7) + 1);
+            const float delta = qh[ib] & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA;
+            for (int l = 0; l < 4; ++l) {
+                const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8)));
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = dl * (grid[j] + delta);
+                }
+                y += 8;
+            }
+            qs += 4;
+        }
+    }
+}
+
+void dequantize_row_iq1_m(const block_iq1_m * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    float delta[4];
+    uint16_t idx[4];
+
+    iq1m_scale_t scale;
+
+    for (int i = 0; i < nb; i++) {
+
+        const uint16_t * sc = (const uint16_t *)x[i].scales;
+        scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+        const float d = GGML_FP16_TO_FP32(scale.f16);
+
+        const uint8_t * qs = x[i].qs;
+        const uint8_t * qh = x[i].qh;
+
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const float dl1 = d * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1);
+            const float dl2 = d * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1);
+
+            idx[0] = qs[0] | ((qh[0] << 8) & 0x700);
+            idx[1] = qs[1] | ((qh[0] << 4) & 0x700);
+            idx[2] = qs[2] | ((qh[1] << 8) & 0x700);
+            idx[3] = qs[3] | ((qh[1] << 4) & 0x700);
+            delta[0] = qh[0] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA;
+            delta[1] = qh[0] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA;
+            delta[2] = qh[1] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA;
+            delta[3] = qh[1] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA;
+            for (int l = 0; l < 2; ++l) {
+                const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]);
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = dl1 * (grid[j] + delta[l]);
+                }
+                y += 8;
+            }
+            for (int l = 2; l < 4; ++l) {
+                const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]);
+                for (int j = 0; j < 8; ++j) {
+                    y[j] = dl2 * (grid[j] + delta[l]);
+                }
+                y += 8;
+            }
+            qs += 4;
+            qh += 2;
+        }
+    }
+}
+
+void dequantize_row_iq4_nl(const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK4_NL == 0);
+    const int64_t nb = k / QK4_NL;
+
+    for (int i = 0; i < nb; i++) {
+
+        const uint8_t * qs = x[i].qs;
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        for (int j = 0; j < QK4_NL/2; ++j) {
+            y[j+       0] = d * kvalues_iq4nl[qs[j] & 0xf];
+            y[j+QK4_NL/2] = d * kvalues_iq4nl[qs[j] >>  4];
+        }
+        y  += QK4_NL;
+        qs += QK4_NL/2;
+    }
+}
+
+void dequantize_row_iq4_xs(const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        const uint8_t * qs = x[i].qs;
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const int ls = ((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4);
+            const float dl = d * (ls - 32);
+            for (int j = 0; j < 16; ++j) {
+                y[j+ 0] = dl * kvalues_iq4nl[qs[j] & 0xf];
+                y[j+16] = dl * kvalues_iq4nl[qs[j] >>  4];
+            }
+            y  += 32;
+            qs += 16;
+        }
+    }
+}
+
+//===================================== Q8_K ==============================================
+
+void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        float max = 0;
+        float amax = 0;
+        for (int j = 0; j < QK_K; ++j) {
+            float ax = fabsf(x[j]);
+            if (ax > amax) {
+                amax = ax; max = x[j];
+            }
+        }
+        if (!amax) {
+            y[i].d = 0;
+            memset(y[i].qs, 0, QK_K);
+            x += QK_K;
+            continue;
+        }
+        //const float iscale = -128.f/max;
+        // We need this change for IQ2_XXS, else the AVX implementation becomes very awkward
+        const float iscale = -127.f/max;
+        for (int j = 0; j < QK_K; ++j) {
+            int v = nearest_int(iscale*x[j]);
+            y[i].qs[j] = MIN(127, v);
+        }
+        for (int j = 0; j < QK_K/16; ++j) {
+            int sum = 0;
+            for (int ii = 0; ii < 16; ++ii) {
+                sum += y[i].qs[j*16 + ii];
+            }
+            y[i].bsums[j] = sum;
+        }
+        y[i].d = 1/iscale;
+        x += QK_K;
+    }
+}
+
+void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    const int64_t nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+        for (int j = 0; j < QK_K; ++j) {
+            *y++ = x[i].d * x[i].qs[j];
+        }
+    }
+}
+
+// ================================ IQ2 quantization =============================================
+
+typedef struct {
+    uint64_t * grid;
+    int      * map;
+    uint16_t * neighbours;
+} iq2_entry_t;
+
+static iq2_entry_t iq2_data[4] = {
+    {NULL, NULL, NULL},
+    {NULL, NULL, NULL},
+    {NULL, NULL, NULL},
+    {NULL, NULL, NULL},
+};
+
+static inline int iq2_data_index(enum ggml_type type) {
+    GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S);
+    return type == GGML_TYPE_IQ2_XXS ? 0 :
+           type == GGML_TYPE_IQ2_XS  ? 1 :
+           type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 2 : 3;
+}
+
+static inline int iq2_grid_size(enum ggml_type type) {
+    GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S);
+    return type == GGML_TYPE_IQ2_XXS ? 256 :
+           type == GGML_TYPE_IQ2_XS  ? 512 :
+           type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? NGRID_IQ1S : 1024;
+}
+
+static int iq2_compare_func(const void * left, const void * right) {
+    const int * l = (const int *)left;
+    const int * r = (const int *)right;
+    return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0;
+}
+
+void iq2xs_init_impl(enum ggml_type type) {
+    const int gindex = iq2_data_index(type);
+    const int grid_size = iq2_grid_size(type);
+    if (iq2_data[gindex].grid) {
+        return;
+    }
+    static const uint16_t kgrid_2bit_256[256] = {
+            0,     2,     5,     8,    10,    17,    20,    32,    34,    40,    42,    65,    68,    80,    88,    97,
+          100,   128,   130,   138,   162,   257,   260,   272,   277,   320,   388,   408,   512,   514,   546,   642,
+         1025,  1028,  1040,  1057,  1060,  1088,  1090,  1096,  1120,  1153,  1156,  1168,  1188,  1280,  1282,  1288,
+         1312,  1350,  1385,  1408,  1425,  1545,  1552,  1600,  1668,  1700,  2048,  2053,  2056,  2068,  2088,  2113,
+         2116,  2128,  2130,  2184,  2308,  2368,  2562,  2580,  4097,  4100,  4112,  4129,  4160,  4192,  4228,  4240,
+         4245,  4352,  4360,  4384,  4432,  4442,  4480,  4644,  4677,  5120,  5128,  5152,  5157,  5193,  5248,  5400,
+         5474,  5632,  5654,  6145,  6148,  6160,  6208,  6273,  6400,  6405,  6560,  6737,  8192,  8194,  8202,  8260,
+         8289,  8320,  8322,  8489,  8520,  8704,  8706,  9217,  9220,  9232,  9280,  9302,  9472,  9537,  9572,  9872,
+        10248, 10272, 10388, 10820, 16385, 16388, 16400, 16408, 16417, 16420, 16448, 16456, 16470, 16480, 16513, 16516,
+        16528, 16640, 16672, 16737, 16768, 16773, 16897, 16912, 16968, 16982, 17000, 17408, 17416, 17440, 17536, 17561,
+        17682, 17700, 17920, 18433, 18436, 18448, 18496, 18501, 18688, 18776, 18785, 18818, 19013, 19088, 20480, 20488,
+        20497, 20505, 20512, 20608, 20616, 20740, 20802, 20900, 21137, 21648, 21650, 21770, 22017, 22100, 22528, 22545,
+        22553, 22628, 22848, 23048, 24580, 24592, 24640, 24680, 24832, 24917, 25112, 25184, 25600, 25605, 25872, 25874,
+        25988, 26690, 32768, 32770, 32778, 32833, 32898, 33028, 33048, 33088, 33297, 33793, 33796, 33808, 33813, 33856,
+        33888, 34048, 34118, 34196, 34313, 34368, 34400, 34818, 35076, 35345, 36868, 36880, 36900, 36928, 37025, 37142,
+        37248, 37445, 37888, 37922, 37956, 38225, 39041, 39200, 40962, 41040, 41093, 41225, 41472, 42008, 43088, 43268,
+    };
+    static const uint16_t kgrid_2bit_512[512] = {
+            0,     2,     5,     8,    10,    17,    20,    22,    25,    32,    34,    37,    40,    65,    68,    70,
+           73,    80,    82,    85,    88,    97,   100,   128,   130,   133,   136,   145,   148,   153,   160,   257,
+          260,   262,   265,   272,   274,   277,   280,   282,   289,   292,   320,   322,   325,   328,   337,   340,
+          352,   360,   385,   388,   400,   512,   514,   517,   520,   529,   532,   544,   577,   580,   592,   597,
+          640,   650,  1025,  1028,  1030,  1033,  1040,  1042,  1045,  1048,  1057,  1060,  1088,  1090,  1093,  1096,
+         1105,  1108,  1110,  1120,  1153,  1156,  1168,  1280,  1282,  1285,  1288,  1297,  1300,  1312,  1345,  1348,
+         1360,  1377,  1408,  1537,  1540,  1552,  1574,  1600,  1602,  1668,  2048,  2050,  2053,  2056,  2058,  2065,
+         2068,  2080,  2085,  2113,  2116,  2128,  2136,  2176,  2208,  2218,  2305,  2308,  2320,  2368,  2433,  2441,
+         2560,  2592,  2600,  2710,  2720,  4097,  4100,  4102,  4105,  4112,  4114,  4117,  4120,  4129,  4132,  4160,
+         4162,  4165,  4168,  4177,  4180,  4192,  4202,  4225,  4228,  4240,  4352,  4354,  4357,  4360,  4369,  4372,
+         4384,  4417,  4420,  4432,  4480,  4500,  4502,  4609,  4612,  4614,  4624,  4672,  4704,  5120,  5122,  5125,
+         5128,  5137,  5140,  5152,  5185,  5188,  5193,  5200,  5220,  5248,  5377,  5380,  5392,  5440,  5632,  5652,
+         5705,  6145,  6148,  6160,  6162,  6208,  6228,  6278,  6400,  6405,  6502,  6737,  6825,  8192,  8194,  8197,
+         8200,  8202,  8209,  8212,  8224,  8257,  8260,  8272,  8320,  8352,  8449,  8452,  8464,  8512,  8520,  8549,
+         8704,  8738,  8832,  8872,  9217,  9220,  9232,  9257,  9280,  9472,  9537,  9554,  9625,  9729,  9754,  9894,
+        10240, 10248, 10250, 10272, 10325, 10376, 10402, 10600, 10640, 10760, 10784, 10882, 10888, 10890, 16385, 16388,
+        16390, 16393, 16400, 16402, 16405, 16408, 16417, 16420, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16480,
+        16485, 16513, 16516, 16528, 16640, 16642, 16645, 16648, 16657, 16660, 16672, 16705, 16708, 16720, 16768, 16773,
+        16802, 16897, 16900, 16912, 16914, 16937, 16960, 17408, 17410, 17413, 17416, 17425, 17428, 17433, 17440, 17473,
+        17476, 17488, 17536, 17556, 17665, 17668, 17680, 17700, 17728, 17818, 17920, 17930, 17988, 18000, 18433, 18436,
+        18448, 18496, 18501, 18516, 18530, 18688, 18705, 18756, 18768, 18793, 18948, 20480, 20482, 20485, 20488, 20497,
+        20500, 20512, 20520, 20545, 20548, 20560, 20608, 20737, 20740, 20752, 20757, 20800, 20802, 20992, 21060, 21162,
+        21505, 21508, 21520, 21537, 21568, 21600, 21633, 21665, 21760, 21768, 21888, 21896, 22049, 22120, 22177, 22528,
+        22548, 22593, 22608, 22681, 22810, 22848, 22850, 23173, 24577, 24580, 24592, 24640, 24660, 24674, 24710, 24745,
+        24832, 25124, 25162, 25234, 25600, 25622, 25872, 25920, 25925, 26020, 26625, 26730, 26917, 27142, 27220, 27234,
+        32768, 32770, 32773, 32776, 32785, 32788, 32800, 32810, 32833, 32836, 32848, 32896, 32898, 32936, 32938, 33025,
+        33028, 33030, 33040, 33088, 33105, 33113, 33280, 33312, 33408, 33410, 33440, 33448, 33793, 33796, 33808, 33810,
+        33813, 33856, 33888, 33929, 34048, 34116, 34213, 34328, 34410, 34816, 34824, 34853, 34906, 34944, 34946, 34984,
+        35078, 35362, 35456, 35464, 35478, 35496, 36865, 36868, 36880, 36928, 36950, 36996, 37120, 37154, 37220, 37462,
+        37513, 37888, 37893, 37956, 37968, 37976, 38185, 38288, 38290, 38465, 38993, 39078, 39241, 39445, 39520, 40960,
+        40962, 40968, 40970, 40992, 41002, 41120, 41297, 41305, 41382, 41472, 41474, 41480, 41514, 41600, 41632, 42048,
+        42133, 42597, 42648, 43018, 43040, 43042, 43048, 43168, 43176, 43268, 43396, 43398, 43560, 43562, 43665, 43690,
+    };
+    static const uint16_t kgrid_1bit_2048[NGRID_IQ1S] = {
+            0,     2,     5,     8,    10,    17,    21,    32,    34,    40,    42,    69,    81,    84,    86,   101,
+          128,   130,   136,   138,   149,   160,   162,   168,   170,   260,   261,   273,   276,   278,   281,   282,
+          293,   321,   326,   329,   338,   341,   346,   353,   356,   358,   360,   389,   401,   404,   406,   421,
+          512,   514,   520,   522,   533,   544,   546,   552,   554,   581,   593,   601,   612,   617,   640,   642,
+          648,   650,   657,   661,   665,   672,   674,   680,   682,  1041,  1044,  1046,  1061,  1089,  1097,  1109,
+         1114,  1124,  1125,  1169,  1177,  1189,  1281,  1284,  1285,  1286,  1301,  1304,  1306,  1321,  1344,  1349,
+         1354,  1360,  1361,  1364,  1365,  1366,  1369,  1376,  1378,  1381,  1384,  1386,  1409,  1425,  1429,  1432,
+         1434,  1441,  1444,  1445,  1446,  1449,  1556,  1561,  1601,  1604,  1616,  1618,  1621,  1624,  1632,  1633,
+         1638,  1641,  1669,  1681,  1684,  1689,  2048,  2050,  2056,  2058,  2069,  2080,  2082,  2088,  2090,  2117,
+         2129,  2134,  2149,  2176,  2178,  2184,  2186,  2197,  2208,  2210,  2216,  2218,  2309,  2321,  2324,  2329,
+         2340,  2341,  2369,  2384,  2385,  2389,  2401,  2404,  2409,  2449,  2452,  2454,  2457,  2469,  2560,  2562,
+         2568,  2570,  2581,  2592,  2594,  2600,  2602,  2629,  2641,  2649,  2657,  2661,  2688,  2690,  2693,  2696,
+         2698,  2709,  2720,  2722,  2728,  2730,  4112,  4113,  4116,  4121,  4132,  4133,  4161,  4164,  4176,  4181,
+         4184,  4193,  4196,  4197,  4201,  4241,  4244,  4246,  4257,  4261,  4353,  4356,  4358,  4361,  4368,  4370,
+         4373,  4376,  4385,  4388,  4393,  4421,  4426,  4432,  4433,  4434,  4436,  4437,  4438,  4441,  4448,  4453,
+         4484,  4498,  4501,  4513,  4516,  4625,  4628,  4630,  4645,  4672,  4678,  4681,  4690,  4693,  4696,  4698,
+         4708,  4710,  4741,  4753,  4756,  4758,  4773,  5121,  5126,  5129,  5140,  5141,  5144,  5145,  5153,  5158,
+         5185,  5189,  5190,  5192,  5194,  5201,  5204,  5205,  5206,  5209,  5218,  5221,  5224,  5252,  5257,  5264,
+         5268,  5269,  5272,  5273,  5274,  5281,  5284,  5285,  5289,  5378,  5381,  5386,  5393,  5396,  5397,  5398,
+         5401,  5408,  5410,  5413,  5416,  5418,  5441,  5444,  5445,  5446,  5457,  5458,  5460,  5461,  5462,  5465,
+         5466,  5473,  5476,  5477,  5478,  5481,  5504,  5506,  5508,  5509,  5512,  5514,  5520,  5521,  5524,  5525,
+         5526,  5529,  5530,  5536,  5538,  5541,  5633,  5636,  5637,  5638,  5653,  5654,  5656,  5658,  5665,  5670,
+         5696,  5698,  5700,  5701,  5704,  5706,  5713,  5717,  5718,  5720,  5721,  5729,  5732,  5733,  5736,  5737,
+         5738,  5766,  5770,  5778,  5781,  5796,  5801,  6161,  6166,  6181,  6209,  6212,  6214,  6217,  6224,  6229,
+         6232,  6234,  6240,  6241,  6244,  6246,  6249,  6277,  6289,  6292,  6309,  6416,  6418,  6421,  6426,  6433,
+         6437,  6466,  6468,  6469,  6472,  6481,  6484,  6485,  6486,  6489,  6490,  6496,  6501,  6506,  6537,  6545,
+         6546,  6549,  6552,  6561,  6566,  6569,  6665,  6678,  6692,  6694,  6724,  6726,  6729,  6736,  6738,  6741,
+         6744,  6753,  6758,  6761,  6789,  6801,  6806,  6810,  8192,  8194,  8200,  8202,  8213,  8224,  8226,  8229,
+         8232,  8234,  8261,  8273,  8281,  8289,  8293,  8320,  8322,  8328,  8330,  8341,  8352,  8354,  8357,  8360,
+         8362,  8453,  8465,  8468,  8473,  8485,  8514,  8516,  8521,  8533,  8536,  8538,  8545,  8548,  8549,  8550,
+         8581,  8592,  8598,  8601,  8613,  8705,  8712,  8714,  8721,  8725,  8736,  8738,  8744,  8746,  8773,  8785,
+         8790,  8793,  8805,  8833,  8840,  8842,  8849,  8853,  8864,  8866,  8872,  8874,  9221,  9236,  9238,  9241,
+         9253,  9284,  9285,  9286,  9289,  9298,  9301,  9304,  9306,  9318,  9349,  9361,  9364,  9369,  9377,  9381,
+         9481,  9493,  9505,  9513,  9536,  9541,  9544,  9553,  9556,  9557,  9561,  9570,  9573,  9576,  9609,  9616,
+         9620,  9621,  9624,  9626,  9633,  9636,  9638,  9641,  9733,  9744,  9746,  9753,  9765,  9793,  9801,  9813,
+         9824,  9825,  9833,  9860,  9862,  9872,  9882, 10240, 10242, 10248, 10250, 10261, 10272, 10274, 10280, 10282,
+        10309, 10321, 10324, 10341, 10368, 10370, 10376, 10378, 10400, 10402, 10408, 10410, 10505, 10513, 10516, 10521,
+        10533, 10566, 10569, 10578, 10581, 10593, 10596, 10598, 10601, 10629, 10640, 10646, 10649, 10660, 10661, 10752,
+        10754, 10760, 10762, 10784, 10786, 10792, 10794, 10821, 10833, 10838, 10841, 10853, 10880, 10882, 10888, 10890,
+        10901, 10912, 10914, 10920, 10922, 16389, 16401, 16406, 16421, 16457, 16466, 16469, 16472, 16474, 16481, 16484,
+        16486, 16532, 16537, 16545, 16550, 16640, 16641, 16644, 16646, 16649, 16658, 16661, 16662, 16664, 16666, 16673,
+        16678, 16681, 16709, 16712, 16714, 16721, 16724, 16725, 16726, 16729, 16730, 16741, 16744, 16746, 16769, 16772,
+        16774, 16784, 16786, 16789, 16800, 16801, 16802, 16901, 16913, 16916, 16918, 16933, 16961, 16978, 16981, 16986,
+        16996, 17001, 17033, 17044, 17061, 17409, 17429, 17433, 17449, 17477, 17480, 17482, 17489, 17492, 17493, 17494,
+        17505, 17506, 17509, 17512, 17514, 17537, 17542, 17545, 17552, 17554, 17557, 17568, 17569, 17577, 17665, 17666,
+        17669, 17674, 17681, 17684, 17685, 17686, 17689, 17696, 17701, 17706, 17729, 17732, 17733, 17734, 17737, 17744,
+        17745, 17748, 17749, 17750, 17752, 17753, 17761, 17764, 17765, 17766, 17769, 17794, 17796, 17797, 17800, 17809,
+        17812, 17813, 17814, 17817, 17818, 17829, 17832, 17834, 17921, 17925, 17929, 17940, 17941, 17944, 17946, 17953,
+        17956, 17961, 17984, 17986, 17989, 17992, 18000, 18001, 18002, 18005, 18006, 18009, 18018, 18021, 18024, 18049,
+        18053, 18058, 18068, 18069, 18081, 18084, 18086, 18437, 18449, 18453, 18458, 18469, 18498, 18505, 18512, 18517,
+        18520, 18529, 18532, 18534, 18537, 18565, 18577, 18580, 18582, 18585, 18597, 18689, 18693, 18694, 18698, 18704,
+        18708, 18709, 18712, 18721, 18724, 18726, 18752, 18757, 18762, 18769, 18770, 18772, 18773, 18774, 18777, 18784,
+        18786, 18789, 18790, 18794, 18822, 18825, 18834, 18837, 18838, 18840, 18849, 18852, 18854, 18857, 18966, 19012,
+        19014, 19017, 19029, 19032, 19034, 19044, 19049, 19092, 19109, 20481, 20484, 20485, 20486, 20489, 20498, 20501,
+        20506, 20513, 20516, 20521, 20544, 20549, 20552, 20561, 20564, 20565, 20566, 20569, 20581, 20584, 20614, 20617,
+        20629, 20632, 20640, 20641, 20646, 20649, 20741, 20744, 20745, 20746, 20753, 20756, 20757, 20758, 20760, 20761,
+        20768, 20773, 20774, 20776, 20778, 20801, 20804, 20805, 20806, 20809, 20816, 20817, 20818, 20820, 20821, 20822,
+        20824, 20825, 20826, 20833, 20836, 20837, 20838, 20841, 20866, 20869, 20881, 20884, 20885, 20886, 20889, 20896,
+        20901, 20906, 20993, 20998, 21010, 21013, 21018, 21025, 21028, 21058, 21061, 21066, 21073, 21076, 21077, 21078,
+        21081, 21090, 21093, 21125, 21136, 21138, 21141, 21145, 21146, 21156, 21508, 21509, 21521, 21524, 21525, 21526,
+        21528, 21529, 21537, 21541, 21544, 21546, 21569, 21572, 21573, 21574, 21577, 21578, 21584, 21585, 21588, 21589,
+        21590, 21592, 21593, 21594, 21601, 21602, 21604, 21605, 21606, 21609, 21632, 21640, 21642, 21649, 21652, 21653,
+        21654, 21657, 21665, 21668, 21669, 21674, 21761, 21762, 21764, 21765, 21766, 21769, 21776, 21777, 21778, 21780,
+        21781, 21782, 21785, 21786, 21793, 21796, 21797, 21798, 21801, 21824, 21825, 21826, 21828, 21829, 21830, 21832,
+        21833, 21840, 21841, 21842, 21844, 21845, 21846, 21848, 21849, 21850, 21856, 21857, 21860, 21861, 21862, 21864,
+        21865, 21866, 21889, 21892, 21893, 21897, 21898, 21904, 21905, 21908, 21909, 21910, 21912, 21913, 21921, 21924,
+        21925, 21926, 21929, 22016, 22017, 22018, 22020, 22022, 22024, 22025, 22033, 22036, 22037, 22040, 22041, 22048,
+        22049, 22050, 22052, 22053, 22054, 22056, 22057, 22081, 22085, 22086, 22088, 22089, 22090, 22096, 22097, 22098,
+        22100, 22101, 22102, 22104, 22105, 22106, 22113, 22116, 22117, 22121, 22146, 22149, 22150, 22152, 22153, 22154,
+        22161, 22165, 22170, 22178, 22181, 22182, 22184, 22185, 22532, 22533, 22534, 22537, 22544, 22549, 22552, 22561,
+        22570, 22597, 22600, 22602, 22609, 22612, 22613, 22614, 22616, 22617, 22624, 22626, 22628, 22629, 22658, 22665,
+        22672, 22674, 22677, 22680, 22689, 22697, 22785, 22786, 22789, 22794, 22801, 22804, 22805, 22806, 22809, 22821,
+        22849, 22852, 22853, 22854, 22857, 22864, 22865, 22866, 22868, 22869, 22870, 22872, 22873, 22874, 22881, 22884,
+        22885, 22886, 22889, 22913, 22917, 22921, 22929, 22932, 22933, 22934, 22936, 22937, 22949, 23044, 23048, 23061,
+        23066, 23072, 23077, 23078, 23081, 23109, 23112, 23113, 23121, 23125, 23126, 23128, 23129, 23138, 23141, 23144,
+        23146, 23169, 23178, 23186, 23189, 23190, 23192, 23194, 23201, 24581, 24596, 24598, 24601, 24613, 24644, 24656,
+        24661, 24662, 24664, 24666, 24673, 24676, 24678, 24681, 24705, 24726, 24741, 24833, 24836, 24838, 24841, 24850,
+        24853, 24865, 24866, 24870, 24873, 24901, 24905, 24913, 24917, 24918, 24921, 24933, 24934, 24938, 24964, 24970,
+        24978, 24981, 24993, 24998, 25001, 25105, 25110, 25113, 25152, 25153, 25158, 25173, 25174, 25176, 25184, 25221,
+        25233, 25238, 25253, 25617, 25618, 25621, 25622, 25626, 25633, 25638, 25641, 25664, 25666, 25669, 25672, 25674,
+        25681, 25684, 25685, 25686, 25689, 25690, 25696, 25698, 25701, 25732, 25733, 25737, 25744, 25746, 25748, 25749,
+        25750, 25752, 25754, 25761, 25764, 25769, 25861, 25864, 25866, 25873, 25877, 25878, 25881, 25924, 25925, 25926,
+        25929, 25936, 25937, 25940, 25941, 25942, 25945, 25953, 25956, 25957, 25958, 25961, 25990, 25993, 25994, 26001,
+        26005, 26006, 26009, 26010, 26018, 26021, 26022, 26024, 26114, 26121, 26133, 26144, 26150, 26152, 26153, 26176,
+        26181, 26184, 26186, 26193, 26196, 26197, 26198, 26200, 26202, 26208, 26213, 26216, 26240, 26242, 26245, 26250,
+        26260, 26262, 26264, 26265, 26272, 26276, 26278, 26282, 26646, 26649, 26661, 26689, 26706, 26709, 26714, 26721,
+        26729, 26757, 26769, 26776, 26790, 26881, 26884, 26896, 26901, 26913, 26916, 26918, 26921, 26944, 26945, 26949,
+        26950, 26952, 26961, 26964, 26965, 26966, 26969, 26976, 26981, 26986, 27010, 27012, 27018, 27029, 27041, 27044,
+        27045, 27049, 27153, 27158, 27160, 27201, 27204, 27209, 27216, 27221, 27224, 27226, 27236, 27237, 27241, 27270,
+        27284, 27288, 27290, 27302, 32768, 32770, 32776, 32778, 32800, 32802, 32808, 32810, 32837, 32848, 32849, 32852,
+        32854, 32857, 32869, 32896, 32898, 32904, 32906, 32917, 32928, 32930, 32936, 32938, 33029, 33041, 33044, 33046,
+        33049, 33061, 33089, 33092, 33097, 33104, 33106, 33109, 33110, 33112, 33113, 33124, 33126, 33129, 33157, 33161,
+        33172, 33174, 33177, 33189, 33280, 33282, 33288, 33290, 33301, 33312, 33314, 33320, 33322, 33361, 33364, 33369,
+        33381, 33408, 33410, 33416, 33418, 33429, 33440, 33442, 33448, 33450, 33812, 33817, 33857, 33860, 33873, 33877,
+        33882, 33889, 33892, 33897, 33940, 33945, 34049, 34057, 34066, 34069, 34074, 34086, 34089, 34112, 34113, 34117,
+        34120, 34129, 34132, 34133, 34134, 34137, 34138, 34149, 34150, 34152, 34154, 34177, 34180, 34182, 34185, 34192,
+        34194, 34197, 34200, 34214, 34321, 34326, 34329, 34341, 34369, 34372, 34377, 34378, 34384, 34389, 34393, 34394,
+        34401, 34406, 34410, 34437, 34449, 34458, 34468, 34816, 34818, 34824, 34826, 34837, 34848, 34850, 34856, 34858,
+        34881, 34885, 34897, 34900, 34905, 34917, 34921, 34944, 34946, 34952, 34954, 34965, 34976, 34978, 34984, 34986,
+        35077, 35078, 35089, 35092, 35094, 35109, 35137, 35140, 35142, 35145, 35152, 35154, 35157, 35162, 35169, 35172,
+        35205, 35222, 35225, 35237, 35328, 35330, 35336, 35338, 35349, 35360, 35362, 35368, 35370, 35397, 35409, 35412,
+        35414, 35456, 35458, 35464, 35466, 35477, 35488, 35490, 35496, 35498, 36869, 36881, 36886, 36888, 36889, 36901,
+        36929, 36934, 36937, 36949, 36952, 36954, 36969, 36970, 36997, 37009, 37012, 37014, 37017, 37029, 37121, 37124,
+        37126, 37129, 37136, 37141, 37144, 37146, 37153, 37156, 37158, 37161, 37184, 37189, 37200, 37201, 37204, 37205,
+        37206, 37209, 37218, 37221, 37252, 37254, 37266, 37269, 37272, 37281, 37284, 37286, 37289, 37381, 37393, 37396,
+        37401, 37413, 37444, 37446, 37449, 37456, 37458, 37461, 37464, 37478, 37481, 37509, 37524, 37526, 37545, 37889,
+        37892, 37894, 37904, 37909, 37912, 37926, 37952, 37962, 37969, 37972, 37973, 37974, 37976, 37977, 37984, 37985,
+        37986, 37989, 38020, 38022, 38034, 38036, 38037, 38040, 38049, 38057, 38144, 38149, 38152, 38154, 38160, 38161,
+        38164, 38165, 38166, 38169, 38177, 38181, 38185, 38186, 38209, 38212, 38213, 38214, 38217, 38224, 38225, 38226,
+        38228, 38229, 38230, 38232, 38233, 38234, 38241, 38244, 38245, 38246, 38249, 38273, 38277, 38280, 38289, 38290,
+        38292, 38293, 38294, 38297, 38298, 38304, 38306, 38309, 38312, 38314, 38401, 38404, 38416, 38421, 38425, 38432,
+        38438, 38441, 38469, 38472, 38473, 38481, 38482, 38485, 38486, 38489, 38501, 38504, 38530, 38532, 38537, 38538,
+        38546, 38548, 38549, 38564, 38566, 38569, 38917, 38934, 38937, 38949, 38977, 38982, 38992, 38994, 38997, 38998,
+        39002, 39012, 39013, 39045, 39057, 39062, 39065, 39077, 39172, 39174, 39177, 39184, 39186, 39189, 39192, 39194,
+        39200, 39201, 39204, 39206, 39232, 39234, 39237, 39240, 39242, 39249, 39252, 39253, 39254, 39257, 39266, 39269,
+        39270, 39274, 39297, 39300, 39312, 39314, 39317, 39322, 39329, 39334, 39429, 39445, 39461, 39492, 39494, 39497,
+        39504, 39509, 39512, 39521, 39557, 39569, 39572, 39573, 39574, 40960, 40962, 40968, 40970, 40981, 40992, 40994,
+        41000, 41002, 41029, 41041, 41044, 41046, 41049, 41088, 41090, 41096, 41098, 41109, 41120, 41122, 41128, 41130,
+        41221, 41225, 41233, 41236, 41238, 41241, 41242, 41286, 41289, 41297, 41301, 41304, 41306, 41313, 41316, 41349,
+        41360, 41362, 41366, 41369, 41474, 41480, 41482, 41488, 41497, 41506, 41512, 41514, 41541, 41553, 41558, 41561,
+        41573, 41600, 41602, 41608, 41610, 41621, 41632, 41634, 41640, 41642, 42009, 42021, 42049, 42052, 42064, 42068,
+        42069, 42072, 42074, 42081, 42085, 42086, 42088, 42089, 42117, 42246, 42249, 42256, 42258, 42261, 42264, 42278,
+        42281, 42306, 42309, 42321, 42324, 42325, 42326, 42329, 42341, 42346, 42369, 42372, 42373, 42374, 42377, 42386,
+        42389, 42392, 42501, 42513, 42518, 42522, 42529, 42533, 42564, 42566, 42570, 42578, 42581, 42582, 42584, 42592,
+        42594, 42630, 42640, 42645, 42646, 42649, 42657, 42660, 42662, 43008, 43010, 43016, 43018, 43040, 43042, 43048,
+        43050, 43089, 43092, 43094, 43097, 43136, 43138, 43144, 43146, 43157, 43168, 43170, 43176, 43178, 43269, 43284,
+        43289, 43297, 43301, 43329, 43344, 43349, 43354, 43361, 43366, 43369, 43408, 43414, 43520, 43522, 43528, 43530,
+        43552, 43554, 43560, 43562, 43601, 43604, 43606, 43648, 43650, 43656, 43658, 43669, 43680, 43682, 43688, 43690,
+    };
+    static const uint16_t kgrid_2bit_1024[1024] = {
+            0,     2,     5,     8,    10,    17,    20,    22,    25,    32,    34,    37,    40,    65,    68,    70,
+           73,    80,    82,    85,    88,    97,   100,   102,   105,   128,   130,   133,   136,   145,   148,   160,
+          165,   170,   257,   260,   262,   265,   272,   274,   277,   280,   289,   292,   320,   322,   325,   328,
+          337,   340,   342,   345,   352,   357,   360,   385,   388,   400,   402,   405,   417,   420,   512,   514,
+          517,   520,   529,   532,   544,   554,   577,   580,   582,   585,   592,   597,   640,   645,   650,   660,
+          674,  1025,  1028,  1030,  1033,  1040,  1042,  1045,  1048,  1057,  1060,  1062,  1065,  1088,  1090,  1093,
+         1096,  1098,  1105,  1108,  1110,  1113,  1120,  1122,  1125,  1153,  1156,  1158,  1161,  1168,  1173,  1176,
+         1185,  1188,  1280,  1282,  1285,  1288,  1290,  1297,  1300,  1302,  1305,  1312,  1317,  1320,  1345,  1348,
+         1350,  1353,  1360,  1362,  1365,  1368,  1377,  1380,  1408,  1410,  1413,  1416,  1425,  1428,  1440,  1537,
+         1540,  1542,  1545,  1552,  1557,  1600,  1605,  1608,  1617,  1620,  1632,  1665,  1668,  1680,  2048,  2050,
+         2053,  2056,  2065,  2068,  2070,  2073,  2080,  2085,  2090,  2113,  2116,  2118,  2121,  2128,  2130,  2133,
+         2136,  2145,  2148,  2176,  2181,  2196,  2218,  2305,  2308,  2320,  2322,  2325,  2328,  2337,  2368,  2373,
+         2376,  2385,  2388,  2400,  2433,  2448,  2560,  2577,  2580,  2594,  2600,  2602,  2640,  2713,  4097,  4100,
+         4102,  4105,  4112,  4114,  4117,  4120,  4129,  4132,  4134,  4160,  4162,  4165,  4168,  4177,  4180,  4182,
+         4185,  4192,  4194,  4197,  4200,  4225,  4228,  4230,  4240,  4245,  4248,  4257,  4260,  4352,  4354,  4357,
+         4360,  4362,  4369,  4372,  4374,  4377,  4384,  4386,  4389,  4392,  4417,  4420,  4422,  4425,  4432,  4434,
+         4437,  4440,  4449,  4452,  4480,  4482,  4485,  4488,  4497,  4500,  4609,  4612,  4617,  4624,  4629,  4641,
+         4644,  4672,  4677,  4689,  4692,  4737,  4740,  4752,  5120,  5122,  5125,  5128,  5137,  5140,  5142,  5145,
+         5152,  5157,  5160,  5185,  5188,  5190,  5193,  5200,  5202,  5205,  5208,  5217,  5220,  5248,  5250,  5253,
+         5256,  5265,  5268,  5280,  5377,  5380,  5382,  5385,  5392,  5394,  5397,  5400,  5409,  5412,  5440,  5442,
+         5445,  5448,  5457,  5460,  5472,  5505,  5508,  5520,  5632,  5637,  5640,  5649,  5652,  5664,  5697,  5700,
+         5712,  5760,  5802,  6145,  6148,  6150,  6153,  6160,  6165,  6168,  6177,  6208,  6210,  6213,  6216,  6225,
+         6228,  6240,  6273,  6276,  6400,  6402,  6405,  6408,  6417,  6420,  6432,  6465,  6468,  6480,  6505,  6562,
+         6660,  6672,  6720,  6742,  8192,  8194,  8197,  8200,  8209,  8212,  8214,  8217,  8224,  8229,  8234,  8257,
+         8260,  8272,  8274,  8277,  8292,  8320,  8330,  8340,  8362,  8449,  8452,  8464,  8466,  8469,  8481,  8512,
+         8514,  8517,  8529,  8532,  8544,  8577,  8580,  8592,  8704,  8714,  8738,  8744,  8746,  8772,  8784,  8840,
+         8842,  8872,  9217,  9220,  9222,  9225,  9232,  9237,  9240,  9249,  9252,  9280,  9282,  9285,  9288,  9297,
+         9300,  9312,  9345,  9348,  9360,  9472,  9477,  9480,  9489,  9492,  9504,  9537,  9540,  9552,  9574,  9600,
+         9729,  9732,  9744,  9792,  9817, 10240, 10245, 10257, 10260, 10305, 10308, 10320, 10378, 10410, 10497, 10500,
+        10512, 10645, 10762, 10786, 10852, 10888, 10890, 16385, 16388, 16390, 16393, 16400, 16402, 16405, 16408, 16410,
+        16417, 16420, 16422, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16470, 16473, 16480, 16482, 16485, 16513,
+        16516, 16528, 16533, 16536, 16545, 16548, 16640, 16642, 16645, 16648, 16657, 16660, 16662, 16665, 16672, 16674,
+        16677, 16705, 16708, 16710, 16713, 16720, 16722, 16725, 16728, 16737, 16740, 16768, 16770, 16773, 16776, 16785,
+        16788, 16800, 16897, 16900, 16912, 16914, 16917, 16920, 16932, 16960, 16965, 16968, 16977, 16980, 16992, 17025,
+        17028, 17408, 17410, 17413, 17416, 17418, 17425, 17428, 17430, 17433, 17440, 17442, 17445, 17448, 17473, 17476,
+        17478, 17481, 17488, 17490, 17493, 17496, 17505, 17508, 17536, 17538, 17541, 17544, 17553, 17556, 17568, 17665,
+        17668, 17670, 17673, 17680, 17682, 17685, 17688, 17697, 17700, 17728, 17730, 17733, 17736, 17745, 17748, 17760,
+        17770, 17793, 17796, 17808, 17920, 17922, 17925, 17928, 17937, 17940, 17952, 17985, 17988, 18000, 18048, 18085,
+        18433, 18436, 18441, 18448, 18450, 18453, 18456, 18465, 18468, 18496, 18498, 18501, 18504, 18513, 18516, 18528,
+        18564, 18576, 18688, 18690, 18693, 18696, 18705, 18708, 18720, 18753, 18756, 18768, 18816, 18838, 18945, 18948,
+        18960, 19008, 20480, 20482, 20485, 20488, 20497, 20500, 20502, 20505, 20512, 20514, 20517, 20520, 20545, 20548,
+        20550, 20553, 20560, 20562, 20565, 20568, 20577, 20580, 20608, 20610, 20613, 20616, 20625, 20628, 20737, 20740,
+        20742, 20745, 20752, 20754, 20757, 20760, 20769, 20772, 20800, 20802, 20805, 20808, 20817, 20820, 20832, 20865,
+        20868, 20880, 20992, 20997, 21000, 21009, 21012, 21024, 21057, 21060, 21072, 21097, 21120, 21505, 21508, 21510,
+        21513, 21520, 21522, 21525, 21528, 21537, 21540, 21568, 21570, 21573, 21576, 21585, 21588, 21600, 21633, 21636,
+        21648, 21760, 21762, 21765, 21768, 21777, 21780, 21792, 21825, 21828, 21840, 21888, 22017, 22020, 22032, 22054,
+        22080, 22528, 22530, 22533, 22536, 22545, 22548, 22560, 22593, 22596, 22608, 22618, 22656, 22785, 22788, 22800,
+        22848, 23040, 23065, 23173, 23208, 24577, 24580, 24582, 24592, 24594, 24597, 24600, 24609, 24612, 24640, 24645,
+        24648, 24657, 24660, 24672, 24708, 24720, 24832, 24834, 24837, 24840, 24849, 24852, 24864, 24897, 24900, 24912,
+        24960, 24985, 25092, 25104, 25152, 25174, 25249, 25600, 25605, 25608, 25617, 25620, 25632, 25665, 25668, 25680,
+        25728, 25857, 25860, 25872, 25920, 25930, 25960, 26002, 26112, 26260, 26625, 26628, 26640, 26725, 26776, 26880,
+        26922, 27202, 27297, 32768, 32770, 32773, 32776, 32785, 32788, 32793, 32800, 32805, 32833, 32836, 32848, 32850,
+        32853, 32856, 32865, 32896, 32901, 32913, 32916, 33025, 33028, 33033, 33040, 33042, 33045, 33048, 33057, 33060,
+        33088, 33090, 33093, 33096, 33105, 33108, 33153, 33156, 33168, 33193, 33280, 33285, 33290, 33297, 33300, 33345,
+        33348, 33360, 33793, 33796, 33798, 33801, 33808, 33810, 33813, 33816, 33825, 33856, 33858, 33861, 33864, 33873,
+        33876, 33888, 33921, 33924, 33936, 34048, 34050, 34053, 34056, 34065, 34068, 34080, 34113, 34116, 34128, 34176,
+        34186, 34305, 34308, 34320, 34345, 34368, 34816, 34821, 34833, 34836, 34881, 34884, 34896, 34978, 35073, 35076,
+        35136, 35173, 35362, 35416, 35418, 35458, 35490, 36865, 36868, 36873, 36880, 36882, 36885, 36888, 36900, 36928,
+        36930, 36933, 36936, 36945, 36948, 36960, 36993, 36996, 37008, 37120, 37125, 37137, 37140, 37185, 37188, 37200,
+        37210, 37377, 37380, 37392, 37440, 37542, 37888, 37890, 37893, 37896, 37905, 37908, 37920, 37953, 37956, 37968,
+        38016, 38038, 38145, 38148, 38160, 38208, 38296, 38305, 38400, 38470, 38500, 38913, 38916, 38928, 38950, 38976,
+        39081, 39168, 39241, 39250, 39568, 40960, 40965, 40970, 40980, 40994, 41002, 41025, 41028, 41040, 41122, 41130,
+        41280, 41317, 41474, 41482, 41506, 41512, 41514, 41602, 41608, 41610, 41640, 41985, 41988, 42000, 42048, 42121,
+        42148, 42240, 42265, 42577, 43018, 43048, 43170, 43348, 43398, 43528, 43530, 43552, 43554, 43560, 43656, 43690,
+    };
+
+    const int kmap_size = 43692;
+    //const int nwant = type == GGML_TYPE_IQ1_S ? 3 : 2;
+    const int nwant = type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 3 : type == GGML_TYPE_IQ2_S ? 1 : 2;
+    const uint16_t * kgrid = type == GGML_TYPE_IQ2_XXS ? kgrid_2bit_256 :
+                             type == GGML_TYPE_IQ2_XS  ? kgrid_2bit_512 :
+                             type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? kgrid_1bit_2048 : kgrid_2bit_1024;
+    uint64_t * kgrid_q2xs;
+    int      * kmap_q2xs;
+    uint16_t * kneighbors_q2xs;
+
+    //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size);
+    uint64_t * the_grid = (uint64_t *)malloc(grid_size*sizeof(uint64_t));
+    for (int k = 0; k < grid_size; ++k) {
+        int8_t * pos = (int8_t *)(the_grid + k);
+        for (int i = 0; i < 8; ++i) {
+            int l = (kgrid[k] >> 2*i) & 0x3;
+            pos[i] = 2*l + 1;
+        }
+    }
+    kgrid_q2xs = the_grid;
+    iq2_data[gindex].grid = the_grid;
+    kmap_q2xs = (int *)malloc(kmap_size*sizeof(int));
+    iq2_data[gindex].map = kmap_q2xs;
+    for (int i = 0; i < kmap_size; ++i) kmap_q2xs[i] = -1;
+    uint64_t aux64;
+    uint8_t * aux8 = (uint8_t *)&aux64;
+    for (int i = 0; i < grid_size; ++i) {
+        aux64 = kgrid_q2xs[i];
+        uint16_t index = 0;
+        for (int k=0; k<8; ++k) {
+            uint16_t q = (aux8[k] - 1)/2;
+            index |= (q << 2*k);
+        }
+        kmap_q2xs[index] = i;
+    }
+    // The neighbour search runs in three passes:
+    //   1. Parallel: for each i, qsort and count its neighbours into n_per_i,
+    //      and reduce the totals (num_neighbors, num_not_in_map).
+    //   2. Serial: prefix-sum n_per_i into offsets[], so each i has a
+    //      pre-assigned slice of kneighbors_q2xs to write into.
+    //   3. Parallel: redo the qsort and write each i's neighbour list at
+    //      offsets[i].
+    int * n_per_i = (int *)malloc(kmap_size*sizeof(int));
+    GGML_ASSERT(n_per_i);
+    int num_neighbors = 0, num_not_in_map = 0;
+#ifdef GGML_USE_OPENMP
+    #pragma omp parallel reduction(+:num_neighbors,num_not_in_map)
+#endif
+    {
+        int * dist2 = (int *)malloc(2*grid_size*sizeof(int));
+        GGML_ASSERT(dist2);
+        int8_t pos[8];
+        int i;
+#ifdef GGML_USE_OPENMP
+        #pragma omp for schedule(dynamic, 64)
+#endif
+        for (i = 0; i < kmap_size; ++i) {
+            if (kmap_q2xs[i] >= 0) {
+                n_per_i[i] = 0;
+                continue;
+            }
+            ++num_not_in_map;
+            for (int k = 0; k < 8; ++k) {
+                int l = (i >> 2*k) & 0x3;
+                pos[k] = 2*l + 1;
+            }
+            for (int j = 0; j < grid_size; ++j) {
+                const int8_t * pg = (const int8_t *)(kgrid_q2xs + j);
+                int d2 = 0;
+                for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]);
+                dist2[2*j+0] = d2;
+                dist2[2*j+1] = j;
+            }
+            qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func);
+            int n = 0; int d2 = dist2[0];
+            int nhave = 1;
+            for (int j = 0; j < grid_size; ++j) {
+                if (dist2[2*j] > d2) {
+                    if (nhave == nwant) break;
+                    d2 = dist2[2*j];
+                    ++nhave;
+                }
+                ++n;
+            }
+            n_per_i[i] = n;
+            num_neighbors += n;
+        }
+        free(dist2);
+    }
+    //printf("%s: %d neighbours in total\n", __func__, num_neighbors);
+    kneighbors_q2xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t));
+    iq2_data[gindex].neighbours = kneighbors_q2xs;
+
+    int * offsets = (int *)malloc(kmap_size*sizeof(int));
+    GGML_ASSERT(offsets);
+    int counter = 0;
+    for (int i = 0; i < kmap_size; ++i) {
+        if (kmap_q2xs[i] >= 0) {
+            offsets[i] = -1;
+            continue;
+        }
+        offsets[i] = counter;
+        counter += 1 + n_per_i[i];
+    }
+
+#ifdef GGML_USE_OPENMP
+    #pragma omp parallel
+#endif
+    {
+        int * dist2 = (int *)malloc(2*grid_size*sizeof(int));
+        GGML_ASSERT(dist2);
+        int8_t pos[8];
+        int i;
+#ifdef GGML_USE_OPENMP
+        #pragma omp for schedule(dynamic, 64)
+#endif
+        for (i = 0; i < kmap_size; ++i) {
+            if (kmap_q2xs[i] >= 0) continue;
+            for (int k = 0; k < 8; ++k) {
+                int l = (i >> 2*k) & 0x3;
+                pos[k] = 2*l + 1;
+            }
+            for (int j = 0; j < grid_size; ++j) {
+                const int8_t * pg = (const int8_t *)(kgrid_q2xs + j);
+                int d2 = 0;
+                for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]);
+                dist2[2*j+0] = d2;
+                dist2[2*j+1] = j;
+            }
+            qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func);
+            int local_counter = offsets[i];
+            kmap_q2xs[i] = -(local_counter + 1);
+            int d2 = dist2[0];
+            uint16_t * start = &kneighbors_q2xs[local_counter++];
+            int n = 0, nhave = 1;
+            for (int j = 0; j < grid_size; ++j) {
+                if (dist2[2*j] > d2) {
+                    if (nhave == nwant) break;
+                    d2 = dist2[2*j];
+                    ++nhave;
+                }
+                kneighbors_q2xs[local_counter++] = dist2[2*j+1];
+                ++n;
+            }
+            *start = n;
+        }
+        free(dist2);
+    }
+    free(offsets);
+    free(n_per_i);
+}
+
+void iq2xs_free_impl(enum ggml_type type) {
+    GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S);
+    const int gindex = iq2_data_index(type);
+    if (iq2_data[gindex].grid) {
+        free(iq2_data[gindex].grid);       iq2_data[gindex].grid = NULL;
+        free(iq2_data[gindex].map);        iq2_data[gindex].map  = NULL;
+        free(iq2_data[gindex].neighbours); iq2_data[gindex].neighbours = NULL;
+    }
+}
+
+static int iq2_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid,
+        const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) {
+    int num_neighbors = neighbours[0];
+    GGML_ASSERT(num_neighbors > 0);
+    float best_d2 = FLT_MAX;
+    int grid_index = -1;
+    for (int j = 1; j <= num_neighbors; ++j) {
+        const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+        float d2 = 0;
+        for (int i = 0; i < 8; ++i) {
+            float q = pg[i];
+            float diff = scale*q - xval[i];
+            d2 += weight[i]*diff*diff;
+        }
+        if (d2 < best_d2) {
+            best_d2 = d2; grid_index = neighbours[j];
+        }
+    }
+    GGML_ASSERT(grid_index >= 0);
+    const int8_t * pg = (const int8_t *)(grid + grid_index);
+    for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2;
+    return grid_index;
+}
+
+static void quantize_row_iq2_xxs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) {
+
+    const int gindex = iq2_data_index(GGML_TYPE_IQ2_XXS);
+
+    const uint64_t * kgrid_q2xs      = iq2_data[gindex].grid;
+    const int      * kmap_q2xs       = iq2_data[gindex].map;
+    const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours;
+
+    GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kgrid_q2xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kmap_q2xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    const int kMaxQ = 3;
+
+    const int64_t nbl = n/QK_K;
+
+    block_iq2_xxs * y = vy;
+
+    float scales[QK_K/32];
+    float weight[32];
+    float xval[32];
+    int8_t L[32];
+    int8_t Laux[32];
+    float  waux[32];
+    uint8_t block_signs[4];
+    uint32_t q2[2*(QK_K/32)];
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        y[ibl].d = GGML_FP32_TO_FP16(0.f);
+        memset(q2, 0, QK_K/4);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const float * xb = xbl + 32*ib;
+            const float * qw = quant_weights + QK_K*ibl + 32*ib;
+            for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]);
+            for (int k = 0; k < 4; ++k) {
+                int nflip = 0;
+                uint8_t s = 0;
+                for (int i = 0; i < 8; ++i) {
+                    if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i];
+                    else {
+                        xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i);
+                    }
+                }
+                if (nflip%2) {
+                    int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin];
+                    for (int i = 1; i < 8; ++i) {
+                        float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i];
+                        if (ax < min) {
+                            min = ax; imin = i;
+                        }
+                    }
+                    xval[8*k+imin] = -xval[8*k+imin];
+                    s ^= (1 << imin);
+                }
+                block_signs[k] = s & 127;
+            }
+            float max = xval[0];
+            for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]);
+            if (max < GROUP_MAX_EPS) {
+                scales[ib] = 0;
+                memset(L, 0, 32);
+                continue;
+            }
+            float scale = make_qp_quants(32, kMaxQ+1, xval, (uint8_t*)L, weight);
+            float eff_max = scale*kMaxQ;
+            if (eff_max <= 0) {
+                scales[ib] = 0;
+                memset(L, 0, 32);
+                continue;
+            }
+            float best = 0;
+            for (int is = -6; is <= 6; ++is) {
+                float id = (2*kMaxQ-1+is*0.1f)/eff_max;
+                float this_scale = 1/id;
+                for (int k = 0; k < 4; ++k) {
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l));
+                    }
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i);
+                    int grid_index = kmap_q2xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 32; ++i) {
+                    float w = weight[i];
+                    float q = 2*Laux[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                    scale = sumqx/sumq2; best = scale*sumqx;
+                    memcpy(L, Laux, 32);
+                }
+            }
+            if (scale > 0) {
+                float id = 1/scale;
+                for (int k = 0; k < 4; ++k) {
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        l = MAX(0, MIN(kMaxQ-1, l));
+                        u |= (l << 2*i);
+                    }
+                    int grid_index = kmap_q2xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k);
+                    }
+                    const int8_t * pg = (const int8_t *)(kgrid_q2xs + grid_index);
+                    for (int i = 0; i < 8; ++i) L[8*k+i] = (pg[i] - 1)/2;
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 32; ++i) {
+                    float w = weight[i];
+                    float q = 2*L[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0) scale = sumqx/sumq2;
+            }
+            if (scale < 0) {
+                // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale)
+                // and correspondingly flip quant signs.
+                scale = -scale;
+                for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127;
+            }
+            for (int k = 0; k < 4; ++k) {
+                uint16_t u = 0;
+                for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i);
+                int grid_index = kmap_q2xs[u];
+                if (grid_index < 0) {
+                    printf("Oops: found point %u not on grid:", u);
+                    for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]);
+                    printf("\n");
+                    GGML_ABORT("fatal error");
+                }
+                q2[2*ib+0] |= ((uint32_t) grid_index << 8*k);
+                q2[2*ib+1] |= (block_signs[k] << 7*k);
+            }
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            memset(y[ibl].qs, 0, QK_K/4);
+            continue;
+        }
+
+        float d = max_scale/31;
+        y[ibl].d = GGML_FP32_TO_FP16(d);
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib]-1));
+            l = MAX(0, MIN(15, l));
+            q2[2*ib+1] |= ((uint32_t)l << 28);
+        }
+        memcpy(y[ibl].qs, q2, QK_K/4);
+    }
+}
+
+static void quantize_row_iq2_xs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) {
+
+    const int gindex = iq2_data_index(GGML_TYPE_IQ2_XS);
+
+    const uint64_t * kgrid_q2xs      = iq2_data[gindex].grid;
+    const int      * kmap_q2xs       = iq2_data[gindex].map;
+    const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours;
+
+    GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kmap_q2xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kgrid_q2xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    const int kMaxQ = 3;
+
+    const int64_t nbl = n/QK_K;
+
+    block_iq2_xs * y = vy;
+
+    float scales[QK_K/16];
+    float weight[16];
+    float xval[16];
+    int8_t L[16];
+    int8_t Laux[16];
+    float  waux[16];
+    bool   is_on_grid[2];
+    bool   is_on_grid_aux[2];
+    uint8_t block_signs[2];
+    uint16_t q2[2*(QK_K/16)];
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        y[ibl].d = GGML_FP32_TO_FP16(0.f);
+        memset(q2, 0, QK_K/4);
+        memset(y[ibl].scales, 0, QK_K/32);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            const float * xb = xbl + 16*ib;
+            const float * qw = quant_weights + QK_K*ibl + 16*ib;
+            for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]);
+            for (int k = 0; k < 2; ++k) {
+                int nflip = 0;
+                uint8_t s = 0;
+                for (int i = 0; i < 8; ++i) {
+                    if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i];
+                    else {
+                        xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i);
+                    }
+                }
+                if (nflip%2) {
+                    int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin];
+                    for (int i = 1; i < 8; ++i) {
+                        float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i];
+                        if (ax < min) {
+                            min = ax; imin = i;
+                        }
+                    }
+                    xval[8*k+imin] = -xval[8*k+imin];
+                    s ^= (1 << imin);
+                }
+                block_signs[k] = s & 127;
+            }
+            float max = xval[0];
+            for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]);
+            memset(L, 0, 16);
+            if (max < GROUP_MAX_EPS) {
+                scales[ib] = 0;
+                continue;
+            }
+            float best = 0;
+            float scale = max/(2*kMaxQ-1);
+            is_on_grid[0] = is_on_grid[1] = true;
+            for (int is = -9; is <= 9; ++is) {
+                float id = (2*kMaxQ-1+is*0.1f)/max;
+                float this_scale = 1/id;
+                for (int k = 0; k < 2; ++k) {
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l));
+                    }
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i);
+                    int grid_index = kmap_q2xs[u];
+                    is_on_grid_aux[k] = true;
+                    if (grid_index < 0) {
+                        is_on_grid_aux[k] = false;
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 16; ++i) {
+                    float w = weight[i];
+                    float q = 2*Laux[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                    scale = sumqx/sumq2; best = scale*sumqx;
+                    for (int i = 0; i < 16; ++i) L[i] = Laux[i];
+                    for (int k = 0; k <  2; ++k) is_on_grid[k] = is_on_grid_aux[k];
+                }
+            }
+            int n_not_ongrid = 0;
+            for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid;
+            if (n_not_ongrid > 0 && scale > 0) {
+                float id = 1/scale;
+                for (int k = 0; k < 2; ++k) {
+                    if (is_on_grid[k]) continue;
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        l = MAX(0, MIN(kMaxQ-1, l));
+                        u |= (l << 2*i);
+                        L[8*k + i] = l;
+                    }
+                    int grid_index = kmap_q2xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 16; ++i) {
+                    float w = weight[i];
+                    float q = 2*L[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0) scale = sumqx/sumq2;
+            }
+            if (scale < 0) {
+                scale = -scale;
+                for (int k = 0; k < 2; ++k) block_signs[k] = (~block_signs[k]) & 127;
+            }
+            for (int k = 0; k < 2; ++k) {
+                uint16_t u = 0;
+                for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i);
+                int grid_index = kmap_q2xs[u];
+                if (grid_index < 0) {
+                    printf("Oops: found point %u not on grid:", u);
+                    for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]);
+                    printf("\n");
+                    GGML_ABORT("fatal error");
+                }
+                q2[2*ib+k] = grid_index | (block_signs[k] << 9);
+            }
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            memset(y[ibl].qs, 0, QK_K/4);
+            continue;
+        }
+
+        float d = max_scale/31;
+        y[ibl].d = GGML_FP32_TO_FP16(d);
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib]-1));
+            l = MAX(0, MIN(15, l));
+            if (ib%2 == 0) y[ibl].scales[ib/2] = l;
+            else y[ibl].scales[ib/2] |= (l << 4);
+        }
+        memcpy(y[ibl].qs, q2, QK_K/4);
+
+    }
+}
+
+size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq2_xxs_impl(src, qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq2_xxs);
+    }
+    return nrow * nblock * sizeof(block_iq2_xxs);
+}
+
+size_t quantize_iq2_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq2_xs_impl(src, qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq2_xs);
+    }
+    return nrow * nblock * sizeof(block_iq2_xs);
+}
+
+//
+// ============================================= 3-bit using D4 lattice
+//
+
+typedef struct {
+    uint32_t * grid;
+    int      * map;
+    uint16_t * neighbours;
+} iq3_entry_t;
+
+static iq3_entry_t iq3_data[2] = {
+    {NULL, NULL, NULL},
+    {NULL, NULL, NULL},
+};
+
+static inline int iq3_data_index(int grid_size) {
+    (void)grid_size;
+    GGML_ASSERT(grid_size == 256 || grid_size == 512);
+    return grid_size == 256 ? 0 : 1;
+}
+
+static int iq3_compare_func(const void * left, const void * right) {
+    const int * l = (const int *)left;
+    const int * r = (const int *)right;
+    return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0;
+}
+
+void iq3xs_init_impl(int grid_size) {
+    const int gindex = iq3_data_index(grid_size);
+    if (iq3_data[gindex].grid) {
+        return;
+    }
+    static const uint16_t kgrid_256[256] = {
+            0,     2,     4,     9,    11,    15,    16,    18,    25,    34,    59,    61,    65,    67,    72,    74,
+           81,    85,    88,    90,    97,   108,   120,   128,   130,   132,   137,   144,   146,   153,   155,   159,
+          169,   175,   189,   193,   199,   200,   202,   213,   248,   267,   287,   292,   303,   315,   317,   321,
+          327,   346,   362,   413,   436,   456,   460,   462,   483,   497,   513,   515,   520,   522,   529,   531,
+          536,   538,   540,   551,   552,   576,   578,   585,   592,   594,   641,   643,   648,   650,   657,   664,
+          698,   704,   706,   720,   729,   742,   758,   769,   773,   808,   848,   852,   870,   889,   901,   978,
+          992,  1024,  1026,  1033,  1035,  1040,  1042,  1046,  1049,  1058,  1089,  1091,  1093,  1096,  1098,  1105,
+         1112,  1139,  1143,  1144,  1152,  1154,  1161,  1167,  1168,  1170,  1183,  1184,  1197,  1217,  1224,  1228,
+         1272,  1276,  1309,  1323,  1347,  1367,  1377,  1404,  1473,  1475,  1486,  1509,  1537,  1544,  1546,  1553,
+         1555,  1576,  1589,  1594,  1600,  1602,  1616,  1625,  1636,  1638,  1665,  1667,  1672,  1685,  1706,  1722,
+         1737,  1755,  1816,  1831,  1850,  1856,  1862,  1874,  1901,  1932,  1950,  1971,  2011,  2032,  2052,  2063,
+         2077,  2079,  2091,  2095,  2172,  2192,  2207,  2208,  2224,  2230,  2247,  2277,  2308,  2345,  2356,  2389,
+         2403,  2424,  2501,  2504,  2506,  2520,  2570,  2593,  2616,  2624,  2630,  2646,  2669,  2700,  2714,  2746,
+         2754,  2795,  2824,  2835,  2839,  2874,  2882,  2905,  2984,  3028,  3042,  3092,  3108,  3110,  3124,  3153,
+         3185,  3215,  3252,  3288,  3294,  3364,  3397,  3434,  3483,  3523,  3537,  3587,  3589,  3591,  3592,  3610,
+         3626,  3670,  3680,  3722,  3749,  3754,  3776,  3789,  3803,  3824,  3857,  3873,  3904,  3906,  3924,  3992,
+    };
+    static const uint16_t kgrid_512[512] = {
+            0,     1,     2,     5,     7,     8,     9,    10,    12,    14,    16,    17,    21,    27,    32,    34,
+           37,    39,    41,    43,    48,    50,    57,    60,    63,    64,    65,    66,    68,    72,    73,    77,
+           80,    83,    87,    89,    93,   100,   113,   117,   122,   128,   129,   133,   135,   136,   139,   142,
+          145,   149,   152,   156,   162,   165,   167,   169,   171,   184,   187,   195,   201,   205,   208,   210,
+          217,   219,   222,   228,   232,   234,   247,   249,   253,   256,   267,   271,   273,   276,   282,   288,
+          291,   297,   312,   322,   324,   336,   338,   342,   347,   353,   357,   359,   374,   379,   390,   393,
+          395,   409,   426,   441,   448,   450,   452,   464,   466,   470,   475,   488,   492,   512,   513,   514,
+          516,   520,   521,   523,   525,   527,   528,   530,   537,   540,   542,   556,   558,   561,   570,   576,
+          577,   579,   582,   584,   588,   593,   600,   603,   609,   616,   618,   632,   638,   640,   650,   653,
+          655,   656,   660,   666,   672,   675,   685,   688,   698,   705,   708,   711,   712,   715,   721,   727,
+          728,   732,   737,   754,   760,   771,   773,   778,   780,   793,   795,   802,   806,   808,   812,   833,
+          840,   843,   849,   856,   858,   873,   912,   916,   919,   932,   934,   961,   963,   968,   970,   977,
+          989,   993,  1010,  1016,  1024,  1025,  1027,  1029,  1031,  1032,  1034,  1036,  1038,  1041,  1043,  1047,
+         1048,  1050,  1057,  1059,  1061,  1064,  1066,  1079,  1080,  1083,  1085,  1088,  1090,  1096,  1099,  1103,
+         1106,  1109,  1113,  1116,  1122,  1129,  1153,  1156,  1159,  1169,  1171,  1176,  1183,  1185,  1195,  1199,
+         1209,  1212,  1216,  1218,  1221,  1225,  1234,  1236,  1241,  1243,  1250,  1256,  1270,  1281,  1287,  1296,
+         1299,  1306,  1309,  1313,  1338,  1341,  1348,  1353,  1362,  1375,  1376,  1387,  1400,  1408,  1410,  1415,
+         1425,  1453,  1457,  1477,  1481,  1494,  1496,  1507,  1512,  1538,  1545,  1547,  1549,  1551,  1554,  1561,
+         1563,  1565,  1570,  1572,  1575,  1577,  1587,  1593,  1601,  1603,  1605,  1612,  1617,  1619,  1632,  1648,
+         1658,  1662,  1664,  1674,  1680,  1690,  1692,  1704,  1729,  1736,  1740,  1745,  1747,  1751,  1752,  1761,
+         1763,  1767,  1773,  1787,  1795,  1801,  1806,  1810,  1817,  1834,  1840,  1844,  1857,  1864,  1866,  1877,
+         1882,  1892,  1902,  1915,  1934,  1953,  1985,  1987,  2000,  2002,  2013,  2048,  2052,  2058,  2064,  2068,
+         2071,  2074,  2081,  2088,  2104,  2114,  2119,  2121,  2123,  2130,  2136,  2141,  2147,  2153,  2157,  2177,
+         2179,  2184,  2189,  2193,  2203,  2208,  2223,  2226,  2232,  2244,  2249,  2251,  2256,  2258,  2265,  2269,
+         2304,  2306,  2324,  2335,  2336,  2361,  2373,  2375,  2385,  2418,  2443,  2460,  2480,  2504,  2509,  2520,
+         2531,  2537,  2562,  2568,  2572,  2578,  2592,  2596,  2599,  2602,  2614,  2620,  2625,  2627,  2629,  2634,
+         2641,  2650,  2682,  2688,  2697,  2707,  2712,  2718,  2731,  2754,  2759,  2760,  2775,  2788,  2793,  2805,
+         2811,  2817,  2820,  2832,  2842,  2854,  2890,  2902,  2921,  2923,  2978,  3010,  3012,  3026,  3081,  3083,
+         3085,  3097,  3099,  3120,  3136,  3152,  3159,  3188,  3210,  3228,  3234,  3245,  3250,  3256,  3264,  3276,
+         3281,  3296,  3349,  3363,  3378,  3392,  3395,  3420,  3440,  3461,  3488,  3529,  3531,  3584,  3588,  3591,
+         3600,  3602,  3614,  3616,  3628,  3634,  3650,  3657,  3668,  3683,  3685,  3713,  3716,  3720,  3726,  3729,
+         3736,  3753,  3778,  3802,  3805,  3819,  3841,  3845,  3851,  3856,  3880,  3922,  3938,  3970,  3993,  4032,
+    };
+
+    const int kmap_size = 4096;
+    const int nwant = grid_size == 256 ? 2 : 3;
+    const uint16_t * kgrid = grid_size == 256 ? kgrid_256 : kgrid_512;
+    uint32_t * kgrid_q3xs;
+    int      * kmap_q3xs;
+    uint16_t * kneighbors_q3xs;
+
+    //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size);
+    uint32_t * the_grid = (uint32_t *)malloc(grid_size*sizeof(uint32_t));
+    for (int k = 0; k < grid_size; ++k) {
+        int8_t * pos = (int8_t *)(the_grid + k);
+        for (int i = 0; i < 4; ++i) {
+            int l = (kgrid[k] >> 3*i) & 0x7;
+            pos[i] = 2*l + 1;
+        }
+    }
+    kgrid_q3xs = the_grid;
+    iq3_data[gindex].grid = the_grid;
+    kmap_q3xs = (int *)malloc(kmap_size*sizeof(int));
+    iq3_data[gindex].map = kmap_q3xs;
+    for (int i = 0; i < kmap_size; ++i) kmap_q3xs[i] = -1;
+    uint32_t aux32;
+    uint8_t * aux8 = (uint8_t *)&aux32;
+    for (int i = 0; i < grid_size; ++i) {
+        aux32 = kgrid_q3xs[i];
+        uint16_t index = 0;
+        for (int k=0; k<4; ++k) {
+            uint16_t q = (aux8[k] - 1)/2;
+            index |= (q << 3*k);
+        }
+        kmap_q3xs[index] = i;
+    }
+    // See explanation of parallelism in iq2xs_init_impl
+    int * n_per_i = (int *)malloc(kmap_size*sizeof(int));
+    GGML_ASSERT(n_per_i);
+    int num_neighbors = 0, num_not_in_map = 0;
+#ifdef GGML_USE_OPENMP
+    #pragma omp parallel reduction(+:num_neighbors,num_not_in_map)
+#endif
+    {
+        int * dist2 = (int *)malloc(2*grid_size*sizeof(int));
+        GGML_ASSERT(dist2);
+        int8_t pos[4];
+        int i;
+#ifdef GGML_USE_OPENMP
+        #pragma omp for schedule(dynamic, 64)
+#endif
+        for (i = 0; i < kmap_size; ++i) {
+            if (kmap_q3xs[i] >= 0) {
+                n_per_i[i] = 0;
+                continue;
+            }
+            ++num_not_in_map;
+            for (int k = 0; k < 4; ++k) {
+                int l = (i >> 3*k) & 0x7;
+                pos[k] = 2*l + 1;
+            }
+            for (int j = 0; j < grid_size; ++j) {
+                const int8_t * pg = (const int8_t *)(kgrid_q3xs + j);
+                int d2 = 0;
+                for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]);
+                dist2[2*j+0] = d2;
+                dist2[2*j+1] = j;
+            }
+            qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func);
+            int n = 0; int d2 = dist2[0];
+            int nhave = 1;
+            for (int j = 0; j < grid_size; ++j) {
+                if (dist2[2*j] > d2) {
+                    if (nhave == nwant) break;
+                    d2 = dist2[2*j];
+                    ++nhave;
+                }
+                ++n;
+            }
+            n_per_i[i] = n;
+            num_neighbors += n;
+        }
+        free(dist2);
+    }
+    //printf("%s: %d neighbours in total\n", __func__, num_neighbors);
+    kneighbors_q3xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t));
+    iq3_data[gindex].neighbours = kneighbors_q3xs;
+
+    int * offsets = (int *)malloc(kmap_size*sizeof(int));
+    GGML_ASSERT(offsets);
+    int counter = 0;
+    for (int i = 0; i < kmap_size; ++i) {
+        if (kmap_q3xs[i] >= 0) {
+            offsets[i] = -1;
+            continue;
+        }
+        offsets[i] = counter;
+        counter += 1 + n_per_i[i];
+    }
+
+#ifdef GGML_USE_OPENMP
+    #pragma omp parallel
+#endif
+    {
+        int * dist2 = (int *)malloc(2*grid_size*sizeof(int));
+        GGML_ASSERT(dist2);
+        int8_t pos[4];
+        int i;
+#ifdef GGML_USE_OPENMP
+        #pragma omp for schedule(dynamic, 64)
+#endif
+        for (i = 0; i < kmap_size; ++i) {
+            if (kmap_q3xs[i] >= 0) continue;
+            for (int k = 0; k < 4; ++k) {
+                int l = (i >> 3*k) & 0x7;
+                pos[k] = 2*l + 1;
+            }
+            for (int j = 0; j < grid_size; ++j) {
+                const int8_t * pg = (const int8_t *)(kgrid_q3xs + j);
+                int d2 = 0;
+                for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]);
+                dist2[2*j+0] = d2;
+                dist2[2*j+1] = j;
+            }
+            qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func);
+            int local_counter = offsets[i];
+            kmap_q3xs[i] = -(local_counter + 1);
+            int d2 = dist2[0];
+            uint16_t * start = &kneighbors_q3xs[local_counter++];
+            int n = 0, nhave = 1;
+            for (int j = 0; j < grid_size; ++j) {
+                if (dist2[2*j] > d2) {
+                    if (nhave == nwant) break;
+                    d2 = dist2[2*j];
+                    ++nhave;
+                }
+                kneighbors_q3xs[local_counter++] = dist2[2*j+1];
+                ++n;
+            }
+            *start = n;
+        }
+        free(dist2);
+    }
+    free(offsets);
+    free(n_per_i);
+}
+
+void iq3xs_free_impl(int grid_size) {
+    GGML_ASSERT(grid_size == 256 || grid_size == 512);
+    const int gindex = iq3_data_index(grid_size);
+    if (iq3_data[gindex].grid) {
+        free(iq3_data[gindex].grid);       iq3_data[gindex].grid = NULL;
+        free(iq3_data[gindex].map);        iq3_data[gindex].map  = NULL;
+        free(iq3_data[gindex].neighbours); iq3_data[gindex].neighbours = NULL;
+    }
+}
+
+static int iq3_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint32_t * GGML_RESTRICT grid,
+        const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) {
+    int num_neighbors = neighbours[0];
+    GGML_ASSERT(num_neighbors > 0);
+    float best_d2 = FLT_MAX;
+    int grid_index = -1;
+    for (int j = 1; j <= num_neighbors; ++j) {
+        const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+        float d2 = 0;
+        for (int i = 0; i < 4; ++i) {
+            float q = pg[i];
+            float diff = scale*q - xval[i];
+            d2 += weight[i]*diff*diff;
+        }
+        if (d2 < best_d2) {
+            best_d2 = d2; grid_index = neighbours[j];
+        }
+    }
+    GGML_ASSERT(grid_index >= 0);
+    const int8_t * pg = (const int8_t *)(grid + grid_index);
+    for (int i = 0; i < 4; ++i) L[i] = (pg[i] - 1)/2;
+    return grid_index;
+}
+
+static void quantize_row_iq3_xxs_impl(int grid_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n,
+        const float * GGML_RESTRICT quant_weights) {
+
+    const int gindex = iq3_data_index(grid_size);
+
+    const uint32_t * kgrid_q3xs      = iq3_data[gindex].grid;
+    const int      * kmap_q3xs       = iq3_data[gindex].map;
+    const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours;
+
+    //GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kgrid_q3xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kmap_q3xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    const int kMaxQ = 8;
+
+    const int64_t nbl = n/QK_K;
+
+    ggml_fp16_t * dh;
+    uint8_t * qs;
+    int block_size;
+    if (grid_size == 256) {
+        block_iq3_xxs * y = vy;
+        dh = &y->d;
+        qs = y->qs;
+        block_size = sizeof(block_iq3_xxs);
+    } else {
+        block_iq3_s * y = vy;
+        dh = &y->d;
+        qs = y->qs;
+        block_size = sizeof(block_iq3_s);
+    }
+    int quant_size = block_size - sizeof(ggml_fp16_t);
+
+    float scales[QK_K/32];
+    float weight[32];
+    float xval[32];
+    int8_t L[32];
+    int8_t Laux[32];
+    float  waux[32];
+    bool   is_on_grid[8];
+    bool   is_on_grid_aux[8];
+    uint8_t block_signs[8];
+    uint8_t q3[3*(QK_K/8)+QK_K/32];
+    uint32_t * scales_and_signs = (uint32_t *)(q3 + QK_K/4);
+    uint8_t  * qh = q3 + 3*(QK_K/8);
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        dh[0] = GGML_FP32_TO_FP16(0.f);
+        memset(q3, 0, 3*QK_K/8+QK_K/32);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            const float * xb = xbl + 32*ib;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*ibl + 32*ib;
+                for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            } else {
+                for (int i = 0; i < 32; ++i) weight[i] = xb[i]*xb[i];
+            }
+            for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]);
+            for (int k = 0; k < 4; ++k) {
+                int nflip = 0;
+                uint8_t s = 0;
+                for (int i = 0; i < 8; ++i) {
+                    if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i];
+                    else {
+                        xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i);
+                    }
+                }
+                if (nflip%2) {
+                    int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin];
+                    for (int i = 1; i < 8; ++i) {
+                        float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i];
+                        if (ax < min) {
+                            min = ax; imin = i;
+                        }
+                    }
+                    xval[8*k+imin] = -xval[8*k+imin];
+                    s ^= (1 << imin);
+                }
+                block_signs[k] = s & 127;
+            }
+            float max = xval[0];
+            for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]);
+            memset(L, 0, 32);
+            if (max < GROUP_MAX_EPS_IQ3_XXS) {
+                scales[ib] = 0;
+                continue;
+            }
+            float best = 0;
+            float scale = max/(2*kMaxQ-1);
+            for (int k = 0; k < 8; ++k) is_on_grid[k] = true;
+            for (int is = -15; is <= 15; ++is) {
+                float id = (2*kMaxQ-1+is*0.2f)/max;
+                float this_scale = 1/id;
+                for (int k = 0; k < 8; ++k) {
+                    for (int i = 0; i < 4; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[4*k+i]-1));
+                        Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l));
+                    }
+                    uint16_t u = 0;
+                    for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i);
+                    int grid_index = kmap_q3xs[u];
+                    is_on_grid_aux[k] = true;
+                    if (grid_index < 0) {
+                        is_on_grid_aux[k] = false;
+                        const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1;
+                        grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 32; ++i) {
+                    float w = weight[i];
+                    float q = 2*Laux[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                    scale = sumqx/sumq2; best = scale*sumqx;
+                    for (int i = 0; i < 32; ++i) L[i] = Laux[i];
+                    for (int k = 0; k <  8; ++k) is_on_grid[k] = is_on_grid_aux[k];
+                }
+            }
+            int n_not_ongrid = 0;
+            for (int k = 0; k < 8; ++k) if (!is_on_grid[k]) ++n_not_ongrid;
+            if (n_not_ongrid > 0 && scale > 0) {
+                float id = 1/scale;
+                for (int k = 0; k < 8; ++k) {
+                    if (is_on_grid[k]) continue;
+                    uint16_t u = 0;
+                    for (int i = 0; i < 4; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[4*k+i]-1));
+                        l = MAX(0, MIN(kMaxQ-1, l));
+                        u |= (l << 3*i);
+                    }
+                    int grid_index = kmap_q3xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1;
+                        grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k);
+                    }
+                    const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index);
+                    for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2;
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 32; ++i) {
+                    float w = weight[i];
+                    float q = 2*L[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0) scale = sumqx/sumq2;
+            }
+            if (scale < 0) {
+                // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale)
+                // and correspondingly flip quant signs.
+                scale = -scale;
+                for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127;
+            }
+            for (int k = 0; k < 8; ++k) {
+                uint16_t u = 0;
+                for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i);
+                int grid_index = kmap_q3xs[u];
+                if (grid_index < 0) {
+                    printf("Oops: found point %u not on grid:", u);
+                    for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]);
+                    printf("\n");
+                    GGML_ABORT("fatal error");
+                }
+                if (grid_size == 256) {
+                    q3[8*ib+k] = grid_index;
+                } else {
+                    q3[8*ib+k] = grid_index & 255;
+                    qh[ib] |= ((grid_index >> 8) << k);
+                }
+
+            }
+            scales_and_signs[ib] = block_signs[0] | (block_signs[1] << 7) | (block_signs[2] << 14) | (block_signs[3] << 21);
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            memset(qs, 0, quant_size);
+            dh += block_size/sizeof(ggml_fp16_t);
+            qs += block_size;
+            continue;
+        }
+
+        float d = max_scale/31;
+        dh[0] = GGML_FP32_TO_FP16(d * 1.0125f);  // small improvement via this fudge factor
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/32; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib]-1));
+            l = MAX(0, MIN(15, l));
+            scales_and_signs[ib] |= ((uint32_t)l << 28);
+        }
+        memcpy(qs, q3, quant_size);
+
+        dh += block_size/sizeof(ggml_fp16_t);
+        qs += block_size;
+
+    }
+}
+
+size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq3_xxs_impl(256, src, qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq3_xxs);
+    }
+    return nrow * nblock * sizeof(block_iq3_xxs);
+}
+
+void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    quantize_row_iq3_xxs_impl(256, x, y, k, NULL);
+}
+
+static void quantize_row_iq3_s_impl(int block_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int n,
+        const float * GGML_RESTRICT quant_weights,
+        float   * scales,
+        float   * weight,
+        float   * xval,
+        int8_t  * L,
+        int8_t  * Laux,
+        float   * waux,
+        bool    * is_on_grid,
+        bool    * is_on_grid_aux,
+        uint8_t * block_signs) {
+
+    const int gindex = iq3_data_index(512);
+
+    const uint32_t * kgrid_q3xs      = iq3_data[gindex].grid;
+    const int      * kmap_q3xs       = iq3_data[gindex].map;
+    const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours;
+
+    //GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kgrid_q3xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kmap_q3xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    const int kMaxQ = 8;
+
+    const int64_t nbl = n/QK_K;
+
+    block_iq3_s * y = vy;
+
+    const int bs4 = block_size/4;
+    const int bs8 = block_size/8;
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        memset(&y[ibl], 0, sizeof(block_iq3_s));
+        y[ibl].d = GGML_FP32_TO_FP16(0.f);
+
+        uint8_t * qs = y[ibl].qs;
+        uint8_t * qh = y[ibl].qh;
+        uint8_t * signs = y[ibl].signs;
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/block_size; ++ib) {
+            const float * xb = xbl + block_size*ib;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*ibl + block_size*ib;
+                for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            } else {
+                for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i];
+            }
+            for (int i = 0; i < block_size; ++i) waux[i] = sqrtf(weight[i]);
+            for (int k = 0; k < bs8; ++k) {
+                uint8_t s = 0;
+                for (int i = 0; i < 8; ++i) {
+                    if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i];
+                    else {
+                        xval[8*k + i] = -xb[8*k + i]; s |= (1 << i);
+                    }
+                }
+                block_signs[k] = s;
+            }
+            float max = xval[0];
+            for (int i = 1; i < block_size; ++i) max = MAX(max, xval[i]);
+            memset(L, 0, block_size);
+            if (!max) {
+                scales[ib] = 0;
+                continue;
+            }
+            float best = 0;
+            float scale = max/(2*kMaxQ-1);
+            for (int k = 0; k < bs4; ++k) is_on_grid[k] = false;
+            for (int is = -9; is <= 9; ++is) {
+                float id = (2*kMaxQ-1+is*0.2f)/max;
+                float this_scale = 1/id;
+                for (int k = 0; k < bs4; ++k) {
+                    for (int i = 0; i < 4; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[4*k+i]-1));
+                        Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l));
+                    }
+                    uint16_t u = 0;
+                    for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i);
+                    int grid_index = kmap_q3xs[u];
+                    is_on_grid_aux[k] = true;
+                    if (grid_index < 0) {
+                        is_on_grid_aux[k] = false;
+                        const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1;
+                        grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < block_size; ++i) {
+                    float w = weight[i];
+                    float q = 2*Laux[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                    scale = sumqx/sumq2; best = scale*sumqx;
+                    for (int i = 0; i < block_size; ++i) L[i] = Laux[i];
+                    for (int k = 0; k < bs4; ++k) is_on_grid[k] = is_on_grid_aux[k];
+                }
+            }
+            int n_not_ongrid = 0;
+            for (int k = 0; k < bs4; ++k) if (!is_on_grid[k]) ++n_not_ongrid;
+            if (n_not_ongrid > 0 && scale > 0) {
+                float id = 1/scale;
+                for (int k = 0; k < bs4; ++k) {
+                    //if (is_on_grid[k]) continue;
+                    uint16_t u = 0;
+                    for (int i = 0; i < 4; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[4*k+i]-1));
+                        l = MAX(0, MIN(kMaxQ-1, l));
+                        u |= (l << 3*i);
+                    }
+                    int grid_index = kmap_q3xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1;
+                        grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k);
+                    }
+                    const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index);
+                    for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2;
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < block_size; ++i) {
+                    float w = weight[i];
+                    float q = 2*L[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0) scale = sumqx/sumq2;
+            }
+            if (scale < 0) {
+                // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale)
+                // and correspondingly flip quant signs.
+                scale = -scale;
+                for (int k = 0; k < bs8; ++k) block_signs[k] = ~block_signs[k];
+            }
+            for (int k = 0; k < bs4; ++k) {
+                uint16_t u = 0;
+                for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i);
+                int grid_index = kmap_q3xs[u];
+                if (grid_index < 0) {
+                    printf("Oops: found point %u not on grid:", u);
+                    for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]);
+                    printf("\n");
+                    GGML_ABORT("fatal error");
+                }
+                qs[k] = grid_index & 255;
+                qh[(ib*bs4+k)/8] |= ((grid_index >> 8) << ((ib*bs4+k)%8));
+            }
+            qs += bs4;
+            for (int k = 0; k < bs8; ++k) signs[k] = block_signs[k];
+            signs += bs8;
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            continue;
+        }
+
+        float d = max_scale/31;
+        y[ibl].d = GGML_FP32_TO_FP16(d * 1.033f);
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/block_size; ib += 2) {
+            int l1 = nearest_int(0.5f*(id*scales[ib+0]-1));
+            l1 = MAX(0, MIN(15, l1));
+            int l2 = nearest_int(0.5f*(id*scales[ib+1]-1));
+            l2 = MAX(0, MIN(15, l2));
+            y[ibl].scales[ib/2] = l1 | (l2 << 4);
+        }
+
+    }
+}
+
+#define IQ3S_BLOCK_SIZE 32
+size_t quantize_iq3_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    float scales[QK_K/IQ3S_BLOCK_SIZE];
+    float weight[IQ3S_BLOCK_SIZE];
+    float xval[IQ3S_BLOCK_SIZE];
+    int8_t L[IQ3S_BLOCK_SIZE];
+    int8_t Laux[IQ3S_BLOCK_SIZE];
+    float  waux[IQ3S_BLOCK_SIZE];
+    bool   is_on_grid[IQ3S_BLOCK_SIZE/4];
+    bool   is_on_grid_aux[IQ3S_BLOCK_SIZE/4];
+    uint8_t block_signs[IQ3S_BLOCK_SIZE/8];
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq3_s_impl(IQ3S_BLOCK_SIZE, src, qrow, n_per_row, quant_weights,
+                scales, weight, xval, L, Laux, waux, is_on_grid, is_on_grid_aux, block_signs);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq3_s);
+    }
+    return nrow * nblock * sizeof(block_iq3_s);
+}
+
+void quantize_row_iq3_s_ref(const float * GGML_RESTRICT x, block_iq3_s * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    quantize_iq3_s(x, y, 1, k, NULL);
+}
+
+
+// =================================== 1.5 bpw ===================================================
+
+static int iq1_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid,
+        const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float * scale, int8_t * GGML_RESTRICT L, int ngrid) {
+    int num_neighbors = neighbours[0];
+    GGML_ASSERT(num_neighbors > 0);
+    float best_score = -FLT_MAX;
+    int grid_index = -1;
+    for (int j = 1; j <= num_neighbors; ++j) {
+        const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+        float sumqx = 0, sumq2 = 0;
+        for (int i = 0; i < 8; ++i) {
+            float q = (pg[i] - 3)/2;
+            float w = weight[i];
+            sumqx += w*q*xval[i];
+            sumq2 += w*q*q;
+        }
+        if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) {
+            *scale = sumqx/sumq2; best_score = *scale * sumqx;
+            grid_index = neighbours[j];
+        }
+    }
+    if (grid_index < 0) {
+        for (int i = 0; i < ngrid; ++i) {
+            const int8_t * grid_i = (const int8_t *)(grid + i);
+            float sumqx = 0, sumq2 = 0;
+            for (int j = 0; j < 8; ++j) {
+                float w = weight[j];
+                float q = (grid_i[j] - 3)/2;
+                sumqx += w*q*xval[j];
+                sumq2 += w*q*q;
+            }
+            if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) {
+                *scale = sumqx/sumq2; best_score = *scale*sumqx;
+                grid_index = i;
+            }
+        }
+    }
+    if (grid_index < 0) {
+        printf("Oops, did not find grid point\n");
+        printf("Have %d neighbours\n", num_neighbors);
+        for (int j = 1; j <= num_neighbors; ++j) {
+            const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+            float sumqx = 0, sumq2 = 0;
+            for (int i = 0; i < 8; ++i) {
+                float q = (pg[i] - 3)/2;
+                float w = weight[i];
+                sumqx += w*q*xval[i];
+                sumq2 += w*q*q;
+            }
+            printf("    neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2);
+        }
+    }
+    GGML_ASSERT(grid_index >= 0);
+    //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+    *scale *= 1.05f;  // This is a fudge factor. Don't ask me why it improves the result.
+    //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+    const int8_t * pg = (const int8_t *)(grid + grid_index);
+    for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2;
+    return grid_index;
+}
+
+static int iq1_find_best_neighbour2(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid,
+        const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, const float * GGML_RESTRICT xg, int8_t * GGML_RESTRICT L, int ngrid) {
+    int num_neighbors = neighbours[0];
+    GGML_ASSERT(num_neighbors > 0);
+    float best_score = FLT_MAX;
+    int grid_index = -1;
+    for (int j = 1; j <= num_neighbors; ++j) {
+        const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+        float d2 = 0;
+        for (int i = 0; i < 8; ++i) {
+            float q = xg[(pg[i] - 1)/2];
+            float w = weight[i];
+            float diff = scale*q - xval[i];
+            d2 += w*diff*diff;
+        }
+        if (d2 < best_score) {
+            best_score = d2;
+            grid_index = neighbours[j];
+        }
+    }
+    if (grid_index < 0) {
+        for (int i = 0; i < ngrid; ++i) {
+            const int8_t * grid_i = (const int8_t *)(grid + i);
+            float d2 = 0;
+            for (int j = 0; j < 8; ++j) {
+                float w = weight[j];
+                float q = xg[(grid_i[j] - 1)/2];
+                float diff = scale*q - xval[i];
+                d2 += w*diff*diff;
+            }
+            if (d2 < best_score) {
+                best_score = d2;
+                grid_index = i;
+            }
+        }
+    }
+    if (grid_index < 0) {
+        printf("Oops, did not find grid point\n");
+        printf("Have %d neighbours\n", num_neighbors);
+        for (int j = 1; j <= num_neighbors; ++j) {
+            const int8_t * pg = (const int8_t *)(grid + neighbours[j]);
+            float sumqx = 0, sumq2 = 0;
+            for (int i = 0; i < 8; ++i) {
+                float q = xg[(pg[i] - 1)/2];
+                float w = weight[i];
+                sumqx += w*q*xval[i];
+                sumq2 += w*q*q;
+            }
+            printf("    neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2);
+        }
+    }
+    GGML_ASSERT(grid_index >= 0);
+    const int8_t * pg = (const int8_t *)(grid + grid_index);
+    for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2;
+    return grid_index;
+}
+
+static int iq1_sort_helper(const void * left, const void * right) {
+    const float * l = left;
+    const float * r = right;
+    return *l < *r ? -1 : *l > *r ? 1 : 0;
+}
+
+#define IQ1S_BLOCK_SIZE 32
+#define IQ1M_BLOCK_SIZE 16
+static void quantize_row_iq1_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights,
+        float    * scales,
+        float    * weight,
+        float    * sumx,
+        float    * sumw,
+        float    * pairs,
+        int8_t   * L,
+        uint16_t * index,
+        int8_t   * shifts) {
+
+    const int gindex = iq2_data_index(GGML_TYPE_IQ1_S);
+
+    const uint64_t * kgrid_q2xs      = iq2_data[gindex].grid;
+    const int      * kmap_q2xs       = iq2_data[gindex].map;
+    const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours;
+
+    GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kgrid_q2xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kmap_q2xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    block_iq1_s * y = vy;
+
+    const int64_t nbl = n/QK_K;
+
+    const int block_size = IQ1S_BLOCK_SIZE;
+
+    const float x_p[3] = {-1 + IQ1S_DELTA,  IQ1S_DELTA, 1 + IQ1S_DELTA};
+    const float x_m[3] = {-1 - IQ1S_DELTA, -IQ1S_DELTA, 1 - IQ1S_DELTA};
+
+
+    int * idx = (int *)(pairs + 1);
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        y[ibl].d = GGML_FP32_TO_FP16(0.f);
+        memset(y[ibl].qs, 0, QK_K/8);
+        memset(y[ibl].qh, 0, QK_K/16);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/block_size; ++ib) {
+            const float * xb = xbl + block_size*ib;
+            const float * qw = quant_weights + QK_K*ibl + block_size*ib;
+            for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            float max = fabsf(xb[0]);
+            for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i]));
+            if (max < GROUP_MAX_EPS_IQ1_S) {
+                scales[ib] = 0;
+                shifts[ib] = 1;
+                memset(L, 1, block_size);
+                continue;
+            }
+            // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem.
+            // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two
+            // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights
+            // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and
+            // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale
+            // for each possible and score for each split.
+            for (int j = 0; j < block_size; ++j) {
+                pairs[2*j] = xb[j];
+                idx[2*j] = j;
+            }
+            qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper);
+            {
+                sumx[0] = sumw[0] = 0;
+                for (int j = 0; j < block_size; ++j) {
+                    int i = idx[2*j];
+                    sumx[j+1] = sumx[j] + weight[i]*xb[i];
+                    sumw[j+1] = sumw[j] + weight[i];
+                }
+            }
+            float best_score = -FLT_MAX, scale = max;
+            int besti1 = -1, besti2 = -1, best_shift = 0;
+            for (int i1 = 0; i1 <= block_size; ++i1) {
+                for (int i2 = i1; i2 <= block_size; ++i2) {
+                    float sumqx = (sumx[i1] - sumx[0])*x_p[0] + (sumx[i2] - sumx[i1])*x_p[1] + (sumx[block_size] - sumx[i2])*x_p[2];
+                    float sumq2 = (sumw[i1] - sumw[0])*x_p[0]*x_p[0] + (sumw[i2] - sumw[i1])*x_p[1]*x_p[1] + (sumw[block_size] - sumw[i2])*x_p[2]*x_p[2];
+                    if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) {
+                        scale = sumqx/sumq2; best_score = scale*sumqx;
+                        besti1 = i1; besti2 = i2; best_shift = 1;
+                    }
+                    sumqx = (sumx[i1] - sumx[0])*x_m[0] + (sumx[i2] - sumx[i1])*x_m[1] + (sumx[block_size] - sumx[i2])*x_m[2];
+                    sumq2 = (sumw[i1] - sumw[0])*x_m[0]*x_m[0] + (sumw[i2] - sumw[i1])*x_m[1]*x_m[1] + (sumw[block_size] - sumw[i2])*x_m[2]*x_m[2];
+                    if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) {
+                        scale = sumqx/sumq2; best_score = scale*sumqx;
+                        besti1 = i1; besti2 = i2; best_shift = -1;
+                    }
+                }
+            }
+            if (besti1 < 0 || besti2 < 0 || best_shift == 0) {
+                scales[ib] = 0;
+                shifts[ib] = 1;
+                memset(L, 1, block_size);
+                continue;
+            }
+            for (int j =      0; j < besti1; ++j) L[idx[2*j]] = 0;
+            for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1;
+            for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2;
+            if (scale < 0) {
+                for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j];
+                scale = -scale; best_shift = -best_shift;
+            }
+            bool all_on_grid = true;
+            const float * xx = best_shift == 1 ? x_p : x_m;
+            for (int k = 0; k < block_size/8; ++k) {
+                uint16_t u = 0;
+                for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j);
+                int grid_index = kmap_q2xs[u];
+                if (grid_index < 0) {
+                    all_on_grid = false;
+                    const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                    grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S);
+                    GGML_ASSERT(grid_index >= 0);
+                }
+                index[k] = grid_index;
+            }
+            if (!all_on_grid) {
+                float sumqx = 0, sumq2 = 0;
+                for (int k = 0; k < block_size/8; ++k) {
+                    const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]);
+                    for (int j = 0; j < 8; ++j) {
+                        float w = weight[8*k + j];
+                        float q = xx[(pg[j] - 1)/2];
+                        sumqx += w*q*xb[8*k+j];
+                        sumq2 += w*q*q;
+                    }
+                }
+                if (sumqx > 0 && sumq2 > 0) scale = sumqx/sumq2;
+            }
+            uint16_t h = 0;
+            for (int k = 0; k < block_size/8; ++k) {
+                y[ibl].qs[(block_size/8)*ib + k] = index[k] & 255;
+                h |= (index[k] >> 8) << 3*k;
+            }
+            y[ibl].qh[ib] = h;
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            shifts[ib] = best_shift;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            continue;
+        }
+
+        float d = max_scale/15;
+        y[ibl].d = GGML_FP32_TO_FP16(d*1.125f); // 1.125f is another fudge factor. Don't ask me why it is needed.
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/block_size; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib]-1));
+            l = MAX(0, MIN(7, l));
+            if (shifts[ib] == -1) l |= 8;
+            y[ibl].qh[ib] |= (l << 12);
+        }
+    }
+}
+
+size_t quantize_iq1_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    float  scales[QK_K/IQ1S_BLOCK_SIZE];
+    float  weight[IQ1S_BLOCK_SIZE];
+    int8_t L[IQ1S_BLOCK_SIZE];
+    float  sumx[IQ1S_BLOCK_SIZE+1];
+    float  sumw[IQ1S_BLOCK_SIZE+1];
+    float  pairs[2*IQ1S_BLOCK_SIZE];
+    uint16_t index[IQ1S_BLOCK_SIZE/8];
+    int8_t shifts[QK_K/IQ1S_BLOCK_SIZE];
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq1_s_impl(src, qrow, n_per_row, quant_weights, scales, weight, sumx, sumw, pairs, L, index, shifts);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq1_s);
+    }
+    return nrow * nblock * sizeof(block_iq1_s);
+}
+
+static void quantize_row_iq1_m_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights,
+        float    * scales,
+        float    * weight,
+        float    * pairs,
+        int8_t   * L,
+        uint16_t * index,
+        int8_t   * shifts) {
+
+    const int gindex = iq2_data_index(GGML_TYPE_IQ1_M);
+
+    const uint64_t * kgrid_q2xs      = iq2_data[gindex].grid;
+    const int      * kmap_q2xs       = iq2_data[gindex].map;
+    const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours;
+
+    //GGML_ASSERT(quant_weights   && "missing quantization weights");
+    GGML_ASSERT(kgrid_q2xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kmap_q2xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    block_iq1_m * y = vy;
+
+    const int64_t nbl = n/QK_K;
+
+    const int block_size = IQ1M_BLOCK_SIZE;
+
+    const float x_p[3] = {-1 + IQ1M_DELTA,  IQ1M_DELTA, 1 + IQ1M_DELTA};
+    const float x_m[3] = {-1 - IQ1M_DELTA, -IQ1M_DELTA, 1 - IQ1M_DELTA};
+    const uint8_t masks[4] = {0x00, 0x80, 0x08, 0x88};
+
+    int * idx = (int *)(pairs + 1);
+
+    float sumqx[4], sumq2[4];
+
+    iq1m_scale_t s;
+    const float * xx;
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+        memset(y[ibl].qs, 0, QK_K/8);
+        memset(y[ibl].qh, 0, QK_K/16);
+        memset(y[ibl].scales, 0, QK_K/32);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/block_size; ++ib) {
+            const float * xb = xbl + block_size*ib;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*ibl + block_size*ib;
+                for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            } else {
+                for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i];
+            }
+            float max = fabsf(xb[0]);
+            for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i]));
+            if (max < GROUP_MAX_EPS_IQ1_M) {
+                scales[ib] = 0;
+                shifts[ib] = 0;
+                memset(L, 1, block_size);
+                continue;
+            }
+            // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem.
+            // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two
+            // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights
+            // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and
+            // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale
+            // for each possible and score for each split.
+            for (int j = 0; j < block_size; ++j) {
+                pairs[2*j] = xb[j];
+                idx[2*j] = j;
+            }
+            qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper);
+            float best_score = -FLT_MAX, scale = max;
+            int besti1 = -1, besti2 = -1, best_k = -1;
+            // 0: +, +
+            // 1: +, -
+            // 2: -, +
+            // 3: -, -
+            // prefix sums are kept per half of the block because each half can use a different sign (x_p or x_m)
+            // since v[0]-v[1] = v[1]-v[2] = -1 for both x_p and x_m, the 3-group sum for a split collapses to T*v[2] - px[i1] - px[i2]
+            {
+                float px[2][IQ1M_BLOCK_SIZE+1];
+                float pw[2][IQ1M_BLOCK_SIZE+1];
+                px[0][0] = px[1][0] = 0;
+                pw[0][0] = pw[1][0] = 0;
+                for (int j = 0; j < block_size; ++j) {
+                    const int i = idx[2*j];
+                    const int h = i < block_size/2 ? 0 : 1;
+                    px[h][j+1]   = px[h][j] + weight[i]*xb[i];
+                    px[1-h][j+1] = px[1-h][j];
+                    pw[h][j+1]   = pw[h][j] + weight[i];
+                    pw[1-h][j+1] = pw[1-h][j];
+                }
+                const float txs[2] = {px[0][block_size], px[1][block_size]}; // total weight*x per half
+                const float tws[2] = {pw[0][block_size], pw[1][block_size]}; // total weight per half
+                const float p2   = x_p[2], m2 = x_m[2];
+                const float cp1  = x_p[0]*x_p[0] - x_p[1]*x_p[1];
+                const float cp2  = x_p[1]*x_p[1] - x_p[2]*x_p[2];
+                const float cm1  = x_m[0]*x_m[0] - x_m[1]*x_m[1];
+                const float cm2  = x_m[1]*x_m[1] - x_m[2]*x_m[2];
+                for (int i1 = 0; i1 <= block_size; ++i1) {
+                    for (int i2 = i1; i2 <= block_size; ++i2) {
+                        float qx_p[2], qx_m[2], q2_p[2], q2_m[2];
+                        for (int h = 0; h < 2; ++h) {
+                            const float sx = px[h][i1] + px[h][i2];
+                            qx_p[h] = txs[h]*p2 - sx;
+                            qx_m[h] = txs[h]*m2 - sx;
+                            q2_p[h] = tws[h]*p2*p2 + pw[h][i1]*cp1 + pw[h][i2]*cp2;
+                            q2_m[h] = tws[h]*m2*m2 + pw[h][i1]*cm1 + pw[h][i2]*cm2;
+                        }
+                        sumqx[0] = qx_p[0] + qx_p[1];
+                        sumqx[1] = qx_p[0] + qx_m[1];
+                        sumqx[2] = qx_m[0] + qx_p[1];
+                        sumqx[3] = qx_m[0] + qx_m[1];
+                        sumq2[0] = q2_p[0] + q2_p[1];
+                        sumq2[1] = q2_p[0] + q2_m[1];
+                        sumq2[2] = q2_m[0] + q2_p[1];
+                        sumq2[3] = q2_m[0] + q2_m[1];
+                        for (int k = 0; k < 4; ++k) {
+                            if (sumq2[k] > 0 && sumqx[k]*sumqx[k] > best_score*sumq2[k]) {
+                                scale = sumqx[k]/sumq2[k]; best_score = scale*sumqx[k];
+                                besti1 = i1; besti2 = i2; best_k = k;
+                            }
+                        }
+                    }
+                }
+            }
+            if (besti1 < 0 || besti2 < 0 || best_k < 0) {
+                scales[ib] = 0;
+                shifts[ib] = 0;
+                memset(L, 1, block_size);
+                continue;
+            }
+            for (int j =      0; j < besti1; ++j) L[idx[2*j]] = 0;
+            for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1;
+            for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2;
+            if (scale < 0) {
+                for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j];
+                scale = -scale;
+                best_k = best_k == 0 ? 3 : best_k == 1 ? 2 : best_k == 2 ? 1 : 0;
+            }
+            bool all_on_grid = true;
+            for (int k = 0; k < block_size/8; ++k) {
+                if (k == 0) xx = best_k < 2 ? x_p : x_m;
+                else xx = best_k%2 == 0 ? x_p : x_m;
+                uint16_t u = 0;
+                for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j);
+                int grid_index = kmap_q2xs[u];
+                if (grid_index < 0) {
+                    all_on_grid = false;
+                    const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                    grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S);
+                    GGML_ASSERT(grid_index >= 0);
+                }
+                index[k] = grid_index;
+            }
+            if (!all_on_grid) {
+                float sumqx_f = 0, sumq2_f = 0;
+                for (int k = 0; k < block_size/8; ++k) {
+                    if (k == 0) xx = best_k < 2 ? x_p : x_m;
+                    else xx = best_k%2 == 0 ? x_p : x_m;
+                    const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]);
+                    for (int j = 0; j < 8; ++j) {
+                        float w = weight[8*k + j];
+                        float q = xx[(pg[j] - 1)/2];
+                        sumqx_f += w*q*xb[8*k+j];
+                        sumq2_f += w*q*q;
+                    }
+                }
+                if (sumqx_f > 0 && sumq2_f > 0) scale = sumqx_f/sumq2_f;
+            }
+            y[ibl].qs[2*ib + 0] = index[0] & 255;
+            y[ibl].qs[2*ib + 1] = index[1] & 255;
+            y[ibl].qh[ib] = (index[0] >> 8) | ((index[1] >> 8) << 4);
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            shifts[ib] = best_k;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            continue;
+        }
+
+        uint16_t * sc = (uint16_t *)y[ibl].scales;
+        float d = max_scale/15;
+        float id = 1/d;
+        float sumqx_f = 0, sumq2_f = 0;
+        for (int ib = 0; ib < QK_K/block_size; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib+0]-1));
+            l = MAX(0, MIN(7, l));
+            sc[ib/4] |= (l << 3*(ib%4));
+            y[ibl].qh[ib] |= masks[shifts[ib]];
+            const float * xb = xbl + block_size*ib;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*ibl + block_size*ib;
+                for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            } else {
+                for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i];
+            }
+            for (int k = 0; k < block_size/8; ++k) {
+                if (k == 0) xx = shifts[ib] < 2 ? x_p : x_m;
+                else xx = shifts[ib]%2 == 0 ? x_p : x_m;
+                const int8_t * pg = (const int8_t *)(kgrid_q2xs + y[ibl].qs[2*ib+k] + ((y[ibl].qh[ib] << (8 - 4*k)) & 0x700));
+                for (int j = 0; j < 8; ++j) {
+                    float w = weight[8*k + j];
+                    float q = xx[(pg[j] - 1)/2]*(2*l+1);
+                    sumqx_f += w*q*xb[8*k+j];
+                    sumq2_f += w*q*q;
+                }
+            }
+        }
+        if (sumq2_f > 0) d = sumqx_f/sumq2_f;
+        s.f16 = GGML_FP32_TO_FP16(d*1.1125f); // 1.1125f is another fudge factor. Don't ask me why it is needed.
+        sc[0] |= ((s.u16 & 0x000f) << 12);
+        sc[1] |= ((s.u16 & 0x00f0) <<  8);
+        sc[2] |= ((s.u16 & 0x0f00) <<  4);
+        sc[3] |= ((s.u16 & 0xf000) <<  0);
+    }
+}
+
+size_t quantize_iq1_m(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    float  scales[QK_K/IQ1M_BLOCK_SIZE];
+    float  weight[IQ1M_BLOCK_SIZE];
+    int8_t L[IQ1M_BLOCK_SIZE];
+    float  pairs[2*IQ1M_BLOCK_SIZE];
+    uint16_t index[IQ1M_BLOCK_SIZE/8];
+    int8_t shifts[QK_K/IQ1M_BLOCK_SIZE];
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq1_m_impl(src, qrow, n_per_row, quant_weights, scales, weight, pairs, L, index, shifts);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq1_m);
+    }
+    return nrow * nblock * sizeof(block_iq1_m);
+}
+
+// ============================ 4-bit non-linear quants
+
+static void quantize_row_iq4_nl_impl(const int super_block_size, const int block_size, const float * GGML_RESTRICT x,
+        ggml_fp16_t * dh, uint8_t * q4, uint16_t * scales_h, uint8_t * scales_l,
+        float * scales, float * weight, uint8_t * L,
+        const int8_t * values,
+        const float * quant_weights,
+        const int ntry) {
+
+    float sigma2 = 0;
+    for (int j = 0; j < super_block_size; ++j) sigma2 += x[j]*x[j];
+    sigma2 *= 2.f/super_block_size;
+
+    memset(q4, 0, super_block_size/2);
+    dh[0] = GGML_FP32_TO_FP16(0.f);
+
+    float max_scale = 0, amax_scale = 0;
+    for (int ib = 0; ib < super_block_size/block_size; ++ib) {
+        const float * xb = x + ib*block_size;
+        uint8_t * Lb = L + ib*block_size;
+        if (quant_weights) {
+            const float * qw = quant_weights + ib*block_size;
+            for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]);
+        } else {
+            for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j];
+        }
+        float amax = 0, max = 0;
+        for (int j = 0; j < block_size; ++j) {
+            float ax = fabsf(xb[j]);
+            if (ax > amax) {
+                amax = ax; max = xb[j];
+            }
+        }
+        if (amax < GROUP_MAX_EPS) {
+            scales[ib] = 0;
+            continue;
+        }
+        float d = ntry > 0 ? -max/values[0] : max/values[0];
+        float id = 1/d;
+        float sumqx = 0, sumq2 = 0;
+        for (int j = 0; j < block_size; ++j) {
+            float al = id*xb[j];
+            int l = best_index_int8(16, values, al);
+            Lb[j] = l;
+            float q = values[l];
+            float w = weight[j];
+            sumqx += w*q*xb[j];
+            sumq2 += w*q*q;
+        }
+        d = sumq2 > 0 ? sumqx/sumq2 : 0.f;
+        float best = d*sumqx;
+        for (int itry = -ntry; itry <= ntry; ++itry) {
+            id = (itry + values[0])/max;
+            sumqx = sumq2 = 0;
+            for (int j = 0; j < block_size; ++j) {
+                float al = id*xb[j];
+                int l = best_index_int8(16, values, al);
+                float q = values[l];
+                float w = weight[j];
+                sumqx += w*q*xb[j];
+                sumq2 += w*q*q;
+            }
+            if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                d = sumqx/sumq2; best = d * sumqx;
+            }
+        }
+        scales[ib] = d;
+        float abs_d = fabsf(d);
+        if (abs_d > amax_scale) {
+            amax_scale = abs_d; max_scale = d;
+        }
+    }
+
+    if (super_block_size/block_size > 1) {
+        int nb = super_block_size/block_size;
+        memset(scales_h, 0, ((nb+7)/8)*sizeof(uint16_t));
+        float d = -max_scale/32;
+        dh[0] = GGML_FP32_TO_FP16(d);
+        float id = d ? 1/d : 0.f;
+        for (int ib = 0; ib < super_block_size/block_size; ++ib) {
+            int l = nearest_int(id*scales[ib]);
+            l = MAX(-32, MIN(31, l));
+            float dl = d * l;
+            float idl = dl ? 1/dl : 0.f;
+            uint8_t * Lb = L + ib*block_size;
+            const float * xb = x + ib*block_size;
+            for (int j = 0; j < block_size; ++j) {
+                Lb[j] = best_index_int8(16, values, idl*xb[j]);
+            }
+            l += 32;
+            uint8_t l_l = l & 0xf;
+            uint8_t l_h = l >>  4;
+            if (ib%2 == 0) scales_l[ib/2] = l_l;
+            else scales_l[ib/2] |= (l_l << 4);
+            scales_h[ib/8] |= (l_h << 2*(ib%8));
+        }
+    } else {
+        dh[0] = GGML_FP32_TO_FP16(scales[0]);
+        if (ntry > 0) {
+            float id = scales[0] ? 1/scales[0] : 0;
+            for (int j = 0; j < super_block_size; ++j) {
+                L[j] = best_index_int8(16, values, id*x[j]);
+            }
+        }
+    }
+
+    for (int i = 0; i < super_block_size/32; ++i) {
+        for (int j = 0; j < 16; ++j) {
+            q4[16*i + j] = L[32*i + j] | (L[32*i + 16 + j] << 4);
+        }
+    }
+}
+
+size_t quantize_iq4_nl(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK4_NL == 0);
+    int64_t nblock = n_per_row/QK4_NL;
+    char * qrow = (char *)dst;
+    uint8_t L[QK4_NL];
+    float weight[QK4_NL];
+    uint16_t unused_h;
+    uint8_t * unused_l = NULL;
+    float scale;
+    for (int64_t row = 0; row < nrow; ++row) {
+        block_iq4_nl * iq4 = (block_iq4_nl *)qrow;
+        for (int ibl = 0; ibl < nblock; ++ibl) {
+            const float * qw = quant_weights ? quant_weights + QK4_NL*ibl : NULL;
+            quantize_row_iq4_nl_impl(QK4_NL, 32, src + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l,
+                    &scale, weight, L, kvalues_iq4nl, qw, 7);
+        }
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq4_nl);
+    }
+    return nrow * nblock * sizeof(block_iq4_nl);
+}
+
+//void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, block_iq4_nl * GGML_RESTRICT y, int64_t k) {
+    GGML_ASSERT(k%QK4_NL == 0);
+    int64_t nblock = k/QK4_NL;
+    uint8_t L[QK4_NL];
+    float weight[QK4_NL];
+    uint16_t unused_h;
+    uint8_t * unused_l = NULL;
+    float scale;
+    block_iq4_nl * iq4 = y;
+    for (int ibl = 0; ibl < nblock; ++ibl) {
+        quantize_row_iq4_nl_impl(QK4_NL, 32, x + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l,
+                &scale, weight, L, kvalues_iq4nl, NULL, -1);
+    }
+}
+
+size_t quantize_iq4_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    uint8_t L[QK_K];
+    float weight[32];
+    float scales[QK_K/32];
+    for (int64_t row = 0; row < nrow; ++row) {
+        block_iq4_xs * iq4 = (block_iq4_xs *)qrow;
+        for (int ibl = 0; ibl < nblock; ++ibl) {
+            const float * qw = quant_weights ? quant_weights + QK_K*ibl : NULL;
+            quantize_row_iq4_nl_impl(QK_K, 32, src + QK_K*ibl, &iq4[ibl].d, iq4[ibl].qs, &iq4[ibl].scales_h, iq4[ibl].scales_l,
+                    scales, weight, L, kvalues_iq4nl, qw, 7);
+        }
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq4_xs);
+    }
+    return nrow * nblock * sizeof(block_iq4_xs);
+}
+
+void quantize_row_iq4_xs_ref(const float * GGML_RESTRICT x, block_iq4_xs * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    quantize_iq4_xs(x, y, 1, k, NULL);
+}
+
+// =============================== 2.5625 bpw
+
+static void quantize_row_iq2_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) {
+
+    const int gindex = iq2_data_index(GGML_TYPE_IQ2_S);
+
+    const uint64_t * kgrid_q2xs      = iq2_data[gindex].grid;
+    const int      * kmap_q2xs       = iq2_data[gindex].map;
+    const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours;
+
+    GGML_ASSERT(kmap_q2xs       && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kgrid_q2xs      && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?");
+    GGML_ASSERT(n%QK_K == 0);
+
+    const int kMaxQ = 3;
+
+    const int64_t nbl = n/QK_K;
+
+    block_iq2_s * y = vy;
+
+    float scales[QK_K/16];
+    float weight[16];
+    float xval[16];
+    int8_t L[16];
+    int8_t Laux[16];
+    float  waux[16];
+    bool   is_on_grid[2];
+    bool   is_on_grid_aux[2];
+    uint8_t block_signs[2];
+
+    for (int ibl = 0; ibl < nbl; ++ibl) {
+
+        memset(&y[ibl], 0, sizeof(block_iq2_s));
+        y[ibl].d = GGML_FP32_TO_FP16(0.f);
+
+        float max_scale = 0;
+
+        const float * xbl = x + QK_K*ibl;
+        float sumx2 = 0;
+        for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i];
+        float sigma2 = 2*sumx2/QK_K;
+
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            const float * xb = xbl + 16*ib;
+            if (quant_weights) {
+                const float * qw = quant_weights + QK_K*ibl + 16*ib;
+                for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]);
+            } else {
+                for (int i = 0; i < 16; ++i) weight[i] = 0.25f*sigma2 + xb[i]*xb[i];
+            }
+            for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]);
+            for (int k = 0; k < 2; ++k) {
+                uint8_t s = 0;
+                for (int i = 0; i < 8; ++i) {
+                    if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i];
+                    else {
+                        xval[8*k + i] = -xb[8*k + i]; s |= (1 << i);
+                    }
+                }
+                block_signs[k] = s;
+            }
+            float max = xval[0];
+            for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]);
+            memset(L, 0, 16);
+            if (max < GROUP_MAX_EPS_IQ2_S) {
+                scales[ib] = 0;
+                continue;
+            }
+            float best = 0;
+            float scale = max/(2*kMaxQ-1);
+            is_on_grid[0] = is_on_grid[1] = true;
+            for (int is = -9; is <= 9; ++is) {
+                float id = (2*kMaxQ-1+is*0.1f)/max;
+                float this_scale = 1/id;
+                for (int k = 0; k < 2; ++k) {
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l));
+                    }
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i);
+                    int grid_index = kmap_q2xs[u];
+                    is_on_grid_aux[k] = true;
+                    if (grid_index < 0) {
+                        is_on_grid_aux[k] = false;
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 16; ++i) {
+                    float w = weight[i];
+                    float q = 2*Laux[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0 && sumqx*sumqx > best*sumq2) {
+                    scale = sumqx/sumq2; best = scale*sumqx;
+                    for (int i = 0; i < 16; ++i) L[i] = Laux[i];
+                    for (int k = 0; k <  2; ++k) is_on_grid[k] = is_on_grid_aux[k];
+                }
+            }
+            int n_not_ongrid = 0;
+            for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid;
+            if (n_not_ongrid > 0 && scale > 0) {
+                float id = 1/scale;
+                for (int k = 0; k < 2; ++k) {
+                    if (is_on_grid[k]) continue;
+                    uint16_t u = 0;
+                    for (int i = 0; i < 8; ++i) {
+                        int l = nearest_int(0.5f*(id*xval[8*k+i]-1));
+                        l = MAX(0, MIN(kMaxQ-1, l));
+                        u |= (l << 2*i);
+                        L[8*k + i] = l;
+                    }
+                    int grid_index = kmap_q2xs[u];
+                    if (grid_index < 0) {
+                        const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1;
+                        grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k);
+                    }
+                }
+                float sumqx = 0, sumq2 = 0;
+                for (int i = 0; i < 16; ++i) {
+                    float w = weight[i];
+                    float q = 2*L[i] + 1;
+                    sumqx += w*xval[i]*q;
+                    sumq2 += w*q*q;
+                }
+                if (sumq2 > 0) scale = sumqx/sumq2;
+            }
+            if (scale < 0) {
+                scale = -scale;
+                for (int k = 0; k < 2; ++k) block_signs[k] = ~block_signs[k];
+            }
+            for (int k = 0; k < 2; ++k) {
+                uint16_t u = 0;
+                for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i);
+                int grid_index = kmap_q2xs[u];
+                if (grid_index < 0) {
+                    printf("Oops: found point %u not on grid:", u);
+                    for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]);
+                    printf("\n");
+                    GGML_ABORT("fatal error");
+                }
+                const int i8 = 2*ib + k;
+                y[ibl].qs[i8] = grid_index & 255;
+                y[ibl].qh[i8/4] |= ((grid_index >> 8) << 2*(i8%4));
+                y[ibl].qs[QK_K/8 + i8] = block_signs[k];
+            }
+            GGML_ASSERT(scale >= 0);
+            scales[ib] = scale;
+            max_scale = MAX(max_scale, scale);
+        }
+
+        if (!max_scale) {
+            continue;
+        }
+
+        float d = max_scale/31;
+        y[ibl].d = GGML_FP32_TO_FP16(d * 0.9875f);
+        float id = 1/d;
+        for (int ib = 0; ib < QK_K/16; ++ib) {
+            int l = nearest_int(0.5f*(id*scales[ib]-1));
+            l = MAX(0, MIN(15, l));
+            if (ib%2 == 0) y[ibl].scales[ib/2] = l;
+            else y[ibl].scales[ib/2] |= (l << 4);
+        }
+    }
+}
+
+size_t quantize_iq2_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    GGML_ASSERT(n_per_row%QK_K == 0);
+    int64_t nblock = n_per_row/QK_K;
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_iq2_s_impl(src, qrow, n_per_row, quant_weights);
+        src += n_per_row;
+        qrow += nblock*sizeof(block_iq2_s);
+    }
+    return nrow * nblock * sizeof(block_iq2_s);
+}
+
+void quantize_row_iq2_s_ref(const float * GGML_RESTRICT x, block_iq2_s * GGML_RESTRICT y, int64_t k) {
+    assert(k % QK_K == 0);
+    quantize_iq2_s(x, y, 1, k, NULL);
+}
+
+// =============================== data validation
+
+static bool validate_float(float f, size_t i) {
+    if (isinf(f)) {
+        fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i);
+        return false;
+    }
+
+    if (isnan(f)) {
+        fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i);
+        return false;
+    }
+
+    return true;
+}
+
+static bool isinf_fp16(ggml_fp16_t f) {
+    return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) == 0;
+}
+
+static bool isnan_fp16(ggml_fp16_t f) {
+    return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) != 0;
+}
+
+static bool validate_fp16(ggml_fp16_t f, size_t i) {
+    if (isinf_fp16(f)) {
+        fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i);
+        return false;
+    }
+
+    if (isnan_fp16(f)) {
+        fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i);
+        return false;
+    }
+
+    return true;
+}
+
+static bool validate_e_e8m0(uint8_t e, size_t i) {
+    if (e == 0xff) {
+        fprintf(stderr, "ggml_validate_row_data: found invalid e value %d at block %zu\n", e, i);
+        return false;
+    }
+
+    return true;
+}
+
+#define VALIDATE_ROW_DATA_D_F16_IMPL(type, data, nb) \
+    const type * q = (const type *) (data); \
+    for (size_t i = 0; i < (nb); ++i) { \
+        if (!validate_fp16(q[i].d, i)) { \
+            return false; \
+        } \
+    }
+
+#define VALIDATE_ROW_DATA_DM_F16_IMPL(type, data, nb, d, m) \
+    const type * q = (const type *) (data); \
+    for (size_t i = 0; i < (nb); ++i) { \
+        if (!validate_fp16(q[i].d, i) || !validate_fp16(q[i].m, i)) { \
+            return false; \
+        } \
+    }
+
+#define VALIDATE_ROW_DATA_E_E8M0_IMPL(type, data, nb) \
+    const type * q = (const type *) (data); \
+    for (size_t i = 0; i < (nb); ++i) { \
+        if (!validate_e_e8m0(q[i].e, i)) { \
+            return false; \
+        } \
+    }
+
+#define VALIDATE_ROW_DATA_DVEC_F16_IMPL(type, data, nb, nr) \
+    const type * q = (const type *) (data); \
+    for (size_t i = 0; i < (nb); ++i) { \
+        for (size_t j = 0; j < (nr); ++j) { \
+            if (!validate_fp16(q[i].d[j], i)) { \
+                return false; \
+            } \
+        } \
+    }
+
+bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes) {
+    if (type < 0 || type >= GGML_TYPE_COUNT) {
+        fprintf(stderr, "%s: invalid type %d\n", __func__, type);
+        return false;
+    }
+
+    if (nbytes % ggml_type_size(type) != 0) {
+        fprintf(stderr, "%s: invalid size %zu for type %s (type size = %zu)\n", __func__, nbytes, ggml_type_name(type), ggml_type_size(type));
+        return false;
+    }
+
+    const size_t nb = nbytes/ggml_type_size(type);
+
+    switch (type) {
+        case GGML_TYPE_BF16:
+            {
+                int nans = 0;
+                int infs = 0;
+                const unsigned short * f = (const unsigned short *) data;
+                for (size_t i = 0; i < nb; ++i) {
+                    nans += (f[i] & 0x7fff) > 0x7f80;
+                    infs += (f[i] & 0x7fff) == 0x7f80;
+                }
+                if (nans) {
+                    fprintf(stderr, "%s: found %d NaNs in row of %zu BF16 values\n", __func__, nans, nb);
+                    return false;
+                }
+                if (infs) {
+                    fprintf(stderr, "%s: found %d infinities in row of %zu BF16 values\n", __func__, infs, nb);
+                    return false;
+                }
+            } break;
+        case GGML_TYPE_F16:
+            {
+                const ggml_fp16_t * f = (const ggml_fp16_t *) data;
+                size_t i = 0;
+#if defined(__AVX2__)
+                for (; i + 15 < nb; i += 16) {
+                    __m256i v = _mm256_loadu_si256((const __m256i *)(f + i));
+                    __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi16(0x7c00));
+                    __m256i cmp = _mm256_cmpeq_epi16(vexp, _mm256_set1_epi16(0x7c00));
+                    int mask = _mm256_movemask_epi8(cmp);
+                    if (mask) {
+                        for (size_t j = 0; j < 16; ++j) {
+                            if (!validate_fp16(f[i + j], i + j)) {
+                                return false;
+                            }
+                        }
+                        GGML_UNREACHABLE();
+                    }
+                }
+#elif defined(__ARM_NEON)
+                for (; i + 7 < nb; i += 8) {
+                    uint16x8_t v = vld1q_u16(f + i);
+                    uint16x8_t vexp = vandq_u16(v, vdupq_n_u16(0x7c00));
+                    uint16x8_t cmp = vceqq_u16(vexp, vdupq_n_u16(0x7c00));
+                    uint64_t mask = vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(cmp, 4)), 0);
+                    if (mask) {
+                        for (size_t j = 0; j < 8; ++j) {
+                            if (!validate_fp16(f[i + j], i + j)) {
+                                return false;
+                            }
+                        }
+                        GGML_UNREACHABLE();
+                    }
+                }
+#endif
+                for (; i < nb; ++i) {
+                    if (!validate_fp16(f[i], i)) {
+                        return false;
+                    }
+                }
+            } break;
+        case GGML_TYPE_F32:
+            {
+                const float * f = (const float *) data;
+                size_t i = 0;
+#if defined(__AVX2__)
+                for (; i + 7 < nb; i += 8) {
+                    __m256i v = _mm256_loadu_si256((const __m256i *)(f + i));
+                    __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi32(0x7f800000));
+                    __m256i cmp = _mm256_cmpeq_epi32(vexp, _mm256_set1_epi32(0x7f800000));
+                    int mask = _mm256_movemask_epi8(cmp);
+                    if (mask) {
+                        for (size_t j = 0; j < 8; ++j) {
+                            if (!validate_float(f[i + j], i + j)) {
+                                return false;
+                            }
+                        }
+                        GGML_UNREACHABLE();
+                    }
+                }
+#elif defined(__ARM_NEON)
+                for (; i + 3 < nb; i += 4) {
+                    uint32x4_t v = vld1q_u32((const uint32_t *)f + i);
+                    uint32x4_t vexp = vandq_u32(v, vdupq_n_u32(0x7f800000));
+                    uint32x4_t cmp = vceqq_u32(vexp, vdupq_n_u32(0x7f800000));
+                    uint64_t mask = vget_lane_u64(vreinterpret_u64_u16(vshrn_n_u32(cmp, 8)), 0);
+                    if (mask) {
+                        for (size_t j = 0; j < 4; ++j) {
+                            if (!validate_float(f[i + j], i + j)) {
+                                return false;
+                            }
+                        }
+                        GGML_UNREACHABLE();
+                    }
+                }
+#endif
+                for (; i < nb; ++i) {
+                    if (!validate_float(f[i], i)) {
+                        return false;
+                    }
+                }
+            } break;
+        case GGML_TYPE_F64:
+            {
+                const double * f = (const double *) data;
+                for (size_t i = 0; i < nb; ++i) {
+                    if (!validate_float(f[i], i)) {
+                        return false;
+                    }
+                }
+            } break;
+        case GGML_TYPE_Q1_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb);
+            } break;
+        case GGML_TYPE_Q2_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q2_0, data, nb);
+            } break;
+        case GGML_TYPE_Q4_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb);
+            } break;
+        case GGML_TYPE_Q4_1:
+            {
+                VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_1, data, nb, d, m);
+            } break;
+        case GGML_TYPE_Q5_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q5_0, data, nb);
+            } break;
+        case GGML_TYPE_Q5_1:
+            {
+                VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_1, data, nb, d, m);
+            } break;
+        case GGML_TYPE_Q8_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q8_0, data, nb);
+            } break;
+        case GGML_TYPE_MXFP4:
+            {
+                VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp4, data, nb);
+            } break;
+        case GGML_TYPE_NVFP4:
+            {
+                // UE4M3 scales are uint8_t — all byte values are valid
+                GGML_UNUSED(data);
+                GGML_UNUSED(nb);
+            } break;
+        case GGML_TYPE_Q2_K:
+            {
+                VALIDATE_ROW_DATA_DM_F16_IMPL(block_q2_K, data, nb, d, dmin);
+            } break;
+        case GGML_TYPE_Q3_K:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q3_K, data, nb);
+            } break;
+        case GGML_TYPE_Q4_K:
+            {
+                VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_K, data, nb, d, dmin);
+            } break;
+        case GGML_TYPE_Q5_K:
+            {
+                VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_K, data, nb, d, dmin);
+            } break;
+        case GGML_TYPE_Q6_K:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q6_K, data, nb);
+            } break;
+        case GGML_TYPE_Q8_K:
+            {
+                const block_q8_K * q = (const block_q8_K *) data;
+                for (size_t i = 0; i < nb; ++i) {
+                    if (!validate_float(q[i].d, i)) {
+                        return false;
+                    }
+                }
+            } break;
+        case GGML_TYPE_TQ1_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_tq1_0, data, nb);
+            } break;
+        case GGML_TYPE_TQ2_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_tq2_0, data, nb);
+            } break;
+        case GGML_TYPE_IQ1_S:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq1_s, data, nb);
+            } break;
+        case GGML_TYPE_IQ1_M:
+            {
+                const block_iq1_m * q = (const block_iq1_m *) data;
+                for (size_t i = 0; i < nb; ++i) {
+                    iq1m_scale_t scale;
+                    const uint16_t * sc = (const uint16_t *)q[i].scales;
+                    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+                    if (!validate_fp16(scale.f16, i)) {
+                        return false;
+                    }
+                }
+            } break;
+        case GGML_TYPE_IQ2_XXS:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xxs, data, nb);
+            } break;
+        case GGML_TYPE_IQ2_XS:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xs, data, nb);
+            } break;
+        case GGML_TYPE_IQ2_S:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_s, data, nb);
+            } break;
+        case GGML_TYPE_IQ3_XXS:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_xxs, data, nb);
+            } break;
+
+        case GGML_TYPE_IQ3_S:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_s, data, nb);
+            } break;
+        case GGML_TYPE_IQ4_XS:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_xs, data, nb);
+            } break;
+        case GGML_TYPE_IQ4_NL:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_nl, data, nb);
+            } break;
+
+        case GGML_TYPE_I8:
+        case GGML_TYPE_I16:
+        case GGML_TYPE_I32:
+        case GGML_TYPE_I64:
+            // nothing to validate
+            break;
+        default:
+            {
+                fprintf(stderr, "%s: invalid type %d\n", __func__, type);
+                return false;
+            }
+    }
+
+    return true;
+}
diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h
new file mode 100644
index 0000000000..75188f1af1
--- /dev/null
+++ b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h
@@ -0,0 +1,115 @@
+#pragma once
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+
+#include "ggml.h"
+
+// GGML internal header
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// NOTE: these functions are defined as GGML_API because they used by the CPU backend
+
+// Quantization
+GGML_API void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k);
+
+GGML_API void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k);
+
+GGML_API void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k);
+
+GGML_API void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k);
+
+GGML_API void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_iq4_nl_ref (const float * GGML_RESTRICT x, block_iq4_nl  * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_iq4_xs_ref (const float * GGML_RESTRICT x, block_iq4_xs  * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_iq3_s_ref  (const float * GGML_RESTRICT x, block_iq3_s   * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_iq2_s_ref  (const float * GGML_RESTRICT x, block_iq2_s   * GGML_RESTRICT y, int64_t k);
+
+// Dequantization
+GGML_API void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+//GGML_API void dequantize_row_q8_1(const block_q8_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+
+GGML_API void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+
+GGML_API void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+
+GGML_API void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+
+GGML_API void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq2_xs (const block_iq2_xs  * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq2_s  (const block_iq2_s   * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq1_s  (const block_iq1_s   * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq1_m  (const block_iq1_m   * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq4_nl (const block_iq4_nl  * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq4_xs (const block_iq4_xs  * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_iq3_s  (const block_iq3_s   * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+
+// Quantization utilizing an importance matrix (a.k.a. "Activation aWare Quantization")
+GGML_API size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq2_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq2_s  (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq1_s  (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq1_m  (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq4_nl (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq4_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_iq3_s  (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+
+GGML_API size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+
+GGML_API size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+
+GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+
+GGML_API void iq2xs_init_impl(enum ggml_type type);
+GGML_API void iq2xs_free_impl(enum ggml_type type);
+GGML_API void iq3xs_init_impl(int grid_size);
+GGML_API void iq3xs_free_impl(int grid_size);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/python/freetoken/kernel/ggml_cpu/__init__.py b/python/freetoken/kernel/ggml_cpu/__init__.py
new file mode 100644
index 0000000000..909c973e3c
--- /dev/null
+++ b/python/freetoken/kernel/ggml_cpu/__init__.py
@@ -0,0 +1,37 @@
+"""ggml-cpu row kernels for GGUF experts on the CPU MoE executor.
+
+setup.py builds one module per x86 variant of llama.cpp's GGML_CPU_ALL_VARIANTS; select() picks the one this CPU runs best.
+"""
+
+from __future__ import annotations
+
+import functools
+import importlib
+import os
+import pkgutil
+from typing import NamedTuple
+
+
+class Variant(NamedTuple):
+    name: str
+    api: int  # address of the variant's kernel table, for _cpu_moe.set_ggml_cpu
+
+
+@functools.cache
+def select() -> Variant | None:
+    """The highest-scoring variant this CPU runs (upstream's rule), or the one FREETOKEN_GGML_CPU_VARIANT names; None when the build has none."""
+    names = sorted(m.name[1:] for m in pkgutil.iter_modules(__path__) if m.name.startswith("_"))
+    forced = os.environ.get("FREETOKEN_GGML_CPU_VARIANT")
+    if forced:
+        if forced not in names:
+            raise RuntimeError(f"FREETOKEN_GGML_CPU_VARIANT={forced}: this build has the ggml-cpu variants {names}")
+        names = [forced]
+    best, best_score = None, 0
+    for name in names:
+        # importing a variant runs no code built for its ISA extensions
+        mod = importlib.import_module(f"{__name__}._{name}")
+        if (score := mod.score()) > best_score:
+            best, best_score = Variant(name, mod.api()), score
+    if forced and best is None:
+        raise RuntimeError(f"FREETOKEN_GGML_CPU_VARIANT={forced}: this CPU lacks an ISA extension the variant is built for")
+    return best
diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py
index dfd545bee6..416c2f4b58 100644
--- a/python/freetoken/layers/quantization/moe/gguf.py
+++ b/python/freetoken/layers/quantization/moe/gguf.py
@@ -142,13 +142,9 @@ def _apply_float(layer, x, topk_weights, topk_ids, view: ExpertView, ggml_type:
 class GGUFMoEMethod(MoEMethod):
     candidates = (GgmlMoEKernel,)
 
-    @property
-    def cpu_format(self) -> str | None:
-        # the CPU executor reads Q4_0 experts in their GGUF rows and BF16 experts as its bf16 banks; other ggml types have no CPU kernel yet
+    def cpu_spec(self) -> tuple[str | None, dict[str, int]]:
         gate_up, down = expert_types(self.cfg)
-        names = {type_name(gate_up), type_name(down)}
-        if names == {"Q4_0"}:
-            return "q4_0"
-        if names == {"BF16"}:
-            return "bf16"
-        return self.kernel.cpu_format
+        # BF16 experts keep the executor's bf16 kernels, which use AVX512-BF16 where the CPU has it; ggml-cpu is built for AVX2 only
+        if float_dtype(gate_up) is torch.bfloat16 and float_dtype(down) is torch.bfloat16:
+            return "bf16", {}
+        return "gguf", {"gate_up_type": gate_up, "down_type": down}
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index 86a03f4175..f7ea5aa307 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -63,7 +63,7 @@
 logger = init_logger(__name__)
 
 # Formats the CPU MoE C++ kernel can compute AND this bench can build banks for; anything
-# else is offload-only here. (The kernel also does q4_0, but this bench has no q4_0 banks.)
+# else is offload-only here.
 _CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4", "fp8_block"})
 # Formats this bench can build synthetic (correctly-sized) banks for.
 _BUILDABLE_FORMATS = frozenset({"bf16", "nvfp4", "fp8_block", "mxfp4_triton", "ds_fp4"})
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index 5d9bd3c9e4..4c0aee7ae5 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -147,6 +147,16 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]:
     return len(reps), list(reps)
 
 
+def _format_label(name: str, opts: tuple) -> str:
+    """One format instance as the startup log names it, e.g. ``gguf(gate_up=IQ2_XXS, down=IQ4_NL)``."""
+    if name == "gguf":
+        from freetoken.layers.quantization.ggml_types import type_name
+
+        gate_up, down = dict(opts)["gate_up_type"], dict(opts)["down_type"]
+        return f"gguf(gate_up={type_name(gate_up)}, down={type_name(down)})"
+    return name + (f"({', '.join(f'{k}={v}' for k, v in opts)})" if opts else "")
+
+
 class CpuMoeExecutor:
     """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks
     (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``)."""
@@ -182,6 +192,15 @@ def __init__(
                     f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are "
                     f"{f!r}; use --moe-strategy offload (GPU-side dequant) instead."
                 )
+            if f == "gguf":
+                from freetoken.kernel import ggml_cpu
+
+                if (variant := ggml_cpu.select()) is None:
+                    raise NotImplementedError(
+                        "--moe-strategy cpu/hybrid cannot compute GGUF experts here: this build has no ggml-cpu kernels (x86-64 only); "
+                        "use --moe-strategy offload instead."
+                    )
+                _cpu_moe.set_ggml_cpu(variant.api)
         if activation not in _ACT_IDS:
             raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}")
         if apply_router_weight_on_input and (output_only := {"mxfp4_triton", "ds_fp4"} & set(names)):
@@ -348,7 +367,7 @@ def __init__(
             self._watchdog.start()
 
 
-        fmt_label = "; ".join(f"{name}{dict(opts) if opts else ''}: {n} layers" for (name, opts), n in layer_counts.items())
+        fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items())
         logger.info_rank0(
             f"CPU MoE executor ready: threads={nthreads} (pinned to cores "
             f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} "
@@ -404,6 +423,8 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
         (see ``_make_table``), not a single bank's -- the C++ side resolves
         ``tbl[layer_id]`` per task.
         """
+        if fmt == "gguf":
+            return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"]))
         if fmt == "bf16":
             gate_up = banks["gate_up"]
             down = banks["down"]
@@ -418,9 +439,6 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
             banks = self._tables(gate_up=gate_up, down=down)
             return banks, {}, (H, I)
 
-        if fmt == "q4_0":
-            return self._resolve_q4_0_banks(banks, rep)
-
         if fmt == "mxfp4_triton":
             return self._resolve_mxfp4_banks(banks, rep)
 
@@ -447,23 +465,18 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
         banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng)
         return banks, {}, (H, I)
 
-    def _resolve_q4_0_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
-        """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble
-        bytes), row-major over K -- the *same* packed banks the GPU offload path streams.
-        gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a
-        row in place (18 bytes / 32 K) and dequantizes weights inside the K-loop."""
+    def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int) -> tuple[dict, dict, tuple[int, int]]:
+        """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype."""
+        from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes, type_name
+
         gate_up, down = banks["gate_up"], banks["down"]
-        assert gate_up[rep].dtype == torch.uint8 and down[rep].dtype == torch.uint8, (
-            gate_up[rep].dtype, down[rep].dtype,
-        )
-        I = int(gate_up[rep].shape[1] // 2)
-        H = int(down[rep].shape[1])
-        assert gate_up[rep].shape[1] == 2 * I
-        assert H % 32 == 0 and I % 32 == 0, (H, I)
-        assert int(gate_up[rep].shape[2]) == (H // 32) * 18, (gate_up[rep].shape, H)
-        assert int(down[rep].shape[2]) == (I // 32) * 18, (down[rep].shape, I)
+        H, I = int(down[rep].shape[1]), int(gate_up[rep].shape[1] // 2)
+        for role, t, ggml_type, k in (("gate_up", gate_up[rep], gate_up_type, H), ("down", down[rep], down_type, I)):
+            dtype = float_dtype(ggml_type) if float_dtype(ggml_type) is not None else torch.uint8
+            if t.dtype != dtype or t.shape[-1] * t.element_size() != row_bytes(ggml_type, k):
+                raise ValueError(f"MoE layer {rep} {role} bank {tuple(t.shape)} {t.dtype} is not {type_name(ggml_type)} rows of {k} values")
         banks = self._tables(gate_up=gate_up, down=down)
-        return banks, {}, (H, I)
+        return banks, {"gate_up_type": gate_up_type, "down_type": down_type}, (H, I)
 
     def _resolve_mxfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales
diff --git a/setup.py b/setup.py
index b968373151..070b8fd5cf 100644
--- a/setup.py
+++ b/setup.py
@@ -2,11 +2,13 @@
 
 import importlib.util
 import os
+import platform
 from pathlib import Path
 
 import sys
 
-from setuptools import setup
+from setuptools import Extension, setup
+from setuptools.command.build_clib import build_clib
 from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension, ROCM_HOME
 
 
@@ -96,8 +98,154 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
 
 _check_toolchain()
 
+# ggml-cpu's row kernels for GGUF experts on the CPU MoE executor: one module per x86 variant of llama.cpp's GGML_CPU_ALL_VARIANTS (ggml/src/CMakeLists.txt), freetoken.kernel.ggml_cpu._<variant>.
+# freetoken.kernel.ggml_cpu.select() imports the best one this CPU runs; each module's entry is built without the variant's ISA flags so that check is safe on any CPU.
+GGML_DIR = "python/freetoken/kernel/csrc/ggml"
+GGML_CPU = platform.machine().lower() in ("x86_64", "amd64")
+_HASWELL = ("SSE42", "AVX", "F16C", "FMA", "AVX2", "BMI2")
+GGML_CPU_VARIANTS = {
+    "x64": (),
+    "sse42": ("SSE42",),
+    "sandybridge": ("SSE42", "AVX"),
+    "ivybridge": ("SSE42", "AVX", "F16C"),
+    "piledriver": ("SSE42", "AVX", "F16C", "FMA"),
+    "haswell": _HASWELL,
+    "skylakex": (*_HASWELL, "AVX512"),
+    "cannonlake": (*_HASWELL, "AVX512", "AVX512_VBMI"),
+    "cascadelake": (*_HASWELL, "AVX512", "AVX512_VNNI"),
+    "icelake": (*_HASWELL, "AVX512", "AVX512_VBMI", "AVX512_VNNI"),
+    "cooperlake": (*_HASWELL, "AVX512", "AVX512_VNNI", "AVX512_BF16"),
+    "zen4": (*_HASWELL, "AVX512", "AVX512_VBMI", "AVX512_VNNI", "AVX512_BF16"),
+    "alderlake": (*_HASWELL, "AVX_VNNI"),
+    "sapphirerapids": (*_HASWELL, "AVX512", "AVX512_VBMI", "AVX512_VNNI", "AVX512_BF16", "AMX_TILE", "AMX_INT8"),
+}
+_GGML_GCC_FLAGS = {
+    "SSE42": ["-msse4.2"], "F16C": ["-mf16c"], "FMA": ["-mfma"], "BMI2": ["-mbmi2"], "AVX": ["-mavx"], "AVX2": ["-mavx2"],
+    "AVX_VNNI": ["-mavxvnni"], "AVX512": ["-mavx512f", "-mavx512cd", "-mavx512vl", "-mavx512dq", "-mavx512bw"],
+    "AVX512_VBMI": ["-mavx512vbmi"], "AVX512_VNNI": ["-mavx512vnni"], "AVX512_BF16": ["-mavx512bf16"],
+    "AMX_TILE": ["-mamx-tile"], "AMX_INT8": ["-mamx-int8"],
+}
+
+
+def _ggml_variant_flags(feats: tuple[str, ...]) -> tuple[list[str], list[str]]:
+    """(ISA flags, macros) of one variant, as upstream's ggml-cpu/CMakeLists.txt sets them; the macros also tell cpu-feats.cpp what to check."""
+    if sys.platform != "win32":
+        return [f for feat in feats for f in _GGML_GCC_FLAGS[feat]], [f"GGML_{feat}" for feat in feats]
+    # cl takes one /arch: level and does not define the macros of the extensions above it
+    if "AVX512" in feats:
+        flags, macros = ["/arch:AVX512"], ["GGML_AVX512"]
+        if "AVX512_VBMI" in feats:
+            macros.append("__AVX512VBMI__")
+        if "AVX512_VNNI" in feats:
+            macros += ["__AVX512VNNI__", "GGML_AVX512_VNNI"]
+    elif "AVX2" in feats:
+        flags, macros = ["/arch:AVX2"], ["GGML_AVX2", "GGML_FMA", "GGML_F16C"]
+    elif "AVX" in feats:
+        flags, macros = ["/arch:AVX"], ["GGML_AVX"]
+    elif "SSE42" in feats:
+        flags, macros = ["/arch:SSE4.2"], ["GGML_SSE42"]
+    else:
+        flags, macros = [], []
+    if "AVX_VNNI" in feats:
+        macros += ["__AVXVNNI__", "GGML_AVX_VNNI"]
+    if "BMI2" in feats:
+        macros += ["__BMI2__", "GGML_BMI2"]
+    return flags, macros
+
+
+def _ggml_cpu_targets() -> tuple[list, list, dict[str, list[str]]]:
+    """(build_clib libraries, variant modules, the static libraries each module links in link order)."""
+    if not GGML_CPU:
+        return [], [], {}
+    if sys.platform == "win32":
+        cflags, c_std, cxx_std, define = ["/O2", "/DNDEBUG", "/utf-8", "/D_CRT_SECURE_NO_WARNINGS"], "/std:c11", "/std:c++20", "/D"
+        # upstream builds these variants only with GCC or Clang
+        variants = {k: v for k, v in GGML_CPU_VARIANTS.items() if k not in ("ivybridge", "piledriver", "cooperlake", "zen4", "sapphirerapids")}
+    else:
+        cflags, c_std, cxx_std, define = ["-O3", "-fPIC", "-fvisibility=hidden", "-DNDEBUG"], "-std=c11", "-std=c++20", "-D"
+        variants = GGML_CPU_VARIANTS
+    build = {
+        "include_dirs": [f"{GGML_DIR}/include", f"{GGML_DIR}/src", f"{GGML_DIR}/ggml-cpu", GGML_DIR],
+        "obj_deps": {"": sorted(str(p.relative_to(ROOT)) for d in ("include", "src", "ggml-cpu", "freetoken") for p in (ROOT / GGML_DIR / d).rglob("*.h"))},
+    }
+    # ggml-base, built once without ISA flags like upstream's; arch/x86/quants.c calls its quantize_row_*_ref
+    libraries = [("freetoken_ggml_base", {"sources": [f"{GGML_DIR}/src/ggml-quants.c"], "cflags": [*cflags, c_std], **build})]
+    modules, module_libs = [], {}
+    for variant, feats in variants.items():
+        flags, macros = _ggml_variant_flags(feats)
+        defines = [define + m for m in macros]
+        libraries += [
+            (f"freetoken_ggml_cpu_{variant}", {
+                "sources": [f"{GGML_DIR}/ggml-cpu/quants.c", f"{GGML_DIR}/ggml-cpu/arch/x86/quants.c", f"{GGML_DIR}/freetoken/cpu_kernels.c"],
+                "cflags": [*cflags, *flags, *defines, c_std],
+                **build,
+            }),
+            # the float row kernels; C++, so a library of their own
+            (f"freetoken_ggml_cpu_vec_{variant}", {"sources": [f"{GGML_DIR}/ggml-cpu/vec.cpp"], "cflags": [*cflags, *flags, *defines, cxx_std], **build}),
+        ]
+        name = f"freetoken.kernel.ggml_cpu._{variant}"
+        module_libs[name] = [f"freetoken_ggml_cpu_{variant}", f"freetoken_ggml_cpu_vec_{variant}", "freetoken_ggml_base"]
+        modules.append(Extension(
+            name=name,
+            sources=[f"{GGML_DIR}/freetoken/variant_module.cpp"],
+            # absolute: ninja compiles from the build directory
+            include_dirs=[str(ROOT / d) for d in (*build["include_dirs"], f"{GGML_DIR}/freetoken")],
+            extra_compile_args=[*cflags, *defines, f"{define}FT_GGML_VARIANT={variant}", cxx_std],
+        ))
+    return libraries, modules, module_libs
+
+
+ggml_cpu_libraries, ggml_cpu_modules, ggml_cpu_module_libs = _ggml_cpu_targets()
+
+
+class BuildClib(build_clib):
+    def build_libraries(self, libraries):
+        # the variants compile the same sources with different flags, so each library gets its own object directory
+        build_temp = self.build_temp
+        try:
+            for lib in libraries:
+                self.build_temp = os.path.join(build_temp, lib[0])
+                super().build_libraries([lib])
+        finally:
+            self.build_temp = build_temp
+
+
+class BuildExt(BuildExtension):
+    def run(self):
+        # build_ext alone (`setup.py build_ext --inplace`) does not run build_clib
+        if self.distribution.has_c_libraries():
+            self.run_command("build_clib")
+        # build_ext would link every build_clib library into every extension; build_extension adds each module's own
+        libraries, self.distribution.libraries = self.distribution.libraries, None
+        try:
+            super().run()
+        finally:
+            self.distribution.libraries = libraries
+
+    def build_extension(self, ext):
+        if ext.name in ggml_cpu_module_libs:
+            build_clib_dir = self.get_finalized_command("build_clib").build_clib
+            if build_clib_dir not in ext.library_dirs:
+                ext.library_dirs.append(build_clib_dir)
+            for name in ggml_cpu_module_libs[ext.name]:
+                if name not in ext.libraries:
+                    ext.libraries.append(name)
+                lib = os.path.join(build_clib_dir, self.compiler.library_filename(name))
+                # relink when only a static library changed
+                if lib not in ext.depends:
+                    ext.depends.append(lib)
+            # every variant module compiles the same source, so each gets its own object directory
+            build_temp, self.build_temp = self.build_temp, os.path.join(self.build_temp, ext.name)
+            try:
+                super().build_extension(ext)
+            finally:
+                self.build_temp = build_temp
+            return
+        super().build_extension(ext)
+
 
 setup(
+    libraries=ggml_cpu_libraries,
     ext_modules=[
         CppExtension(
             name="freetoken.kernel._pinned_tensor",
@@ -126,7 +274,7 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
                     for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp")
                 ),
             ],
-            include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs],
+            include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs, str(ROOT / GGML_DIR / "freetoken")],
             library_dirs=runtime_library_dirs,
             libraries=[runtime_lib],
             extra_compile_args={"cxx": extra_compile + ["-pthread"], "nvcc": device_compile},
@@ -142,6 +290,7 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
                 extra_compile_args=["-O3", "-std=c++20"],
             )
         ] if sys.platform == "linux" else []),
+        *ggml_cpu_modules,
     ],
-    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
+    cmdclass={"build_clib": BuildClib, "build_ext": BuildExt.with_options(use_ninja=True)},
 )
diff --git a/tests/README.md b/tests/README.md
index adc87035d9..85865673e2 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -14,7 +14,7 @@ in its filename (`kvcache/test_dsv4_pool.py`, `models/test_glm4_nvfp4.py`).
 | `tokenizer/` | `freetoken.tokenizer` — tokenize/detokenize request plumbing, thinking-mode resolution |
 | `engine/`    | `freetoken.engine` — cache budget planning, config resolution, backend gating |
 | `kernels/`   | `freetoken.kernel` / `freetoken.layers` — attention, rope, fused ops, JIT cache, pinned memory |
-| `moe/`       | `freetoken.moe` — offload cache, CPU/GPU expert kernels, quantized backends (fp8/nvfp4/mxfp4/q4_0) |
+| `moe/`       | `freetoken.moe` — offload cache, CPU/GPU expert kernels, quantized backends (fp8/nvfp4/mxfp4/gguf) |
 | `models/`    | `freetoken.models` — the registry, and the loading machinery every model shares (sharding, qkv/expert merge, streaming layers into banks) |
 | `e2e/`       | real-server gates: AIME 24/25/26 generation, runtime cache rebuild |
 
diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py
deleted file mode 100644
index 9452f95c7e..0000000000
--- a/tests/moe/test_cpu_moe_q4_0.py
+++ /dev/null
@@ -1,213 +0,0 @@
-"""CPU MoE executor -- native GGUF Q4_0 experts (``--moe-backend cpu``/``hybrid`` for
-the gemma4 GGUF checkpoint).
-
-The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same*
-packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the
-K-loop. We check it against the reference dequant (gguf-py's numpy dequantize) + the
-production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the
-only spread is weight bf16-rounding + reduction order -> tight relative tolerance.
-
-Part 2 covers CUDA-graph capture/replay (the cudaLaunchHostFunc submit/sync nodes
-must recompute from the freshly written pinned routing on each replay).
-"""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-
-def _pack_q4_0(nibbles: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    """Pack [S, OUT, K] uint8 nibble codes (0..15) + [S, OUT, K//32] fp16 scales into
-    the native Q4_0 row layout [S, OUT, K//32*18]: each 32-elem block is a 2-byte fp16
-    scale followed by 16 bytes whose byte j holds elem j (low nibble) and elem j+16
-    (high nibble) -- the exact layout ``dequant_q4_0`` reads back."""
-    S, OUT, K = nibbles.shape
-    nb = K // 32
-    blk = nibbles.reshape(S, OUT, nb, 32)
-    lo = blk[..., :16]      # elems 0..15 -> low nibbles of bytes 0..15
-    hi = blk[..., 16:]      # elems 16..31 -> high nibbles of bytes 0..15
-    packed = (lo | (hi << 4)).to(torch.uint8)                       # [S, OUT, nb, 16]
-    d_bytes = scale.to(torch.float16).view(torch.uint8).reshape(S, OUT, nb, 2)
-    row = torch.cat([d_bytes, packed], dim=-1)                      # [S, OUT, nb, 18]
-    return row.reshape(S, OUT, nb * 18).contiguous()
-
-
-def _dequant_bank(packed: torch.Tensor, K: int, dev) -> torch.Tensor:
-    """[S, OUT, K//32*18] packed Q4_0 -> [S, OUT, K] bf16 (storage order == elem order)."""
-    from freetoken.third_party.gguf import GGMLQuantizationType, quants
-
-    S, OUT, _ = packed.shape
-    values = quants.dequantize(packed.numpy(), GGMLQuantizationType.Q4_0)
-    return torch.from_numpy(values).reshape(S, OUT, K).to(dev, torch.bfloat16)
-
-
-def _make_q4_0_cache(L, E, H, I, seed=0):
-    """Random but valid native Q4_0 banks for the cpu backend (pinned host tensors)."""
-    from freetoken.kernel.pinned import alloc_pinned_tensor
-
-    torch.manual_seed(seed)
-    S = L * E
-
-    def rows(OUT, K):
-        nib = torch.randint(0, 16, (S, OUT, K), dtype=torch.uint8)
-        scale = 0.02 + 0.03 * torch.rand(S, OUT, K // 32)
-        packed = _pack_q4_0(nib, scale)                            # [S, OUT, K//32*18]
-        pinned = alloc_pinned_tensor(*packed.shape, dtype=torch.uint8)
-        pinned.copy_(packed)
-        return pinned
-
-    return SimpleNamespace(
-        quant_format="q4_0",
-        bank_sources={"gate_up": list(rows(2 * I, H).split(E)), "down": list(rows(H, I).split(E))},
-        num_layers=L,
-        num_experts=E,
-        decode_target="cpu",
-        cpu_executor=None,
-    )
-
-
-@pytest.mark.parametrize("bs", [1, 3, 8])
-def test_cpu_decode_q4_0_matches_dequant_then_gpu(bs):
-    """CPU inline-dequant Q4_0 GEMV vs. canonical dequant_q4_0 + bf16 GPU decode."""
-    from freetoken.moe.cpu_executor import CpuMoeExecutor
-    from freetoken.moe.fused import fused_experts_decode_impl
-
-    torch.manual_seed(400 + bs)
-    L, E, H, I, top_k = 3, 16, 2816, 704, 8   # gemma-4-26B-A4B geometry (H,I % 32 == 0)
-    layer = 1
-    dev = torch.device("cuda")
-    cache = _make_q4_0_cache(L, E, H, I)
-
-    ex = CpuMoeExecutor(
-        cache,
-        top_k=top_k,
-        activation="gelu_tanh",
-        apply_router_weight_on_input=False,
-        num_threads=0,
-        max_tokens=bs,
-        device=dev,
-    )
-
-    hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) * 0.5
-    ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32)
-    w = torch.rand(bs, top_k, device=dev, dtype=torch.float32)
-
-    cpu_out = ex.decode(layer, hidden, w, ids).float()
-    torch.cuda.synchronize()
-
-    b = cache.bank_sources
-    gate_up_layer = _dequant_bank(b["gate_up"][layer], H, dev)  # [E, 2I, H]
-    down_layer = _dequant_bank(b["down"][layer], I, dev)        # [E, H, I]
-    gpu_out = fused_experts_decode_impl(
-        hidden, gate_up_layer, down_layer, w, ids.clone(), "gelu_tanh", False
-    ).float()
-
-    rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)
-    assert rel < 2e-2, f"q4_0 bs={bs} rel err {rel.item()}"
-
-
-@pytest.mark.slow
-def test_cpu_decode_q4_0_matches_ggml_mmvq():
-    """Sanity: the CPU W4A16 GEMV lands close to the GPU ggml MMVQ (W4A8) kernel the
-    offload path uses -- a looser tol since MMVQ quantizes activations to int8."""
-    from freetoken.kernel import ggml
-    from freetoken.layers.activation import gated_act_and_mul
-    from freetoken.moe.cpu_executor import CpuMoeExecutor
-    from freetoken.third_party.gguf import GGMLQuantizationType
-
-    torch.manual_seed(77)
-    L, E, H, I, top_k = 2, 16, 2816, 704, 8
-    layer, bs = 1, 4
-    dev = torch.device("cuda")
-    cache = _make_q4_0_cache(L, E, H, I)
-
-    ex = CpuMoeExecutor(
-        cache, top_k=top_k, activation="gelu_tanh",
-        apply_router_weight_on_input=False, num_threads=0, max_tokens=bs, device=dev,
-    )
-
-    hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) * 0.5
-    ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32)
-    w = torch.rand(bs, top_k, device=dev, dtype=torch.float32)
-
-    cpu_out = ex.decode(layer, hidden, w, ids).float()
-    torch.cuda.synchronize()
-
-    b = cache.bank_sources
-    q4_0 = int(GGMLQuantizationType.Q4_0)
-    h = ggml.mul_mat_id(b["gate_up"][layer].to(dev), q4_0, hidden.float(), ids)
-    inter = torch.empty(bs, top_k, I, device=dev)
-    gated_act_and_mul("gelu_tanh", h.view(-1, 2 * I), inter.view(-1, I))
-    y = ggml.mul_mat_id(b["down"][layer].to(dev), q4_0, inter, ids)
-    gpu_out = ggml.moe_weighted_sum(y, w)
-
-    rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)
-    assert rel < 6e-2, f"q4_0 vs ggml-mmvq rel err {rel.item()}"
-
-
-def test_cpu_moe_decode_q4_0_cuda_graph_replay():
-    """Q4_0 CPU path under capture/replay: the host nodes must recompute the GEMV from
-    the freshly written pinned routing on each replay (dep flows through pinned buffers)."""
-    from freetoken.moe.cpu_executor import CpuMoeExecutor
-    from freetoken.moe.fused import fused_experts_decode_impl
-
-    torch.manual_seed(9)
-    L, E, H, I, top_k = 2, 8, 2816, 704, 8
-    layer, bs = 1, 4
-    cache = _make_q4_0_cache(L, E, H, I)
-
-    dev = torch.device("cuda")
-    stream = torch.cuda.Stream()
-    torch.cuda.set_stream(stream)
-
-    ex = CpuMoeExecutor(
-        cache, top_k=top_k, activation="gelu_tanh",
-        apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev,
-    )
-    b = cache.bank_sources
-    gate_up_layer = _dequant_bank(b["gate_up"][layer], H, dev)
-    down_layer = _dequant_bank(b["down"][layer], I, dev)
-
-    def reference(hidden, ids, w):
-        return fused_experts_decode_impl(
-            hidden, gate_up_layer, down_layer, w, ids.clone(), "gelu_tanh", False
-        ).float()
-
-    hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) * 0.5
-    ids = torch.randint(0, E, (bs, top_k), device=dev, dtype=torch.int32)
-    w = torch.rand(bs, top_k, device=dev, dtype=torch.float32)
-
-    ex.decode(layer, hidden, w, ids)  # eager warmup: materialize buffers + task
-    torch.cuda.synchronize()
-
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g, stream=stream):
-        out_cap = ex.decode(layer, hidden, w, ids)
-    torch.cuda.synchronize()
-
-    for it in range(3):
-        torch.manual_seed(600 + it)
-        hidden.copy_(torch.randn(bs, H, dtype=torch.bfloat16) * 0.5)
-        ids.copy_(torch.randint(0, E, (bs, top_k), dtype=torch.int32))
-        w.copy_(torch.rand(bs, top_k, dtype=torch.float32))
-        g.replay()
-        torch.cuda.synchronize()
-        ref = reference(hidden, ids, w)
-        rel = (out_cap.float() - ref).abs().max() / (ref.abs().max() + 1e-6)
-        assert rel < 2e-2, f"q4_0 replay {it} rel err {rel.item()}"
-
-    print("cpu moe cuda graph replay (q4_0) OK")
-
-
-if __name__ == "__main__":
-    for bs in (1, 3, 8):
-        test_cpu_decode_q4_0_matches_dequant_then_gpu(bs)
-        print(f"q4_0 bs={bs} OK")
-    test_cpu_decode_q4_0_matches_ggml_mmvq()
-    print("q4_0 vs ggml-mmvq OK")
-    test_cpu_moe_decode_q4_0_cuda_graph_replay()

From 11c1adb88f58a1a9d0ae47b8cd70ca0c7a4793c4 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 16 Oct 2026 23:56:35 -0700
Subject: [PATCH 118/249] perf(kernel): split triton decode attention by
 context length and sm count (#729)

* perf(kernel): split triton decode attention by context length and sm count

* perf(kernel): merge triton decode splits a tile at a time

- stage 2 loads SPLIT_TILE splits per step and merges them with one
  online-softmax update instead of one split per iteration
- SPLIT_TILE is the launch's split cap rounded up to a power of two,
  at most 8192 fp32 elements and 64 splits; num_warps scales with it
---
 python/freetoken/attention/triton.py        |  65 ++++-----
 python/freetoken/kernel/triton/attention.py | 138 ++++++++++++++------
 tests/kernels/test_triton_attention.py      |  59 +++++++--
 3 files changed, 172 insertions(+), 90 deletions(-)

diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py
index 30bedf4ffa..7f29c481bb 100644
--- a/python/freetoken/attention/triton.py
+++ b/python/freetoken/attention/triton.py
@@ -5,6 +5,7 @@
 
 import torch
 from freetoken.core import Batch, get_global_ctx
+from freetoken.kernel.triton.attention import decode_kv_split_cap
 
 from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
 from .utils import BaseCaptureData
@@ -16,9 +17,9 @@
 @dataclass
 class TritonCaptureData(BaseCaptureData):
     q_to_req: torch.Tensor
+    # flat decode split scratch: each graph batch size views it with its own split count
     attn_logits: torch.Tensor
     attn_lse: torch.Tensor
-    num_kv_splits: torch.Tensor
     swa_page_table: torch.Tensor | None = None
 
     @classmethod
@@ -30,7 +31,7 @@ def create(
         *,
         num_q_heads: int,
         max_head_dim: int,
-        max_kv_splits: int,
+        scratch_rows: int,
         **kwargs,
     ):
         return cls(
@@ -40,22 +41,8 @@ def create(
             cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
             page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
             q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device),
-            attn_logits=torch.empty(
-                (max_bs, num_q_heads, max_kv_splits, max_head_dim),
-                dtype=torch.float32,
-                device=device,
-            ),
-            attn_lse=torch.empty(
-                (max_bs, num_q_heads, max_kv_splits),
-                dtype=torch.float32,
-                device=device,
-            ),
-            num_kv_splits=torch.full(
-                (max_bs,),
-                max_kv_splits,
-                dtype=torch.int32,
-                device=device,
-            ),
+            attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device),
+            attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device),
             **kwargs,
         )
 
@@ -72,7 +59,6 @@ class TritonMetadata(BaseAttnMetadata):
     max_q_len: int
     attn_logits: torch.Tensor | None = None
     attn_lse: torch.Tensor | None = None
-    num_kv_splits: torch.Tensor | None = None
     swa_indices: torch.Tensor | None = None
 
     def get_last_indices(self, bs: int) -> torch.Tensor:
@@ -87,7 +73,6 @@ def __init__(self, config: ModelConfig):
         self.capture: TritonCaptureData | None = None
         self.capture_bs: List[int] = []
         self.max_graph_bs = 0
-        self.max_kv_splits = 8
         self.prefill_tile_min_q = 128
         self.num_q_heads = int(getattr(config, "num_qo_heads", 1))
         kv_groups = getattr(config, "kv_cache_group_specs", lambda: ())()
@@ -95,6 +80,11 @@ def __init__(self, config: ModelConfig):
             (group.head_dim for group in kv_groups),
             default=int(getattr(config, "head_dim", 1)),
         )
+        self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))]
+
+    def _split_cap(self, bs: int) -> int:
+        """The most decode KV splits any layer group takes at this batch size."""
+        return max(decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index) for heads in self.kv_heads)
 
     def _ensure_decode_scratch(
         self,
@@ -103,32 +93,27 @@ def _ensure_decode_scratch(
         num_q_heads: int,
         head_dim: int,
     ) -> None:
+        splits = self._split_cap(bs)
         if (
             metadata.attn_logits is not None
             and metadata.attn_lse is not None
-            and metadata.num_kv_splits is not None
             and metadata.attn_logits.shape[0] >= bs
             and metadata.attn_logits.shape[1] >= num_q_heads
+            and metadata.attn_logits.shape[2] >= splits
             and metadata.attn_logits.shape[3] >= head_dim
         ):
             return
         scratch_head_dim = max(self.max_head_dim, head_dim)
         metadata.attn_logits = torch.empty(
-            (bs, num_q_heads, self.max_kv_splits, scratch_head_dim),
+            (bs, num_q_heads, splits, scratch_head_dim),
             dtype=torch.float32,
             device=self.device,
         )
         metadata.attn_lse = torch.empty(
-            (bs, num_q_heads, self.max_kv_splits),
+            (bs, num_q_heads, splits),
             dtype=torch.float32,
             device=self.device,
         )
-        metadata.num_kv_splits = torch.full(
-            (bs,),
-            self.max_kv_splits,
-            dtype=torch.int32,
-            device=self.device,
-        )
 
     def forward(
         self,
@@ -167,7 +152,6 @@ def forward(
             self._ensure_decode_scratch(metadata, bs, q.shape[1], q.shape[-1])
             assert metadata.attn_logits is not None
             assert metadata.attn_lse is not None
-            assert metadata.num_kv_splits is not None
             return decode_paged_attention(
                 q=q,
                 k_cache=k_cache,
@@ -177,8 +161,6 @@ def forward(
                 q_positions=metadata.q_positions,
                 attn_logits=metadata.attn_logits[:bs],
                 attn_lse=metadata.attn_lse[:bs],
-                num_kv_splits=metadata.num_kv_splits[:bs],
-                max_kv_splits=self.max_kv_splits,
                 sm_scale=scale,
                 sliding_window=spec.sliding_window,
                 sinks=spec.sinks,
@@ -280,7 +262,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
             self.device,
             num_q_heads=max(1, self.num_q_heads),
             max_head_dim=max(1, self.max_head_dim),
-            max_kv_splits=self.max_kv_splits,
+            scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list),
         )
         if self._swa_capture_enabled():
             self.capture.swa_page_table = torch.zeros(
@@ -324,14 +306,20 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None:
         metadata.indices = indices
         metadata.q_to_req = self.capture.q_to_req[: metadata.q_to_req.numel()]
         metadata.q_positions = self.capture.positions[:q_tokens]
-        metadata.attn_logits = self.capture.attn_logits[:bs]
-        metadata.attn_lse = self.capture.attn_lse[:bs]
-        metadata.num_kv_splits = self.capture.num_kv_splits[:bs]
+        metadata.attn_logits, metadata.attn_lse = self._capture_scratch(bs)
+
+    def _capture_scratch(self, bs: int) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self.capture is not None
+        splits, heads, dim = self._split_cap(bs), max(1, self.num_q_heads), max(1, self.max_head_dim)
+        logits = self.capture.attn_logits[: bs * heads * splits * dim].view(bs, heads, splits, dim)
+        lse = self.capture.attn_lse[: bs * heads * splits].view(bs, heads, splits)
+        return logits, lse
 
     def prepare_for_capture(self, batch: Batch) -> None:
         bs = batch.size
         assert bs in self.capture_bs and self.capture is not None
         capture = self.capture
+        attn_logits, attn_lse = self._capture_scratch(bs)
         batch.attn_metadata = TritonMetadata(
             cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1],
             indptr=capture.cu_seqlens_k[: bs + 1],
@@ -341,9 +329,8 @@ def prepare_for_capture(self, batch: Batch) -> None:
             is_decode=True,
             prefix_lens=capture.positions[:bs],
             max_q_len=1,
-            attn_logits=capture.attn_logits[:bs],
-            attn_lse=capture.attn_lse[:bs],
-            num_kv_splits=capture.num_kv_splits[:bs],
+            attn_logits=attn_logits,
+            attn_lse=attn_lse,
             swa_indices=self._capture_swa_indices(),
         )
 
diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index c9087e03c2..d19457cffc 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -7,10 +7,28 @@
 import triton.language as tl
 
 
-_MAX_KV_SPLITS = 8
 _MIN_BLOCK_KV = 32
 
 
+# a long decode request takes splits of at least this many keys, so the partial result each one writes stays small next to the keys it reads
+_MIN_SPLIT_LEN = 128
+# a short request still takes up to this many splits of _MIN_BLOCK_KV keys, so tiny requests stay parallel
+_MIN_SPLITS = 8
+_MAX_KV_SPLITS = 64
+# stage 2 merges at most this many fp32 partial elements per step
+_STAGE2_TILE_ELEMS = 8192
+
+
+def _stage2_launch(max_kv_splits: int, block_dv: int) -> tuple[int, int]:
+    """(splits merged per step, num_warps) for stage 2; the tile is at most the split cap rounded up to a power of two."""
+    tile = min(triton.next_power_of_2(max_kv_splits), max(4, _STAGE2_TILE_ELEMS // block_dv), 64)
+    return tile, max(1, min(8, tile * block_dv // 2048))
+
+
+def _max_request_splits(keys: int) -> int:
+    return max(triton.cdiv(keys, _MIN_SPLIT_LEN), min(_MIN_SPLITS, triton.cdiv(keys, _MIN_BLOCK_KV)), 1)
+
+
 @functools.lru_cache(maxsize=None)
 def _optin_smem_bytes(device_index: int) -> int:
     """Per-block opt-in shared-memory budget for a CUDA device (0 if unavailable)."""
@@ -23,6 +41,12 @@ def _num_sms(device_index: int) -> int:
     return int(torch.cuda.get_device_properties(device_index).multi_processor_count)
 
 
+def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int:
+    """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS]."""
+    head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads)))
+    return max(_MIN_SPLITS, min(_MAX_KV_SPLITS, triton.cdiv(4 * _num_sms(device_index), batch * head_blocks)))
+
+
 def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]:
     """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware.
 
@@ -150,6 +174,13 @@ def _paged_attention_kernel(
 
 
 @triton.jit
+def _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN: tl.constexpr, MIN_SPLITS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr):
+    # same rule as _max_request_splits: long requests split by length, short ones still take up to MIN_SPLITS small splits
+    by_len = tl.maximum(tl.cdiv(effective_len, MIN_SPLIT_LEN), tl.minimum(MIN_SPLITS, tl.cdiv(effective_len, MIN_BLOCK_KV)))
+    return tl.minimum(tl.maximum(by_len, 1), max_kv_splits)
+
+
+@triton.jit(do_not_specialize=["max_kv_splits"])
 def _decode_grouped_stage1_kernel(
     q_ptr,
     k_ptr,
@@ -160,7 +191,7 @@ def _decode_grouped_stage1_kernel(
     q_pos_ptr,
     mid_o_ptr,
     mid_lse_ptr,
-    num_kv_splits_ptr,
+    max_kv_splits,
     stride_qt,
     stride_qh,
     stride_ks,
@@ -181,6 +212,8 @@ def _decode_grouped_stage1_kernel(
     BLOCK_H: tl.constexpr,
     VALID_BLOCK_H: tl.constexpr,
     MIN_BLOCK_KV: tl.constexpr,
+    MIN_SPLIT_LEN: tl.constexpr,
+    MIN_SPLITS: tl.constexpr,
     D: tl.constexpr,
     DV: tl.constexpr,
     SLIDING_WINDOW: tl.constexpr,
@@ -212,7 +245,7 @@ def _decode_grouped_stage1_kernel(
         effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1)
     effective_len = tl.maximum(0, effective_end - effective_start)
 
-    kv_splits = tl.load(num_kv_splits_ptr + batch_id)
+    kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV)
     kv_len_per_split = (
         tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
     )
@@ -275,14 +308,14 @@ def _decode_grouped_stage1_kernel(
         tl.store(mid_lse_ptr + lse_offsets, m_i + tl.log(l_i), mask=mask_h)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["max_kv_splits"])
 def _decode_stage2_kernel(
     mid_o_ptr,
     mid_lse_ptr,
     o_ptr,
     indptr_ptr,
     q_pos_ptr,
-    num_kv_splits_ptr,
+    max_kv_splits,
     sinks_ptr,
     stride_mid_ob,
     stride_mid_oh,
@@ -292,12 +325,14 @@ def _decode_stage2_kernel(
     stride_lse_s,
     stride_ot,
     stride_oh,
-    MAX_KV_SPLITS: tl.constexpr,
     MIN_BLOCK_KV: tl.constexpr,
+    MIN_SPLIT_LEN: tl.constexpr,
+    MIN_SPLITS: tl.constexpr,
     BLOCK_DV: tl.constexpr,
     DV: tl.constexpr,
     SLIDING_WINDOW: tl.constexpr,
     HAS_SINKS: tl.constexpr,
+    SPLIT_TILE: tl.constexpr,
 ):
     batch_id = tl.program_id(0)
     q_head = tl.program_id(1)
@@ -310,13 +345,14 @@ def _decode_stage2_kernel(
         effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1)
     effective_len = tl.maximum(0, effective_end - effective_start)
 
-    kv_splits = tl.load(num_kv_splits_ptr + batch_id)
+    kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV)
     kv_len_per_split = (
         tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
     )
 
     offs_d = tl.arange(0, BLOCK_DV)
     mask_d = offs_d < DV
+    offs_s = tl.arange(0, SPLIT_TILE)
     if HAS_SINKS:
         m_i = tl.load(sinks_ptr + q_head).to(tl.float32)
         l_i = 1.0
@@ -325,26 +361,27 @@ def _decode_stage2_kernel(
         l_i = 0.0
     acc = tl.zeros((BLOCK_DV,), dtype=tl.float32)
 
-    mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh + offs_d
+    mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh
     lse_base = batch_id * stride_lse_b + q_head * stride_lse_h
 
-    for split_id in tl.range(0, MAX_KV_SPLITS, num_stages=2):
-        split_start = kv_len_per_split * split_id
-        split_end = tl.minimum(split_start + kv_len_per_split, effective_len)
-
-        if split_end > split_start:
-            partial = tl.load(
-                mid_o_ptr + mid_base + split_id * stride_mid_os,
-                mask=mask_d,
-                other=0.0,
-            )
-            partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s)
-            m_new = tl.maximum(partial_lse, m_i)
-            alpha = tl.exp(m_i - m_new)
-            beta = tl.exp(partial_lse - m_new)
-            acc = acc * alpha + partial * beta
-            l_i = l_i * alpha + beta
-            m_i = m_new
+    # one load round trip per SPLIT_TILE splits
+    for s0 in range(0, kv_splits, SPLIT_TILE):
+        split_id = s0 + offs_s
+        valid = (split_id < kv_splits) & (kv_len_per_split * split_id < effective_len)
+        partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s, mask=valid, other=-float("inf"))
+        partial = tl.load(
+            mid_o_ptr + mid_base + split_id[:, None] * stride_mid_os + offs_d[None, :],
+            mask=valid[:, None] & mask_d[None, :],
+            other=0.0,
+        )
+        m_new = tl.maximum(m_i, tl.max(partial_lse, axis=0))
+        # a tile of empty splits with no sink leaves m at -inf, where exp(-inf - -inf) is nan
+        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)
+        alpha = tl.exp(m_i - m_safe)
+        beta = tl.exp(partial_lse - m_safe)
+        acc = acc * alpha + tl.sum(partial * beta[:, None], axis=0)
+        l_i = l_i * alpha + tl.sum(beta, axis=0)
+        m_i = m_new
 
     out = tl.where(l_i == 0.0, 0.0, acc / l_i)
     tl.store(
@@ -363,14 +400,17 @@ def decode_paged_attention(
     q_positions: torch.Tensor,
     attn_logits: torch.Tensor,
     attn_lse: torch.Tensor,
-    num_kv_splits: torch.Tensor,
-    max_kv_splits: int,
     sm_scale: float,
     sliding_window: int | None = None,
     sinks: torch.Tensor | None = None,
     out: torch.Tensor | None = None,
+    max_kv_splits: int | None = None,
 ) -> torch.Tensor:
-    """SGLang-style split-k grouped decode attention for one query per request."""
+    """SGLang-style split-k grouped decode attention for one query per request.
+
+    Each request takes ``max(ceil(keys / 128), min(8, ceil(keys / 32)))`` splits, capped by ``max_kv_splits`` (default :func:`decode_kv_split_cap`) and by the split dim of ``attn_logits`` / ``attn_lse``.
+    The kernels compute the split count on the device, so a captured graph serves any context length.
+    """
 
     assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda
     assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3
@@ -382,11 +422,18 @@ def decode_paged_attention(
     assert num_q_heads % num_kv_heads == 0
     assert attn_logits.shape[0] >= batch
     assert attn_logits.shape[1] >= num_q_heads
-    assert attn_logits.shape[2] >= max_kv_splits
     assert attn_logits.shape[3] >= head_dim
     assert attn_lse.shape[0] >= batch
     assert attn_lse.shape[1] >= num_q_heads
-    assert attn_lse.shape[2] >= max_kv_splits
+    if sliding_window is not None and sliding_window <= 0:
+        # the kernels treat a window <= 0 as no window
+        sliding_window = None
+    if max_kv_splits is None:
+        max_kv_splits = decode_kv_split_cap(batch, num_q_heads, num_kv_heads, q.device.index)
+    if sliding_window:
+        # no request attends more than the window, so programs past its splits would only exit
+        max_kv_splits = min(max_kv_splits, _max_request_splits(sliding_window))
+    max_kv_splits = min(max_kv_splits, attn_logits.shape[2], attn_lse.shape[2])
     if sinks is not None:
         assert sinks.is_cuda
         assert sinks.dim() == 1
@@ -404,12 +451,17 @@ def decode_paged_attention(
     block_d = triton.next_power_of_2(head_dim)
     block_dv = triton.next_power_of_2(head_dim)
     grid = (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits)
-    block_n = 32
-    if sliding_window and head_dim <= 256:
+    # a grid that fits in one wave runs about one program per SM, which then needs wider tiles to keep enough loads in flight
+    block_n = 64 if grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index) else _MIN_BLOCK_KV
+    num_stages = 3
+    # stage 1 keeps one K tile and one V tile in shared memory
+    if 2 * block_n * block_d * k_cache.element_size() > 0.8 * _optin_smem_bytes(q.device.index):
+        block_n = _MIN_BLOCK_KV
+    if sliding_window:
         split_len = triton.cdiv(triton.cdiv(sliding_window, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV
-        # A window split of several tiles in a single-wave grid is latency bound; 64-key tiles halve its serial steps.
-        if split_len >= 128 and grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index):
-            block_n = 64
+        if split_len <= block_n:
+            # a full window's split fits in one tile, where a deeper pipeline only adds overhead
+            block_n, num_stages = split_len, 2
 
     _decode_grouped_stage1_kernel[grid](
         q,
@@ -421,7 +473,7 @@ def decode_paged_attention(
         q_positions,
         attn_logits,
         attn_lse,
-        num_kv_splits,
+        max_kv_splits,
         q.stride(0),
         q.stride(1),
         k_cache.stride(0),
@@ -442,19 +494,22 @@ def decode_paged_attention(
         BLOCK_H=block_h,
         VALID_BLOCK_H=valid_block_h,
         MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        MIN_SPLIT_LEN=_MIN_SPLIT_LEN,
+        MIN_SPLITS=_MIN_SPLITS,
         D=head_dim,
         DV=head_dim,
         SLIDING_WINDOW=sliding_window or 0,
         num_warps=4,
-        num_stages=2,
+        num_stages=num_stages,
     )
+    split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv)
     _decode_stage2_kernel[(batch, num_q_heads)](
         attn_logits,
         attn_lse,
         o,
         indptr,
         q_positions,
-        num_kv_splits,
+        max_kv_splits,
         sinks_arg,
         attn_logits.stride(0),
         attn_logits.stride(1),
@@ -464,14 +519,15 @@ def decode_paged_attention(
         attn_lse.stride(2),
         o.stride(0),
         o.stride(1),
-        MAX_KV_SPLITS=max_kv_splits,
         MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        MIN_SPLIT_LEN=_MIN_SPLIT_LEN,
+        MIN_SPLITS=_MIN_SPLITS,
         BLOCK_DV=block_dv,
         DV=head_dim,
         SLIDING_WINDOW=sliding_window or 0,
         HAS_SINKS=sinks is not None,
-        num_warps=4,
-        num_stages=2,
+        SPLIT_TILE=split_tile,
+        num_warps=stage2_warps,
     )
     return o
 
diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py
index a29f621ddc..60a5e27477 100644
--- a/tests/kernels/test_triton_attention.py
+++ b/tests/kernels/test_triton_attention.py
@@ -286,7 +286,6 @@ def test_decode_triton_attention_matches_reference(
         device=device,
     )
     attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device)
-    num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device)
 
     actual = decode_paged_attention(
         q,
@@ -297,8 +296,6 @@ def test_decode_triton_attention_matches_reference(
         q_positions,
         attn_logits,
         attn_lse,
-        num_kv_splits,
-        max_kv_splits,
         sm_scale,
         sliding_window=sliding_window,
     )
@@ -344,11 +341,10 @@ def test_decode_triton_attention_non_pow2_group(num_q_heads: int, num_kv_heads:
         batch, num_q_heads, max_kv_splits, head_dim, dtype=torch.float32, device=device
     )
     attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device)
-    num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device)
 
     actual = decode_paged_attention(
         q, k_cache, v_cache, indptr, indices, q_positions,
-        attn_logits, attn_lse, num_kv_splits, max_kv_splits, sm_scale,
+        attn_logits, attn_lse, sm_scale,
     )
     expected = _reference_paged_attention(
         q, k_cache, v_cache, indptr, indices, q_to_req, q_positions, sm_scale, None,
@@ -392,7 +388,6 @@ def test_decode_triton_attention_with_sinks_matches_reference():
         device=device,
     )
     attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device)
-    num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device)
 
     actual = decode_paged_attention(
         q,
@@ -403,8 +398,6 @@ def test_decode_triton_attention_with_sinks_matches_reference():
         q_positions,
         attn_logits,
         attn_lse,
-        num_kv_splits,
-        max_kv_splits,
         sm_scale,
         sinks=sinks,
     )
@@ -424,6 +417,54 @@ def test_decode_triton_attention_with_sinks_matches_reference():
     torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA")
+@pytest.mark.parametrize(
+    ("seq_lens", "num_kv_heads", "sliding_window", "scratch_splits", "use_sinks"),
+    [
+        ([1, 129, 1000, 8193], 2, None, 64, False),
+        ([1, 129, 1000, 8193], 2, None, 8, False),
+        ([300, 2000, 70], 8, 300, 64, True),
+    ],
+)
+def test_decode_triton_attention_splits_long_requests(
+    seq_lens: list[int],
+    num_kv_heads: int,
+    sliding_window: int | None,
+    scratch_splits: int,
+    use_sinks: bool,
+):
+    """Requests of many splits next to short ones, a scratch smaller than the split cap, and windows that span several splits."""
+    from freetoken.kernel.triton.attention import decode_paged_attention
+
+    torch.manual_seed(5)
+    device = torch.device("cuda")
+    batch = len(seq_lens)
+    num_q_heads = 16
+    head_dim = 256
+    total_kv = sum(seq_lens)
+    q = torch.randn(batch, num_q_heads, head_dim, device=device, dtype=torch.bfloat16)
+    k_cache = torch.randn(total_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16)
+    v_cache = torch.randn(total_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16)
+    indptr = torch.tensor([0, *seq_lens], dtype=torch.int32, device=device).cumsum(0).to(torch.int32)
+    indices = torch.randperm(total_kv, device=device).to(torch.int32)
+    q_positions = torch.tensor([n - 1 for n in seq_lens], dtype=torch.int64, device=device)
+    q_to_req = torch.arange(batch, dtype=torch.int32, device=device)
+    sinks = torch.randn(num_q_heads, dtype=torch.float32, device=device) if use_sinks else None
+    sm_scale = head_dim**-0.5
+    attn_logits = torch.empty(batch, num_q_heads, scratch_splits, head_dim, dtype=torch.float32, device=device)
+    attn_lse = torch.empty(batch, num_q_heads, scratch_splits, dtype=torch.float32, device=device)
+
+    actual = decode_paged_attention(
+        q, k_cache, v_cache, indptr, indices, q_positions,
+        attn_logits, attn_lse, sm_scale,
+        sliding_window=sliding_window, sinks=sinks,
+    )
+    expected = _reference_paged_attention(
+        q, k_cache, v_cache, indptr, indices, q_to_req, q_positions, sm_scale, sliding_window, sinks=sinks,
+    )
+    torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA")
 @pytest.mark.parametrize("use_split_inputs", [False, True])
 @pytest.mark.parametrize(
@@ -837,10 +878,8 @@ def __init__(self, device: torch.device):
     assert replay_metadata.q_positions.data_ptr() == backend.capture.positions.data_ptr()
     assert replay_metadata.attn_logits is not None
     assert replay_metadata.attn_lse is not None
-    assert replay_metadata.num_kv_splits is not None
     assert replay_metadata.attn_logits.data_ptr() == backend.capture.attn_logits.data_ptr()
     assert replay_metadata.attn_lse.data_ptr() == backend.capture.attn_lse.data_ptr()
-    assert replay_metadata.num_kv_splits.data_ptr() == backend.capture.num_kv_splits.data_ptr()
 
     torch.testing.assert_close(backend.capture.cu_seqlens_k[:3], expected_indptr)
     torch.testing.assert_close(backend.capture.page_table.view(-1)[:8], expected_indices)

From 6b8f2004547d202bef16a10f5a20aec14f174f66 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 17 Oct 2026 12:56:19 -0700
Subject: [PATCH 119/249] feat(gguf): serve GLM-4.5, GLM-4.6 and GLM-4.7 GGUF
 checkpoints (#731)

---
 python/freetoken/kernel/aot_models.py         |   1 +
 python/freetoken/models/gguf/config.py        |   1 +
 python/freetoken/models/gguf/tokenizer.py     |  16 +-
 python/freetoken/models/glm4_moe/__init__.py  |   4 +
 python/freetoken/models/glm4_moe/attention.py |  20 +-
 python/freetoken/models/glm4_moe/gguf.py      | 179 ++++++++++++++++++
 python/freetoken/models/glm4_moe/model.py     |  17 +-
 python/freetoken/models/register.py           |   8 +
 8 files changed, 231 insertions(+), 15 deletions(-)
 create mode 100644 python/freetoken/models/glm4_moe/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index f0b70e2024..1c367aa433 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -253,6 +253,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         top_k=8,
         moe_intermediate_size=1536,
         expert_formats=("bf16",),
+        arch_aliases=("Glm4MoeGGUFForCausalLM",),
     ),
     AotModel(
         name="nvidia/GLM-4.7-NVFP4",
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 677adff621..5b780fe0ef 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -18,6 +18,7 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "glm4moe": "Glm4MoeGGUFForCausalLM",
     "llama": "LlamaGGUFForCausalLM",
     "minimax-m2": "MiniMaxM2GGUFForCausalLM",
     "minimax-m3": "MiniMaxM3GGUFForCausalLM",
diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index 7f0bf29a48..3945c05db6 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -22,8 +22,8 @@
     "llama3", "llama-v3", "llama-bpe", "falcon3", "falcon-h1", "pixtral", "midm-2.0", "lfm2", "jina-v5-nano", "tekken", "chameleon",
 ))
 
-# Llama 3 pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it
-_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe"))
+# Llama 3 and GLM pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it
+_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe", "glm4"))
 
 
 def _add_bos(meta: dict) -> bool:
@@ -80,7 +80,17 @@ def load_gguf_tokenizer(model_path: str):
     fast, _ = convert_gguf_tokenizer(arch, tok_dict)
     meta = load_gguf_metadata(model_path)
     fast = _fix_added_tokens(fast, meta)
-    if meta.get("tokenizer.ggml.model") == "gpt2" and meta.get("tokenizer.ggml.pre") in _IGNORE_MERGES_PRE:
+    pre = meta.get("tokenizer.ggml.pre") if meta.get("tokenizer.ggml.model") == "gpt2" else None
+    if pre == "glm4":
+        from tokenizers import Regex, pre_tokenizers
+        from transformers.integrations.gguf.gguf_tokenizer_mapping import GGUF_PRE_TOKENIZER_SPLITS
+
+        # transformers maps no split for glm4, which the HF tokenizer and llama.cpp split like Llama 3
+        fast.pre_tokenizer = pre_tokenizers.Sequence([
+            pre_tokenizers.Split(Regex(GGUF_PRE_TOKENIZER_SPLITS["llama3"]), behavior="isolated"),
+            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
+        ])
+    if pre in _IGNORE_MERGES_PRE:
         fast.model.ignore_merges = True
     # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template
     return PreTrainedTokenizerFast(
diff --git a/python/freetoken/models/glm4_moe/__init__.py b/python/freetoken/models/glm4_moe/__init__.py
index b8748b1d44..7f093c6762 100644
--- a/python/freetoken/models/glm4_moe/__init__.py
+++ b/python/freetoken/models/glm4_moe/__init__.py
@@ -1,4 +1,5 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Glm4MoeForCausalLM
 from .weight import iter_weights, nvfp4_expert_spec
 
@@ -7,4 +8,7 @@
     "Glm4MoeForCausalLM",
     "parse_config",
     "iter_weights",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]
diff --git a/python/freetoken/models/glm4_moe/attention.py b/python/freetoken/models/glm4_moe/attention.py
index 7a4b335ef1..65b70bf155 100644
--- a/python/freetoken/models/glm4_moe/attention.py
+++ b/python/freetoken/models/glm4_moe/attention.py
@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, RMSNorm
+from freetoken.layers import BaseOP, LinearReplicated, RMSNorm
 from freetoken.layers.rotary import get_rope
 from freetoken.utils import nvtx_annotate
 
@@ -27,9 +27,10 @@ class Glm4MoeAttention(BaseOP):
     every ``self_attn`` out of quantization (qkvo stay bf16), so we reproduce that exactly by
     storing the bf16 weights compressed (~10.7 bits/weight) and decoding them bit-for-bit in
     the forward. This fits a 32 GB VRAM target without the accuracy loss fp8 attention would add.
+    From a GGUF file the projections are the file's packed ggml linears instead.
     """
 
-    def __init__(self, config: ModelConfig, layer_id: int):
+    def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         head_dim = config.head_dim
         self.layer_id = layer_id
         self.num_qo_heads = config.num_qo_heads
@@ -38,10 +39,17 @@ def __init__(self, config: ModelConfig, layer_id: int):
         self.kv_attn_dim = self.num_kv_heads * head_dim
         self.head_dim = head_dim
 
-        self.q_proj = LinearDF11(config.hidden_size, self.qo_attn_dim, has_bias=config.has_attn_bias)
-        self.k_proj = LinearDF11(config.hidden_size, self.kv_attn_dim, has_bias=config.has_attn_bias)
-        self.v_proj = LinearDF11(config.hidden_size, self.kv_attn_dim, has_bias=config.has_attn_bias)
-        self.o_proj = LinearDF11(self.qo_attn_dim, config.hidden_size, has_bias=False)
+        gguf = config.quant is not None and config.quant.dialect == "gguf"
+
+        def linear(in_features: int, out_features: int, has_bias: bool, name: str) -> BaseOP:
+            if gguf:
+                return LinearReplicated(in_features, out_features, has_bias, quant_config=config.quant, prefix=f"{prefix}.{name}")
+            return LinearDF11(in_features, out_features, has_bias=has_bias)
+
+        self.q_proj = linear(config.hidden_size, self.qo_attn_dim, config.has_attn_bias, "q_proj")
+        self.k_proj = linear(config.hidden_size, self.kv_attn_dim, config.has_attn_bias, "k_proj")
+        self.v_proj = linear(config.hidden_size, self.kv_attn_dim, config.has_attn_bias, "v_proj")
+        self.o_proj = linear(self.qo_attn_dim, config.hidden_size, False, "o_proj")
 
         if config.use_qk_norm:
             self.q_norm = RMSNorm(head_dim, eps=config.rms_norm_eps)
diff --git a/python/freetoken/models/glm4_moe/gguf.py b/python/freetoken/models/glm4_moe/gguf.py
new file mode 100644
index 0000000000..a477706b86
--- /dev/null
+++ b/python/freetoken/models/glm4_moe/gguf.py
@@ -0,0 +1,179 @@
+"""GLM-4.5 / 4.6 / 4.7 from a llama.cpp ``glm4moe`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
+
+The llama.cpp converter (conversion/glm.py) only renames the tensors and stacks the experts: both sides run NeoX rope on the first ``rope.dimension_count`` dims of a head, so q/k carry no rope permute.
+Files that keep the MTP layer append it as the last block, which is not loaded.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "glm4moe"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+    "mlp.gate": "ffn_gate_inp",
+    "mlp.shared_experts.gate_proj": "ffn_gate_shexp",
+    "mlp.shared_experts.up_proj": "ffn_up_shexp",
+    "mlp.shared_experts.down_proj": "ffn_down_shexp",
+    "mlp.experts.0.gate_proj": "ffn_gate_exps",
+    "mlp.experts.0.up_proj": "ffn_up_exps",
+    "mlp.experts.0.down_proj": "ffn_down_exps",
+}
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+# llama.cpp's LLAMA_EXPERT_GATING_FUNC_TYPE_*: 0 (unset) means sigmoid for this arch
+_SIGMOID_GATING = (0, 2)
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    tensor = _LAYER_LEAVES.get(m.group(2))
+    return (f"blk.{m.group(1)}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    md = shim.metadata
+    if any(k.startswith("prism.hadamard.") for k in md):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    names = sorted(gguf_tensor_index(shim.model_path))
+    if not names:
+        raise NotImplementedError(f"{shim.model_path}: a metadata-only glm4moe GGUF has no tensor table to tell whether the attention has q/k norms and biases")
+    if "blk.0.attn_q.weight" not in names:
+        raise NotImplementedError(f"{shim.model_path}: no decoder layer tensors (blk.0.attn_q), e.g. an MTP-only draft GGUF")
+    scaled = [n for n in names if n.endswith((".scale", ".input_scale"))]
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    fused_qkv = [n for n in names if ".attn_qkv." in n]
+    if fused_qkv:
+        raise NotImplementedError(f"{shim.model_path}: a fused attn_qkv glm4moe GGUF is not supported (e.g. {fused_qkv[0]})")
+    gating = int(md.get(f"{_ARCH}.expert_gating_func", 0))
+    if gating not in _SIGMOID_GATING:
+        raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported, GLM routes with sigmoid")
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    index = gguf_tensor_index(shim.model_path)
+    head_dim = int(kv("attention.key_length"))
+    return SimpleNamespace(
+        model_type="glm4_moe",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        # block_count also counts the MTP layers a GGUF may carry after the decoder
+        num_hidden_layers=int(kv("block_count")) - int(md.get(f"{_ARCH}.nextn_predict_layers", 0)),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=head_dim,
+        rope_theta=float(kv("rope.freq_base")),
+        partial_rotary_factor=int(kv("rope.dimension_count")) / head_dim,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+        n_routed_experts=int(kv("expert_count")),
+        num_experts_per_tok=int(kv("expert_used_count")),
+        moe_intermediate_size=int(kv("expert_feed_forward_length")),
+        n_shared_experts=int(kv("expert_shared_count")),
+        first_k_dense_replace=int(md.get(f"{_ARCH}.leading_dense_block_count", 0)),
+        # llama.cpp neither renormalizes nor scales the router weights when the keys are absent
+        norm_topk_prob=bool(md.get(f"{_ARCH}.expert_weights_norm", False)),
+        routed_scaling_factor=float(md.get(f"{_ARCH}.expert_weights_scale", 0.0)) or 1.0,
+        n_group=max(1, int(md.get(f"{_ARCH}.expert_group_count", 1))),
+        topk_group=max(1, int(md.get(f"{_ARCH}.expert_group_used_count", 1))),
+        # GLM-4.5-Air has no q/k norms; the converter writes no flag for them or for the qkv biases
+        use_qk_norm="blk.0.attn_q_norm.weight" in index,
+        attention_bias="blk.0.attn_q.bias" in index,
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return replace(parse_config(_hf_config(shim)), expert_quant="gguf")
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    def linear(module: str, tensor: str) -> Iterator[tuple[str, torch.Tensor]]:
+        info = t(tensor)
+        yield from fused(module, [info], [matrix(info)], int(info.ne[0]))
+        if f"{tensor}.bias" in index:
+            yield f"{module}.bias", float_tensor(index[f"{tensor}.bias"])
+
+    yield from linear("model.embed_tokens", "token_embd")
+    if not hf.tie_word_embeddings:
+        yield from linear("lm_head", "output")
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.post_attention_norm"))
+        sa = f"{pre}.self_attn"
+        for proj, tensor in (("q_proj", "attn_q"), ("k_proj", "attn_k"), ("v_proj", "attn_v"), ("o_proj", "attn_output")):
+            yield from linear(f"{sa}.{proj}", f"{blk}.{tensor}")
+        if hf.use_qk_norm:
+            yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+            yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+        mlp = f"{pre}.mlp"
+        if layer < hf.first_k_dense_replace:
+            for proj in ("gate", "up", "down"):
+                yield from linear(f"{mlp}.{proj}_proj", f"{blk}.ffn_{proj}")
+            continue
+        yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+        yield f"{mlp}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias"))
+        for proj in ("gate", "up", "down"):
+            yield from linear(f"{mlp}.shared_experts.{proj}_proj", f"{blk}.ffn_{proj}_shexp")
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/glm4_moe/model.py b/python/freetoken/models/glm4_moe/model.py
index b5c647e146..8824dfb54b 100644
--- a/python/freetoken/models/glm4_moe/model.py
+++ b/python/freetoken/models/glm4_moe/model.py
@@ -4,7 +4,7 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused
+from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused, VocabParallelEmbedding
 from freetoken.models.blocks import BaseLLMModel
 from freetoken.utils import nvtx_annotate
 
@@ -19,7 +19,7 @@
 
 class Glm4MoeDecoderLayer(BaseOP):
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
-        self.self_attn = Glm4MoeAttention(config, layer_id)
+        self.self_attn = Glm4MoeAttention(config, layer_id, prefix=f"{prefix}.self_attn")
         if layer_id >= config.first_k_dense_replace:
             self.mlp: BaseOP = Glm4MoeSparseBlock(config, layer_id, prefix=f"{prefix}.mlp")
         else:
@@ -45,10 +45,15 @@ def forward(
 
 class Glm4MoeModel(BaseOP):
     def __init__(self, config: ModelConfig, *, prefix: str = "model"):
-        self.embed_tokens = EmbeddingDF11(
-            num_embeddings=config.vocab_size,
-            embedding_dim=config.hidden_size,
-        )
+        if config.quant is not None and config.quant.dialect == "gguf":
+            self.embed_tokens: BaseOP = VocabParallelEmbedding(
+                config.vocab_size, config.hidden_size, quant_config=config.quant, prefix=f"{prefix}.embed_tokens"
+            )
+        else:
+            self.embed_tokens = EmbeddingDF11(
+                num_embeddings=config.vocab_size,
+                embedding_dim=config.hidden_size,
+            )
         self.layers = OPList(
             [
                 Glm4MoeDecoderLayer(config, layer_id, prefix=f"{prefix}.layers.{layer_id}")
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index bfa0a1d6ce..858e17db3f 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -387,6 +387,14 @@ class ModelSpec:
         "Glm4MoeForCausalLM",
         packed_modules_mapping=_EXPERTS_PACKED,
     ),
+    # llama.cpp glm4moe GGUF (GLM-4.5, 4.5-Air, 4.6, 4.7): the same model classes with every tensor in its packed ggml type
+    "Glm4MoeGGUFForCausalLM": ModelSpec(
+        "freetoken.models.glm4_moe",
+        "Glm4MoeForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        packed_modules_mapping=_EXPERTS_PACKED,
+    ),
     # GLM-5.2 (model_type glm_moe_dsa): DeepSeek-V3.2-class MLA + DSA sparse attention
     # with GLM-4-style sigmoid/noaux_tc MoE routing; NVFP4 routed experts served from
     # the offload cache.

From 31c6ef4f7709a802998260e8a19cf6ef6232946d Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 17 Oct 2026 13:46:50 -0700
Subject: [PATCH 120/249] feat(gguf): apply the weight scales of nvfp4 gguf
 tensors (#732)

---
 python/freetoken/engine/engine.py             |  3 +
 .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp       |  2 +
 .../kernel/csrc/cpu_moe/kernels/gguf.cpp      | 25 +++++--
 python/freetoken/kernel/csrc/ggml/VERSION     |  1 +
 .../kernel/csrc/ggml/freetoken.patch          | 68 ++++++++++++++++++-
 .../kernel/csrc/ggml/freetoken/ggml_ops.cu    | 33 ++++++++-
 .../kernel/csrc/ggml/ggml-cuda/common.cuh     |  3 +
 .../kernel/csrc/ggml/ggml-cuda/mmvq.cu        | 15 ++--
 python/freetoken/kernel/ggml.py               | 23 +++++--
 .../layers/quantization/configs/gguf.py       | 24 ++++---
 .../layers/quantization/embedding/gguf.py     |  7 +-
 .../layers/quantization/ggml_types.py         | 32 ++++++++-
 .../layers/quantization/linear/gguf.py        |  6 ++
 .../freetoken/layers/quantization/moe/gguf.py | 28 ++++++--
 .../freetoken/layers/quantization/scheme.py   |  8 ++-
 python/freetoken/models/gemma4/gguf.py        |  6 +-
 python/freetoken/models/gguf/experts.py       | 37 ++++++++--
 python/freetoken/models/gguf/reader.py        |  4 +-
 python/freetoken/models/gguf/tensors.py       | 52 ++++++++++++--
 python/freetoken/models/llama/gguf.py         |  6 +-
 python/freetoken/models/qwen3/gguf.py         |  6 +-
 python/freetoken/models/qwen3_5_moe/gguf.py   |  6 +-
 python/freetoken/models/qwen3_moe/gguf.py     |  6 +-
 python/freetoken/models/qwen4_exp/gguf.py     |  4 +-
 python/freetoken/moe/cpu_executor.py          | 25 ++++---
 python/freetoken/moe/offload_cache.py         |  2 +
 26 files changed, 356 insertions(+), 76 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 326c467d6d..ee6558014d 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1770,6 +1770,9 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             "--dtype float16 with MXFP8 resident weights is unsupported (the "
             "W8A16 fold is only validated exact in bfloat16); use bfloat16."
         )
+    if _dtype == torch.float16 and getattr(getattr(model_config, "quant", None), "weight_scales", False):
+        # GGUF NVFP4 results are stored in the activation dtype before weight_scale_2 (~1e-4) multiplies them, which overflows fp16.
+        raise ValueError("--dtype float16 with NVFP4 GGUF weight scales is unsupported (the unscaled results overflow float16); use bfloat16.")
     if config.attention_backend == "auto":
         override(
             "attention_backend",
diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
index c4887efbff..8162138ceb 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
@@ -48,6 +48,7 @@ std::unique_ptr<MoeFormat> make_mxfp4_format(const MoeParams& p, const FormatArg
 std::unique_ptr<MoeFormat> make_ds_fp4_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_fp8_block_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a);
 void set_ggml_cpu(uintptr_t addr);
 
 namespace {
@@ -97,6 +98,7 @@ inline MoeTask from_word(uint64_t w) {
 const std::pair<const char*, std::unique_ptr<MoeFormat> (*)(const MoeParams&, const FormatArgs&)> kFormats[] = {
     {"bf16", make_bf16_format},     {"nvfp4", make_nvfp4_format},         {"mxfp4_triton", make_mxfp4_format},
     {"ds_fp4", make_ds_fp4_format}, {"fp8_block", make_fp8_block_format}, {"gguf", make_gguf_format},
+    {"gguf_scaled", make_gguf_scaled_format},
 };
 
 std::unique_ptr<MoeFormat> make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) {
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
index c1641af51d..63c8107dd3 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp
@@ -31,10 +31,11 @@ void set_ggml_cpu(uintptr_t addr) {
 namespace {
 
 // Packed ggml rows of the layers whose experts have this gate_up type and this down type, dotted by ggml-cpu against the activation converted to each type's vec_dot input (q8_0, q8_1 or q8_K, or the float type itself).
+// With scales, each expert's (gate, up, down, 1) NVFP4 weight_scale_2 multiplies its gate and up rows before the glu and its down rows before the routing weight, as llama.cpp's build_moe_ffn does.
 class GgufFormat final : public MoeFormat {
  public:
-  GgufFormat(const MoeParams& p, const FormatArgs& a)
-      : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {
+  GgufFormat(const MoeParams& p, const FormatArgs& a, bool scaled)
+      : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")), scales_(scaled ? a.bank("global") : nullptr) {
     static std::once_flag tables;
     std::call_once(tables, ggml_.init);
     gu_ = kernel(static_cast<int>(a.option("gate_up_type")), p.H, "gate_up");
@@ -63,10 +64,11 @@ class GgufFormat final : public MoeFormat {
   void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, bf16_t* g) const override {
     const uint8_t* w = static_cast<const uint8_t*>(tbl_at(gate_up_, layer)) + static_cast<size_t>(e) * 2 * p.I * gu_.row_bytes;
     const uint8_t* xq = xq_.data() + static_cast<size_t>(tok) * gu_.act_row_bytes;
+    const float* sc = expert_scales(layer, e);
     gated_rows(p, i0, i1, w_in, g, [&](int row) {
       float s;
       gu_.vec_dot(p.H, &s, 0, w + static_cast<size_t>(row) * gu_.row_bytes, 0, xq, 0, 1);
-      return s;
+      return sc == nullptr ? s : s * sc[row < p.I ? 0 : 1];
     });
   }
 
@@ -76,9 +78,19 @@ class GgufFormat final : public MoeFormat {
       for (int j = 0; j < n; ++j)
         dn_.vec_dot(p.I, &out[j * kRowBlock + h - h0], 0, w + (static_cast<size_t>(routes[j].e) * p.H + h) * dn_.row_bytes, 0,
                     gq_.data() + static_cast<size_t>(routes[j].r) * dn_.act_row_bytes, 0, 1);
+    if (scales_ == nullptr) return;
+    for (int j = 0; j < n; ++j) {
+      const float s = expert_scales(layer, routes[j].e)[2];
+      for (int h = h0; h < h1; ++h) out[j * kRowBlock + h - h0] *= s;
+    }
   }
 
  private:
+  // the expert's (gate, up, down, 1) row of the layer's [E, 4] fp32 scale bank, or nullptr without scales
+  const float* expert_scales(int layer, int e) const {
+    return scales_ == nullptr ? nullptr : static_cast<const float*>(tbl_at(scales_, layer)) + static_cast<size_t>(e) * 4;
+  }
+
   ft_ggml_row_kernel kernel(int type, int k, const char* role) const {
     ft_ggml_row_kernel out{};
     const int rc = ggml_.row_kernel(type, k, &out);
@@ -92,6 +104,7 @@ class GgufFormat final : public MoeFormat {
   const ft_ggml_cpu_api& ggml_;
   const uint64_t* gate_up_;
   const uint64_t* down_;
+  const uint64_t* scales_;
   ft_ggml_row_kernel gu_{}, dn_{};
   std::vector<float> xf_;
   // the activation in the vec_dot input type: x per token, g per route
@@ -101,7 +114,11 @@ class GgufFormat final : public MoeFormat {
 }  // namespace
 
 std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a) {
-  return std::make_unique<GgufFormat>(p, a);
+  return std::make_unique<GgufFormat>(p, a, false);
+}
+
+std::unique_ptr<MoeFormat> make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<GgufFormat>(p, a, true);
 }
 
 }  // namespace freetoken::cpu_moe
diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION
index 3c66f9b7d0..f7b5c01705 100644
--- a/python/freetoken/kernel/csrc/ggml/VERSION
+++ b/python/freetoken/kernel/csrc/ggml/VERSION
@@ -10,5 +10,6 @@ Local changes, kept in freetoken.patch (already applied):
 - mmvq.cu/.cuh: MMVQ takes bf16 or fp16 activations, can write bf16 or fp16 results, and can reuse activations already quantized to q8_1 (src1_q8_1_pre).
 - mmq.cu/.cuh, mmid.cu/.cuh: MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold tokens, listed by mm_ids_tile_map; ncols_max > 0 keeps the regular grid.
 - common.cuh: the native fp4 MMA (BLACKWELL_MMA_AVAILABLE, blackwell_mma_available) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build.
+- mmvq.cu, common.cuh: the NVFP4 scale fusion (x_scale, gate_scale) reads strided scale vectors, e.g. columns of a per-expert scale bank.
 
 To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file.
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken.patch b/python/freetoken/kernel/csrc/ggml/freetoken.patch
index 3dc715fe0f..595498dd15 100644
--- a/python/freetoken/kernel/csrc/ggml/freetoken.patch
+++ b/python/freetoken/kernel/csrc/ggml/freetoken.patch
@@ -24,6 +24,16 @@
  }
  
  // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+@@ -1582,6 +1588,9 @@
+     const void * gate_bias = nullptr;
+     const void * x_scale = nullptr;
+     const void * gate_scale = nullptr;
++    // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank
++    uint32_t x_scale_stride = 1;
++    uint32_t gate_scale_stride = 1;
+     ggml_glu_op glu_op;
+     float glu_limit = 0.0f;
+     const void * shared_up = nullptr;
 --- a/ggml/src/ggml-cuda/mmid.cu
 +++ b/ggml/src/ggml-cuda/mmid.cu
 @@ -177,3 +177,45 @@
@@ -572,6 +582,19 @@
          stride_col_dst = fusion.shared_stride_col_dst;
      }
  
+@@ -702,10 +706,10 @@
+             }
+             if constexpr (type == GGML_TYPE_NVFP4) {
+                 if (use_scale) {
+-                    x_scales = x_scale[ids ? channel_x : 0];
++                    x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride];
+                 }
+                 if (use_gate_scale) {
+-                    gate_scales = gate_scale[ids ? channel_x : 0];
++                    gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride];
+                 }
+             }
+         }
 @@ -837,7 +841,7 @@
                          }
                      }
@@ -615,6 +638,24 @@
          stride_col_dst = fusion.shared_stride_col_dst;
      }
  
+@@ -963,7 +969,7 @@
+ 
+             if constexpr (type == GGML_TYPE_NVFP4) {
+                 if (x_scale) {
+-                    result *= x_scale[channel_x];
++                    result *= x_scale[channel_x * fusion.x_scale_stride];
+                 }
+             }
+             if (x_bias) {
+@@ -973,7 +979,7 @@
+                 float gate_value = tmp_gate[threadIdx.x];
+                 if constexpr (type == GGML_TYPE_NVFP4) {
+                     if (gate_scale) {
+-                        gate_value *= gate_scale[channel_x];
++                        gate_value *= gate_scale[channel_x * fusion.gate_scale_stride];
+                     }
+                 }
+                 if (gate_bias) {
 @@ -998,7 +1004,7 @@
                  }
              }
@@ -718,7 +759,30 @@
      GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
  
      GGML_TENSOR_BINARY_OP_LOCALS;
-@@ -1526,12 +1535,22 @@
+@@ -1498,17 +1507,20 @@
+             GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
+             fusion_local.gate_bias = fusion->gate_bias->data;
+         }
++        // FreeToken: a scale may be a strided vector, read with its element stride
+         if (fusion->x_scale) {
+             GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
+-            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
++            GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0);
+             GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
+             fusion_local.x_scale = fusion->x_scale->data;
++            fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float);
+         }
+         if (fusion->gate_scale) {
+             GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
+-            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
++            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0);
+             GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
+             fusion_local.gate_scale = fusion->gate_scale->data;
++            fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float);
+         }
+         fusion_local.glu_op = fusion->glu_op;
+         fusion_local.glu_limit = fusion->glu_limit;
+@@ -1526,12 +1538,22 @@
      }
  
      const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
@@ -744,7 +808,7 @@
      }
  
      const int64_t s01 = src0->nb[1] / ts_src0;
-@@ -1556,11 +1575,23 @@
+@@ -1556,11 +1578,23 @@
  
      const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0;
  
diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
index 69d9322736..170cbcac71 100644
--- a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
+++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu
@@ -115,6 +115,14 @@ int64_t check_weight(TensorView w, ggml_type type, int64_t k) {
     return row_bytes;
 }
 
+// one fp32 scale per expert, e.g. a column of a [E, n] scale bank; the MMVQ epilogue multiplies NVFP4 results by it
+ggml_tensor describe_scale(TensorView s, ggml_type type, int64_t e) {
+    RuntimeCheck(type == GGML_TYPE_NVFP4, "per-expert scales apply to NVFP4 weights, not ", ggml_type_name(type));
+    RuntimeCheck(is_dtype(s.dtype(), kDLFloat, 32) && s.ndim() == 1 && s.size(0) == e && s.stride(0) > 0, "a per-expert scale must be an fp32 [E] view");
+    const size_t nb0 = sizeof(float) * s.stride(0);
+    return describe(GGML_TYPE_F32, s.data_ptr(), {e, 1, 1, 1}, {nb0, nb0 * e, nb0 * e, nb0 * e});
+}
+
 int64_t q8_1_row_bytes(int64_t k) { return GGML_PAD(k, MATRIX_ROW_PADDING) / QK8_1 * (int64_t) sizeof(block_q8_1); }
 
 // MMQ tiles read up to MATRIX_ROW_PADDING values past the last row; the caller keeps this many bytes after it
@@ -234,7 +242,8 @@ void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Opt
 }
 
 // out[T, k, N]: row (t, j) multiplies x[t, j or 0] by expert ids[t, j] of W[E, N, K]; the matrix kernel runs a flat grid over the tiles that hold routes
-void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out) {
+// scale[ids[t, j]] multiplies row (t, j); only the vector kernel takes it
+void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, Optional<TensorView> scale) {
     RuntimeCheck(w.ndim() == 3, "mul_mat_id takes w[E, N, row_bytes]");
     RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
     const DeviceGuard guard(x.device());
@@ -264,8 +273,15 @@ void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids,
     ggml_tensor src1 = describe(a.type, a.data, {k, ky, t, 1}, {a.esz, a.esz * k, a.esz * k * ky, a.esz * k * ky * t});
     const size_t oe = ggml_type_size(otype);
     ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t});
+    RuntimeCheck(mmvq || !scale.has_value(), "mul_mat_id takes a per-expert scale only on the vector kernel, T=", t);
     if (mmvq) {
-        ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst);
+        ggml_cuda_mm_fusion_args_host fusion{};
+        ggml_tensor tscale;
+        if (scale.has_value()) {
+            tscale = describe_scale(scale.value(), type, e);
+            fusion.x_scale = &tscale;
+        }
+        ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst, scale.has_value() ? &fusion : nullptr);
     } else {
         ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0);
     }
@@ -291,7 +307,9 @@ ggml_glu_op fused_glu_op(const String & name) {
 }  // namespace
 
 // out[T, k, I] = glu(x @ gate^T, x @ up^T) for small T in one MMVQ pass; w[E, 2I, row_bytes] holds each expert's gate rows then its up rows; limit is swiglu_clamp's
-void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit) {
+// gate_scale / up_scale multiply the expert's gate / up results before the glu
+void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit,
+                       Optional<TensorView> gate_scale, Optional<TensorView> up_scale) {
     const ggml_glu_op glu_op = fused_glu_op(glu);
     RuntimeCheck(w.ndim() == 3 && w.size(1) % 2 == 0, "mul_mat_id_glu takes w[E, 2I, row_bytes]");
     RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
@@ -325,6 +343,15 @@ void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView i
     fusion.gate = &gate;
     fusion.glu_op = glu_op;
     fusion.glu_limit = (float) limit;
+    ggml_tensor tgate_scale, tup_scale;
+    if (gate_scale.has_value()) {
+        tgate_scale = describe_scale(gate_scale.value(), type, e);
+        fusion.gate_scale = &tgate_scale;
+    }
+    if (up_scale.has_value()) {
+        tup_scale = describe_scale(up_scale.value(), type, e);
+        fusion.x_scale = &tup_scale;
+    }
     ggml_cuda_mul_mat_vec_q(ctx, &up, &src1, &tids, &dst, &fusion);
 }
 
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
index 36912ce4b6..a767c688f5 100644
--- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh
@@ -1588,6 +1588,9 @@ struct ggml_cuda_mm_fusion_args_device {
     const void * gate_bias = nullptr;
     const void * x_scale = nullptr;
     const void * gate_scale = nullptr;
+    // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank
+    uint32_t x_scale_stride = 1;
+    uint32_t gate_scale_stride = 1;
     ggml_glu_op glu_op;
     float glu_limit = 0.0f;
     const void * shared_up = nullptr;
diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
index f3a3ccc683..8487a035c1 100644
--- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
+++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu
@@ -706,10 +706,10 @@ static __global__ void mul_mat_vec_q(
             }
             if constexpr (type == GGML_TYPE_NVFP4) {
                 if (use_scale) {
-                    x_scales = x_scale[ids ? channel_x : 0];
+                    x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride];
                 }
                 if (use_gate_scale) {
-                    gate_scales = gate_scale[ids ? channel_x : 0];
+                    gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride];
                 }
             }
         }
@@ -969,7 +969,7 @@ static __global__ void mul_mat_vec_q_moe(
 
             if constexpr (type == GGML_TYPE_NVFP4) {
                 if (x_scale) {
-                    result *= x_scale[channel_x];
+                    result *= x_scale[channel_x * fusion.x_scale_stride];
                 }
             }
             if (x_bias) {
@@ -979,7 +979,7 @@ static __global__ void mul_mat_vec_q_moe(
                 float gate_value = tmp_gate[threadIdx.x];
                 if constexpr (type == GGML_TYPE_NVFP4) {
                     if (gate_scale) {
-                        gate_value *= gate_scale[channel_x];
+                        gate_value *= gate_scale[channel_x * fusion.gate_scale_stride];
                     }
                 }
                 if (gate_bias) {
@@ -1507,17 +1507,20 @@ void ggml_cuda_mul_mat_vec_q(
             GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
             fusion_local.gate_bias = fusion->gate_bias->data;
         }
+        // FreeToken: a scale may be a strided vector, read with its element stride
         if (fusion->x_scale) {
             GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
-            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
+            GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0);
             GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
             fusion_local.x_scale = fusion->x_scale->data;
+            fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float);
         }
         if (fusion->gate_scale) {
             GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
-            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
+            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0);
             GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
             fusion_local.gate_scale = fusion->gate_scale->data;
+            fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float);
         }
         fusion_local.glu_op = fusion->glu_op;
         fusion_local.glu_limit = fusion->glu_limit;
diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py
index b672d58cf2..7a2765c5c0 100644
--- a/python/freetoken/kernel/ggml.py
+++ b/python/freetoken/kernel/ggml.py
@@ -95,24 +95,37 @@ def uses_mmvq(ggml_type: int, m: int) -> bool:
     return _module().uses_mmvq(ggml_type, m)
 
 
-def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
-    """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``."""
+def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None) -> torch.Tensor:
+    """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``.
+
+    ``scale`` is an fp32 ``[E]`` view (strided is fine) that multiplies each NVFP4 expert's rows; it needs ``T <= mmvq_max_tokens``.
+    """
     t, top_k = ids.shape
     out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device)
     if t > mmvq_max_tokens(ggml_type):
         _check_mmq_tail(w, ggml_type, x.shape[-1])
-    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out)
+    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale)
     return out
 
 
-def mul_mat_id_glu(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, glu: str, limit: float = 0.0) -> torch.Tensor:
+def mul_mat_id_glu(
+    w: torch.Tensor,
+    ggml_type: int,
+    x: torch.Tensor,
+    ids: torch.Tensor,
+    glu: str,
+    limit: float = 0.0,
+    gate_scale: torch.Tensor | None = None,
+    up_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
     """``glu(gate, up)`` per route in one pass for small ``T``: ``w[E, 2I, row_bytes]`` holds gate rows then up rows; returns ``[T, k, I]`` in ``x``'s dtype.
 
     ``glu`` is one of ``swiglu``, ``geglu`` (tanh GELU), ``swiglu_oai`` (alpha 1.702, limit 7) and ``swiglu_clamp`` (alpha 1, clamped at ``limit``).
+    ``gate_scale`` / ``up_scale`` are fp32 ``[E]`` views that multiply each NVFP4 expert's gate / up results before the glu.
     """
     t, top_k = ids.shape
     out = torch.empty(t, top_k, w.shape[1] // 2, dtype=x.dtype, device=x.device)
-    _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit)
+    _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit, gate_scale, up_scale)
     return out
 
 
diff --git a/python/freetoken/layers/quantization/configs/gguf.py b/python/freetoken/layers/quantization/configs/gguf.py
index bbc24f137a..77859cd1f7 100644
--- a/python/freetoken/layers/quantization/configs/gguf.py
+++ b/python/freetoken/layers/quantization/configs/gguf.py
@@ -4,7 +4,7 @@
 
 from typing import Any, Callable, ClassVar
 
-from ..ggml_types import is_quantized, type_name
+from ..ggml_types import has_expert_scales, has_weight_scale, is_quantized, type_name
 from ..names import NameMap
 from ..quant_backend import get_quant_backend
 from ..registry import LayerKind, method_class, register_dialect
@@ -29,19 +29,24 @@ def __init__(
         super().__init__(name_map, unquantized)
         self.tensor_types = tensor_types
         self.to_gguf = to_gguf
+        self.expert_scales = has_expert_scales(tensor_types)
+        self.weight_scales = any(n.endswith(".scale") and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types)
 
     @classmethod
     def claims(cls, q: dict[str, Any]) -> bool:
         return False
 
-    def _part_types(self, name: str) -> tuple[int, ...]:
-        types = []
+    def _part_tensors(self, name: str) -> tuple[str, ...]:
+        keys = []
         for tensor in self.to_gguf(name, self.tensor_types):
             key = f"{tensor}.weight" if f"{tensor}.weight" in self.tensor_types else tensor
             if key not in self.tensor_types:
                 raise KeyError(f"GGUF has no tensor {key!r} for {name!r}")
-            types.append(self.tensor_types[key])
-        return tuple(types)
+            keys.append(key)
+        return tuple(keys)
+
+    def _part_types(self, name: str) -> tuple[int, ...]:
+        return tuple(self.tensor_types[key] for key in self._part_tensors(name))
 
     def scheme_for_name(self, name: str) -> QuantScheme | None:
         types = self._part_types(name)
@@ -55,7 +60,8 @@ def scheme_for(self, prefix: str) -> QuantScheme | None:
         if prefix in self._schemes:
             return self._schemes[prefix]
         names = self.name_map.to_checkpoint(prefix)
-        parts = [() if self.unquantized(n) else self._part_types(n) for n in names]
+        tensors = [() if self.unquantized(n) else self._part_tensors(n) for n in names]
+        parts = [tuple(self.tensor_types[key] for key in ts) for ts in tensors]
         quantized = [any(is_quantized(t) for t in p) for p in parts]
         if not any(quantized):
             scheme = None
@@ -63,7 +69,8 @@ def scheme_for(self, prefix: str) -> QuantScheme | None:
             found = {n: [type_name(t) for t in p] for n, p in zip(names, parts)}
             raise NotImplementedError(f"fused module {prefix!r} needs one GGUF tensor per part, got {found}")
         else:
-            scheme = gguf_scheme(p[0] for p in parts)
+            scaled = any(has_weight_scale(self.tensor_types, ts[0]) for ts in tensors)
+            scheme = gguf_scheme((p[0] for p in parts), scaled=scaled)
         self._schemes[prefix] = scheme
         return scheme
 
@@ -74,7 +81,8 @@ def get_quant_method(self, layer: Any, prefix: str):
         parts = [self._part_types(n) for n in self.name_map.to_checkpoint(prefix)]
         if any(len(p) != 1 for p in parts):
             raise NotImplementedError(f"experts {prefix!r} are stored as {parts} GGUF tensors per projection")
-        scheme = gguf_scheme(p[0] for p in parts)
+        # the scale bank is per file, so every MoE layer keeps one bank set
+        scheme = gguf_scheme((p[0] for p in parts), scaled=self.expert_scales)
         cfg = self.layer_config(layer, LayerKind.MOE, scheme)
         return method_class(QuantKind.GGUF, LayerKind.MOE)(cfg, get_quant_backend().select(LayerKind.MOE, QuantKind.GGUF))
 
diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py
index f41b7b6a4b..56c72c2d6a 100644
--- a/python/freetoken/layers/quantization/embedding/gguf.py
+++ b/python/freetoken/layers/quantization/embedding/gguf.py
@@ -7,7 +7,7 @@
 
 import torch
 
-from ..ggml_types import Segment, pad_rows, row_bytes
+from ..ggml_types import NVFP4, Segment, pad_rows, row_bytes
 from ..linear.gguf import GgmlLinearKernel
 from ..registry import LayerKind, register_method
 from ..scheme import QuantKind, gguf_types
@@ -18,7 +18,10 @@ class GgmlEmbeddingKernel(EmbeddingKernel):
     name = "ggml"
 
     def unusable_reason(self, cfg: EmbeddingConfig) -> str | None:
-        return "GGUF tables are not sharded over TP" if cfg.tp_size > 1 else None
+        if cfg.tp_size > 1:
+            return "GGUF tables are not sharded over TP"
+        # the vendored get_rows has no NVFP4 case, and the table's weight_scale_2 would need applying too
+        return "NVFP4 GGUF tables are not supported" if NVFP4 in gguf_types(cfg.scheme) else None
 
     def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor:
         from freetoken.kernel import ggml
diff --git a/python/freetoken/layers/quantization/ggml_types.py b/python/freetoken/layers/quantization/ggml_types.py
index 8b257cd8c7..262c5668d3 100644
--- a/python/freetoken/layers/quantization/ggml_types.py
+++ b/python/freetoken/layers/quantization/ggml_types.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 
 import torch
@@ -85,4 +86,33 @@ def segment_param(index: int) -> str:
     return "weight" if index == 0 else f"weight_{index}"
 
 
-__all__ = ["MATRIX_ROW_PADDING", "Segment", "float_dtype", "is_quantized", "mmq_tail_bytes", "pad_rows", "row_bytes", "segment_param", "segments", "type_name"]
+NVFP4 = int(gguf.GGMLQuantizationType.NVFP4)
+
+# per-expert fp32 weight scales llama.cpp reads next to separate gate / up / down expert tensors (llama-model.cpp ffn_*_exps_s)
+_EXPERT_SCALE_RE = re.compile(r"^blk\.\d+\.ffn_(gate|up|down)_exps\.scale$")
+
+
+def scale_name(tensor: str) -> str:
+    """``<name>.scale`` of the GGUF tensor ``<name>`` or ``<name>.weight``."""
+    return f"{tensor.removesuffix('.weight')}.scale"
+
+
+def has_weight_scale(tensor_types: dict[str, int], tensor: str) -> bool:
+    """Whether llama.cpp's converter stored an NVFP4 tensor's weight_scale_2 as ``<name>.scale``; other types keep no weight scale."""
+    weight = tensor if tensor.endswith(".weight") else f"{tensor}.weight"
+    return tensor_types.get(weight, tensor_types.get(tensor)) == NVFP4 and scale_name(tensor) in tensor_types
+
+
+def has_expert_scales(tensor_types: dict[str, int]) -> bool:
+    """Whether any NVFP4 expert tensor of the file carries per-expert scales, which then take a scale bank in every MoE layer."""
+    merged = [n for n in tensor_types if n.endswith("ffn_gate_up_exps.scale") and tensor_types.get(n.replace(".scale", ".weight")) == NVFP4]
+    if merged:
+        # llama.cpp scales only separate gate / up expert tensors and would drop this one
+        raise NotImplementedError(f"per-expert scales of merged gate_up NVFP4 experts are not supported (e.g. {merged[0]})")
+    return any(_EXPERT_SCALE_RE.match(n) and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types)
+
+
+__all__ = [
+    "MATRIX_ROW_PADDING", "NVFP4", "Segment", "float_dtype", "has_expert_scales", "has_weight_scale", "is_quantized", "mmq_tail_bytes",
+    "pad_rows", "row_bytes", "scale_name", "segment_param", "segments", "type_name",
+]
diff --git a/python/freetoken/layers/quantization/linear/gguf.py b/python/freetoken/layers/quantization/linear/gguf.py
index ac08dbe287..41f61edf88 100644
--- a/python/freetoken/layers/quantization/linear/gguf.py
+++ b/python/freetoken/layers/quantization/linear/gguf.py
@@ -35,6 +35,9 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
                 ggml.mul_mat(w, s.ggml_type, x2, out=out, xq=xq)
             else:
                 torch.mm(x2, w.t(), out=out)
+        # llama.cpp multiplies an NVFP4 result by the tensor's weight_scale_2 (build_lora_mm)
+        if getattr(layer, "ggml_scale", None) is not None:
+            y.mul_(layer.ggml_scale)
         if layer.bias is not None:
             y = y + layer.bias
         return y.reshape(*x.shape[:-1], y.shape[-1])
@@ -54,3 +57,6 @@ def create_weights(self, layer: Any) -> None:
                 setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(k), row_bytes(s.ggml_type, k), dtype=torch.uint8))
             else:
                 setattr(layer, segment_param(i), torch.empty(s.rows, k))
+        if g.scheme.weight.scale is not None:
+            # one value per output row: each part's weight_scale_2, 1 for parts without one
+            layer.ggml_scale = torch.empty(sum(s.rows for s in segs), dtype=torch.float32)
diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py
index 416c2f4b58..01e754d644 100644
--- a/python/freetoken/layers/quantization/moe/gguf.py
+++ b/python/freetoken/layers/quantization/moe/gguf.py
@@ -6,7 +6,7 @@
 
 import torch
 
-from ..ggml_types import float_dtype, is_quantized, row_bytes, type_name
+from ..ggml_types import NVFP4, float_dtype, is_quantized, row_bytes, type_name
 from ..registry import LayerKind, register_method
 from ..scheme import QuantKind, gguf_types
 from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod, gated_epilogue_reason, is_resident, limit_or_inf
@@ -80,7 +80,11 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         gate_up, down = expert_types(cfg)
         i, h = cfg.local_intermediate, cfg.hidden
-        return {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)}
+        banks = {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)}
+        if cfg.scheme.weight.scale is not None:
+            # each expert's NVFP4 weight_scale_2 as (gate, up, down, 1): a 16-byte row, which the slot cache's copies need
+            banks["global"] = BankSpec((4,), torch.float32)
+        return banks
 
     def pack(self, pieces, cfg: MoEConfig, out):
         i = cfg.local_intermediate
@@ -113,17 +117,29 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor
                 return ggml.mul_mat_id(w, ggml_type, a, ids)
             return float_mul_mat_id(w, a, ids, topk_weights, config)
 
+        # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2, gate and up before the glu, down before the routing weight
+        scales = view.tensors.get("global")
+        gate_up_scales = scales if scales is not None and gate_up_type == NVFP4 else None
+        down_scales = scales if scales is not None and down_type == NVFP4 else None
         glu = fused_glu(layer)
         small = x.shape[0] <= min(ggml.mmvq_max_tokens(t) for t in (gate_up_type, down_type) if is_quantized(t))
         if glu is not None and small and is_quantized(gate_up_type):
-            inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu)
+            gu = (None, None) if gate_up_scales is None else (gate_up_scales[:, 0], gate_up_scales[:, 1])
+            inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu, *gu)
         else:
             h = mul_mat_id(gate_up, gate_up_type, x)
             t, k, two_i = h.shape
+            if gate_up_scales is not None:
+                h.view(t, k, 2, two_i // 2).mul_(gate_up_scales[ids, :2].unsqueeze(-1))
             inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device)
             gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer))
             del h  # free the gate_up output before the down GEMM allocates its own
-        y = mul_mat_id(view.tensors["down"], down_type, inter)
+        if down_scales is not None and x.shape[0] <= ggml.mmvq_max_tokens(down_type):
+            y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids, scale=down_scales[:, 2])
+        else:
+            y = mul_mat_id(view.tensors["down"], down_type, inter)
+            if down_scales is not None:
+                topk_weights = topk_weights * down_scales[ids, 2]
         return ggml.moe_weighted_sum(y, topk_weights)
 
     @staticmethod
@@ -147,4 +163,6 @@ def cpu_spec(self) -> tuple[str | None, dict[str, int]]:
         # BF16 experts keep the executor's bf16 kernels, which use AVX512-BF16 where the CPU has it; ggml-cpu is built for AVX2 only
         if float_dtype(gate_up) is torch.bfloat16 and float_dtype(down) is torch.bfloat16:
             return "bf16", {}
-        return "gguf", {"gate_up_type": gate_up, "down_type": down}
+        # an NVFP4 layer of a file with weight scales reads the layout's global bank too
+        scaled = self.cfg.scheme.weight.scale is not None and NVFP4 in (gate_up, down)
+        return "gguf_scaled" if scaled else "gguf", {"gate_up_type": gate_up, "down_type": down}
diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py
index ab43814e6a..a3aff0bc0f 100644
--- a/python/freetoken/layers/quantization/scheme.py
+++ b/python/freetoken/layers/quantization/scheme.py
@@ -108,9 +108,11 @@ def mxfp4_scheme() -> QuantScheme:
     return QuantScheme(QuantKind.MXFP4, WeightDesc("e2m1", (1, MX_GROUP), "e8m0"), {"weight", "weight_scale"})
 
 
-def gguf_scheme(types: Iterable[int]) -> QuantScheme:
-    """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order)."""
-    return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, None), {"weight"})
+def gguf_scheme(types: Iterable[int], scaled: bool = False) -> QuantScheme:
+    """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order).
+
+    ``scaled``: the NVFP4 tensors carry the fp32 ``<name>.scale`` (weight_scale_2) that llama.cpp's converter writes."""
+    return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, "f32" if scaled else None), {"weight"})
 
 
 def gguf_types(scheme: QuantScheme) -> tuple[int, ...]:
diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py
index 2b86cca127..3d4de920b8 100644
--- a/python/freetoken/models/gemma4/gguf.py
+++ b/python/freetoken/models/gemma4/gguf.py
@@ -24,7 +24,7 @@
     SWAAttentionGroupConfig,
 )
 from freetoken.models.gguf.reader import gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix, weight_scale
 
 if TYPE_CHECKING:
     from freetoken.models.gguf.config import GgufConfigShim
@@ -48,6 +48,7 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int:
 
 
 def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig:
+    check_scales(shim.model_path, known=("ffn_gate_inp.scale", "ffn_down_exps.scale"))
     m = shim.metadata
 
     def g(key: str):
@@ -246,6 +247,9 @@ def t(name: str):
         blk, base = f"blk.{layer}", f"model.layers.{layer}"
         for suffix, rel in scalar_map.items():
             tensor = float_tensor(index[f"{blk}.{suffix}"])
+            if suffix == "ffn_down_exps.scale" and weight_scale(index[f"{blk}.ffn_down_exps.weight"]) is not None:
+                # the converter folds per_expert_scale into the NVFP4 down weight_scale_2, which the expert kernel applies
+                tensor = torch.ones_like(tensor)
             yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor
         if f"{blk}.attn_qkv.weight" in index:
             # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py
index 0da56ceb4f..c21fff9fa2 100644
--- a/python/freetoken/models/gguf/experts.py
+++ b/python/freetoken/models/gguf/experts.py
@@ -1,7 +1,8 @@
 """Routed-expert banks straight from a GGUF's ``blk.N.ffn_{gate,up,down}_exps`` tensors, in their packed ggml rows: pinned host banks for the offload cache, or device banks for resident experts.
 
 Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts keep their dtype ([E, 2I, H] and [E, H, I]).
-A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order."""
+A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order.
+When NVFP4 experts carry per-expert ``.scale`` tensors, every layer also gets ``global`` [E, 4] fp32: the (gate, up, down, 1) weight_scale_2, 1 where a tensor has none."""
 
 from __future__ import annotations
 
@@ -11,10 +12,11 @@
 import torch
 
 from freetoken.layers.quantization import QuantKind
-from freetoken.layers.quantization.ggml_types import float_dtype, mmq_tail_bytes, pad_rows, type_name
+from freetoken.layers.quantization.ggml_types import float_dtype, has_expert_scales, mmq_tail_bytes, pad_rows, type_name
 from freetoken.utils import init_logger
 
 from .reader import GgufTensorInfo, gguf_tensor_index
+from .tensors import weight_scale
 
 logger = init_logger(__name__)
 
@@ -46,14 +48,30 @@ def _device_bank(shape: tuple[int, ...], k: int, dtype: torch.dtype, device: tor
     return flat[: experts * rows].view(shape)
 
 
+def _expert_scales(gate_up: list[GgufTensorInfo], down: GgufTensorInfo, experts: int) -> torch.Tensor:
+    """[E, 4] fp32 (gate, up, down, 1) weight_scale_2 of one layer's experts."""
+    gate, up = gate_up if len(gate_up) == 2 else (gate_up[0], gate_up[0])
+    out = torch.ones(experts, 4, dtype=torch.float32)
+    for col, info in enumerate((gate, up, down)):
+        s = weight_scale(info)
+        if s is None:
+            continue
+        if s.numel() != experts:
+            raise NotImplementedError(f"{info.name}: {s.numel()} scales for {experts} experts")
+        out[:, col] = s
+    return out
+
+
 def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int = 0, dummy: bool = False, device: torch.device | None = None):
     """``ExpertBanks`` with one bank per (layer, role): pinned host banks for the offload cache, filled from the file and pinned while later layers load, or resident banks on ``device``."""
     from freetoken.moe.expert_banks import ExpertBanks
     from freetoken.moe.host_banks import HostBank, PinPipeline
 
     index = gguf_tensor_index(model_path)
-    sources: dict[str, list[torch.Tensor]] = {"gate_up": [], "down": []}
-    tails: dict[str, list[int]] = {"gate_up": [], "down": []}
+    scaled = has_expert_scales({name: info.ggml_type for name, info in index.items()})
+    roles = ("gate_up", "down", "global") if scaled else ("gate_up", "down")
+    sources: dict[str, list[torch.Tensor]] = {role: [] for role in roles}
+    tails: dict[str, list[int]] = {role: [] for role in roles}
     tags = []
     with PinPipeline() if device is None else contextlib.nullcontext() as pins:
         for layer in range(num_layers):
@@ -70,17 +88,26 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int
             tags.append(("ggml", type_name(gate.ggml_type), type_name(down.ggml_type)))
             tails["gate_up"].append(mmq_tail_bytes(gate.ggml_type, gate.ne[0]))
             tails["down"].append(mmq_tail_bytes(down.ggml_type, down.ne[0]))
+            if scaled:
+                tails["global"].append(0)
             rows = [_expert_rows(info) for info in parts]
             d = _expert_rows(down)
             gate_up_shape = (rows[0].shape[0], sum(r.shape[1] for r in rows), rows[0].shape[2])
+            experts = rows[0].shape[0]
             if device is None:
                 host = {"gate_up": HostBank(gate_up_shape, rows[0].dtype), "down": HostBank(tuple(d.shape), d.dtype)}
+                if scaled:
+                    host["global"] = HostBank((experts, 4), torch.float32)
                 banks = {role: bank.tensor for role, bank in host.items()}
             else:
                 banks = {
                     "gate_up": _device_bank(gate_up_shape, gate.ne[0], rows[0].dtype, device),
                     "down": _device_bank(tuple(d.shape), down.ne[0], d.dtype, device),
                 }
+                if scaled:
+                    banks["global"] = torch.empty(experts, 4, dtype=torch.float32, device=device)
+            if scaled:
+                banks["global"].copy_(_expert_scales(parts, down, experts))
             if dummy:
                 banks["gate_up"].zero_()
                 banks["down"].zero_()
@@ -98,7 +125,7 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int
                 sources[role].append(bank)
     total = sum(t.numel() * t.element_size() for ts in sources.values() for t in ts)
     logger.info_rank0(f"GGUF expert banks: {num_layers} layers, {total / 2**30:.1f} GiB {'on ' + str(device) if device is not None else 'in host RAM'}")
-    return ExpertBanks("gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails)
+    return ExpertBanks("gguf_scaled" if scaled else "gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails)
 
 
 __all__ = ["load_gguf_expert_banks"]
diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py
index 2715f12909..ff7350941e 100644
--- a/python/freetoken/models/gguf/reader.py
+++ b/python/freetoken/models/gguf/reader.py
@@ -147,6 +147,8 @@ class GgufTensorInfo:
     nbytes: int
     ggml_type: int
     ne: tuple[int, ...]
+    # the model path the index was built from, which resolves tensors in any split
+    root: str = ""
 
     def data(self) -> np.ndarray:
         """Zero-copy uint8 view of the bytes (memory-mapped)."""
@@ -161,7 +163,7 @@ def gguf_tensor_index(model_path: str) -> dict[str, GgufTensorInfo]:
     for path in gguf_split_paths(model_path):
         for t in _reader(path).tensors:
             index[t.name] = GgufTensorInfo(
-                t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape)
+                t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape), model_path
             )
     return index
 
diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py
index 4cba1d082e..83452d957a 100644
--- a/python/freetoken/models/gguf/tensors.py
+++ b/python/freetoken/models/gguf/tensors.py
@@ -7,9 +7,9 @@
 import numpy as np
 import torch
 
-from freetoken.layers.quantization.ggml_types import is_quantized, pad_rows, segment_param, segments
+from freetoken.layers.quantization.ggml_types import NVFP4, is_quantized, pad_rows, scale_name, segment_param, segments
 
-from .reader import GgufTensorInfo
+from .reader import GgufTensorInfo, gguf_tensor_index
 
 _FLOAT_DTYPES = {0: torch.float32, 1: torch.float16, 30: torch.bfloat16}
 
@@ -29,8 +29,49 @@ def matrix(info: GgufTensorInfo) -> torch.Tensor:
     return packed_rows(info) if is_quantized(info.ggml_type) else float_tensor(info)
 
 
+def check_scales(model_path: str, *, nvfp4: bool = True, known: tuple[str, ...] = ()) -> None:
+    """Refuse ``.scale`` / ``.input_scale`` tensors that nothing reads.
+
+    With ``nvfp4``, an NVFP4 tensor's ``.scale`` (its weight_scale_2) is read and its ``.input_scale`` goes unused, as in llama.cpp.
+    ``known`` lists tensor-name suffixes the family's loader reads itself."""
+    types = {name: info.ggml_type for name, info in gguf_tensor_index(model_path).items()}
+    unread = []
+    for name in types:
+        weight = name.removesuffix(".input_scale") if name.endswith(".input_scale") else name.removesuffix(".scale")
+        if weight == name or name.endswith(known):
+            continue
+        if not (nvfp4 and types.get(f"{weight}.weight", types.get(weight)) == NVFP4):
+            unread.append(name)
+    if unread:
+        what = "scale tensors outside NVFP4 weights" if nvfp4 else "tensors with per-tensor scales"
+        raise NotImplementedError(f"{model_path}: GGUF {what} are not supported (e.g. {sorted(unread)[0]})")
+
+
+def weight_scale(info: GgufTensorInfo) -> torch.Tensor | None:
+    """The fp32 weight_scale_2 llama.cpp's converter writes next to an NVFP4 tensor as ``<name>.scale`` (one value, or one per expert); None if there is none."""
+    if info.ggml_type != NVFP4:
+        return None
+    scale = gguf_tensor_index(info.root or info.path).get(scale_name(info.name))
+    return None if scale is None else float_tensor(scale).float().reshape(-1)
+
+
+def _row_scales(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], pad: int) -> torch.Tensor | None:
+    scales = [weight_scale(i) for i in infos]
+    if all(s is None for s in scales):
+        return None
+    rows = []
+    for info, part, s in zip(infos, parts, scales):
+        # llama.cpp reads a dense tensor's scale as one value
+        if s is not None and s.numel() != 1:
+            raise NotImplementedError(f"{prefix}: {info.name} has {s.numel()} scales, expected one")
+        rows.append(torch.ones(part.shape[0]) if s is None else s.expand(part.shape[0]))
+    return torch.cat([*rows, torch.ones(pad)])
+
+
 def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k: int, pad_to: int = 0) -> Iterator[tuple[str, torch.Tensor]]:
-    """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed)."""
+    """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed).
+
+    NVFP4 parts with a ``.scale`` add ``ggml_scale``: each part's scale on its output rows, fp32."""
     types = [i.ggml_type for i in infos]
     pad = (-sum(p.shape[0] for p in parts)) % pad_to if pad_to else 0
     if not any(is_quantized(t) for t in types):
@@ -51,6 +92,9 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k
             dtype = torch.bfloat16
         tail = [torch.zeros(zeros, rows[0].shape[1], dtype=dtype)] if zeros else []
         yield f"{prefix}.{segment_param(i)}", torch.cat([*rows, *tail], dim=0)
+    scales = _row_scales(prefix, infos, parts, pad)
+    if scales is not None:
+        yield f"{prefix}.ggml_scale", scales
 
 
-__all__ = ["float_tensor", "fused", "matrix", "packed_rows"]
+__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "weight_scale"]
diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py
index 09d02627ae..833a5d66b5 100644
--- a/python/freetoken/models/llama/gguf.py
+++ b/python/freetoken/models/llama/gguf.py
@@ -13,7 +13,7 @@
 import torch
 
 from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix
 
 from .config import parse_config
 
@@ -53,10 +53,8 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 def _check_supported(shim: GgufConfigShim) -> None:
     if any(k.startswith("prism.hadamard.") for k in shim.metadata):
         raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    check_scales(shim.model_path)
     names = sorted(gguf_tensor_index(shim.model_path))
-    scaled = [n for n in names if n.endswith((".scale", ".input_scale"))]
-    if scaled:
-        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
     # the llama arch also carries Mixtral experts and optional biases, which the llama model does not have
     extra = [n for n in names if n.endswith(".bias") or ".ffn_gate_inp." in n]
     if extra:
diff --git a/python/freetoken/models/qwen3/gguf.py b/python/freetoken/models/qwen3/gguf.py
index 06245ce4c2..8e0388fe3c 100644
--- a/python/freetoken/models/qwen3/gguf.py
+++ b/python/freetoken/models/qwen3/gguf.py
@@ -12,7 +12,7 @@
 import torch
 
 from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix
 
 from .config import parse_config
 
@@ -52,9 +52,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 def _check_supported(shim: GgufConfigShim) -> None:
     if any(k.startswith("prism.hadamard.") for k in shim.metadata):
         raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
-    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
-    if scaled:
-        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    check_scales(shim.model_path)
 
 
 def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py
index 85af8ab85f..7f72061e00 100644
--- a/python/freetoken/models/qwen3_5_moe/gguf.py
+++ b/python/freetoken/models/qwen3_5_moe/gguf.py
@@ -14,7 +14,7 @@
 
 from freetoken.layers.quantization.ggml_types import is_quantized
 from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix
 
 from .config import parse_config
 
@@ -67,9 +67,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 def _check_supported(shim: GgufConfigShim) -> None:
     if any(k.startswith("prism.hadamard.") for k in shim.metadata):
         raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
-    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
-    if scaled:
-        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    check_scales(shim.model_path)
 
 
 def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py
index 7bf016292e..3a8ef98a7a 100644
--- a/python/freetoken/models/qwen3_moe/gguf.py
+++ b/python/freetoken/models/qwen3_moe/gguf.py
@@ -13,7 +13,7 @@
 import torch
 
 from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix
 
 from .config import parse_config
 
@@ -56,9 +56,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 def _check_supported(shim: GgufConfigShim) -> None:
     if any(k.startswith("prism.hadamard.") for k in shim.metadata):
         raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
-    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
-    if scaled:
-        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+    check_scales(shim.model_path)
 
 
 def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
diff --git a/python/freetoken/models/qwen4_exp/gguf.py b/python/freetoken/models/qwen4_exp/gguf.py
index 7909324b2c..0c1bf98059 100644
--- a/python/freetoken/models/qwen4_exp/gguf.py
+++ b/python/freetoken/models/qwen4_exp/gguf.py
@@ -14,7 +14,7 @@
 
 from freetoken.layers.quantization.ggml_types import is_quantized
 from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
-from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix
 
 from .config import parse_config
 
@@ -83,6 +83,8 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 
 
 def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    # NVFP4 weight scales are unvalidated on this arch, so refuse them instead of risking wrong output
+    check_scales(shim.model_path, nvfp4=False)
     md = shim.metadata
 
     def kv(key: str):
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index 4c0aee7ae5..bb65146b45 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -149,11 +149,11 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]:
 
 def _format_label(name: str, opts: tuple) -> str:
     """One format instance as the startup log names it, e.g. ``gguf(gate_up=IQ2_XXS, down=IQ4_NL)``."""
-    if name == "gguf":
+    if name in ("gguf", "gguf_scaled"):
         from freetoken.layers.quantization.ggml_types import type_name
 
         gate_up, down = dict(opts)["gate_up_type"], dict(opts)["down_type"]
-        return f"gguf(gate_up={type_name(gate_up)}, down={type_name(down)})"
+        return f"{name}(gate_up={type_name(gate_up)}, down={type_name(down)})"
     return name + (f"({', '.join(f'{k}={v}' for k, v in opts)})" if opts else "")
 
 
@@ -192,7 +192,7 @@ def __init__(
                     f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are "
                     f"{f!r}; use --moe-strategy offload (GPU-side dequant) instead."
                 )
-            if f == "gguf":
+            if f in ("gguf", "gguf_scaled"):
                 from freetoken.kernel import ggml_cpu
 
                 if (variant := ggml_cpu.select()) is None:
@@ -423,8 +423,8 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
         (see ``_make_table``), not a single bank's -- the C++ side resolves
         ``tbl[layer_id]`` per task.
         """
-        if fmt == "gguf":
-            return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"]))
+        if fmt in ("gguf", "gguf_scaled"):
+            return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"]), scaled=fmt == "gguf_scaled")
         if fmt == "bf16":
             gate_up = banks["gate_up"]
             down = banks["down"]
@@ -465,8 +465,10 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
         banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng)
         return banks, {}, (H, I)
 
-    def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int) -> tuple[dict, dict, tuple[int, int]]:
-        """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype."""
+    def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int, scaled: bool = False) -> tuple[dict, dict, tuple[int, int]]:
+        """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype.
+
+        ``scaled`` adds global [E, 4] fp32, each expert's (gate, up, down, 1) NVFP4 weight_scale_2."""
         from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes, type_name
 
         gate_up, down = banks["gate_up"], banks["down"]
@@ -475,8 +477,13 @@ def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_typ
             dtype = float_dtype(ggml_type) if float_dtype(ggml_type) is not None else torch.uint8
             if t.dtype != dtype or t.shape[-1] * t.element_size() != row_bytes(ggml_type, k):
                 raise ValueError(f"MoE layer {rep} {role} bank {tuple(t.shape)} {t.dtype} is not {type_name(ggml_type)} rows of {k} values")
-        banks = self._tables(gate_up=gate_up, down=down)
-        return banks, {"gate_up_type": gate_up_type, "down_type": down_type}, (H, I)
+        types = {"gate_up_type": gate_up_type, "down_type": down_type}
+        if not scaled:
+            return self._tables(gate_up=gate_up, down=down), types, (H, I)
+        scales = banks["global"]
+        if scales[rep].dtype != torch.float32 or tuple(scales[rep].shape) != (gate_up[rep].shape[0], 4) or not scales[rep].is_contiguous():
+            raise ValueError(f"MoE layer {rep} global bank {tuple(scales[rep].shape)} {scales[rep].dtype} is not contiguous [experts, 4] fp32")
+        return self._tables(gate_up=gate_up, down=down, **{"global": scales}), types, (H, I)
 
     def _resolve_mxfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]:
         """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 23ba325106..5291e621f8 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -48,6 +48,8 @@
     "fp8_block": ("gate_up", "gate_up_scale", "down", "down_scale"),
     # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts keep their dtype
     "gguf": ("gate_up", "down"),
+    # GGUF experts whose NVFP4 tensors carry per-expert scales: global [L*E, 4] fp32 (gate, up, down, 1) weight_scale_2
+    "gguf_scaled": ("gate_up", "down", "global"),
     # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes +
     # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct
     # globals, and folding them into the e4m3 block scales would underflow)

From ac6a15de75ea0c80e0d5a206bbaea8218898e54f Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 17 Oct 2026 17:50:59 -0700
Subject: [PATCH 121/249] perf(moe): speed up the nvfp4 moe prefill gemm on rtx
 gpus (#733)

* perf(moe): speed up the nvfp4 moe prefill gemm on rtx gpus

* perf(moe): keep the byte e2m1 decode at every tile size on h100

---------

Co-authored-by: andy-yang-1 <andy_yang@berkeley.edu>
---
 python/freetoken/kernel/gluon/__init__.py     |  35 ++
 .../freetoken/kernel/gluon/nvfp4/moe/sm120.py | 354 ++++++++++++++++++
 python/freetoken/kernel/gluon/nvfp4/utils.py  |  43 +++
 python/freetoken/kernel/nvfp4_moe.py          | 336 +++++++++++++++++
 .../kernel/triton/nvfp4_fused_moe.py          | 134 ++++++-
 python/freetoken/moe/fused_nvfp4.py           | 250 +------------
 tests/moe/test_nvfp4_backends.py              |  47 +++
 7 files changed, 941 insertions(+), 258 deletions(-)
 create mode 100644 python/freetoken/kernel/gluon/__init__.py
 create mode 100644 python/freetoken/kernel/gluon/nvfp4/moe/sm120.py
 create mode 100644 python/freetoken/kernel/gluon/nvfp4/utils.py
 create mode 100644 python/freetoken/kernel/nvfp4_moe.py

diff --git a/python/freetoken/kernel/gluon/__init__.py b/python/freetoken/kernel/gluon/__init__.py
new file mode 100644
index 0000000000..eaf4365ed6
--- /dev/null
+++ b/python/freetoken/kernel/gluon/__init__.py
@@ -0,0 +1,35 @@
+"""Gluon versions of Triton kernels for specific GPUs, written against the Gluon API of Triton :data:`TRITON_VERSION`.
+
+Callers route to a Gluon kernel only when :func:`available` is true and the kernel supports the device and inputs; otherwise the Triton kernel runs.
+"""
+
+from __future__ import annotations
+
+import functools
+
+from freetoken.kernel.backend import is_rocm
+from freetoken.kernel.utils import _env_enabled
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+DISABLE_GLUON_ENV = "FREETOKEN_DISABLE_GLUON"
+# Gluon is experimental and its API changes between Triton releases
+TRITON_VERSION = (3, 8)
+
+
+@functools.cache
+def available() -> bool:
+    """True on CUDA with Triton :data:`TRITON_VERSION`, unless FREETOKEN_DISABLE_GLUON is set; call it only for a GPU that has a Gluon kernel, as another Triton version logs a warning."""
+    if _env_enabled(DISABLE_GLUON_ENV) or is_rocm():
+        return False
+    import triton
+    from packaging.version import Version
+
+    if Version(triton.__version__).release[:2] == TRITON_VERSION:
+        return True
+    logger.warning(
+        "Gluon kernels need Triton %s, found %s; this GPU runs the slower Triton kernels instead (%s=1 silences this)",
+        ".".join(map(str, TRITON_VERSION)), triton.__version__, DISABLE_GLUON_ENV,
+    )
+    return False
diff --git a/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py
new file mode 100644
index 0000000000..0455b237a3
--- /dev/null
+++ b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py
@@ -0,0 +1,354 @@
+"""NVFP4 W4A16 MoE prefill GEMM for sm_120: a loader warp (TMA weights, cp.async scales and tokens), decode warps and mma.sync warps.
+
+Same contract as the Triton prefill GEMM; :mod:`freetoken.kernel.nvfp4_moe` picks it for BLOCK_SIZE_M 64 and 128 and sets its constexprs.
+"""
+
+import torch
+import triton
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language.nvidia.ampere import async_copy, mma_v2
+from triton.experimental.gluon.language.nvidia.hopper import fence_async_shared, mbarrier, tma
+from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
+
+from ..utils import e2m1x8_bf16x2, scale_bf16x2
+
+
+@gluon.jit
+def _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M: gl.constexpr):
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    first_pid_m = (t // num_pid_in_group) * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((t % num_pid_in_group) % group_size_m)
+    pid_n = (t % num_pid_in_group) // group_size_m
+    return pid_m, pid_n
+
+
+@gluon.jit
+def _load_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid,
+                    stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr,
+                    PERSISTENT: gl.constexpr):
+    BN: gl.constexpr = bufs[1].type.shape[1]
+    SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2)
+    if PERSISTENT:
+        # it / jt count chunks and scale blocks across tiles, so the rings keep their phase from one tile to the next
+        it = 0
+        jt = 0
+        for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+            pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M)
+            _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid,
+                       stride_am, ks4, n_rows, top_k, K_TAIL)
+            it += nk
+            jt += gl.cdiv(nk, SPC)
+    else:
+        pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M)
+        _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, 0, 0, num_valid,
+                   stride_am, ks4, n_rows, top_k, K_TAIL)
+
+
+@gluon.jit
+def _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid,
+               stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr):
+    x_bufs, w_bufs, s_bufs, d_bufs = bufs
+    x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars
+    KS: gl.constexpr = w_bufs.type.shape[2] // 8
+    XS: gl.constexpr = x_bufs.type.shape[0]
+    BM: gl.constexpr = x_bufs.type.shape[1] // KS
+    BK: gl.constexpr = KS * 16
+    WS: gl.constexpr = w_bufs.type.shape[0]
+    KB: gl.constexpr = w_bufs.type.shape[2]
+    SS: gl.constexpr = s_bufs.type.shape[0]
+    BN: gl.constexpr = s_bufs.type.shape[1]
+    SPC: gl.constexpr = 256 // BK
+
+    # 16 bytes per cp.async; adjacent lanes read one token's contiguous K
+    NWL: gl.constexpr = gl.num_warps()
+    xl: gl.constexpr = gl.BlockedLayout([1, 1, 8], [KS, 16 // KS, 2], [1, NWL, 1], [2, 0, 1])
+    offs_s = gl.arange(0, KS, layout=gl.SliceLayout(1, gl.SliceLayout(2, xl)))
+    offs_c = gl.arange(0, 16, layout=gl.SliceLayout(0, gl.SliceLayout(1, xl)))
+    cols = gl.expand_dims(gl.expand_dims(offs_s * 16, 1), 2) + gl.expand_dims(gl.expand_dims(offs_c, 0), 1)
+
+    # scale rows are K / 16 bytes: 16-byte copies when K % 256 == 0, else 4-byte copies zeroed past the row end
+    SV: gl.constexpr = 4 if K_TAIL else 16
+    sll: gl.constexpr = gl.BlockedLayout([1, SV], [32 * SV // 16, 16 // SV], [NWL, 1], [1, 0])
+    s_cols = gl.arange(0, 16, layout=gl.SliceLayout(0, sll))
+    # both forms equal K // 16; the second tells the compiler the row is a multiple of 16 bytes
+    row_len = ks4 * 4 if K_TAIL else (ks4 // 4) * 16
+
+    w_row = gl.load(eid_ptr + pid_m) * N + pid_n * BN
+    offs_m = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(2, xl)))
+    tok = gl.load(ids_ptr + offs_m)
+    tok_mask = gl.expand_dims(gl.expand_dims(tok < num_valid, 0), 2)
+    rows = gl.expand_dims(gl.expand_dims((tok // top_k).to(gl.int64) * stride_am, 0), 2)
+    x_ptrs = a_ptr + rows + cols
+    s_rows = w_row + gl.arange(0, BN, layout=gl.SliceLayout(1, sll))
+    s_ptrs = scale_ptr + s_rows.to(gl.int64)[:, None] * row_len + s_cols[None, :]
+    s_row_ok = (s_rows < n_rows)[:, None]
+    for kk in range(nk):
+        if kk % SPC == 0:
+            jl = kk // SPC
+            j = jt + jl
+            ss = j % SS
+            mbarrier.wait(s_empty.index(ss), (j // SS & 1) ^ 1)
+            async_copy.async_load(s_bufs.index(ss), s_ptrs + jl * 16, mask=s_row_ok & (jl * 16 + s_cols < row_len)[None, :])
+            async_copy.mbarrier_arrive(s_ready.index(ss))
+            mbarrier.arrive(s_ready.index(ss))
+        i = it + kk
+        ws = i % WS
+        mbarrier.wait(w_empty.index(ws), (i // WS & 1) ^ 1)
+        mbarrier.expect(w_ready.index(ws), w_desc.block_type.nbytes)
+        tma.async_load(w_desc, [w_row, kk * KB], w_ready.index(ws), w_bufs.index(ws))
+        xs = i % XS
+        mbarrier.wait(x_empty.index(xs), (i // XS & 1) ^ 1)
+        async_copy.async_load(x_bufs.index(xs).reshape([KS, BM, 16]), x_ptrs + kk * BK, mask=tok_mask)
+        async_copy.mbarrier_arrive(x_ready.index(xs))
+        mbarrier.arrive(x_ready.index(xs))
+
+
+@gluon.jit
+def _decode_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid,
+                      stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr,
+                      PERSISTENT: gl.constexpr):
+    SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2)
+    if PERSISTENT:
+        it = 0
+        jt = 0
+        for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+            _decode_tile(bufs, bars, nk, it, jt, K_TAIL)
+            it += nk
+            jt += gl.cdiv(nk, SPC)
+    else:
+        _decode_tile(bufs, bars, nk, 0, 0, K_TAIL)
+
+
+@gluon.jit
+def _decode_tile(bufs, bars, nk, it, jt, K_TAIL: gl.constexpr):
+    x_bufs, w_bufs, s_bufs, d_bufs = bufs
+    x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars
+    WS: gl.constexpr = w_bufs.type.shape[0]
+    BN: gl.constexpr = w_bufs.type.shape[1]
+    KB: gl.constexpr = w_bufs.type.shape[2]
+    BK: gl.constexpr = KB * 2
+    SS: gl.constexpr = s_bufs.type.shape[0]
+    SB: gl.constexpr = BK // 16
+    DS: gl.constexpr = d_bufs.type.shape[0]
+    SPC: gl.constexpr = 256 // BK
+    NW: gl.constexpr = gl.num_warps()
+
+    # one thread per row, so the k16 slice stores of a warp hit 32 consecutive rows
+    ql: gl.constexpr = gl.BlockedLayout([1, 2 * SB], [32, 1], [NW, 1], [1, 0])
+    q3l: gl.constexpr = gl.BlockedLayout([1, SB, 2], [32, 1, 1], [NW, 1, 1], [2, 1, 0])
+    qsl: gl.constexpr = gl.SliceLayout(2, q3l)
+    w32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=w_bufs.type.layout.swizzle_byte_width,
+                                                     element_bitwidth=32, rank=2)
+    d32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=32, element_bitwidth=32, rank=2)
+    w32 = w_bufs._reinterpret(gl.int32, [WS, BN, KB // 4], w32_layout)
+    d32 = d_bufs._reinterpret(gl.int32, [DS, SB * BN, 8], d32_layout)
+
+    # with K % 256 != 0 the last scale block covers fewer than SPC chunks
+    for jb in range(gl.cdiv(nk, SPC)):
+        j = jt + jb
+        ss = j % SS
+        mbarrier.wait(s_ready.index(ss), j // SS & 1)
+        for c in gl.static_range(SPC):
+            kk = jb * SPC + c
+            if not K_TAIL or kk < nk:
+                i = it + kk
+                ws = i % WS
+                mbarrier.wait(w_ready.index(ws), i // WS & 1)
+                q = w32.index(ws).load(ql)
+                s = s_bufs.index(ss).slice(c * SB, SB, dim=1).load(qsl)
+                fence_async_shared()
+                mbarrier.arrive(w_empty.index(ws))
+                q = gl.convert_layout(gl.reshape(q, [BN, SB, 2]), q3l, assert_trivial=True).to(gl.uint32, bitcast=True)
+                w0, w1, w2, w3 = e2m1x8_bf16x2(q, gl.expand_dims(scale_bf16x2(s), 2))
+                # join appends the new axis last, so this flattens to w0, w1, w2, w3; then k16 slice major for the MMA warps
+                w = gl.reshape(gl.join(gl.join(w0, w2), gl.join(w1, w3)), [BN, SB, 8])
+                w = gl.permute(w, (1, 0, 2)).to(gl.int32, bitcast=True)
+
+                ds = i % DS
+                # the MMA warps release x and d slots together on x_empty
+                mbarrier.wait(x_empty.index(ds), (i // DS & 1) ^ 1)
+                d32.index(ds).reshape([SB, BN, 8]).store(w)
+                mbarrier.arrive(d_ready.index(ds))
+        fence_async_shared()
+        mbarrier.arrive(s_empty.index(ss))
+
+
+@gluon.jit
+def _mma_chunks(bufs, bars, acc, it, nk, al: gl.constexpr, bl: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr):
+    x_bufs, w_bufs, s_bufs, d_bufs = bufs
+    x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars
+    KS: gl.constexpr = w_bufs.type.shape[2] // 8
+    XS: gl.constexpr = x_bufs.type.shape[0]
+    for kk in range(nk):
+        i = it + kk
+        xs = i % XS
+        mbarrier.wait(d_ready.index(xs), i // XS & 1)
+        mbarrier.wait(x_ready.index(xs), i // XS & 1)
+        for k16 in gl.static_range(KS):
+            a = d_bufs.index(xs).slice(k16 * BN, BN).load(al)
+            b = x_bufs.index(xs).slice(k16 * BM, BM).permute((1, 0)).load(bl)
+            acc = mma_v2(a, b, acc)
+        mbarrier.arrive(x_empty.index(xs))
+    return acc
+
+
+@gluon.jit
+def _mma_partition(bufs, bars, ids_ptr, eid_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk, num_tiles, num_pid_m,
+                   num_valid, stride_ge, stride_cm, stride_tw, mma: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr,
+                   MUL_ROUTED_WEIGHT: gl.constexpr, GROUP_SIZE_M: gl.constexpr, PERSISTENT: gl.constexpr):
+    al: gl.constexpr = gl.DotOperandLayout(0, mma, 2)
+    bl: gl.constexpr = gl.DotOperandLayout(1, mma, 2)
+    if PERSISTENT:
+        it = 0
+        for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+            acc = _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), it, nk, al, bl, BN, BM)
+            it += nk
+            # the epilogue runs here while the loader and decode warps already fill the rings for the next tile
+            # stored straight from the MMA layout, as a convert_layout needs scratch shared memory on top of the rings
+            pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M)
+            slot = gl.load(eid_ptr + pid_m)
+            offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, mma))
+            offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, mma))
+            _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid,
+                      stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT)
+    else:
+        return _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), 0, nk, al, bl, BN, BM)
+
+
+@gluon.jit
+def _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid,
+              stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT: gl.constexpr):
+    route = gl.load(ids_ptr + offs_t)
+    route_mask = route < num_valid
+    n_mask = offs_n < N
+    g = gl.load(global_ptr + slot.to(gl.int64) * stride_ge + offs_n, mask=n_mask, other=0.0).to(gl.float32)
+    # 256 undoes the 2^-8 that the decode leaves in every weight
+    acc = acc * (g * 256.0)[:, None]
+    if MUL_ROUTED_WEIGHT:
+        wt = gl.load(topk_weights_ptr + route * stride_tw, mask=route_mask, other=0.0)
+        acc = acc * wt[None, :]
+    c_ptrs = c_ptr + route.to(gl.int64)[None, :] * stride_cm + offs_n[:, None]
+    gl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=route_mask[None, :] & n_mask[:, None])
+
+
+SCALE_SMEM_LAYOUT = gl.constexpr(gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=2))
+
+
+@gluon.jit
+def _prefill_nvfp4_moe_kernel(
+    a_ptr, w_desc, scale_ptr, global_ptr, c_ptr, topk_weights_ptr,
+    sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr,
+    N, K, EM, num_valid_tokens, ks4, n_rows,
+    stride_am, stride_ge, stride_cm, stride_tw,
+    BLOCK_SIZE_M: gl.constexpr,
+    BLOCK_SIZE_N: gl.constexpr,
+    BLOCK_SIZE_K: gl.constexpr,
+    GROUP_SIZE_M: gl.constexpr,
+    STAGES: gl.constexpr,
+    W_STAGES: gl.constexpr,
+    PERSISTENT: gl.constexpr,
+    K_TAIL: gl.constexpr,
+    DECODE_WARPS: gl.constexpr,
+    MMA_WARPS_N: gl.constexpr,
+    LOAD_REGS: gl.constexpr,
+    DECODE_REGS: gl.constexpr,
+    MUL_ROUTED_WEIGHT: gl.constexpr,
+    top_k: gl.constexpr,
+):
+    BM: gl.constexpr = BLOCK_SIZE_M
+    BN: gl.constexpr = BLOCK_SIZE_N
+    BK: gl.constexpr = BLOCK_SIZE_K
+    live_m = gl.cdiv(gl.load(num_tokens_post_padded_ptr), BM)
+    if PERSISTENT:
+        num_pid_m = live_m
+    else:
+        # the grid covers the most m-blocks moe_align can produce; each CTA takes the tile of its index
+        num_pid_m = gl.cdiv(EM, BM)
+        pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M)
+        if pid_m >= live_m:
+            return
+    num_tiles = num_pid_m * gl.cdiv(N, BN)
+    nk = K // BK
+
+    # x and the decoded weights are stored as k16 slices, [stage, (k // 16) * rows, 16], and share STAGES slots
+    x_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BM, 16], gl.bfloat16)
+    d_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BN, 16], gl.bfloat16)
+    x_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BM, 16], x_layout)
+    w_bufs = gl.allocate_shared_memory(gl.uint8, [W_STAGES, BN, BK // 2], w_desc.layout)
+    s_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT)
+    d_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BN, 16], d_layout)
+    bufs = (x_bufs, w_bufs, s_bufs, d_bufs)
+
+    x_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    x_empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    w_ready = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout())
+    w_empty = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout())
+    s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    d_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    for i in gl.static_range(STAGES):
+        mbarrier.init(x_ready.index(i), count=1)
+        mbarrier.init(x_empty.index(i), count=1)
+        mbarrier.init(d_ready.index(i), count=1)
+    for i in gl.static_range(W_STAGES):
+        mbarrier.init(w_ready.index(i), count=1)
+        mbarrier.init(w_empty.index(i), count=1)
+    for i in gl.static_range(2):
+        mbarrier.init(s_ready.index(i), count=1)
+        mbarrier.init(s_empty.index(i), count=1)
+    bars = (x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready)
+
+    NW: gl.constexpr = gl.num_warps()
+    mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N],
+                                                  instr_shape=[16, 8])
+    # tuples go inline: a tuple stored in a variable turns its constexprs into runtime values
+    acc = gl.warp_specialize([
+        (_mma_partition, (bufs, bars, sorted_token_ids_ptr, expert_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk,
+                          num_tiles, num_pid_m, num_valid_tokens, stride_ge, stride_cm, stride_tw, mma, BN, BM,
+                          MUL_ROUTED_WEIGHT, GROUP_SIZE_M, PERSISTENT)),
+        (_load_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles,
+                           num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)),
+        (_decode_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles,
+                             num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)),
+    ], [1, DECODE_WARPS], [LOAD_REGS, DECODE_REGS])
+
+    if not PERSISTENT:
+        # a coalesced store layout; the rings are idle by now, so the conversion's scratch fits
+        out_l: gl.constexpr = gl.BlockedLayout([8, 1], [BN // 8, 32 * 8 // BN], [1, NW], [0, 1])
+        pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M)
+        slot = gl.load(expert_ids_ptr + pid_m)
+        offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, out_l))
+        offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, out_l))
+        _epilogue(gl.convert_layout(acc, out_l), sorted_token_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n,
+                  offs_t, N, num_valid_tokens, stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT)
+
+
+def supports(a: torch.Tensor, packed: torch.Tensor, scale: torch.Tensor, c: torch.Tensor) -> bool:
+    """Whether :func:`prefill_gemm` can run this GEMM on an sm_120 GPU."""
+    K = packed.shape[2] * 2
+    # the TMA weight loads and 16-byte cp.async need contiguous, 16-byte aligned rows
+    return (a.dtype == torch.bfloat16 and K % 64 == 0
+            and packed.is_contiguous() and scale.is_contiguous() and c.is_contiguous()
+            and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0
+            and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0)
+
+
+def prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+                 num_valid_tokens, top_k, mul_routed_weight, cfg):
+    """Launch the kernel with the tile / pipeline constexprs in ``cfg``; K must be a multiple of cfg["BLOCK_SIZE_K"]."""
+    S, N, KB = packed.shape
+    K = KB * 2
+    BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"]
+    w_layout = gl.NVMMASharedLayout.get_default_for([BN, BK // 2], gl.uint8)
+    w_desc = TensorDescriptor.from_tensor(packed.view(S * N, KB), [BN, BK // 2], w_layout)
+    max_tiles = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN)
+    if cfg["PERSISTENT"]:
+        max_tiles = min(max_tiles, torch.cuda.get_device_properties(a.device).multi_processor_count)
+    _prefill_nvfp4_moe_kernel[(max_tiles,)](
+        a, w_desc, scale.view(torch.uint8), glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+        N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N,
+        a.stride(0), glob.stride(0), c.stride(1), topk_weights_flat.stride(0),
+        K_TAIL=K % 256 != 0, MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k, **cfg,
+    )
diff --git a/python/freetoken/kernel/gluon/nvfp4/utils.py b/python/freetoken/kernel/gluon/nvfp4/utils.py
new file mode 100644
index 0000000000..430c146d49
--- /dev/null
+++ b/python/freetoken/kernel/gluon/nvfp4/utils.py
@@ -0,0 +1,43 @@
+"""NVFP4 helpers shared by the Gluon kernels."""
+
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def bf16x2_mul(a, b):
+    # sm_90+; sm_80-89 need fma.rn.bf16x2 with a -0.0 addend
+    return gl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [a, b], gl.uint32, True, 1)
+
+
+@gluon.jit
+def prmt(a, b, sel: gl.constexpr):
+    return gl.inline_asm_elementwise(f"prmt.b32 $0, $1, $2, {sel};", "=r,r,r", [a, b], gl.uint32, True, 1)
+
+
+@gluon.jit
+def scale_bf16x2(s):
+    """e4m3 scale bytes -> bf16x2 words holding scale * 2^118 twice."""
+    # 2^118 * the 2^-126 of e2m1x8_bf16x2 leaves 2^-8 for the caller to undo; every step is exact in bf16
+    s = (s.to(gl.float8e4nv, bitcast=True).to(gl.float32) * 3.3230699894622794e35).to(gl.bfloat16)
+    s = s.to(gl.uint16, bitcast=True).to(gl.uint32)
+    return s | (s << 16)
+
+
+@gluon.jit
+def e2m1x8_bf16x2(q, s2):
+    """8 e2m1 codes in a uint32 -> 4 bf16x2 words in k order, each value times the bf16x2 scale s2 and 2^-126."""
+    # each 16-bit half's top nibble lands in a bf16 as sign + exponent/mantissa bits, a value of e2m1 * 2^-126
+    o0 = (q & 0x80008000) | ((q >> 6) & 0x01C001C0)
+    q1 = q << 4
+    o1 = (q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0)
+    q2 = q << 8
+    o2 = (q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0)
+    q3 = q << 12
+    o3 = (q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0)
+    # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4)
+    w0 = bf16x2_mul(prmt(o3, o2, 0x5410), s2)
+    w1 = bf16x2_mul(prmt(o1, o0, 0x5410), s2)
+    w2 = bf16x2_mul(prmt(o3, o2, 0x7632), s2)
+    w3 = bf16x2_mul(prmt(o1, o0, 0x7632), s2)
+    return w0, w1, w2, w3
diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py
new file mode 100644
index 0000000000..a61afda567
--- /dev/null
+++ b/python/freetoken/kernel/nvfp4_moe.py
@@ -0,0 +1,336 @@
+"""Kernel dispatch for the inline-dequant NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch.
+
+The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`);
+:mod:`freetoken.moe.fused_nvfp4` strings the GEMMs into the MoE layer.
+"""
+
+from __future__ import annotations
+
+import functools
+import re
+from typing import Any, Dict, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from triton.tools.tensor_descriptor import TensorDescriptor
+
+from freetoken.kernel.backend import is_rocm
+from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view, e4m3_native
+from freetoken.kernel.triton.nvfp4_fused_moe import (
+    _decode_nvfp4_marlin_kernel,
+    _decode_nvfp4_moe_kernel,
+    _e2m1_lut,
+    _prefill_nvfp4_moe_kernel,
+)
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+# Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune,
+# which benchmarks at run time). Tuned offline against the NVFP4 decode kernels.
+# These drive the original LUT-gather decode (decode_gemm), kept only for A/B.
+_DECODE_BLOCK_N = 64
+_DECODE_BLOCK_KB = 128
+_DECODE_WARPS = 4
+
+# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps.
+# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one.
+_DECODE_MARLIN_BLOCK_N = 16
+_DECODE_MARLIN_BLOCK_KW = 16
+_DECODE_MARLIN_UNROLL = 2
+_DECODE_MARLIN_WARPS = 4
+_DECODE_MARLIN_DEEPK_BLOCK_N = 8
+_DECODE_MARLIN_DEEPK_BLOCK_KW = 32
+_DECODE_MARLIN_DEEPK_UNROLL = 2
+_DECODE_MARLIN_DEEPK_THRESHOLD = 2048
+
+
+def _tl_dtype(dt: torch.dtype):
+    if dt == torch.bfloat16:
+        return tl.bfloat16
+    if dt == torch.float16:
+        return tl.float16
+    return tl.float32
+
+
+def decode_gemm(
+    a: torch.Tensor,
+    packed: torch.Tensor,
+    scale: torch.Tensor,
+    glob: torch.Tensor,
+    c: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    mul_routed_weight: bool,
+    a_row_is_route: bool,
+) -> None:
+    M, top_k = topk_ids.shape
+    N = packed.shape[1]
+    K = packed.shape[2] * 2
+    scale = e4m3_kernel_view(scale)
+    total_routes = M * top_k
+    grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N))
+    _decode_nvfp4_moe_kernel[grid](
+        a, packed, scale, glob, c, topk_weights, topk_ids,
+        _e2m1_lut(a.device.index),
+        total_routes, N, K,
+        a.stride(0), a.stride(1),
+        packed.stride(0), packed.stride(1), packed.stride(2),
+        scale.stride(0), scale.stride(1), scale.stride(2),
+        glob.stride(0), glob.stride(1),
+        c.stride(0), c.stride(1), c.stride(2),
+        topk_weights.stride(0), topk_weights.stride(1),
+        topk_ids.stride(0), topk_ids.stride(1),
+        BLOCK_SIZE_N=_DECODE_BLOCK_N,
+        BLOCK_SIZE_KB=_DECODE_BLOCK_KB,
+        TOP_K=top_k,
+        A_ROW_IS_ROUTE=a_row_is_route,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        compute_type=_tl_dtype(c.dtype),
+        num_warps=_DECODE_WARPS,
+    )
+
+
+def decode_gemm_marlin(
+    a: torch.Tensor,
+    packed: torch.Tensor,
+    scale: torch.Tensor,
+    glob: torch.Tensor,
+    c: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    mul_routed_weight: bool,
+    a_row_is_route: bool,
+) -> None:
+    """Marlin-style decode GEMV: int32 wide loads + deferred reduction
+    (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank;
+    it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4)."""
+    M, top_k = topk_ids.shape
+    N = packed.shape[1]
+    K = packed.shape[2] * 2
+    packed_i32 = packed.view(torch.int32)  # [S, N, K // 8]
+    scale = e4m3_kernel_view(scale)
+    total_routes = M * top_k
+    deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD
+    block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N
+    block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW
+    unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL
+    # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned
+    a_wide = (
+        a.dtype in (torch.bfloat16, torch.float16)
+        and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0
+    )
+    grid = (total_routes, triton.cdiv(N, block_n))
+    _decode_nvfp4_marlin_kernel[grid](
+        a, packed_i32, scale, glob, c, topk_weights, topk_ids,
+        total_routes, N, K,
+        a.stride(0), a.stride(1),
+        packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2),
+        scale.stride(0), scale.stride(1), scale.stride(2),
+        glob.stride(0), glob.stride(1),
+        c.stride(0), c.stride(1), c.stride(2),
+        topk_weights.stride(0), topk_weights.stride(1),
+        topk_ids.stride(0), topk_ids.stride(1),
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_KW=block_kw,
+        UNROLL=unroll,
+        TOP_K=top_k,
+        A_ROW_IS_ROUTE=a_row_is_route,
+        A_WIDE=a_wide,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        compute_type=_tl_dtype(c.dtype),
+        num_warps=_DECODE_MARLIN_WARPS,
+    )
+
+
+# rho = bf16 tensor-core FLOP / fp32 ALU FLOP per SM per clock, with fp32 accumulate (the only accumulator for bf16 MMA).
+# Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M.
+# High rho makes decode the cost, so big blocks win; low rho makes padded MMA the cost, so small blocks win.
+# Spec, FLOP/clk/SM: GeForce RTX 30/40/50 512 / 256 = 2, as GeForce runs fp32 accumulate at half rate.
+# RTX A, RTX Ada Generation, RTX PRO Blackwell, A10, A40, L4, L40S 1024 / 256 = 4; A100 2048 / 128 = 16; H100 4096 / 256 = 16; H20 148 / 44 TFLOP/s = 3.4.
+# Measured (cuBLAS bf16 / Triton fp32 FMA): RTX PRO 6000 4.5, RTX 5090 2.3, RTX 4090 1.9, RTX 3080 Ti 1.6.
+_PRO_GPU_RE = re.compile(
+    r"RTX\s*A\d{3,4}"                        # RTX A400 .. A6000
+    r"|RTX\s*\d{3,4}E?\s*(SFF\s*)?-?\s*Ada"  # RTX 500 .. 6000 Ada Generation
+    r"|RTX\s*PRO\s*\d{3,4}"                  # RTX PRO 500 .. 6000 Blackwell
+    r"|RTX\s*\d{4}D\b"                       # RTX 6000D
+    r"|Quadro|TITAN"
+    r"|^(NVIDIA\s+)?(GRID\s+)?(A2|A10[GM]?|A16|A40|L2|L20|L4|L40S?)\b",
+    re.I,
+)
+# also matches names without "GeForce", e.g. "NVIDIA RTX 4090"
+_GEFORCE_GPU_RE = re.compile(r"GeForce|RTX\s*[2-6]0[5-9]0(?!\d)", re.I)
+
+
+def _tensor_rho(name: str, major: int, minor: int) -> float:
+    if major not in (8, 12) or (major, minor) == (8, 0):
+        return 3.4 if re.search(r"\bH20\b", name) else 16.0
+    if _PRO_GPU_RE.search(name):
+        return 4.0
+    if _GEFORCE_GPU_RE.search(name):
+        return 2.0
+    logger.warning("unknown GPU %r: NVFP4 MoE prefill assumes full-rate fp32-accumulate tensor cores", name)
+    return 4.0
+
+
+# (min rho, tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128), from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti
+_PREFILL_BLOCK_M_STEPS = ((8.0, (12, 24, 64)), (3.0, (16, 36, 150)), (0.0, (16, 150, float("inf"))))
+# BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages)
+_PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)}
+_PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)}
+_PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)}
+# steps while the Gluon kernel runs BLOCK_SIZE_M >= 64, from an RTX PRO 6000 sweep; classes without an entry keep the Triton steps
+_PREFILL_BLOCK_M_STEPS_GLUON = ((3.0, (16, 36, 96)),)
+
+
+# Gluon prefill GEMM tiles: BLOCK_SIZE_N -> (BLOCK_SIZE_K, setmaxnreg of the loader / decode warps)
+_GLUON_PREFILL_TILES = {128: (64, 96, 96), 256: (32, 88, 72)}
+
+
+def _gluon_prefill_config(N: int, block_m: int, mul_routed_weight: bool) -> Dict[str, Any]:
+    # BLOCK_SIZE_M 64 tiles are short (few tokens per expert, often short K): persistent CTAs overlap one tile's epilogue and ring fill with the next tile, and a third x / decoded-weight stage absorbs DRAM latency
+    persistent = block_m == 64
+    # a 256-row tile halves the token loads per weight byte, but a partial last tile still decodes all 256 rows
+    # the persistent epilogue keeps one 64-bit address per element in registers, which spills with the routed-weight loads
+    block_n = 256 if N % 256 == 0 and not (persistent and mul_routed_weight) else 128
+    block_k, load_regs, decode_regs = _GLUON_PREFILL_TILES[block_n]
+    return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_K=block_k, GROUP_SIZE_M=8,
+                STAGES=3 if persistent else 2, W_STAGES=4, PERSISTENT=persistent, DECODE_WARPS=4, MMA_WARPS_N=4,
+                LOAD_REGS=load_regs, DECODE_REGS=decode_regs, num_warps=8)
+
+
+@functools.lru_cache(maxsize=None)
+def _gluon_prefill(device_index: int):
+    """The Gluon module whose prefill GEMM replaces the Triton kernel on this GPU, or None."""
+    from freetoken.kernel import gluon
+
+    if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available():
+        return None
+    from freetoken.kernel.gluon.nvfp4.moe import sm120
+
+    return sm120
+
+
+@functools.lru_cache(maxsize=None)
+def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]:
+    if is_rocm():
+        # TODO: tune for AMD GPUs; they keep the H100 tuning
+        return _PREFILL_BLOCK_M_STEPS[0][1], _PREFILL_TILES_H100
+    props = torch.cuda.get_device_properties(device_index)
+    rho = _tensor_rho(props.name, props.major, props.minor)
+    steps = next(s for min_rho, s in _PREFILL_BLOCK_M_STEPS if rho >= min_rho)
+    if _gluon_prefill(device_index) is not None:
+        steps = next((s for min_rho, s in _PREFILL_BLOCK_M_STEPS_GLUON if rho >= min_rho), steps)
+    # the RTX tables come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 one
+    if (props.major, props.minor) not in ((8, 6), (8, 9), (12, 0)):
+        return steps, _PREFILL_TILES_H100
+    return steps, _PREFILL_TILES_RTX_PRO if rho >= 3.0 else _PREFILL_TILES_GEFORCE
+
+
+@functools.lru_cache(maxsize=None)
+def _tma_device(device_index: int) -> bool:
+    return not is_rocm() and torch.cuda.get_device_capability(device_index) == (12, 0)
+
+
+def _prefill_use_tma(packed: torch.Tensor, scale: torch.Tensor, K: int, device_index: int) -> bool:
+    # a TMA box row is >= 16 bytes and every global stride a multiple of 16 bytes, so K // 16 scale bytes need K % 256 == 0;
+    # the TMA scale path reads the scales as fp8, so the forced e4m3 emulation stays on pointer loads
+    return (_tma_device(device_index) and e4m3_native() and K % 256 == 0
+            and packed.is_contiguous() and scale.is_contiguous()
+            and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0)
+
+
+@functools.lru_cache(maxsize=None)
+def _device_cc(device_index: int) -> Tuple[int, int]:
+    return torch.cuda.get_device_capability(device_index)
+
+
+def _prefill_fast_decode(block_m: int, per_expert: float, device_index: int) -> int:
+    """FAST_DECODE of the prefill kernel: 0 keeps the fp16 pair decode (ROCm: the fast decode is PTX inline asm), 1 words, 2 bytes."""
+    if is_rocm():
+        return 0
+    if _device_cc(device_index) == (8, 9):
+        # RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14%
+        return 0 if per_expert < 48 else 1
+    if _device_cc(device_index) == (9, 0):
+        # H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM)
+        return 2
+    # the word decode halves the loop but stages the decoded tile through shared memory; small tiles keep the byte form
+    return 2 if block_m <= 32 else 1
+
+
+def prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]:
+    # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it.
+    (to32, to64, to128), tiles = _prefill_tuning(device_index)
+    per_expert = M * top_k / num_experts
+    block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16
+    block_n, block_kb, stages = tiles[block_m]
+    return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb,
+                GROUP_SIZE_M=8, num_warps=4, num_stages=stages,
+                FAST_DECODE=_prefill_fast_decode(block_m, per_expert, device_index))
+
+
+def prefill_gemm(
+    a: torch.Tensor,
+    packed: torch.Tensor,
+    scale: torch.Tensor,
+    glob: torch.Tensor,
+    c: torch.Tensor,
+    topk_weights_flat: torch.Tensor,
+    sorted_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    num_valid_tokens: int,
+    kernel_top_k: int,
+    mul_routed_weight: bool,
+    cfg: Dict[str, Any],
+) -> None:
+    gluon = _gluon_prefill(a.device.index)
+    if gluon is not None and cfg["BLOCK_SIZE_M"] >= 64 and gluon.supports(a, packed, scale, c):
+        gluon.prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+                           num_valid_tokens, kernel_top_k, mul_routed_weight,
+                           _gluon_prefill_config(packed.shape[1], cfg["BLOCK_SIZE_M"], mul_routed_weight))
+        return
+    N = packed.shape[1]
+    K = packed.shape[2] * 2
+    EM = sorted_ids.shape[0]
+    scale = e4m3_kernel_view(scale)
+    grid = lambda META: (  # noqa: E731
+        triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+    )
+    cfg = dict(cfg)
+    fast = cfg.pop("FAST_DECODE", 0)
+    if K % 8:
+        fast = 0
+    if fast == 1:
+        packed = packed.view(torch.int32)
+    use_tma = _prefill_use_tma(packed, scale, K, a.device.index)
+    # the 16-byte scale box costs registers that the BLOCK_SIZE_M >= 64 tiles do not have
+    tma_scales = cfg["BLOCK_SIZE_M"] <= 32
+    w_desc = s_desc = None
+    if use_tma:
+        rows = packed.shape[0] * N
+        box_k = cfg["BLOCK_SIZE_KB"] // 4 if fast == 1 else cfg["BLOCK_SIZE_KB"]
+        w_desc = TensorDescriptor.from_tensor(packed.view(rows, packed.shape[2]), [cfg["BLOCK_SIZE_N"], box_k])
+        if tma_scales:
+            s_desc = TensorDescriptor.from_tensor(scale.view(rows, K // 16), [cfg["BLOCK_SIZE_N"], 16])
+    _prefill_nvfp4_moe_kernel[grid](
+        a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids,
+        num_tokens_post_padded,
+        w_desc, s_desc,
+        N, K, EM, num_valid_tokens,
+        a.stride(0), a.stride(1),
+        packed.stride(0), packed.stride(1), packed.stride(2),
+        scale.stride(0), scale.stride(1), scale.stride(2),
+        glob.stride(0), glob.stride(1),
+        c.stride(1), c.stride(2),
+        topk_weights_flat.stride(0),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=kernel_top_k,
+        compute_type=_tl_dtype(c.dtype),
+        USE_TMA=use_tma,
+        TMA_SCALES=tma_scales,
+        FAST_DECODE=fast,
+        **cfg,
+    )
diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
index ca9b9dbe9d..0fa4d4aa9a 100644
--- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py
+++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
@@ -28,6 +28,8 @@
 import torch
 import triton
 import triton.language as tl
+from triton.language import target_info
+from triton.runtime.jit import constexpr_function
 
 from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32
 from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32
@@ -299,6 +301,57 @@ def _e2m1_f16_x2p14(code):
     return bits.to(tl.uint16).to(tl.float16, bitcast=True)
 
 
+@constexpr_function
+def _bf16x2_mul_native():
+    return target_info.cuda_capability_geq(9, 0)
+
+
+@triton.jit
+def _bf16x2_scale(w, s2):
+    if _bf16x2_mul_native():
+        return tl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [w, s2], tl.uint32, True, 1)
+    else:
+        # mul.bf16x2 is sm_90+; fma with a -0.0 addend gives the same product
+        return tl.inline_asm_elementwise("{ .reg .b32 z; mov.b32 z, 0x80008000; fma.rn.bf16x2 $0, $1, $2, z; }",
+                                         "=r,r,r", [w, s2], tl.uint32, True, 1)
+
+
+@triton.jit
+def _bf16x2_halves(w):
+    lo = (w & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
+    hi = (w >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
+    return lo, hi
+
+
+@triton.jit
+def _e2m1x8_bf16(q, s2, BLOCK_N: tl.constexpr, BLOCK_W: tl.constexpr):
+    """[BLOCK_N, BLOCK_W] uint32 words of 8 e2m1 codes -> [BLOCK_N, 8 * BLOCK_W] bf16 in k order, times s2 * 2^-126."""
+    # each 16-bit half's top nibble becomes a bf16 sign + exponent/mantissa: e2m1 * 2^-126, exact incl. e == 0
+    o0 = _bf16x2_scale((q & 0x80008000) | ((q >> 6) & 0x01C001C0), s2)
+    q1 = q << 4
+    o1 = _bf16x2_scale((q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0), s2)
+    q2 = q << 8
+    o2 = _bf16x2_scale((q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0), s2)
+    q3 = q << 12
+    o3 = _bf16x2_scale((q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0), s2)
+    # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4); join appends the new axis last
+    k0, k4 = _bf16x2_halves(o3)
+    k1, k5 = _bf16x2_halves(o2)
+    k2, k6 = _bf16x2_halves(o1)
+    k3, k7 = _bf16x2_halves(o0)
+    w = tl.join(tl.join(tl.join(k0, k4), tl.join(k2, k6)), tl.join(tl.join(k1, k5), tl.join(k3, k7)))
+    return tl.reshape(w, (BLOCK_N, 8 * BLOCK_W))
+
+
+@triton.jit
+def _e2m1x2_bf16(v, s2, BLOCK_N: tl.constexpr, BLOCK_KB: tl.constexpr):
+    """[BLOCK_N, BLOCK_KB] bytes of 2 e2m1 codes -> [BLOCK_N, 2 * BLOCK_KB] bf16 in k order, times s2 * 2^-126."""
+    # low nibble -> low bf16, high nibble -> high bf16: sign to bit 15 / 31, exponent and mantissa to bits 8..6 / 24..22
+    pair = ((v & 0x8) << 12) | ((v & 0x7) << 6) | ((v & 0x80) << 24) | ((v & 0x70) << 18)
+    lo, hi = _bf16x2_halves(_bf16x2_scale(pair.to(tl.uint32, bitcast=True), s2))
+    return tl.reshape(tl.join(lo, hi), (BLOCK_N, 2 * BLOCK_KB))
+
+
 @triton.jit
 def _prefill_nvfp4_moe_kernel(
     a_ptr,             # [M, K] activations
@@ -310,6 +363,8 @@ def _prefill_nvfp4_moe_kernel(
     sorted_token_ids_ptr,
     expert_ids_ptr,    # cache slot per M-block
     num_tokens_post_padded_ptr,
+    w_desc,            # TMA descriptor over packed as [S * N, K // 2], or None
+    s_desc,            # TMA descriptor over scale as [S * N, K // 16], or None
     N,
     K,
     EM,
@@ -327,6 +382,9 @@ def _prefill_nvfp4_moe_kernel(
     MUL_ROUTED_WEIGHT: tl.constexpr,
     top_k: tl.constexpr,
     compute_type: tl.constexpr,
+    USE_TMA: tl.constexpr = False,     # needs K % 256 == 0
+    TMA_SCALES: tl.constexpr = False,  # with USE_TMA: scales by TMA too, else pointer loads
+    FAST_DECODE: tl.constexpr = 0,     # 0: fp16 pair decode, 1: packed_ptr / w_desc are int32 views of 8 codes, 2: bytes
 ):
     # out^T = W @ X^T: wgmma reads its A operand (the decoded weights) from registers, while B must be in shared memory.
     # The token block is the MMA N dimension, so a 16-token block still fills the MMA.
@@ -352,38 +410,82 @@ def _prefill_nvfp4_moe_kernel(
     offs_kb = tl.arange(0, BLOCK_SIZE_KB)
     offs_k = tl.arange(0, 2 * BLOCK_SIZE_KB)
     offs_sb = tl.arange(0, BLOCK_SIZE_KB // 8)
+    KW: tl.constexpr = BLOCK_SIZE_KB // 4
+    offs_kw = tl.arange(0, KW)
     xt_ptrs = a_ptr + (offs_token[None, :] // top_k * stride_am + offs_k[:, None] * stride_ak)
 
     slot = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
-    p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb
+    WORDS: tl.constexpr = FAST_DECODE == 1
+    if WORDS:
+        p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kw[None, :] * stride_pkb
+    else:
+        p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb
     s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk
 
     accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
     K_BYTES = K // 2
+    # rows past N belong to the next slot; their outputs are masked on store
+    w_row = (slot * N + pid_n * BLOCK_SIZE_N).to(tl.int32)
+    SB: tl.constexpr = BLOCK_SIZE_KB // 8
     for kb in range(0, tl.cdiv(K_BYTES, BLOCK_SIZE_KB)):
-        byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES
-        sblk_mask = kb * (BLOCK_SIZE_KB // 8) + offs_sb < K_BYTES // 8
         k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K
-        v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32)
-        if e4m3_native_cx():
+        if USE_TMA and WORDS:
+            v = w_desc.load([w_row, kb * KW])
+        elif USE_TMA:
+            v = w_desc.load([w_row, kb * BLOCK_SIZE_KB]).to(tl.int32)
+        if USE_TMA and TMA_SCALES:
+            # a TMA box row needs >= 16 bytes, so load the 16 scales that hold this step and pick its SB of them
+            s16 = s_desc.load([w_row, (kb * SB) // 16 * 16]).to(tl.uint8, bitcast=True)
+            s3 = tl.reshape(s16, (BLOCK_SIZE_N, 16 // SB, SB))
+            pick = tl.arange(0, 16 // SB) == (kb * SB) % 16 // SB
+            s = tl.sum(tl.where(pick[None, :, None], s3, 0), axis=1).to(tl.uint8)
+            s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+        elif USE_TMA:
+            sblk_mask = kb * SB + offs_sb < K_BYTES // 8
             s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32)
         else:
-            s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0))
-        # one scale per 16 values; the 2^14 undoes the decode's 2^-14
-        s = tl.reshape(tl.broadcast_to((s * 16384.0)[:, :, None], (BLOCK_SIZE_N, BLOCK_SIZE_KB // 8, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB))
-        s = s.to(tl.bfloat16)
-        # exact in bf16: e2m1 has <= 2 significant bits and e4m3 <= 4
-        w_lo = _e2m1_f16_x2p14(v & 0xF).to(tl.bfloat16) * s
-        w_hi = _e2m1_f16_x2p14(v >> 4).to(tl.bfloat16) * s
-        # interleave back to k = 2 * byte + {0, 1} so one dot covers the K step
-        w = tl.reshape(tl.join(w_lo, w_hi), (BLOCK_SIZE_N, 2 * BLOCK_SIZE_KB))
+            sblk_mask = kb * SB + offs_sb < K_BYTES // 8
+            if WORDS:
+                v = tl.load(p_ptrs, mask=(kb * KW + offs_kw < K_BYTES // 4)[None, :], other=0)
+            else:
+                byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES
+                v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32)
+            if e4m3_native_cx():
+                s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32)
+            else:
+                s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0))
+        if FAST_DECODE:
+            # 2^118 here and the decode's 2^-126 leave 2^-8, undone in the epilogue; every step is exact in bf16
+            s = (s * 3.3230699894622794e35).to(tl.bfloat16).to(tl.uint16, bitcast=True).to(tl.uint32)
+            s = s | (s << 16)
+            if WORDS:
+                # one scale per 16 values, two words
+                s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW))
+                w = _e2m1x8_bf16(v.to(tl.uint32, bitcast=True), s, BLOCK_SIZE_N, KW)
+            else:
+                # one scale per 16 values, eight bytes
+                s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB))
+                w = _e2m1x2_bf16(v, s, BLOCK_SIZE_N, BLOCK_SIZE_KB)
+        else:
+            # one scale per 16 values; the 2^14 undoes the decode's 2^-14
+            s = tl.reshape(tl.broadcast_to((s * 16384.0)[:, :, None], (BLOCK_SIZE_N, BLOCK_SIZE_KB // 8, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB))
+            s = s.to(tl.bfloat16)
+            # exact in bf16: e2m1 has <= 2 significant bits and e4m3 <= 4
+            w_lo = _e2m1_f16_x2p14(v & 0xF).to(tl.bfloat16) * s
+            w_hi = _e2m1_f16_x2p14(v >> 4).to(tl.bfloat16) * s
+            # interleave back to k = 2 * byte + {0, 1} so one dot covers the K step
+            w = tl.reshape(tl.join(w_lo, w_hi), (BLOCK_SIZE_N, 2 * BLOCK_SIZE_KB))
         xt = tl.load(xt_ptrs, mask=k_mask[:, None] & token_mask[None, :], other=0.0)
         accumulator = tl.dot(w.to(xt.dtype), xt, accumulator)
         xt_ptrs += 2 * BLOCK_SIZE_KB * stride_ak
-        p_ptrs += BLOCK_SIZE_KB * stride_pkb
-        s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk
+        if not USE_TMA:
+            p_ptrs += (KW if WORDS else BLOCK_SIZE_KB) * stride_pkb
+        if not (USE_TMA and TMA_SCALES):
+            s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk
 
     g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32)
+    if FAST_DECODE:
+        g = g * 256.0
     accumulator = accumulator * g[:, None]
 
     if MUL_ROUTED_WEIGHT:
diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py
index 9202f5f38a..197ef54727 100644
--- a/python/freetoken/moe/fused_nvfp4.py
+++ b/python/freetoken/moe/fused_nvfp4.py
@@ -7,144 +7,11 @@
 
 from __future__ import annotations
 
-import functools
-import re
-from typing import Any, Dict, Tuple
-
 import torch
-import triton
-import triton.language as tl
 
-from freetoken.kernel import moe_sum_reduce_triton
-from freetoken.kernel.backend import is_rocm
-from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view
-from freetoken.kernel.triton.nvfp4_fused_moe import (
-    _decode_nvfp4_marlin_kernel,
-    _decode_nvfp4_moe_kernel,
-    _e2m1_lut,
-    _prefill_nvfp4_moe_kernel,
-)
+from freetoken.kernel import moe_sum_reduce_triton, nvfp4_moe
 from freetoken.layers import gated_act_and_mul
 from freetoken.moe.fused import moe_align_block_size
-from freetoken.utils import init_logger
-
-logger = init_logger(__name__)
-
-# Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune,
-# which benchmarks at run time). Tuned offline against the NVFP4 decode kernels.
-# These drive the original LUT-gather decode (_decode_gemm), kept only for A/B.
-_DECODE_BLOCK_N = 64
-_DECODE_BLOCK_KB = 128
-_DECODE_WARPS = 4
-
-# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps.
-# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one.
-_DECODE_MARLIN_BLOCK_N = 16
-_DECODE_MARLIN_BLOCK_KW = 16
-_DECODE_MARLIN_UNROLL = 2
-_DECODE_MARLIN_WARPS = 4
-_DECODE_MARLIN_DEEPK_BLOCK_N = 8
-_DECODE_MARLIN_DEEPK_BLOCK_KW = 32
-_DECODE_MARLIN_DEEPK_UNROLL = 2
-_DECODE_MARLIN_DEEPK_THRESHOLD = 2048
-
-
-def _tl_dtype(dt: torch.dtype):
-    if dt == torch.bfloat16:
-        return tl.bfloat16
-    if dt == torch.float16:
-        return tl.float16
-    return tl.float32
-
-
-def _decode_gemm(
-    a: torch.Tensor,
-    packed: torch.Tensor,
-    scale: torch.Tensor,
-    glob: torch.Tensor,
-    c: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    mul_routed_weight: bool,
-    a_row_is_route: bool,
-) -> None:
-    M, top_k = topk_ids.shape
-    N = packed.shape[1]
-    K = packed.shape[2] * 2
-    scale = e4m3_kernel_view(scale)
-    total_routes = M * top_k
-    grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N))
-    _decode_nvfp4_moe_kernel[grid](
-        a, packed, scale, glob, c, topk_weights, topk_ids,
-        _e2m1_lut(a.device.index),
-        total_routes, N, K,
-        a.stride(0), a.stride(1),
-        packed.stride(0), packed.stride(1), packed.stride(2),
-        scale.stride(0), scale.stride(1), scale.stride(2),
-        glob.stride(0), glob.stride(1),
-        c.stride(0), c.stride(1), c.stride(2),
-        topk_weights.stride(0), topk_weights.stride(1),
-        topk_ids.stride(0), topk_ids.stride(1),
-        BLOCK_SIZE_N=_DECODE_BLOCK_N,
-        BLOCK_SIZE_KB=_DECODE_BLOCK_KB,
-        TOP_K=top_k,
-        A_ROW_IS_ROUTE=a_row_is_route,
-        MUL_ROUTED_WEIGHT=mul_routed_weight,
-        compute_type=_tl_dtype(c.dtype),
-        num_warps=_DECODE_WARPS,
-    )
-
-
-def _decode_gemm_marlin(
-    a: torch.Tensor,
-    packed: torch.Tensor,
-    scale: torch.Tensor,
-    glob: torch.Tensor,
-    c: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    mul_routed_weight: bool,
-    a_row_is_route: bool,
-) -> None:
-    """Marlin-style decode GEMV: int32 wide loads + deferred reduction
-    (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank;
-    it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4)."""
-    M, top_k = topk_ids.shape
-    N = packed.shape[1]
-    K = packed.shape[2] * 2
-    packed_i32 = packed.view(torch.int32)  # [S, N, K // 8]
-    scale = e4m3_kernel_view(scale)
-    total_routes = M * top_k
-    deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD
-    block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N
-    block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW
-    unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL
-    # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned
-    a_wide = (
-        a.dtype in (torch.bfloat16, torch.float16)
-        and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0
-    )
-    grid = (total_routes, triton.cdiv(N, block_n))
-    _decode_nvfp4_marlin_kernel[grid](
-        a, packed_i32, scale, glob, c, topk_weights, topk_ids,
-        total_routes, N, K,
-        a.stride(0), a.stride(1),
-        packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2),
-        scale.stride(0), scale.stride(1), scale.stride(2),
-        glob.stride(0), glob.stride(1),
-        c.stride(0), c.stride(1), c.stride(2),
-        topk_weights.stride(0), topk_weights.stride(1),
-        topk_ids.stride(0), topk_ids.stride(1),
-        BLOCK_SIZE_N=block_n,
-        BLOCK_SIZE_KW=block_kw,
-        UNROLL=unroll,
-        TOP_K=top_k,
-        A_ROW_IS_ROUTE=a_row_is_route,
-        A_WIDE=a_wide,
-        MUL_ROUTED_WEIGHT=mul_routed_weight,
-        compute_type=_tl_dtype(c.dtype),
-        num_warps=_DECODE_MARLIN_WARPS,
-    )
 
 
 def _fused_experts_decode_nvfp4(
@@ -164,8 +31,8 @@ def _fused_experts_decode_nvfp4(
     act_limit: float = 7.0,
 ) -> torch.Tensor:
     """Shared decode body (gemm1 -> act -> gemm2 -> sum-reduce); ``gemm_fn`` is either
-    the marlin-style int32 GEMV (:func:`_decode_gemm_marlin`) or the original LUT-gather
-    GEMV (:func:`_decode_gemm`), both with the same calling convention."""
+    the marlin-style int32 GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm_marlin`) or the original LUT-gather
+    GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm`), both with the same calling convention."""
     M, H = hidden_states.shape
     top_k = topk_ids.shape[1]
     two_i = gate_up_packed.shape[1]
@@ -206,7 +73,7 @@ def fused_experts_decode_nvfp4_marlin(
 ) -> torch.Tensor:
     """Decode inline-NVFP4 MoE using the Marlin-style int32 wide-load GEMV."""
     return _fused_experts_decode_nvfp4(
-        _decode_gemm_marlin,
+        nvfp4_moe.decode_gemm_marlin,
         hidden_states, gate_up_packed, gate_up_scale, gate_up_global,
         down_packed, down_scale, down_global,
         topk_weights, topk_ids, activation, apply_router_weight_on_input,
@@ -232,7 +99,7 @@ def fused_experts_decode_nvfp4_serial(
     """Original LUT-gather decode (one program per route, full K reduction). Retained for
     A/B benchmarking against the marlin decode path; not on the production decode path."""
     return _fused_experts_decode_nvfp4(
-        _decode_gemm,
+        nvfp4_moe.decode_gemm,
         hidden_states, gate_up_packed, gate_up_scale, gate_up_global,
         down_packed, down_scale, down_global,
         topk_weights, topk_ids, activation, apply_router_weight_on_input,
@@ -240,107 +107,6 @@ def fused_experts_decode_nvfp4_serial(
     )
 
 
-# rho = bf16 tensor-core FLOP / fp32 ALU FLOP per SM per clock, with fp32 accumulate (the only accumulator for bf16 MMA).
-# Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M.
-# High rho makes decode the cost, so big blocks win; low rho makes padded MMA the cost, so small blocks win.
-# Spec, FLOP/clk/SM: GeForce RTX 30/40/50 512 / 256 = 2, as GeForce runs fp32 accumulate at half rate.
-# RTX A, RTX Ada Generation, RTX PRO Blackwell, A10, A40, L4, L40S 1024 / 256 = 4; A100 2048 / 128 = 16; H100 4096 / 256 = 16; H20 148 / 44 TFLOP/s = 3.4.
-# Measured (cuBLAS bf16 / Triton fp32 FMA): RTX PRO 6000 4.5, RTX 5090 2.3, RTX 4090 1.9, RTX 3080 Ti 1.6.
-_PRO_GPU_RE = re.compile(
-    r"RTX\s*A\d{3,4}"                        # RTX A400 .. A6000
-    r"|RTX\s*\d{3,4}E?\s*(SFF\s*)?-?\s*Ada"  # RTX 500 .. 6000 Ada Generation
-    r"|RTX\s*PRO\s*\d{3,4}"                  # RTX PRO 500 .. 6000 Blackwell
-    r"|RTX\s*\d{4}D\b"                       # RTX 6000D
-    r"|Quadro|TITAN"
-    r"|^(NVIDIA\s+)?(GRID\s+)?(A2|A10[GM]?|A16|A40|L2|L20|L4|L40S?)\b",
-    re.I,
-)
-# also matches names without "GeForce", e.g. "NVIDIA RTX 4090"
-_GEFORCE_GPU_RE = re.compile(r"GeForce|RTX\s*[2-6]0[5-9]0(?!\d)", re.I)
-
-
-def _tensor_rho(name: str, major: int, minor: int) -> float:
-    if major not in (8, 12) or (major, minor) == (8, 0):
-        return 3.4 if re.search(r"\bH20\b", name) else 16.0
-    if _PRO_GPU_RE.search(name):
-        return 4.0
-    if _GEFORCE_GPU_RE.search(name):
-        return 2.0
-    logger.warning("unknown GPU %r: NVFP4 MoE prefill assumes full-rate fp32-accumulate tensor cores", name)
-    return 4.0
-
-
-# (min rho, tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128), from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti
-_PREFILL_BLOCK_M_STEPS = ((8.0, (12, 24, 64)), (3.0, (16, 36, 150)), (0.0, (16, 150, float("inf"))))
-# BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages)
-_PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)}
-_PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)}
-_PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)}
-
-
-@functools.lru_cache(maxsize=None)
-def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]:
-    if is_rocm():
-        # TODO: tune for AMD GPUs; they keep the H100 tuning
-        return _PREFILL_BLOCK_M_STEPS[0][1], _PREFILL_TILES_H100
-    props = torch.cuda.get_device_properties(device_index)
-    rho = _tensor_rho(props.name, props.major, props.minor)
-    steps = next(s for min_rho, s in _PREFILL_BLOCK_M_STEPS if rho >= min_rho)
-    # the RTX tables come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 one
-    if (props.major, props.minor) not in ((8, 6), (8, 9), (12, 0)):
-        return steps, _PREFILL_TILES_H100
-    return steps, _PREFILL_TILES_RTX_PRO if rho >= 3.0 else _PREFILL_TILES_GEFORCE
-
-
-def _prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]:
-    # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it.
-    (to32, to64, to128), tiles = _prefill_tuning(device_index)
-    per_expert = M * top_k / num_experts
-    block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16
-    block_n, block_kb, stages = tiles[block_m]
-    return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb,
-                GROUP_SIZE_M=8, num_warps=4, num_stages=stages)
-
-
-def _prefill_gemm(
-    a: torch.Tensor,
-    packed: torch.Tensor,
-    scale: torch.Tensor,
-    glob: torch.Tensor,
-    c: torch.Tensor,
-    topk_weights_flat: torch.Tensor,
-    sorted_ids: torch.Tensor,
-    expert_ids: torch.Tensor,
-    num_tokens_post_padded: torch.Tensor,
-    num_valid_tokens: int,
-    kernel_top_k: int,
-    mul_routed_weight: bool,
-    cfg: Dict[str, Any],
-) -> None:
-    N = packed.shape[1]
-    K = packed.shape[2] * 2
-    EM = sorted_ids.shape[0]
-    scale = e4m3_kernel_view(scale)
-    grid = lambda META: (  # noqa: E731
-        triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
-    )
-    _prefill_nvfp4_moe_kernel[grid](
-        a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids,
-        num_tokens_post_padded,
-        N, K, EM, num_valid_tokens,
-        a.stride(0), a.stride(1),
-        packed.stride(0), packed.stride(1), packed.stride(2),
-        scale.stride(0), scale.stride(1), scale.stride(2),
-        glob.stride(0), glob.stride(1),
-        c.stride(1), c.stride(2),
-        topk_weights_flat.stride(0),
-        MUL_ROUTED_WEIGHT=mul_routed_weight,
-        top_k=kernel_top_k,
-        compute_type=_tl_dtype(c.dtype),
-        **cfg,
-    )
-
-
 def fused_experts_nvfp4(
     hidden_states: torch.Tensor,
     gate_up_packed: torch.Tensor,
@@ -365,14 +131,14 @@ def fused_experts_nvfp4(
     two_i = gate_up_packed.shape[1]
     inter = two_i // 2
     dev, dt = hidden_states.device, hidden_states.dtype
-    cfg = _prefill_config(M, top_k, num_experts, hidden_states.device.index)
+    cfg = nvfp4_moe.prefill_config(M, top_k, num_experts, hidden_states.device.index)
 
     sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
     tw = topk_weights.reshape(-1).contiguous()
     num_valid = topk_ids.numel()
 
     ic1 = torch.empty((M, top_k, two_i), device=dev, dtype=dt)
-    _prefill_gemm(
+    nvfp4_moe.prefill_gemm(
         hidden_states, gate_up_packed, gate_up_scale, gate_up_global, ic1,
         tw, sorted_ids, expert_ids, ntpp, num_valid, top_k,
         apply_router_weight_on_input, cfg,
@@ -380,7 +146,7 @@ def fused_experts_nvfp4(
     ic2 = torch.empty((M * top_k, inter), device=dev, dtype=dt)
     gated_act_and_mul(activation, ic1.view(-1, two_i), ic2, alpha=act_alpha, limit=act_limit)
     ic3 = torch.empty((M, top_k, H), device=dev, dtype=dt)
-    _prefill_gemm(
+    nvfp4_moe.prefill_gemm(
         ic2, down_packed, down_scale, down_global, ic3,
         tw, sorted_ids, expert_ids, ntpp, num_valid, 1,
         not apply_router_weight_on_input, cfg,
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index dbf06bcc6c..29c042ba4a 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -10,6 +10,7 @@
     baseline-kernel equality guard. This is the path used on sm_120 + CUDA 12.x.
   - Marlin (sm_80..sm_99, e.g. H100): prefill + decode + overlap.
   - b12x (sm_120 + CUDA>=13): pack + fused decode, skipped where the kernel cannot run.
+  - Gluon (sm_120): the prefill GEMMs it takes over from the Triton kernel, against that kernel.
 Banks are built the way the loader builds them: the bound kernel object packs native pieces.
 """
 
@@ -366,6 +367,52 @@ def test_triton_swigluoai_matches_dequant_reference():
     _assert_close(out, ref)
 
 
+@cuda
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0), reason="needs sm_120")
+@pytest.mark.parametrize("hidden_size,inter", [(256, 128), (384, 192)])
+@pytest.mark.parametrize("block_m", [64, 128])
+def test_triton_prefill_routes_to_gluon_on_sm120(monkeypatch, hidden_size, inter, block_m):
+    """Both prefill GEMMs go to the Gluon kernel and match the Triton kernel: persistent (BLOCK_SIZE_M 64) and not (128), 256- and 128-row tiles, K % 256 == 0 and not."""
+    import freetoken.moe.fused_nvfp4 as fused_nvfp4
+    from freetoken.kernel import gluon, nvfp4_moe
+
+    if not gluon.available():
+        pytest.skip(f"Gluon is off or needs Triton {'.'.join(map(str, gluon.TRITON_VERSION))}")
+    from freetoken.kernel.gluon.nvfp4.moe import sm120
+
+    # the GPU class picks BLOCK_SIZE_M from tokens per expert; pin it so every sm_120 class runs both kernel modes
+    config = nvfp4_moe.prefill_config
+    monkeypatch.setattr(nvfp4_moe, "prefill_config", lambda *args: {**config(*args), "BLOCK_SIZE_M": block_m})
+
+    device = torch.device("cuda")
+    g = torch.Generator(device=device).manual_seed(0)
+
+    def bank(n, k):
+        packed = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=device, generator=g)
+        scale = (torch.rand(E, n, k // 16, device=device, generator=g) + 0.25).to(torch.float8_e4m3fn)
+        return packed, scale, (torch.rand(E, n, device=device, generator=g) + 0.5).to(torch.float16)
+
+    banks = (*bank(2 * inter, hidden_size), *bank(hidden_size, inter))
+    M = 1024
+    hidden = torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device, generator=g) / 4
+    topk_ids = torch.randint(0, E, (M, TOPK), dtype=torch.int32, device=device, generator=g)
+    topk_weights = torch.rand(M, TOPK, device=device, generator=g)
+
+    calls = []
+    gluon_gemm = sm120.prefill_gemm
+
+    def counted(*args):
+        calls.append(args)
+        return gluon_gemm(*args)
+
+    monkeypatch.setattr(sm120, "prefill_gemm", counted)
+    out = fused_nvfp4.fused_experts_nvfp4(hidden, *banks, topk_weights, topk_ids, E)
+    assert len(calls) == 2
+    monkeypatch.setattr(nvfp4_moe, "_gluon_prefill", lambda device_index: None)
+    ref = fused_nvfp4.fused_experts_nvfp4(hidden, *banks, topk_weights, topk_ids, E)
+    torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max()))
+
+
 def _triton_cache(device, *, cache_size=S, prefill_overlap=False):
     """Native 6-bank NVFP4 cache (no repack), consumed directly by the Triton kernels.
     The banks are not transformed, so ``sources`` doubles as the dequant reference."""

From 8f49e0af98ba444dd261af69e82d5a1841778c68 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sat, 17 Oct 2026 23:25:12 -0700
Subject: [PATCH 122/249] feat(moe): bench hybrid fetch fractions per moe layer
 kind (#734)

* feat(moe): bench hybrid fetch fractions per moe layer kind

The ft bench bw profile (version 6) holds only layer_kinds rows, keyed by CPU format, options, expert shape, top_k and CPU worker threads; older profiles are ignored. ft bench bw --model also takes checkpoint paths and benches each of their layer kinds. Hybrid decode reads each layer's own row for its fetch fraction, else keeps a fixed cap of 1. --moe-strategy auto no longer upgrades to hybrid from the profile.

* fix(bench): fill gguf mxfp4 scales, scope act_block to ds_fp4 and skip cpu-less rows
---
 docs/cli.md                           |  32 ++-
 docs/models.md                        |   8 +-
 python/freetoken/engine/engine.py     | 132 +++++-----
 python/freetoken/moe/bench_profile.py | 127 +++------
 python/freetoken/moe/benchbw.py       | 356 ++++++++++++++++++--------
 python/freetoken/moe/offload_cache.py |   9 +-
 python/freetoken/server/args.py       |   8 +-
 tests/engine/test_cache_budget.py     |   4 +-
 tests/moe/test_hybrid_fetch.py        |  51 ++--
 tests/moe/test_offload.py             |   2 +-
 10 files changed, 392 insertions(+), 337 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 69848f2e44..e109bde0c9 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -11,7 +11,7 @@ ft <command> [args]
 | `ft ctl` | Query and manage a running server over HTTP |
 | `ft launch` | Configure and launch a coding agent against a server |
 | `ft checkpoint` | Convert an HF checkpoint to the FTW fast-load format |
-| `ft bench bw` | Benchmark CPU vs PCIe bandwidth to calibrate the MoE backend |
+| `ft bench bw` | Benchmark CPU vs PCIe bandwidth to calibrate hybrid MoE decode |
 
 `ft --version` prints the installed version (torch-free; nightly wheels carry a
 `+g<sha>` build stamp, tagged releases a bare version). Every command supports
@@ -81,7 +81,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 
 | Flag | Default | Meaning |
 |---|---|---|
-| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling |
+| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling |
 | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name |
 | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=<marlin\|b12x\|triton>` (`flashinfer` means b12x); cannot be combined with `--quant-backend` |
 | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-mem` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / memory such as `12GiB` / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) |
@@ -197,20 +197,24 @@ strategies alike (DeepSeek-V4 experts serve offload only);
 ## ft bench bw
 
 ```bash
-ft bench bw                       # once per GPU
-ft bench bw --dtype nvfp4,bf16    # only the formats you serve
-ft bench bw --gpu 1               # a specific GPU (UUID or nvidia-smi index, as for ft serve)
+ft bench bw --model /path/to/model.gguf   # every MoE layer kind of a checkpoint (GGUF or HF)
+ft bench bw                               # canonical geometry per expert format
+ft bench bw --gpu 1                       # a specific GPU (UUID or nvidia-smi index, as for ft serve)
 ```
 
-Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes a
-profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then read.
-
-- One profile per GPU, at `~/.cache/freetoken/benchbw/<gpu-uuid>.json`.
-- Keyed on expert format + GPU, so a profile from other hardware is ignored rather than
-  misapplied. An older single `benchbw.json` still counts if its GPU name matches.
-- What to measure: `--dtype`, `--model`, `--formats`, `--isa`.
-- `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe
-  by that factor.
+Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes one
+row per MoE layer kind, which `ft serve --moe-strategy hybrid --moe-hybrid-max-fetch -1` reads
+to split each layer's misses.
+
+- One profile per GPU, at `~/.cache/freetoken/benchbw/<gpu-uuid>.json`; a profile from other
+  hardware or another profile version is ignored.
+- A row's key is the layer's CPU format and options, expert shape, top_k and CPU worker
+  threads; a layer without an exact row keeps a fetch cap of 1. Bench with the same
+  `--cpu-threads` as the server's `--moe-cpu-threads`.
+- Each run replaces the rows it benched and keeps the others.
+- What to measure: `--model` (checkpoint paths or presets), `--dtype`, `--formats`, `--isa`.
+- `--threshold` (default 2.0) sets each row's `recommended`: hybrid when CPU bandwidth beats
+  PCIe by that factor.
 
 On NVLink-C2C hosts, see [models.md](models.md#nvlink-c2c-hosts-gh200--gb200)
 for `FREETOKEN_H2D_BLOCKS_PER_BANK` if `PCIe-gather` sits well below the linear H2D ceiling.
diff --git a/docs/models.md b/docs/models.md
index d3f26a086b..4418d92b2b 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -47,10 +47,10 @@ These families accept image input by default; pass `--text-model-only` to skip t
   misses stream over PCIe.
 - **cpu** — misses are computed on the CPU instead of fetched.
 - **hybrid** — per step, fetches some misses over PCIe and computes the rest on
-  CPU, overlapped. Run `ft bench bw` once per machine to calibrate the split.
+  CPU, overlapped. Run `ft bench bw --model <checkpoint>` once per machine and
+  model to calibrate the split.
 - **auto** — dense models always resolve to `fused`; MoE models resolve to
-  `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile
-  recommends it. Unified-memory GPUs differ, see below.
+  `offload`. Unified-memory GPUs differ, see below.
 
 ### Unified-memory GPUs (GB10 / DGX Spark)
 
@@ -107,7 +107,7 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   Bounded output differs from exact by construction and does not depend on the prefill chunk size.
   Under expert offload the prefill time is bounded by streaming each layer's experts, so decoder
   replay saves decoder-layer compute, not prefill time.
-  Run `ft bench bw --model dsv4.1-flash` to measure local PCIe and CPU bandwidth before
+  Run `ft bench bw --model <checkpoint>` to measure local PCIe and CPU bandwidth before
   choosing a MoE strategy; `--moe-strategy hybrid` splits expert misses using those
   measurements. With `--moe-cache-auto`, `--kv-reserve-tokens` sets the token capacity reserved
   before the expert cache is allocated, subject to the attention pool's structural minimum; size it
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index ee6558014d..b93de55c2f 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import collections
 import contextlib
 import errno
 import gc
@@ -853,8 +854,6 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         cache.cpu_layer_ids = cpu_layer_ids
         cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags, row_tail_bytes=banks.row_tail_bytes)
         cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
-        if decode_target == "hybrid":
-            self._resolve_hybrid_fetch(config, cache)
         # Must be set before CUDA graph capture so the (device-side) accumulation ops are
         # captured and re-run on every decode replay.
         cache.collect_stats = config.moe_collect_stats
@@ -862,6 +861,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         assert len(layers) == config.model_config.num_moe_layers
         if cache.decode_target in ("cpu", "hybrid"):
             self._init_cpu_moe_executor(config, cache, layers)
+        if decode_target == "hybrid":
+            self._resolve_hybrid_fetch(config, cache)
         self.ctx.moe_offload_cache = cache
         self.moe_offload_cache = cache
         return cache
@@ -882,30 +883,50 @@ def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None:
         Perfect fetch/compute overlap wants fetched : cpu-computed misses = pcie_bw :
         (cpu_bw - pcie_bw), i.e. fetching a pcie_bw / cpu_bw fraction of each decode
         step's misses -- both sides then finish together instead of one idling. The
-        achieved bandwidths come from the cached `ft bench bw` profile (the same one the
-        auto backend pick reads); without a usable profile the old fixed cap of 1 applies.
+        achieved bandwidths come from the cached `ft bench bw` profile, from the row of each
+        MoE layer's kind (format, options, expert shape, top_k, CPU worker threads); a layer
+        without a row keeps the fixed cap of 1.
         """
         if config.moe_hybrid_max_fetch >= 0:
             return  # explicit fixed cap
-        from freetoken.moe.bench_profile import load_hybrid_fetch_fraction
+        from freetoken.moe.bench_profile import fetch_fraction, layer_kind_key, load_layer_kinds
+        from freetoken.moe.cpu_executor import _format_label
 
+        cache.hybrid_max_fetch = 1
+        methods = offload_layer_methods(self.model)
+        if methods is None:
+            logger.warning_rank0("--moe-hybrid-max-fetch auto: the MoE layers bind no expert method to look up in the `ft bench bw` profile; using a fixed fetch cap of 1")
+            return
         gpu_name, gpu_uuid = _profile_gpu(self.device.index)
-        fraction = load_hybrid_fetch_fraction(
-            cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid
-        )
-        if fraction is None:
-            cache.hybrid_max_fetch = 1
+        rows = load_layer_kinds(gpu_name=gpu_name, gpu_uuid=gpu_uuid) or {}
+        executor = self.cpu_moe_executor
+        keys, labels = [], []
+        for m in methods:
+            fmt, opts = m.cpu_spec()
+            c = m.cfg
+            keys.append(layer_kind_key(fmt, opts, c.hidden, c.local_intermediate, c.top_k, executor.num_threads))
+            labels.append(f"{_format_label(fmt, tuple(sorted(opts.items())))} {c.hidden}x{c.local_intermediate}")
+        fractions = [fetch_fraction(rows.get(key)) for key in keys]
+        cache.hybrid_fetch_fraction = [f or 0.0 for f in fractions]
+        groups = collections.Counter(zip(labels, fractions))
+        fetched = [f"{f:.1%} on {n} {label} layers" for (label, f), n in groups.items() if f is not None]
+        if fetched:
+            logger.info_rank0(
+                "--moe-hybrid-max-fetch auto: each decode step fetches this share of a layer's expert misses over PCIe "
+                f"(benched PCIe/CPU bandwidth ratio) and computes the rest on the CPU: {'; '.join(fetched)}"
+            )
+        if missing := sorted({label for label, f in groups if f is None}):
+            capped = sum(n for (_, f), n in groups.items() if f is None)
             logger.warning_rank0(
-                "--moe-hybrid-max-fetch auto: no usable `ft bench bw` profile for "
-                f"{cache.quant_format!r} experts; using a fixed fetch cap of 1"
+                f"--moe-hybrid-max-fetch auto: no `ft bench bw` row for {capped} MoE layers ({', '.join(missing)}) at "
+                f"{executor.num_threads} CPU threads; using a fixed fetch cap of 1 there. `ft bench bw --model <checkpoint>` benches each layer kind"
+            )
+        benched_isas = {rows[key].get("cpu_moe_isa") for key, f in zip(keys, fractions) if f is not None} - {None}
+        if other := sorted(benched_isas - set(executor.isa.split(","))):
+            logger.warning_rank0(
+                f"--moe-hybrid-max-fetch auto: some `ft bench bw` rows were benched on CPU kernels {other}, but this executor runs "
+                f"{executor.isa}; using them anyway"
             )
-            return
-        cache.hybrid_max_fetch = cache.num_experts  # inert: the fraction is the cap
-        cache.hybrid_fetch_fraction = fraction
-        logger.info_rank0(
-            f"--moe-hybrid-max-fetch auto: fetching {fraction:.1%} of each decode step's "
-            "expert misses over PCIe (benched PCIe/CPU bandwidth ratio), the rest on the CPU"
-        )
 
     def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         """Build the persistent CPU MoE executor (decode-time expert compute).
@@ -1643,24 +1664,34 @@ def _check_layer_formats(banks, layer_methods) -> None:
         )
 
 
+def _meta_model(config: EngineConfig, decode_target: str):
+    from freetoken.models import create_model
+    from freetoken.utils.torch_utils import torch_dtype
+
+    set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend)))
+    object.__setattr__(config.model_config, "moe_strategy", config.moe_strategy)
+    object.__setattr__(config.model_config, "decode_target", decode_target)
+    with torch.device("meta"), torch_dtype(config.dtype):
+        return create_model(config.model_config)
+
+
 def offload_expert_methods(config: EngineConfig) -> tuple:
     """``(method, layer_methods)`` of ``config``'s model, from a meta-device build: the most common expert method and,
     for experts of several formats, one per MoE layer (else None).
 
     For tools that load expert banks without an engine (the FTW converter); they pack for GPU decode."""
-    from freetoken.models import create_model
     from freetoken.moe.expert_banks import is_mixed
-    from freetoken.utils.torch_utils import torch_dtype
 
-    set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend)))
-    object.__setattr__(config.model_config, "moe_strategy", config.moe_strategy)
-    object.__setattr__(config.model_config, "decode_target", "gpu")
-    with torch.device("meta"), torch_dtype(config.dtype):
-        model = create_model(config.model_config)
+    model = _meta_model(config, "gpu")
     layer_methods = offload_layer_methods(model)
     return shared_offload_method(model), (layer_methods if is_mixed(layer_methods) else None)
 
 
+def meta_offload_layer_methods(config: EngineConfig, decode_target: str) -> list | None:
+    """The expert method of each offload MoE layer of ``config``'s model, from a meta-device build for ``decode_target``; see ``offload_layer_methods``."""
+    return offload_layer_methods(_meta_model(config, decode_target))
+
+
 def offload_expert_method(config: EngineConfig):
     """The (most common) offload expert method of ``config``'s model; see ``offload_expert_methods``."""
     return offload_expert_methods(config)[0]
@@ -1789,7 +1820,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     # swigluoai (csrc ActKind; "gpt_oss_swiglu" rides inside the mxfp4 kernel and
     # swigluoai the generic GEMV epilogue). A model with any other expert
     # activation cannot decode on the CPU: reject an explicit cpu/hybrid pick at
-    # config time, and keep auto from upgrading offload -> hybrid off the profile.
+    # config time.
     # hidden_act (the dense activation) stands proxy for the expert activation --
     # true for every in-tree model. mxfp4 experts pass regardless: their act runs
     # inside the mxfp4 kernel, not the generic epilogue.
@@ -1825,8 +1856,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         # GB10 this added a measured ~130 s stall to every request, #369). Resident
         # experts are the safe default here, not the risky one: the model and its banks
         # page from the same pool, so the "wrong guess = weight-load OOM" rationale above
-        # does not apply. The benchbw hybrid upgrade is skipped too: CPU execution adds
-        # no bandwidth when both sides share one memory. Only formats the resident path
+        # does not apply. Only formats the resident path
         # can actually hold take this branch; the rest stay on offload as before.
         unified_memory = _is_unified_memory_gpu()
         if unified_memory and _fused_resident_ok(model_config):
@@ -1847,50 +1877,6 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
                     "(use --moe-strategy offload to keep the offload flags)"
                 )
                 override("moe_cpu_layers", None)
-        # Hardware-adaptive config: a cached `ft bench bw` profile can upgrade
-        # the offload default to hybrid when this machine's CPU MoE bandwidth clears its PCIe
-        # gather bandwidth by the bench threshold (default 2x). hybrid is VRAM-equivalent to
-        # offload -- same auto-sized GPU slot cache (_resolve_auto_moe_cache_size), plus a
-        # host-RAM CPU executor -- so this never raises the OOM risk; with no profile (or one
-        # from different hardware) it stays offload. offload remains the always-safe fallback.
-        # Key the lookup on the real expert format: mxfp4 lives in moe_weight_format when
-        # expert_quant is "none", and "none" with no weight format means plain bf16 experts.
-        moe_wfmt = getattr(model_config, "moe_weight_format", None)
-        bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
-        from freetoken.moe.bench_profile import load_backend_recommendation
-
-        gpu_name, gpu_uuid = _profile_gpu()
-        if (
-            default_backend == "offload"
-            and not unified_memory
-            and load_backend_recommendation(
-                bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid
-            )
-            == "hybrid"
-        ):
-            from freetoken.moe.cpu_executor import compiled_extension_supports
-
-            _act = getattr(model_config, "hidden_act", "silu")
-            if not _cpu_moe_act_ok:
-                logger.info_rank0(
-                    f"benchbw profile recommends hybrid, but the CPU MoE executor does not "
-                    f"support this model's expert activation "
-                    f"{getattr(model_config, 'hidden_act', None)!r}; staying on offload"
-                )
-            elif moe_wfmt != "mxfp4" and not compiled_extension_supports(_act):
-                # Stale prebuilt _cpu_moe.so: an explicit cpu/hybrid pick still
-                # hard-fails in the executor, but a default must not turn into a
-                # post-load crash -- degrade to offload.
-                logger.info_rank0(
-                    f"benchbw profile recommends hybrid, but the compiled _cpu_moe "
-                    f"extension predates activation {_act!r} (rebuild with "
-                    f"`python setup.py build_ext --inplace`); staying on offload"
-                )
-            else:
-                default_backend = "hybrid"
-                logger.info_rank0(
-                    f"benchbw profile recommends hybrid for {bench_fmt!r} experts on this GPU"
-                )
         override("moe_strategy", default_backend)
         logger.info_rank0(f"Auto-selected MoE strategy: {config.moe_strategy}")
 
diff --git a/python/freetoken/moe/bench_profile.py b/python/freetoken/moe/bench_profile.py
index 5fa38b6f9d..c939a947f2 100644
--- a/python/freetoken/moe/bench_profile.py
+++ b/python/freetoken/moe/bench_profile.py
@@ -1,12 +1,8 @@
 """Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw/<gpu-uuid>.json``).
 
-The engine consults this at MoE-backend *auto* resolution (``engine.py``) to make the
-offload-vs-hybrid choice hardware-adaptive without importing the (torch-heavy) benchmark
-itself. ``benchbw.py`` writes the profile; this module only reads it.
-
-The join key is the expert *format*: the CPU-MoE-vs-PCIe-gather bandwidth ratio the choice
-rides on is dominated by ``(format, hardware)``, not by the exact model, so a profile benched
-on one workload transfers to any model with the same expert format on the same GPU.
+The profile holds one row per MoE layer kind (``layer_kind_key``): the CPU MoE and PCIe gather
+bandwidths benched for that kind on this GPU and host. The engine reads it to set each hybrid
+layer's fetch split, without importing the (torch-heavy) benchmark. ``benchbw.py`` writes it.
 """
 
 from __future__ import annotations
@@ -18,20 +14,8 @@
 
 logger = init_logger(__name__)
 
-# Bump when a format gains or loses a CPU MoE path: verdicts and fetch fractions from an older profile are stale.
-PROFILE_VERSION = 5
-
-# Engine ``expert_quant`` (models/config.py) -> benchbw format key (offload_cache._BANK_SCHEMAS
-# / benchbw._offload_bank_specs). Only the offload-family formats with a CPU MoE weight path can
-# ever resolve to hybrid; anything not listed falls through unmapped and finds no profile entry
-# (-> None -> offload), which is the safe default.
-_QUANT_TO_BENCH_FORMAT = {
-    "nvfp4": "nvfp4",
-    "ds_fp4": "ds_fp4",
-    "mxfp4": "mxfp4_triton",
-    "bf16": "bf16",
-    "fp8_block": "fp8_block",
-}
+# Bump when the row key or the measurements change: rows of another version are not comparable.
+PROFILE_VERSION = 6
 
 
 def _cache_dir() -> str:
@@ -104,10 +88,10 @@ def _usable_profile(
             return None
     if not isinstance(prof, dict):
         return None
-    if not isinstance(prof.get("version"), int) or prof["version"] < PROFILE_VERSION:
+    if prof.get("version") != PROFILE_VERSION:
         logger.warning(
-            f"benchbw profile {src} predates this engine (version {prof.get('version')}, "
-            f"need {PROFILE_VERSION}); ignoring it, rerun `ft bench bw`"
+            f"benchbw profile {src} has version {prof.get('version')}, this engine reads "
+            f"{PROFILE_VERSION}; ignoring it, rerun `ft bench bw`"
         )
         return None
     prof_gpu = (prof.get("gpu") or {}).get("name")
@@ -120,81 +104,34 @@ def _usable_profile(
     return prof
 
 
-def load_backend_recommendation(
-    quant_format: str,
-    gpu_name: str | None = None,
-    path: str | None = None,
-    gpu_uuid: str | None = None,
-) -> str | None:
-    """Bench-recommended offload-family backend for ``quant_format`` on this GPU, or ``None``.
+def layer_kind_key(fmt: str, options: dict[str, int], hidden: int, inter: int, top_k: int, cpu_threads: int) -> str:
+    """Row key of one MoE layer kind, e.g. ``gguf[down_type=8,gate_up_type=8]:2816x704:top8:t12``.
 
-    Returns ``"hybrid"`` only when *every* benched workload sharing this expert format
-    recommended hybrid (CPU MoE BW > threshold x PCIe gather BW); a mixed verdict (a
-    near-threshold format) resolves conservatively to ``"offload"``. ``None`` means "no usable
-    profile" (see ``_usable_profile``) or no entry for this format. The caller keeps its own
-    default (offload) on ``None``.
-    """
-    fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
-    prof = _usable_profile(gpu_name, path, gpu_uuid)
-    if prof is None:
-        return None
+    ``cpu_threads`` is the CPU executor's worker count: the per-step pool overhead makes the CPU bandwidth depend on it and on ``top_k``."""
+    opts = ",".join(f"{k}={v}" for k, v in sorted(options.items()))
+    fmt = f"{fmt}[{opts}]" if opts else fmt
+    return f"{fmt}:{hidden}x{inter}:top{top_k}:t{cpu_threads}"
 
-    # Preferred: the per-dtype tuning verdicts (`ft bench bw --dtype`), a direct format->backend
-    # map -- the axis the backend pick is meant to key on.
-    dtypes = prof.get("dtypes")
-    if isinstance(dtypes, dict) and dtypes.get(fmt) in ("hybrid", "offload"):
-        return dtypes[fmt]
 
-    # Fallback: a per-model profile (`ft bench bw --model`). Aggregate the workloads sharing this
-    # format -- unanimous hybrid -> hybrid; any offload (a near-threshold split) -> offload.
-    workloads = prof.get("workloads")
-    if not isinstance(workloads, dict):
-        return None
-    picks = [
-        entry["recommended"]
-        for wl in workloads.values()
-        if isinstance(wl, dict)
-        for entry in [(wl.get("kernels") or {}).get(fmt)]
-        if isinstance(entry, dict) and entry.get("recommended")
-    ]
-    if not picks:
+def fetch_fraction(row: dict | None) -> float | None:
+    """The hybrid fetch split of a row: of a decode step's expert misses, the fraction to fetch over PCIe so the fetch and the CPU compute finish together.
+
+    From the overlapped pair (both sides measured while running concurrently): pcie_ov / (pcie_ov + cpu_ov); without it, pcie / cpu. Clamped to 1."""
+    if not isinstance(row, dict):
         return None
-    return "hybrid" if all(p == "hybrid" for p in picks) else "offload"
-
-
-def load_hybrid_fetch_fraction(
-    quant_format: str,
-    gpu_name: str | None = None,
-    path: str | None = None,
-    gpu_uuid: str | None = None,
-) -> float | None:
-    """Benched hybrid fetch fraction for ``quant_format``, or ``None``.
-
-    The hybrid backend's bandwidth-matched fetch split: of a decode step's expert misses,
-    fetch this fraction over PCIe and compute the rest on the CPU, so both finish together.
-    Preferred source is the *overlapped* pair (CPU MoE and PCIe gather measured while
-    running concurrently -- the real contention regime): fetched/misses = pcie_ov /
-    (pcie_ov + cpu_ov). Older profiles without it fall back to the standalone bandwidths
-    under a full-DRAM-contention assumption (cpu keeps cpu - pcie under DMA), which
-    reduces to pcie/cpu. Per-dtype entry first, then any per-model entry with this format.
-    ``None`` = no usable profile; clamped to [0, 1].
-    """
-    fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
+    cpu_ov, pcie_ov = row.get("cpu_moe_overlap_gbs"), row.get("pcie_gather_overlap_gbs")
+    if cpu_ov and pcie_ov:
+        return min(1.0, pcie_ov / (pcie_ov + cpu_ov))
+    cpu, pcie = row.get("cpu_moe_gbs"), row.get("pcie_gather_gbs")
+    if cpu and pcie:
+        return min(1.0, pcie / cpu)
+    return None
+
+
+def load_layer_kinds(gpu_name: str | None = None, path: str | None = None, gpu_uuid: str | None = None) -> dict | None:
+    """The rows of the usable profile for this GPU (key -> row), or None without one."""
     prof = _usable_profile(gpu_name, path, gpu_uuid)
     if prof is None:
         return None
-    entries = [(prof.get("dtype_kernels") or {}).get(fmt)] + [
-        (wl.get("kernels") or {}).get(fmt)
-        for wl in (prof.get("workloads") or {}).values()
-        if isinstance(wl, dict)
-    ]
-    for entry in entries:
-        if not isinstance(entry, dict):
-            continue
-        cpu_ov, pcie_ov = entry.get("cpu_moe_overlap_gbs"), entry.get("pcie_gather_overlap_gbs")
-        if cpu_ov and pcie_ov:
-            return min(1.0, pcie_ov / (pcie_ov + cpu_ov))
-        cpu, pcie = entry.get("cpu_moe_gbs"), entry.get("pcie_gather_gbs")
-        if cpu and pcie:
-            return min(1.0, pcie / cpu)
-    return None
+    rows = prof.get("layer_kinds")
+    return rows if isinstance(rows, dict) else {}
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index f7ea5aa307..057b86fda8 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -11,23 +11,24 @@
     PCIe on a decode miss). These are what the hybrid-vs-offload choice rides on.
     When both work, the same pair is also measured *overlapped* (running concurrently,
     contending for host DRAM) -- those two achieved bandwidths set the hybrid backend's
-    per-step fetch split (``load_hybrid_fetch_fraction``).
+    per-step fetch split (``bench_profile.fetch_fraction``).
 
 The rule, per (workload, format): recommend ``hybrid`` when the CPU MoE kernel
 bandwidth exceeds ``threshold`` x the PCIe gather bandwidth (default 2x), otherwise
 ``offload``. Formats with no CPU MoE weight path (e.g. block-fp8) can't be computed on
 the CPU, so they always resolve to ``offload``.
 
-The hybrid-vs-offload choice is dtype-dominated, so the default is a **per-dtype tuning bench**
-(``--dtype``): one bench per expert format against a canonical geometry -- the minimal set the
-runtime backend pick matches on. ``--model`` additionally benches specific model geometries for
-per-model detail. Results are written to a JSON file per GPU (default ``$XDG_CACHE_HOME/
-freetoken/benchbw/<gpu-uuid>.json``) so the choice is reproducible and a multi-GPU box keeps
-one profile per card.
+Each bench writes one row per MoE layer kind: format, options, expert shape, top_k and CPU
+worker threads (``bench_profile.layer_kind_key``). ``--model <checkpoint>`` benches every kind
+a real checkpoint has, which is what the hybrid fetch split matches per layer; ``--dtype`` and
+the ``--model`` presets bench canonical geometries. The rows live in a JSON file per GPU
+(default ``$XDG_CACHE_HOME/freetoken/benchbw/<gpu-uuid>.json``); a run replaces the rows it
+benched and keeps the others.
 
     ft bench bw                              # per-dtype tuning bench (default: all formats)
     ft bench bw --dtype nvfp4,bf16           # only these formats
     ft bench bw --model qwen3.6-moe          # per-model detail instead
+    ft bench bw --model /path/to/model.gguf  # every MoE layer kind of a checkpoint
     ft bench bw --dtype all --model all      # both
     ft bench bw --gpu GPU-2f3a...            # bench a specific GPU (UUID or nvidia-smi index)
 """
@@ -38,9 +39,11 @@
 import contextlib
 import json
 import logging
+import math
 import os
 import socket
 import statistics
+import subprocess
 import threading
 import time
 from dataclasses import dataclass
@@ -64,7 +67,7 @@
 
 # Formats the CPU MoE C++ kernel can compute AND this bench can build banks for; anything
 # else is offload-only here.
-_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4", "fp8_block"})
+_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4", "fp8_block", "gguf", "gguf_scaled"})
 # Formats this bench can build synthetic (correctly-sized) banks for.
 _BUILDABLE_FORMATS = frozenset({"bf16", "nvfp4", "fp8_block", "mxfp4_triton", "ds_fp4"})
 # Friendlier CLI/display aliases for the internal quant_format strings.
@@ -110,6 +113,13 @@ class Workload:
     swiglu_alpha: float = 1.702
     swiglu_limit: float | None = None
     act_block: int | None = None  # ds_fp4 W4A8 activation quant block
+    options: tuple[tuple[str, int], ...] = ()  # CPU executor options of the format, e.g. the GGUF types
+
+    def format_options(self, fmt: str) -> dict[str, int]:
+        opts = dict(self.options)
+        if self.act_block is not None and fmt == "ds_fp4":
+            opts.setdefault("act_block", self.act_block)
+        return opts
 
 
 # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py.
@@ -132,13 +142,11 @@ class Workload:
     "minimax-m2.5": Workload("minimax-m2.5", 3072, 1536, 256, 8, ("nvfp4",)),
 }
 
-# Per-dtype canonical geometries for the TUNING bench (`--dtype`). The hybrid-vs-offload choice is
-# dtype-dominated -- the CPU-MoE-vs-PCIe-gather bandwidth ratio rides on the expert *format* and
-# the hardware, not the exact model -- so measuring one representative geometry per format is
-# enough for the runtime backend pick (the engine also matches by dtype). This is the minimal set
-# the Desktop "detect" flow runs; `--model` still benches specific model geometries for per-model
-# detail. Geometries mirror the model families that ship each format (mxfp4 keeps gpt-oss's large
-# experts -- the reason that format sits near the 2x threshold).
+# Per-dtype canonical geometries for the TUNING bench (`--dtype`): one representative geometry per
+# expert format, the set the Desktop "detect" flow runs. A model's layers match these rows only when
+# their kind is the same; `--model <checkpoint>` benches a model's own kinds. Geometries mirror the
+# model families that ship each format (mxfp4 keeps gpt-oss's large experts -- the reason that
+# format sits near the 2x threshold).
 DTYPE_WORKLOADS: dict[str, Workload] = {
     "bf16": Workload("dtype:bf16", 2048, 768, 128, 8, ("bf16",)),
     "nvfp4": Workload("dtype:nvfp4", 3072, 1536, 128, 8, ("nvfp4",)),
@@ -285,9 +293,25 @@ def _timed(copy) -> float:
 # ============================== bank geometry ==============================
 
 
-def _offload_bank_specs(fmt: str, H: int, I: int) -> dict[str, tuple[int, torch.dtype]]:
+def _gguf_bank_shapes(fmt: str, opts: dict, H: int, I: int) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
+    """Per-expert (row shape, dtype) of the GGUF banks: ggml rows as bytes, float rows in their dtype."""
+    from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes
+
+    def rows(ggml_type: int, n: int, k: int):
+        dtype = float_dtype(ggml_type)
+        return ((n, k), dtype) if dtype is not None else ((n, row_bytes(ggml_type, k)), torch.uint8)
+
+    shapes = {"gate_up": rows(opts["gate_up_type"], 2 * I, H), "down": rows(opts["down_type"], H, I)}
+    if fmt == "gguf_scaled":
+        shapes["global"] = ((4,), torch.float32)
+    return shapes
+
+
+def _offload_bank_specs(fmt: str, H: int, I: int, opts: dict | None = None) -> dict[str, tuple[int, torch.dtype]]:
     """Per-bank (elems_per_expert, dtype) for the flat offload host banks (the gather)."""
     u8, bf16, f16, f8 = torch.uint8, torch.bfloat16, torch.float16, torch.float8_e4m3fn
+    if fmt in ("gguf", "gguf_scaled"):
+        return {name: (math.prod(shape), dtype) for name, (shape, dtype) in _gguf_bank_shapes(fmt, opts, H, I).items()}
     if fmt == "bf16":
         return {"gate_up": (2 * I * H, bf16), "down": (H * I, bf16)}
     if fmt == "fp8_block":
@@ -315,8 +339,8 @@ def _offload_bank_specs(fmt: str, H: int, I: int) -> dict[str, tuple[int, torch.
     raise NotImplementedError(fmt)
 
 
-def _expert_bytes(fmt: str, H: int, I: int) -> int:
-    return sum(elems * dt.itemsize for elems, dt in _offload_bank_specs(fmt, H, I).values())
+def _expert_bytes(fmt: str, H: int, I: int, opts: dict | None = None) -> int:
+    return sum(elems * dt.itemsize for elems, dt in _offload_bank_specs(fmt, H, I, opts).values())
 
 
 def _synth_experts(E: int, expert_bytes: int) -> int:
@@ -327,7 +351,7 @@ def _synth_experts(E: int, expert_bytes: int) -> int:
     return min(E, max(1, _SYNTH_BANK_BUDGET // max(1, expert_bytes)))
 
 
-def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int) -> dict:
+def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int, opts: dict | None = None) -> dict:
     """3D pinned banks in the exact layout ``CpuMoeExecutor`` expects for ``fmt``.
 
     Scale banks that decode through an e8m0 exponent (mxfp4/ds_fp4) are set to a unit
@@ -338,6 +362,22 @@ def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int) -> dict:
     def pin(*shape, dtype):
         return alloc_pinned_tensor(*shape, dtype=dtype)
 
+    if fmt in ("gguf", "gguf_scaled"):
+        from freetoken.third_party.gguf import GGML_QUANT_SIZES, GGMLQuantizationType
+
+        b = {}
+        for name, (shape, dtype) in _gguf_bank_shapes(fmt, opts, H, I).items():
+            b[name] = pin(E, *shape, dtype=dtype)
+            if name == "global":
+                b[name].fill_(1.0)
+            elif dtype == torch.uint8:
+                b[name].zero_()  # zero ggml blocks: every fp16 block scale is 0, never denormal
+                if opts["gate_up_type" if name == "gate_up" else "down_type"] == GGMLQuantizationType.MXFP4:
+                    # each block starts with an e8m0 scale byte; 0 decodes to 2^-127, a denormal that slows the CPU kernel
+                    b[name].view(E, -1, GGML_QUANT_SIZES[GGMLQuantizationType.MXFP4][1])[..., 0] = 127
+            else:
+                b[name].fill_(0.02)
+        return b
     if fmt == "bf16":
         gate_up, down = pin(E, 2 * I, H, dtype=torch.bfloat16), pin(E, H, I, dtype=torch.bfloat16)
         gate_up.fill_(0.02)  # uninitialized bf16 can be denormal -> x86 FP slowdown
@@ -408,9 +448,9 @@ def _build_gather_rig(fmt: str, wl: Workload, device: torch.device):
     bench and the overlapped CPU+PCIe bench."""
     from freetoken.moe.offload_cache import OffloadMoeCache
 
-    H, I = wl.hidden, wl.inter
-    E = _synth_experts(wl.experts, _expert_bytes(fmt, H, I))
-    specs = _offload_bank_specs(fmt, H, I)
+    H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt)
+    E = _synth_experts(wl.experts, _expert_bytes(fmt, H, I, opts))
+    specs = _offload_bank_specs(fmt, H, I, opts)
     cache = OffloadMoeCache(num_layers=1, num_experts=E, cache_size=E, device=device, quant_format=fmt)
     total_bytes = 0
     sources = {}
@@ -437,7 +477,7 @@ def measure_pcie_gather_bw(fmt: str, wl: Workload, device: torch.device, iters:
     (all synth experts, scattered sources). Reusing the cache gives the kernel-correct
     int64 ``num_indices``. Timed with CUDA events.
     """
-    eb = _expert_bytes(fmt, wl.hidden, wl.inter)
+    eb = _expert_bytes(fmt, wl.hidden, wl.inter, wl.format_options(fmt))
     cache, total_bytes = _build_gather_rig(fmt, wl, device)
     E = cache.num_experts
 
@@ -476,7 +516,7 @@ def _build_cpu_moe_executor(fmt: str, wl: Workload, banks: dict, num_threads: in
         cache, top_k=wl.top_k, activation=wl.activation,
         apply_router_weight_on_input=False, num_threads=num_threads, max_tokens=1,
         device=torch.device("cuda"), swiglu_alpha=wl.swiglu_alpha,
-        swiglu_limit=wl.swiglu_limit, layer_formats=[(fmt, {"act_block": wl.act_block} if wl.act_block is not None else {})] * cache.num_layers,
+        swiglu_limit=wl.swiglu_limit, layer_formats=[(fmt, wl.format_options(fmt))] * cache.num_layers,
     )
 
 
@@ -528,25 +568,25 @@ def measure_cpu_moe_bw(fmt: str, wl: Workload, iters: int = 64, num_threads: int
     the SAME banks (built once) to populate the diagnostic ``isa_sweep`` breakdown; tiers
     above the CPU's support clamp down and dedupe by the resulting isa name.
     """
-    H, I = wl.hidden, wl.inter
-    eb = _expert_bytes(fmt, H, I)
+    H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt)
+    eb = _expert_bytes(fmt, H, I, opts)
     E = _synth_experts(wl.experts, eb)
-    banks = _cpu_moe_bank_sources(fmt, H, I, E)
+    banks = _cpu_moe_bank_sources(fmt, H, I, E, opts)
 
-    def run(isa: str | None) -> tuple[str, float]:
+    def run(isa: str | None) -> tuple[str, float, list[int]]:
         with _forced_isa(isa):
             ex = _build_cpu_moe_executor(fmt, wl, banks, num_threads, E)
         bw = round(_time_cpu_moe(ex, wl, eb, E, iters), 2)
-        label = ex.isa
+        label, cores = ex.isa, list(ex.core_ids)
         del ex
-        return label, bw
+        return label, bw, cores
 
-    auto_isa, auto_bw = run(None)  # recommendation basis: the tier production would pick
+    auto_isa, auto_bw, cores = run(None)  # recommendation basis: the tier production would pick
     sweep = {auto_isa: auto_bw}
     for isa in isas or []:
-        label, bw = run(isa)
+        label, bw, _ = run(isa)
         sweep[label] = max(sweep.get(label, 0.0), bw)  # dedupe clamped tiers, keep best
-    return {"bw_gbs": auto_bw, "isa": auto_isa, "isa_sweep": sweep,
+    return {"bw_gbs": auto_bw, "isa": auto_isa, "isa_sweep": sweep, "cores": cores,
             "expert_bytes": eb, "synth_experts": E}
 
 
@@ -567,10 +607,10 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device,
     so the DMA is really in flight, not just enqueued). Each side reports bytes / its own
     elapsed; the two windows differ by at most one CPU step + one gather.
     """
-    H, I = wl.hidden, wl.inter
-    eb = _expert_bytes(fmt, H, I)
+    H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt)
+    eb = _expert_bytes(fmt, H, I, opts)
     E = _synth_experts(wl.experts, eb)
-    ex = _build_cpu_moe_executor(fmt, wl, _cpu_moe_bank_sources(fmt, H, I, E), num_threads, E)
+    ex = _build_cpu_moe_executor(fmt, wl, _cpu_moe_bank_sources(fmt, H, I, E, opts), num_threads, E)
     set_ids, run_step = _cpu_moe_step_fns(ex, wl, E)
     cache, gather_bytes = _build_gather_rig(fmt, wl, device)
 
@@ -631,17 +671,15 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
     one format can't kill the run; an unexpected exception is left to propagate.
     """
     entry: dict = {
-        "expert_bytes": None, "synth_experts": None, "cpu_moe_gbs": None, "cpu_moe_isa": None,
-        "isa_sweep": None, "pcie_gather_gbs": None,
+        "cpu_moe_gbs": None, "cpu_moe_isa": None,
+        "cpu_cores": None, "isa_sweep": None, "pcie_gather_gbs": None,
         "cpu_moe_overlap_gbs": None, "pcie_gather_overlap_gbs": None,
-        "ratio": None, "recommended": None,
+        "recommended": None,
         "note": None,
     }
     try:
         g = measure_pcie_gather_bw(fmt, wl, device, pcie_iters)
         entry["pcie_gather_gbs"] = round(g["bw_gbs"], 2)
-        entry["expert_bytes"] = g["expert_bytes"]
-        entry["synth_experts"] = g["synth_experts"]
     except (ImportError, RuntimeError) as e:
         logger.warning(f"benchbw: PCIe gather bench failed for {wl.name}/{fmt}: {e}")
         _note(entry, f"pcie gather unavailable ({e})")
@@ -655,6 +693,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
             c = measure_cpu_moe_bw(fmt, wl, cpu_iters, cpu_threads, isas=isas)
             entry["cpu_moe_gbs"] = c["bw_gbs"]  # auto/best tier (what production runs)
             entry["cpu_moe_isa"] = c["isa"]
+            entry["cpu_cores"] = c["cores"]
             if isas:  # a sweep was requested -> keep the per-tier breakdown
                 entry["isa_sweep"] = c["isa_sweep"]
                 # The tier label comes from the bf16 dot chooser, so it can over-report tiers
@@ -665,18 +704,18 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
                 elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"):
                     _note(entry, "isa tier labels are nominal: avx512bf16 == avx512f here "
                                  "(3 real kernels: scalar/avx2/avx512, no bf16/VNNI variant)")
-            entry["expert_bytes"] = entry["expert_bytes"] or c["expert_bytes"]
-            entry["synth_experts"] = entry["synth_experts"] or c["synth_experts"]
+                elif fmt in ("gguf", "gguf_scaled"):
+                    _note(entry, "isa tier labels are nominal: GGUF rows run on the ggml-cpu variant "
+                                 "picked at load, regardless of the forced tier")
         except (ImportError, RuntimeError) as e:
             logger.warning(f"benchbw: CPU MoE bench failed for {wl.name}/{fmt}: {e}")
             _note(entry, f"cpu moe unavailable ({e})")
 
     cpu_g, pcie_g = entry["cpu_moe_gbs"], entry["pcie_gather_gbs"]
     if cpu_g is not None and pcie_g:  # pcie_g truthy also rules out a div-by-zero
-        entry["ratio"] = round(cpu_g / pcie_g, 3)
         entry["recommended"] = recommend(cpu_g, pcie_g, threshold)
         # Both sides work standalone -> also measure them contended (concurrently). This
-        # pair sets the hybrid backend's fetch split (load_hybrid_fetch_fraction); the
+        # pair sets the hybrid backend's fetch split (bench_profile.fetch_fraction); the
         # hybrid-vs-offload verdict above stays on the standalone numbers.
         try:
             o = measure_overlap_bw(fmt, wl, device, cpu_threads)
@@ -694,6 +733,89 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float
     return entry
 
 
+def _kind(fmt: str, wl: Workload) -> tuple:
+    """A layer kind without the CPU thread count, which the row key adds once the executor has sized its pool."""
+    return (fmt, tuple(sorted(wl.format_options(fmt).items())), wl.hidden, wl.inter, wl.top_k)
+
+
+def checkpoint_layer_kinds(model_path: str, device: torch.device) -> dict[tuple, Workload]:
+    """Each distinct MoE layer kind of a checkpoint as hybrid decode would serve it: kind -> workload."""
+    from freetoken.distributed import DistributedInfo, set_tp_info, try_get_tp_info
+    from freetoken.engine.config import EngineConfig
+    from freetoken.engine.engine import meta_offload_layer_methods
+    from freetoken.layers import set_rope_device
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    tp = try_get_tp_info()
+    cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=torch.bfloat16, moe_strategy="hybrid")
+    if not getattr(cfg.model_config, "is_moe", False):
+        raise RuntimeError(f"{model_path} has no MoE layers")
+    set_rope_device(device)
+    methods = meta_offload_layer_methods(cfg, "hybrid")
+    if methods is None:
+        raise RuntimeError(f"{model_path}: its MoE layers bind no expert method")
+    kinds: dict[tuple, Workload] = {}
+    for m in methods:
+        fmt, opts = m.cpu_spec()
+        c = m.cfg
+        wl = Workload(
+            f"{os.path.basename(model_path.rstrip('/'))}:{fmt}", c.hidden, c.local_intermediate, c.num_experts, c.top_k, (fmt,),
+            activation=c.activation, swiglu_alpha=float(c.alpha), swiglu_limit=c.limit, options=tuple(sorted(opts.items())),
+        )
+        kinds.setdefault(_kind(fmt, wl), wl)
+    return kinds
+
+
+def _freetoken_build() -> dict:
+    """FreeToken version and, in a git checkout, its commit: rows benched on older kernels may be stale."""
+    from freetoken.version import __version__
+
+    try:
+        commit = subprocess.run(
+            ["git", "-C", os.path.dirname(os.path.abspath(__file__)), "rev-parse", "--short=12", "HEAD"],
+            capture_output=True, text=True, timeout=5, check=True,
+        ).stdout.strip() or None
+    except (OSError, subprocess.SubprocessError):
+        commit = None
+    return {"version": __version__, "commit": commit}
+
+
+def _cpu_info() -> dict:
+    model = numa = None
+    try:
+        with open("/proc/cpuinfo") as f:
+            model = next((line.split(":", 1)[1].strip() for line in f if line.startswith("model name")), None)
+    except OSError:
+        pass
+    try:
+        numa = sum(1 for n in os.listdir("/sys/devices/system/node") if n.startswith("node") and n[4:].isdigit()) or None
+    except OSError:
+        pass
+    return {"model": model, "physical_cores": len(physical_core_cpus()), "numa_nodes": numa}
+
+
+def _merge_previous(result: dict, path: str) -> None:
+    """Keep the rows of the profile already at ``path`` (same GPU and version) that this run did not re-bench; a re-benched row keeps its other sources."""
+    try:
+        with open(path) as f:
+            prev = json.load(f)
+    except (OSError, ValueError):
+        return
+    if not isinstance(prev, dict) or prev.get("version") != result["version"] or (prev.get("gpu") or {}).get("uuid") != result["gpu"]["uuid"]:
+        return
+    prev_rows = prev.get("layer_kinds")
+    if not isinstance(prev_rows, dict):
+        return
+    rows = dict(prev_rows)
+    for key, row in result["layer_kinds"].items():
+        old = prev_rows.get(key)
+        if isinstance(old, dict) and isinstance(old.get("sources"), list):
+            row["sources"] = sorted(set(old["sources"]) | set(row["sources"]))
+        rows[key] = row
+    result["layer_kinds"] = rows
+
+
 def _atomic_write_json(path: str, obj: dict) -> None:
     os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
     tmp = f"{path}.tmp.{os.getpid()}"
@@ -717,7 +839,11 @@ def run_benchbw(
     pcie_iters: int = 30,
     kernel_cpu_iters: int = 64,
     kernel_pcie_iters: int = 20,
+    checkpoints: tuple[str, ...] = (),
 ) -> dict:
+    from freetoken.moe.bench_profile import layer_kind_key
+    from freetoken.moe.cpu_executor import _format_label
+
     if not torch.cuda.is_available():
         raise RuntimeError(
             "benchbw needs a CUDA device to measure PCIe bandwidth (both offload and "
@@ -735,12 +861,29 @@ def run_benchbw(
     torch.cuda.set_device(device)
     gpu = gpu_identity(device_index)
 
+    # kind -> (workload, sources); checkpoints are read before any timing so a bad one fails fast
+    units: dict[tuple, tuple[Workload, list[str]]] = {}
+
+    def add(fmt: str, wl: Workload, source: str) -> None:
+        sources = units.setdefault(_kind(fmt, wl), (wl, []))[1]
+        if source not in sources:
+            sources.append(source)
+
+    for fmt in dtypes or ():
+        add(fmt, DTYPE_WORKLOADS[fmt], f"dtype:{_FORMAT_DISPLAY.get(fmt, fmt)}")
+    for mname in models:
+        for fmt in formats or WORKLOADS[mname].formats:
+            add(fmt, WORKLOADS[mname], f"preset:{mname}")
+    for path in checkpoints:
+        logger.info(f"benchbw: reading the MoE layer kinds of {path} ...")
+        for kind, wl in checkpoint_layer_kinds(path, device).items():
+            add(kind[0], wl, path)
+
     # Machine-parseable progress on stdout (opt-in), so the daemon/Desktop can stream feedback
     # while the bench runs -- mirrors ft checkpoint's FTCONVERT lines. `done`/`total` count the
-    # measured units (ceilings + one per benched format); `label` is what's measuring now.
+    # measured units (ceilings + one per layer kind); `label` is what's measuring now.
     prog_on = os.environ.get("FREETOKEN_BENCH_PROGRESS") == "1"
-    n_model_units = sum(len(formats or WORKLOADS[m].formats) for m in models)
-    total = 1 + len(dtypes or ()) + n_model_units
+    total = 1 + len(units)
     done = 0
 
     def _prog(label: str) -> None:
@@ -758,36 +901,24 @@ def _prog(label: str) -> None:
     exec_log = logging.getLogger("freetoken.moe.cpu_executor")
     prev_level = exec_log.level
     exec_log.setLevel(logging.WARNING)
-    dtypes_out: dict = {}
-    workloads_out: dict = {}
+    now = datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")
+    build = _freetoken_build()
+    rows: dict = {}
     try:
-        # dtype tuning: one bench per expert format (canonical geometry). This is what the
-        # runtime backend pick keys off -- matched by dtype in bench_profile.py.
-        for fmt in dtypes or ():
-            wl = DTYPE_WORKLOADS[fmt]
-            _prog(f"dtype:{_FORMAT_DISPLAY.get(fmt, fmt)}")
-            logger.info(f"benchbw: dtype {_FORMAT_DISPLAY.get(fmt, fmt)} real kernels ...")
-            dtypes_out[fmt] = _bench_format(
-                fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas
-            )
+        for (fmt, opts, hidden, inter, top_k), (wl, sources) in units.items():
+            label = f"{_format_label(fmt, opts)} {hidden}x{inter} top{top_k}"
+            _prog(label)
+            logger.info(f"benchbw: {label} real kernels ...")
+            entry = _bench_format(fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas)
             done += 1
-        # optional per-model detail (--model): specific geometries, unchanged.
-        for mname in models:
-            wl = WORKLOADS[mname]
-            fmts = formats or wl.formats
-            kernels = {}
-            for fmt in fmts:
-                _prog(f"{mname}:{_FORMAT_DISPLAY.get(fmt, fmt)}")
-                logger.info(f"benchbw: {mname}/{fmt} real kernels ...")
-                kernels[fmt] = _bench_format(
-                    fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas
-                )
-                done += 1
-            workloads_out[mname] = {
-                "model": {"name": wl.name, "hidden": wl.hidden, "inter": wl.inter,
-                          "experts": wl.experts, "top_k": wl.top_k},
-                "kernels": kernels,
-                "recommended_moe_backend": {f: kernels[f]["recommended"] for f in fmts},
+            if not entry["cpu_cores"]:
+                logger.warning(f"benchbw: no CPU MoE result for {label}, so no row: {entry['note']}")
+                continue
+            threads = len(entry["cpu_cores"])
+            rows[layer_kind_key(fmt, dict(opts), hidden, inter, top_k, threads)] = {
+                "format": fmt, "options": dict(opts), "hidden": hidden, "inter": inter, "top_k": top_k,
+                "cpu_threads": threads, "activation": wl.activation, "sources": sources,
+                "threshold": threshold, "measured_at": now, "freetoken": build, **entry,
             }
     finally:
         exec_log.setLevel(prev_level)
@@ -795,27 +926,27 @@ def _prog(label: str) -> None:
 
     result = {
         "version": PROFILE_VERSION,
-        "timestamp": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"),
+        "timestamp": now,
         "epoch": int(time.time()),
         "host": socket.gethostname(),
         # index is the CUDA ordinal the bench ran on; uuid keys the profile file
         "gpu": {"index": device_index, "name": gpu["name"], "uuid": gpu["uuid"]},
-        "cpu": {"physical_cores": len(physical_core_cpus()), "threads_used": cpu["threads"]},
-        "threshold": threshold,
+        "cpu": _cpu_info(),
         "ceilings": {
             "cpu_stream_read_gbs": round(cpu["bw_gbs"], 2),
+            "cpu_stream_threads": cpu["threads"],
             "pcie_linear_h2d_gbs": round(pcie["h2d_gbs"], 2),
             "pcie_linear_d2h_gbs": round(pcie["d2h_gbs"], 2),
         },
-        # Per-dtype verdicts (the tuning axis) + optional per-model detail.
-        "dtypes": {f: e["recommended"] for f, e in dtypes_out.items()},
-        "dtype_kernels": dtypes_out,
-        "workloads": workloads_out,
+        "layer_kinds": rows,
     }
 
     out_path = os.path.expanduser(out_path or default_out_path(gpu["uuid"]))
+    benched = list(rows)
+    _merge_previous(result, out_path)
     _atomic_write_json(out_path, result)
     result["out_path"] = out_path
+    result["benched"] = benched
     if prog_on:
         print(f"FTBENCH_OUT {out_path}", flush=True)
     return result
@@ -833,8 +964,9 @@ def _print_kernels(kernels: dict, iw: int) -> None:
           f"{'CPU/PCIe':>9}  backend")
     for fmt, e in kernels.items():
         disp = _FORMAT_DISPLAY.get(fmt, fmt)
-        eb = f"{e['expert_bytes'] / 2**20:.2f} MB" if e["expert_bytes"] else "n/a"
-        ratio = f"{e['ratio']:.2f}x" if e["ratio"] is not None else "—"
+        eb = f"{_expert_bytes(fmt, e['hidden'], e['inter'], e['options']) / 2**20:.2f} MB"
+        cpu_g, pcie_g = e["cpu_moe_gbs"], e["pcie_gather_gbs"]
+        ratio = f"{cpu_g / pcie_g:.2f}x" if cpu_g is not None and pcie_g else "—"
         print(f"    {disp:<8} {eb:>9} {_gbs(e['cpu_moe_gbs'])} {_gbs(e['pcie_gather_gbs'])} "
               f"{ratio:>9}  {e['recommended']}")
         c_ov, p_ov = e.get("cpu_moe_overlap_gbs"), e.get("pcie_gather_overlap_gbs")
@@ -851,27 +983,24 @@ def _print_kernels(kernels: dict, iw: int) -> None:
 
 
 def _print_report(r: dict) -> None:
-    c = r["ceilings"]
+    from freetoken.moe.cpu_executor import _format_label
+
+    c, cpu, rows = r["ceilings"], r["cpu"], r["layer_kinds"]
     print()
     print(f"  host {r['host']}   gpu cuda:{r['gpu']['index']} ({r['gpu']['name']})   "
-          f"cpu {r['cpu']['physical_cores']}c/{r['cpu']['threads_used']}t")
+          f"cpu {cpu['model']} ({cpu['physical_cores']}c, {cpu['numa_nodes']} NUMA)")
     print(f"  ceilings: CPU STREAM read {c['cpu_stream_read_gbs']:.1f}  |  "
-          f"PCIe linear H2D {c['pcie_linear_h2d_gbs']:.1f}  D2H {c['pcie_linear_d2h_gbs']:.1f}  GB/s"
-          f"   (threshold {r['threshold']}x)")
-    # one global ISA-name width so every sweep's GB/s lands in the same column; the
-    # redundant "(<fmt>)" tag is dropped -- the format is already in the row header.
-    all_kernels = list((r.get("dtype_kernels") or {}).values()) + [
-        e for wl in r["workloads"].values() for e in wl["kernels"].values()
-    ]
-    iw = max((len(k.split("(")[0]) for e in all_kernels for k in (e.get("isa_sweep") or {})), default=0)
-    if r.get("dtype_kernels"):
-        print("\n  per-dtype (tuning — what the runtime backend pick matches on)")
-        _print_kernels(r["dtype_kernels"], iw)
-    for name, w in r["workloads"].items():
-        m = w["model"]
-        print(f"\n  {name}  H={m['hidden']} I={m['inter']} E={m['experts']} top_k={m['top_k']}")
-        _print_kernels(w["kernels"], iw)
-    print(f"\n  saved: {r['out_path']}\n")
+          f"PCIe linear H2D {c['pcie_linear_h2d_gbs']:.1f}  D2H {c['pcie_linear_d2h_gbs']:.1f}  GB/s")
+    benched = [rows[key] for key in r["benched"]]
+    # one global ISA-name width so every sweep's GB/s lands in the same column
+    iw = max((len(k.split("(")[0]) for e in benched for k in (e.get("isa_sweep") or {})), default=0)
+    for e in benched:
+        label = _format_label(e["format"], tuple(sorted(e["options"].items())))
+        sources = ", ".join(os.path.basename(s.rstrip("/")) for s in e["sources"])
+        print(f"\n  {label}  H={e['hidden']} I={e['inter']} top_k={e['top_k']} threads={e['cpu_threads']} "
+              f" ({sources}; threshold {e['threshold']}x)")
+        _print_kernels({e["format"]: e}, iw)
+    print(f"\n  saved: {r['out_path']} ({len(benched)} rows benched, {len(rows)} in the profile)\n")
 
 
 # ================================ CLI plumbing ================================
@@ -899,13 +1028,14 @@ def _positive_float(s: str) -> float:
 
 
 def _model_list(s: str) -> tuple[str, ...]:
+    """Preset names and checkpoint paths; a path is any item that exists on disk."""
     if s.strip() == "all":
         return tuple(WORKLOADS)
     items = [x.strip() for x in s.split(",") if x.strip()]
-    bad = [x for x in items if x not in WORKLOADS]
+    bad = [x for x in items if x not in WORKLOADS and not os.path.exists(x)]
     if bad or not items:
         raise argparse.ArgumentTypeError(
-            f"'all' or a comma-separated subset of {sorted(WORKLOADS)}, got {s!r}"
+            f"'all', or a comma-separated list of checkpoint paths and presets from {sorted(WORKLOADS)}, got {s!r}"
         )
     return tuple(dict.fromkeys(items))
 
@@ -951,11 +1081,12 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
     p.add_argument("--dtype", type=_dtype_list, default=None,
                    help=f"tuning bench: 'all' (default when no --model), or comma-separated from "
-                        f"{sorted(DTYPE_WORKLOADS)} -- one bench per expert format (what the runtime "
-                        f"backend pick matches on)")
+                        f"{sorted(DTYPE_WORKLOADS)} -- one bench per expert format at a canonical geometry")
     p.add_argument("--model", type=_model_list, default=None,
                    help=f"per-model detail instead of/in addition to --dtype: 'all', or comma-"
-                        f"separated from {sorted(WORKLOADS)}")
+                        f"separated presets from {sorted(WORKLOADS)} and checkpoint paths (a "
+                        f"checkpoint benches each of its MoE layer kinds, which the hybrid fetch "
+                        f"split matches per layer)")
     p.add_argument("--formats", type=_format_list, default=None,
                    help="override a --model's native expert formats (comma-separated)")
     p.add_argument("--isa", type=_isa_list, default=None, dest="isas",
@@ -986,11 +1117,11 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
     except (ValueError, RuntimeError) as e:
         p.error(str(e))
 
-    models = ns.model or ()
+    models = tuple(m for m in ns.model or () if m in WORKLOADS)
+    checkpoints = tuple(m for m in ns.model or () if m not in WORKLOADS)
     dtypes = ns.dtype
-    # Default (no selection): the full per-dtype tuning bench -- the minimal set the runtime
-    # backend pick needs. --model alone skips it; --dtype + --model runs both.
-    if not models and not dtypes:
+    # Default (no selection): the full per-dtype tuning bench. --model alone skips it; --dtype + --model runs both.
+    if not ns.model and not dtypes:
         dtypes = tuple(DTYPE_WORKLOADS)
 
     try:
@@ -999,6 +1130,7 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
             dtypes=dtypes, formats=ns.formats, isas=ns.isas, cpu_threads=ns.cpu_threads,
             cpu_iters=ns.cpu_iters, pcie_bytes=ns.pcie_mib << 20, pcie_iters=ns.pcie_iters,
             kernel_cpu_iters=ns.kernel_cpu_iters, kernel_pcie_iters=ns.kernel_pcie_iters,
+            checkpoints=checkpoints,
         )
     except (RuntimeError, OSError) as e:
         print(f"error: {e}")
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 5291e621f8..6b428c416b 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -152,12 +152,12 @@ class OffloadMoeCache:
     # of that step's misses are computed on the CPU. 0 -> never fetch (CPU does every
     # miss, the GPU cache stays cold); large -> behaves like pure offload.
     hybrid_max_fetch: int = 1
-    # hybrid only: when > 0, replaces the fixed cap with a per-step fraction -- fetch
+    # hybrid only, one per MoE layer: when > 0, replaces the fixed cap with a per-step fraction -- fetch
     # ~fraction * misses experts over PCIe (rounded to whichever integer balances the
     # overlap best), the CPU computes the rest. The engine sets it to the benched
     # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal
     # time (perfect overlap): fetched : cpu = pcie : cpu - pcie.
-    hybrid_fetch_fraction: float = 0.0
+    hybrid_fetch_fraction: list[float] | None = None
     # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots
     layout: dict | None = None
     max_slots: int | None = None
@@ -1095,9 +1095,8 @@ def ensure_experts_hybrid(
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        return ensure_experts_hybrid(
-            self, layer_id, expert_ids, weights, self.hybrid_max_fetch, self.hybrid_fetch_fraction, clamp
-        )
+        fraction = self.hybrid_fetch_fraction[layer_id] if self.hybrid_fetch_fraction is not None else 0.0
+        return ensure_experts_hybrid(self, layer_id, expert_ids, weights, self.hybrid_max_fetch, fraction, clamp)
 
     def materialize_layer(self, layer_id: int) -> None:
         from freetoken.moe.offload_kernels import materialize_layer
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index be244a117c..56d52ee782 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -644,9 +644,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         default=ServerArgs.moe_strategy,
         choices=["auto", *MOE_STRATEGIES],
         help=(
-            "How the routed experts are served. 'auto' resolves a MoE model to the offload family "
-            "(offload, or hybrid when a `ft bench bw` profile recommends it), and to resident "
-            "'fused' experts on unified-memory GPUs (GB10 / DGX Spark)."
+            "How the routed experts are served. 'auto' resolves a MoE model to 'offload', and to "
+            "resident 'fused' experts on unified-memory GPUs (GB10 / DGX Spark)."
         ),
     )
 
@@ -813,7 +812,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
             "For --moe-strategy hybrid: max experts fetched over PCIe per (layer, decode "
             "step); the rest of that step's misses are computed on the CPU, overlapped. "
             "-1 (default) = auto: fetch the benched pcie/cpu bandwidth fraction of each "
-            "step's misses (perfect overlap; needs an `ft bench bw` profile, else 1). "
+            "step's misses (perfect overlap; needs the `ft bench bw` row of the layer's kind, "
+            "from `ft bench bw --model <checkpoint>`, else 1). "
             "0 = never fetch (all misses on CPU); large = behaves like plain offload."
         ),
     )
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 7b37af9e0e..00b314828f 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -409,9 +409,7 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend
     config = _offload_engine_config()
     _adjust_config(config)
 
-    # Which member of the family gets picked is not this test's claim, and is not ours to
-    # decide: a bare "auto" consults ~/.cache/freetoken/benchbw.json, so a box that has run
-    # `ft bench bw` resolves nvfp4 experts to hybrid instead. Assert the family, not the member.
+    # Which member of the family gets picked is not this test's claim. Assert the family, not the member.
     assert is_offload_moe_strategy(config.moe_strategy)
     assert config.moe_cache_auto is True
     assert config.moe_cache_size == 0  # still unresolved -- the scheduler sizes it from VRAM
diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py
index fad4fa7843..5f36ac7b36 100644
--- a/tests/moe/test_hybrid_fetch.py
+++ b/tests/moe/test_hybrid_fetch.py
@@ -1,7 +1,7 @@
 """Hybrid decode's bandwidth-matched fetch split.
 
-Covers the profile reader of --moe-hybrid-max-fetch auto, which turns `ft bench bw`
-kernel bandwidths into a fetch fraction.
+Covers the profile reader of --moe-hybrid-max-fetch auto, which turns the `ft bench bw`
+row of each MoE layer kind into a fetch fraction.
 """
 
 import json
@@ -9,37 +9,36 @@
 
 import pytest
 
-from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction
+from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, fetch_fraction, layer_kind_key, load_layer_kinds
 
 
-def test_load_hybrid_fetch_fraction(tmp_path):
+def test_layer_kind_rows_give_fetch_fractions(tmp_path):
+    bf16 = layer_kind_key("bf16", {}, 2048, 768, 8, 12)
+    gguf = layer_kind_key("gguf", {"gate_up_type": 12, "down_type": 14}, 2048, 768, 8, 12)
+    assert gguf == "gguf[down_type=14,gate_up_type=12]:2048x768:top8:t12"
     prof = {
         "version": PROFILE_VERSION,
         "gpu": {"name": "FAKE GPU"},
-        "dtype_kernels": {
-            "bf16": {"cpu_moe_gbs": 100.0, "pcie_gather_gbs": 40.0},
+        "layer_kinds": {
+            bf16: {"cpu_moe_gbs": 100.0, "pcie_gather_gbs": 40.0},
             # overlapped (contended) pair wins over the standalone numbers when present
-            "nvfp4_x": {"cpu_moe_gbs": 100.0, "pcie_gather_gbs": 40.0,
-                        "cpu_moe_overlap_gbs": 90.0, "pcie_gather_overlap_gbs": 30.0},
-        },
-        "workloads": {
-            "m": {"kernels": {"ds_fp4": {"cpu_moe_gbs": 80.0, "pcie_gather_gbs": 50.0}}}
+            gguf: {"cpu_moe_gbs": 100.0, "pcie_gather_gbs": 40.0, "cpu_moe_overlap_gbs": 90.0, "pcie_gather_overlap_gbs": 30.0},
         },
     }
     path = tmp_path / "benchbw.json"
     path.write_text(json.dumps(prof))
+    rows = load_layer_kinds(path=str(path))
     # standalone fallback: full-contention assumption -> pcie / cpu
-    assert load_hybrid_fetch_fraction("bf16", path=str(path)) == pytest.approx(0.4)
+    assert fetch_fraction(rows[bf16]) == pytest.approx(0.4)
     # overlapped pair preferred: pcie_ov / (pcie_ov + cpu_ov)
-    assert load_hybrid_fetch_fraction("nvfp4_x", path=str(path)) == pytest.approx(0.25)
-    # per-model fallback when there is no per-dtype entry for the format
-    assert load_hybrid_fetch_fraction("ds_fp4", path=str(path)) == pytest.approx(0.625)
-    assert load_hybrid_fetch_fraction("nvfp4", path=str(path)) is None
+    assert fetch_fraction(rows[gguf]) == pytest.approx(0.25)
+    # another thread count is another kind
+    assert layer_kind_key("bf16", {}, 2048, 768, 8, 6) not in rows
     # a profile from different hardware is ignored
-    assert load_hybrid_fetch_fraction("bf16", gpu_name="OTHER", path=str(path)) is None
-    # so is one written before the current set of CPU formats
+    assert load_layer_kinds(gpu_name="OTHER", path=str(path)) is None
+    # so is one of another version
     path.write_text(json.dumps({**prof, "version": PROFILE_VERSION - 1}))
-    assert load_hybrid_fetch_fraction("bf16", path=str(path)) is None
+    assert load_layer_kinds(path=str(path)) is None
 
 
 def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
@@ -47,15 +46,15 @@ def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
     monkeypatch.delenv("FREETOKEN_BENCHBW_PATH", raising=False)
     uuid = "GPU-2f3a9b1c-0000-1111-2222-333344445555"
 
-    def write(path, name, verdict):
+    def write(path, name, row):
         os.makedirs(os.path.dirname(path), exist_ok=True)
         with open(path, "w") as f:
-            json.dump({"version": PROFILE_VERSION, "gpu": {"name": name}, "dtypes": {"bf16": verdict}}, f)
+            json.dump({"version": PROFILE_VERSION, "gpu": {"name": name}, "layer_kinds": {"k": row}}, f)
 
     # legacy single file only: used when the name matches, ignored otherwise
-    write(default_profile_path(), "FAKE GPU", "hybrid")
-    assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "hybrid"
-    assert load_backend_recommendation("bf16", gpu_name="OTHER", gpu_uuid=uuid) is None
+    write(default_profile_path(), "FAKE GPU", "legacy")
+    assert load_layer_kinds(gpu_name="FAKE GPU", gpu_uuid=uuid) == {"k": "legacy"}
+    assert load_layer_kinds(gpu_name="OTHER", gpu_uuid=uuid) is None
     # this card's own file wins over the legacy one
-    write(default_profile_path(uuid), "FAKE GPU", "offload")
-    assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload"
+    write(default_profile_path(uuid), "FAKE GPU", "card")
+    assert load_layer_kinds(gpu_name="FAKE GPU", gpu_uuid=uuid) == {"k": "card"}
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index a5a7adb9f0..5d7d0ab8a5 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -451,7 +451,7 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch):
     from freetoken.moe import is_offload_moe_strategy
 
     assert config.moe_cache_size == 24
-    # Family, not member: a box with a benchbw profile resolves bf16 experts to hybrid.
+    # Family, not member: which offload-family backend auto picks is not this test's claim.
     assert is_offload_moe_strategy(config.moe_strategy)
 
 

From fd410ee6165cdad0a29b00489c571e84c1e4c6c3 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 18 Oct 2026 14:39:56 -0700
Subject: [PATCH 123/249] perf(moe): pipeline the nvfp4 prefill scale loads on
 geforce sm_86 / sm_89 (#735)

The pointer path of the Triton prefill GEMM loads the e4m3 scales one byte per thread, too small for cp.async, so every K step waits on that load and a CTA barrier. Read them as int32 words (4 scales each) on GeForce sm_86 / sm_89 so they join the pipeline.

- SCALE_WORDS constexpr in `_prefill_nvfp4_moe_kernel` (off by default)
- Triton prefill tuning as one record per (GPU class, cc); GeForce sm_86 / sm_89 turn scale words on, with 8 warps when K >= 2048, and sm_86 moves to BLOCK_SIZE_M 64 from 100 tokens per expert
---
 python/freetoken/kernel/nvfp4_moe.py          | 100 ++++++++++++++----
 .../kernel/triton/nvfp4_fused_moe.py          |  27 ++++-
 2 files changed, 99 insertions(+), 28 deletions(-)

diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py
index a61afda567..0aa494222c 100644
--- a/python/freetoken/kernel/nvfp4_moe.py
+++ b/python/freetoken/kernel/nvfp4_moe.py
@@ -8,6 +8,7 @@
 
 import functools
 import re
+from dataclasses import dataclass, replace
 from typing import Any, Dict, Tuple
 
 import torch
@@ -174,12 +175,51 @@ def _tensor_rho(name: str, major: int, minor: int) -> float:
     return 4.0
 
 
-# (min rho, tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128), from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti
-_PREFILL_BLOCK_M_STEPS = ((8.0, (12, 24, 64)), (3.0, (16, 36, 150)), (0.0, (16, 150, float("inf"))))
 # BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages)
 _PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)}
 _PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)}
 _PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)}
+
+
+@dataclass(frozen=True)
+class _PrefillTuning:
+    """Triton prefill GEMM tuning of one GPU class."""
+
+    steps: Tuple[float, float, float]  # tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128
+    tiles: Dict[int, Tuple[int, int, int]]
+    pair_decode_below: float = 0.0  # tokens per expert under which FAST_DECODE stays 0
+    small_tile_decode: int = 2  # FAST_DECODE of BLOCK_SIZE_M <= 32
+    large_tile_decode: int = 1  # FAST_DECODE of BLOCK_SIZE_M >= 64
+    scale_words: bool = False  # SCALE_WORDS on BLOCK_SIZE_M 64 word-decode tiles
+    long_k_warps: int = 4  # num_warps of SCALE_WORDS GEMMs with K >= _LONG_K
+
+
+_INF = float("inf")
+_LONG_K = 2048
+_TUNING_H100 = _PrefillTuning((12, 24, 64), _PREFILL_TILES_H100)
+_TUNING_RTX_PRO = _PrefillTuning((16, 36, 150), _PREFILL_TILES_RTX_PRO)
+_TUNING_GEFORCE = _PrefillTuning((16, 150, _INF), _PREFILL_TILES_GEFORCE)
+# RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14%, also on 32-row tiles
+_SM89_DECODE = dict(pair_decode_below=48, small_tile_decode=1)
+# H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM)
+_SM90_DECODE = dict(small_tile_decode=2, large_tile_decode=2)
+# RTX 4090 / 3080 Ti layer A/B: scale words + 8 warps on the long-K GEMM give 1.03-1.09x / 1.05-1.19x; with them sm_86 gains from 64-row tiles at 100 tokens per expert, sm_89 does not
+_GEFORCE_SCALE_WORDS = dict(scale_words=True, long_k_warps=8)
+# (rho class, compute capability) -> tuning, from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti; (class, None) covers the rest of the class.
+# The RTX tiles come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 tiles.
+_PREFILL_TUNING = {
+    ("wide", (9, 0)): replace(_TUNING_H100, **_SM90_DECODE),
+    ("wide", None): _TUNING_H100,
+    ("pro", (8, 6)): _TUNING_RTX_PRO,
+    ("pro", (8, 9)): replace(_TUNING_RTX_PRO, **_SM89_DECODE),
+    ("pro", (9, 0)): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100, **_SM90_DECODE),
+    ("pro", (12, 0)): _TUNING_RTX_PRO,
+    ("pro", None): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100),
+    ("geforce", (8, 6)): replace(_TUNING_GEFORCE, steps=(16, 100, _INF), **_GEFORCE_SCALE_WORDS),
+    ("geforce", (8, 9)): replace(_TUNING_GEFORCE, **_SM89_DECODE, **_GEFORCE_SCALE_WORDS),
+    ("geforce", (12, 0)): _TUNING_GEFORCE,
+    ("geforce", None): replace(_TUNING_GEFORCE, tiles=_PREFILL_TILES_H100),
+}
 # steps while the Gluon kernel runs BLOCK_SIZE_M >= 64, from an RTX PRO 6000 sweep; classes without an entry keep the Triton steps
 _PREFILL_BLOCK_M_STEPS_GLUON = ((3.0, (16, 36, 96)),)
 
@@ -213,19 +253,29 @@ def _gluon_prefill(device_index: int):
 
 
 @functools.lru_cache(maxsize=None)
-def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]:
+def _device_rho(device_index: int) -> float:
+    props = torch.cuda.get_device_properties(device_index)
+    return _tensor_rho(props.name, props.major, props.minor)
+
+
+@functools.lru_cache(maxsize=None)
+def _prefill_tuning(device_index: int) -> _PrefillTuning:
     if is_rocm():
         # TODO: tune for AMD GPUs; they keep the H100 tuning
-        return _PREFILL_BLOCK_M_STEPS[0][1], _PREFILL_TILES_H100
-    props = torch.cuda.get_device_properties(device_index)
-    rho = _tensor_rho(props.name, props.major, props.minor)
-    steps = next(s for min_rho, s in _PREFILL_BLOCK_M_STEPS if rho >= min_rho)
+        return _TUNING_H100
+    rho = _device_rho(device_index)
+    rho_class = "wide" if rho >= 8.0 else "pro" if rho >= 3.0 else "geforce"
+    cc = torch.cuda.get_device_capability(device_index)
+    return _PREFILL_TUNING.get((rho_class, cc), _PREFILL_TUNING[(rho_class, None)])
+
+
+@functools.lru_cache(maxsize=None)
+def _prefill_block_m_steps(device_index: int) -> Tuple[float, float, float]:
+    steps = _prefill_tuning(device_index).steps
     if _gluon_prefill(device_index) is not None:
+        rho = _device_rho(device_index)
         steps = next((s for min_rho, s in _PREFILL_BLOCK_M_STEPS_GLUON if rho >= min_rho), steps)
-    # the RTX tables come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 one
-    if (props.major, props.minor) not in ((8, 6), (8, 9), (12, 0)):
-        return steps, _PREFILL_TILES_H100
-    return steps, _PREFILL_TILES_RTX_PRO if rho >= 3.0 else _PREFILL_TILES_GEFORCE
+    return steps
 
 
 @functools.lru_cache(maxsize=None)
@@ -241,31 +291,29 @@ def _prefill_use_tma(packed: torch.Tensor, scale: torch.Tensor, K: int, device_i
             and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0)
 
 
-@functools.lru_cache(maxsize=None)
-def _device_cc(device_index: int) -> Tuple[int, int]:
-    return torch.cuda.get_device_capability(device_index)
+def _prefill_use_scale_words(scale: torch.Tensor, K: int, cfg: Dict[str, Any], device_index: int) -> bool:
+    # K // 16 scale bytes per row split into whole int32 words, and a K step reads at least one word
+    return (_prefill_tuning(device_index).scale_words and cfg["BLOCK_SIZE_M"] == 64 and cfg["BLOCK_SIZE_KB"] >= 32
+            and K % 64 == 0 and scale.is_contiguous() and scale.data_ptr() % 4 == 0)
 
 
 def _prefill_fast_decode(block_m: int, per_expert: float, device_index: int) -> int:
     """FAST_DECODE of the prefill kernel: 0 keeps the fp16 pair decode (ROCm: the fast decode is PTX inline asm), 1 words, 2 bytes."""
     if is_rocm():
         return 0
-    if _device_cc(device_index) == (8, 9):
-        # RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14%
-        return 0 if per_expert < 48 else 1
-    if _device_cc(device_index) == (9, 0):
-        # H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM)
-        return 2
+    tuning = _prefill_tuning(device_index)
+    if per_expert < tuning.pair_decode_below:
+        return 0
     # the word decode halves the loop but stages the decoded tile through shared memory; small tiles keep the byte form
-    return 2 if block_m <= 32 else 1
+    return tuning.small_tile_decode if block_m <= 32 else tuning.large_tile_decode
 
 
 def prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]:
     # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it.
-    (to32, to64, to128), tiles = _prefill_tuning(device_index)
+    to32, to64, to128 = _prefill_block_m_steps(device_index)
     per_expert = M * top_k / num_experts
     block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16
-    block_n, block_kb, stages = tiles[block_m]
+    block_n, block_kb, stages = _prefill_tuning(device_index).tiles[block_m]
     return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb,
                 GROUP_SIZE_M=8, num_warps=4, num_stages=stages,
                 FAST_DECODE=_prefill_fast_decode(block_m, per_expert, device_index))
@@ -308,6 +356,11 @@ def prefill_gemm(
     use_tma = _prefill_use_tma(packed, scale, K, a.device.index)
     # the 16-byte scale box costs registers that the BLOCK_SIZE_M >= 64 tiles do not have
     tma_scales = cfg["BLOCK_SIZE_M"] <= 32
+    scale_words = not use_tma and fast == 1 and _prefill_use_scale_words(scale, K, cfg, a.device.index)
+    if scale_words:
+        scale = scale.view(torch.int32)
+        if K >= _LONG_K:
+            cfg["num_warps"] = _prefill_tuning(a.device.index).long_k_warps
     w_desc = s_desc = None
     if use_tma:
         rows = packed.shape[0] * N
@@ -332,5 +385,6 @@ def prefill_gemm(
         USE_TMA=use_tma,
         TMA_SCALES=tma_scales,
         FAST_DECODE=fast,
+        SCALE_WORDS=scale_words,
         **cfg,
     )
diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
index 0fa4d4aa9a..61e867b160 100644
--- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py
+++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
@@ -385,6 +385,7 @@ def _prefill_nvfp4_moe_kernel(
     USE_TMA: tl.constexpr = False,     # needs K % 256 == 0
     TMA_SCALES: tl.constexpr = False,  # with USE_TMA: scales by TMA too, else pointer loads
     FAST_DECODE: tl.constexpr = 0,     # 0: fp16 pair decode, 1: packed_ptr / w_desc are int32 views of 8 codes, 2: bytes
+    SCALE_WORDS: tl.constexpr = False,  # with FAST_DECODE == 1 and pointer loads: scale_ptr is an int32 view of 4 scales
 ):
     # out^T = W @ X^T: wgmma reads its A operand (the decoded weights) from registers, while B must be in shared memory.
     # The token block is the MMA N dimension, so a 16-token block still fills the MMA.
@@ -420,7 +421,13 @@ def _prefill_nvfp4_moe_kernel(
         p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kw[None, :] * stride_pkb
     else:
         p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb
-    s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk
+    tl.static_assert(not SCALE_WORDS or (WORDS and not USE_TMA and BLOCK_SIZE_KB >= 32))
+    SW: tl.constexpr = BLOCK_SIZE_KB // 32
+    if SCALE_WORDS:
+        offs_sw = tl.arange(0, SW)
+        s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sw[None, :] * stride_sblk
+    else:
+        s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk
 
     accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
     K_BYTES = K // 2
@@ -450,7 +457,16 @@ def _prefill_nvfp4_moe_kernel(
             else:
                 byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES
                 v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32)
-            if e4m3_native_cx():
+            if SCALE_WORDS:
+                # 4-byte loads can go through the cp.async pipeline, byte loads cannot; spread each word over the 8 weight words it scales
+                sw = tl.load(s_ptrs, mask=(kb * SW + offs_sw < K_BYTES // 32)[None, :], other=0)
+                sw = tl.reshape(tl.broadcast_to(sw[:, :, None], (BLOCK_SIZE_N, SW, KW // SW)), (BLOCK_SIZE_N, KW))
+                s = ((sw >> ((offs_kw // 2 % 4) * 8)[None, :]) & 0xFF).to(tl.uint8)
+                if e4m3_native_cx():
+                    s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+                else:
+                    s = e4m3_u8_to_f32(s)
+            elif e4m3_native_cx():
                 s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32)
             else:
                 s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0))
@@ -459,8 +475,9 @@ def _prefill_nvfp4_moe_kernel(
             s = (s * 3.3230699894622794e35).to(tl.bfloat16).to(tl.uint16, bitcast=True).to(tl.uint32)
             s = s | (s << 16)
             if WORDS:
-                # one scale per 16 values, two words
-                s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW))
+                if not SCALE_WORDS:
+                    # one scale per 16 values, two words
+                    s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW))
                 w = _e2m1x8_bf16(v.to(tl.uint32, bitcast=True), s, BLOCK_SIZE_N, KW)
             else:
                 # one scale per 16 values, eight bytes
@@ -481,7 +498,7 @@ def _prefill_nvfp4_moe_kernel(
         if not USE_TMA:
             p_ptrs += (KW if WORDS else BLOCK_SIZE_KB) * stride_pkb
         if not (USE_TMA and TMA_SCALES):
-            s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk
+            s_ptrs += (SW if SCALE_WORDS else BLOCK_SIZE_KB // 8) * stride_sblk
 
     g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32)
     if FAST_DECODE:

From e4f69d9e232448fc168389c29f6f9364869cff35 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 18 Oct 2026 20:10:51 -0700
Subject: [PATCH 124/249] perf(moe): prefill nvfp4 moe experts on nvfp4
 activations on sm_120 (#736)

---
 .../kernel/gluon/nvfp4/moe/sm120_w4a4.py      | 536 ++++++++++++++++++
 python/freetoken/kernel/gluon/nvfp4/utils.py  |  39 ++
 python/freetoken/kernel/nvfp4_moe.py          |  74 ++-
 .../kernel/triton/nvfp4_act_quant.py          |  42 ++
 .../layers/quantization/moe/nvfp4.py          |  52 +-
 python/freetoken/models/nvfp4_banks.py        |  54 +-
 python/freetoken/moe/expert_pieces.py         |   4 +-
 python/freetoken/moe/fused_nvfp4.py           |  53 ++
 python/freetoken/moe/legacy_format.py         |   6 +-
 tests/moe/test_nvfp4_backends.py              |  98 +++-
 10 files changed, 939 insertions(+), 19 deletions(-)
 create mode 100644 python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py
 create mode 100644 python/freetoken/kernel/triton/nvfp4_act_quant.py

diff --git a/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py b/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py
new file mode 100644
index 0000000000..2a2c1430ce
--- /dev/null
+++ b/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py
@@ -0,0 +1,536 @@
+"""NVFP4 W4A4 MoE prefill GEMMs for sm_120: NVFP4 weights times NVFP4 activations on mma.sync kind::mxf4nvf4.
+
+The activations are quantized under static global scales (the checkpoint's input_scale), so each GEMM input is e2m1 codes, per-16 e4m3 scales and one fp32 scalar.
+:func:`gate_up_gemm` applies the gated activation in its epilogue and writes the down GEMM's NVFP4 input; :func:`down_gemm` scales by the routing weight and adds each route into its token's output row with bf16 atomics (the sum order, and so the result, varies from run to run).
+Routing follows the W4A16 kernels (moe_align sorted ids, expert per M block, swap-AB accumulator [BN, BM]); :mod:`freetoken.kernel.nvfp4_moe` picks the constexprs.
+"""
+
+import torch
+import triton
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language.nvidia.ampere import async_copy
+from triton.experimental.gluon.language.nvidia.hopper import mbarrier, tma
+from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
+
+from ..utils import mma_nvf4, scale_layout
+
+# gated activations of the gate_up epilogue, as gated_act_and_mul computes them
+ACTIVATIONS = {"silu": 0, "gelu_tanh": 1, "swigluoai": 2, "swiglu_clamp": 3}
+
+SCALE_SMEM_LAYOUT = gl.constexpr(gl.SwizzledSharedLayout(1, 1, 1, [1, 0]))
+
+
+@gluon.jit
+def _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M: gl.constexpr, N_GROUP: gl.constexpr):
+    if N_GROUP > 0:
+        # N-outer strips of N_GROUP tiles: concurrent CTAs add into a narrow output strip that stays in L2
+        per_strip = num_pid_m * N_GROUP
+        strip = t // per_strip
+        rem = t % per_strip
+        width = min(N_GROUP, num_pid_n - strip * N_GROUP)
+        return rem // width, strip * N_GROUP + rem % width
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    first_pid_m = (t // num_pid_in_group) * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    return first_pid_m + ((t % num_pid_in_group) % group_size_m), (t % num_pid_in_group) // group_size_m
+
+
+@gluon.jit
+def _scale_ptrs(ws_ptr, xs_ptr, ids_ptr, w_row, w_row_up, x_div, pid_m, num_valid, ks4, n_rows,
+                BM: gl.constexpr, BN: gl.constexpr, K_TAIL: gl.constexpr, GATED: gl.constexpr, NWL: gl.constexpr):
+    # scale rows are K / 16 bytes: 16-byte copies when K % 256 == 0, else 4-byte copies zeroed past the row end
+    SV: gl.constexpr = 4 if K_TAIL else 16
+    sll: gl.constexpr = gl.BlockedLayout([1, SV], [32 * SV // 16, 16 // SV], [NWL, 1], [1, 0])
+    s_cols = gl.arange(0, 16, layout=gl.SliceLayout(0, sll))
+    # both forms equal K // 16; the second tells the compiler the row is a multiple of 16 bytes
+    row_len = ks4 * 4 if K_TAIL else (ks4 // 4) * 16
+    r = gl.arange(0, BN, layout=gl.SliceLayout(1, sll))
+    if GATED:
+        sw_rows = gl.where(r < BN // 2, w_row + r, w_row_up + r - BN // 2)
+    else:
+        sw_rows = w_row + r
+    sw_ptrs = ws_ptr + sw_rows.to(gl.int64)[:, None] * row_len + s_cols[None, :]
+    sw_ok = (sw_rows < n_rows)[:, None]
+    sx_tok = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(1, sll)))
+    # the x scale rows are contiguous K / 16 bytes too; row_len keeps the divisibility a stride argument loses
+    sx_ptrs = xs_ptr + ((sx_tok // x_div).to(gl.int64) * row_len)[:, None] + s_cols[None, :]
+    sx_ok = (sx_tok < num_valid)[:, None]
+    return sw_ptrs, sw_ok, sx_ptrs, sx_ok, s_cols, row_len
+
+
+@gluon.jit
+def _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid,
+               stride_xq, ks4, n_rows, x_div, K_TAIL: gl.constexpr, GATED: gl.constexpr):
+    x_bufs, w_bufs, sw_bufs, sx_bufs = bufs
+    ready, empty, s_ready, s_empty = bars
+    S: gl.constexpr = x_bufs.type.shape[0]
+    BM: gl.constexpr = x_bufs.type.shape[1]
+    KB: gl.constexpr = x_bufs.type.shape[2]
+    BN: gl.constexpr = w_bufs.type.shape[1]
+    SPC: gl.constexpr = 128 // KB
+    NWL: gl.constexpr = gl.num_warps()
+
+    if GATED:
+        # rows [0, BN/2) of the tile are gate rows, [BN/2, BN) the matching up rows N/2 further down
+        w_row = gl.load(eid_ptr + pid_m) * N + pid_n * (BN // 2)
+        w_row_up = w_row + N // 2
+    else:
+        w_row = gl.load(eid_ptr + pid_m) * N + pid_n * BN
+        w_row_up = w_row
+
+    # x rows: 16 bytes per cp.async, KB // 16 lanes per row
+    xl: gl.constexpr = gl.BlockedLayout([1, 16], [32 * 16 // KB, KB // 16], [NWL, 1], [1, 0])
+    tok = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(1, xl)))
+    x_ptrs = xq_ptr + ((tok // x_div).to(gl.int64) * stride_xq)[:, None] + gl.arange(0, KB, layout=gl.SliceLayout(0, xl))[None, :]
+    x_mask = (tok < num_valid)[:, None]
+    sw_ptrs, sw_ok, sx_ptrs, sx_ok, s_cols, row_len = _scale_ptrs(
+        ws_ptr, xs_ptr, ids_ptr, w_row, w_row_up, x_div, pid_m, num_valid, ks4, n_rows, BM, BN, K_TAIL, GATED, NWL)
+
+    for kk in range(nk):
+        if kk % SPC == 0:
+            jl = kk // SPC
+            j = jt + jl
+            ss = j % 2
+            mbarrier.wait(s_empty.index(ss), (j // 2 & 1) ^ 1)
+            col_ok = (jl * 16 + s_cols < row_len)[None, :]
+            async_copy.async_load(sw_bufs.index(ss), sw_ptrs + jl * 16, mask=sw_ok & col_ok)
+            async_copy.async_load(sx_bufs.index(ss), sx_ptrs + jl * 16, mask=sx_ok & col_ok)
+            async_copy.mbarrier_arrive(s_ready.index(ss))
+            mbarrier.arrive(s_ready.index(ss))
+        c = it + kk
+        st = c % S
+        mbarrier.wait(empty.index(st), (c // S & 1) ^ 1)
+        # the cp.async arrive goes first: it raises the pending count before the TMA expect can complete the phase
+        async_copy.async_load(x_bufs.index(st), x_ptrs + kk * KB, mask=x_mask)
+        async_copy.mbarrier_arrive(ready.index(st))
+        if GATED:
+            mbarrier.expect(ready.index(st), 2 * w_desc.block_type.nbytes)
+            tma.async_load(w_desc, [w_row, kk * KB], ready.index(st), w_bufs.index(st).slice(0, BN // 2))
+            tma.async_load(w_desc, [w_row_up, kk * KB], ready.index(st), w_bufs.index(st).slice(BN // 2, BN // 2))
+        else:
+            mbarrier.expect(ready.index(st), w_desc.block_type.nbytes)
+            tma.async_load(w_desc, [w_row, kk * KB], ready.index(st), w_bufs.index(st))
+
+
+@gluon.jit
+def _load_partition(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, N, nk, num_pid_m, num_valid,
+                    stride_xq, ks4, n_rows, x_div, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr,
+                    GATED: gl.constexpr, N_GROUP: gl.constexpr):
+    BN: gl.constexpr = bufs[1].type.shape[1]
+    pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M, N_GROUP)
+    _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, pid_m, pid_n, N, nk, 0, 0, num_valid,
+               stride_xq, ks4, n_rows, x_div, K_TAIL, GATED)
+
+
+@gluon.jit
+def _mma_partition(bufs, bars, nk, mma: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, K_TAIL: gl.constexpr):
+    x_bufs, w_bufs, sw_bufs, sx_bufs = bufs
+    ready, empty, s_ready, s_empty = bars
+    S: gl.constexpr = x_bufs.type.shape[0]
+    BM: gl.constexpr = x_bufs.type.shape[1]
+    KB: gl.constexpr = x_bufs.type.shape[2]
+    BN: gl.constexpr = w_bufs.type.shape[1]
+    SB: gl.constexpr = KB // 8
+    SPC: gl.constexpr = 128 // KB
+    al: gl.constexpr = gl.DotOperandLayout(0, mma, 4)
+    bl: gl.constexpr = gl.DotOperandLayout(1, mma, 4)
+    acc = gl.zeros([BN, BM], gl.float32, mma)
+    # with K % 256 != 0 the last scale block covers fewer than SPC chunks
+    for jb in range(gl.cdiv(nk, SPC)):
+        ss = jb % 2
+        mbarrier.wait(s_ready.index(ss), jb // 2 & 1)
+        for c in gl.static_range(SPC):
+            kk = jb * SPC + c
+            if not K_TAIL or kk < nk:
+                st = kk % S
+                mbarrier.wait(ready.index(st), kk // S & 1)
+                a = w_bufs.index(st).load(al)
+                b = x_bufs.index(st).permute((1, 0)).load(bl)
+                sw = sw_bufs.index(ss).slice(c * SB, SB, dim=1).load(wsl).to(gl.float8e4nv, bitcast=True)
+                sx = sx_bufs.index(ss).slice(c * SB, SB, dim=1).load(xsl).to(gl.float8e4nv, bitcast=True)
+                mbarrier.arrive(empty.index(st))
+                acc = mma_nvf4(a, sw, b, sx, acc)
+        mbarrier.arrive(s_empty.index(ss))
+    return acc
+
+
+@gluon.jit
+def _e2m1x2(hi, lo):
+    # cvt packs its first source into the high nibble
+    return gl.inline_asm_elementwise("{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $1, $2; cvt.u16.u8 $0, t; }",
+                                     "=h,r,r", [hi, lo], gl.uint16, True, 1)
+
+
+@gluon.jit
+def _gated_epilogue(acc, slot, pid_m, pid_n, wg_ptr, g_in, ids_ptr, q_ptr, qs_ptr, g_out, N, num_valid, stride_ge,
+                    ACT: gl.constexpr, act_alpha, act_limit):
+    """act(gate, up) of the tile's BN / 2 columns as NVFP4 under the global g_out: q [routes, N / 4], qs [routes, N / 32]."""
+    BN: gl.constexpr = acc.type.shape[0]
+    BM: gl.constexpr = acc.type.shape[1]
+    H2: gl.constexpr = BN // 2
+    G: gl.constexpr = H2 // 16
+    NW: gl.constexpr = gl.num_warps()
+    out_l: gl.constexpr = acc.type.layout
+    r = gl.arange(0, BN, layout=gl.SliceLayout(1, out_l))
+    col = gl.where(r < H2, pid_n * H2 + r, N // 2 + pid_n * H2 + r - H2)
+    wg = gl.load(wg_ptr + slot.to(gl.int64) * stride_ge + col).to(gl.float32) * g_in
+    # bf16 like the W4A16 gate_up output; n-contiguous 16-byte runs, swizzled so the lanes of a warp spread over banks
+    h_l: gl.constexpr = gl.SwizzledSharedLayout(8, 1, 8, [0, 1])
+    h_s = gl.allocate_shared_memory(gl.bfloat16, [BN, BM], h_l, (acc * wg[:, None]).to(gl.bfloat16))
+    # one thread holds 16 consecutive columns of one route
+    TG: gl.constexpr = min(32, G)
+    ql2: gl.constexpr = gl.BlockedLayout([16, 1], [TG, 32 // TG], [1, NW], [0, 1])
+    l4: gl.constexpr = gl.BlockedLayout([1, 8, 2, 1], [TG, 1, 1, 32 // TG], [1, 1, 1, NW], [2, 1, 0, 3])
+    gate = gl.convert_layout(gl.reshape(h_s.slice(0, H2).load(ql2), [G, 8, 2, BM]), l4, assert_trivial=True).to(gl.float32)
+    up = gl.convert_layout(gl.reshape(h_s.slice(H2, H2).load(ql2), [G, 8, 2, BM]), l4, assert_trivial=True).to(gl.float32)
+    if ACT == 1:
+        inner = 0.7978845608028654 * (gate + 0.044715 * gate * gate * gate)
+        # tanh(y) = 1 - 2 / (exp(2y) + 1)
+        h = 0.5 * gate * (2.0 - 2.0 / (gl.exp(2.0 * inner) + 1.0)) * up
+    elif ACT == 2 or ACT == 3:
+        gate = gl.minimum(gate, act_limit)
+        up = gl.minimum(gl.maximum(up, -act_limit), act_limit)
+        if ACT == 2:
+            up = up + 1.0
+        h = gate / (1.0 + gl.exp(-gate * act_alpha)) * up
+    else:
+        h = gate / (1.0 + gl.exp(-gate)) * up
+    h = h.to(gl.bfloat16).to(gl.float32)
+    amax = gl.max(gl.max(gl.abs(h), axis=2), axis=1)
+    sc = (amax / (6.0 * g_out)).to(gl.float8e4nv)
+    sf = sc.to(gl.float32) * g_out
+    inv = gl.where(sf > 0, 1.0 / sf, 0.0)
+    v = h * gl.expand_dims(gl.expand_dims(inv, 1), 2)
+    lo, hi = gl.split(gl.permute(v, (0, 1, 3, 2)))
+    q = _e2m1x2(hi, lo).to(gl.uint8)
+    ql: gl.constexpr = q.type.layout
+    g_i = gl.arange(0, G, layout=gl.SliceLayout(1, gl.SliceLayout(2, ql)))
+    b_i = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(2, ql)))
+    rq = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(1, ql))))
+    q_cols = pid_n * (H2 // 2) + g_i[:, None] * 8 + b_i[None, :]
+    q_ptrs = q_ptr + (rq.to(gl.int64) * (N // 4))[None, None, :] + q_cols[:, :, None]
+    gl.store(q_ptrs, q, mask=(rq < num_valid)[None, None, :])
+    sl: gl.constexpr = sc.type.layout
+    sg = gl.arange(0, G, layout=gl.SliceLayout(1, sl))
+    rs = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, sl)))
+    s_ptrs = qs_ptr + (rs.to(gl.int64) * (N // 32))[None, :] + (pid_n * G + sg)[:, None]
+    gl.store(s_ptrs, sc.to(gl.uint8, bitcast=True), mask=(rs < num_valid)[None, :])
+
+
+@gluon.jit
+def _w4a4_moe_kernel(
+    xq_ptr, xs_ptr, g_in_ptr, w_desc, ws_ptr, wg_ptr, out_ptr, topk_weights_ptr,
+    sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr,
+    q_ptr, qs_ptr, g_out_ptr, act_alpha, act_limit,
+    N, K, EM, num_valid_tokens, ks4, n_rows,
+    stride_xq, stride_ge, stride_out, stride_tw,
+    BLOCK_SIZE_M: gl.constexpr, BLOCK_SIZE_N: gl.constexpr, BLOCK_SIZE_K: gl.constexpr,
+    GROUP_SIZE_M: gl.constexpr, N_GROUP: gl.constexpr, STAGES: gl.constexpr, MMA_WARPS_N: gl.constexpr,
+    LOAD_WARPS: gl.constexpr, LOAD_REGS: gl.constexpr, K_TAIL: gl.constexpr, GATED: gl.constexpr,
+    ACT: gl.constexpr, top_k: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr,
+):
+    BM: gl.constexpr = BLOCK_SIZE_M
+    BN: gl.constexpr = BLOCK_SIZE_N
+    KB: gl.constexpr = BLOCK_SIZE_K // 2
+    num_pid_m = gl.cdiv(EM, BM)
+    pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M, N_GROUP)
+    if pid_m >= gl.cdiv(gl.load(num_tokens_post_padded_ptr), BM):
+        return
+    nk = K // BLOCK_SIZE_K
+    # gate_up reads one x row per token, down one row per route
+    x_div = top_k if GATED else 1
+
+    x_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BM, KB], gl.NVMMASharedLayout.get_default_for([BM, KB], gl.uint8))
+    w_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BN, KB], w_desc.layout)
+    sw_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT)
+    sx_bufs = gl.allocate_shared_memory(gl.uint8, [2, BM, 16], SCALE_SMEM_LAYOUT)
+    ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    for i in gl.static_range(STAGES):
+        mbarrier.init(ready.index(i), count=1)
+        mbarrier.init(empty.index(i), count=1)
+    for i in gl.static_range(2):
+        mbarrier.init(s_ready.index(i), count=1)
+        mbarrier.init(s_empty.index(i), count=1)
+
+    NW: gl.constexpr = gl.num_warps()
+    mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N],
+                                                  instr_shape=[16, 8])
+    # tuples go inline: a tuple stored in a variable turns its constexprs into runtime values
+    acc = gl.warp_specialize([
+        (_mma_partition, ((x_bufs, w_bufs, sw_bufs, sx_bufs), (ready, empty, s_ready, s_empty), nk, mma, wsl, xsl, K_TAIL)),
+        (_load_partition, (xq_ptr, xs_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, ws_ptr,
+                           (x_bufs, w_bufs, sw_bufs, sx_bufs), (ready, empty, s_ready, s_empty), N, nk, num_pid_m,
+                           num_valid_tokens, stride_xq, ks4, n_rows, x_div, K_TAIL, GROUP_SIZE_M, GATED, N_GROUP)),
+    ], [LOAD_WARPS], [LOAD_REGS])
+
+    out_l: gl.constexpr = gl.BlockedLayout([8, 1], [BN // 8, 32 * 8 // BN], [1, NW], [0, 1])
+    acc = gl.convert_layout(acc, out_l)
+    slot = gl.load(expert_ids_ptr + pid_m)
+    g_in = gl.load(g_in_ptr)
+    if GATED:
+        _gated_epilogue(acc, slot, pid_m, pid_n, wg_ptr, g_in, sorted_token_ids_ptr, q_ptr, qs_ptr, gl.load(g_out_ptr),
+                        N, num_valid_tokens, stride_ge, ACT, act_alpha, act_limit)
+    else:
+        offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, out_l))
+        n_mask = offs_n < N
+        route = gl.load(sorted_token_ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, out_l)))
+        route_mask = route < num_valid_tokens
+        wg = gl.load(wg_ptr + slot.to(gl.int64) * stride_ge + offs_n, mask=n_mask, other=0.0).to(gl.float32) * g_in
+        tw = gl.load(topk_weights_ptr + route * stride_tw, mask=route_mask, other=0.0)
+        acc = acc * wg[:, None] * tw[None, :]
+        out_ptrs = out_ptr + (route // top_k).to(gl.int64)[None, :] * stride_out + offs_n[:, None]
+        gl.atomic_add(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=route_mask[None, :] & n_mask[:, None], sem="relaxed")
+
+
+# Persistent down GEMM: the 256 x 128 tile and 8 MMA warps of the kernel above, with the routing-weight scale and the atomics on
+# separate epilogue warps. The accumulator is two 128 x 128 halves that go to the epilogue one at a time through one 32 KB buffer,
+# so the next tile's mainloop overlaps the atomics of the last half.
+
+
+@gluon.jit
+def _num_tiles(ntpp_ptr, N, BM: gl.constexpr, BN: gl.constexpr):
+    num_pid_m = gl.cdiv(gl.load(ntpp_ptr), BM)
+    num_pid_n = gl.cdiv(N, BN)
+    return num_pid_m, num_pid_n, num_pid_m * num_pid_n
+
+
+@gluon.jit
+def _down_load_partition(xq_ptr, xs_ptr, ids_ptr, eid_ptr, ntpp_ptr, w_desc, ws_ptr, bufs, bars, N, nk, num_valid,
+                         stride_xq, ks4, n_rows, K_TAIL: gl.constexpr, N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr):
+    BM: gl.constexpr = bufs[0].type.shape[1]
+    KB: gl.constexpr = bufs[0].type.shape[2]
+    BN: gl.constexpr = bufs[1].type.shape[1]
+    SPC: gl.constexpr = 128 // KB
+    num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN)
+    # it / jt count chunks and scale blocks across tiles, so the rings keep their phase from one tile to the next
+    it = 0
+    jt = 0
+    for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+        pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP)
+        _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, (bufs[0], bufs[1], bufs[2], bufs[3]),
+                   (bars[0], bars[1], bars[2], bars[3]), pid_m, pid_n, N, nk, it, jt, num_valid, stride_xq, ks4,
+                   n_rows, 1, K_TAIL, False)
+        it += nk
+        jt += gl.cdiv(nk, SPC)
+
+
+@gluon.jit
+def _down_epi_partition(ids_ptr, ntpp_ptr, out_ptr, bufs, bars, N, num_valid, stride_out, top_k: gl.constexpr,
+                        N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr):
+    h_buf = bufs[4]
+    acc_ready, acc_empty = bars[4], bars[5]
+    HB: gl.constexpr = h_buf.type.shape[0]
+    BM: gl.constexpr = h_buf.type.shape[1]
+    BN: gl.constexpr = 2 * HB
+    CH: gl.constexpr = 64
+    NWE: gl.constexpr = gl.num_warps()
+    # 8 consecutive n of one route per thread: one 16-byte bf16 atomic; 64 rows per pass keep the registers low
+    el: gl.constexpr = gl.BlockedLayout([8, 1], [CH // 8, 32 * 8 // CH], [1, NWE], [0, 1])
+    offs_cl = gl.arange(0, CH, layout=gl.SliceLayout(1, el))
+    offs_tl = gl.arange(0, BM, layout=gl.SliceLayout(0, el))
+    num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN)
+    hc = 0
+    for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+        pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP)
+        route = gl.load(ids_ptr + pid_m * BM + offs_tl)
+        rmask = (route < num_valid)[None, :]
+        rows = (route // top_k).to(gl.int64)[None, :] * stride_out
+        for half in gl.static_range(2):
+            mbarrier.wait(acc_ready.index(0), hc & 1)
+            for ch in gl.static_range(HB // CH):
+                offs_n = pid_n * BN + half * HB + ch * CH + offs_cl
+                h = h_buf.slice(ch * CH, CH).load(el)
+                gl.atomic_add(out_ptr + rows + offs_n[:, None], h, mask=rmask & (offs_n < N)[:, None], sem="relaxed")
+            mbarrier.arrive(acc_empty.index(0))
+            hc += 1
+
+
+@gluon.jit
+def _down_mma_partition(bufs, bars, ids_ptr, eid_ptr, ntpp_ptr, wg_ptr, g_in_ptr, tw_ptr, N, nk, num_valid, stride_ge,
+                        stride_tw, mma: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, K_TAIL: gl.constexpr,
+                        N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr):
+    x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf = bufs
+    ready, empty, s_ready, s_empty, acc_ready, acc_empty = bars
+    S: gl.constexpr = x_bufs.type.shape[0]
+    BM: gl.constexpr = x_bufs.type.shape[1]
+    KB: gl.constexpr = x_bufs.type.shape[2]
+    BN: gl.constexpr = w_bufs.type.shape[1]
+    HB: gl.constexpr = BN // 2
+    SB: gl.constexpr = KB // 8
+    SPC: gl.constexpr = 128 // KB
+    al: gl.constexpr = gl.DotOperandLayout(0, mma, 4)
+    bl: gl.constexpr = gl.DotOperandLayout(1, mma, 4)
+    offs_nl = gl.arange(0, HB, layout=gl.SliceLayout(1, mma))
+    offs_tl = gl.arange(0, BM, layout=gl.SliceLayout(0, mma))
+    g_in = gl.load(g_in_ptr)
+    num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN)
+    it = 0
+    jt = 0
+    hc = 0
+    for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)):
+        pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP)
+        acc0 = gl.zeros([HB, BM], gl.float32, mma)
+        acc1 = gl.zeros([HB, BM], gl.float32, mma)
+        for jb in range(gl.cdiv(nk, SPC)):
+            j = jt + jb
+            ss = j % 2
+            mbarrier.wait(s_ready.index(ss), j // 2 & 1)
+            for c in gl.static_range(SPC):
+                kk = jb * SPC + c
+                if not K_TAIL or kk < nk:
+                    i = it + kk
+                    st = i % S
+                    mbarrier.wait(ready.index(st), i // S & 1)
+                    b = x_bufs.index(st).permute((1, 0)).load(bl)
+                    sx = sx_bufs.index(ss).slice(c * SB, SB, dim=1).load(xsl).to(gl.float8e4nv, bitcast=True)
+                    a0 = w_bufs.index(st).slice(0, HB).load(al)
+                    a1 = w_bufs.index(st).slice(HB, HB).load(al)
+                    sws = sw_bufs.index(ss).slice(c * SB, SB, dim=1)
+                    sw0 = sws.slice(0, HB).load(wsl).to(gl.float8e4nv, bitcast=True)
+                    sw1 = sws.slice(HB, HB).load(wsl).to(gl.float8e4nv, bitcast=True)
+                    mbarrier.arrive(empty.index(st))
+                    acc0 = mma_nvf4(a0, sw0, b, sx, acc0)
+                    acc1 = mma_nvf4(a1, sw1, b, sx, acc1)
+            mbarrier.arrive(s_empty.index(ss))
+        it += nk
+        jt += gl.cdiv(nk, SPC)
+        slot = gl.load(eid_ptr + pid_m)
+        route = gl.load(ids_ptr + pid_m * BM + offs_tl)
+        rmask = route < num_valid
+        xs_ = gl.load(tw_ptr + route * stride_tw, mask=rmask, other=0.0) * g_in
+        wrow = wg_ptr + slot.to(gl.int64) * stride_ge
+        n0 = pid_n * BN + offs_nl
+        wg0 = gl.load(wrow + n0, mask=n0 < N, other=0.0).to(gl.float32)
+        h0 = (acc0 * wg0[:, None] * xs_[None, :]).to(gl.bfloat16)
+        mbarrier.wait(acc_empty.index(0), (hc & 1) ^ 1)
+        h_buf.store(h0)
+        mbarrier.arrive(acc_ready.index(0))
+        hc += 1
+        n1 = n0 + HB
+        wg1 = gl.load(wrow + n1, mask=n1 < N, other=0.0).to(gl.float32)
+        h1 = (acc1 * wg1[:, None] * xs_[None, :]).to(gl.bfloat16)
+        mbarrier.wait(acc_empty.index(0), (hc & 1) ^ 1)
+        h_buf.store(h1)
+        mbarrier.arrive(acc_ready.index(0))
+        hc += 1
+
+
+@gluon.jit
+def _w4a4_down_persistent_kernel(
+    xq_ptr, xs_ptr, g_in_ptr, w_desc, ws_ptr, wg_ptr, out_ptr, topk_weights_ptr,
+    sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr,
+    N, K, num_valid_tokens, ks4, n_rows, stride_xq, stride_ge, stride_out, stride_tw,
+    BLOCK_SIZE_M: gl.constexpr, BLOCK_SIZE_N: gl.constexpr, BLOCK_SIZE_K: gl.constexpr, STAGES: gl.constexpr,
+    K_TAIL: gl.constexpr, N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr, MMA_WARPS_N: gl.constexpr,
+    LOAD_WARPS: gl.constexpr, LOAD_REGS: gl.constexpr, EPI_WARPS: gl.constexpr, EPI_REGS: gl.constexpr,
+    top_k: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr,
+):
+    BM: gl.constexpr = BLOCK_SIZE_M
+    BN: gl.constexpr = BLOCK_SIZE_N
+    KB: gl.constexpr = BLOCK_SIZE_K // 2
+    nk = K // BLOCK_SIZE_K
+    x_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BM, KB], gl.NVMMASharedLayout.get_default_for([BM, KB], gl.uint8))
+    w_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BN, KB], w_desc.layout)
+    sw_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT)
+    sx_bufs = gl.allocate_shared_memory(gl.uint8, [2, BM, 16], SCALE_SMEM_LAYOUT)
+    # n-contiguous and swizzled, so the epilogue's 8-row reads and the MMA-layout stores spread over banks
+    h_buf = gl.allocate_shared_memory(gl.bfloat16, [BN // 2, BM], gl.SwizzledSharedLayout(8, 1, 8, [0, 1]))
+    ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout())
+    s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout())
+    acc_ready = gl.allocate_shared_memory(gl.int64, [1, 1], mbarrier.MBarrierLayout())
+    acc_empty = gl.allocate_shared_memory(gl.int64, [1, 1], mbarrier.MBarrierLayout())
+    for i in gl.static_range(STAGES):
+        mbarrier.init(ready.index(i), count=1)
+        mbarrier.init(empty.index(i), count=1)
+    for i in gl.static_range(2):
+        mbarrier.init(s_ready.index(i), count=1)
+        mbarrier.init(s_empty.index(i), count=1)
+    mbarrier.init(acc_ready.index(0), count=1)
+    mbarrier.init(acc_empty.index(0), count=1)
+
+    NW: gl.constexpr = gl.num_warps()
+    mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N],
+                                                  instr_shape=[16, 8])
+    gl.warp_specialize([
+        (_down_mma_partition, ((x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf), (ready, empty, s_ready, s_empty, acc_ready, acc_empty),
+                               sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, wg_ptr, g_in_ptr,
+                               topk_weights_ptr, N, nk, num_valid_tokens, stride_ge, stride_tw, mma, wsl, xsl, K_TAIL,
+                               N_GROUP, GROUP_SIZE_M)),
+        (_down_load_partition, (xq_ptr, xs_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, w_desc,
+                                ws_ptr, (x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf),
+                                (ready, empty, s_ready, s_empty, acc_ready, acc_empty), N, nk, num_valid_tokens,
+                                stride_xq, ks4, n_rows, K_TAIL, N_GROUP, GROUP_SIZE_M)),
+        (_down_epi_partition, (sorted_token_ids_ptr, num_tokens_post_padded_ptr, out_ptr,
+                               (x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf), (ready, empty, s_ready, s_empty, acc_ready, acc_empty),
+                               N, num_valid_tokens, stride_out, top_k, N_GROUP, GROUP_SIZE_M)),
+    ], [LOAD_WARPS, EPI_WARPS], [LOAD_REGS, EPI_REGS])
+
+
+def supports(H: int, I: int) -> bool:
+    """Whether the GEMMs take a layer of this hidden / intermediate size."""
+    # K is a whole number of 64-wide k blocks for both GEMMs; the gated tile pairs 64 gate rows with 64 up rows at least
+    return H % 64 == 0 and I % 64 == 0
+
+
+def _weight_desc(packed, box_n, BK):
+    S, N, KB = packed.shape
+    layout = gl.NVMMASharedLayout.get_default_for([box_n, BK // 2], gl.uint8)
+    return TensorDescriptor.from_tensor(packed.view(S * N, KB), [box_n, BK // 2], layout)
+
+
+def _scale_layouts(cfg, w_rows):
+    nw, wn = cfg["num_warps"], cfg["MMA_WARPS_N"]
+    kg = cfg["BLOCK_SIZE_K"] // 16
+    return scale_layout(0, w_rows, kg, [wn, nw // wn]), scale_layout(1, cfg["BLOCK_SIZE_M"], kg, [wn, nw // wn])
+
+
+def gate_up_gemm(xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded,
+                 num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg):
+    """q, qs = NVFP4 under g_out of act(gate, up), where [gate; up] = (xq * xs * g_in)[route // top_k] W4A4 packed[expert]^T * glob."""
+    S, N, KB = packed.shape
+    K = KB * 2
+    BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"]
+    assert K % BK == 0 and (N // 2) % (BN // 2) == 0 and xs.stride(0) == K // 16
+    assert q.stride(0) == N // 4 and qs.stride(0) == N // 32
+    wsl, xsl = _scale_layouts(cfg, BN)
+    grid = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN)
+    _w4a4_moe_kernel[(grid,)](
+        xq, xs.view(torch.uint8), g_in, _weight_desc(packed, BN // 2, BK), scale.view(torch.uint8), glob, None, None,
+        sorted_ids, expert_ids, num_tokens_post_padded, q, qs.view(torch.uint8), g_out, act_alpha, act_limit,
+        N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), 0, 0,
+        K_TAIL=K % 256 != 0, GATED=True, ACT=ACTIVATIONS[activation], top_k=top_k, wsl=wsl, xsl=xsl, **cfg,
+    )
+
+
+def down_gemm(xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+              num_valid_tokens, top_k, cfg):
+    """out[route // top_k] += (xq * xs * g_in)[route] W4A4 packed[expert]^T * glob * topk_weights_flat[route]; out starts zeroed."""
+    S, N, KB = packed.shape
+    K = KB * 2
+    cfg = dict(cfg)
+    persistent = cfg.pop("PERSISTENT")
+    BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"]
+    assert K % BK == 0 and xs.stride(0) == K // 16
+    args = (xq, xs.view(torch.uint8), g_in, _weight_desc(packed, BN, BK), scale.view(torch.uint8), glob, out,
+            topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded)
+    if persistent:
+        wsl, xsl = _scale_layouts(cfg, BN // 2)
+        grid = torch.cuda.get_device_properties(out.device).multi_processor_count
+        _w4a4_down_persistent_kernel[(grid,)](
+            *args, N, K, num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), out.stride(0),
+            topk_weights_flat.stride(0), K_TAIL=K % 256 != 0, top_k=top_k, wsl=wsl, xsl=xsl, **cfg,
+        )
+        return
+    wsl, xsl = _scale_layouts(cfg, BN)
+    grid = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN)
+    _w4a4_moe_kernel[(grid,)](
+        *args, None, None, None, 1.0, 1.0,
+        N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), out.stride(0),
+        topk_weights_flat.stride(0), K_TAIL=K % 256 != 0, GATED=False, ACT=0, top_k=top_k, wsl=wsl, xsl=xsl, **cfg,
+    )
diff --git a/python/freetoken/kernel/gluon/nvfp4/utils.py b/python/freetoken/kernel/gluon/nvfp4/utils.py
index 430c146d49..11426ec569 100644
--- a/python/freetoken/kernel/gluon/nvfp4/utils.py
+++ b/python/freetoken/kernel/gluon/nvfp4/utils.py
@@ -2,6 +2,11 @@
 
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language import _core as ttgl
+from triton.experimental.gluon.language._layouts import (
+    DotOperandLayout,
+    NVMMADistributedLayout,
+)
 
 
 @gluon.jit
@@ -41,3 +46,37 @@ def e2m1x8_bf16x2(q, s2):
     w2 = bf16x2_mul(prmt(o3, o2, 0x7632), s2)
     w3 = bf16x2_mul(prmt(o1, o0, 0x7632), s2)
     return w0, w1, w2, w3
+
+
+def scale_layout(op_idx: int, rows: int, kgroups: int, mma_warps: list[int]) -> gl.DistributedLinearLayout:
+    """Layout of a [rows, K / 16] e4m3 scale tile for operand op_idx of an NVFP4 mma_v2; mirrors Triton's getSM120DotScaledScaleLayout."""
+    def bits(n):
+        return n.bit_length() - 1
+
+    w0, w1 = mma_warps
+    reg = [[0, 1 << i] for i in range(bits(kgroups))]
+    if op_idx == 0:
+        lane = [[8, 0], [0, 0], [1, 0], [2, 0], [4, 0]]
+        warp = [[0, 0]] * bits(w1) + [[16 << i, 0] for i in range(bits(w0))]
+        cover = 16 * w0
+    else:
+        lane = [[0, 0], [0, 0], [1, 0], [2, 0], [4, 0]]
+        warp = [[8 << i, 0] for i in range(bits(w1))] + [[0, 0]] * bits(w0)
+        cover = 8 * w1
+    reg += [[cover << i, 0] for i in range(max(0, bits(rows) - bits(cover)))]
+    # bases past the tile's rows broadcast
+    lane = [[b[0] if b[0] < rows else 0, b[1]] for b in lane]
+    warp = [[b[0] if b[0] < rows else 0, b[1]] for b in warp]
+    return gl.DistributedLinearLayout(reg, lane, warp, [], [rows, kgroups])
+
+
+@ttgl.builtin
+def mma_nvf4(a, a_scale, b, b_scale, acc, _semantic=None):
+    """acc += (a * a_scale) @ (b * b_scale) on sm_120 (mma kind::mxf4nvf4, scale_vec::4X).
+
+    a [M, K/2] and b [K/2, N] are packed e2m1 in DotOperandLayout(k_width=4) of acc's mma_v2 layout; the e4m3 scales are [M, K/16] / [N, K/16] in scale_layout.
+    """
+    assert isinstance(acc.type.layout, NVMMADistributedLayout) and acc.type.layout.version == [2, 0]
+    assert isinstance(a.type.layout, DotOperandLayout) and isinstance(b.type.layout, DotOperandLayout)
+    out = _semantic.dot_scaled(a, a_scale, "e2m1", b, b_scale, "e2m1", acc, False, True, True, ttgl.float32)
+    return ttgl.tensor(out.handle, acc.type)
diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py
index 0aa494222c..df89410726 100644
--- a/python/freetoken/kernel/nvfp4_moe.py
+++ b/python/freetoken/kernel/nvfp4_moe.py
@@ -1,6 +1,6 @@
-"""Kernel dispatch for the inline-dequant NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch.
+"""Kernel dispatch for the NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch.
 
-The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`);
+The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120_w4a4`);
 :mod:`freetoken.moe.fused_nvfp4` strings the GEMMs into the MoE layer.
 """
 
@@ -388,3 +388,73 @@ def prefill_gemm(
         SCALE_WORDS=scale_words,
         **cfg,
     )
+
+
+@functools.lru_cache(maxsize=None)
+def _w4a4_prefill(device_index: int):
+    """The Gluon module of the W4A4 prefill GEMMs on this GPU, or None."""
+    from freetoken.kernel import gluon
+
+    if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available():
+        return None
+    from freetoken.kernel.gluon.nvfp4.moe import sm120_w4a4
+
+    return sm120_w4a4
+
+
+def w4a4_available(device_index: int) -> bool:
+    return _w4a4_prefill(device_index) is not None
+
+
+# Below this many tokens per expert both prefills stream the weights at the same speed (RTX PRO 6000 and RTX 5090 sweeps over the
+# NVFP4 MoE shapes); W4A16 keeps the exact activations there
+_W4A4_MIN_TOKENS_PER_EXPERT = 4
+# (min rho, tokens per expert where BLOCK_SIZE_M becomes 128), from the same sweeps
+_W4A4_BLOCK_M_STEPS = ((3.0, 128), (0.0, 64))
+
+
+@functools.lru_cache(maxsize=None)
+def _w4a4_device(device_index: int) -> Tuple[float, int]:
+    props = torch.cuda.get_device_properties(device_index)
+    rho = _tensor_rho(props.name, props.major, props.minor)
+    return next(t for min_rho, t in _W4A4_BLOCK_M_STEPS if rho >= min_rho), props.L2_cache_size
+
+
+def w4a4_prefill_config(M: int, top_k: int, num_experts: int, H: int, I: int, activation: str, dtype: torch.dtype,
+                        device_index: int) -> Dict[str, Any] | None:
+    """Tiles of the W4A4 prefill GEMMs, or None where the W4A16 prefill runs instead."""
+    w4a4 = _w4a4_prefill(device_index)
+    if w4a4 is None or dtype != torch.bfloat16 or activation not in w4a4.ACTIVATIONS or not w4a4.supports(H, I):
+        return None
+    to128, l2_bytes = _w4a4_device(device_index)
+    per_expert = M * top_k / num_experts
+    if per_expert < _W4A4_MIN_TOKENS_PER_EXPERT:
+        return None
+    block_m = 128 if per_expert >= to128 else 64
+    common = dict(BLOCK_SIZE_M=block_m, GROUP_SIZE_M=8, MMA_WARPS_N=4, LOAD_WARPS=4, num_warps=8)
+    # the gated tile holds BN / 2 gate rows and their up rows, so I must split into whole half tiles
+    gate_up = dict(common, BLOCK_SIZE_N=256 if I % 128 == 0 else 128, BLOCK_SIZE_K=128 if H % 128 == 0 else 64,
+                   N_GROUP=0, STAGES=3, LOAD_REGS=88)
+    if I < 1024:
+        # short K: the persistent kernel overlaps a tile's atomics with the next mainloop; N strips keep the [M, H] output
+        # in L2 once it outgrows half of it
+        down = dict(common, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, STAGES=4, N_GROUP=2 if M * H * 2 > l2_bytes // 2 else 0,
+                    LOAD_REGS=56, EPI_WARPS=4, EPI_REGS=72, PERSISTENT=True)
+    else:
+        down = dict(common, BLOCK_SIZE_N=256 if H % 256 == 0 else 128, BLOCK_SIZE_K=128 if I % 128 == 0 else 64,
+                    N_GROUP=0, STAGES=3, LOAD_REGS=88, PERSISTENT=False)
+    return dict(BLOCK_SIZE_M=block_m, gate_up=gate_up, down=down)
+
+
+def w4a4_gate_up_gemm(xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded,
+                      num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg):
+    _w4a4_prefill(xq.device.index).gate_up_gemm(
+        xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded,
+        num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg)
+
+
+def w4a4_down_gemm(xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids,
+                   num_tokens_post_padded, num_valid_tokens, top_k, cfg):
+    _w4a4_prefill(xq.device.index).down_gemm(
+        xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+        num_valid_tokens, top_k, cfg)
diff --git a/python/freetoken/kernel/triton/nvfp4_act_quant.py b/python/freetoken/kernel/triton/nvfp4_act_quant.py
new file mode 100644
index 0000000000..b52cfbe898
--- /dev/null
+++ b/python/freetoken/kernel/triton/nvfp4_act_quant.py
@@ -0,0 +1,42 @@
+"""NVFP4 activation quantization under a static global scale: x ~= e2m1 * e4m3 block scale (per 16) * g."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _e2m1x2(hi, lo):
+    # sm_100+ / sm_120a; cvt packs its first source into the high nibble
+    return tl.inline_asm_elementwise(
+        "{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $1, $2; cvt.u16.u8 $0, t; }",
+        "=h,r,r", [hi, lo], dtype=tl.uint16, is_pure=True, pack=1)
+
+
+@triton.jit
+def _quant_kernel(x_ptr, q_ptr, s_ptr, g_ptr, K, stride_x, BLOCK_K: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK_K)
+    x = tl.load(x_ptr + row * stride_x + offs, mask=offs < K, other=0.0).to(tl.float32)
+    g = tl.load(g_ptr)
+    xb = tl.reshape(x, [BLOCK_K // 16, 16])
+    s = (tl.max(tl.abs(xb), 1) / (6.0 * g)).to(tl.float8e4nv)
+    sf = s.to(tl.float32) * g
+    inv = tl.where(sf > 0, 1.0 / sf, 0.0)
+    lo, hi = tl.split(tl.reshape(xb * inv[:, None], [BLOCK_K // 2, 2]))
+    oq = tl.arange(0, BLOCK_K // 2)
+    tl.store(q_ptr + row * (K // 2) + oq, _e2m1x2(hi, lo).to(tl.uint8), mask=oq < K // 2)
+    os_ = tl.arange(0, BLOCK_K // 16)
+    tl.store(s_ptr + row * (K // 16) + os_, s.to(tl.uint8, bitcast=True), mask=os_ < K // 16)
+
+
+def quantize_nvfp4(x: torch.Tensor, g: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """x [R, K] -> (codes [R, K / 2] uint8, low nibble first; block scales [R, K / 16] e4m3) under the fp32 global g [1]; K % 16 == 0."""
+    R, K = x.shape
+    q = torch.empty(R, K // 2, dtype=torch.uint8, device=x.device)
+    s = torch.empty(R, K // 16, dtype=torch.uint8, device=x.device)
+    if R:
+        _quant_kernel[(R,)](x, q, s, g, K, x.stride(0), BLOCK_K=triton.next_power_of_2(K), num_warps=4)
+    return q, s.view(torch.float8_e4m3fn)
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index 1cd1a1ac73..0dce399561 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -78,6 +78,56 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         return fused_experts_decode_nvfp4_marlin(x, *banks, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, alpha, limit)
 
 
+class W4a4Nvfp4MoEKernel(TritonNvfp4MoEKernel):
+    """The Triton kernels' banks and decode, with the prefill GEMMs on NVFP4 activations (W4A4) where the GPU and the batch favour it.
+
+    Needs a checkpoint whose experts declare static activation scales; the GPU-resident alphas carry them per expert (gate_up_alpha: the larger of gate's and up's input_scale).
+    """
+
+    name = "w4a4"
+
+    def unusable_reason(self, cfg: MoEConfig) -> str | None:
+        reason = super().unusable_reason(cfg)
+        if reason:
+            return reason
+        if cfg.scheme is None or not cfg.scheme.has("input_scale"):
+            return "the experts declare no static activation scale (W4A16 checkpoint)"
+        if cfg.apply_router_weight_on_input:
+            return "the routing weight scales the input, which is quantized once per token"
+        return None
+
+    def worth_it(self, cfg: MoEConfig) -> bool:
+        from freetoken.kernel import nvfp4_moe
+
+        return not backend.is_rocm() and torch.cuda.is_available() and nvfp4_moe.w4a4_available(torch.cuda.current_device())
+
+    def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
+        return {
+            **super().layout(cfg),
+            "gate_up_alpha": BankSpec((), torch.float32, resident=True),
+            "down_alpha": BankSpec((), torch.float32, resident=True),
+        }
+
+    def pack(self, pieces, cfg: MoEConfig, out):
+        super().pack(pieces, cfg, out)
+        if not all(f"{p}_input" in pieces for p in ("gate", "up", "down")):
+            raise ValueError("w4a4 nvfp4 MoE kernel: the checkpoint has no expert input_scale tensors")
+        gate, up, down = (pieces[f"{p}_input"].reshape(len(pieces[f"{p}_input"]), -1)[:, 0].float() for p in ("gate", "up", "down"))
+        return {"gate_up_alpha": torch.maximum(gate, up), "down_alpha": down}
+
+    def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefill: bool):
+        if not is_prefill:
+            return super().apply(layer, x, topk_weights, topk_ids, view, is_prefill=False)
+        from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4_w4a4
+
+        t = view.tensors
+        assert view.alphas is not None
+        return fused_experts_nvfp4_w4a4(
+            x, t["gate_up"], t["gate_up_scale"], t["gate_up_global"], t["down"], t["down_scale"], t["down_global"],
+            view.alphas[0], view.alphas[1], topk_weights, topk_ids, view.n, layer.activation, float(layer.alpha), limit_or_inf(layer),
+        )
+
+
 # ---------------------------------------------------------------------------
 # Marlin: vLLM's fused Marlin MoE (W4A16, dequant in kernel) over Marlin-tiled banks.
 # The Marlin kernels come from vLLM (not sgl-kernel) deliberately: vLLM ships the
@@ -570,4 +620,4 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 
 @register_method(QuantKind.NVFP4, LayerKind.MOE)
 class Nvfp4MoEMethod(MoEMethod):
-    candidates = (TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel)
+    candidates = (W4a4Nvfp4MoEKernel, TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel)
diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py
index 04da050919..e30429c06a 100644
--- a/python/freetoken/models/nvfp4_banks.py
+++ b/python/freetoken/models/nvfp4_banks.py
@@ -39,6 +39,28 @@ def _ingest_global(spec: "Nvfp4ExpertSourceSpec", tensor: torch.Tensor) -> torch
     return tensor.to(torch.float16)
 
 
+def _activation_scale():
+    """The checkpoint tensor of the NVFP4 static activation scale under the installed dialect, or None."""
+    from freetoken.layers.quantization import QuantKind, get_quant_config
+
+    try:
+        quant = get_quant_config()
+    except RuntimeError:
+        return None
+    if quant is None or QuantKind.NVFP4 not in quant.STORAGE:
+        return None
+    return quant.stored_tensors(QuantKind.NVFP4).get("input_scale")
+
+
+def _ingest(spec: "Nvfp4ExpertSourceSpec", act, role: str, tensor: torch.Tensor) -> torch.Tensor:
+    if role.endswith("_global"):
+        return _ingest_global(spec, tensor)
+    if role.endswith("_input"):
+        tensor = tensor.float()
+        return 1.0 / tensor if act.reciprocal else tensor
+    return tensor
+
+
 def _num_moe_layers(config) -> int:
     value = getattr(config, "num_moe_layers", None)
     if value is not None:
@@ -77,7 +99,8 @@ def iter_nvfp4_expert_pieces(
 ):
     """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale``
     (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects,
-    fp16) companions, straight from the safetensors shards.
+    fp16) companions, straight from the safetensors shards. When every expert also stores the
+    dialect's static activation scale, the pieces carry it as fp32 ``_input`` (dequant side).
 
     Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either
     way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land.
@@ -92,6 +115,7 @@ def iter_nvfp4_expert_pieces(
     weight_map = safetensors_weight_map(folder)
 
     keep = None if layers is None else set(layers)
+    act = _activation_scale()
     wanted: dict[str, tuple[int, int, str]] = {}
     for name in weight_map:
         match = spec.key_pattern.match(name)
@@ -107,9 +131,20 @@ def iter_nvfp4_expert_pieces(
         if kind not in ("weight", "weight_scale", "weight_scale_2"):
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
         wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind))
-    expected = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts * 9
-    if len(wanted) != expected:
-        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}")
+        if kind == "weight" and act is not None:
+            act_name = name[: match.start("kind")] + act.name
+            if act_name in weight_map:
+                wanted[act_name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + "_input")
+    experts = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts
+    acts = [name for name, (_, _, role) in wanted.items() if role.endswith("_input")]
+    # a partial set serves no W4A4 kernel; the pieces then carry none
+    if len(acts) != experts * 3:
+        for name in acts:
+            del wanted[name]
+        acts = []
+    per_expert = 12 if acts else 9
+    if len(wanted) != experts * per_expert:
+        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * per_expert}")
 
     def _serial():
         by_shard: dict[str, list[str]] = collections.defaultdict(list)
@@ -121,21 +156,16 @@ def _serial():
             drop(path)
             with safetensors.safe_open(path, framework="pt", device="cpu") as f:
                 for name in by_shard[shard]:
-                    tensor = f.get_tensor(name)
-                    if wanted[name][2].endswith("_global"):
-                        tensor = _ingest_global(spec, tensor)
-                    yield name, tensor
+                    yield name, _ingest(spec, act, wanted[name][2], f.get_tensor(name))
             drop(path)
 
     def _parallel():
         from freetoken.models.weight import iter_expert_tensors_parallel
 
         for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk):
-            if wanted[name][2].endswith("_global"):
-                tensor = _ingest_global(spec, tensor)
-            yield name, tensor
+            yield name, _ingest(spec, act, wanted[name][2], tensor)
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert)
 
 
 __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index 917d84bbcb..74c4a41b0f 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -151,7 +151,7 @@ def per_expert_pieces(
     """Group per-expert tensors into one piece per expert, in whatever order they arrive.
 
     ``locate(name)`` -> ``(bank_layer, expert, role)`` or None to skip; a role tensor is one
-    expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals may be scalars.
+    expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals and activation scales may be scalars.
     """
     pending: dict[tuple[int, int], dict[str, torch.Tensor]] = {}
     for name, tensor in tensors:
@@ -160,7 +160,7 @@ def per_expert_pieces(
             continue
         bank_layer, expert, role = where
         piece = pending.setdefault((bank_layer, expert), {})
-        piece[role] = tensor.reshape(1, -1) if role.endswith("_global") else tensor.unsqueeze(0)
+        piece[role] = tensor.reshape(1, -1) if role.endswith(("_global", "_input")) else tensor.unsqueeze(0)
         if len(piece) == tensors_per_expert:
             del pending[(bank_layer, expert)]
             yield bank_layer, expert, expert + 1, piece
diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py
index 197ef54727..6bedd57b12 100644
--- a/python/freetoken/moe/fused_nvfp4.py
+++ b/python/freetoken/moe/fused_nvfp4.py
@@ -10,6 +10,7 @@
 import torch
 
 from freetoken.kernel import moe_sum_reduce_triton, nvfp4_moe
+from freetoken.kernel.triton.nvfp4_act_quant import quantize_nvfp4
 from freetoken.layers import gated_act_and_mul
 from freetoken.moe.fused import moe_align_block_size
 
@@ -156,8 +157,60 @@ def fused_experts_nvfp4(
     return out
 
 
+def fused_experts_nvfp4_w4a4(
+    hidden_states: torch.Tensor,
+    gate_up_packed: torch.Tensor,
+    gate_up_scale: torch.Tensor,
+    gate_up_global: torch.Tensor,
+    down_packed: torch.Tensor,
+    down_scale: torch.Tensor,
+    down_global: torch.Tensor,
+    gate_up_input_scale: torch.Tensor,
+    down_input_scale: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    activation: str = "silu",
+    act_alpha: float = 1.702,
+    act_limit: float = 7.0,
+) -> torch.Tensor:
+    """Prefill NVFP4 MoE with NVFP4 activations (W4A4) where the GPU and the batch favour it, else :func:`fused_experts_nvfp4`.
+
+    ``*_input_scale`` are the per-expert static activation scales of the ``num_experts`` banks; each GEMM input is quantized under their max, as vLLM does.
+    """
+    M, H = hidden_states.shape
+    top_k = topk_ids.shape[1]
+    inter = gate_up_packed.shape[1] // 2
+    cfg = nvfp4_moe.w4a4_prefill_config(M, top_k, num_experts, H, inter, activation, hidden_states.dtype, hidden_states.device.index)
+    if cfg is None:
+        return fused_experts_nvfp4(
+            hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global,
+            topk_weights, topk_ids, num_experts, activation, False, act_alpha, act_limit,
+        )
+    dev = hidden_states.device
+    g = torch.stack((gate_up_input_scale[:num_experts], down_input_scale[:num_experts])).float().amax(dim=1)
+    sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
+    tw = topk_weights.reshape(-1).contiguous()
+    num_valid = topk_ids.numel()
+
+    xq, xs = quantize_nvfp4(hidden_states, g[0:1])
+    q = torch.empty((num_valid, inter // 2), device=dev, dtype=torch.uint8)
+    qs = torch.empty((num_valid, inter // 16), device=dev, dtype=torch.float8_e4m3fn)
+    nvfp4_moe.w4a4_gate_up_gemm(
+        xq, xs, g[0:1], gate_up_packed, gate_up_scale, gate_up_global, q, qs, g[1:2], sorted_ids, expert_ids, ntpp,
+        num_valid, top_k, activation, act_alpha, act_limit, cfg["gate_up"],
+    )
+    out = torch.zeros_like(hidden_states)
+    nvfp4_moe.w4a4_down_gemm(
+        q, qs, g[1:2], down_packed, down_scale, down_global, out, tw, sorted_ids, expert_ids, ntpp, num_valid, top_k,
+        cfg["down"],
+    )
+    return out
+
+
 __all__ = [
     "fused_experts_decode_nvfp4_marlin",
     "fused_experts_decode_nvfp4_serial",
     "fused_experts_nvfp4",
+    "fused_experts_nvfp4_w4a4",
 ]
diff --git a/python/freetoken/moe/legacy_format.py b/python/freetoken/moe/legacy_format.py
index 973a3dba52..8f347e5c30 100644
--- a/python/freetoken/moe/legacy_format.py
+++ b/python/freetoken/moe/legacy_format.py
@@ -30,12 +30,16 @@ def legacy_bank_names(quant_format: str) -> dict[str, str]:
     (QuantKind.NONE, "fused"): "bf16",
     (QuantKind.FP8_BLOCK, "triton"): "fp8_block",
     (QuantKind.NVFP4, "triton"): "nvfp4",
+    # same banks as triton, plus the resident input scales; a "nvfp4" file reads back as triton
+    (QuantKind.NVFP4, "w4a4"): "nvfp4",
     (QuantKind.NVFP4, "marlin"): "nvfp4_marlin",
     (QuantKind.NVFP4, "b12x"): "nvfp4_b12x",
     (QuantKind.MXFP4, "triton_gptoss"): "mxfp4_triton",
     (QuantKind.MXFP4, "triton"): "ds_fp4",
 }
-_KIND_KERNEL = {fmt: kk for kk, fmt in LEGACY_FORMAT.items()}
+_KIND_KERNEL: dict[str, tuple[QuantKind, str]] = {}
+for _kk, _fmt in LEGACY_FORMAT.items():
+    _KIND_KERNEL.setdefault(_fmt, _kk)
 
 
 def legacy_format_for(kind: QuantKind, kernel: str) -> str:
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 29c042ba4a..fd98c2ee94 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -10,13 +10,15 @@
     baseline-kernel equality guard. This is the path used on sm_120 + CUDA 12.x.
   - Marlin (sm_80..sm_99, e.g. H100): prefill + decode + overlap.
   - b12x (sm_120 + CUDA>=13): pack + fused decode, skipped where the kernel cannot run.
-  - Gluon (sm_120): the prefill GEMMs it takes over from the Triton kernel, against that kernel.
+  - Gluon (sm_120): the prefill GEMMs it takes over from the Triton kernel, against that kernel, and the W4A4 prefill
+    against a torch emulation of its NVFP4 activation rounding.
 Banks are built the way the loader builds them: the bound kernel object packs native pieces.
 """
 
 from __future__ import annotations
 
 import importlib.util
+import math
 
 import pytest
 import torch
@@ -413,6 +415,100 @@ def counted(*args):
     torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max()))
 
 
+def _round_e2m1(v: torch.Tensor) -> torch.Tensor:
+    """Round to the nearest e2m1 value, ties to the even code, saturating at 6 (cvt.rn.satfinite)."""
+    grid = _E2M1[:8].to(v.device)
+    a = v.abs().clamp(max=6.0)
+    hi = torch.searchsorted(grid, a.contiguous()).clamp(max=7)
+    lo = (hi - 1).clamp(min=0)
+    d_lo, d_hi = a - grid[lo], grid[hi] - a
+    code = torch.where(d_lo < d_hi, lo, torch.where(d_hi < d_lo, hi, torch.where(lo % 2 == 0, lo, hi)))
+    return torch.copysign(grid[code], v)
+
+
+def _nvfp4_qdq(x: torch.Tensor, g: float) -> torch.Tensor:
+    """x [R, K] through NVFP4 under the static global g, back to fp32."""
+    xb = x.float().view(x.shape[0], -1, 16)
+    # fp32 scalar math like the kernels: a one-ulp difference in 6 * g can flip an e4m3 block scale
+    g = torch.tensor(g, dtype=torch.float32, device=x.device)
+    sf = ((xb.abs().amax(-1) / (6.0 * g)).to(torch.float8_e4m3fn).float() * g).unsqueeze(-1)
+    inv = torch.where(sf > 0, 1.0 / sf, torch.zeros_like(sf))
+    return (_round_e2m1(xb * inv) * sf).view_as(x)
+
+
+def _gated_ref(h: torch.Tensor, activation: str, inter: int, alpha: float, limit: float) -> torch.Tensor:
+    gate, up = h[:, :inter], h[:, inter:]
+    if activation == "gelu_tanh":
+        return torch.nn.functional.gelu(gate, approximate="tanh") * up
+    if activation in ("swiglu_clamp", "swigluoai"):
+        gate = gate.clamp(max=limit)
+        return gate * torch.sigmoid(alpha * gate) * (up.clamp(-limit, limit) + (activation == "swigluoai"))
+    return torch.nn.functional.silu(gate) * up
+
+
+@cuda
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0), reason="needs sm_120")
+@pytest.mark.parametrize("hidden_size,inter,activation,tokens", [
+    (256, 128, "silu", 256), (320, 192, "gelu_tanh", 1024), (512, 1024, "swiglu_clamp", 256), (256, 128, "swigluoai", 256),
+])
+def test_w4a4_prefill_matches_nvfp4_activation_reference(hidden_size, inter, activation, tokens):
+    """The w4a4 kernel from checkpoint pieces to the prefill output, against a torch emulation of the same NVFP4 activation rounding.
+
+    Covers 256- / 128-row gated tiles, K % 256 != 0 in both GEMMs, the persistent down GEMM (I < 1024) and the other one, BLOCK_SIZE_M 64 and 128.
+    """
+    from types import SimpleNamespace
+
+    from freetoken.kernel import nvfp4_moe
+    from freetoken.layers.quantization.moe.base import ExpertView, MoEConfig
+    from freetoken.layers.quantization.moe.nvfp4 import W4a4Nvfp4MoEKernel
+
+    device = torch.device("cuda")
+    if not nvfp4_moe.w4a4_available(device.index or 0):
+        pytest.skip("the W4A4 Gluon kernels need Triton's Gluon on sm_120")
+    alpha, limit = (1.5, 2.0) if activation in ("swiglu_clamp", "swigluoai") else (1.0, None)
+    g = torch.Generator(device=device).manual_seed(0)
+
+    def bank(n, k):
+        packed = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=device, generator=g)
+        scale = (torch.rand(E, n, k // 16, device=device, generator=g) + 0.25).to(torch.float8_e4m3fn)
+        return packed, scale, (torch.rand(E, n, device=device, generator=g) * 0.02 + 0.01).to(torch.float16)
+
+    gu, gus, gug = bank(2 * inter, hidden_size)
+    dn, dns, dng = bank(hidden_size, inter)
+    hidden = torch.randn(tokens, hidden_size, dtype=torch.bfloat16, device=device, generator=g)
+    topk_ids = torch.randint(0, E, (tokens, TOPK), dtype=torch.int32, device=device, generator=g)
+    topk_weights = torch.rand(tokens, TOPK, device=device, generator=g)
+
+    # static scales near a calibration's amax / (6 * 448), off by an irrational factor: with amax / 2688 itself many values sit exactly on an
+    # e2m1 midpoint, where the kernel's approximate fp32 reciprocal and torch's exact one round apart; per expert, below the max for some experts
+    g1 = float(hidden.float().abs().max()) / 2688 * math.pi / 3
+    xq = _nvfp4_qdq(hidden, g1)
+    w_gu = torch.stack([_dequant_ref(gu[e], gus[e], gug[e]) for e in range(E)])
+    w_dn = torch.stack([_dequant_ref(dn[e], dns[e], dng[e]) for e in range(E)])
+    routes = topk_ids.reshape(-1).long()
+    h = torch.bmm(w_gu[routes], xq.repeat_interleave(TOPK, 0).unsqueeze(-1)).squeeze(-1).bfloat16().float()
+    act = _gated_ref(h, activation, inter, alpha, limit or 0.0).bfloat16().float()
+    g2 = float(act.abs().max()) / 2688 * math.pi / 3
+    ref = torch.bmm(w_dn[routes], _nvfp4_qdq(act, g2).unsqueeze(-1)).squeeze(-1)
+    ref = (ref * topk_weights.reshape(-1, 1)).view(tokens, TOPK, hidden_size).sum(1)
+
+    def input_scale(top):
+        return torch.full((E, 1), top, device=device).index_fill_(0, torch.tensor([1, 2], device=device), top / 2)
+
+    pieces = {"gate_up": gu, "gate_up_scale": gus, "gate_up_global": gug, "down": dn, "down_scale": dns, "down_global": dng,
+              "gate_input": input_scale(g1), "up_input": input_scale(g1).roll(1, 0), "down_input": input_scale(g2)}
+    cfg = MoEConfig(num_experts=E, hidden=hidden_size, intermediate=inter, top_k=TOPK, activation=activation, alpha=alpha, limit=limit)
+    kernel = W4a4Nvfp4MoEKernel()
+    layout = kernel.layout(cfg)
+    banks = {role: torch.empty((E, *spec.shape), dtype=spec.dtype, device=device) for role, spec in layout.items() if not spec.resident}
+    alphas = kernel.pack(pieces, cfg, banks)
+    view = ExpertView(banks, n=E, alphas=(alphas["gate_up_alpha"], alphas["down_alpha"]))
+    layer = SimpleNamespace(activation=activation, alpha=alpha, limit=limit, apply_router_weight_on_input=False)
+    out = kernel.apply(layer, hidden, topk_weights, topk_ids, view, is_prefill=True)
+    assert torch.equal(alphas["gate_up_alpha"], torch.maximum(pieces["gate_input"], pieces["up_input"]).squeeze(1))
+    torch.testing.assert_close(out.float(), ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max()))
+
+
 def _triton_cache(device, *, cache_size=S, prefill_overlap=False):
     """Native 6-bank NVFP4 cache (no repack), consumed directly by the Triton kernels.
     The banks are not transformed, so ``sources`` doubles as the dequant reference."""

From 30dd715a9328147141480aa1bd7a4230b147ec0d Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 18 Oct 2026 23:16:37 -0700
Subject: [PATCH 125/249] fix(muse_glimmer)!: keep the attention gate out of
 the qkv fusion (#741)

The NVFP4 release quantizes the attention gate with a different weight global scale than q/k/v, so
one fused weight cannot carry a single global for all four parts.

BREAKING CHANGE: an FTW checkpoint converted with the fused qkvg_proj no longer loads; reconvert it
with `ft checkpoint`.
---
 docs/models.md                                |  2 +
 .../models/muse_glimmer/attention.py          | 29 ++++++++------
 python/freetoken/models/muse_glimmer/gguf.py  |  3 +-
 .../freetoken/models/muse_glimmer/weight.py   | 13 +++----
 python/freetoken/models/register.py           |  8 +---
 tests/models/test_muse_glimmer.py             | 38 ++++++++++---------
 tests/models/test_quant_config.py             |  3 +-
 7 files changed, 50 insertions(+), 46 deletions(-)

diff --git a/docs/models.md b/docs/models.md
index 4418d92b2b..fd124ea3d2 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -84,6 +84,8 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   its experts as dense weights and no longer loads; reconvert it with `ft checkpoint`.
 - FTW files converted by builds before the quantization refactor may fail to load;
   see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool.
+- Muse-Glimmer keeps the attention gate as its own linear next to `qkv_proj`; an FTW converted
+  when the gate was fused into `qkvg_proj` no longer loads, reconvert it with `ft checkpoint`.
 - An FTW converted before its family served images holds no vision encoder: `ft serve`
   refuses it unless started with `--text-model-only` (or `--mm-disable vision`); reconvert it
   with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md).
diff --git a/python/freetoken/models/muse_glimmer/attention.py b/python/freetoken/models/muse_glimmer/attention.py
index 7c814d4c1b..a4af6e5210 100644
--- a/python/freetoken/models/muse_glimmer/attention.py
+++ b/python/freetoken/models/muse_glimmer/attention.py
@@ -18,16 +18,16 @@
 class MuseGlimmerAttention(BaseOP):
     """Gated GQA for one sliding-window or full-attention (NoPE) layer.
 
-        q, k, v, gate = split(qkvg_proj(x))
+        q, k, v = split(qkv_proj(x))
         q = qk_norm(q); k = qk_norm(k)          # weightless per-head RMSNorm
         q, k = rope(q, k)                        # sliding layers only; full layers are NoPE
         attn = paged_attention(q, k, v)          # sm_scale carries qk_scale_factor/sqrt(d)
-        out = o_proj(attn * sigmoid(gate))
+        out = o_proj(attn * sigmoid(gate_proj(x)))
 
     The reference multiplies q by ``qk_scale_factor`` before rope; rope is a rotation, so
     the factor commutes and is folded into ``AttentionSpec.sm_scale`` instead (parse_config
-    sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate is computed from
-    the same normed layer input as q/k/v, so it rides the fused projection.
+    sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate stays its own
+    linear: the NVFP4 release quantizes it with a different global scale than q/k/v.
     """
 
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
@@ -41,10 +41,14 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         self.qo_attn_dim = self.num_q * head_dim
         self.kv_attn_dim = self.num_kv * head_dim
 
-        self._qkvg_split = [self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim, self.qo_attn_dim]
-        self.qkvg_proj = LinearColParallelMerged(
-            config.hidden_size, self._qkvg_split, has_bias=False,
-            quant_config=config.quant, prefix=f"{prefix}.qkvg_proj",
+        self._qkv_split = [self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim]
+        self.qkv_proj = LinearColParallelMerged(
+            config.hidden_size, self._qkv_split, has_bias=False,
+            quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
+        )
+        self.gate_proj = LinearColParallelMerged(
+            config.hidden_size, [self.qo_attn_dim], has_bias=False,
+            quant_config=config.quant, prefix=f"{prefix}.gate_proj",
         )
         self.o_proj = LinearReplicated(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
@@ -76,12 +80,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         ctx = get_global_ctx()
         T = x.shape[0]
 
-        qkvg = self.qkvg_proj.forward(x)
-        q, k, v, gate = torch.split(qkvg, self._qkvg_split, dim=-1)
-        del qkvg
+        gate = self.gate_proj.forward(x)
+        qkv = self.qkv_proj.forward(x)
+        q, k, v = torch.split(qkv, self._qkv_split, dim=-1)
+        del qkv
         q = self.qk_norm.forward(q.contiguous().view(T, self.num_q, self.head_dim))
         k = self.qk_norm.forward(k.contiguous().view(T, self.num_kv, self.head_dim))
-        v = v.contiguous()  # the split view has the qkvg row stride; the KV store needs contiguous
+        v = v.contiguous()  # the split view has the qkv row stride; the KV store needs contiguous
 
         q = q.reshape(T, self.qo_attn_dim)
         k = k.reshape(T, self.kv_attn_dim)
diff --git a/python/freetoken/models/muse_glimmer/gguf.py b/python/freetoken/models/muse_glimmer/gguf.py
index 7d5cee3b33..cad5e633d8 100644
--- a/python/freetoken/models/muse_glimmer/gguf.py
+++ b/python/freetoken/models/muse_glimmer/gguf.py
@@ -171,8 +171,9 @@ def t(name: str) -> GgufTensorInfo:
         else:
             qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
             q, k, v = (matrix(p) for p in qkv)
+        yield from fused(f"{sa}.qkv_proj", qkv, [_unpermute(q, heads), _unpermute(k, kv_heads), v], hidden)
         gate = t(f"{blk}.attn_gate")
-        yield from fused(f"{sa}.qkvg_proj", [*qkv, gate], [_unpermute(q, heads), _unpermute(k, kv_heads), v, matrix(gate)], hidden)
+        yield from fused(f"{sa}.gate_proj", [gate], [matrix(gate)], hidden)
         out = t(f"{blk}.attn_output")
         yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
         mlp = f"{pre}.mlp"
diff --git a/python/freetoken/models/muse_glimmer/weight.py b/python/freetoken/models/muse_glimmer/weight.py
index 8b9bb21317..3b8eb6d212 100644
--- a/python/freetoken/models/muse_glimmer/weight.py
+++ b/python/freetoken/models/muse_glimmer/weight.py
@@ -26,13 +26,10 @@
 _VISION_QKV_PARTS = (".attn.q_proj", ".attn.k_proj", ".attn.v_proj")
 
 # Fused projections, concatenated on the output dim in this exact order to match the
-# model's merged-linear splits. The attention gate rides the q/k/v fusion (it is computed
-# from the same layer input); ``.self_attn.gate_proj`` and the SwiGLU ``.mlp.gate_proj``
-# are disambiguated by the full suffix.
+# model's merged-linear splits. The full suffix keeps the standalone attention gate
+# ``.self_attn.gate_proj`` out of the SwiGLU ``.mlp.gate_proj`` fusion.
 _FUSIONS: dict[str, tuple[str, ...]] = {
-    ".self_attn.qkvg_proj": (
-        ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.gate_proj",
-    ),
+    ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"),
     ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"),
 }
 
@@ -134,8 +131,8 @@ def _iter_weights_compressed_tensors(
 
     Every text Linear (q/k/v/o, the attention gate, the MLP) is kept native NVFP4
     (W4A16): ``.weight`` (uint8 packed) + ``.weight_scale`` (fp8 block) + ``.weight_global``
-    (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v/gate fuse into
-    ``qkvg_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping
+    (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v fuse into
+    ``qkv_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping
     its own scales, so the fused FP4 weights are exact. Embeddings, norms, lm_head and the
     vision tower are bf16 (the checkpoint's ignore list). Scale lookups go through the
     shard-map reader: they can land in a different shard than their weight_packed."""
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 858e17db3f..42c6f3f548 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -80,10 +80,6 @@ class ModelSpec:
 _GEMMA4_UNIFIED_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4UnifiedMMProcessor"
 _GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MUSE_GLIMMER_PROCESSOR = "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor"
-_MUSE_GLIMMER_PACKED = (
-    ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")),
-    ("gate_up_proj", ("gate_proj", "up_proj")),
-)
 _MUSE_GLIMMER_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MINIMAX_M3_PACKED = _DENSE_PACKED + (
     ("index_qk_proj", ("index_q_proj", "index_k_proj")),
@@ -316,7 +312,7 @@ class ModelSpec:
         "freetoken.models.muse_glimmer",
         "MuseGlimmerForConditionalGeneration",
         checkpoint_roots=_LANGUAGE_MODEL_ROOT,
-        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
+        packed_modules_mapping=_DENSE_PACKED,
         mm_processor=_MUSE_GLIMMER_PROCESSOR,
         encoders=_MUSE_GLIMMER_ENCODERS,
     ),
@@ -326,7 +322,7 @@ class ModelSpec:
         "MuseGlimmerForCausalLM",
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
-        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
+        packed_modules_mapping=_DENSE_PACKED,
     ),
     "MistralForCausalLM": ModelSpec(
         "freetoken.models.mistral",
diff --git a/tests/models/test_muse_glimmer.py b/tests/models/test_muse_glimmer.py
index d35ff503a9..0a09a2840f 100644
--- a/tests/models/test_muse_glimmer.py
+++ b/tests/models/test_muse_glimmer.py
@@ -252,13 +252,12 @@ def test_weight_rename_and_fusion():
     ((key, weight),) = _vision_tensors("vision_tower.layers.0.attn.proj.weight", torch.zeros(2, 2), vbuf)
     assert key == "vision_tower.layers.0.attn.proj.weight" and torch.equal(weight, torch.zeros(2, 2))
 
-    # q/k/v + the attention gate fuse into qkvg_proj in declaration order.
+    # q/k/v fuse into qkv_proj in declaration order; the attention gate stays out of every fusion.
     buf: dict = {}
     parts = {
         "q_proj": torch.full((4, 2), 0.0),
         "k_proj": torch.full((2, 2), 1.0),
         "v_proj": torch.full((2, 2), 2.0),
-        "gate_proj": torch.full((4, 2), 3.0),
     }
     fused = None
     for name, tensor in parts.items():
@@ -268,10 +267,10 @@ def test_weight_rename_and_fusion():
             fused = out[0]
     assert fused is not None and not buf
     key, tensor = fused
-    assert key == "model.layers.0.self_attn.qkvg_proj.weight"
-    assert tensor.shape == (12, 2)
+    assert key == "model.layers.0.self_attn.qkv_proj.weight"
+    assert tensor.shape == (8, 2)
     assert tensor[0, 0] == 0.0 and tensor[4, 0] == 1.0 and tensor[6, 0] == 2.0
-    assert tensor[8, 0] == 3.0
+    assert ct_bf16_fuse("model.layers.0.self_attn.gate_proj", torch.zeros(4, 2), buf, _FUSIONS) is None
 
     # The MLP's own gate_proj is a different fusion (gate|up), not the attention one.
     buf2: dict = {}
@@ -365,7 +364,7 @@ def _bf16_checkpoint_tensors(hf) -> dict:
 @pytest.mark.parametrize("include_vision", [False, True])
 def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch, include_vision):
     """The BF16 loader must produce exactly the model's state-dict keys with the
-    right shapes (rename + qkvg / gate_up fusion, norms raw); the image path only
+    right shapes (rename + qkv / gate_up fusion, norms raw); the image path only
     when an encoder is built, then with q/k/v fused into attn.qkv."""
     import torch
 
@@ -399,14 +398,15 @@ def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch, inclu
         assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.weight"], torch.cat([tensors[v + f"{n}_proj.weight"] for n in "qkv"]))
         assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.bias"], torch.cat([tensors[v + f"{n}_proj.bias"] for n in "qkv"]))
         assert torch.equal(loaded["vision_tower.projection.weight"], tensors["model.vision_projection.weight"])
-    # fusion order [q, k, v, gate] against the raw parts
-    fused = loaded["model.layers.0.self_attn.qkvg_proj.weight"]
+    # fusion order [q, k, v] against the raw parts; the attention gate loads on its own
+    fused = loaded["model.layers.0.self_attn.qkv_proj.weight"]
     p = "model.language_model.layers.0.self_attn."
     q_dim = tensors[p + "q_proj.weight"].shape[0]
     kv_dim = tensors[p + "k_proj.weight"].shape[0]
     assert torch.equal(fused[:q_dim], tensors[p + "q_proj.weight"])
     assert torch.equal(fused[q_dim : q_dim + kv_dim], tensors[p + "k_proj.weight"])
-    assert torch.equal(fused[-q_dim:], tensors[p + "gate_proj.weight"])
+    assert torch.equal(fused[-kv_dim:], tensors[p + "v_proj.weight"])
+    assert torch.equal(loaded["model.layers.0.self_attn.gate_proj.weight"], tensors[p + "gate_proj.weight"])
 
 
 def test_iter_weights_nvfp4_cross_shard_scales(tmp_path, monkeypatch):
@@ -481,18 +481,19 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None):
     loaded = dict(
         iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True)
     )
-    qkvg = "model.layers.0.self_attn.qkvg_proj"
-    assert loaded[qkvg + ".weight"].shape == (2 * q + 2 * kv, H // 2)
-    assert loaded[qkvg + ".weight_scale"].shape == (2 * q + 2 * kv, H // 16)
-    assert loaded[qkvg + ".weight_global"].shape == (2 * q + 2 * kv,)
+    qkv = "model.layers.0.self_attn.qkv_proj"
+    assert loaded[qkv + ".weight"].shape == (q + 2 * kv, H // 2)
+    assert loaded[qkv + ".weight_scale"].shape == (q + 2 * kv, H // 16)
+    assert loaded[qkv + ".weight_global"].shape == (q + 2 * kv,)
     # global is the reciprocal of the stored quant-side scale
-    assert loaded[qkvg + ".weight_global"][0].item() == pytest.approx(0.5)
+    assert loaded[qkv + ".weight_global"][0].item() == pytest.approx(0.5)
+    assert loaded["model.layers.0.self_attn.gate_proj.weight"].shape == (q, H // 2)
     dp = "model.layers.0.mlp.down_proj"
     assert loaded[dp + ".weight"].shape == (H, I // 2)  # cross-shard scales resolved
     assert loaded[dp + ".weight_global"][0].item() == pytest.approx(0.25)
     # the quant-side activation global lands as the dequant-side input_scale the layer declares
     assert loaded[dp + ".input_scale"].item() == pytest.approx(1.0)
-    assert loaded[qkvg + ".input_scale"].shape == ()
+    assert loaded[qkv + ".input_scale"].shape == ()
     assert not any(k.endswith(".input_global_scale") for k in loaded)
     expected_vision = {"vision_tower." + k for k in MuseGlimmerVisionModel(parse_vision_config(hf)).state_dict()}
     assert {k for k in loaded if k.startswith("vision_tower.")} == expected_vision
@@ -503,7 +504,7 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None):
     text_only = dict(iter_weights(
         str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True, include_vision=False,
     ))
-    assert not any(k.startswith("vision_tower.") for k in text_only) and qkvg + ".weight" in text_only
+    assert not any(k.startswith("vision_tower.") for k in text_only) and qkv + ".weight" in text_only
 
 
 def test_raw_config_shim_serves_unknown_model_type(tmp_path):
@@ -551,7 +552,8 @@ def test_model_state_dict_matches_loader_keys():
     keys = set(model.state_dict().keys())
     layer0 = {k for k in keys if k.startswith("model.layers.0.")}
     assert layer0 == {
-        "model.layers.0.self_attn.qkvg_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.gate_proj.weight",
         "model.layers.0.self_attn.o_proj.weight",
         "model.layers.0.mlp.gate_up_proj.weight",
         "model.layers.0.mlp.down_proj.weight",
@@ -588,7 +590,7 @@ def test_model_state_dict_matches_loader_keys():
     qmodel = MuseGlimmerForCausalLM(qcfg)
     attn = qmodel.model.layers.op_list[0].self_attn
     mlp = qmodel.model.layers.op_list[0].mlp
-    for proj in (attn.qkvg_proj, attn.o_proj, mlp.gate_up_proj, mlp.down_proj):
+    for proj in (attn.qkv_proj, attn.gate_proj, attn.o_proj, mlp.gate_up_proj, mlp.down_proj):
         assert type(proj.quant_method) is Nvfp4LinearMethod
     assert type(qmodel.lm_head.quant_method) is UnquantizedLinearMethod  # lm_head stays bf16
     del model, qmodel, torch
diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py
index 0f31d6dee7..9c18b74907 100644
--- a/tests/models/test_quant_config.py
+++ b/tests/models/test_quant_config.py
@@ -257,7 +257,8 @@ def _gemma_tied_head(model, idx):
         "model.layers.{lin}.mlp.down_proj": FP8B, "lm_head": BF16,
     }, strategy="auto") for name in ("Qwen/Qwen3.8-27B-FP8", "Qwen/Qwen3.6-27B-FP8")],
     Case("RedHatAI/Muse-Glimmer-30B-NVFP4", CompressedTensorsConfig, {
-        "model.layers.0.self_attn.qkvg_proj": NVFP4, "model.layers.0.self_attn.o_proj": NVFP4,
+        "model.layers.0.self_attn.qkv_proj": NVFP4, "model.layers.0.self_attn.gate_proj": NVFP4,
+        "model.layers.0.self_attn.o_proj": NVFP4,
         "model.layers.0.mlp.gate_up_proj": NVFP4, "model.layers.0.mlp.down_proj": NVFP4, "lm_head": BF16,
     }, strategy="auto"),
     Case("nvidia/Gemma-4-31B-IT-NVFP4", ModelOptConfig, {

From 4d2e2e2104ab104fa17aa41710cd11afbf604588 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Sun, 18 Oct 2026 23:55:07 -0700
Subject: [PATCH 126/249] perf(linear): run nvfp4 linears on nvfp4 activations
 on sm_120 (#737)

* perf(linear): run nvfp4 linears on nvfp4 activations on sm_120

* fix(linear): guard the b12x nvfp4 linear against inputs it mishandles

* fix(linear): use the largest per-part global in the w4a4 nvfp4 linears
---
 .../kernel/gluon/nvfp4/linear/sm120_gemv.py   | 230 ++++++++++++++++++
 .../kernel/triton/nvfp4_act_quant.py          |  56 +++--
 .../layers/quantization/linear/nvfp4.py       | 182 +++++++++++++-
 tests/kernels/test_nvfp4_w4a4_linear.py       | 158 ++++++++++++
 4 files changed, 610 insertions(+), 16 deletions(-)
 create mode 100644 python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py
 create mode 100644 tests/kernels/test_nvfp4_w4a4_linear.py

diff --git a/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py b/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py
new file mode 100644
index 0000000000..cb63149a81
--- /dev/null
+++ b/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py
@@ -0,0 +1,230 @@
+"""NVFP4 W4A16 GEMV and skinny GEMM for sm_120 on the row-major weight: [N, K/2] e2m1 codes, flat 128x4 e4m3 block scales, one fp32 global.
+
+:func:`gemv_fma` decodes each code word once on the CUDA cores and reuses it for every token; :func:`gemv_mma` runs bf16 mma.sync with the
+weight rows on M and the tokens on N. Both keep K a constexpr with the K loop unrolled and load the activation with evict_last, so it
+stays in L2 while the weight stream passes; the weight keeps the default policy, as evict_first on it slows the kernels after it in a
+model step (lines it would have evicted stay in L2 for them to evict). :func:`gemv_config` picks the kernel and tile for a shape.
+"""
+
+import functools
+
+import torch
+import triton
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language.nvidia.ampere import mma_v2
+
+from ..utils import e2m1x8_bf16x2, scale_bf16x2
+
+# one code word (k = 8w + 0..7) times its f16x2 scale pair -> 8 fp32; e2m1 times e4m3 is exact in f16
+_DECODE = gl.constexpr("""
+{
+.reg .b8 b0, b1, b2, b3;
+.reg .b32 h0, h1, h2, h3;
+.reg .b16 l0, l1, l2, l3, l4, l5, l6, l7;
+mov.b32 {b0, b1, b2, b3}, $8;
+cvt.rn.f16x2.e2m1x2 h0, b0;
+cvt.rn.f16x2.e2m1x2 h1, b1;
+cvt.rn.f16x2.e2m1x2 h2, b2;
+cvt.rn.f16x2.e2m1x2 h3, b3;
+mul.rn.f16x2 h0, h0, $9;
+mul.rn.f16x2 h1, h1, $9;
+mul.rn.f16x2 h2, h2, $9;
+mul.rn.f16x2 h3, h3, $9;
+mov.b32 {l0, l1}, h0;
+mov.b32 {l2, l3}, h1;
+mov.b32 {l4, l5}, h2;
+mov.b32 {l6, l7}, h3;
+cvt.f32.f16 $0, l0;
+cvt.f32.f16 $1, l1;
+cvt.f32.f16 $2, l2;
+cvt.f32.f16 $3, l3;
+cvt.f32.f16 $4, l4;
+cvt.f32.f16 $5, l5;
+cvt.f32.f16 $6, l6;
+cvt.f32.f16 $7, l7;
+}
+""")
+
+
+@gluon.jit
+def _scale_row(row, K_TILES: gl.constexpr):
+    # 128x4 atoms: [row tile, k tile, row % 32, (row % 128) // 32, group % 4]
+    return (row // 128).to(gl.int64) * (K_TILES * 512) + (row % 32) * 16 + (row % 128) // 32 * 4
+
+
+@gluon.jit
+def _gemv_fma_kernel(x_ptr, w_ptr, s_ptr, g_ptr, out_ptr, T, stride_x, stride_o,
+                     K: gl.constexpr, T_C: gl.constexpr, R: gl.constexpr, V: gl.constexpr):
+    W: gl.constexpr = gl.num_warps()
+    BN: gl.constexpr = R * W
+    BKW: gl.constexpr = 32 * V
+    KW: gl.constexpr = K // 8
+    K_TILES: gl.constexpr = (K // 16 + 3) // 4
+    # a warp owns R whole rows; each lane holds V code words and their 8 k per row and per token
+    L4: gl.constexpr = gl.BlockedLayout([T_C, R, V, 8], [1, 1, 32, 1], [1, W, 1, 1], [3, 2, 1, 0])
+    La: gl.constexpr = gl.SliceLayout(3, L4)
+    Lw: gl.constexpr = gl.SliceLayout(0, La)
+    Lx: gl.constexpr = gl.SliceLayout(1, L4)
+    Lxj: gl.constexpr = gl.SliceLayout(1, La)
+
+    # 4 consecutive CTA rows are rows r, r + 32, r + 64, r + 96: one 16-byte chunk of every 128x4 scale atom
+    flat = gl.program_id(0) * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, Lw))
+    row = (flat // 128) * 128 + (flat % 128) // 4 + (flat % 4) * 32
+    wcol = gl.arange(0, BKW, layout=gl.SliceLayout(0, Lw))
+    w_base = w_ptr + row.to(gl.int64)[:, None] * KW + wcol[None, :]
+    s_row = _scale_row(row, K_TILES)
+    xt = gl.arange(0, T_C, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lx)))
+    xw = gl.arange(0, BKW, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lx)))
+    xj = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lx)))
+    x_off = xt[:, None, None] * stride_x + (xw * 8)[None, :, None] + xj[None, None, :]
+    x_ok = (xt < T)[:, None, None]
+
+    acc = gl.zeros([T_C, BN, BKW], gl.float32, La)
+    for ph in gl.static_range(KW // BKW):
+        words = gl.load(w_base + ph * BKW)
+        grp = (ph * BKW + wcol) // 2
+        sc = gl.load(s_ptr + s_row[:, None] + (grp // 4 * 512 + grp % 4)[None, :])
+        s16 = sc.to(gl.float8e4nv, bitcast=True).to(gl.float16).to(gl.uint16, bitcast=True).to(gl.uint32)
+        s2 = (s16 | (s16 << 16)).to(gl.int32, bitcast=True)
+        v0, v1, v2, v3, v4, v5, v6, v7 = gl.inline_asm_elementwise(_DECODE, "=r,=r,=r,=r,=r,=r,=r,=r,r,r", [words, s2],
+                                                                   dtype=(gl.float32,) * 8, is_pure=True, pack=1)
+        xa = gl.load(x_ptr + ph * BKW * 8 + x_off, mask=x_ok, other=0.0, eviction_policy="evict_last").to(gl.float32)
+        # in-thread split of the 8 k of each word: x0..x7 line up with v0..v7
+        ev, od = gl.split(gl.reshape(xa, [T_C, BKW, 4, 2]))
+        e0, e1 = gl.split(gl.reshape(ev, [T_C, BKW, 2, 2]))
+        o0, o1 = gl.split(gl.reshape(od, [T_C, BKW, 2, 2]))
+        x0, x4 = gl.split(e0)
+        x2, x6 = gl.split(e1)
+        x1, x5 = gl.split(o0)
+        x3, x7 = gl.split(o1)
+        acc += v0[None] * gl.convert_layout(x0, Lxj)[:, None, :]
+        acc += v1[None] * gl.convert_layout(x1, Lxj)[:, None, :]
+        acc += v2[None] * gl.convert_layout(x2, Lxj)[:, None, :]
+        acc += v3[None] * gl.convert_layout(x3, Lxj)[:, None, :]
+        acc += v4[None] * gl.convert_layout(x4, Lxj)[:, None, :]
+        acc += v5[None] * gl.convert_layout(x5, Lxj)[:, None, :]
+        acc += v6[None] * gl.convert_layout(x6, Lxj)[:, None, :]
+        acc += v7[None] * gl.convert_layout(x7, Lxj)[:, None, :]
+    res = gl.sum(acc, axis=2) * gl.load(g_ptr)
+    Lo: gl.constexpr = gl.SliceLayout(2, La)
+    ot = gl.arange(0, T_C, layout=gl.SliceLayout(1, Lo))
+    orow = gl.convert_layout(row, gl.SliceLayout(0, Lo))
+    gl.store(out_ptr + ot[:, None] * stride_o + orow[None, :], res.to(out_ptr.dtype.element_ty), mask=(ot < T)[:, None])
+
+
+@gluon.jit
+def _to_bf16(q):
+    lo = (q & 0xFFFF).to(gl.uint16).to(gl.bfloat16, bitcast=True)
+    hi = (q >> 16).to(gl.uint16).to(gl.bfloat16, bitcast=True)
+    return gl.join(lo, hi)
+
+
+@gluon.jit
+def _kperm(t, A: gl.constexpr, B: gl.constexpr):
+    # k_phys = 32 lid + 4 s + 2 h + e -> k_mma = 16 s + 8 h + 2 lid + e: a lane's contiguous 32 k become its fragment slots
+    t = gl.reshape(t, [A, B, 4, 8, 2, 2])
+    t = gl.permute(t, (0, 1, 3, 4, 2, 5))
+    return gl.reshape(t, [A, B, 128])
+
+
+@gluon.jit
+def _gemv_mma_kernel(x_ptr, w_ptr, s_ptr, g_ptr, out_ptr, T, stride_x, stride_o,
+                     K: gl.constexpr, T_C: gl.constexpr, BM: gl.constexpr, Lc: gl.constexpr, Lx: gl.constexpr, mma: gl.constexpr):
+    KW: gl.constexpr = gl.num_warps()
+    NB: gl.constexpr = K // 128
+    K_TILES: gl.constexpr = (K // 16 + 3) // 4
+    al: gl.constexpr = gl.DotOperandLayout(0, mma, 2)
+    bl: gl.constexpr = gl.DotOperandLayout(1, mma, 2)
+
+    # batched [KW, BM, T_C] mma: warp w takes the 128-k blocks w, w + KW, ...; the warps' partials are summed at the end
+    row0 = gl.program_id(0) * BM
+    cw = gl.arange(0, KW, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lc)))
+    cm = gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lc)))
+    cj = gl.arange(0, 16, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lc)))
+    row = row0 + cm
+    c_off = (row.to(gl.int64) * (K // 8))[None, :, None] + (cw * 16)[:, None, None] + cj[None, None, :]
+    grp = cw[:, None, None] * 8 + (cj // 2)[None, None, :]
+    s_off = _scale_row(row, K_TILES)[None, :, None] + (grp // 4 * 512 + grp % 4)
+    xw = gl.arange(0, KW, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lx)))
+    xt = gl.arange(0, T_C, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lx)))
+    xk = gl.arange(0, 128, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lx)))
+    x_off = (xt * stride_x)[None, :, None] + (xw * 128)[:, None, None] + xk[None, None, :]
+    x_ok = (xt < T)[None, :, None]
+
+    acc = gl.zeros([KW, BM, T_C], gl.float32, mma)
+    for it in gl.static_range((NB + KW - 1) // KW):
+        b0 = it * KW
+        if (it + 1) * KW <= NB:
+            q = gl.load(w_ptr + c_off + b0 * 16)
+            sc = gl.load(s_ptr + s_off + b0 * 1024)
+            xa = gl.load(x_ptr + x_off + b0 * 128, mask=x_ok, other=0.0, eviction_policy="evict_last")
+        else:
+            ok = (cw + b0 < NB)[:, None, None]
+            q = gl.load(w_ptr + c_off + b0 * 16, mask=ok, other=0)
+            sc = gl.load(s_ptr + s_off + b0 * 1024, mask=ok, other=0)
+            xa = gl.load(x_ptr + x_off + b0 * 128, mask=x_ok & (xw + b0 < NB)[:, None, None], other=0.0, eviction_policy="evict_last")
+        w0, w1, w2, w3 = e2m1x8_bf16x2(q.to(gl.uint32, bitcast=True), scale_bf16x2(sc))
+        # join appends its axis last: this order flattens to the k pairs w0, w1, w2, w3
+        p = gl.reshape(gl.join(gl.join(w0, w2), gl.join(w1, w3)), [KW, BM, 64])
+        a = gl.convert_layout(_kperm(gl.reshape(_to_bf16(p), [KW, BM, 128]), KW, BM), al)
+        b = gl.convert_layout(gl.permute(_kperm(xa, KW, T_C), (0, 2, 1)), bl)
+        acc = mma_v2(a, b, acc)
+    # e2m1x8_bf16x2 leaves a factor 2^-8 in the decoded weight
+    res = gl.sum(acc, axis=0) * (gl.load(g_ptr) * 256.0)
+    Lo: gl.constexpr = gl.SliceLayout(0, mma)
+    om = gl.arange(0, BM, layout=gl.SliceLayout(1, Lo))
+    ot = gl.arange(0, T_C, layout=gl.SliceLayout(0, Lo))
+    gl.store(out_ptr + ot[None, :] * stride_o + (row0 + om)[:, None], res.to(out_ptr.dtype.element_ty), mask=(ot < T)[None, :])
+
+
+@functools.cache
+def _mma_layouts(kw: int, t_c: int, bm: int):
+    warps = [[1 << i, 0, 0] for i in range(kw.bit_length() - 1)]
+    lanes = [[0, 0, 32], [0, 0, 64], [0, 1, 0], [0, 2, 0], [0, 4, 0]]
+    # codes [KW, BM, 16 words]: lane (gid, lid) loads words 4 lid .. 4 lid + 3 of rows gid, gid + 8, gid + 16, ...
+    codes = gl.DistributedLinearLayout([[0, 0, 1], [0, 0, 2], [0, 8, 0]] + [[0, 16 << i, 0] for i in range(bm.bit_length() - 5)],
+                                       [[0, 0, 4], [0, 0, 8], [0, 1, 0], [0, 2, 0], [0, 4, 0]], warps, [], [kw, bm, 16])
+    # activations [KW, T_C, 128]: lane (gid, lid) loads k 32 lid .. 32 lid + 31 of tokens gid, gid + 8
+    acts = gl.DistributedLinearLayout([[0, 0, 1 << i] for i in range(5)] + [[0, 8 << i, 0] for i in range(t_c.bit_length() - 4)],
+                                      lanes, warps, [], [kw, t_c, 128])
+    return codes, acts, gl.NVMMADistributedLayout([2, 0], [kw, 1, 1], [1, 16, 8])
+
+
+def gemv_fma(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, g: torch.Tensor, R: int, V: int, num_warps: int) -> torch.Tensor:
+    """x [T, K] bf16 row-contiguous, weight [N, K/2] uint8, scale flat 128x4 e4m3, g fp32 -> [T, N]; see :func:`gemv_config` for the shape rules."""
+    T, K = x.shape
+    N = weight.shape[0]
+    out = torch.empty(T, N, dtype=x.dtype, device=x.device)
+    _gemv_fma_kernel[(N // (R * num_warps),)](x, weight.view(torch.int32), scale.view(torch.uint8), g, out, T, x.stride(0), out.stride(0),
+                                              K=K, T_C=triton.next_power_of_2(T), R=R, V=V, num_warps=num_warps)
+    return out
+
+
+def gemv_mma(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, g: torch.Tensor, bm: int, num_warps: int) -> torch.Tensor:
+    """Same contract as :func:`gemv_fma`; T <= 16."""
+    T, K = x.shape
+    N = weight.shape[0]
+    t_c = max(8, triton.next_power_of_2(T))
+    codes, acts, mma = _mma_layouts(num_warps, t_c, bm)
+    out = torch.empty(T, N, dtype=x.dtype, device=x.device)
+    _gemv_mma_kernel[(N // bm,)](x, weight.view(torch.int32), scale.view(torch.uint8), g, out, T, x.stride(0), out.stride(0),
+                                 K=K, T_C=t_c, BM=bm, Lc=codes, Lx=acts, mma=mma, num_warps=num_warps)
+    return out
+
+
+def gemv_config(N: int, K: int, T: int, max_tokens: int = 16):
+    """(kernel, tile) for a [T, K] x [N, K] product, or None when neither kernel covers it (T > max_tokens <= 16, or an unsupported shape)."""
+    if T == 1 and K < 16384 or T == 2 and N >= 16384:
+        words = K // 8
+        if N % 128 == 0:
+            # wide layers take more rows per CTA; narrow ones longer words per lane
+            if N >= 16384 and words % 64 == 0:
+                return gemv_fma, {"R": 4, "V": 2, "num_warps": 4}
+            for V, warps in ((4, 4), (2, 8), (1, 8)):
+                if words % (32 * V) == 0 and not (N >= 16384 and V == 4):
+                    return gemv_fma, {"R": 2, "V": V, "num_warps": warps}
+    if T > max_tokens or K % 128:
+        return None
+    tile = {"bm": 16, "num_warps": 8} if T <= 8 or N % 32 else {"bm": 32, "num_warps": 8}
+    return (gemv_mma, tile) if N % tile["bm"] == 0 else None
diff --git a/python/freetoken/kernel/triton/nvfp4_act_quant.py b/python/freetoken/kernel/triton/nvfp4_act_quant.py
index b52cfbe898..6f56a46068 100644
--- a/python/freetoken/kernel/triton/nvfp4_act_quant.py
+++ b/python/freetoken/kernel/triton/nvfp4_act_quant.py
@@ -16,27 +16,53 @@ def _e2m1x2(hi, lo):
 
 
 @triton.jit
-def _quant_kernel(x_ptr, q_ptr, s_ptr, g_ptr, K, stride_x, BLOCK_K: tl.constexpr):
-    row = tl.program_id(0).to(tl.int64)
-    offs = tl.arange(0, BLOCK_K)
-    x = tl.load(x_ptr + row * stride_x + offs, mask=offs < K, other=0.0).to(tl.float32)
+def _quant_kernel(x_ptr, q_ptr, s_ptr, g_ptr, R, K, stride_x, k_tiles,
+                  SWIZZLED: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr):
+    # with a static global every 16-wide block quantizes on its own, so a row splits over programs
+    rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    cols = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K)
+    row_ok = rows < R
+    r64 = rows.to(tl.int64)
+    x = tl.load(x_ptr + r64[:, None] * stride_x + cols[None, :], mask=row_ok[:, None] & (cols < K)[None, :], other=0.0)
     g = tl.load(g_ptr)
-    xb = tl.reshape(x, [BLOCK_K // 16, 16])
-    s = (tl.max(tl.abs(xb), 1) / (6.0 * g)).to(tl.float8e4nv)
+    xb = tl.reshape(x.to(tl.float32), [BLOCK_M, BLOCK_K // 16, 16])
+    s = (tl.max(tl.abs(xb), 2) / (6.0 * g)).to(tl.float8e4nv)
     sf = s.to(tl.float32) * g
     inv = tl.where(sf > 0, 1.0 / sf, 0.0)
-    lo, hi = tl.split(tl.reshape(xb * inv[:, None], [BLOCK_K // 2, 2]))
-    oq = tl.arange(0, BLOCK_K // 2)
-    tl.store(q_ptr + row * (K // 2) + oq, _e2m1x2(hi, lo).to(tl.uint8), mask=oq < K // 2)
-    os_ = tl.arange(0, BLOCK_K // 16)
-    tl.store(s_ptr + row * (K // 16) + os_, s.to(tl.uint8, bitcast=True), mask=os_ < K // 16)
+    lo, hi = tl.split(tl.reshape(xb * inv[:, :, None], [BLOCK_M, BLOCK_K // 2, 2]))
+    oq = tl.program_id(1) * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2)
+    tl.store(q_ptr + r64[:, None] * (K // 2) + oq[None, :], _e2m1x2(hi, lo).to(tl.uint8),
+             mask=row_ok[:, None] & (oq < K // 2)[None, :])
+    sc = tl.program_id(1) * (BLOCK_K // 16) + tl.arange(0, BLOCK_K // 16)
+    if SWIZZLED:
+        # cuBLAS / CUTLASS 128x4 block-scale atoms: [row tile, k tile, row % 32, (row % 128) // 32, col % 4]
+        s_off = ((r64 // 128)[:, None] * k_tiles + (sc // 4)[None, :]) * 512 \
+            + ((rows % 32) * 16 + (rows % 128) // 32 * 4)[:, None] + (sc % 4)[None, :]
+    else:
+        s_off = r64[:, None] * (K // 16) + sc[None, :]
+    tl.store(s_ptr + s_off, s.to(tl.uint8, bitcast=True), mask=row_ok[:, None] & (sc < K // 16)[None, :])
 
 
-def quantize_nvfp4(x: torch.Tensor, g: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """x [R, K] -> (codes [R, K / 2] uint8, low nibble first; block scales [R, K / 16] e4m3) under the fp32 global g [1]; K % 16 == 0."""
+def quantize_nvfp4(x: torch.Tensor, g: torch.Tensor, swizzled: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    """x [R, K] -> (codes [R, K / 2] uint8, low nibble first; e4m3 block scales) under the fp32 global g [1]; K % 16 == 0.
+
+    The scales are [R, K / 16], or with ``swizzled`` the flat cuBLAS / CUTLASS 128x4 layout that ``torch.nn.functional.scaled_mm`` takes,
+    rows padded to 128 and K / 16 to 4 with zeros.
+    """
     R, K = x.shape
+    assert K % 16 == 0 and x.stride(1) == 1
     q = torch.empty(R, K // 2, dtype=torch.uint8, device=x.device)
-    s = torch.empty(R, K // 16, dtype=torch.uint8, device=x.device)
+    k_tiles = triton.cdiv(K // 16, 4)
+    if swizzled:
+        n = triton.cdiv(R, 128) * 128 * k_tiles * 4
+        exact = R % 128 == 0 and (K // 16) % 4 == 0
+        s = (torch.empty if exact else torch.zeros)(n, dtype=torch.uint8, device=x.device)
+    else:
+        s = torch.empty(R, K // 16, dtype=torch.uint8, device=x.device)
     if R:
-        _quant_kernel[(R,)](x, q, s, g, K, x.stride(0), BLOCK_K=triton.next_power_of_2(K), num_warps=4)
+        block_k = min(triton.next_power_of_2(K), 1024)
+        block_m = max(1, 1024 // block_k)
+        grid = (triton.cdiv(R, block_m), triton.cdiv(K, block_k))
+        _quant_kernel[grid](x, q, s, g, R, K, x.stride(0), k_tiles, SWIZZLED=swizzled, BLOCK_M=block_m, BLOCK_K=block_k,
+                            num_warps=4)
     return q, s.view(torch.float8_e4m3fn)
diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py
index 80f5b78d3c..6a13165318 100644
--- a/python/freetoken/layers/quantization/linear/nvfp4.py
+++ b/python/freetoken/layers/quantization/linear/nvfp4.py
@@ -2,18 +2,24 @@
 
 from __future__ import annotations
 
+import functools
 from types import SimpleNamespace
 from typing import Any
 
 import torch
 
 from freetoken.kernel import backend
+from freetoken.utils import init_logger
 
 from ..registry import LayerKind, register_method
 from ..scheme import NVFP4_GROUP as GROUP, QuantKind
 from .base import LinearConfig, LinearKernel, LinearMethod
 
+logger = init_logger(__name__)
+
 FP8 = torch.float8_e4m3fn
+# fused parts whose weight globals are further apart than this ratio get a load-time warning
+_PART_GLOBAL_WARN_RATIO = 1.05
 
 
 class TritonNvfp4LinearKernel(LinearKernel):
@@ -32,6 +38,180 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
         return nvfp4_dense_linear_t(x, layer.weight, layer.weight_scale, layer.weight_global, layer.bias)
 
 
+@functools.cache
+def _warn_part_globals() -> None:
+    logger.warning_rank0(
+        "an NVFP4 linear fuses parts whose weight global scales differ by more than %d%%; the W4A4 kernels use the largest "
+        "for every part, which scales the other parts' outputs by their ratio; --quant-backend linear.nvfp4=triton keeps "
+        "each part's global", round((_PART_GLOBAL_WARN_RATIO - 1) * 100),
+    )
+
+
+def _tensor_global(layer: Any) -> torch.Tensor:
+    """The one weight global the W4A4 GEMMs and the small-batch kernels take: the largest per-part global, the block scales unchanged."""
+    g = layer.weight_global.float()
+    g_max = g.max()
+    if g_max > _PART_GLOBAL_WARN_RATIO * g.min():
+        _warn_part_globals()
+    return g_max
+
+
+def _swizzle_128x4(scale: torch.Tensor) -> torch.Tensor:
+    """[N, K/16] block scales -> the flat cuBLAS / CUTLASS 128x4 layout, rows padded to 128 and columns to 4 with zeros."""
+    rows, cols = scale.shape
+    rb, cb = -(-rows // 128), -(-cols // 4)
+    padded = torch.zeros(rb * 128, cb * 4, dtype=scale.dtype, device=scale.device)
+    padded[:rows, :cols] = scale
+    blocks = padded.view(rb, 128, cb, 4).permute(0, 2, 1, 3)
+    return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1).contiguous()
+
+
+@functools.cache
+def _sm120_gemv(device_index: int):
+    """The row-major W4A16 GEMV / skinny GEMM module for this GPU, or None."""
+    from freetoken.kernel import gluon
+
+    if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available():
+        return None
+    from freetoken.kernel.gluon.nvfp4.linear import sm120_gemv
+
+    return sm120_gemv
+
+
+def _small_batch(layer: Any, x: torch.Tensor, max_tokens: int) -> torch.Tensor | None:
+    """W4A16 on the W4A4 kernels' own weight for decode-sized batches, where the W4A4 GEMMs leave bandwidth unused; None to run the GEMM."""
+    gemv = _sm120_gemv(x.device.index)
+    if gemv is None or x.dtype != torch.bfloat16 or x.shape[0] == 0:
+        return None
+    route = gemv.gemv_config(layer.weight.shape[0], x.shape[1], x.shape[0], max_tokens)
+    if route is None:
+        return None
+    kernel, tile = route
+    return kernel(x if x.stride(1) == 1 else x.contiguous(), layer.weight, layer.weight_scale, layer.weight_global, **tile)
+
+
+@functools.cache
+def _flashinfer_cuda_major() -> int:
+    """The CUDA version flashinfer's b12x GEMM checks at call time: nvcc's, else torch's."""
+    from flashinfer.jit.cpp_ext import get_cuda_version
+
+    return get_cuda_version().major
+
+
+def _small_batch_usable() -> bool:
+    return torch.cuda.is_available() and _sm120_gemv(torch.cuda.current_device()) is not None
+
+
+class CublasNvfp4LinearKernel(LinearKernel):
+    """W4A4: activations quantized under the checkpoint's input_scale, then torch's NVFP4 scaled_mm (cuBLASLt) on the row-major weight."""
+
+    name = "cublas"
+
+    def unusable_reason(self, cfg: LinearConfig) -> str | None:
+        if backend.is_rocm() or not torch.cuda.is_available():
+            return "needs CUDA"
+        if cfg.scheme is None or not cfg.scheme.has("input_scale"):
+            return "the weight declares no static activation scale (W4A16 checkpoint)"
+        if backend.device_capability() < (10, 0):
+            return "NVFP4 tensor cores need sm_100+"
+        if not hasattr(torch.nn.functional, "scaled_mm") or not hasattr(torch, "float4_e2m1fn_x2"):
+            return f"torch {torch.__version__} has no NVFP4 scaled_mm"
+        if cfg.out_features % 16 or cfg.in_features % 32:
+            return f"scaled_mm needs N % 16 == 0 and K % 32 == 0, got N={cfg.out_features} K={cfg.in_features}"
+        return None
+
+    def worth_it(self, cfg: LinearConfig) -> bool:
+        # cuBLAS alone runs decode-sized batches up to 1.4x slower than the triton GEMV; the small-batch kernels cover those on sm_120
+        return _small_batch_usable()
+
+    def finalize(self, layer: Any) -> None:
+        g_w = _tensor_global(layer)
+        layer.weight_scale = _swizzle_128x4(layer.weight_scale)
+        layer.weight_global = g_w.reshape(())
+        layer.input_scale = layer.input_scale.float().reshape(())
+
+    def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
+        from torch.nn import functional as F
+
+        from freetoken.kernel.triton.nvfp4_act_quant import quantize_nvfp4
+
+        *lead, K = x.shape
+        out = _small_batch(layer, x.reshape(-1, K), 16)
+        if out is not None:
+            out = out.reshape(*lead, layer.weight.shape[0])
+            return out + layer.bias.to(out.dtype) if layer.bias is not None else out
+        xq, xs = quantize_nvfp4(x.reshape(-1, K), layer.input_scale.reshape(1), swizzled=True)
+        block, tensor = F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise
+        swizzle = [F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE]
+        out = F.scaled_mm(
+            xq.view(torch.float4_e2m1fn_x2), layer.weight.view(torch.float4_e2m1fn_x2).t(),
+            [xs, layer.input_scale], [block, tensor], [layer.weight_scale, layer.weight_global], [block, tensor],
+            swizzle_a=swizzle, swizzle_b=swizzle, output_dtype=x.dtype,
+        ).reshape(*lead, layer.weight.shape[0])
+        return out + layer.bias.to(out.dtype) if layer.bias is not None else out
+
+
+class B12xNvfp4LinearKernel(LinearKernel):
+    """W4A4 through flashinfer alone: its activation quant and its SM12x CuTe-DSL GEMM (mm_fp4 backend b12x) on the row-major weight."""
+
+    name = "b12x"
+
+    def unusable_reason(self, cfg: LinearConfig) -> str | None:
+        if backend.is_rocm() or not torch.cuda.is_available():
+            return "needs CUDA"
+        if cfg.scheme is None or not cfg.scheme.has("input_scale"):
+            return "the weight declares no static activation scale (W4A16 checkpoint)"
+        cc = backend.device_capability()
+        if cc[0] != 12:
+            return f"b12x requires sm_120 / sm_121, got sm_{cc[0]}{cc[1]}"
+        if not backend.is_flashinfer_installed():
+            return "flashinfer is not installed"
+        drv = backend.driver_cuda_version()
+        if drv is not None and drv < 13000:
+            return f"b12x JIT-compiles through the driver and needs CUDA 13+, the driver supports {drv // 1000}.{drv % 1000 // 10}"
+        if cfg.in_features % 32 or cfg.out_features % 8:
+            # flashinfer checks only K; N % 8 != 0 returns wrong output
+            return f"b12x needs K % 32 == 0 and N % 8 == 0, got K={cfg.in_features} N={cfg.out_features}"
+        try:
+            from flashinfer import fp4_quantize, mm_fp4, nvfp4_block_scale_interleave  # noqa: F401
+
+            cuda_major = _flashinfer_cuda_major()
+        except Exception as exc:
+            return f"flashinfer FP4 GEMM symbols are unusable ({exc!r})"
+        if cuda_major < 13:
+            return f"flashinfer's b12x GEMM needs a CUDA 13 toolkit (nvcc, else torch), got CUDA {cuda_major}"
+        return None
+
+    def worth_it(self, cfg: LinearConfig) -> bool:
+        return _small_batch_usable()
+
+    def finalize(self, layer: Any) -> None:
+        from flashinfer import nvfp4_block_scale_interleave
+
+        g_w = _tensor_global(layer)
+        layer.weight_scale = nvfp4_block_scale_interleave(layer.weight_scale.view(torch.uint8)).view(FP8)
+        # flashinfer takes the quant-side activation global and alpha = 1 / (that * the weight's quant-side global)
+        layer.input_scale_inv = (1.0 / layer.input_scale.float()).reshape(1)
+        layer.alpha = (layer.input_scale.float() * g_w).reshape(1)
+        layer.weight_global = g_w.reshape(())
+
+    def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
+        from flashinfer import fp4_quantize, mm_fp4
+
+        *lead, K = x.shape
+        x2 = x.reshape(-1, K)
+        if x2.shape[0] == 0:
+            return x.new_empty(*lead, layer.weight.shape[0])
+        # b12x's W4A4 GEMM beats the W4A16 kernels from 9 tokens on GeForce, where bf16 mma accumulates in fp32 at half rate
+        out = _small_batch(layer, x2, 8)
+        if out is not None:
+            out = out.reshape(*lead, layer.weight.shape[0])
+            return out + layer.bias.to(out.dtype) if layer.bias is not None else out
+        xq, xs = fp4_quantize(x2.contiguous(), layer.input_scale_inv, is_sf_swizzled_layout=True)
+        out = mm_fp4(xq, layer.weight.T, xs, layer.weight_scale, layer.alpha, x.dtype, backend="b12x").reshape(*lead, layer.weight.shape[0])
+        return out + layer.bias.to(out.dtype) if layer.bias is not None else out
+
+
 class MarlinNvfp4LinearKernel(LinearKernel):
     """vLLM's fused FP4 Marlin GEMM (W4A16); needs one global scale for the whole weight."""
 
@@ -113,7 +293,7 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
 
 @register_method(QuantKind.NVFP4, LayerKind.LINEAR)
 class Nvfp4LinearMethod(LinearMethod):
-    candidates = (TritonNvfp4LinearKernel, MarlinNvfp4LinearKernel, EmulationNvfp4LinearKernel)
+    candidates = (B12xNvfp4LinearKernel, CublasNvfp4LinearKernel, TritonNvfp4LinearKernel, MarlinNvfp4LinearKernel, EmulationNvfp4LinearKernel)
 
     def create_weights(self, layer: Any) -> None:
         g = self.cfg
diff --git a/tests/kernels/test_nvfp4_w4a4_linear.py b/tests/kernels/test_nvfp4_w4a4_linear.py
new file mode 100644
index 0000000000..1840d30f5a
--- /dev/null
+++ b/tests/kernels/test_nvfp4_w4a4_linear.py
@@ -0,0 +1,158 @@
+"""The W4A4 NVFP4 linears: cublas (our activation quant with 128x4-swizzled scales feeding torch's NVFP4 scaled_mm) and b12x (flashinfer only),
+and the sm_120 W4A16 GEMV / skinny GEMM both run for small batches on the same row-major weight."""
+
+from __future__ import annotations
+
+import math
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), reason="needs NVFP4 tensor cores (sm_100+)"
+)
+
+_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
+
+
+def _round_e2m1(v: torch.Tensor) -> torch.Tensor:
+    """Round to the nearest e2m1 value, ties to the even code, saturating at 6 (cvt.rn.satfinite)."""
+    grid = _E2M1[:8].to(v.device)
+    a = v.abs().clamp(max=6.0)
+    hi = torch.searchsorted(grid, a.contiguous()).clamp(max=7)
+    lo = (hi - 1).clamp(min=0)
+    d_lo, d_hi = a - grid[lo], grid[hi] - a
+    code = torch.where(d_lo < d_hi, lo, torch.where(d_hi < d_lo, hi, torch.where(lo % 2 == 0, lo, hi)))
+    return torch.copysign(grid[code], v)
+
+
+def _qdq(x: torch.Tensor, g: float) -> torch.Tensor:
+    xb = x.float().view(x.shape[0], -1, 16)
+    g = torch.tensor(g, dtype=torch.float32, device=x.device)
+    sf = ((xb.abs().amax(-1) / (6.0 * g)).to(torch.float8_e4m3fn).float() * g).unsqueeze(-1)
+    inv = torch.where(sf > 0, 1.0 / sf, torch.zeros_like(sf))
+    return (_round_e2m1(xb * inv) * sf).view_as(x)
+
+
+def _dequant(packed: torch.Tensor, scale: torch.Tensor, row_global: torch.Tensor) -> torch.Tensor:
+    n, k2 = packed.shape
+    codes = torch.stack([packed & 0xF, packed >> 4], dim=-1).view(n, 2 * k2).long()
+    return _E2M1.to(packed.device)[codes] * scale.float().repeat_interleave(16, dim=1) * row_global.float().unsqueeze(1)
+
+
+def _small_batch_route(N, K, M, max_tokens=16):
+    if torch.cuda.get_device_capability() != (12, 0):
+        return None
+    from freetoken.kernel import gluon
+
+    if not gluon.available():
+        return None
+    from freetoken.kernel.gluon.nvfp4.linear.sm120_gemv import gemv_config
+
+    return gemv_config(N, K, M, max_tokens)
+
+
+@pytest.mark.parametrize("rows,K", [(1, 2048), (7, 2080), (300, 5376)])
+def test_swizzled_scales_are_the_cublas_128x4_layout(rows, K):
+    """Same codes either way; the swizzled scales are the linear ones in cuBLAS's [row tile, k tile, 32, 4, 4] order, zero padded."""
+    from freetoken.kernel.triton.nvfp4_act_quant import quantize_nvfp4
+    from freetoken.layers.quantization.linear.nvfp4 import _swizzle_128x4
+
+    x = torch.randn(rows, K, device="cuda", dtype=torch.bfloat16)
+    g = (x.float().abs().max() / 2688).reshape(1)
+    q, s = quantize_nvfp4(x, g)
+    q_swz, s_swz = quantize_nvfp4(x, g, swizzled=True)
+    assert torch.equal(q, q_swz)
+    assert torch.equal(_swizzle_128x4(s).view(torch.uint8), s_swz.view(torch.uint8))
+
+
+@pytest.mark.parametrize("kernel_name", ["cublas", "b12x"])
+@pytest.mark.parametrize("M", [1, 2, 7, 16, 300])
+@pytest.mark.parametrize("N,K", [(512, 1024), (528, 2080), (12, 2048)])
+@pytest.mark.parametrize("parts_differ", [False, True])
+def test_w4a4_linear_matches_nvfp4_activation_reference(kernel_name, M, N, K, parts_differ):
+    """finalize + apply against a torch emulation: W4A4, or W4A16 where the sm_120 small-batch kernels take the batch (M <= 16 for cublas,
+    M <= 8 for b12x, at 512 x 1024; 528 x 2080 fits neither and keeps the GEMM; N = 12 must be refused, b12x returns garbage there);
+    parts_differ gives the halves of a fused weight different globals, and the larger one applies to both."""
+    from freetoken.layers.quantization.linear.base import LinearConfig
+    from freetoken.layers.quantization.linear.nvfp4 import Nvfp4LinearMethod
+    from freetoken.layers.quantization.scheme import nvfp4_scheme
+
+    kernel = next(c for c in Nvfp4LinearMethod.candidates if c.name == kernel_name)()
+    reason = kernel.unusable_reason(LinearConfig(K, N, (N // 2, N // 2), nvfp4_scheme(input_scale=True)))
+    if reason:
+        pytest.skip(reason)
+    g = torch.Generator(device="cuda").manual_seed(0)
+    weight = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda", generator=g)
+    scale = (torch.rand(N, K // 16, device="cuda", generator=g) + 0.25).to(torch.float8_e4m3fn)
+    row_global = torch.full((N,), 1e-3, dtype=torch.float16, device="cuda")
+    if parts_differ:
+        row_global[N // 2:] = 4e-3
+    bias = torch.randn(N, device="cuda", dtype=torch.bfloat16)
+    x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16, generator=g)
+    # off amax / 2688 by an irrational factor: with amax / 2688 many values sit exactly on an e2m1 midpoint, where the kernel's
+    # approximate fp32 reciprocal and torch's exact one round apart
+    g_in = float(x.float().abs().max()) / 2688 * math.pi / 3
+    # one tensor-wise global: the largest part's, applied to every row
+    w_ref = _dequant(weight, scale, torch.full_like(row_global, float(row_global.max())))
+    small = _small_batch_route(N, K, M, 16 if kernel_name == "cublas" else 8) is not None
+    ref = (x.float() if small else _qdq(x, g_in)) @ w_ref.T + bias.float()
+
+    layer = SimpleNamespace(weight=weight.clone(), weight_scale=scale.clone(), weight_global=row_global.clone(),
+                            input_scale=torch.tensor(g_in, device="cuda"), bias=bias)
+    kernel.finalize(layer)
+    out = kernel.apply(layer, x)
+    torch.testing.assert_close(out.float(), ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max()))
+
+
+@pytest.mark.parametrize("N,K,T", [
+    (16384, 1024, 1),  # fma, 4 rows per warp
+    (16384, 768, 2),  # fma, one word per lane
+    (640, 2048, 1),  # fma, four words per lane
+    (640, 17408, 1),  # long K: mma
+    (640, 4224, 12),  # mma bm32, masked last K block
+    (256, 16512, 4),  # mma, masked last K block
+    (640, 2048, 5),  # mma bm16, padded tokens
+])
+def test_sm120_gemv_tiles_match_dequant_reference(N, K, T):
+    """Each tile gemv_config hands out, against x @ dequant(W)^T on the flat 128x4 scales."""
+    route = _small_batch_route(N, K, T)
+    if route is None:
+        pytest.skip("needs sm_120 with Gluon")
+    from freetoken.layers.quantization.linear.nvfp4 import _swizzle_128x4
+
+    g = torch.Generator(device="cuda").manual_seed(0)
+    weight = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda", generator=g)
+    scale = (torch.rand(N, K // 16, device="cuda", generator=g) + 0.25).to(torch.float8_e4m3fn)
+    x = torch.randn(T, K, device="cuda", dtype=torch.bfloat16, generator=g)
+    g_w = torch.tensor(1e-3, device="cuda")
+    ref = x.float() @ _dequant(weight, scale, torch.full((N,), 1e-3, device="cuda")).T
+    kernel, tile = route
+    out = kernel(x, weight, _swizzle_128x4(scale), g_w, **tile)
+    torch.testing.assert_close(out.float(), ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max()))
+
+
+@pytest.mark.parametrize("kernel_name", ["cublas", "b12x"])
+@pytest.mark.parametrize("M", [0, 4, 64])
+def test_w4a4_linear_takes_row_strided_and_empty_inputs(kernel_name, M):
+    """A row-strided slice and an empty batch give the same output as a contiguous input, on both the small-batch and the GEMM route."""
+    from freetoken.layers.quantization.linear.base import LinearConfig
+    from freetoken.layers.quantization.linear.nvfp4 import Nvfp4LinearMethod
+    from freetoken.layers.quantization.scheme import nvfp4_scheme
+
+    N, K = 512, 1024
+    kernel = next(c for c in Nvfp4LinearMethod.candidates if c.name == kernel_name)()
+    reason = kernel.unusable_reason(LinearConfig(K, N, (N,), nvfp4_scheme(input_scale=True)))
+    if reason:
+        pytest.skip(reason)
+    g = torch.Generator(device="cuda").manual_seed(0)
+    layer = SimpleNamespace(weight=torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda", generator=g),
+                            weight_scale=(torch.rand(N, K // 16, device="cuda", generator=g) + 0.25).to(torch.float8_e4m3fn),
+                            weight_global=torch.full((N,), 1e-3, dtype=torch.float16, device="cuda"),
+                            input_scale=torch.tensor(0.01, device="cuda"), bias=None)
+    kernel.finalize(layer)
+    buf = torch.randn(M, K + 64, device="cuda", dtype=torch.bfloat16, generator=g)
+    out = kernel.apply(layer, buf[:, :K])
+    assert out.shape == (M, N)
+    torch.testing.assert_close(out, kernel.apply(layer, buf[:, :K].contiguous()), rtol=0, atol=0)

From 0db5023b60ec8a72e7f9c87be5119956e298f7b9 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 12:33:43 -0700
Subject: [PATCH 127/249] fix(quant)!: keep nvfp4 global scales in fp32 (#739)

NVFP4 weight global scales (amax / 2688, 2e-5 to 4e-4 in the served checkpoints) sit at fp16's
normal minimum; stored as fp16, up to 97% of a checkpoint's expert globals were subnormal, with
up to 1.7e-3 relative error per row.

BREAKING CHANGE: an FTW checkpoint that stores NVFP4 global scales as fp16 is refused at load;
reconvert it with `ft checkpoint`.
---
 docs/ftw-hotfix.md                            |  1 +
 docs/models.md                                |  2 ++
 python/freetoken/checkpoint/ftw.py            | 11 ++++++++
 python/freetoken/kernel/aot_models.py         |  6 ++---
 python/freetoken/kernel/csrc/cpu_moe/common.h | 27 -------------------
 .../kernel/csrc/cpu_moe/kernels/nvfp4.cpp     | 16 +++++------
 .../freetoken/kernel/triton/nvfp4_dequant.py  |  2 +-
 .../kernel/triton/nvfp4_fused_moe.py          | 10 +++----
 .../freetoken/kernel/triton/nvfp4_linear.py   | 14 +++++-----
 .../layers/quantization/linear/nvfp4.py       |  6 ++---
 .../freetoken/layers/quantization/moe/base.py |  4 +--
 .../layers/quantization/moe/nvfp4.py          |  6 ++---
 python/freetoken/models/gemma4/weight.py      |  6 ++---
 python/freetoken/models/glm4_moe/weight.py    |  4 +--
 python/freetoken/models/loader.py             |  4 +--
 .../freetoken/models/muse_glimmer/weight.py   |  2 +-
 python/freetoken/models/nvfp4_banks.py        |  4 +--
 python/freetoken/models/qwen3_5_moe/weight.py |  2 +-
 python/freetoken/moe/benchbw.py               | 12 ++++-----
 python/freetoken/moe/cpu_executor.py          |  4 +--
 python/freetoken/moe/offload_cache.py         |  4 +--
 scripts/ftw_hotfix.py                         |  6 ++++-
 tests/kernels/test_e4m3_compat.py             |  8 +++---
 tests/kernels/test_nvfp4_w4a4_linear.py       |  4 +--
 tests/models/qwen4_exp/test_weight_ckpt.py    |  6 ++---
 tests/models/test_qwen3_5_moe_weight.py       | 12 ++++-----
 tests/moe/test_cpu_moe.py                     |  4 +--
 tests/moe/test_nvfp4_backends.py              | 12 ++++-----
 28 files changed, 95 insertions(+), 104 deletions(-)

diff --git a/docs/ftw-hotfix.md b/docs/ftw-hotfix.md
index 299d522d6e..a88c0066d8 100644
--- a/docs/ftw-hotfix.md
+++ b/docs/ftw-hotfix.md
@@ -32,6 +32,7 @@ byte range, never the whole checkpoint; `--revision` pins the Hub revision to re
 | `KeyError: 'model.embed.weight'` | deepseek-ai/DeepSeek-V4-Flash-0731 | rename the index entries | none |
 | `RuntimeError: Unexpected keys ... .weight_scale` | nvidia/GLM-5.2-NVFP4 | dequantize the old runtime-fp8 weights back to bf16 | none |
 | `PLE shard indices are not contiguous 0..N-1: []` | Qwen3.8-Flash-Next (FTWs converted before #420) | write the PLE table as `ple-table-*.safetensors` | 47.7 GiB |
+| `ValueError: ... stores NVFP4 global scales as fp16` | NVFP4 FTWs converted by older builds | none: the fp16 rounding cannot be undone, so the script refuses; reconvert with `ft checkpoint` | - |
 | `ValueError: ... holds no vision encoder tensors` | image-input families (Qwen3-VL, Qwen3.6, Qwen3.8-Flash-Next, Gemma-4, GLM-5.3-Flash, Muse-Glimmer, MiniMax-M3) converted before the family served images | add the vision encoder: the tower's checkpoint tensors are fetched and the family's own weight reader emits them as the converter would | 1 to 4 GiB |
 
 The script decides by itself: it builds the current model from the FTW's `config.json`, compares
diff --git a/docs/models.md b/docs/models.md
index fd124ea3d2..9c35ece2fd 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -86,6 +86,8 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool.
 - Muse-Glimmer keeps the attention gate as its own linear next to `qkv_proj`; an FTW converted
   when the gate was fused into `qkvg_proj` no longer loads, reconvert it with `ft checkpoint`.
+- NVFP4 global scales are kept in fp32. An NVFP4 FTW that stores them as fp16 (converted by an
+  older build) is refused at load; reconvert it with `ft checkpoint`.
 - An FTW converted before its family served images holds no vision encoder: `ft serve`
   refuses it unless started with `--text-model-only` (or `--mm-disable vision`); reconvert it
   with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md).
diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index ae97a98d5c..ca9c33dc8c 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -102,6 +102,11 @@ def _elsize(dt: torch.dtype) -> int:
     return torch.empty((), dtype=dt).element_size()
 
 
+def _is_nvfp4_global(name: str) -> bool:
+    m = _LAYER_ENTRY_RE.match(name)
+    return (m["base"] if m else name).endswith(("gate_up_global", "down_global", ".weight_global"))
+
+
 def is_ftw_checkpoint(path: str) -> bool:
     """True if ``path`` is a directory holding a FreeToken Weight (FTW) index."""
     return os.path.isfile(os.path.join(path, INDEX_NAME))
@@ -221,6 +226,12 @@ def __init__(self, path: str):
         self.dir = path
         self.shards = sorted(self.index["shards"], key=lambda s: s["global_off"])
         self.tensors = {t["name"]: t for t in self.index["tensors"]}
+        fp16_globals = [n for n, t in self.tensors.items() if t["dtype"] == "float16" and _is_nvfp4_global(n)]
+        if fp16_globals:
+            raise ValueError(
+                f"{path!r} stores NVFP4 global scales as fp16 (e.g. {fp16_globals[0]!r}) but FreeToken reads them as fp32; "
+                "reconvert it from the source checkpoint with `ft checkpoint`"
+            )
         self._fds: dict[str, int] = {}
         self._maps: dict[str, tuple[mmap.mmap, memoryview]] = {}
         # O_DIRECT (DMA straight from disk, bypassing the page cache) is the fast path but a
diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index 1c367aa433..fd67d56541 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -92,7 +92,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
             "down_scale": (H // B) * fp8_block_scale_pad(H // B, I // B) * 2,
         }
     if fmt in ("nvfp4", "nvfp4_marlin", "nvfp4_b12x"):
-        # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp16
+        # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp32
         # per-row globals; marlin/b12x repacks are byte-identical with the globals
         # folded into GPU-resident alphas (layers/quantization/moe/nvfp4.py), so no global banks.
         banks = {
@@ -102,8 +102,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
             "down_scale": H * (I // 16),
         }
         if fmt == "nvfp4":
-            banks["gate_up_global"] = 2 * I * 2
-            banks["down_global"] = H * 2
+            banks["gate_up_global"] = 2 * I * 4
+            banks["down_global"] = H * 4
         return banks
     if fmt == "mxfp4_triton":
         # gpt-oss mxfp4 expert banks: transposed split-K blocks/scales + bf16 bias
diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h
index adf7e14584..0edee0150a 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/common.h
+++ b/python/freetoken/kernel/csrc/cpu_moe/common.h
@@ -92,33 +92,6 @@ constexpr int PF_AHEAD = 512;
 const float kE2M1[16] = {0.0f,  0.5f,  1.0f,  1.5f,  2.0f,  3.0f,  4.0f,  6.0f,
                          -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f};
 
-inline float fp16_to_f32(uint16_t h) {
-  const uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
-  uint32_t exp = (h >> 10) & 0x1Fu;
-  uint32_t man = h & 0x3FFu;
-  uint32_t f;
-  if (exp == 0) {
-    if (man == 0) {
-      f = sign;
-    } else {
-      exp = 127 - 15 + 1;
-      while ((man & 0x400u) == 0) {
-        man <<= 1;
-        --exp;
-      }
-      man &= 0x3FFu;
-      f = sign | (exp << 23) | (man << 13);
-    }
-  } else if (exp == 0x1Fu) {
-    f = sign | 0x7F800000u | (man << 13);
-  } else {
-    f = sign | ((exp + (127 - 15)) << 23) | (man << 13);
-  }
-  float out;
-  std::memcpy(&out, &f, sizeof(out));
-  return out;
-}
-
 // e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT.
 inline float e4m3_decode(uint8_t v) {
   const float sign = (v & 0x80u) ? -1.0f : 1.0f;
diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
index 429955faf3..3c5a56d51f 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp
@@ -7,7 +7,7 @@ namespace {
 
 // --------------------------- NVFP4 (W4A16) dequant ---------------------------
 // Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in
-// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels
+// fp8-e4m3, per-output-row global scale in fp32. Dequant matches the GPU kernels
 // (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global.
 // Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop.
 
@@ -405,7 +405,7 @@ class Nvfp4Format final : public MoeFormat {
     const int H = p.H;
     const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(gate_up_, layer));
     const uint8_t* scale = static_cast<const uint8_t*>(tbl_at(gu_scale_, layer));
-    const uint16_t* global = static_cast<const uint16_t*>(tbl_at(gu_global_, layer));
+    const float* global = static_cast<const float*>(tbl_at(gu_global_, layer));
     const float* xe = xe_.data() + (size_t)tok * (H / 2);
     const float* xo = xo_.data() + (size_t)tok * (H / 2);
     const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr;
@@ -414,9 +414,9 @@ class Nvfp4Format final : public MoeFormat {
       const size_t r = (size_t)e * (2 * p.I) + row;
       if (nvi8dot_ != nullptr)
         return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16),
-                        fp16_to_f32(global[r]), xi8, H, e4m3_lut_, xas);
+                        global[r], xi8, H, e4m3_lut_, xas);
       return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16),
-                    fp16_to_f32(global[r]), xe, xo, H, e2m1_lut_, e4m3_lut_);
+                    global[r], xe, xo, H, e2m1_lut_, e4m3_lut_);
     });
   }
 
@@ -424,15 +424,15 @@ class Nvfp4Format final : public MoeFormat {
     const int I = p.I;
     const uint8_t* packed = static_cast<const uint8_t*>(tbl_at(down_, layer));
     const uint8_t* scale = static_cast<const uint8_t*>(tbl_at(dn_scale_, layer));
-    const uint16_t* global = static_cast<const uint16_t*>(tbl_at(dn_global_, layer));
+    const float* global = static_cast<const float*>(tbl_at(dn_global_, layer));
     for (int h = h0; h < h1; ++h) {
       for (int j = 0; j < n; ++j) {
         const size_t r = (size_t)routes[j].r;
         const size_t row = (size_t)routes[j].e * p.H + h;
         out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr
-            ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]),
+            ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row],
                        gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16))
-            : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]),
+            : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row],
                      ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_);
       }
     }
@@ -441,7 +441,7 @@ class Nvfp4Format final : public MoeFormat {
  private:
   nvdot_fn nvdot_;
   nvi8dot_fn nvi8dot_;  // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot)
-  // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp16 row globals [E,2I]; down likewise with K = I.
+  // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp32 row globals [E,2I]; down likewise with K = I.
   const uint64_t* gate_up_ = nullptr;
   const uint64_t* gu_scale_ = nullptr;
   const uint64_t* gu_global_ = nullptr;
diff --git a/python/freetoken/kernel/triton/nvfp4_dequant.py b/python/freetoken/kernel/triton/nvfp4_dequant.py
index a949e7fa0a..214e6ab71e 100644
--- a/python/freetoken/kernel/triton/nvfp4_dequant.py
+++ b/python/freetoken/kernel/triton/nvfp4_dequant.py
@@ -26,7 +26,7 @@ def _e2m1_lut(device_index: int) -> torch.Tensor:
 def _dequant_nvfp4_kernel(
     packed_ptr,      # [S, OUT, IN // 2] uint8
     scale_ptr,       # [S, OUT, IN // 16] fp8-e4m3 per-block scale (checkpoint native)
-    global_ptr,      # [S, OUT] fp16 per-output-row global scale (weight_scale_2)
+    global_ptr,      # [S, OUT] fp32 per-output-row global scale (weight_scale_2)
     slots_ptr,       # [N] int32 -> cache slot for each output expert
     out_ptr,         # [N, OUT, IN] compute dtype
     lut_ptr,         # [16] float32 E2M1 values
diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
index 61e867b160..08f2f77702 100644
--- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py
+++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py
@@ -1,14 +1,14 @@
 """Inline-dequant NVFP4 fused-MoE Triton kernels.
 
 These kernels read the NVFP4 expert cache (packed FP4 codes + fp8 block scales +
-fp16 per-row global scale) *directly* and dequantize inside the GEMM K-loop, so the
+fp32 per-row global scale) *directly* and dequantize inside the GEMM K-loop, so the
 grouped MoE GEMM never materializes a BF16 copy of the experts (no separate dequant
 pass, no HBM round-trip of the 4x larger BF16 weights).
 
 Layout, for an expert weight ``W[N, K]``:
   - ``packed[slot, n, k//2]`` uint8: two FP4 codes per byte, low nibble = even k.
   - ``scale[slot, n, k//16]`` fp8-e4m3: per 16-wide block scale.
-  - ``global[slot, n]`` fp16: per-output-row scale (``weight_scale_2``).
+  - ``global[slot, n]`` fp32: per-output-row scale (``weight_scale_2``).
   - ``W[n, k] = E2M1[code] * scale[n, k//16] * global[n]``.
 
 The global scale is constant along K, so it is applied once after the K-loop.
@@ -52,7 +52,7 @@ def _decode_nvfp4_moe_kernel(
     a_ptr,             # [M, K] activations (compute dtype)
     packed_ptr,        # [S, N, K // 2] uint8
     scale_ptr,         # [S, N, K // 16] fp8-e4m3
-    global_ptr,        # [S, N] fp16
+    global_ptr,        # [S, N] fp32
     c_ptr,             # [M, TOP_K, N] output (compute dtype)
     topk_weights_ptr,  # [M, TOP_K] fp32
     topk_ids_ptr,      # [M, TOP_K] int32 -> cache slot
@@ -163,7 +163,7 @@ def _decode_nvfp4_marlin_kernel(
     a_ptr,             # [M, K] activations (compute dtype)
     packed_ptr,        # [S, N, K // 8] int32 (8 fp4 codes per word, nibble j -> k=8*w+j)
     scale_ptr,         # [S, N, K // 16] fp8-e4m3
-    global_ptr,        # [S, N] fp16
+    global_ptr,        # [S, N] fp32
     c_ptr,             # [M, TOP_K, N] output (compute dtype)
     topk_weights_ptr,  # [M, TOP_K] fp32
     topk_ids_ptr,      # [M, TOP_K] int32 -> cache slot
@@ -357,7 +357,7 @@ def _prefill_nvfp4_moe_kernel(
     a_ptr,             # [M, K] activations
     packed_ptr,        # [S, N, K // 2] uint8
     scale_ptr,         # [S, N, K // 16] fp8-e4m3 (uint8 view before sm_89)
-    global_ptr,        # [S, N] fp16
+    global_ptr,        # [S, N] fp32
     c_ptr,             # [num_valid_tokens, N] output (flat over M*top_k)
     topk_weights_ptr,
     sorted_token_ids_ptr,
diff --git a/python/freetoken/kernel/triton/nvfp4_linear.py b/python/freetoken/kernel/triton/nvfp4_linear.py
index 53abdedc71..38f148f0ff 100644
--- a/python/freetoken/kernel/triton/nvfp4_linear.py
+++ b/python/freetoken/kernel/triton/nvfp4_linear.py
@@ -158,7 +158,7 @@ def _nvfp4_gemv_kernel(
     a_ptr,        # [K] activation (compute dtype), contiguous
     packed_ptr,   # [N, K // 8] int32 (8 fp4 codes per word, nibble j -> k = 8*w + j)
     scale_ptr,    # [N, K // 16] fp8-e4m3 per-16 block scale
-    gscale_ptr,   # [N] fp16 per-output-row global scale (weight_scale_2)
+    gscale_ptr,   # [N] fp32 per-output-row global scale (weight_scale_2)
     out_ptr,      # [N] out dtype
     N, K,
     stride_pn, stride_pkw,
@@ -226,7 +226,7 @@ def _nvfp4_gemv_splitk_kernel(
     a_ptr,        # [K] activation (compute dtype), contiguous
     packed_ptr,   # [N, K // 8] int32
     scale_ptr,    # [N, K // 16] fp8-e4m3
-    gscale_ptr,   # [N] fp16 per-output-row global scale
+    gscale_ptr,   # [N] fp32 per-output-row global scale
     part_ptr,     # [SPLIT_K, N] fp32 partial sums (pre-global-scale)
     counter_ptr,  # [cdiv(N, BLOCK_N)] int32 arrival counters, zeroed (self-resetting)
     out_ptr,      # [N] out dtype
@@ -311,7 +311,7 @@ def _gemv(a: torch.Tensor, packed_i32: torch.Tensor, scale: torch.Tensor,
           gscale: torch.Tensor, out_dtype: torch.dtype, transposed: bool) -> torch.Tensor:
     """M==1 W4A16 GEMV. ``a`` [K] compute-dtype; ``packed_i32`` logical [N, K//8] int32
     (row-major or a K-major transposed view); ``scale`` logical [N, K//16] fp8; ``gscale``
-    [N] fp16. Picks split-K so small ``N`` (shared expert) still fills the SMs; large ``N``
+    [N] fp32. Picks split-K so small ``N`` (shared expert) still fills the SMs; large ``N``
     (lm_head) stays single-pass (split_k=1)."""
     out_tl = _TL_DTYPE[out_dtype if out_dtype in _TL_DTYPE else torch.bfloat16]
     N = packed_i32.shape[0]
@@ -389,7 +389,7 @@ def _nvfp4_gemm_kernel(
     a_ptr,        # [M, K] activations
     packed_ptr,   # logical [N, K // 8] int32 (row- or K-major via strides)
     scale_ptr,    # logical [N, K // 16] fp8-e4m3
-    gscale_ptr,   # [N] fp16
+    gscale_ptr,   # [N] fp32
     c_ptr,        # [M, N] output (written when SPLIT_K == 1)
     part_ptr,     # [SPLIT_K, M, N] fp32 partials (written when SPLIT_K > 1)
     M, N, K, K_WORDS, tiles_per,
@@ -612,7 +612,7 @@ def _gemm_inkernel(a: torch.Tensor, packed_i32: torch.Tensor, scale: torch.Tenso
 def _nvfp4_dequant_rows_kernel(
     packed_ptr,   # [N, K // 8] int32
     scale_ptr,    # [N, K // 16] fp8-e4m3
-    gscale_ptr,   # [N] fp16
+    gscale_ptr,   # [N] fp32
     out_ptr,      # [N, K] compute dtype
     K_WORDS,
     stride_pn, stride_pkw,
@@ -646,7 +646,7 @@ def _nvfp4_dequant_rows_kernel(
 def _nvfp4_dequant_cols_kernel(
     packed_ptr,   # K-major: element (n, kw) at kw * stride_pkw + n (stride_pn == 1)
     scale_ptr,
-    gscale_ptr,   # [N] fp16
+    gscale_ptr,   # [N] fp32
     out_ptr,      # [n_rows, K] compute dtype, row-major
     N, K_WORDS,
     stride_pn, stride_pkw,
@@ -781,7 +781,7 @@ def nvfp4_dense_linear(
     weight_global: torch.Tensor, bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """``y = x @ dequant(weight)^T`` on the checkpoint-native row-major layout.
-    ``weight`` [N, K//2] uint8, ``weight_scale`` [N, K//16] fp8-e4m3, ``weight_global`` [N] fp16."""
+    ``weight`` [N, K//2] uint8, ``weight_scale`` [N, K//16] fp8-e4m3, ``weight_global`` [N] fp32."""
     return _linear_impl(
         x, weight.view(torch.int32), weight_scale, weight_global, bias, x.dtype, transposed=False
     )
diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py
index 6a13165318..1002a70b3c 100644
--- a/python/freetoken/layers/quantization/linear/nvfp4.py
+++ b/python/freetoken/layers/quantization/linear/nvfp4.py
@@ -1,4 +1,4 @@
-"""e2m1 weight packed two per byte, e4m3 scales per 16 input elements, one fp16 global per output row."""
+"""e2m1 weight packed two per byte, e4m3 scales per 16 input elements, one fp32 global per output row."""
 
 from __future__ import annotations
 
@@ -256,7 +256,7 @@ def finalize(self, layer: Any) -> None:
         dev = layer.weight.device
         layer.weight = torch.empty(0, dtype=torch.uint8, device=dev)
         layer.weight_scale = torch.empty(0, dtype=FP8, device=dev)
-        layer.weight_global = torch.empty(0, dtype=torch.float16, device=dev)
+        layer.weight_global = torch.empty(0, dtype=torch.float32, device=dev)
 
     def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor:
         from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import apply_fp4_marlin_linear
@@ -301,6 +301,6 @@ def create_weights(self, layer: Any) -> None:
             raise ValueError(f"nvfp4 needs in_features divisible by {GROUP}, got {g.in_features}")
         layer.weight = torch.empty(g.out_features, g.in_features // 2, dtype=torch.uint8)
         layer.weight_scale = torch.empty(g.out_features, g.in_features // GROUP, dtype=FP8)
-        layer.weight_global = torch.empty(g.out_features, dtype=torch.float16)
+        layer.weight_global = torch.empty(g.out_features, dtype=torch.float32)
         # no W4A16 kernel reads it; declared so a W4A4 checkpoint loads complete and a W4A4 kernel finds it in place
         layer.input_scale = torch.empty((), dtype=torch.float32) if self.scheme.has("input_scale") else None
diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py
index 5ae89c97fd..b17a91359b 100644
--- a/python/freetoken/layers/quantization/moe/base.py
+++ b/python/freetoken/layers/quantization/moe/base.py
@@ -96,8 +96,8 @@ def fused_piece(pieces: dict[str, torch.Tensor], role: str) -> torch.Tensor:
 
 
 def global_rows(piece: torch.Tensor, rows: int) -> torch.Tensor:
-    """Per-expert global scale as one fp16 value per output row (a scalar per expert is broadcast)."""
-    flat = piece.reshape(piece.shape[0], -1).to(torch.float16)
+    """Per-expert global scale as one fp32 value per output row (a scalar per expert is broadcast)."""
+    flat = piece.reshape(piece.shape[0], -1).to(torch.float32)
     if flat.shape[1] == rows:
         return flat
     assert flat.shape[1] == 1, (flat.shape, rows)
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index 0dce399561..2a5e8bacd0 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/python/freetoken/layers/quantization/moe/nvfp4.py
@@ -52,10 +52,10 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         return {
             "gate_up": BankSpec((2 * i, h // 2), torch.uint8),
             "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8),
-            "gate_up_global": BankSpec((2 * i,), torch.float16),
+            "gate_up_global": BankSpec((2 * i,), torch.float32),
             "down": BankSpec((h, i // 2), torch.uint8),
             "down_scale": BankSpec((h, i // GROUP), FP8),
-            "down_global": BankSpec((h,), torch.float16),
+            "down_global": BankSpec((h,), torch.float32),
         }
 
     def pack(self, pieces, cfg: MoEConfig, out):
@@ -183,7 +183,7 @@ def _marlin_pack_proj(
     """Pack one expert projection (on GPU) into the Marlin layout.
 
     ``packed`` ``[N, K//2]`` uint8 e2m1 pairs, ``scale`` ``[N, K//16]`` e4m3,
-    ``row_global`` ``[N]`` fp16 per-row global (rows of a merged gate_up may carry
+    ``row_global`` ``[N]`` fp32 per-row global (rows of a merged gate_up may carry
     w1's and w3's different globals). Returns ``(qweight [K//16, 2N] int32,
     scales [K//16, N] e4m3-coded, global scalar bf16)``.
 
diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py
index 08293c7eab..7797c7842e 100644
--- a/python/freetoken/models/gemma4/weight.py
+++ b/python/freetoken/models/gemma4/weight.py
@@ -66,15 +66,15 @@
 
 
 def _nvfp4_dense_parts(reader: ShardReader, raw_base: str):
-    """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp16 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none."""
+    """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp32 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none."""
     w = reader.get_tensor(raw_base + ".weight")
     s = reader.get_tensor(raw_base + ".weight_scale")
-    g = reader.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float16)
+    g = reader.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float32)
     g = g.expand(w.shape[0]).contiguous()
     assert (
         w.dtype is torch.uint8
         and s.dtype is torch.float8_e4m3fn
-        and g.dtype is torch.float16
+        and g.dtype is torch.float32
     ), f"unexpected NVFP4 dense dtypes at {raw_base}: {w.dtype}/{s.dtype}/{g.dtype}"
     a = reader.get_tensor(raw_base + ".input_scale").reshape(()).to(torch.float32) if reader.has(raw_base + ".input_scale") else None
     return w, s, g, a
diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py
index aa6538dcf9..0ca79efbd8 100644
--- a/python/freetoken/models/glm4_moe/weight.py
+++ b/python/freetoken/models/glm4_moe/weight.py
@@ -78,12 +78,12 @@ def _iter_nvfp4_resident(
     """Yield native NVFP4 buffers for an always-resident Linear.
 
     ``weight`` (packed fp4) + ``weight_scale`` (fp8 block scale) verbatim; per-tensor
-    ``weight_scale_2`` broadcast to per-row fp16 ``weight_global`` as the NVFP4 linear
+    ``weight_scale_2`` broadcast to per-row fp32 ``weight_global`` as the NVFP4 linear
     method / dequant_nvfp4 expect. lossless vs checkpoint; same dequant math as routed experts.
     """
     packed = reader.get(f"{src_prefix}.weight")  # [OUT, IN//2] uint8
     scale = reader.get(f"{src_prefix}.weight_scale")  # [OUT, IN//16] fp8-e4m3
-    g = reader.get(f"{src_prefix}.weight_scale_2").reshape(()).to(torch.float16)  # scalar
+    g = reader.get(f"{src_prefix}.weight_scale_2").reshape(()).to(torch.float32)  # scalar
     yield f"{dst_prefix}.weight", packed
     yield f"{dst_prefix}.weight_scale", scale
     yield f"{dst_prefix}.weight_global", g.expand(packed.shape[0]).contiguous()
diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py
index 0f310ccbb2..4a9ac94c73 100644
--- a/python/freetoken/models/loader.py
+++ b/python/freetoken/models/loader.py
@@ -208,11 +208,11 @@ def close(self) -> None:
 
 
 def nvfp4_parts_ct(f, raw_base: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], per-output-row global fp16 [O], dequant-side input_scale fp32 scalar or None)`` for the NVFP4 linear method."""
+    """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], per-output-row global fp32 [O], dequant-side input_scale fp32 scalar or None)`` for the NVFP4 linear method."""
     w = f.get_tensor(raw_base + ".weight_packed")
     s = f.get_tensor(raw_base + ".weight_scale")
     wg = f.get_tensor(raw_base + ".weight_global_scale").reshape(1).to(torch.float32)
-    g = (1.0 / wg).to(torch.float16).expand(w.shape[0]).contiguous()
+    g = (1.0 / wg).expand(w.shape[0]).contiguous()
     a = None
     if f.has(raw_base + ".input_global_scale"):
         a = (1.0 / f.get_tensor(raw_base + ".input_global_scale").reshape(()).to(torch.float32))
diff --git a/python/freetoken/models/muse_glimmer/weight.py b/python/freetoken/models/muse_glimmer/weight.py
index 3b8eb6d212..11a1817222 100644
--- a/python/freetoken/models/muse_glimmer/weight.py
+++ b/python/freetoken/models/muse_glimmer/weight.py
@@ -131,7 +131,7 @@ def _iter_weights_compressed_tensors(
 
     Every text Linear (q/k/v/o, the attention gate, the MLP) is kept native NVFP4
     (W4A16): ``.weight`` (uint8 packed) + ``.weight_scale`` (fp8 block) + ``.weight_global``
-    (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v fuse into
+    (fp32 per-row, the reciprocal of the stored quant-side global). q/k/v fuse into
     ``qkv_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping
     its own scales, so the fused FP4 weights are exact. Embeddings, norms, lm_head and the
     vision tower are bf16 (the checkpoint's ignore list). Scale lookups go through the
diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py
index e30429c06a..5382294573 100644
--- a/python/freetoken/models/nvfp4_banks.py
+++ b/python/freetoken/models/nvfp4_banks.py
@@ -36,7 +36,7 @@ def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str:
 def _ingest_global(spec: "Nvfp4ExpertSourceSpec", tensor: torch.Tensor) -> torch.Tensor:
     if spec.global_reciprocal:
         tensor = 1.0 / tensor.float()
-    return tensor.to(torch.float16)
+    return tensor.to(torch.float32)
 
 
 def _activation_scale():
@@ -99,7 +99,7 @@ def iter_nvfp4_expert_pieces(
 ):
     """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale``
     (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects,
-    fp16) companions, straight from the safetensors shards. When every expert also stores the
+    fp32) companions, straight from the safetensors shards. When every expert also stores the
     dialect's static activation scale, the pieces carry it as fp32 ``_input`` (dequant side).
 
     Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index 724ad15312..bf3cecf052 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -227,7 +227,7 @@ def _check(self, target: str, scheme: QuantScheme, part: dict[str, torch.Tensor]
             g = part["weight_global"].reshape(-1).to(torch.float32)
             if g.numel() != 1:
                 raise ValueError(f"{target}: weight_global has {g.numel()} elements, expected one per-tensor scale")
-            out["weight_global"] = g.to(torch.float16).expand(rows).contiguous()
+            out["weight_global"] = g.expand(rows).contiguous()
         if "input_scale" in part:
             out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32)
         return out
diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py
index 057b86fda8..0225cc4fa2 100644
--- a/python/freetoken/moe/benchbw.py
+++ b/python/freetoken/moe/benchbw.py
@@ -309,7 +309,7 @@ def rows(ggml_type: int, n: int, k: int):
 
 def _offload_bank_specs(fmt: str, H: int, I: int, opts: dict | None = None) -> dict[str, tuple[int, torch.dtype]]:
     """Per-bank (elems_per_expert, dtype) for the flat offload host banks (the gather)."""
-    u8, bf16, f16, f8 = torch.uint8, torch.bfloat16, torch.float16, torch.float8_e4m3fn
+    u8, bf16, f32, f8 = torch.uint8, torch.bfloat16, torch.float32, torch.float8_e4m3fn
     if fmt in ("gguf", "gguf_scaled"):
         return {name: (math.prod(shape), dtype) for name, (shape, dtype) in _gguf_bank_shapes(fmt, opts, H, I).items()}
     if fmt == "bf16":
@@ -322,8 +322,8 @@ def _offload_bank_specs(fmt: str, H: int, I: int, opts: dict | None = None) -> d
     if fmt == "nvfp4":
         return {
             "gate_up_packed": (2 * I * (H // 2), u8), "gate_up_scale": (2 * I * (H // 16), u8),
-            "gate_up_global": (2 * I, f16), "down_packed": (H * (I // 2), u8),
-            "down_scale": (H * (I // 16), u8), "down_global": (H, f16),
+            "gate_up_global": (2 * I, f32), "down_packed": (H * (I // 2), u8),
+            "down_scale": (H * (I // 16), u8), "down_global": (H, f32),
         }
     if fmt == "mxfp4_triton":
         return {
@@ -403,13 +403,13 @@ def nb(n):
         b = {
             "gate_up_packed": pin(E, 2 * I, H // 2, dtype=torch.uint8),
             "gate_up_scale": pin(E, 2 * I, H // 16, dtype=torch.uint8),
-            "gate_up_global": pin(E, 2 * I, dtype=torch.float16),
+            "gate_up_global": pin(E, 2 * I, dtype=torch.float32),
             "down_packed": pin(E, H, I // 2, dtype=torch.uint8),
             "down_scale": pin(E, H, I // 16, dtype=torch.uint8),
-            "down_global": pin(E, H, dtype=torch.float16),
+            "down_global": pin(E, H, dtype=torch.float32),
         }
         b["gate_up_global"].fill_(1.0)
-        b["down_global"].fill_(1.0)  # e4m3 scales decode to normal float32; globals are fp16
+        b["down_global"].fill_(1.0)  # e4m3 scales decode to normal float32
         return b
     if fmt == "mxfp4_triton":  # transposed split-K layout
         b = {
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index bb65146b45..1f275749d7 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -448,12 +448,12 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None
         if fmt == "fp8_block":
             return self._resolve_fp8_block_banks(banks, rep)
 
-        # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals.
+        # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp32 row globals.
         gup, gus, gug = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_global"]
         dnp, dns, dng = banks["down"], banks["down_scale"], banks["down_global"]
         assert gup[rep].dtype == torch.uint8 and dnp[rep].dtype == torch.uint8, (gup[rep].dtype, dnp[rep].dtype)
         assert gus[rep].element_size() == 1 and dns[rep].element_size() == 1, "block scales must be 1 byte"
-        assert gug[rep].dtype == torch.float16 and dng[rep].dtype == torch.float16, (gug[rep].dtype, dng[rep].dtype)
+        assert gug[rep].dtype == torch.float32 and dng[rep].dtype == torch.float32, (gug[rep].dtype, dng[rep].dtype)
         I = int(gup[rep].shape[1] // 2)
         H = int(gup[rep].shape[2] * 2)
         assert gup[rep].shape[1] == 2 * I
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 6b428c416b..0d702d8a68 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -51,7 +51,7 @@
     # GGUF experts whose NVFP4 tensors carry per-expert scales: global [L*E, 4] fp32 (gate, up, down, 1) weight_scale_2
     "gguf_scaled": ("gate_up", "down", "global"),
     # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes +
-    # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct
+    # fp8-e4m3 per-16 block scales + per-output-row fp32 globals (w1/w3 carry distinct
     # globals, and folding them into the e4m3 block scales would underflow)
     "nvfp4": (
         "gate_up_packed",
@@ -94,7 +94,7 @@
         (2 * I // 128) * fp8_block_scale_pad(2 * I // 128, H // 128)
         + (H // 128) * fp8_block_scale_pad(H // 128, I // 128)
     ) * 2,
-    "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2),
+    "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 4) + H * (I // 2 + I // 16 + 4),
     "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2),
     "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32),
 }
diff --git a/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py
index 12574651e2..7b4d2257a2 100644
--- a/scripts/ftw_hotfix.py
+++ b/scripts/ftw_hotfix.py
@@ -38,7 +38,7 @@
 
 import torch
 
-from freetoken.checkpoint.ftw import ALIGN, FORMAT_TAG, INDEX_NAME, _align_up, _dtype_of, _dtype_str, _SHARD_FMT
+from freetoken.checkpoint.ftw import ALIGN, FORMAT_TAG, INDEX_NAME, _align_up, _dtype_of, _dtype_str, _is_nvfp4_global, _SHARD_FMT
 from freetoken.models.config import VISION_KEY_PREFIXES
 from freetoken.distributed.info import set_tp_info, try_get_tp_info
 from freetoken.engine.config import EngineConfig
@@ -340,6 +340,10 @@ def validate_ftw(ftw_dir: str, index: dict, expected: dict, renames: dict[str, s
         # the loader casts between float types; anything else has to match exactly
         if dt != exp[1] and not (dt in _FLOAT_DTYPES and exp[1] in _FLOAT_DTYPES):
             problems.append(f"{name}: FTW dtype {e['dtype']} is not loadable as {exp[1]}")
+    fp16_globals = [e["name"] for e in index["tensors"] if e["dtype"] == "float16" and _is_nvfp4_global(e["name"])]
+    if fp16_globals:
+        problems.append(f"{len(fp16_globals)} NVFP4 global scale entries are fp16 (e.g. {fp16_globals[0]}) but FreeToken reads them as fp32; "
+                        "this script cannot repair that, reconvert the FTW with `ft checkpoint`")
     blocks = sorted((e["global_off"], e["global_off"] + _align_up(e["nbytes"]), e["name"]) for e in index["tensors"])
     for (_, a1, an), (b0, _, bn) in zip(blocks, blocks[1:]):
         if b0 < a1:
diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py
index fa8a1f1d32..1532feb406 100644
--- a/tests/kernels/test_e4m3_compat.py
+++ b/tests/kernels/test_e4m3_compat.py
@@ -217,7 +217,7 @@ def _emit_all(path: str) -> None:
     N, K = 128, 256
     pk = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device=dev)
     sc = (torch.rand(N, K // 16, device=dev) * 2 + 0.1).to(FP8)
-    g = (torch.rand(N, device=dev) * 0.05 + 0.01).to(torch.float16)
+    g = torch.rand(N, device=dev) * 0.05 + 0.01
     xa = torch.randn(1, K, dtype=torch.bfloat16, device=dev)
     xb = torch.randn(8, K, dtype=torch.bfloat16, device=dev)
     xc = torch.randn(200, K, dtype=torch.bfloat16, device=dev)
@@ -231,7 +231,7 @@ def _emit_all(path: str) -> None:
     # nvfp4 bank dequant
     pkc = torch.randint(0, 256, (E, 64, 128), dtype=torch.uint8, device=dev)
     scc = (torch.rand(E, 64, 16, device=dev) + 0.1).to(FP8)
-    gc = (torch.rand(E, 64, device=dev) * 0.05 + 0.01).to(torch.float16)
+    gc = torch.rand(E, 64, device=dev) * 0.05 + 0.01
     out["nv_dequant"] = f32(dequant_nvfp4(
         pkc, scc, gc, torch.tensor([1, 3], dtype=torch.int32, device=dev)))
 
@@ -239,10 +239,10 @@ def _emit_all(path: str) -> None:
     S, H4, I4 = 4, 256, 128
     gup = torch.randint(0, 256, (S, 2 * I4, H4 // 2), dtype=torch.uint8, device=dev)
     gus4 = (torch.rand(S, 2 * I4, H4 // 16, device=dev) + 0.1).to(FP8)
-    gug = (torch.rand(S, 2 * I4, device=dev) * 0.05 + 0.01).to(torch.float16)
+    gug = torch.rand(S, 2 * I4, device=dev) * 0.05 + 0.01
     dnp = torch.randint(0, 256, (S, H4, I4 // 2), dtype=torch.uint8, device=dev)
     dns4 = (torch.rand(S, H4, I4 // 16, device=dev) + 0.1).to(FP8)
-    dng = (torch.rand(S, H4, device=dev) * 0.05 + 0.01).to(torch.float16)
+    dng = torch.rand(S, H4, device=dev) * 0.05 + 0.01
     hid4 = torch.randn(2, H4, dtype=torch.bfloat16, device=dev)
     out["nv_moe_marlin"] = f32(fused_experts_decode_nvfp4_marlin(
         hid4, gup, gus4, gug, dnp, dns4, dng, tw, tids))
diff --git a/tests/kernels/test_nvfp4_w4a4_linear.py b/tests/kernels/test_nvfp4_w4a4_linear.py
index 1840d30f5a..03fd2c7e63 100644
--- a/tests/kernels/test_nvfp4_w4a4_linear.py
+++ b/tests/kernels/test_nvfp4_w4a4_linear.py
@@ -86,7 +86,7 @@ def test_w4a4_linear_matches_nvfp4_activation_reference(kernel_name, M, N, K, pa
     g = torch.Generator(device="cuda").manual_seed(0)
     weight = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda", generator=g)
     scale = (torch.rand(N, K // 16, device="cuda", generator=g) + 0.25).to(torch.float8_e4m3fn)
-    row_global = torch.full((N,), 1e-3, dtype=torch.float16, device="cuda")
+    row_global = torch.full((N,), 1e-3, dtype=torch.float32, device="cuda")
     if parts_differ:
         row_global[N // 2:] = 4e-3
     bias = torch.randn(N, device="cuda", dtype=torch.bfloat16)
@@ -149,7 +149,7 @@ def test_w4a4_linear_takes_row_strided_and_empty_inputs(kernel_name, M):
     g = torch.Generator(device="cuda").manual_seed(0)
     layer = SimpleNamespace(weight=torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda", generator=g),
                             weight_scale=(torch.rand(N, K // 16, device="cuda", generator=g) + 0.25).to(torch.float8_e4m3fn),
-                            weight_global=torch.full((N,), 1e-3, dtype=torch.float16, device="cuda"),
+                            weight_global=torch.full((N,), 1e-3, dtype=torch.float32, device="cuda"),
                             input_scale=torch.tensor(0.01, device="cuda"), bias=None)
     kernel.finalize(layer)
     buf = torch.randn(M, K + 64, device="cuda", dtype=torch.bfloat16, generator=g)
diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py
index 605c40f18b..4395a86394 100644
--- a/tests/models/qwen4_exp/test_weight_ckpt.py
+++ b/tests/models/qwen4_exp/test_weight_ckpt.py
@@ -263,7 +263,7 @@ def test_ple_table_loads_pinned_and_matches_the_checkpoint(reader):
 
     # The two pinned host allocations the engine must budget for.
     total = PLE_BYTES + NUM_LAYERS * EXPERT_LAYER_BYTES
-    assert abs(total / 2**30 - 111.14) < 0.05
+    assert abs(total / 2**30 - 111.32) < 0.05
 
 
 @pytest.fixture(scope="module")
@@ -312,14 +312,14 @@ def test_sampled_experts_match_the_checkpoint(layer0_expert_banks, reader):
             scale = reader.get(f"{base}.{proj}.weight_scale")
             assert torch.equal(banks[bank][0][expert][rows].reshape(-1).view(torch.uint8),
                                scale.reshape(-1).view(torch.uint8))
-        gate_g = reader.get(f"{base}.gate_proj.weight_scale_2").to(torch.float16)
+        gate_g = reader.get(f"{base}.gate_proj.weight_scale_2").to(torch.float32)
         assert torch.equal(banks["gate_up_global"][0][expert, :I], gate_g.reshape(1).expand(I))
 
 
 def test_expert_bank_bytes_match_the_aot_row_table(layer0_expert_banks):
     measured = sum(t[0].numel() * t[0].element_size() for t in layer0_expert_banks.values())
     assert measured == EXPERT_LAYER_BYTES
-    assert abs(NUM_LAYERS * EXPERT_LAYER_BYTES / 2**30 - 63.46) < 0.05
+    assert abs(NUM_LAYERS * EXPERT_LAYER_BYTES / 2**30 - 63.63) < 0.05
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="dummy banks are pinned")
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index 16f439a9ec..c502067ee4 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -431,7 +431,7 @@ def test_modelopt_fp8_scales_broadcast_per_part_and_input_scale_is_the_max(check
     assert _same(loaded[f"{fused}.weight"][:I], raw[f"{shared}.gate_proj.weight"])
     assert _same(loaded[f"{fused}.weight_scale"][I:], raw[f"{shared}.up_proj.weight_scale"])
     glob = loaded[f"{fused}.weight_global"]
-    assert glob.dtype is torch.float16 and torch.equal(glob[:I], raw[f"{shared}.gate_proj.weight_scale_2"].to(torch.float16).expand(I))
+    assert glob.dtype is torch.float32 and torch.equal(glob[:I], raw[f"{shared}.gate_proj.weight_scale_2"].float().expand(I))
     # W4A16: no activation quantizer, so the layer has no input_scale even when the export ships one
     assert f"{fused}.input_scale" not in loaded and "lm_head.input_scale" not in loaded
     assert loaded["lm_head.weight"].dtype is torch.uint8 and loaded["lm_head.weight_global"].shape == (V,)
@@ -444,8 +444,8 @@ def test_compressed_tensors_nvfp4_global_is_the_reciprocal(checkpoint):
     loaded = _load(folder)
     attn = f"{LM}.layers.1.self_attn"
     glob = loaded["model.layers.1.self_attn.qkv_proj.weight_global"]
-    expected = torch.cat([(1.0 / raw[f"{attn}.{p}_proj.weight_global_scale"]).to(torch.float16).expand(raw[f"{attn}.{p}_proj.weight_packed"].shape[0]) for p in "qkv"])
-    assert glob.dtype is torch.float16 and torch.equal(glob, expected)
+    expected = torch.cat([(1.0 / raw[f"{attn}.{p}_proj.weight_global_scale"].float()).expand(raw[f"{attn}.{p}_proj.weight_packed"].shape[0]) for p in "qkv"])
+    assert glob.dtype is torch.float32 and torch.equal(glob, expected)
     assert _same(loaded["model.layers.1.self_attn.o_proj.weight_scale"], raw[f"{attn}.o_proj.weight_scale"])
     # ``dynamic: local`` ships the quant-side activation global too; it lands as the dequant-side input_scale
     assert torch.equal(loaded["model.layers.1.self_attn.o_proj.input_scale"], (1.0 / raw[f"{attn}.o_proj.input_global_scale"]).reshape(()))
@@ -514,11 +514,11 @@ def test_nvfp4_expert_pieces_read_either_dialect_from_a_single_file(checkpoint):
     base = f"{LM}.layers.1.mlp.experts.2.gate_proj"
     if name.startswith("modelopt_mixed"):
         assert _same(piece["gate"][0], raw[f"{base}.weight"])
-        assert torch.equal(piece["gate_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).to(torch.float16))
+        assert torch.equal(piece["gate_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).float())
     else:
         assert _same(piece["gate"][0], raw[f"{base}.weight_packed"])
         assert _same(piece["down_scale"][0], raw[f"{LM}.layers.1.mlp.experts.2.down_proj.weight_scale"])
-        assert torch.equal(piece["gate_global"].reshape(-1), (1.0 / raw[f"{base}.weight_global_scale"]).to(torch.float16))
+        assert torch.equal(piece["gate_global"].reshape(-1), 1.0 / raw[f"{base}.weight_global_scale"].float())
 
 
 @pytest.mark.parametrize("parallel", [False, True])
@@ -563,7 +563,7 @@ def test_a_quantized_router_is_dequantized_because_the_family_serves_it_bf16(tmp
     for layer, gate in enumerate(gates):
         got = loaded[f"model.layers.{layer}.mlp.gate.weight"]
         assert got.dtype is torch.bfloat16 and got.shape == (E, H)
-        expected = _nvfp4_reference(raw[f"{gate}.weight_packed"], raw[f"{gate}.weight_scale"], (1.0 / raw[f"{gate}.weight_global_scale"]).to(torch.float16).expand(E))
+        expected = _nvfp4_reference(raw[f"{gate}.weight_packed"], raw[f"{gate}.weight_scale"], (1.0 / raw[f"{gate}.weight_global_scale"].float()).expand(E))
         assert torch.equal(got, expected)
 
 
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index 56c6c18aa4..edd9ddc9df 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -90,9 +90,9 @@ def _make_nvfp4_cache(L, E, H, I, seed=0):
     def rows(OUT, IN):
         codes = torch.randint(0, 16, (S, OUT, IN), dtype=torch.uint8)
         packed = _pack_nvfp4(codes)  # [S, OUT, IN//2]
-        # block scales in fp8-e4m3 (1 byte), ~[0.5, 1.5]; row globals in fp16, small.
+        # block scales in fp8-e4m3 (1 byte), ~[0.5, 1.5]; row globals in fp32, small.
         scale = (0.5 + torch.rand(S, OUT, IN // 16)).to(torch.float8_e4m3fn)
-        glob = (0.02 + 0.01 * torch.rand(S, OUT)).to(torch.float16)
+        glob = 0.02 + 0.01 * torch.rand(S, OUT)
         return packed, scale.contiguous(), glob.contiguous()
 
     gup, gus, gug = rows(2 * I, H)
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index fd98c2ee94..58a55bfa64 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -64,9 +64,9 @@ def rand_u8(*shape):
     def rand_scale(*shape):
         return (torch.rand(*shape, generator=g) * 1.5 + 0.25).to(torch.float8_e4m3fn)
 
-    gate_up_global = torch.full((total, 2 * I), 1.0, dtype=torch.float16)
+    gate_up_global = torch.full((total, 2 * I), 1.0, dtype=torch.float32)
     gate_up_global[:, I:] = 0.5  # w3 global != w1 global: exercises the alpha fold
-    down_global = torch.full((total, H), 0.75, dtype=torch.float16)
+    down_global = torch.full((total, H), 0.75, dtype=torch.float32)
     flat = {
         "gate_up_packed": rand_u8(total, 2 * I, H // 2),
         "gate_up_scale": rand_scale(total, 2 * I, H // 16),
@@ -392,7 +392,7 @@ def test_triton_prefill_routes_to_gluon_on_sm120(monkeypatch, hidden_size, inter
     def bank(n, k):
         packed = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=device, generator=g)
         scale = (torch.rand(E, n, k // 16, device=device, generator=g) + 0.25).to(torch.float8_e4m3fn)
-        return packed, scale, (torch.rand(E, n, device=device, generator=g) + 0.5).to(torch.float16)
+        return packed, scale, torch.rand(E, n, device=device, generator=g) + 0.5
 
     banks = (*bank(2 * inter, hidden_size), *bank(hidden_size, inter))
     M = 1024
@@ -471,7 +471,7 @@ def test_w4a4_prefill_matches_nvfp4_activation_reference(hidden_size, inter, act
     def bank(n, k):
         packed = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=device, generator=g)
         scale = (torch.rand(E, n, k // 16, device=device, generator=g) + 0.25).to(torch.float8_e4m3fn)
-        return packed, scale, (torch.rand(E, n, device=device, generator=g) * 0.02 + 0.01).to(torch.float16)
+        return packed, scale, torch.rand(E, n, device=device, generator=g) * 0.02 + 0.01
 
     gu, gus, gug = bank(2 * inter, hidden_size)
     dn, dns, dng = bank(hidden_size, inter)
@@ -621,10 +621,10 @@ def test_dummy_nvfp4_banks_match_the_triton_layout():
     expected = {
         "gate_up": ((E, 2 * I, H // 2), torch.uint8),
         "gate_up_scale": ((E, 2 * I, H // 16), torch.float8_e4m3fn),
-        "gate_up_global": ((E, 2 * I), torch.float16),
+        "gate_up_global": ((E, 2 * I), torch.float32),
         "down": ((E, H, I // 2), torch.uint8),
         "down_scale": ((E, H, I // 16), torch.float8_e4m3fn),
-        "down_global": ((E, H), torch.float16),
+        "down_global": ((E, H), torch.float32),
     }
     assert banks.sources.keys() == expected.keys()
     for name, (shape, dtype) in expected.items():

From 5596e6b3904e6d8c64e2bc9e67ee347b2ab7eedc Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 14:35:22 -0700
Subject: [PATCH 128/249] feat(gguf): add scripts that check which gguf files
 load (#756)

---
 scripts/gguf_check/README.md        |  42 ++
 scripts/gguf_check/check.py         | 593 ++++++++++++++++++++++++++++
 scripts/gguf_check/fetch_headers.py | 193 +++++++++
 scripts/gguf_check/probe_kernels.py | 111 ++++++
 scripts/gguf_check/tok_compare.py   |  90 +++++
 5 files changed, 1029 insertions(+)
 create mode 100644 scripts/gguf_check/README.md
 create mode 100644 scripts/gguf_check/check.py
 create mode 100644 scripts/gguf_check/fetch_headers.py
 create mode 100644 scripts/gguf_check/probe_kernels.py
 create mode 100644 scripts/gguf_check/tok_compare.py

diff --git a/scripts/gguf_check/README.md b/scripts/gguf_check/README.md
new file mode 100644
index 0000000000..a1f369f7ab
--- /dev/null
+++ b/scripts/gguf_check/README.md
@@ -0,0 +1,42 @@
+# gguf_check
+
+Checks for GGUF support of a model family: does FreeToken load most of the GGUF files people publish for it, and does it read them right.
+Run from the repo root with `PYTHONPATH=python` and transformers 5.18.
+
+| Script | Needs | Checks |
+|---|---|---|
+| `fetch_headers.py` | network | nothing, writes header-only stubs of HF GGUF files |
+| `probe_kernels.py` | GPU | which ggml ops this build runs per ggml type |
+| `check.py` | CPU | config, dry load, expert banks, kernel coverage, tokenizer, chat template per file |
+| `tok_compare.py` | CPU | GGUF tokenizer vs the model's HF `tokenizer.json` |
+
+## Coverage: most published files load
+
+```bash
+# header-only stubs of every GGUF in the repos (about 15 MB each, weights read as zeros)
+python scripts/gguf_check/fetch_headers.py ~/gguf-stubs --search GLM-4.5 --search GLM-4.7 unsloth/GLM-4.7-GGUF > fetch.jsonl
+
+# once per build and GPU architecture
+python scripts/gguf_check/probe_kernels.py kernels.json
+
+python scripts/gguf_check/check.py --kernels kernels.json --out glm.jsonl ~/gguf-stubs
+python scripts/gguf_check/check.py --report glm.jsonl
+```
+
+`check.py` builds the model on the meta device and runs the family's weight reader on meta matrices, so a worker needs no RAM for the weights; it falls back to real reads only for a reader that needs the values.
+It prints `passed/total` per served GGUF architecture, the files of architectures FreeToken does not serve (draft models, ...), and the failures grouped by their first problem.
+Every failure should be one of:
+
+- the family's loader refuses the file with a clear error (a type or layout we do not serve), or
+- the file is broken.
+
+List the refused kinds under "Not covered" in the PR.
+Passing means the server builds the model and finds every tensor and kernel. It does not read the weights or check the model fits in memory.
+
+## Correctness: the files load right
+
+On one or two real files per family, smallest and most common quant:
+
+- `tok_compare.py <gguf> <HF model dir>`: 0 differing texts. The reference is the model's own tokenizer, not llama.cpp.
+- `ft serve --model-path <file>`: a short eval (e.g. AIME problems) against llama.cpp on the same file, greedy outputs against llama.cpp, 48 staggered requests at concurrency 16 with no empty or corrupted output, and bs=1 tok/s.
+- `pytest tests/models tests/moe tests/kernels/test_ggml.py tests/server`.
diff --git a/scripts/gguf_check/check.py b/scripts/gguf_check/check.py
new file mode 100644
index 0000000000..390a52b41a
--- /dev/null
+++ b/scripts/gguf_check/check.py
@@ -0,0 +1,593 @@
+"""Check that FreeToken loads GGUF files, on the CPU and without reading their weights.
+
+    python scripts/gguf_check/check.py --kernels kernels.json --out results.jsonl <gguf file or dir> ...
+    python scripts/gguf_check/check.py --report results.jsonl
+
+Per file, through the code the server runs:
+
+- config: the family's ``parse_gguf_config`` accepts the file (unsupported files must fail here with a clear error)
+- dry load: the model builds on the meta device, and the family's weight reader yields exactly its state dict (names, shapes, packed vs float); the expert banks build for every MoE layer with the shapes the model needs
+- kernels: every quantized ggml type a GGUF linear, embedding or expert bank uses has the ops it needs in the ``probe_kernels.py`` table
+- tokenizer: it builds, the chat template renders, and the file names an EOS token
+
+The weight reader runs on meta tensors, so no weight bytes are read or copied. A reader that needs real values (a numpy dequant) gets a second run on the real bytes, once enough RAM is free.
+The files can be real or ``fetch_headers.py`` stubs. Passing does not prove the weights are right or that the model fits in memory.
+Results append to ``--out``; a rerun skips the files already there, except crashed ones (and failed ones with ``--retry-failed``).
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures as cf
+import contextlib
+import functools
+import glob
+import hashlib
+import json
+import logging
+import multiprocessing as mp
+import os
+import re
+import struct
+import sys
+import time
+import traceback
+import warnings
+from collections import Counter
+
+MOE_TOKENS = ("mmid1", "mmid8", "mmid16", "mmid512")
+# ops each layer kind runs over the token counts the server sends; fused-glu experts also need glu<T> up to the type's mmvq_max_tokens
+REQUIRED_OPS = {
+    "linear": ("mm1", "mm8", "mm16", "mm512"),
+    "embedding": ("get_rows",),
+    "experts gate_up": MOE_TOKENS,
+    "experts down": MOE_TOKENS,
+}
+GLU_TOKENS = (1, 8)
+_SPLIT = re.compile(r"-(\d{5})-of-(\d{5})\.gguf$")
+_SCALAR = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", 7: "?", 10: "Q", 11: "q", 12: "d"}
+_STRING, _ARRAY = 8, 9
+_UINT32 = 4
+
+
+class _Field:
+    __slots__ = ("value",)
+
+    def __init__(self, value):
+        self.value = value
+
+    def contents(self):
+        return self.value
+
+
+class _Tensor:
+    __slots__ = ("name", "shape", "tensor_type", "data_offset", "n_bytes")
+
+
+class FastReader:
+    """The parts of gguf-py's ``GGUFReader`` FreeToken reads, parsed in pure Python, with the same checks.
+
+    gguf-py builds numpy views for every field and takes seconds and about 0.7 GB per large file, which a scan over many files cannot afford."""
+
+    def __init__(self, path: str):
+        import numpy as np
+
+        from freetoken.third_party.gguf import GGML_QUANT_SIZES, GGMLQuantizationType
+
+        self.data = data = np.memmap(path, dtype=np.uint8, mode="r")
+        pos = 0
+
+        def unpack(fmt):
+            nonlocal pos
+            vals = struct.unpack_from("<" + fmt, data, pos)
+            pos += struct.calcsize("<" + fmt)
+            return vals
+
+        def string():
+            nonlocal pos
+            (n,) = unpack("Q")
+            pos += n
+            return data[pos - n : pos].tobytes().decode("utf-8")
+
+        def value(kind):
+            if kind in _SCALAR:
+                return unpack(_SCALAR[kind])[0]
+            if kind == _STRING:
+                return string()
+            if kind != _ARRAY:
+                raise ValueError(f"{path}: unknown GGUF value type {kind}")
+            elem, n = unpack("I")[0], unpack("Q")[0]
+            if elem in _SCALAR:
+                return list(unpack(f"{n}{_SCALAR[elem]}"))
+            return [value(elem) for _ in range(n)]
+
+        if data[:4].tobytes() != b"GGUF":
+            raise ValueError(f"{path} is not a GGUF file")
+        pos = 4
+        _, n_tensors, n_kv = unpack("IQQ")
+        self.fields = {}
+        align = 32
+        for _ in range(n_kv):
+            key = string()
+            if key in self.fields:
+                raise KeyError(f"{path}: duplicate GGUF key {key}")
+            kind = unpack("I")[0]
+            self.fields[key] = _Field(value(kind))
+            if key == "general.alignment":
+                align = self.fields[key].value
+                if kind != _UINT32 or align <= 0 or align & (align - 1):
+                    raise ValueError(f"{path}: bad general.alignment {align}")
+        infos = []
+        for _ in range(n_tensors):
+            name = string()
+            (nd,) = unpack("I")
+            dims = unpack(f"{nd}Q")
+            kind, offset = unpack("IQ")
+            infos.append((name, dims, kind, offset))
+        start = (pos + align - 1) // align * align
+        self.tensors, names = [], set()
+        for name, dims, kind, offset in infos:
+            if name in names:
+                raise ValueError(f"{path}: duplicate tensor {name}")
+            names.add(name)
+            block, size = GGML_QUANT_SIZES[GGMLQuantizationType(kind)]
+            t = _Tensor()
+            t.name, t.shape, t.tensor_type, t.data_offset = name, list(dims), kind, start + offset
+            t.n_bytes = _prod(dims) // block * size
+            if t.data_offset + t.n_bytes > len(data):
+                raise ValueError(f"{path}: tensor {name} runs past the end of the file")
+            self.tensors.append(t)
+
+
+def _prod(dims) -> int:
+    n = 1
+    for d in dims:
+        n *= d
+    return n
+
+
+def _clear_caches() -> None:
+    """Drop the caches that hold a file's metadata, config or rope tables."""
+    import freetoken.models.gguf.reader as reader
+    from freetoken.layers import rotary
+    from freetoken.utils import hf
+
+    for fn in (reader._reader, reader.load_gguf_metadata, reader.gguf_split_paths, reader.gguf_tensor_index,
+               rotary.get_rope, rotary._cos_sin_table, hf._load_hf_config):
+        fn.cache_clear()
+
+
+_KERNELS: dict | None = None
+_TOKENIZERS: dict = {}
+_IN_FLIGHT = None
+_MIN_FREE_GB = 32.0
+
+
+def _init_worker(kernels: dict | None, in_flight, min_free_gb: float) -> None:
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    warnings.filterwarnings("ignore")
+    logging.disable(logging.WARNING)
+    import torch
+
+    import freetoken.models.gguf.reader as reader
+    from freetoken.distributed import set_tp_info
+    from freetoken.layers import rotary
+
+    reader._reader = functools.cache(FastReader)
+    set_tp_info(rank=0, size=1)
+    rotary.set_rope_device(torch.device("cpu"))
+    global _KERNELS, _IN_FLIGHT, _MIN_FREE_GB
+    _KERNELS, _IN_FLIGHT, _MIN_FREE_GB = kernels, in_flight, min_free_gb
+
+
+def _build_model(path: str):
+    """The engine's config and the model it builds for ``path`` (experts offloaded), on the meta device."""
+    import torch
+
+    from freetoken.distributed import try_get_tp_info
+    from freetoken.engine.config import EngineConfig
+    from freetoken.engine.engine import _decode_target
+    from freetoken.mm.config import ENCODER_KINDS, MultimodalConfig
+    from freetoken.models import create_model
+    from freetoken.utils.torch_utils import torch_dtype
+
+    config = EngineConfig(model_path=path, tp_info=try_get_tp_info(), dtype=torch.bfloat16, moe_strategy="offload",
+                          mm=MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS)))
+    object.__setattr__(config.model_config, "moe_strategy", "offload")
+    object.__setattr__(config.model_config, "decode_target", _decode_target(config))
+    with torch.device("meta"), torch_dtype(torch.bfloat16):
+        return config, create_model(config.model_config)
+
+
+def _gguf_layers(model) -> tuple[list, bool]:
+    """(kind, ggml types) of every GGUF linear and embedding, and whether any GGUF MoE layer runs the fused-glu kernel."""
+    from freetoken.layers.base import BaseOP
+    from freetoken.layers.quantization import QuantKind
+    from freetoken.layers.quantization.moe.gguf import fused_glu
+    from freetoken.layers.quantization.registry import LayerKind
+    from freetoken.layers.quantization.scheme import gguf_types
+
+    kinds = {LayerKind.LINEAR: "linear", LayerKind.EMBEDDING: "embedding"}
+    seen, out, glu = set(), [], False
+
+    def walk(op):
+        nonlocal glu
+        if id(op) in seen:
+            return
+        seen.add(id(op))
+        method = getattr(op, "quant_method", None)
+        if method is not None and method.kind is QuantKind.GGUF:
+            if method.layer_kind in kinds:
+                out.append((kinds[method.layer_kind], gguf_types(method.scheme)))
+            elif method.layer_kind is LayerKind.MOE:
+                glu = glu or fused_glu(op) is not None
+        for value in vars(op).values():
+            for item in value if isinstance(value, (list, tuple)) else (value,):
+                if isinstance(item, BaseOP):
+                    walk(item)
+
+    walk(model)
+    return out, glu
+
+
+def _check_kernels(layers, tied: bool, glu: bool, expert_tags, problems: list[str]) -> dict:
+    from freetoken.layers.quantization.ggml_types import is_quantized, type_name
+    from freetoken.third_party.gguf import GGMLQuantizationType
+
+    used: dict[str, Counter] = {}
+    if tied:
+        layers = layers + [("linear", types) for kind, types in layers if kind == "embedding"]
+    for kind, types in layers:
+        for t in types:
+            used.setdefault(kind, Counter())[type_name(t)] += 1
+    for _, gate_up, down in expert_tags:
+        used.setdefault("experts gate_up", Counter())[gate_up] += 1
+        used.setdefault("experts down", Counter())[down] += 1
+    if _KERNELS is not None:
+        for kind, names in used.items():
+            for name in names:
+                if not is_quantized(int(GGMLQuantizationType[name])):
+                    continue
+                probe = _KERNELS.get(name, {})
+                need = list(REQUIRED_OPS[kind])
+                if kind == "experts gate_up" and glu:
+                    need += [f"glu{t}" for t in GLU_TOKENS if t <= probe.get("mmvq_max_tokens", 0)]
+                missing = [op for op in need if probe.get(op) != "ok"]
+                if missing:
+                    problems.append(f"{kind} {name}: no {', '.join(missing)} kernel")
+    return {k: dict(v) for k, v in used.items()}
+
+
+@functools.cache
+def _meta_readers():
+    """The GGUF tensor readers, and stand-ins that give matrices as meta tensors."""
+    import torch
+
+    from freetoken.models.gguf import tensors
+    from freetoken.models.gguf.tensors import _FLOAT_DTYPES
+
+    real_float = tensors.float_tensor
+
+    def float_tensor(info):
+        # 1-D tensors (norms, biases, rope_freqs) are small, and config parses read their values
+        if len(info.ne) == 1:
+            return real_float(info)
+        return torch.empty(tuple(reversed(info.ne)), dtype=_FLOAT_DTYPES[info.ggml_type], device="meta")
+
+    def packed_rows(info):
+        rows = _prod(info.ne[1:])
+        return torch.empty(rows, info.nbytes // rows, dtype=torch.uint8, device="meta")
+
+    fakes = {"float_tensor": float_tensor, "packed_rows": packed_rows}
+    real = {name: getattr(tensors, name) for name in fakes}
+    return real, fakes
+
+
+@functools.cache
+def _meta_wins():
+    """A torch function mode that moves CPU operands to meta when an op mixes the two, so a reader can combine meta weights with the small tensors it builds."""
+    import torch
+    from torch.overrides import TorchFunctionMode
+    from torch.utils._pytree import tree_flatten, tree_unflatten
+
+    class MetaWins(TorchFunctionMode):
+        def __torch_function__(self, func, types, args=(), kwargs=None):
+            flat, spec = tree_flatten((args, kwargs or {}))
+            if any(isinstance(x, torch.Tensor) and x.is_meta for x in flat):
+                flat = [x.to("meta") if isinstance(x, torch.Tensor) and not x.is_meta else x for x in flat]
+            args, kwargs = tree_unflatten(flat, spec)
+            return func(*args, **kwargs)
+
+    return MetaWins
+
+
+def _swap_readers(use_fakes: bool) -> None:
+    """Point every loaded binding of the readers, including the family modules' imported names, at the fakes or back."""
+    real, fakes = _meta_readers()
+    src, dst = (real, fakes) if use_fakes else (fakes, real)
+    for module in list(sys.modules.values()):
+        if getattr(module, "__name__", "").startswith("freetoken.models"):
+            for name, fn in src.items():
+                if getattr(module, name, None) is fn:
+                    setattr(module, name, dst[name])
+
+
+def _wait_for_ram(gb: float, timeout: float = 600.0) -> bool:
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        with open("/proc/meminfo") as f:
+            avail = next(int(line.split()[1]) for line in f if line.startswith("MemAvailable:")) / 2**20
+        if avail >= gb:
+            return True
+        time.sleep(2)
+    return False
+
+
+def _dry_load(path: str, config, model, *, meta: bool) -> list[str]:
+    import torch
+
+    from freetoken.models.register import _load_attr
+
+    spec = config.model_spec
+    expected = model.state_dict()
+    problems, seen = [], set()
+    _swap_readers(meta)
+    try:
+        with _meta_wins()() if meta else contextlib.nullcontext():
+            for key, t in _load_attr(spec.module, spec.iter_weights)(path, torch.device("cpu"), include_moe_experts=False, include_non_moe=True):
+                ref = expected.get(key)
+                if ref is None:
+                    problems.append(f"unexpected tensor {key}")
+                    continue
+                if key in seen:
+                    problems.append(f"tensor {key} yielded twice")
+                seen.add(key)
+                if tuple(t.shape) != tuple(ref.shape) or (t.dtype == torch.uint8) != (ref.dtype == torch.uint8):
+                    problems.append(f"{key}: got {tuple(t.shape)} {t.dtype}, model has {tuple(ref.shape)} {ref.dtype}")
+    finally:
+        _swap_readers(False)
+    missing = sorted(set(expected) - seen)
+    if missing:
+        problems.append(f"{len(missing)} tensors never loaded, e.g. {missing[:3]}")
+    return problems
+
+
+def _expert_banks(path: str, cfg, problems: list[str]) -> list:
+    """Build the expert banks on the meta device and check their shapes; returns each MoE layer's (format, gate_up type, down type)."""
+    import torch
+
+    if not cfg.num_experts:
+        return []
+    from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes
+    from freetoken.models.gguf.experts import load_gguf_expert_banks
+    from freetoken.third_party.gguf import GGMLQuantizationType
+
+    banks = load_gguf_expert_banks(path, cfg.num_moe_layers, first_layer=cfg.first_k_dense_replace, dummy=True, device=torch.device("meta"))
+    e, inter, hidden = cfg.num_experts, cfg.moe_intermediate_size, cfg.hidden_size
+    for layer, (_, gate_up, down) in enumerate(banks.layer_tags):
+        for role, name, rows, k in (("gate_up", gate_up, 2 * inter, hidden), ("down", down, hidden, inter)):
+            t = int(GGMLQuantizationType[name])
+            want = (e, rows, k if float_dtype(t) is not None else row_bytes(t, k))
+            got = tuple(banks.sources[role][layer].shape)
+            if got != want:
+                problems.append(f"MoE layer {layer} {role} bank is {got}, the model needs {want}")
+    return banks.layer_tags
+
+
+def _tokenizer_key(meta: dict) -> str:
+    blob = json.dumps(sorted((k, v) for k, v in meta.items() if k.startswith("tokenizer.")))
+    return hashlib.sha1(blob.encode()).hexdigest()[:12]
+
+
+def _check_tokenizer(path: str) -> dict:
+    from freetoken.models.gguf.tokenizer import gguf_eos_token_ids, load_gguf_tokenizer
+
+    problems = []
+    tok = load_gguf_tokenizer(path)
+    eos = sorted(gguf_eos_token_ids(path))
+    if not eos:
+        problems.append("tokenizer: no EOS token")
+    tail = None
+    if tok.chat_template is None:
+        problems.append("tokenizer: no chat template")
+    else:
+        try:
+            tail = tok.apply_chat_template([{"role": "user", "content": "hi"}], tokenize=False, add_generation_prompt=True)[-60:]
+        except Exception as e:
+            problems.append(f"chat template: {type(e).__name__}: {str(e)[:150]}")
+    return {"problems": problems, "eos": eos, "template_tail": tail}
+
+
+def check(path: str) -> dict:
+    from freetoken.models.gguf.config import GGUF_ARCH_TO_REGISTRY
+    from freetoken.models.gguf.reader import gguf_architecture, load_gguf_metadata
+
+    if _IN_FLIGHT is not None:
+        _IN_FLIGHT[os.getpid()] = path
+    t0 = time.time()
+    out = {"path": path, "problems": []}
+    problems = out["problems"]
+    try:
+        out["arch"] = gguf_architecture(path)
+        if out["arch"] not in GGUF_ARCH_TO_REGISTRY:
+            out["status"] = "unsupported"
+            return out
+        config, model = _build_model(path)
+        cfg = config.model_config
+        out["config"] = {"layers": cfg.num_layers, "hidden": cfg.hidden_size, "experts": cfg.num_experts, "vocab": cfg.vocab_size, "tied": cfg.tie_word_embeddings}
+        out["reads"] = "meta"
+        try:
+            load_problems = _dry_load(path, config, model, meta=True)
+        except Exception:
+            out["reads"] = "real"
+            if not _wait_for_ram(_MIN_FREE_GB):
+                raise MemoryError(f"the weight reader needs real bytes and free RAM stayed under {_MIN_FREE_GB} GB") from None
+            load_problems = _dry_load(path, config, model, meta=False)
+        problems.extend(load_problems)
+        tags = _expert_banks(path, cfg, problems)
+        layers, glu = _gguf_layers(model)
+        out["types"] = _check_kernels(layers, cfg.tie_word_embeddings, glu, tags, problems)
+        meta = load_gguf_metadata(path)
+        key = _tokenizer_key(meta)
+        if key not in _TOKENIZERS:
+            try:
+                _TOKENIZERS[key] = _check_tokenizer(path)
+            except Exception as e:
+                _TOKENIZERS[key] = {"problems": [f"tokenizer: {type(e).__name__}: {str(e)[:200]}"]}
+        tok = _TOKENIZERS[key]
+        problems.extend(tok["problems"])
+        out["tokenizer"] = {"hash": key, "pre": meta.get("tokenizer.ggml.pre"), "eos": tok.get("eos"), "template_tail": tok.get("template_tail")}
+    except Exception as e:
+        problems.append(f"{type(e).__name__}: {str(e)[:300]}")
+        out["traceback"] = traceback.format_exc()[-1500:]
+    finally:
+        _clear_caches()
+        out["seconds"] = round(time.time() - t0, 1)
+        if _IN_FLIGHT is not None:
+            _IN_FLIGHT.pop(os.getpid(), None)
+    out.setdefault("status", "fail" if problems else "pass")
+    return out
+
+
+def gguf_files(paths: list[str]) -> list[str]:
+    """Every model GGUF under ``paths``, once; a split model is checked through its first file."""
+    files = []
+    for p in paths:
+        found = sorted(glob.glob(os.path.join(p, "**", "*.gguf"), recursive=True)) if os.path.isdir(p) else [p]
+        for f in found:
+            m = _SPLIT.search(f)
+            if os.path.isfile(f) and (m is None or m.group(1) == "00001") and not any(s in os.path.basename(f).lower() for s in ("mmproj", "imatrix")):
+                files.append(os.path.abspath(f))
+    return list(dict.fromkeys(files))
+
+
+def read_results(path: str) -> list[dict]:
+    """The latest record per file; a line cut short by a killed run is dropped."""
+    latest = {}
+    if os.path.exists(path):
+        with open(path) as f:
+            for line in f:
+                try:
+                    rec = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                latest[rec["path"]] = rec
+    return list(latest.values())
+
+
+def _problem_kind(problem: str) -> str:
+    """A problem with tensor names, layer numbers and sizes masked, to group the same failure across files."""
+    problem = re.sub(r"blk\.\d+", "blk.N", problem)
+    problem = re.sub(r"model\.layers\.\d+", "model.layers.N", problem)
+    problem = re.sub(r"/\S+\.gguf", "<file>", problem)
+    return re.sub(r"\b\d+\b", "N", problem)[:160]
+
+
+def report(results: list[dict]) -> None:
+    supported = [r for r in results if r["status"] != "unsupported"]
+    by_arch: dict[str, list[int]] = {}
+    for r in supported:
+        stats = by_arch.setdefault(r.get("arch") or "?", [0, 0])
+        stats[0] += r["status"] == "pass"
+        stats[1] += 1
+    print(f"{sum(r['status'] == 'pass' for r in supported)}/{len(supported)} files of served architectures pass")
+    for arch, (ok, n) in sorted(by_arch.items()):
+        print(f"  {arch:16s} {ok}/{n}")
+    unsupported = Counter(r.get("arch") for r in results if r["status"] == "unsupported")
+    if unsupported:
+        print("not checked, architecture not served: " + ", ".join(f"{a} {n}" for a, n in unsupported.most_common()))
+    real = sum(r.get("reads") == "real" for r in supported)
+    if real:
+        print(f"{real} files needed real weight reads")
+    kinds = Counter(_problem_kind(r["problems"][0]) for r in supported if r["status"] != "pass")
+    if kinds:
+        print("failures, by first problem:")
+        for kind, n in kinds.most_common():
+            print(f"  {n:5d}  {kind}")
+
+
+class _Pool:
+    """Worker processes that check files, and know which files were running when a worker died."""
+
+    def __init__(self, kernels: dict | None, min_free_gb: float):
+        self.ctx = mp.get_context("spawn")
+        self.manager = self.ctx.Manager()
+        self.in_flight = self.manager.dict()
+        self.args = (kernels, self.in_flight, min_free_gb)
+
+    def run(self, files: list[str], workers: int):
+        """Yields (path, record), or (path, None) for a file that was running when a worker died."""
+        pending = list(files)
+        while pending:
+            self.in_flight.clear()
+            # no max_tasks_per_child: on Python 3.12 the pool hangs once every worker has retired
+            with cf.ProcessPoolExecutor(min(workers, len(pending)), mp_context=self.ctx, initializer=_init_worker, initargs=self.args) as pool:
+                futures = {pool.submit(check, f): f for f in pending}
+                pending, broken = [], False
+                for fut in cf.as_completed(futures):
+                    try:
+                        yield futures[fut], fut.result()
+                    except cf.process.BrokenProcessPool:
+                        broken = True
+                        pending.append(futures[fut])
+            if broken:
+                # a worker that died before taking a file leaves none in flight: blame every file left
+                running = set(self.in_flight.values()) & set(pending) or set(pending)
+                for path in running:
+                    yield path, None
+                pending = [p for p in pending if p not in running]
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("paths", nargs="*", help="GGUF files or directories to search")
+    ap.add_argument("--out", help="results JSONL, appended to")
+    ap.add_argument("--kernels", help="kernel table from probe_kernels.py; without it the kernel check is skipped")
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--min-free-gb", type=float, default=32.0, help="free RAM a worker waits for before reading real weight bytes")
+    ap.add_argument("--retry-failed", action="store_true", help="check again the files that failed in --out")
+    ap.add_argument("--report", metavar="JSONL", help="only print the summary of an earlier run")
+    args = ap.parse_args()
+
+    if args.report:
+        report(read_results(args.report))
+        return 0
+    if not args.out or not args.paths:
+        ap.error("give --out and at least one path")
+    kernels = None
+    if args.kernels:
+        with open(args.kernels) as f:
+            kernels = json.load(f)["types"]
+    else:
+        print("no --kernels table: the kernel check is skipped", file=sys.stderr)
+    keep = {"pass", "unsupported"} if args.retry_failed else {"pass", "unsupported", "fail"}
+    done = {r["path"] for r in read_results(args.out) if r["status"] in keep}
+    files = [f for f in gguf_files(args.paths) if f not in done]
+    print(f"{len(files)} files to check, {len(done)} already in {args.out}", file=sys.stderr, flush=True)
+
+    pool = _Pool(kernels, args.min_free_gb)
+    with open(args.out, "a") as out:
+
+        def write(rec):
+            mark = {"pass": "ok  ", "unsupported": "skip"}.get(rec["status"], "FAIL")
+            print(f"{mark} {rec['path']}" + (f"  {rec['problems'][0][:120]}" if rec["problems"] else ""), file=sys.stderr, flush=True)
+            out.write(json.dumps(rec) + "\n")
+            out.flush()
+
+        crashed = []
+        for path, rec in pool.run(files, args.workers):
+            if rec is None:
+                crashed.append(path)
+            else:
+                write(rec)
+        # each file that was running when a worker died runs alone, to find the one that kills it
+        for path in crashed:
+            for _, rec in pool.run([path], 1):
+                write(rec or {"path": path, "status": "crash", "problems": ["worker process died (out of memory or a crash)"]})
+    results = read_results(args.out)
+    report(results)
+    return 0 if all(r["status"] in ("pass", "unsupported") for r in results) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/gguf_check/fetch_headers.py b/scripts/gguf_check/fetch_headers.py
new file mode 100644
index 0000000000..fdcbd33704
--- /dev/null
+++ b/scripts/gguf_check/fetch_headers.py
@@ -0,0 +1,193 @@
+"""Write header-only stubs of the GGUF files in HF repos: the KV metadata and tensor table, in a sparse file of the real size.
+
+    python scripts/gguf_check/fetch_headers.py <out_dir> [<repo> ...] [--search <query>]
+
+The weights stay holes that read as zeros, so a stub takes a few MB of disk and ``check.py`` reads it like the real file.
+Stubs land at ``<out_dir>/<repo>/<path in repo>``; a stub already there with the right size is kept, so a rerun only fetches what is missing.
+Prints one JSON line per file. Gated repos read ``HF_TOKEN``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures as cf
+import json
+import os
+import struct
+import sys
+import threading
+
+import requests
+from huggingface_hub import HfApi, hf_hub_url
+from huggingface_hub.utils import build_hf_headers
+
+_FIRST_READ, _NEXT_READ = 24 << 20, 16 << 20
+_SCALAR = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", 7: "?", 10: "Q", 11: "q", 12: "d"}
+_STRING, _ARRAY = 8, 9
+_F32 = 0
+# tensors whose values a parse_gguf_config reads, so their bytes go into the stub
+VALUE_TENSORS = ("rope_freqs.weight",)
+
+_local = threading.local()
+
+
+def _get_range(url: str, start: int, n: int) -> requests.Response:
+    """Bytes [start, start + n) of ``url``, or fewer at the end of the file; refuses a reply that is not exactly that range."""
+    r = _session().get(url, headers={"Range": f"bytes={start}-{start + n - 1}"}, timeout=300)
+    r.raise_for_status()
+    got = r.headers.get("Content-Range", "")
+    if r.status_code != 206 or not got.startswith(f"bytes {start}-"):
+        raise ValueError(f"{url}: asked for bytes {start}-{start + n - 1}, got status {r.status_code} {got!r}")
+    end = int(got.split("-", 1)[1].split("/", 1)[0])
+    if len(r.content) != end - start + 1:
+        raise ValueError(f"{url}: short read at byte {start}")
+    return r
+
+
+def _session() -> requests.Session:
+    if not hasattr(_local, "session"):
+        _local.session = requests.Session()
+        _local.session.headers.update(build_hf_headers())
+    return _local.session
+
+
+class _Remote:
+    """Reads a remote file front to back through range requests, keeping every byte read so far."""
+
+    def __init__(self, url: str):
+        self.url, self.buf, self.pos, self.size = url, bytearray(), 0, None
+
+    def _need(self, n: int) -> None:
+        while len(self.buf) < self.pos + n:
+            start = len(self.buf)
+            if self.size is not None and start >= self.size:
+                raise EOFError(f"{self.url}: header runs past the end of the file")
+            r = _get_range(self.url, start, max(_FIRST_READ if start == 0 else _NEXT_READ, self.pos + n - start))
+            if self.size is None:
+                self.size = int(r.headers["Content-Range"].rsplit("/", 1)[1])
+            self.buf += r.content
+
+    def read(self, n: int) -> bytes:
+        self._need(n)
+        self.pos += n
+        return bytes(self.buf[self.pos - n : self.pos])
+
+    def unpack(self, fmt: str) -> tuple:
+        return struct.unpack("<" + fmt, self.read(struct.calcsize("<" + fmt)))
+
+    def string(self) -> str:
+        (n,) = self.unpack("Q")
+        return self.read(n).decode("utf-8", "replace")
+
+    def skip_value(self, kind: int) -> None:
+        if kind in _SCALAR:
+            self.read(struct.calcsize(_SCALAR[kind]))
+        elif kind == _STRING:
+            self.string()
+        elif kind == _ARRAY:
+            elem, n = self.unpack("I")[0], self.unpack("Q")[0]
+            if elem in _SCALAR:
+                self.read(struct.calcsize(_SCALAR[elem]) * n)
+            else:
+                for _ in range(n):
+                    self.skip_value(elem)
+        else:
+            raise ValueError(f"{self.url}: unknown GGUF value type {kind}")
+
+
+def fetch(repo: str, path: str, size: int, out_dir: str) -> dict:
+    dst = os.path.join(out_dir, repo, path)
+    rec = {"repo": repo, "path": path, "stub": dst, "size": size, "value_tensors": list(VALUE_TENSORS)}
+    # the sidecar marks a finished stub and which tensors it holds the values of
+    sidecar = dst + ".json"
+    if os.path.isfile(dst) and os.path.getsize(dst) == size and os.path.isfile(sidecar):
+        with open(sidecar) as f:
+            if json.load(f).get("value_tensors") == rec["value_tensors"]:
+                return {**rec, "ok": True, "cached": True}
+    url = hf_hub_url(repo, path)
+    try:
+        f = _Remote(url)
+        if f.read(4) != b"GGUF":
+            raise ValueError("not a GGUF file")
+        _, n_tensors, n_kv = f.unpack("IQQ")
+        align = 32
+        for _ in range(n_kv):
+            key, kind = f.string(), f.unpack("I")[0]
+            if key == "general.alignment" and kind in (4, 5):
+                align = f.unpack("I")[0]
+            else:
+                f.skip_value(kind)
+        tensors = {}
+        for _ in range(n_tensors):
+            name = f.string()
+            (nd,) = f.unpack("I")
+            dims = f.unpack(f"{nd}Q")
+            kind, offset = f.unpack("IQ")
+            tensors[name] = (dims, kind, offset)
+        header_end = f.pos
+        data_start = (header_end + align - 1) // align * align
+        tmp = dst + ".part"
+        os.makedirs(os.path.dirname(dst), exist_ok=True)
+        with open(tmp, "wb") as out:
+            out.write(f.buf[:header_end])
+            out.truncate(f.size)
+            for name in VALUE_TENSORS:
+                if name not in tensors:
+                    continue
+                dims, kind, offset = tensors[name]
+                if kind != _F32 or len(dims) != 1:
+                    raise ValueError(f"{name} is type {kind} with dims {dims}, expected 1-D F32")
+                start, n = data_start + offset, dims[0] * 4
+                r = _get_range(url, start, n)
+                if len(r.content) != n:
+                    raise ValueError(f"{name} runs past the end of the file")
+                out.seek(start)
+                out.write(r.content)
+        os.replace(tmp, dst)
+        rec.update(header_bytes=header_end, tensors=n_tensors)
+        with open(sidecar, "w") as f:
+            json.dump(rec, f)
+        return {**rec, "ok": True}
+    except Exception as e:
+        return {**rec, "ok": False, "error": f"{type(e).__name__}: {e}"[:300]}
+
+
+def list_gguf(api: HfApi, repo: str) -> list[tuple[str, int]]:
+    """(path, size) of every model GGUF in the repo; vision projectors and imatrix files are left out."""
+    files = []
+    for entry in api.list_repo_tree(repo, recursive=True):
+        path = entry.path
+        if path.endswith(".gguf") and not any(s in path.lower() for s in ("mmproj", "imatrix")):
+            files.append((path, entry.size))
+    return files
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("out_dir")
+    ap.add_argument("repos", nargs="*")
+    ap.add_argument("--search", action="append", default=[], help="also take every HF GGUF repo whose name matches this query (repeatable)")
+    ap.add_argument("--workers", type=int, default=12)
+    args = ap.parse_args()
+
+    api = HfApi()
+    repos = list(args.repos)
+    for query in args.search:
+        repos += [m.id for m in api.list_models(search=query, filter="gguf") if m.id not in repos]
+    jobs = []
+    for repo in repos:
+        try:
+            jobs += [(repo, path, size) for path, size in list_gguf(api, repo)]
+        except Exception as e:
+            print(json.dumps({"repo": repo, "ok": False, "error": f"{type(e).__name__}: {e}"[:300]}), flush=True)
+    print(f"{len(jobs)} GGUF files in {len(repos)} repos", file=sys.stderr, flush=True)
+    failed = 0
+    with cf.ThreadPoolExecutor(args.workers) as ex:
+        for rec in ex.map(lambda job: fetch(*job, args.out_dir), jobs):
+            failed += not rec["ok"]
+            print(json.dumps(rec), flush=True)
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/gguf_check/probe_kernels.py b/scripts/gguf_check/probe_kernels.py
new file mode 100644
index 0000000000..4b53cd06ef
--- /dev/null
+++ b/scripts/gguf_check/probe_kernels.py
@@ -0,0 +1,111 @@
+"""Record which ggml ops this build runs for each quantized ggml type, as the kernel table ``check.py --kernels`` reads.
+
+    python scripts/gguf_check/probe_kernels.py <out.json>
+
+Runs every op on random packed weights on the current GPU, at the token counts that pick different kernels (vector, small-batch and MMQ).
+An op counts as covered when it runs; the values are not checked. The table holds for this build and GPU architecture only.
+Each type runs in its own process, so a CUDA fault in one type leaves the others' results intact.
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures as cf
+import json
+import multiprocessing as mp
+import sys
+
+K, N, EXPERTS, TOP_K = 1024, 256, 8, 2
+TOKENS = (1, 8, 16, 512)
+GLU_TOKENS = (1, 8)
+# not weight types: plain integers and doubles, and the q8 blocks ggml quantizes activations to
+_NOT_WEIGHTS = ("F64", "I8", "I16", "I32", "I64", "Q8_1", "Q8_K")
+
+
+def _run(fn) -> str:
+    import torch
+
+    try:
+        fn()
+        torch.cuda.synchronize()
+        return "ok"
+    except Exception as e:
+        return f"{type(e).__name__}: {e}"[:200]
+
+
+def probe(ggml_type: int) -> dict:
+    import torch
+
+    from freetoken.kernel import ggml
+    from freetoken.layers.quantization.ggml_types import row_bytes
+
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    row = row_bytes(ggml_type, K)
+    w = torch.randint(0, 256, (N, row), dtype=torch.uint8, device="cuda", generator=gen)
+    banks = torch.randint(0, 256, (EXPERTS, N, row), dtype=torch.uint8, device="cuda", generator=gen)
+    out = {"mmvq_max_tokens": ggml.mmvq_max_tokens(ggml_type)}
+    for t in TOKENS:
+        x = torch.randn(t, K, dtype=torch.bfloat16, device="cuda", generator=gen)
+        # distinct experts per token, as a router picks them: MMQ faults on a token that routes to one expert twice
+        ids = torch.rand(t, EXPERTS, device="cuda", generator=gen).argsort(dim=1)[:, :TOP_K].to(torch.int32).contiguous()
+        out[f"mm{t}"] = _run(lambda: ggml.mul_mat(w, ggml_type, x))
+        out[f"mmid{t}"] = _run(lambda: ggml.mul_mat_id(banks, ggml_type, x, ids))
+        if t in GLU_TOKENS:
+            out[f"glu{t}"] = _run(lambda: ggml.mul_mat_id_glu(banks, ggml_type, x, ids, "swiglu"))
+    ids = torch.arange(4, dtype=torch.int32, device="cuda")
+    out["get_rows"] = _run(lambda: ggml.get_rows(w, ggml_type, K, ids))
+    out["dequant"] = _run(lambda: ggml.dequantize(w, ggml_type, K))
+    return out
+
+
+def weight_types() -> list[str]:
+    from freetoken.layers.quantization.ggml_types import is_quantized, row_bytes
+    from freetoken.third_party.gguf import GGML_QUANT_SIZES, GGMLQuantizationType
+
+    names = []
+    for qtype in GGMLQuantizationType:
+        if qtype not in GGML_QUANT_SIZES or not is_quantized(int(qtype)) or qtype.name in _NOT_WEIGHTS:
+            continue
+        try:
+            row_bytes(int(qtype), K)
+        except ValueError:
+            continue
+        names.append(qtype.name)
+    return names
+
+
+def _probe_by_name(name: str) -> dict:
+    from freetoken.third_party.gguf import GGMLQuantizationType
+
+    return probe(int(GGMLQuantizationType[name]))
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("out")
+    args = ap.parse_args()
+
+    import torch
+
+    from freetoken.kernel import ggml
+
+    # build the module once up front: a failed build would otherwise show up as every op missing
+    ggml.load_module()
+    ctx = mp.get_context("spawn")
+    table = {}
+    for name in weight_types():
+        with cf.ProcessPoolExecutor(1, mp_context=ctx) as pool:
+            try:
+                table[name] = pool.submit(_probe_by_name, name).result()
+            except cf.process.BrokenProcessPool:
+                table[name] = {"error": "probe process died"}
+        bad = sorted(op for op, r in table[name].items() if op != "mmvq_max_tokens" and r != "ok")
+        print(f"{name:10s} {'ok' if not bad else 'missing ' + ' '.join(bad)}", flush=True)
+    cc = torch.cuda.get_device_capability()
+    with open(args.out, "w") as f:
+        json.dump({"gpu": torch.cuda.get_device_name(), "sm": f"{cc[0]}{cc[1]}", "types": table}, f, indent=1)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/gguf_check/tok_compare.py b/scripts/gguf_check/tok_compare.py
new file mode 100644
index 0000000000..4635fb1266
--- /dev/null
+++ b/scripts/gguf_check/tok_compare.py
@@ -0,0 +1,90 @@
+"""Compare the tokenizer FreeToken builds from a GGUF with the model's own HF tokenizer.
+
+    python scripts/gguf_check/tok_compare.py <gguf> <HF model dir or tokenizer.json> [--combos 20000]
+
+Encodes, without special tokens, a set of edge strings, every normal vocabulary token, and random joins of 2-4 vocabulary tokens, and counts the texts whose ids differ.
+The HF tokenizer is the reference: a stub works, since only the GGUF metadata is read.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+
+EDGE = [
+    "Hello world",
+    "  leading spaces",
+    "trailing spaces   ",
+    "tabs\t\tand\nnewlines\n\n\n",
+    "I'm you're it's we'll they'd DON'T",
+    "1234567890 3.14159 1,000,000",
+    "def f(x):\n    return x ** 2  # square\n",
+    "数学题：求 x 的值。日本語のテキスト。한국어 문장.",
+    "emoji 😀👍🏽 and combining é ä",
+    "Ünïcödé àccents fi ligature",
+    "<|im_start|> <s> </s> [INST] literal special-looking text",
+    "$$\\frac{a}{b} + \\sqrt{x}$$",
+    "URL https://example.com/a?b=c&d=e",
+]
+_NORMAL = 1
+
+
+def vocab_words(meta: dict) -> list[str]:
+    """The text of every normal vocabulary token, for byte-level BPE and SentencePiece vocabularies."""
+    from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+    tokens, types = meta["tokenizer.ggml.tokens"], meta.get("tokenizer.ggml.token_type") or [_NORMAL] * len(meta["tokenizer.ggml.tokens"])
+    model = meta.get("tokenizer.ggml.model")
+    byte_of = {c: b for b, c in bytes_to_unicode().items()}
+    words = []
+    for token, kind in zip(tokens, types):
+        if kind != _NORMAL:
+            continue
+        if model == "gpt2":
+            try:
+                words.append(bytes(byte_of[c] for c in token).decode())
+            except (KeyError, UnicodeDecodeError):
+                pass
+        elif model == "llama":
+            words.append(token.replace("\u2581", " "))
+        else:
+            words.append(token)
+    return words
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("gguf")
+    ap.add_argument("hf", help="HF model dir or its tokenizer.json")
+    ap.add_argument("--combos", type=int, default=20000)
+    ap.add_argument("--examples", type=int, default=3)
+    args = ap.parse_args()
+
+    from tokenizers import Tokenizer
+
+    from freetoken.models.gguf.reader import load_gguf_metadata
+    from freetoken.models.gguf.tokenizer import load_gguf_tokenizer
+
+    hf = Tokenizer.from_file(os.path.join(args.hf, "tokenizer.json") if os.path.isdir(args.hf) else args.hf)
+    ours = load_gguf_tokenizer(args.gguf).backend_tokenizer
+    meta = load_gguf_metadata(args.gguf)
+    words = vocab_words(meta)
+    rng = random.Random(0)
+    combos = [rng.choice(["", " ", "\n"]).join(rng.choice(words) for _ in range(rng.randint(2, 4))) for _ in range(args.combos)]
+    out = {"gguf": args.gguf, "pre": meta.get("tokenizer.ggml.pre")}
+    differ = 0
+    for name, texts in (("edge", EDGE), ("vocab", words), ("combos", combos)):
+        ref = [e.ids for e in hf.encode_batch(texts, add_special_tokens=False)]
+        got = [e.ids for e in ours.encode_batch(texts, add_special_tokens=False)]
+        bad = [(t, r, g) for t, r, g in zip(texts, ref, got) if r != g]
+        differ += len(bad)
+        out[name] = {"texts": len(texts), "differ": len(bad), "examples": [{"text": t, "hf": r[:16], "gguf": g[:16]} for t, r, g in bad[: args.examples]]}
+    print(json.dumps(out, ensure_ascii=False, indent=1))
+    return 1 if differ else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From 163d1de5c57d3427898f4dea77ac1252e36397e1 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 20:18:34 -0700
Subject: [PATCH 129/249] fix(qwen4_exp): read weights under the checkpoint's
 quant config (#759)

* fix(qwen4_exp): read weights under the checkpoint's quant config

- route each dense Linear tensor by its scheme's stored names, so per-channel fp8 scales (llm-compressor) load and fused parts concatenate per role
- name the routed NVFP4 experts by the checkpoint's dialect, so compressed-tensors weight_packed / weight_global_scale load
- accept a bf16 PLE table without a scale in the pinned and disk backends

* fix(quant): accept weight_scale as the modelopt mxfp8 scale name

ModelOpt exports store the same e8m0 codes as weight_scale_inv (MiniMax-M3) or weight_scale (Qwen3.8-Flash-Next-FP4-Local); Stored gains aliases and the scheme-driven readers match any of them.
---
 .../layers/quantization/configs/base.py       |   7 +-
 .../layers/quantization/configs/modelopt.py   |   3 +-
 python/freetoken/models/qwen3_5_moe/weight.py |   2 +-
 python/freetoken/models/qwen4_exp/ple_disk.py |  26 +-
 python/freetoken/models/qwen4_exp/weight.py   | 283 +++++++++++-------
 tests/models/qwen4_exp/test_ple_disk.py       |  21 +-
 tests/models/qwen4_exp/test_weight.py         | 126 +++++++-
 tests/models/qwen4_exp/test_weight_ckpt.py    |   7 +-
 tests/models/test_quant_config.py             |   4 +-
 9 files changed, 344 insertions(+), 135 deletions(-)

diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py
index 78818e8091..3dd63ecc96 100644
--- a/python/freetoken/layers/quantization/configs/base.py
+++ b/python/freetoken/layers/quantization/configs/base.py
@@ -40,10 +40,15 @@ def quantization_config_of(hf_config: Any) -> dict[str, Any] | None:
 
 @dataclass(frozen=True)
 class Stored:
-    """One checkpoint tensor behind a role: its suffix, and whether it holds the quant-side scale whose reciprocal the layer wants."""
+    """One checkpoint tensor behind a role: its suffix, other suffixes exports of the same dialect use for it, and whether it holds the quant-side scale whose reciprocal the layer wants."""
 
     name: str
     reciprocal: bool = False
+    aliases: tuple[str, ...] = ()
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return (self.name, *self.aliases)
 
 
 class QuantConfig(ABC):
diff --git a/python/freetoken/layers/quantization/configs/modelopt.py b/python/freetoken/layers/quantization/configs/modelopt.py
index 6183dfc96b..3547f186d4 100644
--- a/python/freetoken/layers/quantization/configs/modelopt.py
+++ b/python/freetoken/layers/quantization/configs/modelopt.py
@@ -27,7 +27,8 @@ class ModelOptConfig(QuantConfig):
     STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {
         QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"},
         QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"},
-        QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"},
+        # same e8m0 codes under either name: MiniMax-M3 exports weight_scale_inv, Qwen3.8-Flash-Next weight_scale
+        QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": Stored("weight_scale_inv", aliases=("weight_scale",))},
         QuantKind.NVFP4: {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2", "input_scale": "input_scale"},
     }
 
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index bf3cecf052..e7c4676dd4 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -149,7 +149,7 @@ def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]]
         if module.rpartition(".")[2] not in _LINEAR_LEAVES:
             return None
         stored = self.stored(module)
-        roles = {"weight": "weight"} if stored is None else {e.name: r for r, e in self.quant.storage(stored).items()}
+        roles = {"weight": "weight"} if stored is None else {n: r for r, e in self.quant.storage(stored).items() for n in e.names}
         role = roles.get(suffix)
         if role is None:
             if suffix in _DROPPED_SUFFIXES:
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index bb71da2ff7..536853be3b 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -1,4 +1,4 @@
-"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant.
+"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 or bf16 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant.
 
 Hash windows are pure functions of ``req.input_ids`` + ``device_len`` (prefix hits, restores and COW forks need no bookkeeping); the decode input token lives device-side under overlap scheduling and is read back here.
 """
@@ -22,7 +22,7 @@
 from .weight import (
     _PLE_SCALE_SUFFIX,
     _PLE_SHARD_RE,
-    _PLE_ST_DTYPE,
+    _PLE_ST_DTYPES,
     _ple_table_files,
     _safetensors_header,
 )
@@ -50,9 +50,10 @@ class PleRowSource:
     row_bytes: int
     row_stride: int
     scale: float
-    # packed ggml rows (a GGUF table) of ``row_values`` values each; None for one fp8 byte per value
+    # packed ggml rows (a GGUF table) of ``row_values`` values each; None for float rows of ``row_dtype``
     ggml_type: int | None = None
     row_values: int | None = None
+    row_dtype: torch.dtype = torch.float8_e4m3fn
 
     @property
     def total_rows(self) -> int:
@@ -62,6 +63,7 @@ def total_rows(self) -> int:
 def source_from_safetensors(folder: str) -> PleRowSource:
     """Map the checkpoint's ``ngram_embedding.shard_<i>`` tensors in place: one extent per shard, no copy."""
     rows = cols = 0
+    dtype: torch.dtype | None = None
     scale: torch.Tensor | None = None
     paths: list[str] = []
     path_idx: dict[str, int] = {}
@@ -78,8 +80,10 @@ def source_from_safetensors(folder: str) -> PleRowSource:
             match = _PLE_SHARD_RE.search(key)
             if match is None:
                 continue
-            if meta["dtype"] != _PLE_ST_DTYPE:
-                raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}")
+            shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"])
+            if shard_dtype is None or dtype not in (None, shard_dtype):
+                raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard")
+            dtype = shard_dtype
             if rows and tuple(meta["shape"]) != (rows, cols):
                 raise ValueError(f"PLE shard {key} is {meta['shape']}, expected {[rows, cols]}")
             rows, cols = meta["shape"]
@@ -93,9 +97,13 @@ def source_from_safetensors(folder: str) -> PleRowSource:
     if sorted(shards) != list(range(len(shards))) or not shards:
         raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(shards)[:8]}")
     if scale is None:
-        raise ValueError("PLE table has no weight_scale")
+        if dtype is torch.float8_e4m3fn:
+            raise ValueError("PLE table has no weight_scale")
+        scale = torch.tensor(1.0)
     order = [shards[i] for i in range(len(shards))]
-    return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, cols, cols, float(scale))
+    row_bytes = cols * dtype.itemsize
+    return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, row_bytes, row_bytes, float(scale),
+                        row_values=cols, row_dtype=dtype)
 
 
 def source_from_gguf(model_path: str) -> PleRowSource:
@@ -137,11 +145,11 @@ def __init__(
         from freetoken.layers.quantization.ggml_types import float_dtype
 
         self.num_rows = source.total_rows
-        self.head_dim = source.row_values or source.row_bytes  # fp8: one byte per element
+        self.head_dim = source.row_values or source.row_bytes
         self._ggml_type = source.ggml_type
         self._row_bytes = source.row_bytes
         # float rows (fp8, or F32 / F16 / BF16 in a GGUF) are only cast; quantized ggml rows go through ggml
-        self._row_dtype = torch.float8_e4m3fn if source.ggml_type is None else float_dtype(source.ggml_type)
+        self._row_dtype = source.row_dtype if source.ggml_type is None else float_dtype(source.ggml_type)
         self.dtype = dtype
         self.heads = int(hash_constants["num_ngram_heads"])
         self.scale = source.scale
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 691651d4cc..6c62d141cb 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -1,10 +1,10 @@
-"""Qwen3.8-Flash-Next checkpoint reader (the NVFP4 and the official block-fp8 releases).
+"""Qwen3.8-Flash-Next checkpoint reader (bf16, block-fp8, ModelOpt and llm-compressor exports).
 
 Three separate paths, because the checkpoint's three weight classes live in different places:
 
-* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseFuser``.
-* :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`.
-* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named, for the offload cache's expert reader.
+* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseReader``.
+* :func:`load_ple_table` -- the n-gram table (47.7 GiB fp8 or 95.4 GiB bf16), 128 checkpoint shards concatenated into one pinned :class:`HostBank`.
+* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named under the checkpoint's dialect, for the offload cache's expert reader.
 
 Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``); ``model.visual.*`` is kept only when the model built the tower.
 """
@@ -28,29 +28,35 @@
 from freetoken.models.nvfp4_banks import (
     Nvfp4ExpertSourceSpec,
 )
-from freetoken.layers.quantization import get_quant_config
+from freetoken.layers.quantization import QuantConfig, QuantKind, QuantScheme, get_quant_config
 from freetoken.models.register import get_model_spec
 from freetoken.moe.host_banks import HostBank, read_range_into
 from freetoken.utils import cached_load_hf_config, download_hf_weight
 from freetoken.utils.progress import byte_bar
 from tqdm import tqdm
 
-# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW
-# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's
-# stacked ``mtp.layers.N.mlp.experts.*`` tensors.
-_EXPERT_KEY_RE = re.compile(
+# Routed NVFP4 experts: per-expert, un-fused. Matched against the RAW weight_map key in nvfp4_banks.
+# The ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*`` tensors.
+_EXPERT_KEY_RE = (
     r"^model\.language_model\.layers\.(?P<layer>\d+)\.mlp\.experts\.(?P<expert>\d+)\."
-    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>weight|weight_scale|weight_scale_2)$"
+    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>{kinds})$"
 )
 _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.")
-_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec(
-    key_pattern=_EXPERT_KEY_RE,
-    proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"},
-    layer_to_bank=lambda layer, config: layer,  # every layer is MoE
-    desc="Qwen3.8-Flash-Next NVFP4 experts",
-)
-# Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused.
-_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale")
+# role -> the expert bank reader's canonical (ModelOpt) tensor kind
+_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"}
+
+# leaves the model builds as quant-aware Linear layers: their tensors are read under the QuantConfig
+_LINEAR_LEAVES = frozenset({
+    "q_proj", "k_proj", "v_proj", "o_proj", "index_qk_proj",
+    "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj",
+    "gate_proj", "up_proj", "down_proj",
+    "input_mix_weight_down", "block_inject_weight", "input_mix_weight_up",
+    "key_proj", "value_proj", "lm_head",
+})
+# activation scales of modules whose scheme carries no input_scale role
+_DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"})
+_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn, "e2m1": torch.uint8}
+_QUANT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2, torch.uint8, torch.int8)
 
 # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table.
 _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding."
@@ -79,9 +85,6 @@
 # The top-level hyper_connection_mixer has no injection and never fuses.
 _PAD_TO = {"input_mix_weight_down_block_inject": 16}
 _HC_WITH_INJECT = (".attn_hyper_connection", ".mlp_hyper_connection")
-_KIND_SUFFIXES = (".weight_scale_inv", ".weight")
-_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
-_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn}
 
 
 def _rename(raw_name: str) -> str | None:
@@ -92,100 +95,157 @@ def _rename(raw_name: str) -> str | None:
         return None  # n-gram table + its scale: load_ple_table
     if _EXPERT_RE.search(raw_name):
         return None  # routed experts: offload source banks
-    if raw_name.endswith(_SCALE_SUFFIXES):
+    # static KV-cache scales of the quantizers; the KV cache runs in the engine's dtype
+    if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
         return None
     return rename_vl_prefix(raw_name)
 
 
-def _split_kind(name: str) -> tuple[str, str]:
-    """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale."""
-    for suffix in _KIND_SUFFIXES:
-        if name.endswith(suffix):
-            return name[: -len(suffix)], suffix
-    return name, ""
+def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor:
+    """Per-tensor scalar or per-channel ``[rows, 1]`` scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows."""
+    flat = scale.reshape(-1).to(torch.float32)
+    if flat.numel() == 1:
+        return flat.expand(rows).contiguous()
+    if flat.numel() != rows:
+        raise ValueError(f"scale has {flat.numel()} elements for a weight with {rows} output rows (shape {tuple(scale.shape)}); expected 1 or {rows}")
+    return flat.contiguous()
 
 
-class _DenseFuser:
-    """Concatenates checkpoint projection parts into the model's merged buffers, per kind (weight / block scale).
+class _DenseReader:
+    """Routes each Linear tensor to the buffer its module's scheme declares; packed projections are concatenated per role once every part is in.
 
-    The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from.
+    The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and, through its dialect's storage names, which checkpoint tensor fills which role.
     """
 
-    def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None:
+    def __init__(self, quant: QuantConfig | None, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None:
         self.quant = quant
         self.groups = {fused: parts for fused, parts in packed if fused != "experts"}  # experts: bank reader
         self.by_part: dict[str, list[tuple[str, int]]] = {}
         for fused, parts in self.groups.items():
             for idx, part in enumerate(parts):
                 self.by_part.setdefault(part, []).append((fused, idx))
-        self.buf: dict[tuple[str, str], dict[int, torch.Tensor]] = {}
+        # target module -> (part count, {part: {role: tensor}}, {part: the roles its module stores}, scheme)
+        self.pending: dict[str, tuple[int, dict[int, dict[str, torch.Tensor]], dict[int, set[str]], QuantScheme | None]] = {}
 
-    def scheme(self, module: str):
+    def scheme(self, module: str) -> QuantScheme | None:
         return None if self.quant is None else self.quant.scheme_for(module)
 
-    def _target(self, parent: str, leaf: str) -> tuple[str, int] | None:
+    def target(self, module: str) -> tuple[str, int, int]:
+        """``(fused module, part index, part count)``; a standalone linear is its own single-part target."""
+        parent, _, leaf = module.rpartition(".")
         candidates = self.by_part.get(leaf)
         if not candidates:
-            return None
+            return module, 0, 1
         if len(candidates) > 1:
-            # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py
+            # GDN: quantized checkpoints split qkv|z from b|a; same test as gdn.py
             split = self.scheme(f"{parent}.in_proj_qkvz") is not None
             keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"}
             candidates = [c for c in candidates if c[0] in keep]
-            if not candidates:
-                raise ValueError(f"{parent}.{leaf}: no merged projection for the {'split' if split else 'fused'} GDN layout")
         fused, idx = candidates[0]
         if fused in _PAD_TO and not parent.endswith(_HC_WITH_INJECT):
-            return None
-        return f"{parent}.{fused}", idx
+            return module, 0, 1
+        return f"{parent}.{fused}", idx, len(self.groups[fused])
 
-    def check(self, module: str, name: str, tensor: torch.Tensor) -> None:
-        """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from."""
+    def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None:
+        """Take one tensor; the emitted ``[(name, tensor)]`` once its target module is complete, ``[]`` before, None if ``name`` is not a Linear's tensor."""
+        module, _, suffix = name.rpartition(".")
+        if module.rpartition(".")[2] not in _LINEAR_LEAVES:
+            return None
         scheme = self.scheme(module)
-        if name.endswith(".weight_scale_inv"):
-            if scheme is None or not scheme.has("weight_scale_inv"):
-                raise ValueError(f"{name}: {module} has no block scale in the checkpoint's quant config ({scheme})")
-            return
-        is_fp8 = tensor.dtype in _FP8_DTYPES
+        roles = {"weight": "weight"} if scheme is None else {n: r for r, e in self.quant.storage(scheme).items() for n in e.names}
+        role = roles.get(suffix)
+        if role is None:
+            if suffix in _DROPPED_SUFFIXES:
+                return []
+            raise ValueError(f"{name}: the checkpoint's quant config declares {module} {scheme or 'unquantized'}, stored as {sorted(roles)}")
+        if role == "weight":
+            self._check_weight(name, module, scheme, tensor)
+        target, idx, count = self.target(module)
+        _, parts, expected, _ = self.pending.setdefault(target, (count, {}, {}, scheme))
+        part = parts.setdefault(idx, {})
+        if role in part:
+            raise ValueError(f"{name}: {module} already has its {role} from another checkpoint tensor")
+        part[role] = tensor
+        expected[idx] = set(roles.values())
+        if scheme is not None and set(part) == expected[idx]:
+            parts[idx] = self._finish(module, scheme, part)
+        if len(parts) < count or any(set(parts[i]) != expected[i] for i in parts):
+            return []
+        del self.pending[target]
+        return self._emit(target, [parts[i] for i in range(count)], scheme)
+
+    def missing(self) -> list[str]:
+        """One line per incomplete module: the roles its parts still lack."""
+        lines = []
+        for target, (count, parts, expected, _scheme) in sorted(self.pending.items()):
+            lacking = sorted(set().union(*(expected[i] - set(parts[i]) for i in parts)))
+            if len(parts) < count:
+                lacking.append(f"{count - len(parts)} of {count} fused parts")
+            lines.append(f"{target}: missing {lacking}")
+        return lines
+
+    @staticmethod
+    def _check_weight(name: str, module: str, scheme: QuantScheme | None, tensor: torch.Tensor) -> None:
         if scheme is None:
-            if is_fp8:
+            if tensor.dtype in _QUANT_DTYPES:
                 raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized")
             return
-        expected = _ELEM_DTYPES.get(scheme.weight.elem)
-        if expected is not None and tensor.dtype is not expected:
+        elem = _ELEM_DTYPES.get(scheme.weight.elem)
+        if elem is None:
+            raise ValueError(f"{name}: qwen4_exp has no dense reader for {scheme}")
+        if tensor.dtype is not elem:
             raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} {scheme}")
-        rows, cols = (scheme.weight.group or (1, 1))
-        if rows > 1 and tensor.shape[0] % rows or cols > 1 and tensor.shape[1] % cols:
-            raise ValueError(f"{name}: {tuple(tensor.shape)} is not a multiple of the {rows}x{cols} scale block of {module}")
-
-    def check_unfused(self, name: str, tensor: torch.Tensor) -> None:
-        module, kind = _split_kind(name)
-        if kind == ".weight_scale_inv" or (kind == ".weight" and tensor.dtype in _FP8_DTYPES):
-            self.check(module, name, tensor)
-
-    def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None:
-        """Buffer a part; return the merged ``[(name, tensor)]`` once its kind is complete, ``[]`` while incomplete, ``None`` if ``name`` is not a part."""
-        module, kind = _split_kind(name)
-        if not kind:
-            return None
-        parent, _, leaf = module.rpartition(".")
-        hit = self._target(parent, leaf)
-        if hit is None:
-            return None
-        fused, idx = hit
-        self.check(fused, name, tensor)
-        slots = self.buf.setdefault((fused, kind), {})
-        slots[idx] = tensor
-        parts = self.groups[fused.rpartition(".")[2]]
-        if len(slots) < len(parts):
-            return []
-        del self.buf[(fused, kind)]
-        rows = [slots[i] for i in range(len(parts))]
-        pad_to = _PAD_TO.get(fused.rpartition(".")[2], 0) if kind == ".weight" else 0
-        pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0
-        if pad:
-            rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device))
-        return [(fused + kind, torch.cat(rows, dim=0))]
+        rows, cols = tensor.shape[0], tensor.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1)
+        block_rows, block_cols = scheme.weight.group or (1, 1)
+        if block_rows > 1 and rows % block_rows or block_cols > 1 and cols % block_cols:
+            raise ValueError(f"{name}: {rows}x{cols} is not a multiple of the {block_rows}x{block_cols} scale block of {module}")
+
+    def _finish(self, module: str, scheme: QuantScheme, part: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """Validate one complete part against ``scheme`` and put its scales in the layer's form."""
+        storage = self.quant.storage(scheme)
+        part = {role: 1.0 / t.to(torch.float32) if storage[role].reciprocal else t for role, t in part.items()}
+        weight = part["weight"]
+        rows, cols = weight.shape[0], weight.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1)
+        block_rows, block_cols = scheme.weight.group or (1, 1)
+        scale_role = "weight_scale_inv" if "weight_scale_inv" in part else "weight_scale"
+        out = dict(part)
+        if block_cols < 0:
+            try:
+                out[scale_role] = _per_row_scale(part[scale_role], rows)
+            except ValueError as exc:
+                raise ValueError(f"{module}.{scale_role}: {exc}") from None
+        else:
+            want = (rows // block_rows, cols // block_cols)
+            if tuple(part[scale_role].shape) != want:
+                raise ValueError(f"{module}: {scale_role} is {tuple(part[scale_role].shape)}, expected {want} for {scheme}")
+            if scheme.weight.scale == "e4m3" and part[scale_role].dtype is not torch.float8_e4m3fn:
+                raise ValueError(f"{module}: {scale_role} is {part[scale_role].dtype} but {scheme} stores e4m3 scales")
+        if "weight_global" in part:
+            g = part["weight_global"].reshape(-1).to(torch.float32)
+            if g.numel() != 1:
+                raise ValueError(f"{module}: weight_global has {g.numel()} elements, expected one per-tensor scale")
+            out["weight_global"] = g.expand(rows).contiguous()
+        if "input_scale" in part:
+            out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32)
+        return out
+
+    @staticmethod
+    def _emit(target: str, parts: list[dict[str, torch.Tensor]], scheme: QuantScheme | None) -> list[tuple[str, torch.Tensor]]:
+        pad_to = _PAD_TO.get(target.rpartition(".")[2], 0)
+        if pad_to and scheme is not None:
+            raise ValueError(f"{target}: qwen4_exp pads this fusion to {pad_to} rows and reads it bf16 only, the checkpoint's quant config declares {scheme}")
+        out = []
+        for role in parts[0]:
+            tensors = [part[role] for part in parts]
+            if role == "input_scale":
+                # fused parts read the same activation, so one range covers them: max is exact then and safe if they drift
+                out.append((f"{target}.{role}", torch.stack(tensors).max()))
+                continue
+            pad = (-sum(t.shape[0] for t in tensors)) % pad_to if pad_to else 0
+            if pad:
+                tensors.append(torch.zeros(pad, *tensors[0].shape[1:], dtype=tensors[0].dtype, device=tensors[0].device))
+            out.append((f"{target}.{role}", tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)))
+        return out
 
 
 def iter_weights(
@@ -199,8 +259,8 @@ def iter_weights(
     """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers.
 
     Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts.
-    A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections.
-    Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``.
+    Each Linear is read under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, per-channel / per-tensor fp8, block-fp8, MXFP8 and NVFP4 projections in any mix land as the model's state dict.
+    Fusions, per role: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``.
     ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader.
     """
     if get_tp_info().size > 1:
@@ -210,7 +270,7 @@ def iter_weights(
 
     hf_config = cached_load_hf_config(model_path)
     spec = get_model_spec(hf_config.architectures[0])
-    fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping)
+    reader = _DenseReader(get_quant_config(), spec.packed_modules_mapping)
     for file in tqdm(
         iter_weight_files(model_path),
         desc="Loading weights",
@@ -224,14 +284,16 @@ def iter_weights(
                 if not include_vision and name.startswith(VISION_KEY_PREFIXES):
                     continue
                 tensor = f.get_tensor(raw_name)
-                fused = fuser.fuse(name, tensor)
-                if fused is None:
-                    fuser.check_unfused(name, tensor)
+                emitted = None if name.startswith(VISION_KEY_PREFIXES) else reader.add(name, tensor)
+                if emitted is None:
                     yield name, tensor
                 else:
-                    yield from fused
+                    yield from emitted
 
-    assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}"
+    if reader.pending:
+        lines = reader.missing()
+        shown = "\n  ".join(lines[:8]) + (f"\n  ... {len(lines) - 8} more" if len(lines) > 8 else "")
+        raise ValueError(f"checkpoint is missing tensors the quant config declares for {len(lines)} modules:\n  {shown}")
 
 
 def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
@@ -251,18 +313,18 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple
 
 @dataclass(frozen=True)
 class PleTable:
-    """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor FP8 scale."""
+    """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor scale (1 for a bf16 table)."""
 
     bank: HostBank
     weight_scale: torch.Tensor  # scalar, checkpoint dtype (bf16)
 
     @property
     def tensor(self) -> torch.Tensor:
-        """``[total_rows, ngram_head_dim]`` float8_e4m3fn view of the bank."""
+        """``[total_rows, ngram_head_dim]`` view of the bank in the stored dtype."""
         return self.bank.tensor
 
 
-_PLE_ST_DTYPE = "F8_E4M3"
+_PLE_ST_DTYPES = {"F8_E4M3": torch.float8_e4m3fn, "BF16": torch.bfloat16}
 
 
 def _safetensors_header(path: str) -> tuple[dict, int]:
@@ -322,12 +384,14 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True,
     The checkpoint splits the table into ``split_ngram_parts`` equal row blocks named by shard
     index and scattered over the ``model-plefp8-*`` shards in header (lexicographic) order, so the
     bank is filled shard by shard at ``shard_index * rows_per_shard``. Each read is O_DIRECT: the
-    table is ~47.7 GiB and must not also sit in the page cache while the bank holds the same bytes.
+    table is ~47.7 GiB (fp8) and must not also sit in the page cache while the bank holds the same bytes.
+    An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled.
     """
     folder = download_hf_weight(model_path)
     parts: dict[int, tuple[str, int, int]] = {}  # shard index -> (path, file offset, bytes)
     scale: torch.Tensor | None = None
     rows = cols = 0
+    dtype: torch.dtype | None = None
     for path in _ple_table_files(folder):
         header, base = _safetensors_header(path)
         for key, meta in header.items():
@@ -340,8 +404,10 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True,
             match = _PLE_SHARD_RE.search(key)
             if match is None:
                 continue
-            if meta["dtype"] != _PLE_ST_DTYPE:
-                raise ValueError(f"PLE table shard {key} has unsupported dtype {meta['dtype']}")
+            shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"])
+            if shard_dtype is None or dtype not in (None, shard_dtype):
+                raise ValueError(f"PLE table shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard")
+            dtype = shard_dtype
             shape = meta["shape"]
             if rows and tuple(shape) != (rows, cols):
                 raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}")
@@ -357,10 +423,12 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True,
     if cols != qwen4_args.ngram_head_dim:
         raise ValueError(f"PLE table row is {cols} wide, config says {qwen4_args.ngram_head_dim}")
     if scale is None:
-        raise ValueError("PLE table has no weight_scale")
+        if dtype is torch.float8_e4m3fn:
+            raise ValueError("PLE table has no weight_scale")
+        scale = torch.tensor(1.0, dtype=torch.bfloat16)
 
-    bank = HostBank((expected * rows, cols), torch.float8_e4m3fn)
-    shard_bytes = rows * cols
+    bank = HostBank((expected * rows, cols), dtype)
+    shard_bytes = rows * cols * dtype.itemsize
     bar = byte_bar(expected * shard_bytes, "Loading PLE table")
     try:
         buf = bank.memoryview()
@@ -382,8 +450,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True,
 # ======================================================================================
 
 
-def nvfp4_expert_spec(model_path: str, config):
-    return _NVFP4_SOURCE_SPEC
+def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec:
+    """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor)."""
+    quant = get_quant_config()
+    stored = quant.stored_tensors(QuantKind.NVFP4)
+    kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()}
+    return Nvfp4ExpertSourceSpec(
+        key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))),
+        proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"},
+        layer_to_bank=lambda layer, config: layer,  # every layer is MoE
+        desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})",
+        kind_map=kind_map,
+        global_reciprocal=stored["weight_global"].reciprocal,
+    )
 
 
 __all__ = [
diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py
index ba9cd98c63..8381720740 100644
--- a/tests/models/qwen4_exp/test_ple_disk.py
+++ b/tests/models/qwen4_exp/test_ple_disk.py
@@ -68,25 +68,31 @@ def _fill(store, args, window, tokens):
 
 
 def _write_checkpoint(tmp_path, table, n_shards):
+    """uint8 ``table`` bytes are stored as a scaled fp8 table, any other dtype as an unscaled table of that dtype."""
     from safetensors.torch import save_file
 
+    fp8 = table.dtype is torch.uint8
     per = table.shape[0] // n_shards
     tensors = {
-        f"{_KEY_PREFIX}.shard_{i}.weight": table[i * per : (i + 1) * per].view(torch.float8_e4m3fn)
+        f"{_KEY_PREFIX}.shard_{i}.weight": table[i * per : (i + 1) * per].view(torch.float8_e4m3fn) if fp8 else table[i * per : (i + 1) * per].clone()
         for i in range(n_shards)
     }
-    tensors[f"{_KEY_PREFIX}.weight_scale"] = torch.tensor(0.03125, dtype=torch.bfloat16)
+    if fp8:
+        tensors[f"{_KEY_PREFIX}.weight_scale"] = torch.tensor(0.03125, dtype=torch.bfloat16)
     save_file(tensors, str(tmp_path / "model.safetensors"))
 
 
-def _make_table(tmp_path):
+def _make_table(tmp_path, dtype=torch.float8_e4m3fn):
     from freetoken.models.qwen4_exp.ple_disk import DiskRowTable, source_from_safetensors
 
     args = parse_config(toy_hf_config()).qwen4_args
     multipliers, sizes, offsets = hash_constants(args)
     total_rows = int(offsets[-1] + sizes[-1])
     gen = torch.Generator().manual_seed(9)
-    table = torch.randint(0, 256, (total_rows, args.ngram_head_dim), dtype=torch.uint8, generator=gen)
+    if dtype is torch.float8_e4m3fn:
+        table = torch.randint(0, 256, (total_rows, args.ngram_head_dim), dtype=torch.uint8, generator=gen)
+    else:
+        table = torch.randn((total_rows, args.ngram_head_dim), generator=gen).to(dtype)
     n_shards = next(k for k in (4, 2, 1) if total_rows % k == 0)
     _write_checkpoint(tmp_path, table, n_shards)
     constants = {
@@ -97,7 +103,7 @@ def _make_table(tmp_path):
         "eos_token_id": EOS,
     }
     disk = DiskRowTable(source_from_safetensors(str(tmp_path)), constants)
-    oracle = GpuResidentTable(table.cuda().view(torch.float8_e4m3fn), scale=0.03125)
+    oracle = GpuResidentTable(table.cuda().view(torch.float8_e4m3fn), scale=0.03125) if dtype is torch.float8_e4m3fn else GpuResidentTable(table.cuda())
     return disk, oracle, args
 
 
@@ -255,8 +261,9 @@ def run(store):
 
 
 @requires_cuda
-def test_disk_table_matches_oracle(tmp_path):
-    disk, oracle, args = _make_table(tmp_path)
+@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16], ids=["fp8", "bf16"])
+def test_disk_table_matches_oracle(tmp_path, dtype):
+    disk, oracle, args = _make_table(tmp_path, dtype)
     emb = _embedding()
 
     # prefill: two segments, one fresh and one mid-sequence window
diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py
index 1a9a100edb..e066acbc2c 100644
--- a/tests/models/qwen4_exp/test_weight.py
+++ b/tests/models/qwen4_exp/test_weight.py
@@ -18,9 +18,10 @@
 from freetoken.kernel.aot_models import SUPPORTED_MODELS, expert_bank_row_bytes
 from freetoken.models.qwen4_exp.weight import (
     _ZERO_CENTERED_NORM_SUFFIXES,
-    _DenseFuser,
+    _DenseReader,
     iter_weights,
     load_ple_table,
+    nvfp4_expert_spec,
 )
 from freetoken.models.register import get_model_spec
 from freetoken.moe.host_banks import HostBank, read_range_into
@@ -64,10 +65,12 @@ def _fp8_scale(weight: torch.Tensor) -> torch.Tensor:
     return torch.rand(weight.shape[0] // BLOCK, weight.shape[1] // BLOCK) + 0.5
 
 
-def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]:
+def _raw_checkpoint(dense_fp8: bool = False, channel_fp8: bool = False, mxfp8_scale: str | None = None) -> dict[str, torch.Tensor]:
     """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise.
 
     ``dense_fp8`` stores the attention and GDN qkv|z / out projections as 128x128 block-fp8 (e4m3 ``.weight`` + fp32 ``.weight_scale_inv``) like the community NVFP4-FP8 requants.
+    ``channel_fp8`` stores those and the shared experts as per-channel fp8 (e4m3 ``.weight`` + bf16 ``[rows, 1]`` ``.weight_scale``) like the llm-compressor requants.
+    ``mxfp8_scale`` stores the attention and GDN qkv|z / out projections as MXFP8 (e4m3 ``.weight`` + uint8 e8m0 codes per 32 inputs) under that scale suffix.
     """
     lm = "model.language_model"
     raw: dict[str, torch.Tensor] = {
@@ -146,12 +149,63 @@ def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]:
             weight = raw[f"{module}.weight"]
             raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn)
             raw[f"{module}.weight_scale_inv"] = _fp8_scale(weight)
+    if channel_fp8:
+        for module in _channel_fp8_modules(gdn_layers=(0,), attn_layers=(1,), mlp_layers=(0, 1)):
+            weight = raw[f"{module}.weight"]
+            raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn)
+            raw[f"{module}.weight_scale"] = (torch.rand(weight.shape[0], 1) + 0.5).to(torch.bfloat16)
+    if mxfp8_scale is not None:
+        for module in (f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z", f"{gdn}.out_proj", *(f"{attn}.{p}_proj" for p in "qkvo")):
+            weight = raw[f"{module}.weight"]
+            raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn)
+            raw[f"{module}.{mxfp8_scale}"] = torch.randint(120, 135, (weight.shape[0], weight.shape[1] // 32), dtype=torch.uint8)
     return raw
 
 
 FP8_DENSE_QUANT = mixed_precision_quant(gdn_layers=(0,), attn_layers=(1,), moe_layers=(0, 1))
 
 
+# ModelOpt MIXED_PRECISION with MXFP8 attention / GDN projections, as in nvidia/Qwen3.8-Flash-Next-FP4-Local
+MXFP8_DENSE_QUANT = {
+    **FP8_DENSE_QUANT,
+    "quantized_layers": {
+        k: ({"quant_algo": "MXFP8", "group_size": 32} if v["quant_algo"] == "FP8_PB_WO" else v)
+        for k, v in FP8_DENSE_QUANT["quantized_layers"].items()
+    },
+}
+
+
+def _channel_fp8_modules(gdn_layers, attn_layers, mlp_layers) -> list[str]:
+    return (
+        [f"{LM}.layers.{i}.linear_attn.{p}" for i in gdn_layers for p in ("in_proj_qkv", "in_proj_z", "out_proj")]
+        + [f"{LM}.layers.{i}.self_attn.{p}_proj" for i in attn_layers for p in "qkvo"]
+        + [f"{LM}.layers.{i}.mlp.shared_expert.{p}_proj" for i in mlp_layers for p in ("gate", "up", "down")]
+    )
+
+
+# llm-compressor mixed-precision as in orcarouter/Qwen3.8-Flash-Next-Uncensored-NVFP4: per-channel fp8 dense, weight-only NVFP4 experts
+CT_CHANNEL_FP8_QUANT = {
+    "quant_method": "compressed-tensors",
+    "format": "mixed-precision",
+    "quantization_status": "compressed",
+    "config_groups": {
+        "group_0": {
+            "format": "naive-quantized",
+            "targets": _channel_fp8_modules(gdn_layers=(0,), attn_layers=(1,), mlp_layers=(0, 1)),
+            "weights": {"num_bits": 8, "type": "float", "strategy": "channel", "symmetric": True, "dynamic": False},
+            "input_activations": None,
+        },
+        "group_1": {
+            "format": "nvfp4-pack-quantized",
+            "targets": [r"re:.*mlp\.experts\..*proj$"],
+            "weights": {"num_bits": 4, "type": "float", "strategy": "tensor_group", "group_size": 16, "symmetric": True, "dynamic": False},
+            "input_activations": None,
+        },
+    },
+    "ignore": ["lm_head"],
+}
+
+
 def _config_json(quantization_config) -> dict:
     cfg = hf_config(
         num_layers=2, head_dim=AHD, num_q=QH, num_kv=KVH, index_head_dim=IHD, index_heads=2,
@@ -227,6 +281,14 @@ def loaded_fp8(checkpoint_fp8) -> dict[str, torch.Tensor]:
     return _load(checkpoint_fp8[0])
 
 
+@pytest.fixture(scope="module")
+def checkpoint_fp8_channel(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]:
+    torch.manual_seed(3)
+    return _write_checkpoint(
+        tmp_path_factory.mktemp("qwen4_exp_fp8_channel_ckpt"), _raw_checkpoint(channel_fp8=True), CT_CHANNEL_FP8_QUANT
+    )
+
+
 def test_tower_keys_come_out_under_the_prefix_load_weight_filters(loaded):
     assert {n for n in loaded if "visual" in n} == {"visual.blocks.0.attn.qkv.weight", "visual.merger.norm.weight"}
 
@@ -348,6 +410,15 @@ def test_load_ple_table_concatenates_shards_in_index_order(checkpoint):
     assert float(table.weight_scale) == 0.125
 
 
+def test_load_ple_table_reads_an_unscaled_bf16_table(tmp_path):
+    prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding"
+    shards = {f"{prefix}.shard_{i}.weight": _bf16(NGRAM_ROWS, NGRAM_DIM) for i in range(NGRAM_SHARDS)}
+    save_file(shards, str(tmp_path / "model.safetensors"))
+    table = load_ple_table(str(tmp_path), SimpleNamespace(split_ngram_parts=NGRAM_SHARDS, ngram_head_dim=NGRAM_DIM), pin=False)
+    assert torch.equal(table.tensor, torch.cat([shards[f"{prefix}.shard_{i}.weight"] for i in range(NGRAM_SHARDS)]))
+    assert float(table.weight_scale) == 1.0
+
+
 def test_load_ple_table_rejects_a_shard_count_mismatch(checkpoint):
     folder, _raw = checkpoint
     args = SimpleNamespace(split_ngram_parts=NGRAM_SHARDS + 1, ngram_head_dim=NGRAM_DIM)
@@ -429,11 +500,11 @@ def test_every_registry_architecture_is_claimed_by_an_aot_entry():
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs cuda")
 def test_fusion_pad_rides_the_tensor_device():
     """safetensors loads straight to cuda; a cpu-allocated pad row would break torch.cat."""
-    fuser = _DenseFuser(None, get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping)
+    reader = _DenseReader(None, get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping)
     down = torch.randn(320, 64, device="cuda", dtype=torch.bfloat16)
     inject = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16)
-    assert fuser.fuse("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down) == []
-    [(key, fused)] = fuser.fuse("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject)
+    assert reader.add("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down) == []
+    [(key, fused)] = reader.add("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject)
     assert key == "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight"
     assert fused.device.type == "cuda" and fused.shape[0] == 336
     assert torch.equal(fused[324:], torch.zeros(12, 64, device="cuda", dtype=torch.bfloat16))
@@ -457,20 +528,55 @@ def checkpoint_nvfp4(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]:
 )
 
 
-@pytest.mark.parametrize("fixture", ["checkpoint", "checkpoint_nvfp4", "checkpoint_fp8"])
+@pytest.mark.parametrize("fixture", ["checkpoint", "checkpoint_nvfp4", "checkpoint_fp8", "checkpoint_fp8_channel"])
 def test_emitted_keys_are_the_model_state_dict(fixture, request):
-    """The reader fills exactly the buffers the engine builds from the same config, block-fp8 ones with the stored dtypes."""
+    """The reader fills exactly the buffers the engine builds from the same config, fp8 ones with the stored weight dtype and fp32 scales."""
     folder, _raw = request.getfixturevalue(fixture)
     loaded, state = _load(folder, vision=False), meta_state_dict(folder)
     assert set(loaded) == set(state)
-    if fixture != "checkpoint_fp8":
+    if fixture in ("checkpoint", "checkpoint_nvfp4"):
         assert loaded["model.layers.0.linear_attn.in_proj.weight"].dtype is torch.bfloat16
         return
+    scale = ".weight_scale_inv" if fixture == "checkpoint_fp8" else ".weight_scale"
     for module in FP8_MODULES:
-        for kind in (".weight", ".weight_scale_inv"):
+        for kind in (".weight", scale):
             assert loaded[module + kind].shape == state[module + kind].shape, module + kind
         assert loaded[module + ".weight"].dtype is state[module + ".weight"].dtype is torch.float8_e4m3fn
-        assert loaded[module + ".weight_scale_inv"].dtype is torch.float32  # the engine casts it to the bf16 buffer at load
+        assert loaded[module + scale].dtype is torch.float32  # the engine casts it to the buffer dtype at load
+
+
+def test_channel_fp8_scales_fuse_as_fp32_rows(checkpoint_fp8_channel):
+    folder, raw = checkpoint_fp8_channel
+    loaded = _load(folder)
+    gdn, shared = f"{LM}.layers.0.linear_attn", f"{LM}.layers.1.mlp.shared_expert"
+    for fused, parts in (
+        ("model.layers.0.linear_attn.in_proj_qkvz", [f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z"]),
+        ("model.layers.1.mlp.shared_expert.gate_up_proj", [f"{shared}.gate_proj", f"{shared}.up_proj"]),
+    ):
+        assert torch.equal(loaded[f"{fused}.weight_scale"], torch.cat([raw[f"{p}.weight_scale"].reshape(-1).float() for p in parts]))
+        assert torch.equal(loaded[f"{fused}.weight"].view(torch.uint8), torch.cat([raw[f"{p}.weight"].view(torch.uint8) for p in parts]))
+
+
+@pytest.mark.parametrize("scale_name", ["weight_scale_inv", "weight_scale"])
+def test_mxfp8_scales_load_under_either_modelopt_name(tmp_path, scale_name):
+    folder, raw = _write_checkpoint(tmp_path, _raw_checkpoint(mxfp8_scale=scale_name), MXFP8_DENSE_QUANT)
+    loaded, state = _load(folder, vision=False), meta_state_dict(folder)
+    assert set(loaded) == set(state)
+    gdn = f"{LM}.layers.0.linear_attn"
+    assert torch.equal(
+        loaded["model.layers.0.linear_attn.in_proj_qkvz.weight_scale_inv"],
+        torch.cat([raw[f"{gdn}.{p}.{scale_name}"] for p in ("in_proj_qkv", "in_proj_z")]),
+    )
+
+
+def test_expert_spec_reads_the_llm_compressor_names(checkpoint_fp8_channel):
+    folder, _raw = checkpoint_fp8_channel
+    install_quant_config(folder)
+    spec = nvfp4_expert_spec(folder, None)
+    base = f"{LM}.layers.1.mlp.experts.2.down_proj"
+    kinds = {spec.kind_map[spec.key_pattern.match(f"{base}.{k}")["kind"]] for k in ("weight_packed", "weight_scale", "weight_global_scale")}
+    assert kinds == {"weight", "weight_scale", "weight_scale_2"}
+    assert spec.global_reciprocal
 
 
 def _assert_fused_per_kind(loaded, raw, fused: str, parts: list[str]) -> None:
diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py
index 4395a86394..1c073792a5 100644
--- a/tests/models/qwen4_exp/test_weight_ckpt.py
+++ b/tests/models/qwen4_exp/test_weight_ckpt.py
@@ -22,10 +22,10 @@
 from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces
 from freetoken.models.qwen4_exp.config import parse_config
 from freetoken.models.qwen4_exp.weight import (
-    _NVFP4_SOURCE_SPEC,
     _ZERO_CENTERED_NORM_SUFFIXES,
     iter_weights,
     load_ple_table,
+    nvfp4_expert_spec,
 )
 from freetoken.moe.expert_banks import build_expert_banks
 from freetoken.moe.host_banks import HostResidency
@@ -170,7 +170,7 @@ def dense_pass() -> tuple[list[str], dict[str, torch.Tensor]]:
     sampled: dict[str, torch.Tensor] = {}
     install_quant_config(MODEL_PATH)
     for name, tensor in iter_weights(
-        MODEL_PATH, torch.device("cpu"), include_moe_experts=True, include_non_moe=True
+        MODEL_PATH, torch.device("cpu"), include_moe_experts=True, include_non_moe=True, include_vision=False
     ):
         names.append(name)
         if name in wanted or name.endswith(_ZERO_CENTERED_NORM_SUFFIXES):
@@ -271,8 +271,9 @@ def layer0_expert_banks():
     """The real NVFP4 source-bank loader, restricted to layer 0 (1.32 GiB instead of 63.5)."""
     if not torch.cuda.is_available():
         pytest.skip("expert bank pinning needs CUDA")
+    install_quant_config(MODEL_PATH)
     spec = dataclasses.replace(
-        _NVFP4_SOURCE_SPEC, layer_to_bank=lambda layer, config: 0 if layer == 0 else None
+        nvfp4_expert_spec(MODEL_PATH, None), layer_to_bank=lambda layer, config: 0 if layer == 0 else None
     )
     config = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I,
                              num_moe_layers=1)
diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py
index 9c18b74907..6b9ecd810c 100644
--- a/tests/models/test_quant_config.py
+++ b/tests/models/test_quant_config.py
@@ -569,9 +569,11 @@ def shape(suffix):
         if tensors:
             if dtype("weight") == "U8":
                 return {QuantKind.MXFP4}
+            if dtype("weight_scale") == "U8":
+                return {QuantKind.MXFP8}
             sh = shape("weight_scale") or []
             return {QuantKind.FP8_BLOCK} if len(sh) == 2 and sh[0] > 1 and sh[1] > 1 else {QuantKind.FP8_TENSOR}
-        return {QuantKind.FP8_TENSOR, QuantKind.FP8_BLOCK, QuantKind.MXFP4}
+        return {QuantKind.FP8_TENSOR, QuantKind.FP8_BLOCK, QuantKind.MXFP4, QuantKind.MXFP8}
     return {QuantKind.NONE}
 
 

From 05e7cb7177951ddacd79ea0ba5988c605abd77d3 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 20:24:04 -0700
Subject: [PATCH 130/249] fix(qwen3_5_moe)!: apply the zero-centred norm
 weights' +1 in fp32 at runtime (#755)

BREAKING CHANGE: an FTW of Qwen3.5 / Qwen3.6 / Qwen3.8-27B converted by an older build
stores the norms with the 1 added and gives wrong output; reconvert it with `ft checkpoint`.
---
 docs/ftw-hotfix.md                            |  1 +
 docs/models.md                                |  3 +++
 .../freetoken/models/qwen3_5_moe/attention.py |  8 +++----
 python/freetoken/models/qwen3_5_moe/gguf.py   | 15 +++++++------
 python/freetoken/models/qwen3_5_moe/model.py  | 21 +++++++------------
 python/freetoken/models/qwen3_5_moe/weight.py | 13 ------------
 tests/models/test_qwen3_5_moe_weight.py       |  6 +++---
 7 files changed, 27 insertions(+), 40 deletions(-)

diff --git a/docs/ftw-hotfix.md b/docs/ftw-hotfix.md
index a88c0066d8..0c6b6b9feb 100644
--- a/docs/ftw-hotfix.md
+++ b/docs/ftw-hotfix.md
@@ -33,6 +33,7 @@ byte range, never the whole checkpoint; `--revision` pins the Hub revision to re
 | `RuntimeError: Unexpected keys ... .weight_scale` | nvidia/GLM-5.2-NVFP4 | dequantize the old runtime-fp8 weights back to bf16 | none |
 | `PLE shard indices are not contiguous 0..N-1: []` | Qwen3.8-Flash-Next (FTWs converted before #420) | write the PLE table as `ple-table-*.safetensors` | 47.7 GiB |
 | `ValueError: ... stores NVFP4 global scales as fp16` | NVFP4 FTWs converted by older builds | none: the fp16 rounding cannot be undone, so the script refuses; reconvert with `ft checkpoint` | - |
+| none: loads, but the output is wrong | Qwen3.5 / Qwen3.6 / Qwen3.8-27B FTWs converted by older builds (norm weights stored as bf16(1 + w)) | none: the script does not detect it; reconvert with `ft checkpoint` | - |
 | `ValueError: ... holds no vision encoder tensors` | image-input families (Qwen3-VL, Qwen3.6, Qwen3.8-Flash-Next, Gemma-4, GLM-5.3-Flash, Muse-Glimmer, MiniMax-M3) converted before the family served images | add the vision encoder: the tower's checkpoint tensors are fetched and the family's own weight reader emits them as the converter would | 1 to 4 GiB |
 
 The script decides by itself: it builds the current model from the FTW's `config.json`, compares
diff --git a/docs/models.md b/docs/models.md
index 9c35ece2fd..18f73fac54 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -88,6 +88,9 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   when the gate was fused into `qkvg_proj` no longer loads, reconvert it with `ft checkpoint`.
 - NVFP4 global scales are kept in fp32. An NVFP4 FTW that stores them as fp16 (converted by an
   older build) is refused at load; reconvert it with `ft checkpoint`.
+- Qwen3.5 / Qwen3.6 / Qwen3.8-27B load their zero-centred norm weights as stored and add the 1 in
+  fp32 at runtime. An FTW of these models converted by an older build stores the norms with the 1
+  already added: it loads without an error but gives wrong output; reconvert it with `ft checkpoint`.
 - An FTW converted before its family served images holds no vision encoder: `ft serve`
   refuses it unless started with `--text-model-only` (or `--mm-disable vision`); reconvert it
   with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md).
diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py
index 7a9e6a333a..0a736a1042 100644
--- a/python/freetoken/models/qwen3_5_moe/attention.py
+++ b/python/freetoken/models/qwen3_5_moe/attention.py
@@ -4,7 +4,7 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated
 from freetoken.layers.rotary import get_rope
 from freetoken.utils import nvtx_annotate
 
@@ -41,10 +41,8 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
             config.hidden_size, self._qkv_split, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
         )
-        # Qwen3.5 uses Gemma-style (1+weight) RMSNorm; the weight loader bakes the +1
-        # into the stored weight (GemmaRMSNorm scales by the raw weight).
-        self.q_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps)
-        self.k_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps)
+        self.q_norm = GemmaPlusOneRMSNorm(head_dim, eps=config.rms_norm_eps)
+        self.k_norm = GemmaPlusOneRMSNorm(head_dim, eps=config.rms_norm_eps)
         self.rotary = get_rope(
             head_dim=head_dim,
             rotary_dim=config.rotary_config.rotary_dim,
diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py
index 7f72061e00..caf1268008 100644
--- a/python/freetoken/models/qwen3_5_moe/gguf.py
+++ b/python/freetoken/models/qwen3_5_moe/gguf.py
@@ -1,6 +1,6 @@
 """Qwen3.5 / 3.6 / 3.8 dense and MoE from a llama.cpp ``qwen35`` / ``qwen35moe`` GGUF: config from the GGUF metadata, tensor names, and the converter's transforms undone.
 
-The llama.cpp converter (conversion/qwen.py) stores ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, and the GDN value heads in tiled order (k-head minor). Its zero-centred norms are stored as ``w + 1``, which is what this family loads. Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead.
+The llama.cpp converter (conversion/qwen.py) stores ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, and the GDN value heads in tiled order (k-head minor). Its zero-centred norms are stored as ``w + 1``; the loader subtracts the 1 in fp32 and the norms add it back in fp32 at runtime. Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead.
 """
 
 from __future__ import annotations
@@ -162,15 +162,18 @@ def iter_gguf_weights(
     def t(name: str) -> GgufTensorInfo:
         return index[name if name in index else f"{name}.weight"]
 
+    def norm(name: str) -> torch.Tensor:
+        return float_tensor(t(name)).float() - 1.0
+
     yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
     if not config.tie_word_embeddings:
         yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
-    yield "model.norm.weight", float_tensor(t("output_norm"))
+    yield "model.norm.weight", norm("output_norm")
 
     for layer in range(config.num_layers):
         blk, pre = f"blk.{layer}", f"model.layers.{layer}"
-        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
-        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.post_attention_norm"))
+        yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm")
+        yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm")
         if config.is_linear_layer(layer):
             la = f"{pre}.linear_attn"
             nk, dk = g.num_key_heads, g.key_head_dim
@@ -199,8 +202,8 @@ def t(name: str) -> GgufTensorInfo:
             yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden)
             out = t(f"{blk}.attn_output")
             yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
-            yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
-            yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+            yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm")
+            yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm")
 
         mlp = f"{pre}.mlp"
         if config.moe_enabled:
diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py
index 3fe373ca3d..88c6ae1345 100644
--- a/python/freetoken/models/qwen3_5_moe/model.py
+++ b/python/freetoken/models/qwen3_5_moe/model.py
@@ -6,7 +6,7 @@
 from freetoken.core import get_global_ctx
 from freetoken.layers import (
     BaseOP,
-    GemmaRMSNorm,
+    GemmaPlusOneRMSNormFused,
     OPList,
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -58,20 +58,15 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
             if config.moe_enabled
             else Qwen3_5DenseMLP(config, prefix=f"{prefix}.mlp")
         )
-        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps)
 
     @nvtx_annotate("Layer_{}", layer_id_field="_layer_id")
     def forward(self, hidden: torch.Tensor, residual: torch.Tensor | None):
-        # Residual-stream form: fuse each residual-add into the next RMSNorm
-        # (GemmaRMSNorm.forward_add_residual) so add + norm are one kernel per sublayer.
-        if residual is None:
-            residual = hidden
-            hidden = self.input_layernorm.forward(hidden)
-        else:
-            hidden, residual = self.input_layernorm.forward_add_residual(hidden, residual)
+        # Residual-stream form: fuse each residual-add into the next RMSNorm so add + norm are one kernel per sublayer.
+        hidden, residual = self.input_layernorm.forward(hidden, residual)
         hidden = self.linear_attn.forward(hidden) if self._is_linear else self.self_attn.forward(hidden)
-        hidden, residual = self.post_attention_layernorm.forward_add_residual(hidden, residual)
+        hidden, residual = self.post_attention_layernorm.forward(hidden, residual)
         hidden = self.mlp.forward(hidden)
         return hidden, residual
 
@@ -90,14 +85,14 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
                 for layer_id in range(config.num_layers)
             ]
         )
-        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
         x = embed_input_ids(self.embed_tokens, input_ids, get_global_ctx().batch)
         residual: torch.Tensor | None = None
         for layer in self.layers.op_list:
             x, residual = layer.forward(x, residual)
-        x, _ = self.norm.forward_add_residual(x, residual)
+        x, _ = self.norm.forward(x, residual)
         return x
 
 
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index e7c4676dd4..afd32eb2c0 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -35,13 +35,6 @@
 # role -> the expert bank reader's canonical (ModelOpt) tensor kind
 _BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"}
 
-# Gemma-style (1+weight) RMSNorm weights; the GDN gated norm (linear_attn.norm) is a plain weight*x norm
-_GEMMA_NORM_SUFFIXES = (
-    ".input_layernorm.weight",
-    ".post_attention_layernorm.weight",
-    ".self_attn.q_norm.weight",
-    ".self_attn.k_norm.weight",
-)
 # leaves the model builds as Linear layers: only their tensors are read under the QuantConfig, the rest passes through as stored
 _LINEAR_LEAVES = frozenset({
     "q_proj", "k_proj", "v_proj", "o_proj", "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj",
@@ -63,10 +56,6 @@ def _rename(raw_name: str) -> str | None:
     return rename_vl_prefix(raw_name)
 
 
-def _is_gemma_norm(name: str) -> bool:
-    return name == "model.norm.weight" or name.endswith(_GEMMA_NORM_SUFFIXES)
-
-
 def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor:
     """Per-tensor scalar or per-channel ``[rows, 1]`` fp8 scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows."""
     flat = scale.reshape(-1).to(torch.float32)
@@ -274,8 +263,6 @@ def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | N
                 emitted = reader.add(name, tensor)
                 if emitted is not None:
                     yield from emitted
-                elif _is_gemma_norm(name):
-                    yield name, tensor + 1.0  # (1 + weight) baked into the stored norm weight
                 else:
                     yield name, tensor
     if reader is not None and reader.pending:
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index c502067ee4..958d6daf66 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -358,7 +358,7 @@ def _same(a: torch.Tensor, b: torch.Tensor) -> bool:
     return a.dtype is b.dtype and torch.equal(a.view(torch.uint8), b.view(torch.uint8))
 
 
-def test_bf16_fusions_slice_back_and_norms_get_plus_one(checkpoint):
+def test_bf16_fusions_slice_back_and_norms_load_as_stored(checkpoint):
     name, folder, raw = checkpoint
     if name != "bf16":
         pytest.skip("bf16 layout only")
@@ -371,8 +371,8 @@ def test_bf16_fusions_slice_back_and_norms_get_plus_one(checkpoint):
     merged = loaded["model.layers.1.mlp.shared_expert.gate_up_proj.weight"]
     assert _same(merged[:I], raw[f"{LM}.layers.1.mlp.shared_expert.gate_proj.weight"])
     assert _same(merged[I:], raw[f"{LM}.layers.1.mlp.shared_expert.up_proj.weight"])
-    assert torch.equal(loaded["model.norm.weight"], raw[f"{LM}.norm.weight"] + 1.0)
-    assert torch.equal(loaded["model.layers.1.self_attn.q_norm.weight"], raw[f"{attn}.q_norm.weight"] + 1.0)
+    assert torch.equal(loaded["model.norm.weight"], raw[f"{LM}.norm.weight"])
+    assert torch.equal(loaded["model.layers.1.self_attn.q_norm.weight"], raw[f"{attn}.q_norm.weight"])
     assert torch.equal(loaded["model.layers.0.linear_attn.norm.weight"], raw[f"{gdn}.norm.weight"])
     assert loaded["model.layers.0.linear_attn.A_log"].dtype is torch.float32
 

From 64b85c344c7a10862cc1d25ea50facd1d59d8f3b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 21:14:02 -0700
Subject: [PATCH 131/249] fix(qwen4_exp): load gguf files with a quantized
 indexer (#760)

- split index_qk_proj into the indexer's q and k GGUF tensors for the quant lookup
- refuse an MTP-only draft GGUF with a clear error instead of a KeyError
---
 python/freetoken/models/qwen4_exp/attention.py | 2 +-
 python/freetoken/models/qwen4_exp/gguf.py      | 7 ++++++-
 python/freetoken/models/register.py            | 4 +++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py
index 8f6b938870..0ac45ef1ee 100644
--- a/python/freetoken/models/qwen4_exp/attention.py
+++ b/python/freetoken/models/qwen4_exp/attention.py
@@ -85,7 +85,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
         self.eps = config.rms_norm_eps
         self._split = [self.num_heads * self.head_dim, self.num_kv_heads * self.head_dim]
         self.index_qk_proj = LinearReplicated(
-            args.hidden_size, sum(self._split), has_bias=False,
+            args.hidden_size, sum(self._split), has_bias=False, output_sizes=self._split,
             quant_config=config.quant, prefix=f"{prefix}.index_qk_proj",
         )
         self.q_layernorm = GemmaPlusOneRMSNorm(self.head_dim, eps=self.eps)
diff --git a/python/freetoken/models/qwen4_exp/gguf.py b/python/freetoken/models/qwen4_exp/gguf.py
index 0c1bf98059..f29f65337e 100644
--- a/python/freetoken/models/qwen4_exp/gguf.py
+++ b/python/freetoken/models/qwen4_exp/gguf.py
@@ -45,7 +45,8 @@
     "self_attn.k_proj": ("attn_k",),
     "self_attn.v_proj": ("attn_v",),
     "self_attn.o_proj": ("attn_output",),
-    "self_attn.indexer.index_qk_proj": ("indexer.q_proj", "indexer.k_proj"),
+    "self_attn.indexer.index_q_proj": ("indexer.q_proj",),
+    "self_attn.indexer.index_k_proj": ("indexer.k_proj",),
     "mlp.gate": ("ffn_gate_inp",),
     "mlp.shared_expert_gate": ("ffn_gate_inp_shexp",),
     "mlp.shared_expert.gate_proj": ("ffn_gate_shexp",),
@@ -85,6 +86,10 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
 def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
     # NVFP4 weight scales are unvalidated on this arch, so refuse them instead of risking wrong output
     check_scales(shim.model_path, nvfp4=False)
+    # a metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table to check
+    index = gguf_tensor_index(shim.model_path)
+    if index and "blk.0.hc_attn_down.weight" not in index:
+        raise NotImplementedError(f"{shim.model_path}: no decoder layer tensors (blk.0.hc_attn_down), e.g. an MTP-only draft GGUF")
     md = shim.metadata
 
     def kv(key: str):
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 42c6f3f548..80b732cd3f 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -63,6 +63,8 @@ class ModelSpec:
 _QWEN4_EXP_PACKED = _QWEN3_5_PACKED + (
     ("input_mix_weight_down_block_inject", ("input_mix_weight_down", "block_inject_weight")),
 )
+# the llama.cpp converter splits the indexer projection into q and k
+_QWEN4_EXP_GGUF_PACKED = _QWEN4_EXP_PACKED + (("index_qk_proj", ("index_q_proj", "index_k_proj")),)
 # GLM-5.3-Flash's KDA fuses its six input projections.
 _GLM5_NEXT_PACKED = _EXPERTS_PACKED + (
     ("in_proj", ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj")),
@@ -288,7 +290,7 @@ class ModelSpec:
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
         checkpoint_roots=_LANGUAGE_MODEL_ROOT,
-        packed_modules_mapping=_QWEN4_EXP_PACKED,
+        packed_modules_mapping=_QWEN4_EXP_GGUF_PACKED,
     ),
     # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the
     # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the

From 70550f8f9e368d62155a98c59a10b91519ac94eb Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 19 Oct 2026 21:30:00 -0700
Subject: [PATCH 132/249] feat(scripts): add checkpoint diagnosis and fix
 scripts (#762)

---
 AGENTS.md                                     |    3 +
 README.md                                     |    1 +
 docs/models.md                                |    1 +
 scripts/checkpoint_diagnose/README.md         |  177 +++
 scripts/checkpoint_diagnose/ckpt_edit.py      |  512 ++++++++
 scripts/checkpoint_diagnose/diagnose.py       | 1040 +++++++++++++++++
 scripts/checkpoint_diagnose/fetch_headers.py  |  137 +++
 scripts/checkpoint_diagnose/fix_config.py     |  102 ++
 scripts/checkpoint_diagnose/rename_tensors.py |  123 ++
 9 files changed, 2096 insertions(+)
 create mode 100644 scripts/checkpoint_diagnose/README.md
 create mode 100644 scripts/checkpoint_diagnose/ckpt_edit.py
 create mode 100644 scripts/checkpoint_diagnose/diagnose.py
 create mode 100644 scripts/checkpoint_diagnose/fetch_headers.py
 create mode 100644 scripts/checkpoint_diagnose/fix_config.py
 create mode 100644 scripts/checkpoint_diagnose/rename_tensors.py

diff --git a/AGENTS.md b/AGENTS.md
index a2d263db5b..0d0878d9eb 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -34,6 +34,8 @@ benchmarks/            end-to-end and micro benchmarks, see benchmarks/README.md
 docs/                  install, quickstart, CLI and model docs
 freetoken-kernel-cache/ companion wheel of prebuilt kernels, see its README
 scripts/               wheel build and release scripts
+  checkpoint_diagnose/ why an HF safetensors checkpoint does not load, and fixes for non-standard ones
+  gguf_check/          which GGUF files of a family load
 ```
 
 ## Development
@@ -54,6 +56,7 @@ Put a new test in the `tests/` directory that mirrors the module it protects, an
 
 - Search existing issues and PRs before starting. Items on the [Roadmap](https://github.com/FlashML-org/FreeToken/issues/79) are discussed with maintainers before implementation; features not on it start as an issue.
 - When helping the user draft an issue, follow the matching template in `.github/ISSUE_TEMPLATE/` (engine bug, model checkpoint, feature request) and fill in every required field: hardware, driver, FreeToken version, checkpoint ID, exact command, and the full log.
+- When a checkpoint fails to load, follow [scripts/checkpoint_diagnose/README.md](scripts/checkpoint_diagnose/README.md) before drafting a model checkpoint issue: it tells a FreeToken gap from a non-standard checkpoint, and repairs the latter.
 - One change per PR, linked to its issue, with the hardware, checkpoint ID and exact command it was tested with.
 
 ## Code comments
diff --git a/README.md b/README.md
index 6d9ab4ae47..e72a5cd0d6 100644
--- a/README.md
+++ b/README.md
@@ -56,6 +56,7 @@ For More details:
 - [Supported models](https://github.com/FlashML-org/FreeToken/blob/main/docs/models.md)
 - [CLI reference](https://github.com/FlashML-org/FreeToken/blob/main/docs/cli.md)
 - [Repairing old FTW checkpoints](https://github.com/FlashML-org/FreeToken/blob/main/docs/ftw-hotfix.md)
+- [Diagnosing a checkpoint that fails to load](https://github.com/FlashML-org/FreeToken/blob/main/scripts/checkpoint_diagnose/README.md)
 
 ## Citation
 
diff --git a/docs/models.md b/docs/models.md
index 18f73fac54..f6d99af998 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -78,6 +78,7 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
 
 ## Notes
 
+- A checkpoint that fails to load: from a source checkout, run `PYTHONPATH=python python scripts/checkpoint_diagnose/diagnose.py <checkpoint dir or HF repo id>`. It reads the headers, not the weights, says whether FreeToken or the checkpoint is at fault, and for a local checkpoint gives the commands that fix a non-standard one; see [its README](../scripts/checkpoint_diagnose/README.md). GGUF files are checked with [scripts/gguf_check](../scripts/gguf_check/README.md).
 - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into
   FreeToken's fast-load format, and `ft serve --model` auto-detects the result.
 - An FTW converted by an older build with `--moe-backend fused` or `triton` keeps
diff --git a/scripts/checkpoint_diagnose/README.md b/scripts/checkpoint_diagnose/README.md
new file mode 100644
index 0000000000..578952889e
--- /dev/null
+++ b/scripts/checkpoint_diagnose/README.md
@@ -0,0 +1,177 @@
+# checkpoint_diagnose
+
+Find out why FreeToken does not load a checkpoint, sort the failure into one kind, and repair the checkpoints whose own files are non-standard.
+Run everything from the repo root with `PYTHONPATH=python`. Nothing here needs a GPU or the checkpoint's weights.
+
+> [!WARNING]
+> The diagnosis is heuristic and can be wrong: the failure kind, the findings and the suggested fixes are a best guess from the headers and configs, not a proof. Check them against the problems it prints before acting on them.
+
+> [!IMPORTANT]
+> For agents: explain the diagnosis to the user in plain terms first: what fails, which kind it is, and what the fix would change. Ask the user before you run a fix script or open an issue. Do not change the user's checkpoint on your own.
+
+| Script | Does |
+|---|---|
+| `diagnose.py` | checks a checkpoint dir or HF repo id, prints the failure kind and the next step |
+| `fix_config.py` | level 1 fix: edits only `config.json` / `hf_quant_config.json` |
+| `rename_tensors.py` | level 2 fix: renames tensors (shard headers, index, quant skip lists); tensor data is kept byte for byte |
+| `fetch_headers.py` | header-only stub of an HF repo; `diagnose.py` calls it for a repo id |
+| `ckpt_edit.py` | the edit primitives both fix scripts use, for a case they do not cover |
+
+## Step 1: which format
+
+`diagnose.py` routes by itself, but know where each format goes:
+
+| The checkpoint has | Go to |
+|---|---|
+| `*.gguf` files | `scripts/gguf_check/` (its README), then the GGUF section below |
+| `freetoken_weight.json` (FTW, made by `ft checkpoint`) | diagnose the HF checkpoint it was converted from; FTW repairs are in `docs/ftw-hotfix.md` |
+| root `*.safetensors` shards | this directory, from step 2 |
+| only `*.bin` / `*.pt`, or MLX weights | not served: FreeToken reads safetensors and GGUF |
+
+### GGUF: three kinds of failure
+
+`scripts/gguf_check/check.py` prints the problems of each file; the first one tells the kind.
+
+| Kind | check.py says | Do |
+|---|---|---|
+| a ggml type has no kernel | `<layer kind> <type>: no <ops> kernel` (needs `--kernels` from `probe_kernels.py`) | the build or GPU lacks kernels for that quant type: pick a file with another type, or report it with the `probe_kernels.py` table |
+| the model's loader does not read the file | `unexpected tensor ...`, `N tensors never loaded`, `tensor ... yielded twice`, `<tensor>: got ..., model has ...`, `MoE layer N ... bank is ..., the model needs ...`, or a `NotImplementedError` from the family's config parse | a FreeToken gap for that layout: report it with the check.py record |
+| the GGUF file itself | `is not a GGUF file`, `duplicate GGUF key / tensor`, `runs past the end of the file`, `bad general.alignment`, `tokenizer: no EOS token`, `tokenizer: no chat template`, `chat template: <error>` | the file is broken or incomplete: download it again, or use another file of the repo |
+
+A file whose architecture FreeToken does not serve shows as `skip`, not as a failure.
+
+## Step 2: diagnose a safetensors checkpoint
+
+```bash
+python scripts/checkpoint_diagnose/diagnose.py <ckpt dir | HF repo id[@revision]> --out report.jsonl
+```
+
+- A local dir is read as is: the headers and JSON files, so a 400 GB checkpoint takes seconds. A weight reader that needs tensor values gets the real bytes it asks for.
+- A repo id is fetched as a header-only stub into `~/.cache/freetoken/checkpoint-stubs/` (a few MB). Nothing is downloaded beyond that, and no `.py` file: a repo whose config is its own code (`auto_map`) is skipped, diagnose a local download of it instead.
+- A remote repo is only diagnosed: its record has no fixes, and the fix scripts refuse a stub. When it says `fixable`, download the checkpoint and diagnose the local directory to get the fix commands.
+- Add `--text-model-only` when the checkpoint will be served with `--text-model-only`.
+- `--out` appends one JSON record per checkpoint. The fix scripts read their fixes from it with `--from-report`.
+
+It first checks the files themselves (data offsets cover each shard exactly, the index agrees with the shard headers), then runs the server's own load path on the CPU with meta tensors: config parse, model build, dense weight load through the model's `load_state_dict`, and expert-bank fill for MoE models with the reader the engine picks.
+When that fails, it also checks the tensor names against transformers' own model class and its checkpoint conversion rules, and lints the files for the known non-standard forms.
+
+Output, per checkpoint: `PASS`, or `FAIL at <stage>: [<kind>] ...` with the problems, what transformers makes of the names, the findings, and the next step.
+
+## Step 3: the kind of failure
+
+Each failure gets one kind. `diagnose.py` tries, in this order: broken files (`checkpoint`), `arch`, `quant_method`, a broken prefix (`checkpoint`), `quant_model`, names transformers cannot map or another non-standard finding (`checkpoint`), `packing`; anything left is `loader`.
+
+| Kind | Means | Do |
+|---|---|---|
+| `arch` | the architecture is not registered in `python/freetoken/models/register.py` | not served; a new architecture starts as a feature request |
+| `quant_method` | FreeToken has no reader for the quantization format (GPTQ, AWQ, MLX, bitsandbytes, an unsupported compressed-tensors scheme, ...) | not served; use a release in a format `python/freetoken/layers/quantization/configs/` reads |
+| `quant_model` | the format is read, but this model or layer kind has no method for it (`no quant method for moe.<kind>`), or the reader does not produce the scales the method needs | a FreeToken gap: report it |
+| `loader` | the model and format are served, and transformers loads these names, but FreeToken's reader does not | a FreeToken gap: report it. A finding marked `loader` may offer a fix that works around it |
+| `packing` | the expert banks do not fill, and transformers repacks these expert tensors (for example one tensor per expert where the family's release stacks them) | a FreeToken gap: report it. Repacking tensor data is not offered here |
+| `checkpoint` | the checkpoint's own files are broken or non-standard: an incomplete shard or index, names transformers would not load either, or a config that contradicts what the shards store | fix it, step 4 |
+
+To report a FreeToken gap, use the "model checkpoint" issue template and paste the `diagnose.py` output.
+
+How to tell `loader` / `packing` from `checkpoint`: the line `transformers ... would not load it either` means transformers' model class and conversion rules cannot map the names either, so the checkpoint is at fault. `transformers loads these names` means the names are a form transformers accepts, so FreeToken should read them too.
+
+## Step 4: fix a non-standard checkpoint
+
+Fixes apply to a local copy of the checkpoint only, never to a remote repo or its stub.
+
+Rules:
+
+1. Fix into a new directory (`--out <new dir>`). Unchanged files there are symlinks to the original, so it costs only the rewritten shards. Use `--in-place` only when the user asks for it. It keeps a backup, and `--restore` undoes the latest in-place edit.
+2. One level per round. If the record has a `rename` fix, apply that first, then run `diagnose.py` on the new directory, then apply the config fixes of that new record. A rename changes the names the config findings talk about.
+3. Read the plan before applying it. Every run prints the plan first and changes nothing until it is confirmed. Check that:
+   - each rule renames the tensors you expect, and the example `old -> new` lines are right;
+   - nothing is renamed that should not be. A `warning: ... matches are inside the name` line means the pattern is not anchored;
+   - the config edits name the right modules;
+   - the shard costs are what you expect. `rewritten, copies N GiB` needs that much free disk.
+
+   If anything in the plan is not what you expected, do not apply it: change the rule or the edit and plan again.
+4. Confirming: in a terminal the script asks `[y/N]`. Without a terminal (an agent), pass `--plan <file>`. The first run saves the plan there and changes nothing. Read it, show it to the user, and run the exact same command again to apply it only after the user agrees. If the command or the checkpoint changed in between, the second run saves the new plan instead of applying anything.
+5. After every fix, run `diagnose.py` on the result, until it passes or nothing fixable is left.
+
+The usual round, for an agent:
+
+```bash
+D=scripts/checkpoint_diagnose
+python $D/diagnose.py /models/x --out r0.jsonl
+python $D/rename_tensors.py /models/x --out /models/x-fixed --from-report r0.jsonl --plan plan.json   # prints and saves the plan
+python $D/rename_tensors.py /models/x --out /models/x-fixed --from-report r0.jsonl --plan plan.json   # after the user agrees to the plan: applies it
+python $D/diagnose.py /models/x-fixed --out r1.jsonl
+python $D/fix_config.py /models/x-fixed --in-place --from-report r1.jsonl --plan plan.json            # twice, as above
+python $D/diagnose.py /models/x-fixed
+```
+
+`--from-report` applies the fixes `diagnose.py` found for that checkpoint; anything else is a manual edit, below.
+
+### Fix levels
+
+| Level | Script | Changes | Typical use |
+|---|---|---|---|
+| 1 | `fix_config.py` | `config.json`, `hf_quant_config.json` | the quantization skip lists, compressed-tensors targets, a config key |
+| 2 | `rename_tensors.py` | shard headers, `model.safetensors.index.json`, and the skip-list entries / targets that are module names or globs (a `re:` entry is not rewritten; the plan warns about every entry left naming the old modules) | broken prefixes, legacy tensor names |
+| - | not offered | tensor data | repacking experts, dropping tensors, dequantizing: re-export the checkpoint, or report the FreeToken gap |
+
+### Manual edits
+
+When no automatic fix fits, write the edit yourself. It goes through the same plan and confirmation.
+
+```bash
+# rename: re.sub(PATTERN, REPL, name), rules in order; anchor with ^ and escape dots
+python $D/rename_tensors.py /models/x --out /models/x-fixed --rule '^model\.language_model\.visual\.' 'model.visual.'
+
+# config key: dotted path, JSON value
+python $D/fix_config.py /models/x --out /models/x-fixed --set text_config.tie_word_embeddings false
+
+# quantization skip lists, in the dialect's syntax:
+#   ModelOpt: glob (model.language_model.layers.*.mlp.gate)
+#   compressed-tensors: module name or re:<regex> (re:^model\.visual\..*)
+#   fp8: module name (model.language_model.layers.0.linear_attn.in_proj_a)
+python $D/fix_config.py /models/x --out /models/x-fixed --skip-add 'model.language_model.layers.*.mlp.gate'
+python $D/fix_config.py /models/x --out /models/x-fixed --skip-remove 'lm_head'
+```
+
+A rule that renames nothing, and a rename that gives two tensors the same name, are refused. So is an `--out` directory that already has files in it.
+
+### Known non-standard forms
+
+The findings `diagnose.py` prints, and what they fix. `kind` is the failure kind the finding points to.
+
+| Finding | The checkpoint | Level | Fix |
+|---|---|---|---|
+| `prefix` | nests the text tower more than once (`model.language_model.language_model.`) | 2 | rename to `model.language_model.` |
+| `prefix` | stores the vision tower under the text tower (`model.language_model.visual.`) | 2 | rename to `model.visual.` |
+| `prefix` | uses vLLM's internal names (`language_model.model.`, `language_model.lm_head.`) | 2 | rename to `model.language_model.` / `lm_head.` |
+| `quant_declared_but_dense` | stores modules in bf16 that its quantization_config does not skip (vision tower, lm_head) | 1 | add them to the skip lists |
+| `quant_skipped_but_quantized` | stores modules quantized that its skip list covers, often through a container entry (`...linear_attn`) | 1 | remove the entry, and list its dense children instead |
+| `ct_target_vllm_names` | writes compressed-tensors targets with vLLM's names, so they match no tensor | 1 | rewrite the targets to `model.language_model.` |
+| `quant_unlisted_dense_router` (kind `loader`) | stores in bf16, without listing them as skipped, modules the family serves in bf16 anyway (the Qwen3.5 routers) | 1 | list them; works around the reader |
+| `kv_output_scale` (kind `loader`) | stores fp8 KV scales as `k_proj.output_scale` / `v_proj.output_scale` (only reported for the families whose reader drops `k_scale`) | 2 | rename to `k_scale` / `v_scale`, which the reader drops |
+| `expert_tensors_incomplete` | lacks some expert's tensors in some layer | - | download it again, or ask its uploader |
+| (stage `files`) | has a shard whose tensors do not cover the file, or an index that disagrees with the shards | - | download it again |
+| `no_vision_tower` | is text-only under a multimodal architecture | runtime | serve with `--text-model-only` |
+| `no_processor_config` | has a vision tower but no `preprocessor_config.json` (image requests fail with 400) | file | copy that file from the base model |
+| `fp8_without_config` | stores fp8 weights with no quantization_config | - | not fixable here; the scales and block size are unknown |
+| `mlx` (kind `quant_method`) | is MLX-quantized (`.scales` / `.biases`) | - | not served; use a non-MLX release |
+
+## Step 5: check the fixed checkpoint
+
+`diagnose.py` passing means: the names, shapes, quantization methods and expert-bank layout all line up with what the model builds.
+It does not read weight values. On a real checkpoint, finish with `ft serve --model-path <fixed dir>` and a short request. For a checkpoint with a vision tower, also send an image request.
+
+## Using the primitives
+
+```python
+import sys; sys.path.insert(0, "scripts/checkpoint_diagnose")
+from ckpt_edit import Checkpoint, confirm, plan_lines_for_save
+
+ck = Checkpoint("/models/x")
+ck.rename_regex([(r"^old\.", "new.")])
+ck.quant_skip_add(["model.visual.*"])
+print("\n".join(ck.log + plan_lines_for_save(ck, "/models/x-fixed")))
+ck.save(out="/models/x-fixed")      # ck.save() edits in place, with a backup
+```
+
+`ck.tensors()` gives every tensor's shard, dtype and shape from the headers. `ck.quant_configs()` gives every quantization config the checkpoint carries.
diff --git a/scripts/checkpoint_diagnose/ckpt_edit.py b/scripts/checkpoint_diagnose/ckpt_edit.py
new file mode 100644
index 0000000000..5ec5215f55
--- /dev/null
+++ b/scripts/checkpoint_diagnose/ckpt_edit.py
@@ -0,0 +1,512 @@
+"""Edit primitives for HF safetensors checkpoints: config keys, quantization skip lists, tensor names. Tensor data is never changed.
+
+    from ckpt_edit import Checkpoint
+    ck = Checkpoint("/path/to/ckpt")
+    ck.set_config("text_config.tie_word_embeddings", False)
+    ck.quant_skip_add(["model.language_model.layers.*.mlp.gate"])
+    ck.rename(lambda name: name.replace("model.language_model.language_model.", "model.language_model."))
+    ck.save(out="/path/to/fixed")      # or ck.save() to edit in place
+
+``save(out=...)`` writes a new directory: changed files are written there, every other file is a symlink to the original.
+``save()`` edits in place and first copies each changed JSON file and shard header to ``.ckpt_edit_backup/<time>/``; ``restore()`` puts the latest backup back.
+A renamed shard keeps its data bytes: its header is patched in place when the new one fits in the old one's bytes, else the shard is rewritten through a temp file.
+A shard that is a symlink (hub cache) or has other hard links is never written through: it becomes a private copy.
+"""
+
+from __future__ import annotations
+
+import copy
+import errno
+import fnmatch
+import hashlib
+import json
+import os
+import re
+import shutil
+import struct
+import sys
+import time
+from typing import Any, Callable, Iterable
+
+INDEX = "model.safetensors.index.json"
+CONFIG = "config.json"
+SIDECAR = "hf_quant_config.json"
+BACKUP_DIR = ".ckpt_edit_backup"
+STUB_MARKER = ".stub.json"
+# keys that list modules a quantization_config leaves unquantized
+SKIP_KEYS = ("ignore", "exclude_modules", "modules_to_not_convert")
+_DEFAULT_SKIP_KEY = {"modelopt": "ignore", "compressed-tensors": "ignore", "fp8": "modules_to_not_convert", "mxfp4": "modules_to_not_convert"}
+
+
+def header_len(path: str) -> int:
+    with open(path, "rb") as f:
+        return struct.unpack("<Q", f.read(8))[0]
+
+
+def read_header(path: str) -> tuple[dict, int]:
+    """(header, header byte length) of a safetensors file."""
+    with open(path, "rb") as f:
+        (n,) = struct.unpack("<Q", f.read(8))
+        return json.loads(f.read(n)), n
+
+
+def encode_header(header: dict, fit: int | None = None) -> bytes:
+    """The header JSON padded with spaces to ``fit`` bytes when it fits, else to a multiple of 8."""
+    raw = json.dumps(header, separators=(",", ":")).encode()
+    size = fit if fit is not None and len(raw) <= fit else (len(raw) + 7) // 8 * 8
+    return raw + b" " * (size - len(raw))
+
+
+def _copy_range(src_fd: int, dst_fd: int, src_off: int, dst_off: int, count: int) -> None:
+    fast = hasattr(os, "copy_file_range")
+    while count:
+        n = 0
+        if fast:
+            try:
+                n = os.copy_file_range(src_fd, dst_fd, count, src_off, dst_off)
+            except OSError:
+                fast = False
+        if n <= 0:
+            buf = os.pread(src_fd, min(count, 64 << 20), src_off)
+            if not buf:
+                raise EOFError("source shard ended early")
+            n = os.pwrite(dst_fd, buf, dst_off)
+        count, src_off, dst_off = count - n, src_off + n, dst_off + n
+
+
+def stub_data_end(path: str, n: int) -> int | None:
+    """For a fetch_headers.py stub, the end of the block that holds the header; past it the file is a hole. None for anything else.
+
+    A shard counts as a stub only next to the stub marker and when SEEK_DATA finds no data past the header's block."""
+    if not os.path.isfile(os.path.join(os.path.dirname(os.path.abspath(path)), STUB_MARKER)):
+        return None
+    st = os.stat(path)
+    end = min(st.st_size, (8 + n + st.st_blksize - 1) // st.st_blksize * st.st_blksize)
+    if end >= st.st_size:
+        return None
+    fd = os.open(path, os.O_RDONLY)
+    try:
+        os.lseek(fd, end, os.SEEK_DATA)
+    except OSError as e:
+        return end if e.errno == errno.ENXIO else None
+    finally:
+        os.close(fd)
+    return None
+
+
+def _fsync_dir(path: str) -> None:
+    fd = os.open(os.path.dirname(os.path.abspath(path)) or ".", os.O_RDONLY)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def write_shard(src: str, dst: str, header: dict | bytes, *, private: bool = False) -> str:
+    """Write ``header`` (a dict, or the exact header bytes) with the data bytes of ``src`` to ``dst``; returns ``"in place"`` or ``"rewritten"``.
+
+    ``private`` forces a rewrite through a temp file even when ``dst`` is ``src``, so a symlink or a hard link is replaced instead of written through."""
+    n_old = header_len(src)
+    size = os.path.getsize(src)
+    same = os.path.realpath(src) == os.path.realpath(dst)
+    data = header if isinstance(header, bytes) else encode_header(header, fit=n_old if same else None)
+    if same and not private and len(data) == n_old:
+        fd = os.open(dst, os.O_WRONLY)
+        try:
+            os.pwrite(fd, data, 8)
+            os.fsync(fd)
+        finally:
+            os.close(fd)
+        return "in place"
+    tmp = dst + ".ckpt_edit.part"
+    head = struct.pack("<Q", len(data)) + data
+    stub_end = stub_data_end(src, n_old)
+    src_fd, dst_fd = os.open(src, os.O_RDONLY), os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
+    try:
+        os.pwrite(dst_fd, head, 0)
+        if stub_end is not None:
+            # a stub's data is a hole: copy the bytes up to the end of the header's block, keep the rest a hole
+            _copy_range(src_fd, dst_fd, 8 + n_old, len(head), stub_end - 8 - n_old)
+            os.ftruncate(dst_fd, len(head) + size - 8 - n_old)
+        else:
+            _copy_range(src_fd, dst_fd, 8 + n_old, len(head), size - 8 - n_old)
+        os.fsync(dst_fd)
+    except BaseException:
+        os.close(dst_fd)
+        os.remove(tmp)
+        raise
+    finally:
+        os.close(src_fd)
+    os.close(dst_fd)
+    os.replace(tmp, dst)
+    _fsync_dir(dst)
+    return "rewritten"
+
+
+def _shared(path: str) -> bool:
+    """A symlink or a file with other hard links: writing through it would change another copy."""
+    return os.path.islink(path) or os.stat(path).st_nlink > 1
+
+
+def get_path(cfg: dict, path: str) -> Any:
+    for key in path.split("."):
+        if not isinstance(cfg, dict) or key not in cfg:
+            return None
+        cfg = cfg[key]
+    return cfg
+
+
+def set_path(cfg: dict, path: str, value: Any) -> None:
+    *parents, leaf = path.split(".")
+    for key in parents:
+        cfg = cfg.setdefault(key, {})
+    cfg[leaf] = value
+
+
+def entry_names(entry: str, module: str) -> bool:
+    """Whether a skip-list entry or target covers ``module``, read loosely across dialects (literal name or parent, glob, ``re:``)."""
+    if entry.startswith("re:"):
+        try:
+            return re.match(entry[3:], module) is not None
+        except re.error:
+            return False
+    names = [module]
+    while "." in names[-1]:
+        names.append(names[-1].rpartition(".")[0])
+    return any(n == entry or fnmatch.fnmatchcase(n, entry) for n in names)
+
+
+class Checkpoint:
+    def __init__(self, path: str):
+        self.path = os.path.abspath(path)
+        self.json: dict[str, dict] = {}
+        for name in (CONFIG, SIDECAR, INDEX):
+            p = os.path.join(self.path, name)
+            if os.path.isfile(p):
+                with open(p, encoding="utf-8") as f:
+                    self.json[name] = json.load(f)
+        self._orig_json = copy.deepcopy(self.json)
+        self.headers: dict[str, dict] = {}
+        self.header_len: dict[str, int] = {}
+        digest = hashlib.sha256(json.dumps(self.json, sort_keys=True).encode())
+        for name in sorted(os.listdir(self.path)):
+            if name.endswith(".safetensors"):
+                self.headers[name], self.header_len[name] = read_header(os.path.join(self.path, name))
+                digest.update(name.encode() + json.dumps(self.headers[name], sort_keys=True).encode())
+        # the state the edits are planned against; a plan made on another state is refused
+        self.fingerprint = digest.hexdigest()
+        self._changed_shards: set[str] = set()
+        self.log: list[str] = []
+
+    # --- reading ---
+
+    def tensors(self) -> dict[str, tuple[str, dict]]:
+        """name -> (shard, {dtype, shape, data_offsets})."""
+        return {k: (shard, v) for shard, h in self.headers.items() for k, v in h.items() if k != "__metadata__"}
+
+    @property
+    def config(self) -> dict:
+        return self.json.setdefault(CONFIG, {})
+
+    def quant_configs(self) -> list[tuple[str, dict]]:
+        """(where, dict) of every quantization config the checkpoint carries: config.json top level, its text_config, and hf_quant_config.json."""
+        found = []
+        for where, q in ((CONFIG, get_path(self.config, "quantization_config")),
+                         (f"{CONFIG}:text_config", get_path(self.config, "text_config.quantization_config")),
+                         (SIDECAR, get_path(self.json.get(SIDECAR, {}), "quantization"))):
+            if isinstance(q, dict):
+                found.append((where, q))
+        return found
+
+    def _method(self, where: str, q: dict) -> str:
+        return "modelopt" if where == SIDECAR else str(q.get("quant_method") or "").lower()
+
+    def quant_entries(self) -> list[tuple[str, str, list]]:
+        """(where, key, list) of every skip list and compressed-tensors targets list."""
+        lists = []
+        for where, q in self.quant_configs():
+            lists += [(where, k, q[k]) for k in SKIP_KEYS if isinstance(q.get(k), list)]
+            for gname, group in (q.get("config_groups") or {}).items():
+                if isinstance(group, dict) and isinstance(group.get("targets"), list):
+                    lists.append((where, f"config_groups.{gname}.targets", group["targets"]))
+        return lists
+
+    # --- config edits ---
+
+    def set_config(self, path: str, value: Any) -> None:
+        old = get_path(self.config, path)
+        set_path(self.config, path, value)
+        self.log.append(f"config.json {path}: {json.dumps(old)} -> {json.dumps(value)}")
+
+    def quant_skip_add(self, entries: Iterable[str]) -> None:
+        """Add module names or patterns to the skip lists of every quantization config, creating the dialect's list if it has none.
+
+        An empty list next to a non-empty one is left alone (ModelOpt readers take ``ignore`` only when it is non-empty); an fp8 list that is created or empty first gets the ``lm_head`` the dialect skips by default."""
+        entries = list(entries)
+        for where, q in self.quant_configs():
+            method = self._method(where, q)
+            keys = [k for k in SKIP_KEYS if isinstance(q.get(k), list)]
+            if not keys:
+                keys = ["exclude_modules" if where == SIDECAR else _DEFAULT_SKIP_KEY.get(method, "ignore")]
+                q[keys[0]] = []
+            if any(q[k] for k in keys):
+                keys = [k for k in keys if q[k]]
+            for key in keys:
+                lst = q[key]
+                implicit = ["lm_head"] if method == "fp8" and key == "modules_to_not_convert" and not lst else []
+                new = list(dict.fromkeys(e for e in implicit + entries if e not in lst))
+                lst.extend(new)
+                if new:
+                    self.log.append(f"{where} {key}: added {new}")
+
+    def quant_skip_remove(self, entries: Iterable[str]) -> None:
+        entries = set(entries)
+        for where, q in self.quant_configs():
+            for key in SKIP_KEYS:
+                lst = q.get(key)
+                if isinstance(lst, list):
+                    gone = [e for e in lst if e in entries]
+                    lst[:] = [e for e in lst if e not in entries]
+                    if gone:
+                        self.log.append(f"{where} {key}: removed {gone}")
+
+    def quant_sub(self, pattern: str, repl: str) -> None:
+        """re.sub over the text of every skip-list entry and compressed-tensors target, ``re:`` entries included."""
+        self._sub_entries(lambda e: re.sub(pattern, repl, e))
+
+    def quant_rename(self, pattern: str, repl: str) -> None:
+        """Apply a tensor rename rule to the skip-list entries and targets that are module names or globs; ``re:`` entries are left alone.
+
+        Each entry is matched as ``entry + "."``, so a rule written for tensor names also renames a container entry."""
+        rx = re.compile(pattern)
+        self._sub_entries(lambda e: e if e.startswith("re:") else rx.sub(repl, e + ".").removesuffix("."))
+
+    def _sub_entries(self, fn: Callable[[str], str]) -> None:
+        for where, key, lst in self.quant_entries():
+            new = [fn(e) for e in lst]
+            if new != lst:
+                self.log.append(f"{where} {key}: {[e for e in lst if e not in new]} -> {[e for e in new if e not in lst]}")
+                lst[:] = new
+
+    def stale_quant_entries(self, old_modules: Iterable[str], new_modules: Iterable[str]) -> list[str]:
+        """Skip-list entries and targets that named some of ``old_modules`` and name none of ``new_modules``."""
+        old, new = list(old_modules), list(new_modules)
+        stale = []
+        for where, key, lst in self.quant_entries():
+            for e in lst:
+                if e != "Linear" and any(entry_names(e, m) for m in old) and not any(entry_names(e, m) for m in new):
+                    stale.append(f"{where} {key}: {e!r}")
+        return stale
+
+    # --- tensor names ---
+
+    def rename(self, fn: Callable[[str], str]) -> dict[str, str]:
+        """Rename tensors by ``fn(name)``; returns {old: new} of the names that changed. Refuses a rename that makes two tensors share a name."""
+        changes = {}
+        for shard, h in self.headers.items():
+            for k in h:
+                if k != "__metadata__" and fn(k) != k:
+                    changes[k] = (shard, fn(k))
+        if not changes:
+            return {}
+        final = [changes[k][1] if k in changes else k for k in self.tensors()]
+        if len(final) != len(set(final)):
+            dup = sorted({n for n in final if final.count(n) > 1})[:5]
+            raise ValueError(f"the rename gives two tensors the same name, e.g. {dup}")
+        for old, (shard, new) in changes.items():
+            h = self.headers[shard]
+            h[new] = h.pop(old)
+            self._changed_shards.add(shard)
+        index = self.json.get(INDEX)
+        if index is not None:
+            wm = index["weight_map"]
+            index["weight_map"] = {(changes[k][1] if k in changes else k): v for k, v in wm.items()}
+        self.log.append(f"renamed {len(changes)} tensors in {len({s for s, _ in changes.values()})} shards")
+        return {old: new for old, (_, new) in changes.items()}
+
+    def rename_regex(self, rules: Iterable[tuple[str, str]]) -> dict[str, str]:
+        """Apply ``re.sub(pattern, repl, name)`` for each rule in order."""
+        compiled = [(re.compile(p), r) for p, r in rules]
+
+        def fn(name: str) -> str:
+            for rx, repl in compiled:
+                name = rx.sub(repl, name)
+            return name
+
+        return self.rename(fn)
+
+    # --- writing ---
+
+    def _changed_json(self) -> list[str]:
+        return [n for n, v in self.json.items() if v != self._orig_json.get(n)]
+
+    def pending(self, out: str | None) -> dict:
+        """What ``save(out)`` would write: {"json": [file names], "shards": {shard: how}}."""
+        shards = {}
+        for shard in sorted(self._changed_shards):
+            p, n = os.path.join(self.path, shard), self.header_len[shard]
+            fits = len(encode_header(self.headers[shard], fit=n)) == n
+            if out is None and fits and not _shared(p):
+                shards[shard] = "header patched in place"
+            elif stub_data_end(p, n) is not None:
+                shards[shard] = "rewritten (header-only stub, no data copied)"
+            else:
+                why = " (a shared symlink / hard link becomes a private copy)" if out is None and _shared(p) else ""
+                shards[shard] = f"rewritten, copies {os.path.getsize(p) / 2**30:.1f} GiB{why}"
+        return {"json": self._changed_json(), "shards": shards}
+
+    def save(self, out: str | None = None) -> list[str]:
+        """Write the edits; returns what was written. ``out=None`` edits in place (with a backup)."""
+        changed_json, written = self._changed_json(), []
+        if out is None:
+            if not changed_json and not self._changed_shards:
+                return []
+            backup = os.path.join(self.path, BACKUP_DIR, time.strftime("%Y%m%d-%H%M%S") + f"-{time.time_ns() % 10**9:09d}")
+            os.makedirs(backup)
+            manifest = {"links": {}, "created": []}
+            for name in changed_json:
+                p = os.path.join(self.path, name)
+                if os.path.lexists(p):
+                    if os.path.islink(p):
+                        manifest["links"][name] = os.readlink(p)
+                    shutil.copyfile(p, os.path.join(backup, name))
+                else:
+                    manifest["created"].append(name)
+            for shard in self._changed_shards:
+                p = os.path.join(self.path, shard)
+                if os.path.islink(p):
+                    manifest["links"][shard] = os.readlink(p)
+                with open(p, "rb") as f:
+                    (n,) = struct.unpack("<Q", f.read(8))
+                    head = f.read(n)
+                with open(os.path.join(backup, shard + ".header"), "wb") as f:
+                    f.write(head)
+            with open(os.path.join(backup, "manifest.json"), "w") as f:
+                json.dump(manifest, f)
+            for shard in sorted(self._changed_shards):
+                p = os.path.join(self.path, shard)
+                written.append(f"{shard} ({write_shard(p, p, self.headers[shard], private=_shared(p))})")
+            target = self.path
+        else:
+            target = os.path.realpath(out)
+            if target == os.path.realpath(self.path) or target.startswith(os.path.realpath(self.path) + os.sep):
+                raise ValueError("--out must be outside the checkpoint directory")
+            os.makedirs(target, exist_ok=True)
+            for name in sorted(os.listdir(self.path)):
+                src, dst = os.path.join(self.path, name), os.path.join(target, name)
+                if name == BACKUP_DIR or name in changed_json or name in self._changed_shards or os.path.lexists(dst):
+                    continue
+                os.symlink(os.path.realpath(src), dst)
+            for shard in sorted(self._changed_shards):
+                write_shard(os.path.join(self.path, shard), os.path.join(target, shard), self.headers[shard])
+                written.append(f"{shard} (rewritten)")
+        for name in changed_json:
+            dst = os.path.join(target, name)
+            tmp = dst + ".ckpt_edit.part"
+            with open(tmp, "w", encoding="utf-8") as f:
+                json.dump(self.json[name], f, indent=2, ensure_ascii=False)
+                f.write("\n")
+                f.flush()
+                os.fsync(f.fileno())
+            # replaces a symlink without writing through it
+            os.replace(tmp, dst)
+            written.append(name)
+        self._orig_json = copy.deepcopy(self.json)
+        self._changed_shards.clear()
+        return written
+
+
+def confirm(plan: dict, lines: list[str], plan_file: str | None) -> bool:
+    """Print the plan, then decide whether to apply it.
+
+    With ``plan_file`` (no prompt): a run whose plan is not already in the file writes it there and changes nothing; the next run of the same command applies it only if it computes the same plan from the same checkpoint state.
+    Without: ask [y/N], and change nothing when stdin is not a terminal."""
+    print("\n".join(lines))
+    if plan_file:
+        saved = None
+        if os.path.exists(plan_file):
+            try:
+                with open(plan_file) as f:
+                    saved = json.load(f)
+            except (OSError, ValueError):
+                saved = None
+            if not isinstance(saved, dict) or "fingerprint" not in saved:
+                raise SystemExit(f"--plan {plan_file} exists and is not a plan saved by these scripts; pick another path")
+        if saved == plan:
+            print(f"\nthis is the plan saved in {plan_file}: applying it")
+            return True
+        with open(plan_file, "w") as f:
+            json.dump(plan, f, indent=1)
+        what = "differs from the plan saved there (the command or the checkpoint changed); saved the new one" if saved else "saved"
+        print(f"\nplan {what} in {plan_file}. Nothing was changed: check the plan, then run the same command again to apply it.")
+        return False
+    if not sys.stdin.isatty():
+        print("\nstdin is not a terminal, so nothing was changed. Pass --plan <file>: the first run saves the plan, a second identical run applies it.")
+        return False
+    return input("\napply this plan? [y/N] ").strip().lower() in ("y", "yes")
+
+
+def plan_lines_for_save(ck: "Checkpoint", out: str | None) -> list[str]:
+    todo = ck.pending(out)
+    lines = [f"  write: {', '.join(todo['json']) or 'no JSON file'}"]
+    lines += [f"  shard {s}: {how}" for s, how in todo["shards"].items()]
+    if out is None:
+        lines.append(f"  in place: the originals of these files go to {os.path.join(ck.path, BACKUP_DIR)}/ first (--restore puts them back)")
+    else:
+        lines.append(f"  into {os.path.realpath(out)}: every other file is a symlink to the original")
+    return lines
+
+
+def refuse_stub(path: str) -> None:
+    """The fix scripts edit real checkpoints only: a fetch_headers.py stub of a remote repo is for diagnosis."""
+    if os.path.isfile(os.path.join(path, STUB_MARKER)):
+        raise SystemExit(f"{path} is a header-only stub of an HF repo, which is only diagnosed: download the checkpoint, run diagnose.py on the local directory, and fix that")
+
+
+def check_out_dir(ck: "Checkpoint", out: str | None) -> None:
+    if out is None:
+        return
+    target, root = os.path.realpath(out), os.path.realpath(ck.path)
+    if target == root or target.startswith(root + os.sep):
+        raise SystemExit("--out must be outside the checkpoint directory")
+    if os.path.isdir(target) and os.listdir(target):
+        raise SystemExit(f"--out {target} exists and is not empty; pick a new directory")
+
+
+def restore(path: str) -> list[str]:
+    """Put back the latest in-place backup of ``path``; returns what was restored."""
+    root = os.path.join(path, BACKUP_DIR)
+    stamps = sorted(os.listdir(root)) if os.path.isdir(root) else []
+    if not stamps:
+        raise FileNotFoundError(f"no backup under {root}")
+    backup, done = os.path.join(root, stamps[-1]), []
+    manifest_path = os.path.join(backup, "manifest.json")
+    manifest = {"links": {}, "created": []}
+    if os.path.isfile(manifest_path):
+        with open(manifest_path) as f:
+            manifest = json.load(f)
+    for name in sorted(os.listdir(backup)):
+        if name == "manifest.json":
+            continue
+        target = name.removesuffix(".header")
+        p = os.path.join(path, target)
+        if target in manifest["links"]:
+            # the edit made a private copy of a shared file: point the name back at the original
+            tmp = p + ".ckpt_edit.link"
+            os.symlink(manifest["links"][target], tmp)
+            os.replace(tmp, p)
+        elif name.endswith(".safetensors.header"):
+            if os.path.islink(p):
+                raise SystemExit(f"{p} is a symlink now; the backup was taken of a regular file, restore it by hand")
+            with open(os.path.join(backup, name), "rb") as f:
+                write_shard(p, p, f.read())
+        else:
+            shutil.copyfile(os.path.join(backup, name), p + ".ckpt_edit.part")
+            os.replace(p + ".ckpt_edit.part", p)
+        done.append(target)
+    for name in manifest["created"]:
+        if os.path.exists(os.path.join(path, name)):
+            os.remove(os.path.join(path, name))
+            done.append(f"{name} (removed, the edit created it)")
+    shutil.rmtree(backup)
+    return done
diff --git a/scripts/checkpoint_diagnose/diagnose.py b/scripts/checkpoint_diagnose/diagnose.py
new file mode 100644
index 0000000000..0b1b1e4465
--- /dev/null
+++ b/scripts/checkpoint_diagnose/diagnose.py
@@ -0,0 +1,1040 @@
+"""Diagnose why FreeToken does not load an HF safetensors checkpoint, without reading its weights.
+
+    python scripts/checkpoint_diagnose/diagnose.py <ckpt dir or HF repo[@revision]> ... [--out report.jsonl] [--text-model-only]
+
+A repo id is fetched as a header-only stub (``fetch_headers.py``) into ~/.cache/freetoken/checkpoint-stubs first. GGUF and FTW checkpoints are pointed elsewhere.
+Per checkpoint, through the code the server runs, on the CPU and the meta device:
+
+- config: the architecture is registered and the quantization_config parses
+- build: the model builds with a quant method for every layer
+- dense load: the family's weight reader yields exactly the model's state dict (names and shapes)
+- experts: every MoE layer's expert banks fill, packed by the layer's expert method
+
+The readers get meta tensors with the header's shapes and dtypes; a reader that needs values gets a second run on the real bytes (zeros in a stub).
+On a failure it also checks the tensor names against transformers' own model and its checkpoint conversion rules, and lints the files for the known non-standard forms.
+The failure is then put in one category (see README.md), with the fixes the fix scripts can apply; ``--out`` keeps them for ``--from-report``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures as cf
+import contextlib
+import functools
+import json
+import logging
+import multiprocessing as mp
+import os
+import re
+import struct
+import sys
+import time
+import traceback
+import warnings
+from collections import Counter
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+STUB_DIR = os.path.expanduser("~/.cache/freetoken/checkpoint-stubs")
+
+CATEGORIES = {
+    "arch": "architecture not served",
+    "quant_method": "quantization method not supported",
+    "quant_model": "quantization method supported, but not for this model or layer",
+    "loader": "model supported, the loader does not read this checkpoint",
+    "packing": "the checkpoint packs tensors in a form the loader does not read",
+    "checkpoint": "the checkpoint itself is non-standard",
+}
+FIX_SCRIPTS = {"config": "fix_config.py", "rename": "rename_tensors.py"}
+# tensors stored next to a quantized weight; they fold into the module's .weight when the names are compared with transformers
+QUANT_COMPANIONS = (
+    "weight_scale", "weight_scale_2", "weight_scale_inv", "weight_global_scale", "input_scale", "input_global_scale",
+    "weight_zero_point", "k_scale", "v_scale", "q_scale", "prob_scale", "output_scale", "weight_shape",
+)
+_QUANT_DTYPES = ("U8", "I8", "F8_E4M3", "F8_E5M2", "I32", "F4")
+_DIGITS = re.compile(r"\.\d+(?=\.|$)")
+# known non-standard prefixes: (pattern, replacement, what is wrong, checkpoint root the family expects)
+PREFIX_RULES = (
+    (r"^model\.language_model\.(?:language_model\.)+", "model.language_model.", "the text tower is nested more than once under language_model", "model.language_model"),
+    (r"^model\.language_model\.visual\.", "model.visual.", "the vision tower is stored under the text tower", "model.visual"),
+    (r"^language_model\.model\.", "model.language_model.", "the text tower uses the older language_model.model. layout", "model.language_model"),
+    (r"^language_model\.lm_head\.", "lm_head.", "lm_head uses the older language_model.lm_head layout", "model.language_model"),
+)
+# fp8 KV calibration scales under the projection's name; vLLM reads them as attn.k_scale, which FreeToken drops
+KV_SCALE_RULE = (r"\.([kv])_proj\.output_scale$", r".\1_scale")
+# the families whose weight reader drops .k_scale / .v_scale tensors; for the others the rename only moves the unexpected key
+KV_SCALE_DROPPERS = ("freetoken.models.qwen3_5_moe",)
+_PER_EXPERT = re.compile(r"^(?P<pre>.*\.layers\.)(?P<layer>\d+)(?P<mid>\..*\.experts\.)(?P<e>\d+)\.(?P<proj>[^.]+)\.")
+# compressed-tensors targets written with vLLM's internal module names
+TARGET_SWAPS = ((r"language_model\[\.\]model\[\.\]", "model[.]language_model[.]"), (r"language_model\\\.model\\\.", r"model\\.language_model\\."))
+# per-module tensors a quant method's reader produces next to the weight
+_SCALE_ROLES = (".weight_scale", ".weight_scale_inv", ".weight_scale_2", ".weight_global", ".input_scale")
+
+
+def pattern_of(name: str) -> str:
+    return _DIGITS.sub(".N", name)
+
+
+def group_names(names, limit: int = 4) -> str:
+    """Names grouped by their pattern (layer and expert numbers as N), most common first."""
+    c = Counter(pattern_of(n) for n in names)
+    parts = [f"{p} x{n}" if n > 1 else p for p, n in c.most_common(limit)]
+    return ", ".join(parts) + (f", ... ({len(c) - limit} more patterns)" if len(c) > limit else "")
+
+
+def shard_files(path: str) -> list[str]:
+    """The root shards the loaders read: consolidated.safetensors only when it is the only one (loader.iter_weight_files)."""
+    files = sorted(f for f in os.listdir(path) if f.endswith(".safetensors"))
+    return [f for f in files if f != "consolidated.safetensors"] or files
+
+
+def read_tensor_table(path: str) -> dict[str, dict]:
+    """name -> {dtype, shape, shard} from every shard header."""
+    table = {}
+    for f in shard_files(path):
+        for k, v in _header(os.path.join(path, f)).items():
+            if k != "__metadata__":
+                table[k] = {"dtype": v["dtype"], "shape": v["shape"], "shard": f}
+    return table
+
+
+_ITEM_BITS = {"F64": 64, "I64": 64, "U64": 64, "F32": 32, "I32": 32, "U32": 32, "F16": 16, "BF16": 16, "I16": 16, "U16": 16,
+              "I8": 8, "U8": 8, "BOOL": 8, "F8_E4M3": 8, "F8_E5M2": 8, "F8_E8M0": 8, "F4": 4}
+
+
+def check_files(path: str) -> list[str]:
+    """What safetensors and the loaders would refuse in the files themselves: data offsets that leave gaps or do not cover the file, a tensor whose bytes do not match its shape, an index that disagrees with the shard headers."""
+    problems, names = [], {}
+    for f in shard_files(path):
+        p = os.path.join(path, f)
+        with open(p, "rb") as fh:
+            (n,) = struct.unpack("<Q", fh.read(8))
+        header = _header(p)
+        spans = sorted((v["data_offsets"][0], v["data_offsets"][1], k) for k, v in header.items() if k != "__metadata__")
+        pos = 0
+        for start, end, k in spans:
+            if start != pos:
+                problems.append(f"{f}: tensor {k} starts at byte {start} of the data, expected {pos} (gap or overlap)")
+                break
+            bits = _ITEM_BITS.get(header[k]["dtype"])
+            numel = 1
+            for d in header[k]["shape"]:
+                numel *= d
+            if bits is not None and end - start != numel * bits // 8:
+                problems.append(f"{f}: tensor {k} has {end - start} bytes, its shape {header[k]['shape']} {header[k]['dtype']} needs {numel * bits // 8}")
+                break
+            pos = end
+        size = os.path.getsize(p) - 8 - n
+        if not problems and pos != size:
+            problems.append(f"{f}: the tensors cover {pos} data bytes but the file has {size} (an incomplete download?)")
+        names.update({k: f for k in header if k != "__metadata__"})
+    index = os.path.join(path, "model.safetensors.index.json")
+    if os.path.isfile(index):
+        with open(index) as fh:
+            wm = json.load(fh).get("weight_map", {})
+        read = set(shard_files(path))
+        # entries for shards the loaders skip (a subdir, an unread consolidated.safetensors) say nothing about the read ones
+        wm = {k: v for k, v in wm.items() if "/" not in v and (v in read or v != "consolidated.safetensors")}
+        absent = [k for k in wm if k not in names]
+        unlisted = [k for k in names if k not in wm]
+        moved = [k for k in wm if k in names and wm[k] != names[k]]
+        if absent:
+            problems.append(f"the index lists {len(absent)} tensors no shard holds: {group_names(absent)}")
+        if unlisted:
+            problems.append(f"the shards hold {len(unlisted)} tensors the index does not list: {group_names(unlisted)}")
+        if moved:
+            problems.append(f"the index puts {len(moved)} tensors in another shard than the one that holds them: {group_names(moved)}")
+    return problems
+
+
+def split_module(name: str) -> tuple[str, str]:
+    """(module, leaf) of a tensor name: model.layers.0.mlp.gate.weight_scale -> (model.layers.0.mlp.gate, weight_scale)."""
+    module, _, leaf = name.rpartition(".")
+    return module, leaf
+
+
+def stored_modules(table: dict) -> dict[str, tuple[bool, str]]:
+    """module -> (stored quantized, weight dtype) for every module with a 2-D weight or a packed weight."""
+    out = {}
+    by_module: dict[str, dict[str, dict]] = {}
+    for name, meta in table.items():
+        module, leaf = split_module(name)
+        by_module.setdefault(module, {})[leaf] = meta
+    for module, leaves in by_module.items():
+        w = leaves.get("weight_packed") or leaves.get("weight")
+        if w is None or (len(w["shape"]) != 2 and "weight_packed" not in leaves):
+            continue
+        quantized = "weight_packed" in leaves or w["dtype"] in _QUANT_DTYPES
+        out[module] = (quantized, w["dtype"])
+    return out
+
+
+# --- format routing (main process, no torch) ---
+
+
+def route(path: str) -> dict | None:
+    """A record for a checkpoint this script does not check (GGUF, FTW, no safetensors), else None."""
+    files = os.listdir(path)
+    if "freetoken_weight.json" in files:
+        return {"path": path, "status": "skip", "format": "ftw",
+                "hint": "an FTW checkpoint converted by `ft checkpoint`: diagnose its source HF checkpoint; docs/ftw-hotfix.md covers FTW repairs"}
+    stub = _stub_record(path)
+    gguf = any(f.endswith(".gguf") for f in files) or bool(stub and stub.get("gguf"))
+    shards = [f for f in files if f.endswith(".safetensors")]
+    if gguf and not shards:
+        return {"path": path, "status": "skip", "format": "gguf",
+                "hint": "a GGUF checkpoint: check it with scripts/gguf_check/ (README.md there); the GGUF section of this README sorts its failures"}
+    code = _missing_config_code(path) if stub else []
+    if code:
+        return {"path": path, "status": "skip", "format": "remote code",
+                "hint": f"the repo ships its own config code ({', '.join(code)}), which a header-only stub does not fetch and FreeToken runs with trust_remote_code; download the repo and diagnose the local directory"}
+    if not shards:
+        other = sorted({os.path.splitext(f)[1] for f in files if f.endswith((".bin", ".pt", ".pth", ".npz"))} | set((stub or {}).get("other_weights") or ()))
+        return {"path": path, "status": "skip", "format": "other",
+                "hint": f"no root *.safetensors shard ({', '.join(other) or 'no weights found'}); FreeToken reads safetensors and GGUF only"}
+    return None
+
+
+def _missing_config_code(path: str) -> list[str]:
+    """The .py files the config's auto_map names for AutoConfig that are not in ``path``."""
+    try:
+        with open(os.path.join(path, "config.json")) as f:
+            cfg = json.load(f)
+    except (OSError, ValueError):
+        return []
+    files = []
+    for c in (cfg, cfg.get("text_config") or {}):
+        ref = (c.get("auto_map") or {}).get("AutoConfig") if isinstance(c, dict) else None
+        if isinstance(ref, str) and "--" not in ref:
+            files.append(ref.rsplit(".", 1)[0] + ".py")
+    return [f for f in dict.fromkeys(files) if not os.path.isfile(os.path.join(path, f))]
+
+
+def _stub_record(path: str) -> dict | None:
+    p = os.path.join(path, ".stub.json")
+    if not os.path.isfile(p):
+        return None
+    with open(p) as f:
+        return json.load(f)
+
+
+# --- meta-device readers (worker) ---
+
+
+@functools.cache
+def _header(path: str) -> dict:
+    with open(path, "rb") as f:
+        (n,) = struct.unpack("<Q", f.read(8))
+        return json.loads(f.read(n))
+
+
+def _meta_tensor(meta: dict):
+    import torch
+    from freetoken.models.weight import _ST_DTYPE
+
+    return torch.empty(meta["shape"], dtype=_ST_DTYPE[meta["dtype"]], device="meta")
+
+
+class _MetaSlice:
+    def __init__(self, meta: dict):
+        self.meta = meta
+
+    def get_shape(self):
+        return list(self.meta["shape"])
+
+    def get_dtype(self):
+        return self.meta["dtype"]
+
+    def __getitem__(self, idx):
+        return _meta_tensor(self.meta)[idx]
+
+
+class _MetaSafeOpen:
+    """safetensors.safe_open whose tensors are meta tensors with the header's shape and dtype."""
+
+    def __init__(self, filename, framework="pt", device="cpu"):
+        self._h = _header(os.fspath(filename))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+    def keys(self):
+        return [k for k in self._h if k != "__metadata__"]
+
+    def offset_keys(self):
+        return self.keys()
+
+    def metadata(self):
+        return self._h.get("__metadata__")
+
+    def get_tensor(self, name):
+        return _meta_tensor(self._h[name])
+
+    def get_slice(self, name):
+        return _MetaSlice(self._h[name])
+
+
+def _meta_iter_expert_tensors_parallel(model_path, is_expert, **_kwargs):
+    from freetoken.models.loader import safetensors_weight_map
+    from freetoken.utils.hf import download_hf_weight
+
+    folder = download_hf_weight(model_path)
+    for name, shard in safetensors_weight_map(folder).items():
+        if is_expert(name):
+            yield name, _meta_tensor(_header(os.path.join(folder, shard))[name])
+
+
+@functools.cache
+def _meta_wins():
+    """A torch function mode that moves CPU operands to meta when an op mixes the two."""
+    import torch
+    from torch.overrides import TorchFunctionMode
+    from torch.utils._pytree import tree_flatten, tree_unflatten
+
+    class MetaWins(TorchFunctionMode):
+        def __torch_function__(self, func, types, args=(), kwargs=None):
+            flat, spec = tree_flatten((args, kwargs or {}))
+            if any(isinstance(x, torch.Tensor) and x.is_meta for x in flat):
+                flat = [x.to("meta") if isinstance(x, torch.Tensor) and not x.is_meta else x for x in flat]
+            args, kwargs = tree_unflatten(flat, spec)
+            return func(*args, **kwargs)
+
+    return MetaWins
+
+
+@contextlib.contextmanager
+def meta_reads(enabled: bool):
+    """Point safetensors.safe_open and every binding of iter_expert_tensors_parallel at the meta readers."""
+    if not enabled:
+        yield
+        return
+    import safetensors
+
+    import freetoken.models.weight as weight
+
+    real_open, real_parallel = safetensors.safe_open, weight.iter_expert_tensors_parallel
+    swapped = []
+    safetensors.safe_open = _MetaSafeOpen
+    for module in list(sys.modules.values()):
+        if getattr(module, "__name__", "").startswith("freetoken.") and getattr(module, "iter_expert_tensors_parallel", None) is real_parallel:
+            module.iter_expert_tensors_parallel = _meta_iter_expert_tensors_parallel
+            swapped.append(module)
+    try:
+        with _meta_wins()():
+            yield
+    finally:
+        safetensors.safe_open = real_open
+        for module in swapped:
+            module.iter_expert_tensors_parallel = real_parallel
+
+
+def _needs_values(exc: BaseException) -> bool:
+    return "meta" in f"{type(exc).__name__}: {exc}".lower()
+
+
+# --- load stages (worker) ---
+
+
+def _init_worker() -> None:
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    # the readers' progress bars; errors reach the main process in the record
+    sys.stderr = open(os.devnull, "w")
+    warnings.filterwarnings("ignore")
+    logging.disable(logging.WARNING)
+    import torch
+
+    from freetoken.distributed import set_tp_info
+    from freetoken.layers import rotary
+
+    set_tp_info(rank=0, size=1)
+    rotary.set_rope_device(torch.device("cpu"))
+
+
+def _engine_config(path: str, text_only: bool):
+    import torch
+
+    from freetoken.distributed import try_get_tp_info
+    from freetoken.engine.config import EngineConfig
+    from freetoken.mm.config import ENCODER_KINDS, MultimodalConfig
+
+    mm = MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS) if text_only else frozenset())
+    return EngineConfig(model_path=path, tp_info=try_get_tp_info(), dtype=torch.bfloat16, moe_strategy="offload", mm=mm)
+
+
+def _build_model(config):
+    import torch
+
+    from freetoken.engine.engine import _decode_target
+    from freetoken.layers.quantization import QuantBackend, set_quant_backend
+    from freetoken.models import create_model
+    from freetoken.utils.torch_utils import torch_dtype
+
+    set_quant_backend(QuantBackend.parse(config.quant_backend))
+    object.__setattr__(config.model_config, "moe_strategy", "offload")
+    object.__setattr__(config.model_config, "decode_target", _decode_target(config))
+    with torch.device("meta"), torch_dtype(torch.bfloat16):
+        return create_model(config.model_config)
+
+
+class _Recorder(dict):
+    """The state dict model.load_state_dict consumes: it records the keys a layer asks for and does not find, and the shapes that differ, instead of stopping at the first.
+
+    The layers' own load_state_dict decide what they take (a tied lm_head pops lm_head.weight itself); what is left at the end is unexpected."""
+
+    def __init__(self, expected: dict):
+        super().__init__()
+        self.expected, self.missing, self.shapes = expected, [], []
+
+    def pop(self, key, *default):
+        import torch
+
+        if key not in self:
+            if default:
+                return default[0]
+            self.missing.append(key)
+            return torch.empty_like(self.expected[key], device="meta")
+        t = super().pop(key)
+        ref = self.expected.get(key)
+        if ref is None:
+            return t
+        if tuple(t.shape) != tuple(ref.shape):
+            self.shapes.append(f"{key}: checkpoint {tuple(t.shape)}, model {tuple(ref.shape)}")
+            return torch.empty_like(ref, device="meta")
+        return t.to(dtype=ref.dtype)
+
+
+_LOSSY = {("float8_e4m3fn", "bfloat16"), ("float8_e5m2", "bfloat16"), ("bfloat16", "float8_e4m3fn"), ("float32", "float8_e4m3fn"), ("bfloat16", "uint8")}
+
+
+def dense_load(path: str, config, model, *, meta: bool) -> tuple[list[str], list[str], dict]:
+    """(problems, warnings, facts) of the engine's dense weight load: what load_weight yields against the model's state dict."""
+    import torch
+
+    from freetoken.models import load_weight
+
+    expected = model.state_dict()
+    is_moe = bool(getattr(config.model_config, "is_moe", False))
+    loaded, dups, casts = _Recorder(expected), [], Counter()
+    with meta_reads(meta):
+        for key, t in load_weight(path, torch.device("cpu"), include_moe_experts=not is_moe, include_vision=bool(config.active_encoders)):
+            ref = expected.get(key)
+            # as _materialize_loaded_weight_state_dict: drop an input_scale the layer declares none for, cast to the parameter dtype
+            if ref is None and key.endswith(".input_scale"):
+                continue
+            if key in loaded:
+                dups.append(key)
+            if ref is not None and (str(t.dtype).removeprefix("torch."), str(ref.dtype).removeprefix("torch.")) in _LOSSY:
+                casts[f"{t.dtype} -> {ref.dtype}"] += 1
+            # the check keeps no tensor bytes
+            dict.__setitem__(loaded, key, t if t.is_meta else torch.empty_like(t, device="meta"))
+    try:
+        model.load_state_dict(loaded)
+    except RuntimeError as e:
+        if not str(e).startswith("Unexpected keys"):
+            raise
+    unexpected, missing, shapes = sorted(loaded), sorted(loaded.missing), loaded.shapes
+    problems = []
+    if unexpected:
+        problems.append(f"{len(unexpected)} unexpected tensors: {group_names(unexpected)}")
+    if missing:
+        problems.append(f"{len(missing)} model tensors never loaded: {group_names(missing)}")
+    if shapes:
+        problems.append(f"{len(shapes)} shape mismatches, e.g. {shapes[0]}")
+    warns = [f"{n} tensors cast {k} at load (lossy)" for k, n in casts.items()]
+    if dups:
+        warns.append(f"{len(dups)} tensors are read twice and the engine keeps the last: {group_names(dups)}")
+    facts = {"missing_only_scales": bool(missing) and not unexpected and all(m.endswith(_SCALE_ROLES) for m in missing)}
+    return problems, warns, facts
+
+
+def expert_load(path: str, config, model, *, meta: bool) -> list[str]:
+    """Problems of filling every MoE layer's expert banks, as build_expert_banks fills them; nothing is kept."""
+    import torch
+
+    from freetoken.engine.engine import offload_layer_methods
+    from freetoken.moe.expert_banks import method_groups
+    from freetoken.models.weight import experts_scattered
+    from freetoken.moe.expert_pieces import iter_expert_pieces
+
+    methods = offload_layer_methods(model)
+    if methods is None:
+        return []
+    mc = config.model_config
+    num_layers, num_experts = mc.num_moe_layers, mc.num_experts
+    written = torch.zeros(num_layers, num_experts, dtype=torch.int32)
+    groups = method_groups(methods)
+    device = torch.device("meta" if meta else "cpu")
+    # the reader load_expert_banks picks: the parallel one (names from the index) for experts stored as many small tensors
+    parallel = experts_scattered(path)
+
+    def fill(method, layers, par: bool) -> str | None:
+        keep = set(layers)
+        layout = {role: spec for role, spec in method.layout().items() if not spec.resident}
+        pieces = iter_expert_pieces(path, mc, method.kind, layers=layers if len(groups) > 1 else None, parallel=par)
+        for layer, e0, e1, piece in pieces:
+            if not (0 <= layer < num_layers and 0 <= e0 < e1 <= num_experts):
+                return f"expert piece out of range: layer {layer}, experts {e0}:{e1} of {num_layers} x {num_experts}"
+            if layer not in keep:
+                continue
+            if written[layer, e0:e1].any():
+                return f"expert rows read more than once: layer {layer}, experts {e0}:{e1}"
+            written[layer, e0:e1] = 1
+            out = {role: torch.empty((e1 - e0, *spec.shape), dtype=spec.dtype, device=device) for role, spec in layout.items()}
+            method.pack(piece, out)
+        return None
+
+    with meta_reads(meta):
+        for method, layers in groups:
+            try:
+                problem = fill(method, layers, parallel)
+            except NotImplementedError:
+                # load_expert_banks rebuilds serially when a quant has no parallel reader
+                if not parallel:
+                    raise
+                written[layers] = 0
+                problem = fill(method, layers, False)
+            if problem:
+                return [problem]
+    missing = (written == 0).nonzero().tolist()
+    if missing:
+        return [f"expert banks were not filled: {len(missing)} (layer, expert) rows missing (first {missing[:4]})"]
+    return []
+
+
+# --- transformers cross-check (worker) ---
+
+
+def hf_check(path: str, table: dict) -> dict:
+    """Whether transformers' own model takes these tensor names, through its checkpoint conversion rules.
+
+    Quantization companions fold into their module's ``.weight``. Returns the names it does not map, the model keys nothing maps to, and the names a WeightConverter repacks."""
+    import torch
+    import transformers
+    from transformers.conversion_mapping import get_model_conversion_mapping
+    from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key
+
+    transformers.utils.logging.disable_progress_bar()
+    cfg = transformers.AutoConfig.from_pretrained(path)
+    arch = (getattr(cfg, "architectures", None) or [None])[0]
+    cls = getattr(transformers, arch or "", None)
+    if cls is None:
+        return {"skipped": f"transformers has no class {arch}"}
+    with torch.device("meta"):
+        model = cls._from_config(cfg)
+    meta_sd = model.state_dict()
+    transforms = get_model_conversion_mapping(model)
+    renamings = [t for t in transforms if isinstance(t, WeightRenaming)]
+    converters = [t for t in transforms if isinstance(t, WeightConverter)]
+    ignore_unexpected = [re.compile(p) for p in (getattr(model, "_keys_to_ignore_on_load_unexpected", None) or [])]
+    ignore_missing = [re.compile(p) for p in (getattr(model, "_keys_to_ignore_on_load_missing", None) or [])]
+    names = set()
+    for name in table:
+        module, leaf = split_module(name)
+        if leaf in QUANT_COMPANIONS:
+            continue
+        names.add(module + ".weight" if leaf == "weight_packed" else name)
+    mapped, converted = {}, set()
+    for name in names:
+        new, conv = rename_source_key(name, renamings, converters, model.base_model_prefix, meta_sd)
+        mapped[name] = new
+        if conv is not None:
+            converted.add(name)
+    unmapped = sorted(n for n, k in mapped.items() if k not in meta_sd and not any(rx.search(k) for rx in ignore_unexpected))
+    tied = bool(getattr(cfg.get_text_config(), "tie_word_embeddings", False))
+    targets = set(mapped.values())
+    missing = sorted(k for k in meta_sd if k not in targets and not any(rx.search(k) for rx in ignore_missing)
+                     and not (tied and k.endswith("lm_head.weight")))
+    renamed = sorted(n for n, k in mapped.items() if k != n and n not in converted and k in meta_sd)
+    return {"arch": arch, "unmapped": unmapped, "missing": missing, "converted": sorted(converted), "renamed": renamed}
+
+
+# --- lints (worker) ---
+
+
+def _skip_entry(dialect: str, modules: list[str], all_modules) -> list[str]:
+    """Skip-list entries in the dialect's syntax that cover exactly ``modules`` of ``all_modules``.
+
+    A pattern (layer numbers as a wildcard) stands for its modules only when every module of ``all_modules`` with that pattern is in ``modules``; the others are listed by name."""
+    chosen = set(modules)
+    by_pattern: dict[str, list[str]] = {}
+    for m in all_modules:
+        by_pattern.setdefault(pattern_of(m), []).append(m)
+    entries = []
+    for p in sorted({pattern_of(m) for m in chosen}):
+        whole = all(m in chosen for m in by_pattern.get(p, ())) and "N" in p.split(".")
+        if whole and dialect == "modelopt":
+            entries.append(p.replace(".N", ".*"))
+        elif whole and dialect == "compressed-tensors":
+            entries.append("re:^" + re.escape(p).replace(r"\.N", r"\.\d+") + "$")
+        else:
+            entries += sorted(m for m in by_pattern.get(p, [p]) if m in chosen)
+    return entries
+
+
+def _entry_matches(dialect: str, entry: str, name: str) -> bool:
+    from freetoken.layers.quantization.names import ct_set, name_set, substr_set
+
+    if dialect == "compressed-tensors":
+        return ct_set((entry,), class_names=False)(name)
+    if dialect == "fp8":
+        return name_set((entry,))(name) or substr_set((entry,))(name)
+    return name_set((entry,))(name)
+
+
+def _raw_skip_entries(path: str) -> list[str]:
+    from ckpt_edit import SKIP_KEYS, Checkpoint
+
+    entries = []
+    for _, q in Checkpoint(path).quant_configs():
+        for key in SKIP_KEYS:
+            entries += [e for e in q.get(key) or () if isinstance(e, str)]
+    return list(dict.fromkeys(entries))
+
+
+def lint_quant(path: str, table: dict, spec, hf_config) -> list[dict]:
+    """Modules whose stored form contradicts what the quantization_config declares for them, with the skip-list edits that make the two agree."""
+    from freetoken.models.register import checkpoint_quant_config
+
+    quant = checkpoint_quant_config(path, hf_config, spec)
+    dialect = getattr(quant, "dialect", "none")
+    if dialect in ("none", "gguf"):
+        return []
+    declared_q_stored_f, declared_f_stored_q, exempt_dense = [], [], []
+    stored = stored_modules(table)
+    for module, (quantized, _) in stored.items():
+        # embedding tables (and PLE n-gram tables) are read dense or by their own method whatever the config says
+        if any("embed" in seg for seg in module.split(".")) or module.endswith("conv1d"):
+            continue
+        declared = quant.scheme_for_name(module) is not None
+        if quant.unquantized(module):
+            if declared and not quantized:
+                exempt_dense.append(module)
+        elif declared and not quantized:
+            declared_q_stored_f.append(module)
+        elif quantized and not declared:
+            declared_f_stored_q.append(module)
+    findings = []
+    if exempt_dense:
+        findings.append({
+            "id": "quant_unlisted_dense_router", "level": "config", "category": "loader",
+            "message": f"{len(exempt_dense)} modules FreeToken serves in bf16 anyway are stored bf16 and not listed as skipped: {group_names(exempt_dense)}; listing them works around a reader that wants the stored form to match the config",
+            "fix": {"op": "skip_add", "entries": _skip_entry(dialect, exempt_dense, stored), "auto": True},
+        })
+    targets = _lint_ct_targets(path, table, dialect)
+    if declared_q_stored_f:
+        findings.append({
+            "id": "quant_declared_but_dense", "level": "config", "category": "checkpoint",
+            "message": f"{len(declared_q_stored_f)} modules are stored unquantized but the quantization_config does not list them as skipped: {group_names(declared_q_stored_f)}",
+            "fix": {"op": "skip_add", "entries": _skip_entry(dialect, declared_q_stored_f, stored), "auto": True},
+        })
+    # targets that name no module explain the quantized modules nothing covers; recheck after fixing them
+    if declared_f_stored_q and not targets:
+        entries = _raw_skip_entries(path)
+        remove, add = [], []
+        for e in entries:
+            hits = [m for m in declared_f_stored_q if _entry_matches(dialect, e, m)]
+            if not hits:
+                continue
+            remove.append(e)
+            # an entry naming a container also covers its dense children: list those alone instead
+            dense = [m for m, (q, _) in stored.items() if not q and _entry_matches(dialect, e, m)]
+            add += _skip_entry(dialect, dense, stored) if dense else []
+        msg = f"{len(declared_f_stored_q)} modules are stored quantized but the quantization_config skips them or no target covers them: {group_names(declared_f_stored_q)}"
+        finding = {"id": "quant_skipped_but_quantized", "level": "config", "category": "checkpoint", "message": msg}
+        if remove:
+            finding["message"] += f"; skip entries that cover them: {remove[:6]}"
+            finding["fixes"] = [{"op": "skip_remove", "entries": remove, "auto": True}] + ([{"op": "skip_add", "entries": add, "auto": True}] if add else [])
+        findings.append(finding)
+    return findings + targets
+
+
+def _lint_ct_targets(path: str, table: dict, dialect: str) -> list[dict]:
+    if dialect != "compressed-tensors":
+        return []
+    from ckpt_edit import Checkpoint
+
+    modules = list(stored_modules(table))
+    dead: dict[tuple[str, str], list[str]] = {}
+    for _, q in Checkpoint(path).quant_configs():
+        for group in (q.get("config_groups") or {}).values():
+            for target in group.get("targets") or ():
+                if target == "Linear" or any(_entry_matches(dialect, target, m) for m in modules):
+                    continue
+                for pat, repl in TARGET_SWAPS:
+                    swapped = re.sub(pat, repl, target)
+                    if swapped != target and any(_entry_matches(dialect, swapped, m) for m in modules):
+                        dead.setdefault((pat, repl), []).append(target)
+                        break
+    return [{
+        "id": "ct_target_vllm_names", "level": "config", "category": "checkpoint",
+        "message": f"{len(ts)} compressed-tensors targets name modules as vLLM does (language_model.model.) and match no tensor, e.g. {ts[0]!r}",
+        "fix": {"op": "quant_sub", "pattern": pat, "repl": repl, "auto": True},
+    } for (pat, repl), ts in dead.items()]
+
+
+def rename_rules(findings: list[dict]) -> list[tuple[str, str]]:
+    return [tuple(r) for f in findings if f["level"] == "rename" for r in f["fix"]["rules"]]
+
+
+def renamed_table(table: dict, rules) -> dict:
+    """The tensor table as it reads after the rename rules."""
+    out = {}
+    for name, meta in table.items():
+        for pat, repl in rules:
+            name = re.sub(pat, repl, name)
+        out[name] = meta
+    return out
+
+
+def lint_kv_scales(table: dict, spec) -> list[dict]:
+    hits = [n for n in table if re.search(KV_SCALE_RULE[0], n)]
+    if not hits or spec.module not in KV_SCALE_DROPPERS:
+        return []
+    return [{
+        "id": "kv_output_scale", "level": "rename", "category": "loader",
+        "message": f"{len(hits)} fp8 KV-cache scales stored as <k|v>_proj.output_scale ({group_names(hits, 2)}); vLLM reads them as k_scale / v_scale, FreeToken drops k_scale / v_scale but not output_scale",
+        "fix": {"op": "rename", "rules": [list(KV_SCALE_RULE)], "auto": True},
+    }]
+
+
+def lint_prefixes(table: dict, spec) -> list[dict]:
+    roots = {ckpt for _, ckpt in spec.checkpoint_roots}
+    names = list(table)
+    findings = []
+    for pat, repl, why, root in PREFIX_RULES:
+        if root not in roots:
+            continue
+        rx = re.compile(pat)
+        hits = [n for n in names if rx.search(n)]
+        if not hits:
+            continue
+        renamed = {rx.sub(repl, n) for n in hits}
+        if renamed & (set(names) - set(hits)):
+            continue
+        findings.append({
+            "id": "prefix", "level": "rename", "category": "checkpoint",
+            "message": f"{why}: {len(hits)} tensors {rx.search(hits[0]).group(0)} -> {repl}",
+            "fix": {"op": "rename", "rules": [[pat, repl]], "quant_too": True, "auto": True},
+        })
+    return findings
+
+
+def lint_experts(table: dict, num_experts: int | None) -> list[dict]:
+    """Per-expert tensors (experts.E.<proj>) where some layer lacks some expert's projection: a checkpoint that lost tensors, not a packing form."""
+    if not num_experts:
+        return []
+    seen: dict[tuple[str, str, str], set[int]] = {}
+    for name in table:
+        m = _PER_EXPERT.match(name)
+        if m:
+            seen.setdefault((m["pre"] + m["layer"], m["mid"], m["proj"]), set()).add(int(m["e"]))
+    gaps = [(f"{pre}{mid}N.{proj}", sorted(set(range(num_experts)) - es)) for (pre, mid, proj), es in seen.items() if len(es) < num_experts]
+    if not gaps:
+        return []
+    shown = "; ".join(f"{where} lacks experts {missing[:6]}" for where, missing in gaps[:3])
+    return [{"id": "expert_tensors_incomplete", "level": "none", "category": "checkpoint",
+             "message": f"{len(gaps)} (layer, projection) groups of per-expert tensors are incomplete: {shown}; the checkpoint is missing tensors (download it again, or ask its uploader)"}]
+
+
+def lint_vision(path: str, table: dict, spec, text_only: bool) -> list[dict]:
+    from freetoken.models.config import VISION_KEY_PREFIXES
+
+    if text_only or not spec.encoders:
+        return []
+    # checkpoint names put the tower under a root (model.visual., model.vision_tower.)
+    segments = tuple("." + p for p in VISION_KEY_PREFIXES)
+    if not any(("." + n).count(s) for n in table for s in segments):
+        return [{"id": "no_vision_tower", "level": "runtime", "category": "checkpoint",
+                 "message": "the architecture has a vision tower but the checkpoint stores none: a text-only checkpoint under a multimodal architecture; serve it with --text-model-only (and diagnose it with --text-model-only)"}]
+    if spec.mm_processor and not any(os.path.isfile(os.path.join(path, f)) for f in ("preprocessor_config.json", "processor_config.json")):
+        return [{"id": "no_processor_config", "level": "file", "category": "checkpoint",
+                 "message": "the checkpoint has a vision tower but no preprocessor_config.json, so image requests fail with a 400; copy that file from the base model the checkpoint was tuned from"}]
+    return []
+
+
+def lint_misc(table: dict, hf_config) -> list[dict]:
+    from freetoken.layers.quantization.configs import quantization_config_of
+
+    out = []
+    dtypes = Counter(m["dtype"] for m in table.values())
+    if quantization_config_of(hf_config) is None and (dtypes["F8_E4M3"] or dtypes["F8_E5M2"]):
+        out.append({"id": "fp8_without_config", "level": "none", "category": "checkpoint",
+                    "message": f"{dtypes['F8_E4M3'] + dtypes['F8_E5M2']} fp8 tensors but no quantization_config says how to read them (scales, block size)"})
+    if any(n.endswith((".scales", ".biases")) for n in table) and any(m["dtype"] == "U32" for m in table.values()):
+        out.append({"id": "mlx", "level": "none", "category": "quant_method", "message": "MLX-quantized weights (.scales / .biases, U32 packing); FreeToken does not read MLX"})
+    return out
+
+
+def expert_layout(table: dict) -> str | None:
+    per_expert = any(re.search(r"\.experts\.\d+\.(gate|up|down|w[123])", n) for n in table)
+    stacked = any(re.search(r"\.experts\.(gate_up_proj|down_proj|w13_weight|w2_weight)$", n) for n in table)
+    if per_expert and stacked:
+        return "mixed: some layers per expert, some stacked"
+    return "per expert (experts.N.<proj>)" if per_expert else "stacked (experts.gate_up_proj / down_proj)" if stacked else None
+
+
+# --- classification ---
+
+_QUANT_METHOD = re.compile(r"quantization method .* is not supported|quant_algo .* is not supported|compressed-tensors .*scheme .* is not supported|MIXED_PRECISION without|fp8 checkpoint with weight_block_size|fp8 block size .* is not supported|target class .* is not supported")
+_QUANT_MODEL = re.compile(r"no quant method for|provides no (expert reader|nvfp4_expert_spec|iter_weights_parallel)|needs in/out sizes divisible|is not supported for|cannot use the .* kernel")
+
+
+def classify(rec: dict) -> str:
+    """The category of a failed record (README.md, "Safetensors: six kinds of failure"), first match wins."""
+    first = (rec["problems"] or [""])[0]
+    findings = rec["findings"]
+    hf = rec.get("transformers") or {}
+    if rec["stage"] == "files":
+        return "checkpoint"
+    if rec["stage"] == "config" and "architecture" in first.lower() and "not supported" in first:
+        return "arch"
+    if any(f["category"] == "quant_method" for f in findings) or _QUANT_METHOD.search(first):
+        return "quant_method"
+    # a broken prefix misnames every module under it, so whatever fails next is its symptom
+    if any(f["id"] == "prefix" and f["category"] == "checkpoint" for f in findings):
+        return "checkpoint"
+    if _QUANT_MODEL.search(first) or rec.get("missing_only_scales"):
+        return "quant_model"
+    if hf.get("unmapped") or hf.get("missing"):
+        return "checkpoint"
+    if any(f["category"] == "checkpoint" for f in findings):
+        return "checkpoint"
+    if rec["stage"] == "experts" and hf.get("converted"):
+        return "packing"
+    return "loader"
+
+
+def _fix_list(findings: list[dict]) -> list[dict]:
+    fixes = []
+    for f in findings:
+        for fix in f.get("fixes") or ([f["fix"]] if "fix" in f else []):
+            fixes.append({"level": f["level"], "finding": f["id"], **fix})
+    return fixes
+
+
+def run_lints(path: str, table: dict, spec, hf_config, text_only: bool, warns: list[str], num_experts: int | None = None) -> list[dict]:
+    """Every lint; the quantization lints see the names as they read after the rename fixes, which a rename applies first."""
+    def safe(fn, *args):
+        try:
+            return fn(*args)
+        except Exception as e:
+            warns.append(f"lint {fn.__name__} failed: {type(e).__name__}: {str(e)[:200]}")
+            return []
+
+    if spec is None:
+        return safe(lint_misc, table, hf_config) if hf_config is not None else []
+    findings = safe(lint_prefixes, table, spec) + safe(lint_kv_scales, table, spec)
+    after = renamed_table(table, rename_rules(findings))
+    findings += safe(lint_experts, after, num_experts) + safe(lint_vision, path, after, spec, text_only) + safe(lint_misc, after, hf_config)
+    return findings + safe(lint_quant, path, after, spec, hf_config)
+
+
+class _Stop(Exception):
+    """A stage found problems of its own; the later stages would only repeat them."""
+
+
+def check(path: str, text_only: bool) -> dict:
+    rec = {"path": path, "stage": "config", "status": "fail", "problems": [], "warnings": [], "findings": []}
+    stub = _stub_record(path)
+    if stub:
+        rec["repo"], rec["revision"] = stub["repo"], stub["revision"]
+    t0 = time.time()
+    hf_config = spec = table = None
+    try:
+        from freetoken.models.register import get_model_spec
+        from freetoken.utils import cached_load_hf_config
+
+        rec["stage"] = "files"
+        table = read_tensor_table(path)
+        file_problems = check_files(path)
+        if file_problems:
+            rec["problems"] += file_problems
+            raise _Stop
+        rec["stage"] = "config"
+        hf_config = cached_load_hf_config(path)
+        rec["arch"] = (getattr(hf_config, "architectures", None) or [None])[0]
+        spec = get_model_spec(rec["arch"])
+        rec["stage"] = "quant"
+        config = _engine_config(path, text_only)
+        mc = config.model_config
+        rec["dialect"] = getattr(getattr(mc, "quant", None), "dialect", None)
+        rec["moe"] = {"layers": mc.num_moe_layers, "experts": mc.num_experts, "layout": expert_layout(table)} if getattr(mc, "is_moe", False) else None
+        rec["encoders"] = [e.kind for e in config.active_encoders]
+        rec["stage"] = "build"
+        model = _build_model(config)
+        for stage, fn in (("dense", dense_load), ("experts", expert_load)):
+            rec["stage"] = stage
+            try:
+                out = fn(path, config, model, meta=True)
+            except Exception as e:
+                if not _needs_values(e):
+                    raise
+                rec.setdefault("real_reads", []).append(stage)
+                out = fn(path, config, model, meta=False)
+            problems, warns, facts = out if stage == "dense" else (out, [], {})
+            rec["warnings"] += warns
+            rec.update(facts)
+            if problems:
+                rec["problems"] += problems
+                break
+        else:
+            rec["stage"], rec["status"] = "done", "pass"
+    except _Stop:
+        pass
+    except Exception as e:
+        rec["problems"].append(f"{type(e).__name__}: {str(e)[:600]}")
+        rec["traceback"] = traceback.format_exc()[-2500:]
+    if table is not None:
+        rec["findings"] = run_lints(path, table, spec, hf_config, text_only, rec["warnings"], (rec.get("moe") or {}).get("experts"))
+    if rec["status"] == "fail" and table is not None and rec["stage"] in ("build", "dense", "experts"):
+        try:
+            rec["transformers"] = hf_check(path, table)
+        except Exception as e:
+            rec["transformers"] = {"skipped": f"{type(e).__name__}: {str(e)[:200]}"}
+        hf = rec["transformers"]
+        if "skipped" not in hf and not hf["unmapped"] and not hf["missing"]:
+            # transformers reads these names, so a rename is a workaround for the loader rather than a repair
+            for f in rec["findings"]:
+                if f["id"] == "prefix":
+                    f["category"] = "loader"
+    if rec["status"] == "fail":
+        rec["category"] = classify(rec)
+    # nothing to repair in the files of a checkpoint whose format or architecture is not served
+    fixes = [] if rec.get("category") in ("arch", "quant_method") else _fix_list(rec["findings"])
+    # a remote repo is only diagnosed: fixes are made on a local download, diagnosed again
+    rec["fixes"] = [] if stub else fixes
+    rec["fixable_locally"] = bool(stub and fixes)
+    rec["seconds"] = round(time.time() - t0, 1)
+    return rec
+
+
+# --- main process ---
+
+
+def resolve(target: str, stub_dir: str) -> tuple[str | None, dict | None]:
+    """(local path, None), or (None, error record) when a repo cannot be fetched."""
+    if os.path.isdir(target):
+        return os.path.abspath(target), None
+    from fetch_headers import fetch_repo
+
+    try:
+        return fetch_repo(target, stub_dir)["stub"], None
+    except Exception as e:
+        return None, {"path": target, "status": "skip", "format": "unknown", "hint": f"not a local directory, and fetching it from HF failed: {type(e).__name__}: {str(e)[:300]}"}
+
+
+def _short(rec: dict) -> str:
+    return rec.get("repo") or rec["path"]
+
+
+def print_record(rec: dict, report_path: str | None) -> None:
+    print(f"\n== {_short(rec)}" + (" (header-only stub)" if rec.get("repo") else ""))
+    if rec["status"] == "skip":
+        print(f"   {rec['hint']}")
+        return
+    bits = [rec.get("arch") or "?", f"quant {rec.get('dialect') or 'none'}"]
+    if rec.get("moe"):
+        bits.append(f"MoE {rec['moe']['layers']} layers x {rec['moe']['experts']} experts, {rec['moe']['layout'] or 'experts not found by name'}")
+    print("   " + ", ".join(bits))
+    if rec["status"] == "pass":
+        print("   PASS: config, build, dense load" + (", expert banks" if rec.get("moe") else ""))
+    else:
+        cat = rec["category"]
+        print(f"   FAIL at {rec['stage']}: [{cat}] {CATEGORIES[cat]}")
+        for p in rec["problems"]:
+            print(f"     - {p}")
+    hf = rec.get("transformers") or {}
+    if hf.get("skipped"):
+        print(f"   transformers check skipped: {hf['skipped']}")
+    elif hf:
+        if hf["unmapped"] or hf["missing"]:
+            parts = []
+            if hf["unmapped"]:
+                parts.append(f"{len(hf['unmapped'])} checkpoint names it cannot map ({group_names(hf['unmapped'], 3)})")
+            if hf["missing"]:
+                parts.append(f"{len(hf['missing'])} model keys nothing fills ({group_names(hf['missing'], 3)})")
+            print(f"   transformers {hf['arch']} would not load it either: " + "; ".join(parts))
+        else:
+            extra = []
+            if hf["renamed"]:
+                extra.append(f"renames {len(hf['renamed'])} ({group_names(hf['renamed'], 2)})")
+            if hf["converted"]:
+                extra.append(f"repacks {len(hf['converted'])} ({group_names(hf['converted'], 2)})")
+            print("   transformers loads these names" + (": " + "; ".join(extra) if extra else ""))
+    for w in rec["warnings"]:
+        print(f"   warning: {w}")
+    for f in rec["findings"]:
+        print(f"   {'note' if rec['status'] == 'pass' else 'finding'} [{f['level']}, {f['category']}]: {f['message']}")
+    if rec.get("fixable_locally") and rec["status"] == "fail":
+        print("   fixable: download the checkpoint and run diagnose.py on the local directory to get the fix commands (a remote repo is only diagnosed)")
+    fixes = rec.get("fixes") or []
+    if not fixes or rec["status"] != "fail":
+        return
+    target, report_arg = rec["path"], report_path or "<report.jsonl>"
+    # one level per round: a rename changes the names the config findings are about, so diagnose again after it
+    level = "rename" if any(f["level"] == "rename" and f.get("auto") for f in fixes) else "config"
+    if any(f["level"] == level and f.get("auto") for f in fixes):
+        if not report_path:
+            print("   fixes: rerun with --out <report.jsonl> to get the fix command")
+            return
+        print(f"   next step ({level} fix; it prints a plan and asks before writing):")
+        print(f"     python scripts/checkpoint_diagnose/{FIX_SCRIPTS[level]} {target} --out <new dir> --from-report {report_arg}")
+        print("     then diagnose <new dir> again")
+
+
+def summary(records: list[dict]) -> None:
+    status = Counter(r["status"] for r in records)
+    print(f"\n{status['pass']} pass, {status['fail']} fail, {status['skip']} not checked")
+    cats = Counter(r["category"] for r in records if r["status"] == "fail")
+    for cat in CATEGORIES:
+        if cats[cat]:
+            print(f"  {cats[cat]:4d}  [{cat}] {CATEGORIES[cat]}")
+            for r in records:
+                if r["status"] == "fail" and r["category"] == cat:
+                    fixable = [lv for lv in FIX_SCRIPTS if any(f["level"] == lv and f.get("auto") for f in r.get("fixes", []))]
+                    note = f"  (fixes: {', '.join(fixable)})" if fixable else "  (fixable on a local download)" if r.get("fixable_locally") else ""
+                    print(f"          {_short(r)}{note}")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("targets", nargs="+", help="checkpoint directories or HF repo ids")
+    ap.add_argument("--out", metavar="JSONL", help="append the records here; the fix scripts read them with --from-report")
+    ap.add_argument("--text-model-only", action="store_true", help="check as `ft serve --text-model-only` loads: no vision tower")
+    args = ap.parse_args()
+
+    records = []
+    out = open(args.out, "a") if args.out else None
+    ctx = mp.get_context("spawn")
+    pool = None
+    for target in args.targets:
+        path, rec = resolve(target, STUB_DIR)
+        rec = rec or route(path)
+        if rec is None:
+            # one worker process, reused; a crash takes only the checkpoint it was checking
+            pool = pool or cf.ProcessPoolExecutor(1, mp_context=ctx, initializer=_init_worker)
+            try:
+                rec = pool.submit(check, path, args.text_model_only).result()
+            except cf.process.BrokenProcessPool:
+                pool = None
+                rec = {"path": path, "status": "fail", "stage": "?", "category": "loader", "problems": ["the check process died (out of host memory, or a crash)"],
+                       "warnings": [], "findings": [], "fixes": []}
+        records.append(rec)
+        print_record(rec, args.out)
+        if out:
+            out.write(json.dumps(rec) + "\n")
+            out.flush()
+    if pool:
+        pool.shutdown()
+    if out:
+        out.close()
+    if len(records) > 1:
+        summary(records)
+    return 0 if all(r["status"] != "fail" for r in records) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/checkpoint_diagnose/fetch_headers.py b/scripts/checkpoint_diagnose/fetch_headers.py
new file mode 100644
index 0000000000..da2ae71673
--- /dev/null
+++ b/scripts/checkpoint_diagnose/fetch_headers.py
@@ -0,0 +1,137 @@
+"""Write header-only stubs of HF safetensors checkpoints: the config files and every shard's header, in sparse files of the real size.
+
+    python scripts/checkpoint_diagnose/fetch_headers.py <out_dir> <repo>[@<revision>] ...
+
+The tensor data stays holes that read as zeros, so a stub takes a few MB of disk and ``diagnose.py`` reads it like the real checkpoint.
+Stubs land at ``<out_dir>/<repo>``; a rerun keeps the files that already have the right size.
+Python files are never fetched. Gated repos read ``HF_TOKEN``. Prints one JSON line per repo.
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures as cf
+import json
+import os
+import struct
+import sys
+import threading
+import time
+
+import requests
+from huggingface_hub import HfApi, hf_hub_url
+from huggingface_hub.utils import build_hf_headers
+
+STUB_MARKER = ".stub.json"
+CONFIG_EXT = (".json", ".jinja")
+# large tokenizer files: the load check never reads them
+SKIP = {"tokenizer.json", "vocab.json", "merges.txt"}
+_FIRST_READ = 1 << 20
+
+_local = threading.local()
+
+
+def _session() -> requests.Session:
+    if not hasattr(_local, "session"):
+        _local.session = requests.Session()
+        _local.session.headers.update(build_hf_headers())
+    return _local.session
+
+
+def _get(url: str, start: int | None = None, end: int | None = None) -> bytes:
+    """The file, or bytes [start, end] of it; retries rate limits and server errors."""
+    headers = {} if start is None else {"Range": f"bytes={start}-{end}"}
+    for attempt in range(6):
+        r = _session().get(url, headers=headers, timeout=300)
+        if r.status_code == 429 or r.status_code >= 500:
+            time.sleep(5 * (attempt + 1))
+            continue
+        r.raise_for_status()
+        if start is not None and r.status_code != 206:
+            raise ValueError(f"{url}: asked for a byte range, got status {r.status_code}")
+        return r.content
+    r.raise_for_status()
+    raise RuntimeError(f"{url}: status {r.status_code}")
+
+
+def _fetch_shard(url: str, dst: str, size: int) -> None:
+    if os.path.isfile(dst) and os.path.getsize(dst) == size:
+        return
+    head = _get(url, 0, min(size, _FIRST_READ) - 1)
+    (n,) = struct.unpack("<Q", head[:8])
+    if 8 + n > len(head):
+        head += _get(url, len(head), 8 + n - 1)
+    head = head[: 8 + n]
+    json.loads(head[8:])
+    tmp = dst + ".part"
+    with open(tmp, "wb") as f:
+        f.write(head)
+        f.truncate(size)
+    os.replace(tmp, dst)
+
+
+def _fetch_file(url: str, dst: str, size: int | None) -> None:
+    if os.path.isfile(dst) and size is not None and os.path.getsize(dst) == size:
+        return
+    data = _get(url)
+    with open(dst + ".part", "wb") as f:
+        f.write(data)
+    os.replace(dst + ".part", dst)
+
+
+def fetch_repo(repo: str, out_dir: str, *, workers: int = 8) -> dict:
+    """Write the stub of ``repo`` (``name`` or ``name@revision``) under ``out_dir``; returns its record, with ``stub`` the directory."""
+    name, _, revision = repo.partition("@")
+    info = HfApi().model_info(name, revision=revision or None, files_metadata=True)
+    files = [(s.rfilename, s.size) for s in info.siblings or []]
+    root = [(f, size) for f, size in files if "/" not in f]
+    shards = [(f, size) for f, size in root if f.endswith(".safetensors")]
+    configs = [(f, size) for f, size in root if f.endswith(CONFIG_EXT) and f not in SKIP]
+    dst = os.path.join(out_dir, name)
+    rec = {"repo": name, "revision": info.sha, "stub": dst, "shards": len(shards),
+           "gguf": sum(f.endswith(".gguf") for f, _ in files),
+           "other_weights": sorted({os.path.splitext(f)[1] for f, _ in files if f.endswith((".bin", ".pt", ".pth", ".npz"))}),
+           "files": [f for f, _ in configs + shards]}
+    marker = os.path.join(dst, STUB_MARKER)
+    if os.path.isfile(marker):
+        with open(marker) as f:
+            old = json.load(f)
+        if old.get("revision") != info.sha:
+            # another revision: none of its files can be kept
+            for f in old.get("files", []):
+                if os.path.isfile(os.path.join(dst, f)):
+                    os.remove(os.path.join(dst, f))
+            os.remove(marker)
+    os.makedirs(dst, exist_ok=True)
+    jobs = [(_fetch_file, f, size) for f, size in configs] + [(_fetch_shard, f, size) for f, size in shards]
+    with cf.ThreadPoolExecutor(workers) as ex:
+        futures = [ex.submit(fn, hf_hub_url(name, f, revision=info.sha), os.path.join(dst, f), size) for fn, f, size in jobs]
+        for fut in futures:
+            fut.result()
+    with open(os.path.join(dst, STUB_MARKER), "w") as f:
+        json.dump(rec, f, indent=1)
+    return rec
+
+
+def is_stub(path: str) -> bool:
+    return os.path.isfile(os.path.join(path, STUB_MARKER))
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("out_dir")
+    ap.add_argument("repos", nargs="+", help="HF repo ids, optionally with @revision")
+    args = ap.parse_args()
+    failed = 0
+    for repo in args.repos:
+        try:
+            rec = {**fetch_repo(repo, args.out_dir), "ok": True}
+        except Exception as e:
+            rec = {"repo": repo, "ok": False, "error": f"{type(e).__name__}: {e}"[:300]}
+            failed += 1
+        print(json.dumps(rec), flush=True)
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/checkpoint_diagnose/fix_config.py b/scripts/checkpoint_diagnose/fix_config.py
new file mode 100644
index 0000000000..93bda8af02
--- /dev/null
+++ b/scripts/checkpoint_diagnose/fix_config.py
@@ -0,0 +1,102 @@
+"""Level 1 fix: edit only config.json / hf_quant_config.json of a local HF safetensors checkpoint (a remote repo's stub is refused).
+
+    python scripts/checkpoint_diagnose/fix_config.py <ckpt> --out <new dir> --from-report report.jsonl
+    python scripts/checkpoint_diagnose/fix_config.py <ckpt> --out <new dir> --set text_config.tie_word_embeddings false
+    python scripts/checkpoint_diagnose/fix_config.py <ckpt> --in-place --skip-add 'model.layers.*.mlp.gate' --plan plan.json
+    python scripts/checkpoint_diagnose/fix_config.py <ckpt> --restore
+
+Edits, applied in this order:
+    --from-report JSONL     the automatic config fixes diagnose.py --out found for this checkpoint
+    --set PATH VALUE        a config.json key; PATH is dotted, VALUE is JSON (false, 4, "text")
+    --skip-add ENTRY        add a module to the quantization skip lists (ignore / exclude_modules / modules_to_not_convert)
+    --skip-remove ENTRY     remove an entry from those lists
+
+Write a skip entry in the checkpoint dialect's syntax: a glob for ModelOpt (model.layers.*.mlp.gate), a module name or ``re:<regex>`` for compressed-tensors, a module name for fp8.
+Every run prints the plan first and asks [y/N]; with ``--plan FILE`` the first run only saves the plan and the second identical run applies it.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from ckpt_edit import Checkpoint, check_out_dir, confirm, plan_lines_for_save, refuse_stub, restore  # noqa: E402
+
+
+def report_fixes(report: str, ckpt: str, level: str) -> list[dict]:
+    """The automatic ``level`` fixes in the latest diagnose.py record of ``ckpt`` (matched by path or repo id)."""
+    with open(report) as f:
+        recs = [json.loads(line) for line in f if line.strip()]
+    mine = [r for r in recs if os.path.realpath(r.get("path", "")) == os.path.realpath(ckpt) or r.get("repo") == ckpt]
+    if not mine:
+        raise SystemExit(f"{report} has no record of {ckpt}: run diagnose.py {ckpt} --out {report} first")
+    if mine[-1].get("status") != "fail":
+        raise SystemExit(f"{ckpt} passed in {report}: its notes are not fixes to apply")
+    return [f for f in mine[-1].get("fixes", []) if f["level"] == level and f.get("auto")]
+
+
+def apply_fix(ck: Checkpoint, fix: dict) -> None:
+    op = fix["op"]
+    if op == "set":
+        ck.set_config(fix["path"], fix["value"])
+    elif op == "skip_add":
+        ck.quant_skip_add(fix["entries"])
+    elif op == "skip_remove":
+        ck.quant_skip_remove(fix["entries"])
+    elif op == "quant_sub":
+        ck.quant_sub(fix["pattern"], fix["repl"])
+    else:
+        raise ValueError(f"unknown config op {op!r}")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("ckpt")
+    where = ap.add_mutually_exclusive_group(required=True)
+    where.add_argument("--out", help="write the fixed checkpoint into this new directory")
+    where.add_argument("--in-place", action="store_true", help="edit the checkpoint itself, keeping a backup")
+    where.add_argument("--restore", action="store_true", help="undo the latest --in-place edit")
+    ap.add_argument("--from-report", metavar="JSONL")
+    ap.add_argument("--set", nargs=2, action="append", default=[], metavar=("PATH", "VALUE"))
+    ap.add_argument("--skip-add", action="append", default=[], metavar="ENTRY")
+    ap.add_argument("--skip-remove", action="append", default=[], metavar="ENTRY")
+    ap.add_argument("--plan", metavar="FILE", help="no prompt: the first run saves the plan here, the second identical run applies it")
+    args = ap.parse_args()
+
+    refuse_stub(args.ckpt)
+    if args.restore:
+        for name in restore(args.ckpt):
+            print(f"restored {name}")
+        return 0
+    fixes = report_fixes(args.from_report, args.ckpt, "config") if args.from_report else []
+    if args.from_report and not fixes and not (args.set or args.skip_add or args.skip_remove):
+        raise SystemExit(f"the record of {args.ckpt} in {args.from_report} has no automatic config fix")
+    fixes += [{"op": "set", "path": p, "value": json.loads(v)} for p, v in args.set]
+    fixes += [{"op": "skip_add", "entries": args.skip_add}] if args.skip_add else []
+    fixes += [{"op": "skip_remove", "entries": args.skip_remove}] if args.skip_remove else []
+    if not fixes:
+        ap.error("give --from-report, --set, --skip-add or --skip-remove")
+    ck = Checkpoint(args.ckpt)
+    check_out_dir(ck, args.out)
+    for fix in fixes:
+        apply_fix(ck, fix)
+    if not ck.log:
+        print("nothing to change: the config already reads this way")
+        return 0
+    lines = [f"plan: edit the config of {ck.path}"] + [f"  {line}" for line in ck.log] + plan_lines_for_save(ck, args.out)
+    plan = {"tool": "fix_config", "ckpt": ck.path, "fingerprint": ck.fingerprint, "out": args.out and os.path.abspath(args.out), "log": ck.log}
+    if not confirm(plan, lines, args.plan):
+        return 0
+    for name in ck.save(out=None if args.in_place else args.out):
+        print(f"wrote {name}")
+    if args.plan:
+        os.remove(args.plan)
+    print("done: run diagnose.py on the result to check it")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/checkpoint_diagnose/rename_tensors.py b/scripts/checkpoint_diagnose/rename_tensors.py
new file mode 100644
index 0000000000..74cc040d28
--- /dev/null
+++ b/scripts/checkpoint_diagnose/rename_tensors.py
@@ -0,0 +1,123 @@
+"""Level 2 fix: rename tensors of a local HF safetensors checkpoint (a remote repo's stub is refused). Only the shard headers, the index and the quantization skip lists change; tensor data is kept byte for byte.
+
+    python scripts/checkpoint_diagnose/rename_tensors.py <ckpt> --out <new dir> --rule PATTERN REPL [--rule ...]
+    python scripts/checkpoint_diagnose/rename_tensors.py <ckpt> --out <new dir> --from-report report.jsonl
+    python scripts/checkpoint_diagnose/rename_tensors.py <ckpt> --in-place ... --plan plan.json
+    python scripts/checkpoint_diagnose/rename_tensors.py <ckpt> --restore
+
+A rule is ``re.sub(PATTERN, REPL, name)``; rules apply in order. They are also applied to the quantization skip-list entries and compressed-tensors targets that are module names or globs; ``re:`` entries are not rewritten.
+Every run prints the plan first: per rule how many tensors it renames with examples, the skip entries that still name the old modules, and what each shard write costs. Then it asks [y/N].
+With ``--plan FILE`` it never asks: the first run saves the plan to FILE and changes nothing, the second identical run applies it.
+Refused: a rule that renames nothing, a rename that gives two tensors one name, an ``--out`` that is not empty.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+import sys
+from collections import Counter
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from ckpt_edit import Checkpoint, check_out_dir, confirm, plan_lines_for_save, refuse_stub, restore  # noqa: E402
+from fix_config import report_fixes  # noqa: E402
+
+_DIGITS = re.compile(r"\.\d+(?=\.|$)")
+
+
+def _pattern(name: str) -> str:
+    return _DIGITS.sub(".N", name)
+
+
+def rule_effects(names: list[str], rules: list[tuple[str, str]]) -> list[dict]:
+    """Per rule, applied in order: the names it changes as pattern pairs, and how many of its matches are not at the start of the name."""
+    effects, current = [], list(names)
+    for pat, repl in rules:
+        rx = re.compile(pat)
+        pairs, mid = Counter(), 0
+        for i, name in enumerate(current):
+            new = rx.sub(repl, name)
+            if new != name:
+                pairs[(_pattern(name), _pattern(new))] += 1
+                m = rx.search(name)
+                mid += m is not None and m.start() > 0
+                current[i] = new
+        anchored = pat.startswith("^") or pat.endswith("$")
+        effects.append({"rule": [pat, repl], "renamed": sum(pairs.values()), "pairs": pairs, "mid": 0 if anchored else mid})
+    return effects
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("ckpt")
+    where = ap.add_mutually_exclusive_group(required=True)
+    where.add_argument("--out", help="write the renamed checkpoint into this new directory")
+    where.add_argument("--in-place", action="store_true", help="edit the checkpoint itself, keeping a backup")
+    where.add_argument("--restore", action="store_true", help="undo the latest --in-place edit")
+    ap.add_argument("--rule", nargs=2, action="append", default=[], metavar=("PATTERN", "REPL"))
+    ap.add_argument("--from-report", metavar="JSONL", help="the automatic rename fixes diagnose.py --out found for this checkpoint")
+    ap.add_argument("--plan", metavar="FILE", help="no prompt: the first run saves the plan here, the second identical run applies it")
+    args = ap.parse_args()
+
+    refuse_stub(args.ckpt)
+    if args.restore:
+        for name in restore(args.ckpt):
+            print(f"restored {name}")
+        return 0
+    rules = []
+    if args.from_report:
+        for fix in report_fixes(args.from_report, args.ckpt, "rename"):
+            rules += [tuple(r) for r in fix["rules"]]
+        if not rules and not args.rule:
+            raise SystemExit(f"the record of {args.ckpt} in {args.from_report} has no automatic rename fix")
+    rules += [tuple(r) for r in args.rule]
+    if not rules:
+        ap.error("give --rule PATTERN REPL or --from-report")
+    ck = Checkpoint(args.ckpt)
+    check_out_dir(ck, args.out)
+
+    names = list(ck.tensors())
+    effects = rule_effects(names, rules)
+    lines = [f"plan: rename tensors of {ck.path} ({len(names)} tensors)"]
+    for i, e in enumerate(effects, 1):
+        pat, repl = e["rule"]
+        lines.append(f"  rule {i}: s/{pat}/{repl}/ renames {e['renamed']} tensors")
+        for (old, new), n in e["pairs"].most_common(6):
+            lines.append(f"      {old}  ->  {new}" + (f"   (x{n})" if n > 1 else ""))
+        if len(e["pairs"]) > 6:
+            lines.append(f"      ... {len(e['pairs']) - 6} more name patterns")
+        if e["mid"]:
+            lines.append(f"      warning: {e['mid']} of these matches are inside the name, not at its start; anchor the pattern with ^ or $ if that is not meant")
+    dead = [i for i, e in enumerate(effects, 1) if not e["renamed"]]
+    if dead:
+        print("\n".join(lines))
+        raise SystemExit(f"refused: rule {', '.join(map(str, dead))} renames no tensor (check the pattern against the names above)")
+
+    try:
+        changes = ck.rename_regex(rules)
+    except ValueError as e:
+        raise SystemExit(f"refused: {e}")
+    for pat, repl in rules:
+        ck.quant_rename(pat, repl)
+    lines += [f"  {line}" for line in ck.log]
+    lines.append(f"  {len(changes)} tensors renamed, {len(names) - len(changes)} keep their names")
+    old_modules = {n.rpartition(".")[0] for n in changes}
+    new_modules = {n.rpartition(".")[0] for n in ck.tensors()}
+    for entry in ck.stale_quant_entries(old_modules, new_modules):
+        lines.append(f"  warning: {entry} still names the old modules and none of the new ones; fix it with fix_config.py after this rename")
+    lines += plan_lines_for_save(ck, args.out)
+    plan = {"tool": "rename_tensors", "ckpt": ck.path, "fingerprint": ck.fingerprint, "out": args.out and os.path.abspath(args.out),
+            "rules": [list(r) for r in rules], "renamed": len(changes), "log": ck.log}
+    if not confirm(plan, lines, args.plan):
+        return 0
+    for name in ck.save(out=None if args.in_place else args.out):
+        print(f"wrote {name}")
+    if args.plan:
+        os.remove(args.plan)
+    print("done: run diagnose.py on the result to check it")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From bb15d1d387c8bb45441e2cef96115e9449bebc87 Mon Sep 17 00:00:00 2001
From: Your Open Source Tooling Specialist
 <121191527+bettercallcaleb@users.noreply.github.com>
Date: Tue, 20 Oct 2026 00:40:55 -0400
Subject: [PATCH 133/249] fix(server): accept localhost authority in local file
 URIs (#758)

Resolve the parsed URI path without prepending localhost, while
rejecting non-local authorities and retaining the media root check.
Add regression cases for localhost authority spellings and remote hosts.

Refs FlashML-org/FreeToken#757
---
 python/freetoken/mm/media.py  |  5 ++++-
 tests/server/test_mm_input.py | 10 ++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py
index 5c3c4ef789..40c782eed4 100644
--- a/python/freetoken/mm/media.py
+++ b/python/freetoken/mm/media.py
@@ -66,7 +66,10 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes:
     if not root:
         raise ValueError("cannot load local files without --allowed-local-media-path")
     spec = urlparse(url)
-    filepath = Path(url2pathname((spec.netloc or "") + (spec.path or "")))
+    # RFC 8089: empty authority and localhost identify the local machine.
+    if spec.netloc and spec.netloc.lower().rstrip(".") != "localhost":
+        raise ValueError("file:// URLs must have no host or use localhost")
+    filepath = Path(url2pathname(spec.path or ""))
     # resolve() follows symlinks, so a link inside the root escaping it is rejected too
     resolved = filepath.resolve()
     if Path(root).resolve() not in resolved.parents:
diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py
index 19d9f86529..bfc9b16d65 100644
--- a/tests/server/test_mm_input.py
+++ b/tests/server/test_mm_input.py
@@ -119,6 +119,16 @@ def test_local_media_requires_allowlisted_root(tmp_path):
     config = _config(allowed_local_media_path=str(tmp_path))
     assert asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config)) == [b"fakepng"]
 
+    # RFC 8089 local authority refers to the same file as file:///... .
+    for authority in ("localhost", "LOCALHOST", "localhost."):
+        local_url = f"file://{authority}{img}"
+        assert asyncio.run(fetch_image_bytes([{"kind": "url", "data": local_url}], config)) == [b"fakepng"]
+
+    # Non-local authorities must not be interpreted as filesystem paths.
+    remote_url = f"file://remote.invalid{img}"
+    with pytest.raises(ValueError, match="no host or use localhost"):
+        asyncio.run(fetch_image_bytes([{"kind": "url", "data": remote_url}], config))
+
     # a path outside the root is rejected even with the gate on
     with pytest.raises(ValueError, match="subpath"):
         asyncio.run(fetch_image_bytes([{"kind": "url", "data": "file:///etc/hostname"}], config))

From cc56b6bdcf9188eb632197d6ec719c61bff11798 Mon Sep 17 00:00:00 2001
From: Mukkesh Ganesh <mukmckenzie@gmail.com>
Date: Mon, 19 Oct 2026 23:20:40 -0700
Subject: [PATCH 134/249] feat(gguf): serve GLM-5.3-Flash GGUF checkpoints on
 glm5-next (#721)

* feat(gguf): serve GLM-5.3-Flash GGUF checkpoints on glm5-next

Adds the llama.cpp glm5-next GGUF reader (unsloth GLM-5.3-Flash-GGUF, both
the glm5-next and the older glm5next arch spellings) on the GGUF quant
kind from #690:

- parse_gguf_config derives the text config from the GGUF KV metadata,
  with per-block attention.head_count_kv laying out the KDA/DSA hybrid,
  the trailing MTP block dropped like the HF loader drops it, and the
  MLA head widths from the *_length_mla keys.
- iter_gguf_weights maps every tensor to the glm5_next state dict and
  undoes the converter transforms: A_log stored as -exp(A_log) is
  inverted, the per-head attn_k_b/attn_v_b pair (k transposed) fuses
  into the dense kv_b_proj layout the MLA absorption reads, and the
  q|k|v conv kernels concatenate on the channel axis. The mixed-type
  fused KDA in_proj (Q8_0 q|k|v + BF16 b|f_a|g_a) rides #702 segments;
  the routed experts load as GGUF expert banks at their shipped types
  (unsloth dynamic quants: IQ3_S gate/up + IQ4_XS down, Q8_0 outlier
  layers pass through natively).
- passes quant_config to embed_tokens so the Q8_0 token_embd serves
  packed; registers Glm5NextGGUFForCausalLM and claims it in the AOT
  table.

Verified end to end against unsloth GLM-5.3-Flash-UD-IQ4_XS (5 split
shards): the full engine on a 32GB GPU serves coherent output with the
experts off host banks. Text-only, TP=1.

Part of #671.

* fix(gguf): refuse unread scale tensors in glm5-next ggufs

check_scales at config parse like the llama, qwen3 and gemma4 loaders
(review on #721): a .scale / .input_scale outside an NVFP4 weight now
fails with a clear error instead of loading silently unread.

* fix(gguf): serve the quantized tensors of unsloth dynamic glm5-next quants

unsloth's UD quants quantize tensors the AesSedAI-style files keep float:
the MLA kv_b halves (Q8_0), the mHC mixers and the indexer's wk /
weights_proj / compressor gate. kv_b_proj joins the registry's
unquantized_modules (glm_moe_dsa keeps its kv_b dense for the MLA
absorption for the same reason) and the reader dequantizes through the
vendored gguf-py reference on every float-expecting read; wq_b keeps
serving its checkpoint quant packed (attention.py keeps wk and
weights_proj dense). Regression test stores the kv_b halves Q8_0.

* fix(gguf): treat the glm5 pre-tokenizer spelling like glm4

llama-vocab.cpp maps both spellings to CHATGLM4 with ignore_merges;
unsloth writes glm4 and AesSedAI's writer glm5 for the same vocabulary.
Without the mapping a glm5 file loads but silently mis-tokenizes
(tok_compare: 22,840/153,725 vocab words differ from the HF reference;
0 everywhere with it).

* fix(gguf): size the mmq tail padding by ggml type and k

pad_rows added one zero row of k columns, which covers the bytes MMQ
reads past the last row only when k is at least half a padding unit
(512); glm5-next's KDA f_b / g_b (k = linear head dim 128, Q8_0 in
dynamic quants) are the first packed tensors below that and died at the
first mul_mat. pad_rows now takes the ggml type and returns enough whole
rows to cover mmq_tail_bytes for any (type, k), consistently in the
reader's fused(), the gguf linear and embedding params and the expert
host banks. Found serving the real unsloth UD-IQ4_XS: the dry load
never runs a forward.

* fix(gguf): take the older glm5-next writer variants in stride

Two long-tail writer variants the full check sweep found (the GSQ-RCO
lineage alone is ~57k downloads/month):

- the conv kernels with a trailing unit dim: llama.cpp's glm5-next
  loader tries that 4-D shape before the 3-D one, so trailing unit dims
  in ne are dropped on float reads (they carry no elements).
- attention.head_count_kv as the generic scalar (all blocks): wrong for
  the hybrid -- llama.cpp broadcasts it and then fails tensor creation
  -- so only the list is trusted and a scalar falls to the tensor-name
  census like the pre-list writer.

* fix(gguf): glm5-next review cleanups

- refuse instead of guessing metadata the model cannot express: the
  scalar attention.head_count_kv spelling (llama.cpp broadcasts it and
  then fails tensor creation; every public file carries the list),
  missing *_length_mla keys through need(), non-uniform
  swiglu_clamp_exp / swiglu_clamp_shexp, indexer.types with non-full
  layers, expert_gating_func other than sigmoid, and rope dims as a
  NotImplementedError instead of an assert.
- call check_scales directly; aot_models comment on its own line.
- drop the GGUF test files; coverage will be handled across families.
---
 python/freetoken/kernel/aot_models.py         |   3 +-
 .../layers/quantization/embedding/gguf.py     |   2 +-
 .../layers/quantization/ggml_types.py         |   5 +-
 .../layers/quantization/linear/gguf.py        |   2 +-
 python/freetoken/models/gguf/config.py        |   3 +
 python/freetoken/models/gguf/experts.py       |   8 +-
 python/freetoken/models/gguf/tensors.py       |   2 +-
 python/freetoken/models/gguf/tokenizer.py     |   7 +-
 python/freetoken/models/glm5_next/__init__.py |   4 +
 python/freetoken/models/glm5_next/gguf.py     | 367 ++++++++++++++++++
 python/freetoken/models/glm5_next/model.py    |   2 +
 python/freetoken/models/register.py           |  13 +
 12 files changed, 405 insertions(+), 13 deletions(-)
 create mode 100644 python/freetoken/models/glm5_next/gguf.py

diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py
index fd67d56541..a68c868ea7 100644
--- a/python/freetoken/kernel/aot_models.py
+++ b/python/freetoken/kernel/aot_models.py
@@ -281,7 +281,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         # KDA conv/recurrent state lives in the LinearStatePool, not paged KV.
         name="RedHatAI/GLM-5.3-Flash-NVFP4",
         architecture="Glm5NextForCausalLM",
-        arch_aliases=("Glm5NextForConditionalGeneration",),
+        # GGUF expert rows differ per layer and reach the fused copy as runtime sizes
+        arch_aliases=("Glm5NextForConditionalGeneration", "Glm5NextGGUFForCausalLM"),
         hidden_size=4096,
         kv_groups=(),
         top_k=8,
diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py
index 56c72c2d6a..98f12dd9cc 100644
--- a/python/freetoken/layers/quantization/embedding/gguf.py
+++ b/python/freetoken/layers/quantization/embedding/gguf.py
@@ -45,4 +45,4 @@ class GGUFEmbeddingMethod(EmbeddingMethod):
     def create_weights(self, layer: Any) -> None:
         (ggml_type,) = gguf_types(self.cfg.scheme)
         h = self.cfg.embedding_dim
-        layer.weight = torch.empty(self.cfg.num_embeddings + pad_rows(h), row_bytes(ggml_type, h), dtype=torch.uint8)
+        layer.weight = torch.empty(self.cfg.num_embeddings + pad_rows(ggml_type, h), row_bytes(ggml_type, h), dtype=torch.uint8)
diff --git a/python/freetoken/layers/quantization/ggml_types.py b/python/freetoken/layers/quantization/ggml_types.py
index 262c5668d3..95c2779311 100644
--- a/python/freetoken/layers/quantization/ggml_types.py
+++ b/python/freetoken/layers/quantization/ggml_types.py
@@ -44,9 +44,10 @@ def row_bytes(ggml_type: int, k: int) -> int:
     return k // block * size
 
 
-def pad_rows(k: int) -> int:
+def pad_rows(ggml_type: int, k: int) -> int:
     """Zero rows to append after a packed [rows, K] matrix so ggml's tail reads stay inside it."""
-    return 0 if k % MATRIX_ROW_PADDING == 0 else 1
+    tail = mmq_tail_bytes(ggml_type, k)
+    return 0 if tail == 0 else max(1, -(-tail // row_bytes(ggml_type, k)))
 
 
 def mmq_tail_bytes(ggml_type: int, k: int) -> int:
diff --git a/python/freetoken/layers/quantization/linear/gguf.py b/python/freetoken/layers/quantization/linear/gguf.py
index 41f61edf88..badbe92b49 100644
--- a/python/freetoken/layers/quantization/linear/gguf.py
+++ b/python/freetoken/layers/quantization/linear/gguf.py
@@ -54,7 +54,7 @@ def create_weights(self, layer: Any) -> None:
         k = g.in_features
         for i, s in enumerate(segs):
             if is_quantized(s.ggml_type):
-                setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(k), row_bytes(s.ggml_type, k), dtype=torch.uint8))
+                setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(s.ggml_type, k), row_bytes(s.ggml_type, k), dtype=torch.uint8))
             else:
                 setattr(layer, segment_param(i), torch.empty(s.rows, k))
         if g.scheme.weight.scale is not None:
diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py
index 5b780fe0ef..5767d4eb1c 100644
--- a/python/freetoken/models/gguf/config.py
+++ b/python/freetoken/models/gguf/config.py
@@ -19,6 +19,9 @@
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
     "glm4moe": "Glm4MoeGGUFForCausalLM",
+    # glm5-next: the llama.cpp spelling; glm5next: unsloth's first release
+    "glm5-next": "Glm5NextGGUFForCausalLM",
+    "glm5next": "Glm5NextGGUFForCausalLM",
     "llama": "LlamaGGUFForCausalLM",
     "minimax-m2": "MiniMaxM2GGUFForCausalLM",
     "minimax-m3": "MiniMaxM3GGUFForCausalLM",
diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py
index c21fff9fa2..0efb5982fb 100644
--- a/python/freetoken/models/gguf/experts.py
+++ b/python/freetoken/models/gguf/experts.py
@@ -40,10 +40,10 @@ def _drop_page_cache(info: GgufTensorInfo) -> None:
         pass
 
 
-def _device_bank(shape: tuple[int, ...], k: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+def _device_bank(shape: tuple[int, ...], ggml_type: int, k: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
     """An [E, rows, row] device bank; a zero row follows the last one when ggml's MMQ tiles read past it."""
     experts, rows, row = shape
-    flat = torch.empty(experts * rows + pad_rows(k), row, dtype=dtype, device=device)
+    flat = torch.empty(experts * rows + pad_rows(ggml_type, k), row, dtype=dtype, device=device)
     flat[experts * rows :].zero_()
     return flat[: experts * rows].view(shape)
 
@@ -101,8 +101,8 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int
                 banks = {role: bank.tensor for role, bank in host.items()}
             else:
                 banks = {
-                    "gate_up": _device_bank(gate_up_shape, gate.ne[0], rows[0].dtype, device),
-                    "down": _device_bank(tuple(d.shape), down.ne[0], d.dtype, device),
+                    "gate_up": _device_bank(gate_up_shape, gate.ggml_type, gate.ne[0], rows[0].dtype, device),
+                    "down": _device_bank(tuple(d.shape), down.ggml_type, down.ne[0], d.dtype, device),
                 }
                 if scaled:
                     banks["global"] = torch.empty(experts, 4, dtype=torch.float32, device=device)
diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py
index 83452d957a..f2cb0e0038 100644
--- a/python/freetoken/models/gguf/tensors.py
+++ b/python/freetoken/models/gguf/tensors.py
@@ -85,7 +85,7 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k
         rows = parts[seg.first_part : seg.first_part + seg.num_parts]
         zeros = pad if i == len(segs) - 1 else 0
         if is_quantized(seg.ggml_type):
-            zeros += pad_rows(k)
+            zeros += pad_rows(seg.ggml_type, k)
             dtype = torch.uint8
         else:
             rows = [r.to(torch.bfloat16) for r in rows]
diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py
index 3945c05db6..26e447aad2 100644
--- a/python/freetoken/models/gguf/tokenizer.py
+++ b/python/freetoken/models/gguf/tokenizer.py
@@ -23,7 +23,7 @@
 ))
 
 # Llama 3 and GLM pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it
-_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe", "glm4"))
+_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe", "glm4", "glm5"))
 
 
 def _add_bos(meta: dict) -> bool:
@@ -81,11 +81,12 @@ def load_gguf_tokenizer(model_path: str):
     meta = load_gguf_metadata(model_path)
     fast = _fix_added_tokens(fast, meta)
     pre = meta.get("tokenizer.ggml.pre") if meta.get("tokenizer.ggml.model") == "gpt2" else None
-    if pre == "glm4":
+    if pre in ("glm4", "glm5"):
+        # llama.cpp spells the same GLM pre-tokenizer "glm4" and "glm5" (llama-vocab.cpp maps both
+        # to CHATGLM4 with ignore_merges); transformers maps no split for either
         from tokenizers import Regex, pre_tokenizers
         from transformers.integrations.gguf.gguf_tokenizer_mapping import GGUF_PRE_TOKENIZER_SPLITS
 
-        # transformers maps no split for glm4, which the HF tokenizer and llama.cpp split like Llama 3
         fast.pre_tokenizer = pre_tokenizers.Sequence([
             pre_tokenizers.Split(Regex(GGUF_PRE_TOKENIZER_SPLITS["llama3"]), behavior="isolated"),
             pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
diff --git a/python/freetoken/models/glm5_next/__init__.py b/python/freetoken/models/glm5_next/__init__.py
index 7eee33dfee..0c548da7f6 100644
--- a/python/freetoken/models/glm5_next/__init__.py
+++ b/python/freetoken/models/glm5_next/__init__.py
@@ -1,4 +1,5 @@
 from .config import VisionConfig, parse_config, parse_vision_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
 from .vision import Glm5NextVisionModel
 from .weight import iter_expert_pieces, iter_vision_weights, iter_weights, nvfp4_expert_spec
@@ -10,8 +11,11 @@
     "Glm5NextVisionModel",
     "VisionConfig",
     "parse_config",
+    "parse_gguf_config",
     "parse_vision_config",
     "iter_vision_weights",
+    "iter_gguf_weights",
     "iter_weights",
+    "to_gguf",
     "iter_expert_pieces",
 ]
diff --git a/python/freetoken/models/glm5_next/gguf.py b/python/freetoken/models/glm5_next/gguf.py
new file mode 100644
index 0000000000..73eb0a82ce
--- /dev/null
+++ b/python/freetoken/models/glm5_next/gguf.py
@@ -0,0 +1,367 @@
+"""GLM-5.3-Flash from a llama.cpp ``glm5-next`` GGUF (unsloth ``GLM-5.3-Flash-GGUF``): config from the GGUF metadata, tensor names, and the converter's transforms undone.
+
+The GGUF geometry is identical to the HF checkpoint (34 KDA linear-attention layers,
+11 NoPE-MLA/DSA layers with a kpool-compressed indexer on a 3+1 hybrid block layout,
+mHC x4 residual streams, 288-expert sigmoid/noaux_tc MoE with a shared expert and
+clamped SwiGLU), so ``parse_gguf_config`` derives a text-config-shaped namespace from
+the GGUF KV metadata and reuses ``glm5_next.config.parse_config`` wholesale.
+
+The llama.cpp converter (unsloth's glm5-next writer; tensor names identical across
+the ``glm5-next`` and the older ``glm5next`` arch spellings) stores KDA's ``A_log``
+pre-transformed as ``-exp(A_log)`` -- inverted at load -- and splits the MLA
+``kv_b_proj`` into per-head 3-D ``attn_k_b`` / ``attn_v_b`` tensors (k stored
+transposed, [heads, kv_lora, nope] next to [heads, v, kv_lora]); the halves may be
+quantized (dynamic quants store them Q8_0), so they dequantize and fuse into the
+dense layout the module's MLA absorption reads. The trailing MTP block (``blk.N.nextn.*``, ``nextn_predict_layers == 1``) is
+dropped exactly like the HF loader drops the never-read MTP layer; routed experts
+serve from the GGUF expert banks (``freetoken.models.gguf.experts``).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import numpy as np
+import torch
+
+from freetoken.layers.quantization.ggml_types import is_quantized
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names
+from freetoken.models.gguf.tensors import check_scales, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+# HF config fields the GGUF metadata does not carry (GLM-5.3-Flash's config.json values).
+_FAMILY_DEFAULTS = dict(
+    hidden_act="silu",
+    attention_bias=False,
+    # NoPE means no rope on the main attention; the indexer's interleaved rope is a
+    # served-model fact the GGUF KV predates, and getting it wrong scores the indexer
+    # keys incorrectly (DSA retrieval degrades to garbage).
+    indexer_rope_interleave=True,
+)
+
+# checkpoint module leaf (glm5_next HF naming, language-model rooted) -> GGUF tensors
+_LAYER_LEAVES: dict[str, tuple[str, ...]] = {
+    # KDA input projections (the six leaves of the fused in_proj)
+    "self_attn.q_proj": ("attn_q",),
+    "self_attn.k_proj": ("attn_k",),
+    "self_attn.v_proj": ("attn_v",),
+    "self_attn.b_proj": ("ssm_beta",),
+    "self_attn.f_a_proj": ("ssm_f_a",),
+    "self_attn.g_a_proj": ("ssm_g_a",),
+    "self_attn.f_b_proj": ("ssm_f_b",),
+    "self_attn.g_b_proj": ("ssm_g_b",),
+    # both layer kinds name their output projection the same
+    "self_attn.o_proj": ("attn_output",),
+    # DSA (MLA) projections
+    "self_attn.q_a_proj": ("attn_q_a",),
+    "self_attn.q_b_proj": ("attn_q_b",),
+    "self_attn.kv_a_proj_with_mqa": ("attn_kv_a_mqa",),
+    # two tensors, dequantized and fused at load: the module keeps its dense bf16 kv_b
+    # (exempted in the registry's unquantized_modules, so quantized halves are fine too)
+    "self_attn.kv_b_proj": ("attn_k_b", "attn_v_b"),
+    # kpool indexer
+    "self_attn.indexer.wq_b": ("indexer.attn_q_b",),
+    "self_attn.indexer.wk": ("indexer.attn_k",),
+    "self_attn.indexer.weights_proj": ("indexer.proj",),
+    # dense MLP of the leading blocks
+    "mlp.gate_proj": ("ffn_gate",),
+    "mlp.up_proj": ("ffn_up",),
+    "mlp.down_proj": ("ffn_down",),
+    # shared expert of the MoE blocks
+    "mlp.shared_experts.gate_proj": ("ffn_gate_shexp",),
+    "mlp.shared_experts.up_proj": ("ffn_up_shexp",),
+    "mlp.shared_experts.down_proj": ("ffn_down_shexp",),
+    # routed experts (GGUF expert banks)
+    "mlp.experts.0.gate_proj": ("ffn_gate_exps",),
+    "mlp.experts.0.up_proj": ("ffn_up_exps",),
+    "mlp.experts.0.down_proj": ("ffn_down_exps",),
+}
+_TOP_LEAVES: dict[str, tuple[str, ...]] = {
+    # checkpoint-space names: the language-model root is already applied when to_gguf runs
+    "model.language_model.embed_tokens": ("token_embd",),
+    "lm_head": ("output",),
+}
+_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return _TOP_LEAVES[name]
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ()))
+
+
+def _layer_types(shim: "GgufConfigShim", num_layers: int) -> tuple[str, ...]:
+    """Per-decoder-layer kinds: DSA where the MLA projections sit, KDA elsewhere, from the
+    per-block ``attention.head_count_kv`` list (0 on the KDA blocks, 1 on the MLA ones).
+
+    Every public glm5-next GGUF carries the list (the generic scalar spelling cannot express
+    the hybrid; llama.cpp broadcasts it and then fails tensor creation), so anything else is
+    refused rather than guessed at."""
+    from .args import DSA_LAYER, KDA_LAYER
+
+    per_block = shim.metadata.get(f"{shim.model_type}.attention.head_count_kv")
+    if not isinstance(per_block, list) or len(per_block) < num_layers:
+        raise ValueError(
+            f"{shim.model_path}: glm5-next GGUF lacks the per-block attention.head_count_kv list "
+            "the KDA/MLA hybrid layout needs (reconvert with a current llama.cpp)"
+        )
+    return tuple(DSA_LAYER if int(kv) else KDA_LAYER for kv in per_block[:num_layers])
+
+
+def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig:
+    md = shim.metadata
+    prefix = shim.model_type
+
+    check_scales(shim.model_path)
+
+    def kv(key: str, default=None):
+        return md.get(f"{prefix}.{key}", default)
+
+    def need(key: str):
+        val = kv(key)
+        if val is None:
+            raise KeyError(f"missing GGUF metadata key {prefix}.{key}")
+        return val
+
+    # block_count also counts the MTP block a GGUF carries after the decoder
+    n_mtp = int(kv("nextn_predict_layers", 0) or 0)
+    num_layers = int(need("block_count")) - n_mtp
+    layer_types = _layer_types(shim, num_layers)
+    first_dense = int(need("leading_dense_block_count"))
+
+    heads = int(need("attention.head_count"))
+    qk_nope, v_head = int(need("attention.key_length_mla")), int(need("attention.value_length_mla"))
+    # NoPE: the main attention carries no rotary dims (rope survives only in the indexer)
+    if int(kv("rope.dimension_count", 0) or 0) != 0:
+        raise NotImplementedError("glm5-next GGUF: rotary dims on the main attention are not served")
+
+    # KDA head count: one q/k/v head set of kda.head_dim each; the conv channels or the
+    # beta rows count them, falling back to the shared head_count (the two coincide on
+    # every release) when the tensor table is gone (metadata-only GGUF).
+    kda_head_dim = int(need("kda.head_dim"))
+    index = gguf_tensor_index(shim.model_path)
+    conv_q = next((i for i in index.values() if i.name.endswith("ssm_conv1d_q.weight")), None)
+    kda_heads = conv_q.ne[2] // kda_head_dim if conv_q is not None else heads
+
+    kpool = int(kv("attention.indexer.kpool", 1) or 1)
+    # the compressor tensors exist exactly when the indexer pools; a metadata-only GGUF
+    # keeps the fact only through kpool > 1
+    names = gguf_tensor_names(shim.model_path)
+    kpool_compress = any(n.endswith("indexer_compressor_gate.weight") for n in names) if names else kpool > 1
+    indexer_types_kv = kv("attention.indexer.types")
+    if indexer_types_kv is not None and not all(indexer_types_kv):
+        # llama.cpp reads False as "no full indexer on this layer", which the model cannot express
+        raise NotImplementedError(f"{shim.model_path}: glm5-next GGUF with layers that lack a full indexer is not served")
+    indexer_types = ("full",) * num_layers
+
+    # the clamp is per-layer in the GGUF (and the shared expert carries its own array); the model
+    # takes one value, so anything non-uniform is refused instead of silently taking the first
+    clamp = need("swiglu_clamp_exp")
+    clamp_vals = list(clamp) if isinstance(clamp, list) else [clamp]
+    clamp_shexp = kv("swiglu_clamp_shexp")
+    shexp_vals = list(clamp_shexp) if isinstance(clamp_shexp, list) else [clamp_shexp] if clamp_shexp is not None else []
+    if len(set(map(float, clamp_vals + shexp_vals))) != 1:
+        raise NotImplementedError(f"{shim.model_path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served")
+    if int(kv("expert_gating_func", 2) or 2) != 2:
+        raise NotImplementedError(f"{shim.model_path}: expert_gating_func other than sigmoid (2) is not served")
+    text = SimpleNamespace(
+        **_FAMILY_DEFAULTS,
+        num_hidden_layers=num_layers,
+        layer_types=tuple(layer_types),
+        mlp_layer_types=("dense",) * first_dense + ("sparse",) * (num_layers - first_dense),
+        first_k_dense_replace=first_dense,
+        hidden_size=int(need("embedding_length")),
+        num_attention_heads=heads,
+        vocab_size=int(shim.vocab_size),
+        intermediate_size=int(need("feed_forward_length")),
+        rms_norm_eps=float(need("attention.layer_norm_rms_epsilon")),
+        max_position_embeddings=int(need("context_length")),
+        q_lora_rank=int(need("attention.q_lora_rank")),
+        kv_lora_rank=int(need("attention.kv_lora_rank")),
+        qk_nope_head_dim=qk_nope,
+        qk_rope_head_dim=0,
+        v_head_dim=v_head,
+        mla_use_nope=True,
+        index_n_heads=int(need("attention.indexer.head_count")),
+        index_head_dim=int(need("attention.indexer.key_length")),
+        index_topk=int(need("attention.indexer.top_k")),
+        indexer_types=indexer_types,
+        index_kpool=kpool,
+        index_kpool_compress=kpool_compress,
+        # old converters predate the key; both released checkpoints set it true
+        index_kpool_always_select_tail=bool(kv("attention.indexer.kpool_select_tail", True)),
+        linear_attn_config={
+            "num_heads": kda_heads,
+            "head_dim": kda_head_dim,
+            "short_conv_kernel_size": int(need("ssm.conv_kernel")),
+            "gate_lower_bound": float(need("kda.gate_lower_bound")),
+        },
+        mhc=True,
+        hc_mult=int(need("hyper_connection.count")),
+        hc_sinkhorn_iters=int(need("hyper_connection.sinkhorn_iterations")),
+        hc_eps=float(need("hyper_connection.epsilon")),
+        n_routed_experts=int(need("expert_count")),
+        n_shared_experts=int(kv("expert_shared_count", 1) or 1),
+        num_experts_per_tok=int(need("expert_used_count")),
+        moe_intermediate_size=int(need("expert_feed_forward_length")),
+        norm_topk_prob=bool(need("expert_weights_norm")),
+        routed_scaling_factor=float(need("expert_weights_scale")),
+        n_group=int(kv("expert_group_count", 1) or 1),
+        topk_group=int(kv("expert_group_used_count", 1) or 1),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        swiglu_limit=float(clamp_vals[0]),
+    )
+    hf_like = SimpleNamespace(
+        text_config=text,
+        model_type="glm5_next",
+        architectures=list(shim.architectures),
+        vision_config=None,
+        image_token_id=None,
+    )
+    config = parse_config(hf_like)
+    return replace(config, expert_quant="gguf")
+
+
+# ----------------------------------------------------------------------------------------------
+# Weights
+# ----------------------------------------------------------------------------------------------
+
+def _float(info: GgufTensorInfo) -> torch.Tensor:
+    """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize
+    tensors the module consumes as floats (the kv_b halves, the mHC mixers, the compressor gate).
+    Trailing unit dims in ``ne`` are dropped -- some writers ship the conv kernels 4-D
+    (llama.cpp's glm5-next loader tries that shape first) and a trailing 1 carries no elements."""
+    ne = info.ne
+    while len(ne) > 1 and ne[-1] == 1:
+        ne = ne[:-1]
+    if not is_quantized(info.ggml_type):
+        from freetoken.models.gguf.tensors import float_tensor
+
+        return float_tensor(info).reshape(tuple(reversed(ne)))
+    from freetoken.third_party.gguf import quants
+    from freetoken.third_party.gguf.constants import GGMLQuantizationType
+
+    flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type))
+    return torch.from_numpy(flat).view(*reversed(ne))
+
+
+def _kv_b(kb: GgufTensorInfo, vb: GgufTensorInfo, num_heads: int, kv_lora: int) -> torch.Tensor:
+    """Per-head k [kv_lora, nope] (stored transposed) + v [v, kv_lora] -> the module's fused dense [heads * (nope + v), kv_lora]."""
+    k = _float(kb).transpose(1, 2)  # [heads, nope, kv_lora]
+    fused_w = torch.cat([k, _float(vb)], dim=1)
+    assert fused_w.shape[0] == num_heads, (fused_w.shape, num_heads)
+    return fused_w.reshape(num_heads * fused_w.shape[1], kv_lora).to(torch.bfloat16).contiguous()
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+    include_vision: bool = True,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``)."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    config = parse_gguf_config(cached_load_hf_config(model_path))
+    args = config.glm5_args
+    index = gguf_tensor_index(model_path)
+    hidden = config.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    def bf16(name: str) -> torch.Tensor:
+        return _float(t(name)).to(torch.bfloat16)
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", bf16("output_norm")
+
+    for layer in range(config.num_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", bf16(f"{blk}.attn_norm")
+        yield f"{pre}.post_attention_layernorm.weight", bf16(f"{blk}.ffn_norm")
+        # mHC mixing tensors, fp32 on every layer
+        for sub in ("attn", "ffn"):
+            for part in ("fn", "base", "scale"):
+                yield f"{pre}.hc_{sub}_{part}", _float(t(f"{blk}.hc_{sub}_{part}"))
+
+        if args.is_kda_layer(layer):
+            sa = f"{pre}.self_attn"
+            p, h, d = args.linear_num_heads * args.linear_head_dim, args.linear_num_heads, args.linear_head_dim
+            group = [t(f"{blk}.{n}") for n in ("attn_q", "attn_k", "attn_v", "ssm_beta", "ssm_f_a", "ssm_g_a")]
+            yield from fused(f"{sa}.in_proj", group, [matrix(i) for i in group], hidden)
+            # the GGUF conv kernels already carry the [channels, 1, kernel] singleton dim
+            conv = torch.cat([_float(t(f"{blk}.ssm_conv1d_{q}")) for q in ("q", "k", "v")], dim=0)
+            yield f"{sa}.conv1d.weight", conv.to(torch.bfloat16)
+            for proj, src, k in (("f_b_proj", "ssm_f_b", d), ("g_b_proj", "ssm_g_b", d), ("o_proj", "attn_output", p)):
+                info = t(f"{blk}.{src}")
+                yield from fused(f"{sa}.{proj}", [info], [matrix(info)], k)
+            # the converter stores A_log pre-transformed as -exp(A_log); the kernels want the raw log
+            yield f"{sa}.A_log", torch.log(-_float(t(f"{blk}.ssm_a")))
+            yield f"{sa}.dt_bias", _float(t(f"{blk}.ssm_dt.bias"))
+            yield f"{sa}.o_norm.weight", bf16(f"{blk}.ssm_norm")
+        else:
+            sa = f"{pre}.self_attn"
+            for proj, src, k in (
+                ("q_a_proj", "attn_q_a", hidden),
+                ("q_b_proj", "attn_q_b", args.q_lora_rank),
+                ("kv_a_proj_with_mqa", "attn_kv_a_mqa", hidden),
+                ("o_proj", "attn_output", args.num_heads * args.v_head_dim),
+            ):
+                info = t(f"{blk}.{src}")
+                yield from fused(f"{sa}.{proj}", [info], [matrix(info)], k)
+            yield f"{sa}.kv_b_proj.weight", _kv_b(
+                t(f"{blk}.attn_k_b"), t(f"{blk}.attn_v_b"), args.num_heads, args.kv_lora_rank
+            )
+            yield f"{sa}.q_a_layernorm.weight", bf16(f"{blk}.attn_q_a_norm")
+            yield f"{sa}.kv_a_layernorm.weight", bf16(f"{blk}.attn_kv_a_norm")
+            # wq_b serves its checkpoint quant (packed); wk / weights_proj are dense modules
+            # (glm_moe_dsa's indexer reasoning), dequantized at load when the file quantized them
+            yield from fused(f"{sa}.indexer.wq_b", [t(f"{blk}.indexer.attn_q_b")], [matrix(t(f"{blk}.indexer.attn_q_b"))], args.q_lora_rank)
+            yield f"{sa}.indexer.wk.weight", _float(t(f"{blk}.indexer.attn_k")).to(torch.bfloat16)
+            yield f"{sa}.indexer.weights_proj.weight", _float(t(f"{blk}.indexer.proj")).to(torch.bfloat16)
+            yield f"{sa}.indexer.k_norm.weight", bf16(f"{blk}.indexer.k_norm")
+            yield f"{sa}.indexer.k_norm.bias", bf16(f"{blk}.indexer.k_norm.bias")
+            yield f"{sa}.indexer.index_kpool_compress_gate", bf16(f"{blk}.indexer_compressor_gate")
+            yield f"{sa}.indexer.index_kpool_compress_ape", _float(t(f"{blk}.indexer_compressor_ape"))
+
+        mlp = f"{pre}.mlp"
+        if layer < config.first_k_dense_replace:
+            for proj, src, k in (
+                ("gate_proj", "ffn_gate", hidden),
+                ("up_proj", "ffn_up", hidden),
+                ("down_proj", "ffn_down", config.intermediate_size),
+            ):
+                info = t(f"{blk}.{src}")
+                yield from fused(f"{mlp}.{proj}", [info], [matrix(info)], k)
+        else:
+            yield f"{mlp}.gate.weight", bf16(f"{blk}.ffn_gate_inp")
+            yield f"{mlp}.e_score_correction_bias", _float(t(f"{blk}.exp_probs_b.bias"))
+            shared_i = config.moe_intermediate_size * max(1, config.n_shared_experts)
+            for proj, src, k in (
+                ("gate_proj", "ffn_gate_shexp", hidden),
+                ("up_proj", "ffn_up_shexp", hidden),
+                ("down_proj", "ffn_down_shexp", shared_i),
+            ):
+                info = t(f"{blk}.{src}")
+                yield from fused(f"{mlp}.shared_experts.{proj}", [info], [matrix(info)], k)
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
diff --git a/python/freetoken/models/glm5_next/model.py b/python/freetoken/models/glm5_next/model.py
index d2bfa684c1..d92977dea2 100644
--- a/python/freetoken/models/glm5_next/model.py
+++ b/python/freetoken/models/glm5_next/model.py
@@ -142,6 +142,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [Glm5NextDecoderLayer(config, i, prefix=f"{prefix}.layers.{i}") for i in range(config.num_layers)]
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 80b732cd3f..4d268ffb50 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -420,6 +420,19 @@ class ModelSpec:
         "Glm5NextForCausalLM",
         packed_modules_mapping=_GLM5_NEXT_PACKED,
     ),
+    # llama.cpp glm5-next GGUF (unsloth GLM-5.3-Flash-GGUF; text only): the same model
+    # classes with every tensor in its packed ggml type
+    "Glm5NextGGUFForCausalLM": ModelSpec(
+        "freetoken.models.glm5_next",
+        "Glm5NextForCausalLM",
+        parse_config="parse_gguf_config",
+        iter_weights="iter_gguf_weights",
+        checkpoint_roots=_LANGUAGE_MODEL_ROOT,
+        packed_modules_mapping=_GLM5_NEXT_PACKED,
+        # the MLA kv_b ships as two per-head tensors (attn_k_b transposed + attn_v_b), quantized in
+        # dynamic quants (unsloth UD); the reader dequantizes and fuses them into the dense module
+        unquantized_modules=("*.self_attn.kv_b_proj",),
+    ),
 }
 
 

From 7a8b7dd8c2220b0594fc58ec60b2d1e32f3b0dbe Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Mon, 14 Sep 2026 18:23:04 -0400
Subject: [PATCH 135/249] test(kvcache): give the qsa fp8 pool fakes the mrope
 flag main now reads

Assisted-by: Claude Opus 5
(cherry picked from commit 150c95d6f39b83e0516f5c9c42fef54c679eae09)
(cherry picked from commit 4db8b858375b01ff7edf37031dc81671b0181853)
(cherry picked from commit 495770e33dcb5c3aa6979c6515b61dfade5e29b3)
---
 tests/kvcache/test_qsa_pool_fp8.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py
index c0a91f862a..02b525f679 100644
--- a/tests/kvcache/test_qsa_pool_fp8.py
+++ b/tests/kvcache/test_qsa_pool_fp8.py
@@ -74,7 +74,7 @@ def _spec():
 
 
 def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3):
-    mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True)
+    mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False)
     mc.kv_cache_group_specs = lambda: (_spec(),)
     return SimpleNamespace(
         model_config=mc,
@@ -200,7 +200,7 @@ def test_factory_threads_kv_quant_into_the_qsa_pool():
 
     mc = SimpleNamespace(
         num_layers=8, has_swa_attention=False, has_linear_attention=True,
-        num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None,
+        num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, model_is_mrope=False,
     )
     mc.kv_cache_group_specs = lambda: (_spec(),)
     pool = create_kvcache_pool(

From 940c26afa4755f6f3ec6b20b62ac5970d6e7f034 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Mon, 14 Sep 2026 23:02:28 -0400
Subject: [PATCH 136/249] fix(models): restore the os import and truthy set the
 fp8 env helpers read

Assisted-by: Claude Opus 5
(cherry picked from commit 82a21ae2432fb891175c8c881f52cb4b710241bc)
(cherry picked from commit 0369ed94d7d085ece40ffdcc0adf93a5c26a6b4e)
(cherry picked from commit b5c7812fceb68b4b7835a80ab7c533e0385d9577)
---
 python/freetoken/models/config.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index bfa3b6cbbb..0d4faae420 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -1,4 +1,6 @@
 from __future__ import annotations
+
+import os
 from dataclasses import dataclass
 from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias
 
@@ -7,12 +9,14 @@
 # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only.
 VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.")
 
+_ENV_TRUE = {"1", "true", "yes", "on"}
+
 
 def fp8_dense_enabled() -> bool:
     """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF):
     per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs
     (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``."""
-    return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE
+    return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE
 
 
 def fp8_lmhead_enabled() -> bool:
@@ -20,7 +24,7 @@ def fp8_lmhead_enabled() -> bool:
     for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits:
     a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own
     quality gate. ``FREETOKEN_FP8_LMHEAD=1``."""
-    return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE
+    return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE
 
 
 def detect_expert_quant(hf_config: Any) -> str:

From 8808762bcfdf80f04cafc2435a07f6b931aa3a7c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Mon, 14 Sep 2026 23:00:44 -0400
Subject: [PATCH 137/249] test(qwen4_exp): restore the rope device and clear
 get_rope after the cpu build test

Assisted-by: Claude Opus 5
(cherry picked from commit 27ad92e0a554be016d6e533683c40a2eb67c6164)
(cherry picked from commit e152f665cc1f35f073fbb82daedece77d02f91e1)
(cherry picked from commit 333f7f9f0501116694ede067202070353faf341b)
---
 tests/models/qwen4_exp/test_build_cpu.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py
index 0cf2027d14..58ad0c4dc3 100644
--- a/tests/models/qwen4_exp/test_build_cpu.py
+++ b/tests/models/qwen4_exp/test_build_cpu.py
@@ -10,7 +10,7 @@
 
 import torch
 from freetoken.distributed import set_tp_info, try_get_tp_info
-from freetoken.layers import set_rope_device
+from freetoken.layers import rotary, set_rope_device
 from freetoken.models.qwen4_exp.config import parse_config
 from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM
 
@@ -20,9 +20,15 @@
 def _build():
     if try_get_tp_info() is None:
         set_tp_info(rank=0, size=1)
+    previous = rotary._ROPE_DEVICE
     set_rope_device(torch.device("cpu"))
-    with torch.device("meta"):
-        return Qwen4ExpForCausalLM(parse_config(toy_hf_config()))
+    try:
+        with torch.device("meta"):
+            return Qwen4ExpForCausalLM(parse_config(toy_hf_config()))
+    finally:
+        # get_rope is cached per geometry: a CPU cos/sin table left behind feeds later CUDA tests
+        rotary.get_rope.cache_clear()
+        set_rope_device(previous)
 
 
 def test_model_constructs_on_meta_without_cuda():

From c601ccda40ffcfa584a93c924ae0c62b8de667c8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Mon, 14 Sep 2026 23:19:07 -0400
Subject: [PATCH 138/249] test(kernels): give the block_ends backend fake the
 kv scale accessors the quantized-kv backend reads

Assisted-by: Claude Opus 5
(cherry picked from commit 600819fd7ce18c813cd4b8400cf534bff38f6f7f)
(cherry picked from commit dd2ebdfda3e1c53e5269bd9adfbee7291cc8ffc2)
(cherry picked from commit 12754d41c5fe646ccc6081ce26f05fdf7b0a57b9)
---
 tests/kernels/test_triton_attention.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py
index f9e8e09794..57a84118cd 100644
--- a/tests/kernels/test_triton_attention.py
+++ b/tests/kernels/test_triton_attention.py
@@ -845,6 +845,12 @@ def k_cache(self, layer_id):
         def v_cache(self, layer_id):
             return self.v
 
+        def k_scale(self, layer_id):  # 16-bit pool: no fp8 scales
+            return None
+
+        def v_scale(self, layer_id):
+            return None
+
     device = torch.device("cuda")
     head_dim = 256
     kv_cache = FakeKVCache(device, head_dim)

From 8241a1eea68fa7c9f2d4b62ca6a767ef76167b70 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:58 -0400
Subject: [PATCH 139/249] merge fork #18: fix(moe): report the routing top-k
 figure as a static bound, not an oracle

---
 python/freetoken/engine/engine.py     | 19 +++++++----------
 python/freetoken/moe/offload_cache.py | 12 +++++------
 tests/moe/test_moe_collect_stats.py   | 30 +++++++++++++++++++++++----
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index df5009a2e3..f97feaa732 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -920,9 +920,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         cache.collect_stats = config.moe_collect_stats
         # The routing histogram rides the same switch: on its own the miss rate says how
         # often we fetch, but not whether a smarter policy could have avoided the fetch.
-        # decode_routing_stats turns it into an oracle hit rate -- the ceiling any policy
-        # holding this many slots could reach on the observed routing -- which is the number
-        # worth having before anyone rewrites eviction.
+        # decode_routing_stats turns it into routing-skew numbers, including the hit rate of
+        # the best fixed expert set per layer (not a ceiling: LRU's temporal locality can beat it).
         cache.collect_decode_freq = config.moe_collect_stats
         # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's
         # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4).
@@ -1312,8 +1311,8 @@ def _emit_moe_stats(self) -> None:
         Accumulation is device-side and captured into the decode graph, so it is free to
         leave running; reading it is not (the counters have to come back to the host), which
         is why this only fires every MOE_STATS_INTERVAL decode steps. The routing histogram
-        is deliberately *not* reset -- the oracle bound wants the whole run's distribution,
-        not one window's.
+        is deliberately *not* reset -- it describes the whole run's distribution, not one
+        window's, so the static top-k figure flattens as a session covers more topics.
         """
         cache = self.moe_offload_cache
         agg = cache.decode_miss_stats()
@@ -1339,14 +1338,12 @@ def _emit_moe_stats(self) -> None:
             )
         routing = cache.decode_routing_stats()
         if routing:
-            # oracle_hit_at_slots is the upper bound on hit rate for *any* policy with this
-            # many slots per layer. If it sits near the realized hit rate, the cache is
-            # already doing as well as the routing allows and the win has to come from
-            # somewhere else (more slots, more bandwidth); if it sits far above, eviction
-            # policy is leaving something on the table.
+            # static_topk_hit is what pinning each layer's most frequent experts would score over
+            # the whole run. A realized hit rate above it means the cache is winning on temporal
+            # locality; it says nothing about how close to optimal eviction is.
             logger.info_rank0(
                 "MoE routing: "
-                f"oracle_hit={routing['oracle_hit_at_slots']:.3f} "
+                f"static_topk_hit={routing['static_topk_hit_at_slots']:.3f} "
                 f"(realized {1.0 - agg['miss_rate']:.3f}), "
                 f"slots/layer={routing['slots_per_layer']:.1f}, "
                 f"working_set={routing['working_set_mean']:.1f}"
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 0d702d8a68..48bba1a725 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -1189,10 +1189,10 @@ def decode_miss_stats_per_layer(self) -> dict:
     def decode_routing_stats(self) -> dict:
         """Per-layer decode routing concentration, for cache-skew analysis.
 
-        Uses the histogram from ``collect_decode_freq``. The ``oracle_hit`` is the best a
-        per-layer LRU holding ``cache_size/num_layers`` slots could achieve on the observed
-        (stationary) routing distribution -- i.e. an upper bound on hit rate that depends
-        purely on how skewed routing is, independent of any LRU/LFU dynamics.
+        Uses the histogram from ``collect_decode_freq``. ``static_topk_hit_at_slots`` is the
+        hit rate of the best *fixed* expert set: each layer's ``cache_size/num_layers`` (rounded)
+        most frequent experts over the whole histogram. It measures routing skew only. It is
+        not an upper bound on a dynamic cache: LRU exploits temporal locality and can beat it.
         """
         freq = self.decode_freq.float()
         total = freq.sum(dim=1)
@@ -1202,7 +1202,7 @@ def decode_routing_stats(self) -> dict:
         slots_per_layer = self.cache_size / self.num_layers
         C = max(1, int(round(slots_per_layer)))
         sorted_f, _ = torch.sort(freq, dim=1, descending=True)
-        oracle_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item()
+        static_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item()
         ws = (freq > 0).sum(dim=1).float()
         cdf = torch.cumsum(sorted_f, dim=1) / total.clamp(min=1).unsqueeze(1)
         cover90 = ((cdf < 0.9).sum(dim=1).float() + 1)[valid]
@@ -1214,7 +1214,7 @@ def decode_routing_stats(self) -> dict:
             "working_set_mean": ws[valid].mean().item(),
             "working_set_max": int(ws[valid].max().item()),
             "experts_for_90pct": cover90.mean().item(),
-            "oracle_hit_at_slots": oracle_hit,
+            "static_topk_hit_at_slots": static_hit,
             "norm_entropy": norm_ent,
         }
 
diff --git a/tests/moe/test_moe_collect_stats.py b/tests/moe/test_moe_collect_stats.py
index f960abbdc0..bd1c4e9b3f 100644
--- a/tests/moe/test_moe_collect_stats.py
+++ b/tests/moe/test_moe_collect_stats.py
@@ -10,7 +10,10 @@
 import io
 from types import SimpleNamespace
 
+import torch
+
 from freetoken.engine.engine import MOE_STATS_INTERVAL, Engine
+from freetoken.moe.offload_cache import OffloadMoeCache
 from freetoken.server.args import ServerArgs, parse_args
 
 
@@ -63,7 +66,7 @@ def decode_routing_stats(self):
             "working_set_mean": 173.1,
             "working_set_max": 243,
             "experts_for_90pct": 92.3,
-            "oracle_hit_at_slots": 0.764,
+            "static_topk_hit_at_slots": 0.764,
             "norm_entropy": 0.813,
         }
 
@@ -82,9 +85,9 @@ def test_emit_reports_and_resets_the_window(caplog):
     cache = _StubCache()
     out = _emit(cache, caplog)
     assert "miss_rate=0.250" in out
-    # The oracle bound is the whole point of the report: it says how much room a different
-    # eviction policy could possibly have.
-    assert "oracle_hit=0.764" in out
+    # Labelled as the fixed-set figure it is; "oracle" implied an upper bound it is not.
+    assert "static_topk_hit=0.764" in out
+    assert "oracle" not in out
     assert "(realized 0.750)" in out
     # Ranked worst-first, and the layer that never ran is left out entirely.
     assert "L0=0.500, L1=0.100" in out
@@ -107,3 +110,22 @@ def test_idle_window_emits_nothing_and_keeps_counters(caplog):
 
 def test_interval_is_a_sane_window():
     assert MOE_STATS_INTERVAL >= 1
+
+
+def test_static_topk_hit_is_not_an_upper_bound_on_lru():
+    """One slot, four experts each routed in a run of four tokens: A A A A B B B B ...
+
+    The best fixed expert catches 4 of 16 tokens; an LRU holding one slot misses only on
+    each run's first token. The figure is routing skew, not a ceiling on a dynamic cache.
+    """
+    trace = [e for e in range(4) for _ in range(4)]
+    cached, lru_hits = None, 0
+    for e in trace:
+        lru_hits += e == cached
+        cached = e
+    freq = torch.bincount(torch.tensor(trace), minlength=4).unsqueeze(0)
+    cache = SimpleNamespace(decode_freq=freq, cache_size=1, num_layers=1, num_experts=4)
+    stats = OffloadMoeCache.decode_routing_stats(cache)
+    assert stats["static_topk_hit_at_slots"] == 0.25
+    assert "oracle_hit_at_slots" not in stats
+    assert lru_hits / len(trace) == 0.75

From aeea9000f9db6fa856aa5097141cafc53f77004d Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:58 -0400
Subject: [PATCH 140/249] merge fork #19: fix(qwen4_exp): refuse an ftw
 checkpoint whose fp8 split differs from the fp8 env flags

---
 python/freetoken/models/qwen4_exp/__init__.py |  2 +
 python/freetoken/models/qwen4_exp/weight.py   | 26 +++++++
 python/freetoken/models/weight.py             | 10 ++-
 tests/models/qwen4_exp/test_fp8_dense.py      | 69 +++++++++++++++++++
 4 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py
index 15f866cad2..a623985636 100644
--- a/python/freetoken/models/qwen4_exp/__init__.py
+++ b/python/freetoken/models/qwen4_exp/__init__.py
@@ -13,6 +13,7 @@
 from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration
 from .weight import (
+    check_ftw_weights,
     ftw_side_files,
     nvfp4_expert_spec,
     iter_vision_weights,
@@ -25,6 +26,7 @@
 from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces
 
 __all__ = [
+    "check_ftw_weights",
     "ftw_side_files",
     "nvfp4_expert_spec",
     "Qwen4ExpForCausalLM",
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 863ea864d6..9b4247ce13 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -479,6 +479,32 @@ def _ple_table_files(folder: str) -> list[str]:
     return sorted(os.path.join(folder, shard) for shard in files)
 
 
+def check_ftw_weights(config, model_path: str, names: set[str]) -> None:
+    """Refuse an FTW checkpoint whose stored FP8 split disagrees with this process's FP8 flags.
+
+    The FTW holds the tensors ``iter_weights`` emitted at conversion, so a mismatch otherwise
+    surfaces only after the dense read, as a KeyError on the first fused GDN projection."""
+    from .config import use_fp8_lmhead
+
+    stored = {
+        "FREETOKEN_FP8_DENSE": any(n.endswith(".in_proj_qkvz.weight_scale") for n in names),
+        "FREETOKEN_FP8_LMHEAD": "lm_head.weight_scale" in names,
+    }
+    wanted = {
+        "FREETOKEN_FP8_DENSE": config.attn_quant == "fp8_dynamic",
+        "FREETOKEN_FP8_LMHEAD": use_fp8_lmhead(config),
+    }
+    wrong = [var for var in stored if stored[var] != wanted[var]]
+    if wrong:
+        was = ", ".join(f"{var}={int(stored[var])}" for var in wrong)
+        now = ", ".join(f"{var}={int(wanted[var])}" for var in wrong)
+        raise ValueError(
+            f"FTW checkpoint {model_path} stores its dense weights as converted with {was}, "
+            f"but this process has {now}; serve it with the conversion-time setting, or "
+            "reconvert with ft checkpoint under the setting you want"
+        )
+
+
 def ftw_side_files(model_path: str, out_dir: str) -> list[str]:
     """Write the PLE n-gram table tensors, and only those, into ``ple-table-*.safetensors`` next to an FTW checkpoint.
 
diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py
index b03266fb67..514d73705b 100644
--- a/python/freetoken/models/weight.py
+++ b/python/freetoken/models/weight.py
@@ -218,12 +218,20 @@ def load_weight(
     # decided at conversion (offload -> experts live in banks, not here); a backend mismatch
     # fails loudly in load_state_dict (strict missing/unexpected expert keys), so the reader
     # just yields the stored weight tensors regardless of the include_moe_experts flag.
-    from freetoken.checkpoint.ftw import is_ftw_checkpoint, iter_ftw_weights
+    from freetoken.checkpoint.ftw import FTWReader, is_ftw_checkpoint, iter_ftw_weights
     from freetoken.models.config import VISION_KEY_PREFIXES
 
     # a text-only engine never built the tower, so its tensors are not even read
     keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES))
     if is_ftw_checkpoint(model_path):
+        spec = get_model_spec(cached_load_hf_config(model_path).architectures[0])
+        # a family whose stored tensors depend on load-time settings checks them against the index before the read
+        check = _model_override(spec, "check_ftw_weights")
+        if check is not None:
+            reader = FTWReader(model_path)
+            names = {e["name"] for e in reader.entries("weight")}
+            reader.close()
+            check(_spec_for_model_path(model_path)[0], model_path, names)
         weights = iter_ftw_weights(model_path, keep=keep)
     else:
         _config, spec = _spec_for_model_path(model_path)
diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
index 058515b183..74da973341 100644
--- a/tests/models/qwen4_exp/test_fp8_dense.py
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -151,3 +151,72 @@ def test_reader_gates_lm_head_behind_its_own_flag():
     passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1,
                                   dense=False, lm_head=True))
     assert list(passthrough) == ["x.self_attn.qkv_proj.weight"]
+
+
+def _ftw_with(tmp_path, names):
+    from freetoken.checkpoint.ftw import FTWWriter
+
+    writer = FTWWriter(str(tmp_path))
+    for name in names:
+        writer.add_tensor(name, torch.zeros(4, dtype=torch.bfloat16))
+    writer.finalize({})
+    return str(tmp_path)
+
+
+_BF16_NAMES = ["model.layers.0.linear_attn.in_proj.weight", "lm_head.weight"]
+_FP8_NAMES = [
+    "model.layers.0.linear_attn.in_proj_qkvz.weight",
+    "model.layers.0.linear_attn.in_proj_qkvz.weight_scale",
+    "model.layers.0.linear_attn.in_proj_ba.weight",
+    "lm_head.weight",
+    "lm_head.weight_scale",
+]
+
+
+@pytest.mark.parametrize(
+    "names, env, bad",
+    [
+        (_BF16_NAMES, {"FREETOKEN_FP8_DENSE": "1", "FREETOKEN_FP8_LMHEAD": "1"},
+         "FREETOKEN_FP8_DENSE=0, FREETOKEN_FP8_LMHEAD=0"),
+        (_FP8_NAMES, {}, "FREETOKEN_FP8_DENSE=1, FREETOKEN_FP8_LMHEAD=1"),
+        (_FP8_NAMES, {"FREETOKEN_FP8_DENSE": "1"}, "converted with FREETOKEN_FP8_LMHEAD=1, but"),
+    ],
+)
+def test_ftw_load_refuses_fp8_flags_that_differ_from_conversion(tmp_path, monkeypatch, names, env, bad):
+    """#8: the FTW replays conversion-time tensors, so the FP8 flags must match before any byte is read."""
+    from freetoken.models import weight as model_weight
+
+    for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"):
+        monkeypatch.delenv(var, raising=False)
+    for var, value in env.items():
+        monkeypatch.setenv(var, value)
+    path = _ftw_with(tmp_path, names)
+    spec = SimpleNamespace(module="freetoken.models.qwen4_exp")
+    config = SimpleNamespace(
+        attn_quant="fp8_dynamic" if env.get("FREETOKEN_FP8_DENSE") else "none",
+        tie_word_embeddings=False,
+        quant=None,
+    )
+    monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"]))
+    monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec)
+    monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec))
+    with pytest.raises(ValueError, match=bad):
+        next(model_weight.load_weight(path, torch.device("cpu")))
+
+
+@pytest.mark.parametrize("names, fp8", [(_BF16_NAMES, False), (_FP8_NAMES, True)])
+def test_ftw_load_accepts_matching_fp8_flags(tmp_path, monkeypatch, names, fp8):
+    from freetoken.models import weight as model_weight
+
+    for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"):
+        if fp8:
+            monkeypatch.setenv(var, "1")
+        else:
+            monkeypatch.delenv(var, raising=False)
+    path = _ftw_with(tmp_path, names)
+    spec = SimpleNamespace(module="freetoken.models.qwen4_exp")
+    config = SimpleNamespace(attn_quant="fp8_dynamic" if fp8 else "none", tie_word_embeddings=False, quant=None)
+    monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"]))
+    monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec)
+    monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec))
+    assert [name for name, _ in model_weight.load_weight(path, torch.device("cpu"))] == names

From f07663db8cfe5bc644dd8c2d5f5ed6f5ae8f5a0f Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:58 -0400
Subject: [PATCH 141/249] merge fork #20: fix(engine): let moe-cache-auto drop
 prefill overlap when its slot floor does not fit

---
 python/freetoken/engine/cache_budget.py | 16 +++++++++--
 python/freetoken/engine/engine.py       |  5 ++++
 tests/engine/test_cache_budget.py       | 37 +++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py
index 891ac2ac11..00a50b86ad 100644
--- a/python/freetoken/engine/cache_budget.py
+++ b/python/freetoken/engine/cache_budget.py
@@ -61,6 +61,7 @@ def plan_cache_budget(
     prefill_overlap: bool,
     kv_reserve_pages: int,
     max_slots: int,
+    page_size: int = 1,
 ) -> tuple[int, int, bool]:
     """Split ``budget_bytes`` MoE-first into (moe_cache_size, num_pages, prefill_overlap).
 
@@ -68,7 +69,8 @@ def plan_cache_budget(
     weights + fixed_cache_size; the (1-memory_ratio) remainder is the graph headroom).
     Experts greedily fill the budget after reserving ``kv_reserve_pages`` for KV, clamped
     to ``[floor, min(total_experts, max_slots)]`` (floor is ``2*num_experts`` when prefill
-    overlap is feasible else ``num_experts``); KV pages take whatever remains.
+    overlap is feasible, in slots and in bytes, else ``num_experts``); KV pages take whatever
+    remains. ``page_size`` (tokens per page) only feeds the error message.
     """
     assert per_expert_bytes > 0, "per_expert_bytes must be positive"
     assert cache_per_page > 0, "cache_per_page must be positive (owned-KV models unsupported here)"
@@ -81,6 +83,10 @@ def plan_cache_budget(
     assert hi >= lo, f"slot cap {hi} below the minimum {lo} slots"
 
     kv_reserve_bytes = kv_reserve_pages * cache_per_page
+    # Overlap is a prefill speedup, not a requirement: when its 2*num_experts floor leaves no
+    # room for the KV reserve (and the 2 pages asserted below), plan without it.
+    if overlap and 2 * num_experts * per_expert_bytes + max(kv_reserve_pages, 2) * cache_per_page > budget_bytes:
+        overlap, lo = False, num_experts
     # MoE-priority: reserve KV first, then experts greedily take the remaining budget.
     raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes
     moe_cache_size = max(lo, min(raw, hi))
@@ -93,10 +99,13 @@ def plan_cache_budget(
     # the reserve (or negative), yielding a plan that exceeds budget_bytes. Reject here so
     # --moe-cache-auto fails in arithmetic instead of OOMing in a later CUDA allocation.
     total = moe_cache_size * per_expert_bytes + num_pages * cache_per_page
+    fit_pages = max(0, (budget_bytes - moe_cache_size * per_expert_bytes) // cache_per_page)
     assert total <= budget_bytes, (
         f"cache budget too small: minimum plan (moe={moe_cache_size} slots, "
-        f"kv={num_pages} pages) needs {total} B > budget {budget_bytes} B "
-        "(raise memory_ratio, lower kv_reserve_tokens, or free GPU memory)"
+        f"kv={num_pages} pages) needs {total / 2**30:.2f} GiB > budget {budget_bytes / 2**30:.2f} GiB; "
+        f"beside {moe_cache_size} slots at most {fit_pages} KV pages ({fit_pages * page_size} tokens) fit. "
+        "Lower --kv-reserve-tokens, use a smaller --kv-cache-dtype, raise --memory-ratio, "
+        "or set --moe-cache-size and --max-seq-len-override explicitly"
     )
     assert num_pages > 1, "not enough memory for KV cache after MoE allocation"
     return moe_cache_size, num_pages, overlap
@@ -137,4 +146,5 @@ def resolve_moe_cache_auto(
         prefill_overlap=prefill_overlap,
         kv_reserve_pages=kv_reserve_pages,
         max_slots=max_slots,
+        page_size=page_size,
     )
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index f97feaa732..cccaa140fb 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -836,6 +836,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         size_cap = sum(p.max_num_slots for p in pools)
         if config.moe_cache_auto:
             size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap)
+            if config.moe_prefill_overlap and not overlap:
+                logger.info_rank0(
+                    f"--moe-cache-auto: prefill overlap disabled, its {2 * config.model_config.num_experts} "
+                    "slot floor does not fit beside the KV reserve"
+                )
             object.__setattr__(config, "moe_cache_size", size)
             object.__setattr__(config, "moe_prefill_overlap", overlap)
             if config.num_page_override is None:
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 00b314828f..1479b148b0 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -80,6 +80,43 @@ def test_budget_too_small_for_min_moe_plus_reserve_raises():
         )  # min moe = 4 slots (400 B) + reserve (10 pages = 100 B) = 500 B > 300 B budget
 
 
+def test_overlap_floor_that_does_not_fit_falls_back_to_num_experts():
+    # #4: 2*num_experts slots (800 B) + the KV reserve (100 B) exceed the 700 B budget, but
+    # num_experts slots do fit. Plan without overlap instead of refusing to start.
+    size, pages, overlap = plan_cache_budget(
+        budget_bytes=700, per_expert_bytes=100, cache_per_page=10,
+        num_experts=4, total_experts=50, prefill_overlap=True,
+        kv_reserve_pages=10, max_slots=50,
+    )
+    assert overlap is False
+    assert size == 6  # (700 - 100) // 100, above the num_experts floor
+    assert pages == 10
+    assert size * 100 + pages * 10 <= 700
+
+
+def test_overlap_kept_when_its_floor_fits_the_budget():
+    size, pages, overlap = plan_cache_budget(
+        budget_bytes=900, per_expert_bytes=100, cache_per_page=10,
+        num_experts=4, total_experts=50, prefill_overlap=True,
+        kv_reserve_pages=10, max_slots=50,
+    )
+    assert (size, pages, overlap) == (8, 10, True)
+
+
+def test_budget_too_small_message_names_what_fits_and_the_flags():
+    # num_experts slots (400 B) + 10 reserved pages (100 B) > 450 B: 5 pages fit beside them.
+    with pytest.raises(AssertionError) as err:
+        plan_cache_budget(
+            budget_bytes=450, per_expert_bytes=100, cache_per_page=10,
+            num_experts=4, total_experts=50, prefill_overlap=True,
+            kv_reserve_pages=10, max_slots=50, page_size=64,
+        )
+    msg = str(err.value)
+    assert "beside 4 slots at most 5 KV pages (320 tokens) fit" in msg
+    for flag in ("--kv-reserve-tokens", "--kv-cache-dtype", "--memory-ratio", "--moe-cache-size"):
+        assert flag in msg
+
+
 def test_prefill_overlap_false_is_honored():
     # Even when the cache could fit 2*num_experts, an explicit False stays False.
     size, pages, overlap = plan_cache_budget(

From 3afc08f6acf53db2b244c51bf3f39156e49a673d Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:59 -0400
Subject: [PATCH 142/249] merge fork #21: fix(triton): keep the float glm dsa
 sparse launches two-stage to fit sm_120 shared memory

---
 python/freetoken/kernel/triton/glm_dsa_sparse.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py
index d73ece8cf9..bac2cd7eed 100644
--- a/python/freetoken/kernel/triton/glm_dsa_sparse.py
+++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py
@@ -498,9 +498,9 @@ def glm_dsa_sparse_attn(
             D_V=d_v, D_R=d_r,
             BLOCK_H=BLOCK_H, BLOCK_T=block_t,
             HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits,
-            # The 512-wide latent accumulator exceeds the 99 KiB shared-memory
-            # limit of consumer Blackwell GPUs with a two-stage pipeline.
-            num_warps=4, num_stages=1,
+            # Consumer Blackwell caps kernel shared memory at 99 KiB. Single-stage, the float path
+            # needs 100 KiB; two stages (its setting before quantized KV) fit.
+            num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2,
         )
         grid2 = (m, b, h)
         _glm_dsa_merge_kernel[grid2](
@@ -527,7 +527,8 @@ def glm_dsa_sparse_attn(
         D_V=d_v, D_R=d_r,
         BLOCK_H=BLOCK_H, BLOCK_T=block_t,
         HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4,
-        num_warps=4, num_stages=1,
+        # same shared-memory split as the split-k launch above
+        num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2,
     )
     return o
 

From 9d60f844b6a5fec5b569de2f4b3c449dce18065c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:59 -0400
Subject: [PATCH 143/249] merge fork #23 (upstream FlashML-org/FreeToken#17
 rebased): fix(server): resolve Hugging Face repo ids to a local checkpoint
 dir

# Conflicts:
#	python/freetoken/utils/hf.py
---
 python/freetoken/server/args.py            | 20 +++--
 python/freetoken/utils/__init__.py         |  2 +
 python/freetoken/utils/hf.py               | 45 ++++++++++-
 tests/server/test_model_path_resolution.py | 87 ++++++++++++++++++++++
 4 files changed, 148 insertions(+), 6 deletions(-)
 create mode 100644 tests/server/test_model_path_resolution.py

diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 8dc7718271..a3d1a5d588 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -988,16 +988,26 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
     if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag:
         kwargs["moe_cache_auto"] = True
 
-    if kwargs["model_source"] == "modelscope":
-        model_path = kwargs["model_path"]
-        if not os.path.isdir(model_path):
+    # Resolve a hub repo id ("org/model") to a local checkpoint directory. This has to
+    # happen here, after served_model_name and the parser cascade have read the repo id
+    # (they want "DeepSeek-V4-Flash-0731", not a snapshot hash) and before anything opens
+    # the checkpoint: config parsing reads model-specific files straight off disk, so a
+    # bare repo id reaches `open()` as a relative path and dies with a FileNotFoundError.
+    model_path = kwargs["model_path"]
+    if not os.path.isdir(model_path):
+        if kwargs["model_source"] == "modelscope":
             from modelscope import snapshot_download
 
             ignore_patterns = []
             if kwargs["use_dummy_weight"]:
                 ignore_patterns = ["*.bin", "*.safetensors", "*.pt", "*.ckpt"]
-            model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns)
-            kwargs["model_path"] = model_path
+            kwargs["model_path"] = snapshot_download(model_path, ignore_patterns=ignore_patterns)
+        else:
+            from freetoken.utils import download_hf_checkpoint
+
+            kwargs["model_path"] = download_hf_checkpoint(
+                model_path, dummy_weight=kwargs["use_dummy_weight"]
+            )
     del kwargs["model_source"]
 
     # "auto" (or an unspecified dtype) resolves to the checkpoint's dtype. Multimodal /
diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py
index af68bf5816..a1dd10b7c3 100644
--- a/python/freetoken/utils/__init__.py
+++ b/python/freetoken/utils/__init__.py
@@ -9,6 +9,7 @@
 )
 from .hf import (
     cached_load_hf_config,
+    download_hf_checkpoint,
     download_hf_weight,
     load_eos_token_ids,
     load_generation_sampling,
@@ -30,6 +31,7 @@
 
 __all__ = [
     "cached_load_hf_config",
+    "download_hf_checkpoint",
     "download_hf_weight",
     "load_eos_token_ids",
     "load_generation_sampling",
diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py
index 5a7b7e6935..ce2b574a26 100644
--- a/python/freetoken/utils/hf.py
+++ b/python/freetoken/utils/hf.py
@@ -2,6 +2,7 @@
 import functools
 import json
 import os
+import re
 from typing import Any, FrozenSet, Mapping
 
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -296,4 +297,46 @@ def download_hf_weight(model_path: str) -> str:
     except Exception as e:
         raise ValueError(
             f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}"
-        )
\ No newline at end of file
+        )
+
+
+# Weight formats FreeToken never loads -- it reads safetensors (plus GGUF for Gemma-4).
+# Plenty of repos still ship a duplicate .bin/.pth copy of the same tensors, so skipping
+# them can halve a first-time pull.
+_UNLOADABLE_WEIGHTS = ("*.bin", "*.pt", "*.pth", "*.h5", "*.msgpack", "*.onnx")
+
+
+# A hub repo id is "org/name" and nothing else. Anything with a leading separator, a
+# "~", or more than one "/" is a filesystem path the user got wrong, and turning that
+# into a hub lookup replaces "no such directory" with a confusing download error.
+_HUB_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+def looks_like_hub_repo_id(model_path: str) -> bool:
+    """Whether ``model_path`` should be resolved against the Hub rather than the disk."""
+    return bool(_HUB_REPO_ID.match(model_path))
+
+
+def download_hf_checkpoint(model_path: str, *, dummy_weight: bool = False) -> str:
+    """Resolve a Hugging Face repo id to a complete local checkpoint directory.
+
+    ``download_hf_weight`` fetches only ``*.safetensors``, which is all the weight
+    loaders need once the engine is running. Config parsing happens earlier and wants
+    more than that: the tokenizer, and the model-specific subdirs some architectures
+    read directly off disk -- DeepSeek-V4 fails without ``inference/config.json``
+    (see ``models/deepseek_v4/args.py``). So this pulls the whole snapshot.
+
+    A path that is already a local directory is returned untouched.
+    """
+    if os.path.isdir(model_path) or not looks_like_hub_repo_id(model_path):
+        return model_path
+    ignore_patterns = list(_UNLOADABLE_WEIGHTS)
+    if dummy_weight:
+        # Config and tokenizer are still needed; the tensors are not.
+        ignore_patterns += ["*.safetensors", "*.gguf"]
+    try:
+        return snapshot_download(model_path, ignore_patterns=ignore_patterns)
+    except Exception as e:
+        raise ValueError(
+            f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}"
+        )
diff --git a/tests/server/test_model_path_resolution.py b/tests/server/test_model_path_resolution.py
new file mode 100644
index 0000000000..76016e6869
--- /dev/null
+++ b/tests/server/test_model_path_resolution.py
@@ -0,0 +1,87 @@
+"""``--model org/name``: a hub repo id resolves to a local checkpoint directory.
+
+``ft serve --model deepseek-ai/DeepSeek-V4-Flash-0731`` used to pass the repo id through
+untouched as a filesystem path. Config parsing then opened it relative to cwd:
+
+    FileNotFoundError: No DeepSeek-V4 ModelArgs JSON found under
+    deepseek-ai/DeepSeek-V4-Flash-0731 (looked for inference/config.json)
+
+Only the modelscope branch ever downloaded, and ``download_hf_weight`` (used later, by the
+weight loaders) fetches ``*.safetensors`` alone -- so it would never have produced
+``inference/config.json`` either.
+
+These pin the resolve for both sources, and pin the two things that must keep reading the
+*repo id* rather than the resolved snapshot path: the served model name and the parser
+cascade. Resolving too early renames the served model to a snapshot hash.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+from unittest.mock import patch
+
+import pytest
+
+from freetoken.server.args import parse_args
+
+REPO_ID = "deepseek-ai/DeepSeek-V4-Flash-0731"
+SNAPSHOT = "/cache/hub/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01"
+
+# Explicit dtype and parsers keep parse_args off the network: all three otherwise read the
+# checkpoint config, which is exactly what cannot be loaded until the path resolves.
+BASE_ARGS = ["--dtype", "bfloat16", "--tool-call-parser", "llama3", "--reasoning-parser", "off"]
+
+
+def _parse(model_path: str, *extra: str):
+    return parse_args([*BASE_ARGS, "--model", model_path, *extra])[0]
+
+
+@pytest.fixture
+def hf_download():
+    with patch("freetoken.utils.hf.snapshot_download", return_value=SNAPSHOT) as m:
+        yield m
+
+
+def test_repo_id_resolves_to_local_snapshot(hf_download):
+    assert _parse(REPO_ID).model_path == SNAPSHOT
+    assert hf_download.call_args.args[0] == REPO_ID
+
+
+def test_served_model_name_keeps_the_repo_id(hf_download):
+    """Not the snapshot hash -- this is the id clients send in the `model` field."""
+    assert _parse(REPO_ID).served_model_name == "DeepSeek-V4-Flash-0731"
+
+
+def test_local_dir_is_never_downloaded(tmp_path, hf_download):
+    assert _parse(str(tmp_path)).model_path == str(tmp_path)
+    hf_download.assert_not_called()
+
+
+def test_formats_freetoken_cannot_load_are_skipped(hf_download):
+    """Many repos ship a duplicate .bin copy; the loaders only read safetensors/gguf."""
+    _parse(REPO_ID)
+    ignored = hf_download.call_args.kwargs["ignore_patterns"]
+    assert "*.bin" in ignored
+    assert "*.safetensors" not in ignored
+
+
+def test_dummy_weight_skips_the_tensors(hf_download):
+    """--dummy-weight still needs config + tokenizer, but not 156 GiB of weights."""
+    _parse(REPO_ID, "--dummy-weight")
+    ignored = hf_download.call_args.kwargs["ignore_patterns"]
+    assert "*.safetensors" in ignored
+    assert "*.gguf" in ignored
+
+
+def test_modelscope_source_still_uses_modelscope(monkeypatch):
+    """The isdir check moved out of the modelscope branch; it must still short-circuit."""
+    fake = types.ModuleType("modelscope")
+    calls = []
+    fake.snapshot_download = lambda repo, **kw: calls.append((repo, kw)) or SNAPSHOT
+    monkeypatch.setitem(sys.modules, "modelscope", fake)
+
+    with patch("freetoken.utils.hf.snapshot_download") as hf:
+        assert _parse(REPO_ID, "--model-source", "modelscope").model_path == SNAPSHOT
+        hf.assert_not_called()
+    assert calls[0][0] == REPO_ID

From 99ffa2b9375435fdffb613e3ef4114521415f9bd Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:17:59 -0400
Subject: [PATCH 144/249] merge fork #47 (upstream FlashML-org/FreeToken#393
 rebased): OpenAI protocol coverage

---
 docs/openai_api.md                        | 101 +++++
 python/freetoken/core.py                  |  33 +-
 python/freetoken/engine/sample.py         | 204 ++++++++--
 python/freetoken/message/frontend.py      |   3 +
 python/freetoken/message/tokenizer.py     |   4 +
 python/freetoken/scheduler/scheduler.py   |   9 +-
 python/freetoken/server/api_models.py     |  58 +++
 python/freetoken/server/api_server.py     |  10 +-
 python/freetoken/server/control_api.py    |  18 +
 python/freetoken/server/generation.py     |  58 ++-
 python/freetoken/server/openai_api.py     | 471 ++++++++++++++++------
 python/freetoken/server/stats.py          |  41 ++
 python/freetoken/tokenizer/detokenize.py  |  54 ++-
 python/freetoken/tokenizer/server.py      |  24 +-
 python/freetoken/tokenizer/tokenize.py    |   8 +-
 tests/engine/test_logits_processors.py    | 155 +++++++
 tests/server/test_openai_extras.py        | 383 ++++++++++++++++++
 tests/tokenizer/test_detokenize_extras.py |  81 ++++
 18 files changed, 1538 insertions(+), 177 deletions(-)
 create mode 100644 docs/openai_api.md
 create mode 100644 tests/engine/test_logits_processors.py
 create mode 100644 tests/server/test_openai_extras.py
 create mode 100644 tests/tokenizer/test_detokenize_extras.py

diff --git a/docs/openai_api.md b/docs/openai_api.md
new file mode 100644
index 0000000000..0b1cc9d93c
--- /dev/null
+++ b/docs/openai_api.md
@@ -0,0 +1,101 @@
+# OpenAI-compatible API: endpoints and parameters
+
+What `ft serve` accepts on its OpenAI-shaped routes, what it ignores, and what it refuses,
+next to the same surface in vLLM, SGLang and llama.cpp's server. "honoured" means the value
+changes the generation; "accepted" means the request is not rejected but the value has no
+effect; "400" means the server refuses the request with a clear message.
+
+## Endpoints
+
+| route | FreeToken | vLLM | SGLang | llama.cpp |
+|---|---|---|---|---|
+| `POST /v1/chat/completions` | yes | yes | yes | yes |
+| `POST /v1/completions` | yes | yes | yes | yes |
+| `GET /v1/models`, `GET /v1/models/{id}` | yes | yes | yes | yes |
+| `POST /v1/responses` (+ get, cancel) | yes | yes | yes | yes |
+| `POST /v1/messages` (Anthropic), `count_tokens` | yes | yes | no | yes |
+| `POST /tokenize`, `POST /detokenize` (also under `/v1/`) | yes | yes | yes | yes (`/tokenize`, `/detokenize`) |
+| `GET /metrics` (Prometheus) | yes | yes | yes | yes |
+| `GET /version` | yes | yes | yes | no (`/props`) |
+| `GET /health` | yes | yes | yes | yes |
+| `GET /v1/stats`, `/v1/cache/*`, `/v1/requests`, `/v1/admin/*` | yes (FreeToken control plane) | `/metrics` only | `/get_server_info`, `/flush_cache` | `/slots`, `/props`, `/metrics` |
+| `POST /v1/embeddings`, `/rerank`, `/score`, `/pooling` | no (generative models only) | embedding / reranker models | embedding / reranker models | yes, any model with pooling |
+| `POST /v1/audio/*`, `/v1/files`, `/v1/batches` | no | transcriptions, batch runner (offline) | transcriptions, files, batches | no |
+| `POST /infill`, `/apply-template`, `/completion` (non-OpenAI) | `suffix` on `/v1/completions`; `/tokenize` with `messages` renders the template | no | no | yes |
+
+## Chat completions: request
+
+| parameter | FreeToken | vLLM | SGLang | llama.cpp |
+|---|---|---|---|---|
+| `model`, `messages` (text + `image_url` parts) | honoured | honoured | honoured | honoured |
+| `max_tokens`, `max_completion_tokens` | honoured (default 32k) | honoured | honoured | `max_tokens` |
+| `temperature`, `top_p`, `top_k` | honoured | honoured | honoured | honoured |
+| `min_p` | honoured (prob. below `min_p` x top prob dropped) | honoured | honoured | honoured |
+| `presence_penalty`, `frequency_penalty` | honoured, over generated tokens | honoured | honoured | honoured |
+| `repetition_penalty` | honoured, over prompt + generated tokens (HF semantics) | honoured | honoured | `repeat_penalty` |
+| `logit_bias` | honoured (token id -> bias, clamped to [-100, 100]) | honoured | honoured | honoured |
+| `min_tokens` | honoured: no EOS / stop token before this many generated tokens | honoured | `min_new_tokens` | no |
+| `stop` (strings), `stop_token_ids` | honoured | honoured | honoured | `stop` only |
+| `include_stop_str_in_output` | honoured | honoured | `no_stop_trim` | no |
+| `skip_special_tokens` | honoured; default off here, the reasoning and tool parsers read the tags | honoured (default on) | honoured (default on) | n/a |
+| `n` | honoured, 1..16, one generation per choice (the prefix cache shares the prompt) | honoured | honoured | `n_cmpl` |
+| `stream`, `stream_options.include_usage` | honoured | honoured | honoured | honoured |
+| `logprobs`, `top_logprobs` | 400 on this route (see #224); `/v1/completions` `logprobs` 400 | honoured | honoured | `n_probs` |
+| `tools`, `tool_choice` (`auto` / `none` / `required` / named), `parallel_tool_calls` | honoured | honoured | honoured | honoured |
+| `reasoning_effort`, `thinking`, `chat_template_kwargs` | honoured | `chat_template_kwargs` | honoured | honoured |
+| `continue_final_message` | honoured (the last assistant message is continued, no generation prompt) | honoured | honoured | no |
+| `response_format` `json_object` / `json_schema`, guided decoding | 400 (no constrained decoding) | honoured (xgrammar / guidance) | honoured (xgrammar / outlines / llguidance) | `grammar`, `json_schema` |
+| `seed` | accepted, not honoured (one batched sampling kernel, no per-request generator) | honoured | `sampling_seed` | honoured |
+| `user`, `metadata`, `store`, `service_tier` | accepted, ignored | accepted, ignored | accepted, ignored | accepted, ignored |
+| `request_id` | honoured: becomes the response id | honoured | `rid` | `id_slot` (different meaning) |
+| `function_call` (legacy) | 400 | 400 | 400 | no |
+| `best_of`, `use_beam_search`, `length_penalty`, `prompt_logprobs`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly |
+| `ignore_eos` | honoured (benchmarking) | honoured | honoured | honoured |
+
+## Legacy completions: request
+
+`prompt` as a string or a list of strings (token-id prompts are refused); the same sampling
+parameters as above; plus:
+
+| parameter | FreeToken | vLLM | SGLang | llama.cpp |
+|---|---|---|---|---|
+| `echo` | honoured (the prompt leads the text; not with `logprobs`) | honoured | honoured | no |
+| `suffix` | honoured on models with fill-in-the-middle tokens (`<|fim_prefix|>` ... `<|fim_middle|>`): the prompt becomes prefix-suffix-middle | 400 | 400 | `/infill` |
+| `n` | honoured, prompt-major choice order (`prompt i`, sample `k` -> `index i*n+k`) | honoured | honoured | no |
+| `logprobs` | 400 (see #224 for the sampled-token logprobs) | honoured | honoured | `n_probs` |
+
+## Responses
+
+| field | FreeToken | notes |
+|---|---|---|
+| `id` | `chatcmpl-<uid>` / `cmpl-<uuid>`, or `<prefix>-<request_id>` when the client sent one | |
+| `created`, `model`, `object` | yes | |
+| `system_fingerprint` | `null` | vLLM and SGLang also send `null` |
+| `choices[].message.reasoning_content` | yes, when the model reasoned | vLLM `reasoning_content`, llama.cpp `reasoning_content` |
+| `choices[].message.tool_calls` | yes | |
+| `choices[].finish_reason` | `stop`, `length`, `tool_calls` | |
+| `usage.prompt_tokens`, `completion_tokens`, `total_tokens` | yes | for `n > 1` the prompt is counted once, completions summed |
+| `usage.prompt_tokens_details.cached_tokens` | with `--enable-cache-report` | vLLM, SGLang: always |
+| `usage.completion_tokens_details.reasoning_tokens` | yes, when the model emitted a reasoning end tag | vLLM, SGLang |
+| streaming: one `[DONE]`, an optional final usage chunk | yes | |
+| llama.cpp `timings` | no; `/v1/stats` and `/metrics` carry the server-side rates | |
+
+## Tokenize / detokenize
+
+```
+POST /tokenize      {"prompt": "text", "add_special_tokens": true, "return_token_strs": false}
+POST /tokenize      {"messages": [...], "add_generation_prompt": true, "chat_template_kwargs": {}}
+                    -> {"count": N, "max_model_len": L, "tokens": [...], "token_strs": [...]?}
+POST /detokenize    {"tokens": [...], "skip_special_tokens": false} -> {"prompt": "text"}
+```
+
+`messages` are rendered through the model's chat template exactly as a generation would
+render them, so `count` equals the `prompt_tokens` that generation reports.
+
+## Not implemented, on purpose
+
+- Constrained decoding (`response_format` json / json_schema, regex, grammar): needs a
+  grammar engine on the sampling path.
+- Prompt logprobs and `echo` with `logprobs`: need the prefill logits.
+- Per-request `seed`: the batched sampling kernel draws from one generator.
+- Embeddings, reranking, scoring, audio: not generation.
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 812f9a4349..a160beaf81 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -24,13 +24,45 @@ class SamplingParams:
     # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one
     # appears in the decoded output; the matched substring (and anything after) is trimmed.
     stop_strs: list[str] = field(default_factory=list)
+    # ---- logits processors (engine/sample.py), all off by default ----
+    # min_p: drop tokens whose probability is below min_p x the top probability.
+    min_p: float = 0.0
+    # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the
+    # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off).
     presence_penalty: float = 0.0
     frequency_penalty: float = 0.0
+    repetition_penalty: float = 1.0
+    # OpenAI logit_bias as [[token id, bias], ...] (clamped to [-100, 100] by the API). A list,
+    # not a dict: the messages cross process boundaries as msgpack, which refuses int map keys.
+    logit_bias: list[list[float]] | None = None
+    # No EOS / stop token before this many generated tokens (vLLM min_tokens, SGLang
+    # min_new_tokens). The scheduler fills min_tokens_stop_ids with the ids to mask.
+    min_tokens: int = 0
+    min_tokens_stop_ids: list[int] = field(default_factory=list)
+    # Extra token ids that end generation with finish_reason "stop" (vLLM stop_token_ids).
+    stop_token_ids: list[int] = field(default_factory=list)
+    # Keep the matched stop string in the output instead of trimming it (vLLM
+    # include_stop_str_in_output, SGLang no_stop_trim).
+    include_stop_str_in_output: bool = False
+    # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers
+    # consume the <think>/tool tags from the decoded text.
+    skip_special_tokens: bool = False
 
     @property
     def is_greedy(self) -> bool:
         return self.temperature <= 0.0 or self.top_k == 1
 
+    @property
+    def needs_logits_processing(self) -> bool:
+        return bool(
+            self.min_p > 0.0
+            or self.presence_penalty != 0.0
+            or self.frequency_penalty != 0.0
+            or self.repetition_penalty != 1.0
+            or self.logit_bias
+            or self.min_tokens > 0
+        )
+
 
 @dataclass(eq=False)
 class Req:
@@ -68,7 +100,6 @@ class Req:
     # handler must not free resources under an in-flight forward; it sets this flag and
     # _process_last_data frees the request when the batch drains (after copy_done.synchronize).
     aborted: bool = False
-    output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         assert self.input_ids.is_cpu
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index c6fb2f40ee..71dd79b844 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, List
 
 import torch
@@ -10,19 +10,114 @@
     from freetoken.core import Batch
 
 
+@dataclass
+class LogitsPlan:
+    """Inputs of the logits processors for the rows of one batch that asked for any
+    (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p).
+
+    Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and
+    token histories, applied on the device by ``apply_logits_processors`` before the
+    sampling kernel. ``None`` when no request in the batch uses a processor, so the
+    default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1,
+    in the order of ``rows``); ``rows`` maps them back to the batch.
+    """
+
+    rows: torch.Tensor  # [r] int64 batch rows
+    temps: torch.Tensor  # [r, 1] float32 temperatures (min_p works on probabilities)
+    presence: torch.Tensor | None = None  # [r, 1] float32
+    frequency: torch.Tensor | None = None  # [r, 1] float32
+    repetition: torch.Tensor | None = None  # [r, 1] float32 (1.0 = off)
+    # Generated ids per row, right-padded with vocab_size (a scratch column that is
+    # dropped), for the presence / frequency counts and the repetition mask.
+    hist_ids: torch.Tensor | None = None  # [r, L] int64
+    # Prompt ids per row, same padding: repetition_penalty also covers the prompt.
+    prompt_ids: torch.Tensor | None = None  # [r, P] int64
+    bias_rows: torch.Tensor | None = None  # [b] int64 local rows
+    bias_ids: torch.Tensor | None = None  # [b] int64
+    bias_vals: torch.Tensor | None = None  # [b] float32
+    # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must
+    # not sample yet.
+    min_rows: torch.Tensor | None = None  # [m] int64 local rows
+    min_ids: torch.Tensor | None = None  # [m, S] int64
+    min_p: torch.Tensor | None = None  # [r, 1] float32 (0 = off for that row)
+
+
 @dataclass
 class BatchSamplingArgs:
     temperatures: torch.Tensor | None
     top_k: torch.Tensor | None = None
     top_p: torch.Tensor | None = None
     greedy_mask: torch.Tensor | None = None
-    penalties: list[tuple[int, torch.Tensor, float, float]] = field(default_factory=list)
+    plan: LogitsPlan | None = None
 
 
 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    if device.type != "cuda":  # the unit tests build plans on the CPU
+        return torch.tensor(data, dtype=dtype, device=device)
     return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True)
 
 
+def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor:
+    """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the
+    histories can be long (prompt + output) and a pinned buffer must not be freed under
+    an in-flight non_blocking copy."""
+    width = max(1, max((len(r) for r in rows), default=1))
+    out = torch.full((len(rows), width), pad, dtype=torch.int64)
+    for i, r in enumerate(rows):
+        if r:
+            out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64)
+    return out.to(device)
+
+
+def apply_logits_processors(
+    logits: torch.Tensor, plan: LogitsPlan, vocab_size: int
+) -> torch.Tensor:
+    """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order:
+    repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then
+    min_p. Rows outside the plan are copied unchanged."""
+    out = logits.to(torch.float32, copy=True)
+    sub = out[plan.rows]  # [r, V] (advanced indexing copies)
+    r = sub.shape[0]
+    device = sub.device
+    neg_inf = float("-inf")
+
+    if plan.hist_ids is not None:
+        counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device)
+        counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32))
+        counts = counts[:, :vocab_size]
+        if plan.repetition is not None:
+            seen = counts > 0
+            if plan.prompt_ids is not None:
+                in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device)
+                in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool))
+                seen = seen | in_prompt[:, :vocab_size]
+            # HF / vLLM semantics: divide positive logits, multiply negative ones.
+            penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition)
+            sub = torch.where(seen, penalized, sub)
+        if plan.presence is not None:
+            sub = sub - plan.presence * (counts > 0).to(sub.dtype)
+        if plan.frequency is not None:
+            sub = sub - plan.frequency * counts
+
+    if plan.bias_ids is not None:
+        sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True)
+
+    if plan.min_ids is not None:
+        block = torch.zeros((plan.min_ids.shape[0], vocab_size + 1), dtype=torch.bool, device=device)
+        block.scatter_(1, plan.min_ids, torch.ones_like(plan.min_ids, dtype=torch.bool))
+        rows = sub[plan.min_rows]
+        sub[plan.min_rows] = torch.where(block[:, :vocab_size], torch.full_like(rows, neg_inf), rows)
+
+    if plan.min_p is not None:
+        probs = torch.softmax(sub / plan.temps, dim=-1)
+        keep = probs >= plan.min_p * probs.amax(dim=-1, keepdim=True)
+        keep = keep | (plan.min_p <= 0.0)  # rows without min_p stay untouched
+        sub = torch.where(keep, sub, torch.full_like(sub, neg_inf))
+
+    out[plan.rows] = sub
+    return out
+
+
 def sample_impl(
     logits: torch.Tensor,
     temperatures: torch.Tensor,
@@ -57,23 +152,88 @@ class Sampler:
     device: torch.device
     vocab_size: int
 
+    def _plan(self, batch: Batch, params) -> LogitsPlan | None:
+        """The LogitsPlan for this batch, or None when no request asks for a processor.
+        Token histories come from the host-side ``req.input_ids``; under overlap
+        scheduling the previous step's token may not be appended yet, so a penalty sees
+        the history one token late. That is the accepted cost of keeping the counts off
+        the device."""
+        MIN_T = 1e-6
+        rows = [i for i, p in enumerate(params) if p.needs_logits_processing]
+        if not rows:
+            return None
+        picked = [(batch.reqs[i], params[i]) for i in rows]
+        temps = [[max(0.0 if p.is_greedy else p.temperature, MIN_T)] for _, p in picked]
+        plan = LogitsPlan(
+            rows=make_device_tensor(rows, torch.int64, self.device),
+            temps=make_device_tensor(temps, torch.float32, self.device),
+        )
+        pad = self.vocab_size
+
+        want_hist = any(
+            p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0
+            for _, p in picked
+        )
+        if want_hist:
+            hist: list[list[int]] = []
+            prompts: list[list[int]] = []
+            want_prompt = False
+            for req, p in picked:
+                prompt_len = req.max_device_len - req.output_len
+                ids = req.input_ids
+                penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0
+                hist.append(ids[prompt_len:].tolist() if penalized else [])
+                if p.repetition_penalty != 1.0:
+                    want_prompt = True
+                    prompts.append(ids[:prompt_len].tolist())
+                else:
+                    prompts.append([])
+            plan.hist_ids = _padded_rows(hist, pad, self.device)
+            if want_prompt:
+                plan.prompt_ids = _padded_rows(prompts, pad, self.device)
+            if any(p.presence_penalty != 0.0 for _, p in picked):
+                plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device)
+            if any(p.frequency_penalty != 0.0 for _, p in picked):
+                plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device)
+            if any(p.repetition_penalty != 1.0 for _, p in picked):
+                plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device)
+
+        bias_rows: list[int] = []
+        bias_ids: list[int] = []
+        bias_vals: list[float] = []
+        for local, (_, p) in enumerate(picked):
+            for tid, val in p.logit_bias or ():
+                if 0 <= int(tid) < self.vocab_size:
+                    bias_rows.append(local)
+                    bias_ids.append(int(tid))
+                    bias_vals.append(float(val))
+        if bias_ids:
+            plan.bias_rows = make_device_tensor(bias_rows, torch.int64, self.device)
+            plan.bias_ids = make_device_tensor(bias_ids, torch.int64, self.device)
+            plan.bias_vals = make_device_tensor(bias_vals, torch.float32, self.device)
+
+        min_rows: list[int] = []
+        min_ids: list[list[int]] = []
+        for local, (req, p) in enumerate(picked):
+            if p.min_tokens > 0 and p.min_tokens_stop_ids:
+                generated = len(req.input_ids) - (req.max_device_len - req.output_len)
+                if generated < p.min_tokens:
+                    min_rows.append(local)
+                    min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size])
+        if min_rows:
+            plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device)
+            plan.min_ids = _padded_rows(min_ids, pad, self.device)
+
+        if any(p.min_p > 0.0 for _, p in picked):
+            plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device)
+        return plan
+
     def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
         is_greedy = [p.is_greedy for p in params]
-        penalties = []
-        for row, req in enumerate(batch.reqs):
-            p = req.sampling_params
-            if not (p.presence_penalty or p.frequency_penalty) or not req.can_decode:
-                continue
-            if req.output_token_counts is None:
-                req.output_token_counts = torch.zeros(
-                    self.vocab_size, dtype=torch.int32, device=self.device
-                )
-            penalties.append(
-                (row, req.output_token_counts, p.presence_penalty, p.frequency_penalty)
-            )
+        plan = self._plan(batch, params)
         if all(is_greedy):
-            return BatchSamplingArgs(temperatures=None, penalties=penalties)
+            return BatchSamplingArgs(temperatures=None, plan=plan)
 
         MIN_P = MIN_T = 1e-6
         # Greedy outputs are selected explicitly in sample(); use neutral sampling
@@ -97,16 +257,14 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None
         )
         return BatchSamplingArgs(
-            temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, penalties=penalties
+            temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, plan=plan
         )
 
     @nvtx_annotate("Sampler")
     def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
         with torch.cuda.nvtx.range("Sampler"):
-            if args.penalties:
-                logits = logits.float().clone()
-                for row, counts, presence, frequency in args.penalties:
-                    logits[row] -= frequency * counts + presence * (counts > 0)
+            if args.plan is not None:
+                logits = apply_logits_processors(logits, args.plan, self.vocab_size)
             if args.temperatures is None:  # greedy sampling
                 tokens = torch.argmax(logits, dim=-1)
             else:
@@ -116,10 +274,4 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
                     # greedy rows must follow argmax's deterministic tie-breaking.
                     greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
                     tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
-            # Update on the sampling stream: overlapped scheduling can prepare the next
-            # batch before the previous token reaches Req.input_ids on the CPU.
-            for row, counts, _, _ in args.penalties:
-                counts.scatter_add_(
-                    0, tokens[row : row + 1].long(), counts.new_ones(1)
-                )
             return tokens
diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py
index 24725567b9..f3200650d8 100644
--- a/python/freetoken/message/frontend.py
+++ b/python/freetoken/message/frontend.py
@@ -56,6 +56,9 @@ class UserReply(BaseFrontendMsg):
     finish_reason: str | None = None
     # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence').
     matched_stop: str | None = None
+    # On the finished reply: tokens generated up to and including the reasoning end tag
+    # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens.
+    reasoning_tokens: int = 0
 
 
 @dataclass
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index d52d658e63..0d2116038f 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -36,6 +36,10 @@ class DetokenizeMsg(BaseTokenizerMsg):
     # The request's stop strings (None when it has none), so the detokenizer can hold back
     # a trailing partial-stop prefix instead of streaming it and then needing to retract.
     stop_strs: list[str] | None = None
+    # Keep the matched stop string in the output (SamplingParams.include_stop_str_in_output).
+    keep_stop_str: bool = False
+    # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens).
+    skip_special_tokens: bool = False
     # KV page-pool usage snapshot at this step (not-evictable used/total), passed
     # through to the frontend for the shell status bar. 0/0 for owned-KV models.
     kv_used_pages: int = 0
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index d057d3b908..f508a9dc89 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -353,7 +353,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                 hit_length = req.input_ids.numel() >= req.max_device_len
                 hit_eos = (
                     not req.sampling_params.ignore_eos and next_token in self.eos_token_ids
-                )
+                ) or next_token in req.sampling_params.stop_token_ids
                 matched_stop = (
                     self._match_stop_str(req)
                     if not hit_eos and req.sampling_params.stop_strs
@@ -379,6 +379,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         finish_reason=finish_reason,
                         matched_stop=matched_stop,
                         stop_strs=req.sampling_params.stop_strs or None,
+                        keep_stop_str=req.sampling_params.include_stop_str_in_output,
+                        skip_special_tokens=req.sampling_params.skip_special_tokens,
                     )
                 )
 
@@ -542,6 +544,11 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None:
                 logger.warning_rank0(
                     f"Adjust max_tokens to {max_output_len} for request {msg.uid}."
                 )
+            sp = msg.sampling_params
+            if sp.min_tokens > 0:
+                # min_tokens: the sampler masks these ids while the output is shorter.
+                sp.min_tokens = min(sp.min_tokens, sp.max_tokens)
+                sp.min_tokens_stop_ids = sorted(set(self.eos_token_ids) | set(sp.stop_token_ids))
             self.prefill_manager.add_one_req(msg)
         elif isinstance(msg, AbortBackendMsg):
             logger.debug_rank0("Aborting request %d", msg.uid)
diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 6c3ade0924..474af6b961 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -90,11 +90,30 @@ class ChatCompletionRequest(BaseModel):
     function_call: Any | None = None
     logit_bias: dict[str, float] | None = None
     response_format: dict[str, Any] | None = None
+    # The last message is an assistant prefix the model continues (no generation prompt).
+    continue_final_message: bool = False
+    # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ----
+    min_p: float | None = None
+    repetition_penalty: float | None = None
+    min_tokens: int = 0
+    stop_token_ids: list[int] | None = None
+    include_stop_str_in_output: bool = False
+    skip_special_tokens: bool | None = None
+    # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no
+    # per-request generator (seed); user / metadata / store are opaque to the server.
+    seed: int | None = None
+    user: str | None = None
+    # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid).
+    request_id: str | None = None
 
     @model_validator(mode="after")
     def _sync_max_completion_tokens(self) -> "ChatCompletionRequest":
         if self.max_completion_tokens is not None:
             self.max_tokens = self.max_completion_tokens
+        if self.continue_final_message and (
+            not self.messages or self.messages[-1].role != "assistant"
+        ):
+            raise ValueError("continue_final_message needs a final assistant message")
         return self
 
 
@@ -120,6 +139,19 @@ class CompletionRequest(BaseModel):
     suffix: str | None = None
     logit_bias: dict[str, float] | None = None
     response_format: dict[str, Any] | None = None
+    # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ----
+    min_p: float | None = None
+    repetition_penalty: float | None = None
+    min_tokens: int = 0
+    stop_token_ids: list[int] | None = None
+    include_stop_str_in_output: bool = False
+    skip_special_tokens: bool | None = None
+    # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no
+    # per-request generator (seed); user / metadata / store are opaque to the server.
+    seed: int | None = None
+    user: str | None = None
+    # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid).
+    request_id: str | None = None
 
     @model_validator(mode="after")
     def _sync_max_completion_tokens(self) -> "CompletionRequest":
@@ -128,6 +160,32 @@ def _sync_max_completion_tokens(self) -> "CompletionRequest":
         return self
 
 
+MAX_N = 16  # choices per request; each is a full generation, so keep the fan-out bounded
+
+
+class TokenizeRequest(BaseModel):
+    """``POST /tokenize`` (the vLLM / SGLang shape): a raw prompt, or chat messages that
+    are rendered through the model's chat template first."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = None
+    prompt: str | None = None
+    messages: list[Message] | None = None
+    add_special_tokens: bool = True
+    add_generation_prompt: bool = True
+    return_token_strs: bool = False
+    chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)
+
+
+class DetokenizeRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = None
+    tokens: list[int]
+    skip_special_tokens: bool = False
+
+
 class ModelCard(BaseModel):
     id: str
     object: str = "model"
diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index a9b16adfd7..f74750e691 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -361,16 +361,20 @@ async def stream_generate(self, uid: int):
         yield b"data: [DONE]\n\n"
         logger.debug("Finished streaming response for user %s", uid)
 
-    async def stream_with_cancellation(self, generator, request: Request, uid: int):
+    async def stream_with_cancellation(self, generator, request: Request, uid):
+        """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams
+        several generations into one response): a disconnect aborts every one."""
+        uids = list(uid) if isinstance(uid, (list, tuple)) else [uid]
         try:
             async for chunk in generator:
                 # detect if the client has disconnected
                 if await request.is_disconnected():
-                    logger.info("Client disconnected for user %s", uid)
+                    logger.info("Client disconnected for user %s", uids)
                     raise asyncio.CancelledError
                 yield chunk
         except asyncio.CancelledError:
-            asyncio.create_task(self.abort_user(uid))
+            for one in uids:
+                asyncio.create_task(self.abort_user(one))
             raise
 
     async def abort_user(self, uid: int):
diff --git a/python/freetoken/server/control_api.py b/python/freetoken/server/control_api.py
index 7158e4e5fa..cef632306e 100644
--- a/python/freetoken/server/control_api.py
+++ b/python/freetoken/server/control_api.py
@@ -68,6 +68,24 @@ async def list_requests(since: int = 0, limit: int = 100):
 
     from .stats import build_stats
 
+    @app.get("/metrics")
+    async def metrics():
+        """Prometheus text exposition of /v1/stats (the llama.cpp / vLLM convention)."""
+        from fastapi.responses import PlainTextResponse
+
+        from .stats import prometheus_text
+
+        doc = build_stats(
+            get_state(), request_ring.requests_p95_ms(), request_ring.requests_ttft_mean_ms()
+        )
+        return PlainTextResponse(prometheus_text(doc), media_type="text/plain; version=0.0.4")
+
+    @app.get("/version")
+    async def version():
+        from freetoken import __version__
+
+        return {"version": __version__}
+
     @app.get("/v1/stats")
     async def stats():
         doc = build_stats(
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 2e342ff899..c7d658688c 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -115,6 +115,7 @@ class GenDone:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    reasoning_tokens: int = 0
 
 
 GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone
@@ -130,6 +131,7 @@ class GenResult:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    reasoning_tokens: int = 0
 
 
 @dataclass
@@ -169,11 +171,19 @@ def resolve_sampling(
     model_sampling: dict[str, Any],
     stop: str | list[str] | None = None,
     default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    min_p: float | None = None,
     presence_penalty: float = 0.0,
     frequency_penalty: float = 0.0,
+    repetition_penalty: float | None = None,
+    logit_bias: dict[str, float] | dict[int, float] | None = None,
+    min_tokens: int = 0,
+    stop_token_ids: list[int] | None = None,
+    include_stop_str_in_output: bool = False,
+    skip_special_tokens: bool | None = None,
 ) -> SamplingParams:
     """Map a protocol's sampling fields onto the engine's neutral SamplingParams,
-    filling unspecified fields from the checkpoint's recommended defaults."""
+    filling unspecified fields from the checkpoint's recommended defaults. Range errors
+    raise ValueError (the adapters answer 400)."""
 
     def pick(value, key, framework):
         return value if value is not None else model_sampling.get(key, framework)
@@ -185,16 +195,50 @@ def pick(value, key, framework):
     # non-positive value is a client error.
     if max_tokens is not None and max_tokens < 1:
         raise ValueError(f"max_tokens must be at least 1, got {max_tokens}")
-    return SamplingParams(
+    resolved_min_p = float(pick(min_p, "min_p", 0.0))
+    resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0))
+    if not 0.0 <= resolved_min_p <= 1.0:
+        raise ValueError(f"min_p must be in [0, 1], got {resolved_min_p}")
+    for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)):
+        if not -2.0 <= float(value) <= 2.0:
+            raise ValueError(f"{name} must be in [-2, 2], got {value}")
+    if resolved_rep <= 0.0:
+        raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}")
+    if min_tokens < 0:
+        raise ValueError(f"min_tokens must be >= 0, got {min_tokens}")
+    bias: list[list[float]] | None = None
+    if logit_bias:
+        bias = []
+        for key, value in logit_bias.items():
+            try:
+                tid = int(key)
+            except (TypeError, ValueError):
+                raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None
+            if tid < 0:
+                raise ValueError(f"logit_bias token id must be >= 0, got {tid}")
+            bias.append([tid, max(-100.0, min(100.0, float(value)))])  # the OpenAI range
+    ids = [int(t) for t in (stop_token_ids or [])]
+    if any(t < 0 for t in ids):
+        raise ValueError("stop_token_ids must be non-negative token ids")
+    params = SamplingParams(
         ignore_eos=ignore_eos,
         max_tokens=default_max_tokens if max_tokens is None else max_tokens,
         temperature=pick(temperature, "temperature", 0.0),
         top_k=pick(top_k, "top_k", -1),
         top_p=pick(top_p, "top_p", 1.0),
         stop_strs=[s for s in stop_list if s],  # drop empty strings (would match everything)
-        presence_penalty=presence_penalty,
-        frequency_penalty=frequency_penalty,
+        min_p=resolved_min_p,
+        presence_penalty=float(presence_penalty),
+        frequency_penalty=float(frequency_penalty),
+        repetition_penalty=resolved_rep,
+        logit_bias=bias,
+        min_tokens=int(min_tokens),
+        stop_token_ids=ids,
+        include_stop_str_in_output=bool(include_stop_str_in_output),
     )
+    if skip_special_tokens is not None:
+        params.skip_special_tokens = bool(skip_special_tokens)
+    return params
 
 
 def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
@@ -607,6 +651,7 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
+    reasoning_tokens = 0
     pending = ""
     parse_tools = spec.parse_tools
     reasoning_parser = _make_reasoning_parser(spec, state)
@@ -702,6 +747,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
+        reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
         content_delta = ack.incremental_output
         if reasoning_parser is not None and content_delta:
             reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta)
@@ -784,6 +830,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
         matched_stop=engine_matched_stop, cached_tokens=cached_tokens,
+        reasoning_tokens=reasoning_tokens,
     )
 
 
@@ -794,6 +841,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
+    reasoning_tokens = 0
     engine_finish_reason: str | None = None
     engine_matched_stop: str | None = None
     async for ack in state.wait_for_ack(uid):
@@ -802,6 +850,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
+        reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
         full_content += ack.incremental_output
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
@@ -828,4 +877,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens=completion_tokens,
         matched_stop=engine_matched_stop,
         cached_tokens=cached_tokens,
+        reasoning_tokens=reasoning_tokens,
     )
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 99a4283549..f1bb0888f0 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -14,6 +14,9 @@
 from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS
 
 from .api_models import (
+    MAX_N,
+    DetokenizeRequest,
+    TokenizeRequest,
     ChatCompletionRequest,
     CompletionRequest,
     ModelCard,
@@ -68,20 +71,11 @@ def chat_request_to_genspec(
     thinking_type = _thinking_type(req)
     if req.reasoning_effort or thinking_type:
         ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type)
+    if req.continue_final_message:
+        ctk = {**ctk, "continue_final_message": True}
     return GenSpec(
         messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]),
-        sampling_params=resolve_sampling(
-            temperature=req.temperature,
-            top_k=req.top_k,
-            top_p=req.top_p,
-            max_tokens=req.max_tokens,
-            ignore_eos=req.ignore_eos,
-            model_sampling=model_sampling,
-            stop=req.stop,
-            default_max_tokens=default_max_tokens,
-            presence_penalty=req.presence_penalty,
-            frequency_penalty=req.frequency_penalty,
-        ),
+        sampling_params=_resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens),
         chat_template_kwargs=ctk,
         template_tools=_tools_for_template(req),
         parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None),
@@ -133,6 +127,16 @@ async def v1_completions(req: CompletionRequest, request: Request):
             return gate
         return await handle_completion(req, request, state, get_model_sampling())
 
+    @app.post("/tokenize")
+    @app.post("/v1/tokenize")
+    async def tokenize(req: TokenizeRequest):
+        return await handle_tokenize(req, get_state())
+
+    @app.post("/detokenize")
+    @app.post("/v1/detokenize")
+    async def detokenize(req: DetokenizeRequest):
+        return await handle_detokenize(req, get_state())
+
     @app.get("/v1/models")
     async def v1_models():
         state = get_state()
@@ -157,15 +161,13 @@ async def handle_chat_completion(
 ):
     if req.function_call is not None:
         return create_error_response("function_call is not supported; use tools/tool_choice instead")
-    if req.logit_bias is not None:
-        return create_error_response("logit_bias is not supported")
     if _response_format_unsupported(req.response_format):
         return create_error_response(
             "response_format json_object/json_schema is not supported (no constrained decoding)",
             param="response_format",
         )
-    if req.n != 1:
-        return create_error_response("Only n=1 is supported", param="n")
+    if not 1 <= req.n <= MAX_N:
+        return create_error_response(f"n must be between 1 and {MAX_N}", param="n")
     # Case/whitespace and the "off" disable synonym stay accepted here because
     # effort_toggle_kwargs normalizes and honors them downstream.
     effort = req.reasoning_effort.strip().lower() if isinstance(req.reasoning_effort, str) else None
@@ -198,43 +200,65 @@ async def handle_chat_completion(
         if err is not None:
             return create_error_response(str(err), code=err.code)
 
+    # n > 1: one generation per choice, submitted together (the prefix cache serves the
+    # shared prompt after the first prefill).
+    uids: list[int] = []
     try:
-        uid = await submit_generation(spec, state)
+        for _ in range(req.n):
+            uids.append(await submit_generation(spec, state))
     except GenerationError as exc:
+        for u in uids:  # choices already queued before a later one failed
+            await state.abort_user(u)
         return create_error_response(str(exc), code=exc.code)
+    uid = uids[0]
 
     if req.stream:
-        chunks = stream_chat_completion_chunks(uid, req, state, spec)
+        if req.n == 1:
+            chunks = stream_chat_completion_chunks(uid, req, state, spec)
+        else:
+            chunks = _merge_streams(
+                [
+                    stream_chat_completion_chunks(u, req, state, spec, index=i, terminal=False)
+                    for i, u in enumerate(uids)
+                ],
+                lambda: _chat_chunk(req, uid, []),
+                include_usage=bool(req.stream_options and req.stream_options.include_usage),
+                state=state,
+            )
         if request is not None:
-            chunks = state.stream_with_cancellation(chunks, request, uid)
+            chunks = state.stream_with_cancellation(chunks, request, uids)
         return StreamingResponse(chunks, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/chat/completions")
+        results = await asyncio.gather(
+            *(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids)
+        )
     except GenerationError as exc:
+        for u in uids:
+            await state.abort_user(u)
         return create_error_response(str(exc), code=exc.code)
-    message: dict[str, Any] = {"role": "assistant", "content": result.content}
-    if result.reasoning:
-        message["reasoning_content"] = result.reasoning
-    if result.tool_calls:
-        message["tool_calls"] = _tool_calls_to_openai(result.tool_calls)
-
+    choices: list[dict[str, Any]] = []
+    for index, result in enumerate(results):
+        message: dict[str, Any] = {"role": "assistant", "content": result.content}
+        if result.reasoning:
+            message["reasoning_content"] = result.reasoning
+        if result.tool_calls:
+            message["tool_calls"] = _tool_calls_to_openai(result.tool_calls)
+        choices.append({"index": index, "message": message, "finish_reason": result.finish_reason})
+
+    first = results[0]
     return {
-        "id": f"chatcmpl-{uid}",
+        "id": _response_id("chatcmpl", req, uid),
         "object": "chat.completion",
         "created": int(time.time()),
         "model": req.model,
-        "choices": [
-            {
-                "index": 0,
-                "message": message,
-                "finish_reason": result.finish_reason,
-            }
-        ],
+        "system_fingerprint": None,
+        "choices": choices,
         "usage": _usage(
-            result.prompt_tokens,
-            result.completion_tokens,
-            _reported_cached(state, result.cached_tokens),
+            first.prompt_tokens,
+            sum(r.completion_tokens for r in results),
+            _reported_cached(state, max(r.cached_tokens for r in results)),
+            reasoning_tokens=sum(r.reasoning_tokens for r in results),
         ),
     }
 
@@ -244,21 +268,28 @@ async def stream_chat_completion_chunks(
     req: ChatCompletionRequest,
     state: Any,
     spec: GenSpec | None = None,
+    index: int = 0,
+    terminal: bool = True,
 ) -> AsyncIterator[bytes]:
-    """Format generate_events() into the OpenAI chat.completion.chunk SSE stream."""
+    """Format generate_events() into the OpenAI chat.completion.chunk SSE stream.
+
+    ``index`` is the choice index (``n > 1`` runs one of these per choice); with
+    ``terminal=False`` the usage chunk and ``[DONE]`` are left to the merger, which reads
+    this stream's usage from the ``_StreamUsage`` yielded last."""
     if spec is None:
         spec = chat_request_to_genspec(req, {})
     yield _sse(
         _chat_chunk(
             req,
             uid,
-            [{"delta": {"role": "assistant", "content": ""}, "index": 0, "finish_reason": None}],
+            [{"delta": {"role": "assistant", "content": ""}, "index": index, "finish_reason": None}],
         )
     )
 
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
+    reasoning_tokens = 0
     tool_calls_sent = 0
     open_tool: dict[str, Any] | None = None
     events = generate_events(uid, spec, state, source="/v1/chat/completions")
@@ -279,7 +310,7 @@ async def stream_chat_completion_chunks(
                 _chat_chunk(
                     req,
                     uid,
-                    [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}],
+                    [{"delta": {"reasoning_content": ev.text}, "index": index, "finish_reason": None}],
                 )
             )
         elif isinstance(ev, ContentDelta):
@@ -287,7 +318,7 @@ async def stream_chat_completion_chunks(
                 _chat_chunk(
                     req,
                     uid,
-                    [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}],
+                    [{"delta": {"content": ev.text}, "index": index, "finish_reason": None}],
                 )
             )
         elif isinstance(ev, ToolCallStart):
@@ -307,7 +338,7 @@ async def stream_chat_completion_chunks(
                             "type": "function",
                             "function": {"name": ev.name, "arguments": ""},
                         }]},
-                        "index": 0, "finish_reason": None,
+                        "index": index, "finish_reason": None,
                     }],
                 )
             )
@@ -324,7 +355,7 @@ async def stream_chat_completion_chunks(
                                 "index": open_tool["index"],
                                 "function": {"arguments": ev.fragment},
                             }]},
-                            "index": 0, "finish_reason": None,
+                            "index": index, "finish_reason": None,
                         }],
                     )
                 )
@@ -348,7 +379,7 @@ async def stream_chat_completion_chunks(
                                         "index": open_tool["index"],
                                         "function": {"arguments": remainder},
                                     }]},
-                                    "index": 0, "finish_reason": None,
+                                    "index": index, "finish_reason": None,
                                 }],
                             )
                         )
@@ -360,7 +391,7 @@ async def stream_chat_completion_chunks(
                     yield _sse(
                         _chat_chunk(
                             req, uid,
-                            [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}],
+                            [{"delta": {"tool_calls": [delta]}, "index": index, "finish_reason": None}],
                         )
                     )
                 tool_calls_sent += 1
@@ -368,21 +399,18 @@ async def stream_chat_completion_chunks(
             prompt_tokens = ev.prompt_tokens
             completion_tokens = ev.completion_tokens
             cached_tokens = ev.cached_tokens
-            yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}]))
+            reasoning_tokens = ev.reasoning_tokens
+            yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": index, "finish_reason": ev.finish_reason}]))
 
+    usage = _usage(
+        prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens),
+        reasoning_tokens=reasoning_tokens,
+    )
+    if not terminal:
+        yield _StreamUsage(usage)
+        return
     if req.stream_options and req.stream_options.include_usage:
-        yield _sse(
-            {
-                "id": f"chatcmpl-{uid}",
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": req.model,
-                "choices": [],
-                "usage": _usage(
-                    prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)
-                ),
-            }
-        )
+        yield _sse({**_chat_chunk(req, uid, []), "usage": usage})
 
     yield b"data: [DONE]\n\n"
 
@@ -396,47 +424,63 @@ async def handle_completion(
     unsupported = _completion_unsupported_reason(req)
     if unsupported is not None:
         return create_error_response(unsupported)
-    try:  # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker
+    if not 1 <= req.n <= MAX_N:
+        return create_error_response(f"n must be between 1 and {MAX_N}", param="n")
+    try:  # surfaces an out-of-range value as a 400 rather than a 500 from the worker
         default_max_tokens = (
             getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS
         )
-        _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens)
+        sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens)
     except ValueError as exc:
-        return create_error_response(str(exc), param="max_tokens")
+        return create_error_response(str(exc))
 
     prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
     assert isinstance(prompts, list)
+    if req.suffix is not None:
+        markers = await _fim_markers(state)
+        if markers is None:
+            return create_error_response(
+                "suffix needs a model with fill-in-the-middle tokens "
+                "(<|fim_prefix|> / <|fim_suffix|> / <|fim_middle|>)",
+                param="suffix",
+            )
+        prompts = [_fim_prompt(p, req.suffix, markers) for p in prompts]
+    echo_texts = list(prompts) if req.echo else [""] * len(prompts)
+    if req.suffix is not None and req.echo:
+        # echo returns what the client sent, not the FIM-wrapped prompt
+        echo_texts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt)
+
+    # choices are ordered prompt-major: prompt i, sample k -> index i * n + k (OpenAI)
+    uids: list[int] = []
+    for prompt in prompts:
+        for _ in range(req.n):
+            uid = state.new_user()
+            await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling))
+            uids.append(uid)
+
     if req.stream:
-        if len(prompts) != 1:
-            return create_error_response("Streaming completions only support a single text prompt")
-        uid = state.new_user()
-        await state.send_one(
-            TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling(
-                req, model_sampling, default_max_tokens=default_max_tokens
-            ))
-        )
-        chunks = stream_completion_chunks(uid, req, state)
+        if len(uids) == 1:
+            chunks = stream_completion_chunks(uids[0], req, state, echo_text=echo_texts[0])
+        else:
+            chunks = _merge_streams(
+                [
+                    stream_completion_chunks(
+                        u, req, state, index=i, terminal=False, echo_text=echo_texts[i // req.n]
+                    )
+                    for i, u in enumerate(uids)
+                ],
+                lambda: _completion_chunk(req, uids[0], []),
+                include_usage=bool(req.stream_options and req.stream_options.include_usage),
+                state=state,
+            )
         if request is not None:
-            chunks = state.stream_with_cancellation(chunks, request, uid)
+            chunks = state.stream_with_cancellation(chunks, request, uids)
         return StreamingResponse(chunks, media_type="text/event-stream")
 
-    choices: list[dict[str, Any]] = []
-    prompt_tokens = 0
-    completion_tokens = 0
-    cached_tokens = 0
-    for index, prompt in enumerate(prompts):
-        uid = state.new_user()
-        await state.send_one(
-            TokenizeMsg(
-                uid=uid,
-                text=prompt,
-                sampling_params=_resolve_sampling(
-                    req, model_sampling, default_max_tokens=default_max_tokens
-                ),
-            )
-        )
+    async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
         text = ""
         finish_reason = "stop"
+        prompt_tokens = completion_tokens = cached_tokens = 0
         async for ack in state.wait_for_ack(uid):
             if getattr(ack, "error", None):
                 return create_error_response(ack.error)
@@ -447,77 +491,225 @@ async def handle_completion(
             if ack.finished:
                 finish_reason = getattr(ack, "finish_reason", None) or "stop"
                 break
-        choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None})
+        return {
+            "index": index,
+            "text": echo_texts[index // req.n] + text,
+            "finish_reason": finish_reason,
+            "logprobs": None,
+            "_usage": (prompt_tokens, completion_tokens, cached_tokens),
+        }
+
+    drained = await asyncio.gather(*(_drain_one(u, i) for i, u in enumerate(uids)))
+    choices: list[dict[str, Any]] = []
+    prompt_tokens = completion_tokens = cached_tokens = 0
+    for choice in drained:
+        if isinstance(choice, JSONResponse):
+            return choice
+        p, c, cached = choice.pop("_usage")
+        # the prompt is counted once per prompt, not once per sample of it
+        if choice["index"] % req.n == 0:
+            prompt_tokens += p
+        completion_tokens += c
+        cached_tokens = max(cached_tokens, cached)
+        choices.append(choice)
 
     return {
-        "id": f"cmpl-{uuid.uuid4().hex}",
+        "id": _response_id("cmpl", req, uuid.uuid4().hex),
         "object": "text_completion",
         "created": int(time.time()),
         "model": req.model,
+        "system_fingerprint": None,
         "choices": choices,
         "usage": _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)),
     }
 
 
-async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]:
+def _completion_chunk(req: CompletionRequest, uid: Any, choices: list[dict[str, Any]]) -> dict[str, Any]:
+    return {
+        "id": _response_id("cmpl", req, uid),
+        "object": "text_completion.chunk",
+        "created": int(time.time()),
+        "model": req.model,
+        "choices": choices,
+    }
+
+
+async def stream_completion_chunks(
+    uid: int,
+    req: CompletionRequest,
+    state: Any,
+    index: int = 0,
+    terminal: bool = True,
+    echo_text: str = "",
+) -> AsyncIterator[bytes]:
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
     finish_reason = "stop"
+    if echo_text:
+        # echo: the prompt leads the stream (vLLM sends it in the first chunk)
+        yield _sse(_completion_chunk(req, uid, [{"text": echo_text, "index": index, "finish_reason": None, "logprobs": None}]))
     async for ack in state.wait_for_ack(uid):
         if getattr(ack, "error", None):
             yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}})
-            yield b"data: [DONE]\n\n"
+            if terminal:
+                yield b"data: [DONE]\n\n"
             return
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         if ack.incremental_output:
             yield _sse(
-                {
-                    "id": f"cmpl-{uid}",
-                    "object": "text_completion.chunk",
-                    "created": int(time.time()),
-                    "model": req.model,
-                    "choices": [
-                        {
-                            "text": ack.incremental_output,
-                            "index": 0,
-                            "finish_reason": None,
-                            "logprobs": None,
-                        }
-                    ],
-                }
+                _completion_chunk(
+                    req, uid,
+                    [{"text": ack.incremental_output, "index": index, "finish_reason": None, "logprobs": None}],
+                )
             )
         if ack.finished:
             finish_reason = getattr(ack, "finish_reason", None) or "stop"
             break
 
-    yield _sse(
-        {
-            "id": f"cmpl-{uid}",
-            "object": "text_completion.chunk",
-            "created": int(time.time()),
-            "model": req.model,
-            "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}],
-        }
-    )
+    yield _sse(_completion_chunk(req, uid, [{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}]))
+    usage = _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens))
+    if not terminal:
+        yield _StreamUsage(usage)
+        return
     if req.stream_options and req.stream_options.include_usage:
-        yield _sse(
-            {
-                "id": f"cmpl-{uid}",
-                "object": "text_completion.chunk",
-                "created": int(time.time()),
-                "model": req.model,
-                "choices": [],
-                "usage": _usage(
-                    prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)
-                ),
-            }
+        yield _sse({**_completion_chunk(req, uid, []), "usage": usage})
+    yield b"data: [DONE]\n\n"
+
+
+class _StreamUsage:
+    """Yielded last by a non-terminal sub-stream: its usage dict for the merger."""
+
+    __slots__ = ("usage",)
+
+    def __init__(self, usage: dict[str, Any]) -> None:
+        self.usage = usage
+
+
+async def _merge_streams(
+    streams: list[AsyncIterator[Any]],
+    envelope: Callable[[], dict[str, Any]],
+    *,
+    include_usage: bool,
+    state: Any,
+) -> AsyncIterator[bytes]:
+    """Interleave the SSE chunks of several choice streams into one response (``n > 1``),
+    then one usage chunk (prompt counted once, completions summed) and one ``[DONE]``."""
+    queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue()
+    usages: dict[int, dict[str, Any]] = {}
+
+    async def pump(slot: int, stream: AsyncIterator[Any]) -> None:
+        try:
+            async for item in stream:
+                if isinstance(item, _StreamUsage):
+                    usages[slot] = item.usage
+                else:
+                    await queue.put((slot, item))
+        finally:
+            await queue.put((slot, None))
+
+    tasks = [asyncio.create_task(pump(i, st)) for i, st in enumerate(streams)]
+    try:
+        open_slots = len(streams)
+        while open_slots:
+            slot, item = await queue.get()
+            if item is None:
+                open_slots -= 1
+                continue
+            yield item
+    finally:
+        for task in tasks:
+            task.cancel()
+    if include_usage and usages:
+        first = usages[min(usages)]
+        merged = _usage(
+            first["prompt_tokens"],
+            sum(u["completion_tokens"] for u in usages.values()),
+            max(u.get("prompt_tokens_details", {}).get("cached_tokens", 0) for u in usages.values()),
+            reasoning_tokens=sum(
+                u.get("completion_tokens_details", {}).get("reasoning_tokens", 0) for u in usages.values()
+            ),
         )
+        yield _sse({**envelope(), "usage": merged})
     yield b"data: [DONE]\n\n"
 
 
+_FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")
+
+
+async def _fim_markers(state: Any) -> tuple[str, str, str] | None:
+    """The model's fill-in-the-middle marker strings when its vocabulary has them
+    (Qwen / DeepSeek-Coder style), else None."""
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is None:
+        return None
+    try:
+        manager = await asyncio.to_thread(build)
+        tok = manager.tokenizer
+        unk = getattr(tok, "unk_token_id", None)
+        for name in _FIM_TOKENS:
+            tid = tok.convert_tokens_to_ids(name)
+            if not isinstance(tid, int) or tid < 0 or tid == unk:
+                return None
+    except Exception:  # noqa: BLE001 -- no tokenizer here means no FIM
+        return None
+    return _FIM_TOKENS
+
+
+def _fim_prompt(prompt: str, suffix: str, markers: tuple[str, str, str]) -> str:
+    """The prefix-suffix-middle prompt of OpenAI's ``suffix`` (llama.cpp's /infill)."""
+    prefix_tok, suffix_tok, middle_tok = markers
+    return f"{prefix_tok}{prompt}{suffix_tok}{suffix}{middle_tok}"
+
+
+async def handle_tokenize(req: TokenizeRequest, state: Any):
+    """``POST /tokenize``: the vLLM / SGLang shape. A raw ``prompt`` is encoded as is; chat
+    ``messages`` are rendered through the model's chat template first (the exact prompt a
+    generation would tokenize)."""
+    if (req.prompt is None) == (req.messages is None):
+        return create_error_response("pass exactly one of prompt or messages")
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is None:
+        return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error")
+    try:
+        manager = await asyncio.to_thread(build)
+        if req.prompt is not None:
+            ids = await asyncio.to_thread(
+                manager.tokenizer.encode, req.prompt, add_special_tokens=req.add_special_tokens
+            )
+        else:
+            messages = render_messages([m.model_dump(exclude_none=True) for m in req.messages])
+            ctk = dict(req.chat_template_kwargs)
+            if not req.add_generation_prompt:
+                ctk["continue_final_message"] = True
+            msg = TokenizeMsg(uid=-1, text=messages, sampling_params=SamplingParams(), chat_template_kwargs=ctk)
+            text = await asyncio.to_thread(manager.render_prompt, msg)
+            ids = await asyncio.to_thread(manager.tokenizer.encode, text, add_special_tokens=False)
+    except ValueError as exc:
+        return create_error_response(str(exc))
+    body: dict[str, Any] = {
+        "count": len(ids),
+        "max_model_len": _model_context_length(state),
+        "tokens": [int(t) for t in ids],
+    }
+    if req.return_token_strs:
+        body["token_strs"] = manager.tokenizer.convert_ids_to_tokens(ids)
+    return body
+
+
+async def handle_detokenize(req: DetokenizeRequest, state: Any):
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is None:
+        return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error")
+    manager = await asyncio.to_thread(build)
+    text = await asyncio.to_thread(
+        manager.tokenizer.decode, req.tokens, skip_special_tokens=req.skip_special_tokens
+    )
+    return {"prompt": text}
+
+
 def create_error_response(
     message: str,
     status_code: int = 400,
@@ -552,11 +744,25 @@ def _resolve_sampling(
         model_sampling=model_sampling,
         stop=req.stop,
         default_max_tokens=default_max_tokens,
+        min_p=req.min_p,
         presence_penalty=req.presence_penalty,
         frequency_penalty=req.frequency_penalty,
+        repetition_penalty=req.repetition_penalty,
+        logit_bias=req.logit_bias,
+        min_tokens=req.min_tokens,
+        stop_token_ids=req.stop_token_ids,
+        include_stop_str_in_output=req.include_stop_str_in_output,
+        skip_special_tokens=req.skip_special_tokens,
     )
 
 
+def _response_id(prefix: str, req: Any, fallback: Any) -> str:
+    """``request_id`` from the client when given (vLLM request_id / SGLang rid), else the
+    server's own id."""
+    rid = getattr(req, "request_id", None)
+    return f"{prefix}-{rid}" if rid else f"{prefix}-{fallback}"
+
+
 def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None:
     if not _should_parse_tools(req):
         return None
@@ -617,10 +823,11 @@ def _tool_call_id(name: str | None, index: int) -> str:
 
 def _chat_chunk(req: ChatCompletionRequest, uid: int, choices: list[dict[str, Any]]) -> dict[str, Any]:
     return {
-        "id": f"chatcmpl-{uid}",
+        "id": _response_id("chatcmpl", req, uid),
         "object": "chat.completion.chunk",
         "created": int(time.time()),
         "model": req.model,
+        "system_fingerprint": None,
         "choices": choices,
     }
 
@@ -634,7 +841,9 @@ def _reported_cached(state: Any, cached_tokens: int) -> int:
     return cached_tokens if getattr(state.config, "enable_cache_report", False) else 0
 
 
-def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> dict[str, Any]:
+def _usage(
+    prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0
+) -> dict[str, Any]:
     usage: dict[str, Any] = {
         "prompt_tokens": prompt_tokens,
         "completion_tokens": completion_tokens,
@@ -644,6 +853,10 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -
     # disabled report and a 0-token hit serialize identically.
     if cached_tokens > 0:
         usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens}
+    # OpenAI's completion_tokens_details.reasoning_tokens: tokens up to and including the
+    # reasoning end tag, only when the model emitted one.
+    if reasoning_tokens > 0:
+        usage["completion_tokens_details"] = {"reasoning_tokens": reasoning_tokens}
     return usage
 
 
@@ -657,12 +870,6 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
         return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead"
     if req.logprobs is not None:
         return "logprobs is not supported"
-    if req.echo:
-        return "echo is not supported"
-    if req.suffix is not None:
-        return "suffix is not supported"
-    if req.logit_bias is not None:
-        return "logit_bias is not supported"
     if _response_format_unsupported(req.response_format):
         return "response_format json_object/json_schema is not supported (no constrained decoding)"
     return None
diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py
index 3d48ee96a0..c9d0212c32 100644
--- a/python/freetoken/server/stats.py
+++ b/python/freetoken/server/stats.py
@@ -174,3 +174,44 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict:
             "completion_tokens_total": tr.completion_tokens_total,
         },
     }
+
+
+
+def prometheus_text(doc: dict) -> str:
+    """The /v1/stats document as Prometheus text exposition (gauges and counters)."""
+    lines: list[str] = []
+
+    def gauge(name: str, value, help_text: str, kind: str = "gauge") -> None:
+        if value is None:
+            return
+        lines.append(f"# HELP {name} {help_text}")
+        lines.append(f"# TYPE {name} {kind}")
+        lines.append(f"{name} {value}")
+
+    model = doc.get("model") or {}
+    model_id = str(model.get("id") or model.get("name") or "").replace('"', "'")
+    lines.append("# HELP freetoken_info Server identity.")
+    lines.append("# TYPE freetoken_info gauge")
+    lines.append(f'freetoken_info{{instance_id="{doc.get("instance_id")}",model="{model_id}"}} 1')
+    gauge("freetoken_uptime_seconds", doc.get("uptime_s"), "Seconds since the server became ready.")
+    req = doc.get("requests") or {}
+    gauge("freetoken_requests_active", req.get("active"), "Requests in flight.")
+    gauge("freetoken_requests_completed_total", req.get("completed"), "Requests completed.", "counter")
+    gauge("freetoken_request_latency_p95_milliseconds", req.get("p95_ms"), "p95 request latency over the recent window.")
+    gauge("freetoken_ttft_mean_milliseconds", req.get("ttft_mean_ms"), "Mean time to first token over the recent window.")
+    gauge("freetoken_prompt_tokens_total", req.get("prompt_tokens_total"), "Prompt tokens processed.", "counter")
+    gauge("freetoken_completion_tokens_total", req.get("completion_tokens_total"), "Tokens generated.", "counter")
+    tp = doc.get("throughput") or {}
+    gauge("freetoken_decode_tokens_per_second", tp.get("decode_tps"), "Decode throughput, sliding window.")
+    gauge("freetoken_prefill_tokens_per_second", tp.get("prefill_tps"), "Prefill throughput, sliding window.")
+    for pool in ("kv", "swa"):
+        p = doc.get(pool)
+        if p:
+            gauge(f"freetoken_{pool}_pages_used", p.get("used_pages"), f"{pool} pages in use.")
+            gauge(f"freetoken_{pool}_pages_total", p.get("total_pages"), f"{pool} pages in the pool.")
+    mamba = doc.get("mamba")
+    if mamba:
+        gauge("freetoken_mamba_slots_used", mamba.get("used_slots"), "GDN state slots in use.")
+        gauge("freetoken_mamba_slots_total", mamba.get("total_slots"), "GDN state slots in the pool.")
+    gauge("freetoken_vram_bytes", doc.get("vram_bytes"), "GPU memory in use by the engine.")
+    return "\n".join(lines) + "\n"
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index d138ac1a68..27dd3b8034 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -71,11 +71,16 @@ class DecodeStatus:
     read_offset: int  # length of read ids
     surr_offset: int  # length of surr ids
     sent_offset: int  # length of sent out string
+    skip_special: bool = False  # decode this request with skip_special_tokens
+    think_end_at: int | None = None  # generated tokens up to and incl. the reasoning end tag
 
 
 class DetokenizeManager:
     def __init__(
-        self, tokenizer: PreTrainedTokenizerBase, eos_token_ids: FrozenSet[int] | None = None
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        eos_token_ids: FrozenSet[int] | None = None,
+        think_end_id: int | None = None,
     ) -> None:
         # uid -> DecodeStatus
         self.decode_map: Dict[int, DecodeStatus] = {}
@@ -83,6 +88,9 @@ def __init__(
         self.eos_token_ids = (
             eos_token_ids if eos_token_ids is not None else frozenset({tokenizer.eos_token_id})
         )
+        # The reasoning end tag's id (e.g. ``</think>``): tokens generated up to and
+        # including it are reported as usage.completion_tokens_details.reasoning_tokens.
+        self.think_end_id = think_end_id if think_end_id is None else int(think_end_id)
 
     def discard(self, uid: int) -> None:
         """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored
@@ -91,8 +99,26 @@ def discard(self, uid: int) -> None:
         self.decode_map.pop(uid, None)
 
     def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
+        return self.detokenize_with_meta(msgs)[0]
+
+    def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]:
+        """batch_decode honouring a per-request skip_special_tokens flag (one call per
+        distinct flag value, results back in order)."""
+        out: List[str | None] = [None] * len(ids)
+        for flag in (False, True):
+            idx = [i for i, f in enumerate(skip) if f == flag]
+            if idx:
+                texts = self.tokenizer.batch_decode([ids[i] for i in idx], skip_special_tokens=flag)
+                for i, t in zip(idx, texts, strict=True):
+                    out[i] = t
+        return [t if t is not None else "" for t in out]
+
+    def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]:
+        """Incremental text per message plus, for a finished message, its reasoning token
+        count (0 otherwise or when the model emitted no reasoning end tag)."""
         read_ids: List[List[int]] = []
         surr_ids: List[List[int]] = []
+        skip: List[bool] = []
         for msg in msgs:
             if msg.uid not in self.decode_map:
                 self.decode_map[msg.uid] = DecodeStatus(
@@ -101,17 +127,26 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
                     read_offset=0,
                     surr_offset=0,
                     sent_offset=0,
+                    skip_special=bool(getattr(msg, "skip_special_tokens", False)),
                 )
             s = self.decode_map[msg.uid]
             if not (msg.finished and msg.next_token in self.eos_token_ids):
                 s.decoded_ids.append(msg.next_token)
+            if (
+                self.think_end_id is not None
+                and s.think_end_at is None
+                and msg.next_token == self.think_end_id
+            ):
+                s.think_end_at = len(s.decoded_ids)
             read_ids.append(s.decoded_ids[s.surr_offset :])
             surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset])
+            skip.append(s.skip_special)
 
-        read_texts = self.tokenizer.batch_decode(read_ids)
-        surr_texts = self.tokenizer.batch_decode(surr_ids)
+        read_texts = self._batch_decode(read_ids, skip)
+        surr_texts = self._batch_decode(surr_ids, skip)
 
         incremental_strs: List[str] = []
+        reasoning_tokens: List[int] = []
         for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True):
             s = self.decode_map[msg.uid]
             new_text = read_str[len(surr_str) :]
@@ -127,10 +162,16 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
 
             prev_sent = s.sent_offset
             if msg.finished:
-                # Generation is over: flush everything, trimming at the matched stop string.
+                # Generation is over: flush everything, trimming at the matched stop string
+                # (or keeping it, for include_stop_str_in_output).
                 if msg.matched_stop:
                     cut = output_str.find(msg.matched_stop)
-                    emit_end = cut if cut >= 0 else len(output_str)
+                    if cut < 0:
+                        emit_end = len(output_str)
+                    elif getattr(msg, "keep_stop_str", False):
+                        emit_end = cut + len(msg.matched_stop)
+                    else:
+                        emit_end = cut
                 else:
                     emit_end = len(output_str)
             elif msg.stop_strs:
@@ -141,7 +182,8 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
             incremental_output = output_str[prev_sent:emit_end] if emit_end > prev_sent else ""
             s.sent_offset = max(prev_sent, emit_end)
             incremental_strs.append(incremental_output)
+            reasoning_tokens.append((s.think_end_at or 0) if msg.finished else 0)
             if msg.finished:
                 del self.decode_map[msg.uid]
 
-        return incremental_strs
+        return incremental_strs, reasoning_tokens
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index d3365bfecd..2050fea3f6 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -81,6 +81,19 @@ def _send_generation_replies(
     _put_user_replies(send_frontend, terminal_errors)
 
 
+def _reasoning_end_id(tokenizer) -> int | None:
+    """The id of the reasoning end tag when the vocabulary has one (``</think>`` for the
+    Qwen / DeepSeek / GLM families), else None. Only used to count reasoning tokens."""
+    for tag in ("</think>", "<|end_of_thinking|>"):
+        try:
+            tid = tokenizer.convert_tokens_to_ids(tag)
+        except Exception:  # noqa: BLE001 -- a tokenizer without the tag answers unk or raises
+            continue
+        if isinstance(tid, int) and tid >= 0 and tid != getattr(tokenizer, "unk_token_id", -1):
+            return tid
+    return None
+
+
 def _tokenize_requests(
     tokenize_manager: Any,
     messages: List[TokenizeMsg],
@@ -151,7 +164,9 @@ def tokenize_worker(
 
     tokenize_manager = TokenizeManager(tokenizer, get_mm_processor(tokenizer_path, mm))
     detokenize_manager = DetokenizeManager(
-        tokenizer, load_eos_token_ids(tokenizer_path, tokenizer)
+        tokenizer,
+        load_eos_token_ids(tokenizer_path, tokenizer),
+        think_end_id=_reasoning_end_id(tokenizer),
     )
 
     if ack_queue is not None:
@@ -209,7 +224,7 @@ def tokenize_worker(
             )
             sampled_replies: List[UserReply] = []
             if len(detokenize_msg) > 0:
-                replies = detokenize_manager.detokenize(detokenize_msg)
+                replies, reasoning_counts = detokenize_manager.detokenize_with_meta(detokenize_msg)
                 sampled_replies = [
                     UserReply(
                         uid=msg.uid,
@@ -217,6 +232,7 @@ def tokenize_worker(
                         finished=msg.finished,
                         finish_reason=msg.finish_reason,
                         matched_stop=msg.matched_stop,
+                        reasoning_tokens=reasoning_tokens,
                         completion_tokens_delta=1,
                         kv_used_pages=msg.kv_used_pages,
                         kv_total_pages=msg.kv_total_pages,
@@ -226,7 +242,9 @@ def tokenize_worker(
                         swa_total_tokens=msg.swa_total_tokens,
                         gpu_mem_bytes=msg.gpu_mem_bytes,
                     )
-                    for msg, reply in zip(detokenize_msg, replies, strict=True)
+                    for msg, reply, reasoning_tokens in zip(
+                        detokenize_msg, replies, reasoning_counts, strict=True
+                    )
                 ]
 
             # An error reply and a client abort are both terminal for their uid, and neither
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index fc66391cf2..f512930dee 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -176,10 +176,16 @@ def _render(
         if tools is not None:
             chat_template_kwargs = {**chat_template_kwargs, "tools": tools}
         messages = _map_developer_role(messages, getattr(self.tokenizer, "chat_template", None))
+        # continue_final_message (an OpenAI-compatible extra, as in vLLM / SGLang): the last
+        # message is an assistant prefix the model must continue, so no generation prompt.
+        chat_template_kwargs = dict(chat_template_kwargs)
+        continue_final = bool(chat_template_kwargs.pop("continue_final_message", False))
+        if continue_final:
+            chat_template_kwargs["continue_final_message"] = True
         prompt = self.tokenizer.apply_chat_template(
             messages,
             tokenize=False,
-            add_generation_prompt=True,
+            add_generation_prompt=not continue_final,
             **chat_template_kwargs,
         )
         assert isinstance(prompt, str)
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
new file mode 100644
index 0000000000..f051dec51f
--- /dev/null
+++ b/tests/engine/test_logits_processors.py
@@ -0,0 +1,155 @@
+"""The logits processors behind min_p, the penalties, logit_bias and min_tokens: pure
+torch on the CPU, so the math is checked without a GPU or the sampling kernels."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+from freetoken.core import SamplingParams
+from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors
+
+V = 8
+CPU = torch.device("cpu")
+
+
+def _plan(rows, **fields) -> LogitsPlan:
+    plan = LogitsPlan(
+        rows=torch.tensor(rows, dtype=torch.int64),
+        temps=torch.ones((len(rows), 1), dtype=torch.float32),
+    )
+    for k, v in fields.items():
+        setattr(plan, k, v)
+    return plan
+
+
+def _ids(rows: list[list[int]]) -> torch.Tensor:
+    width = max(len(r) for r in rows)
+    out = torch.full((len(rows), width), V, dtype=torch.int64)
+    for i, r in enumerate(rows):
+        out[i, : len(r)] = torch.tensor(r)
+    return out
+
+
+def test_presence_and_frequency_penalties_count_generated_tokens():
+    logits = torch.zeros((2, V))
+    plan = _plan(
+        [1],
+        hist_ids=_ids([[3, 3, 5]]),
+        presence=torch.tensor([[0.5]]),
+        frequency=torch.tensor([[0.25]]),
+    )
+    out = apply_logits_processors(logits, plan, V)
+    assert torch.equal(out[0], logits[0])  # row outside the plan untouched
+    assert out[1, 3].item() == -(0.5 + 2 * 0.25)
+    assert out[1, 5].item() == -(0.5 + 0.25)
+    assert out[1, 0].item() == 0.0
+
+
+def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics():
+    logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]])
+    plan = _plan(
+        [0],
+        hist_ids=_ids([[0, 1]]),  # generated: 0 (positive), 1 (negative)
+        prompt_ids=_ids([[2, 3]]),  # prompt: 2 (positive), 3 (negative)
+        repetition=torch.tensor([[2.0]]),
+    )
+    out = apply_logits_processors(logits, plan, V)
+    assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0  # divided
+    assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0  # multiplied
+    assert out[0, 4].item() == 1.0  # unseen
+
+
+def test_logit_bias_adds_per_row():
+    logits = torch.zeros((2, V))
+    plan = _plan(
+        [0, 1],
+        bias_rows=torch.tensor([0, 1, 1]),
+        bias_ids=torch.tensor([2, 2, 6]),
+        bias_vals=torch.tensor([5.0, -100.0, 1.5]),
+    )
+    out = apply_logits_processors(logits, plan, V)
+    assert (
+        out[0, 2].item() == 5.0
+        and out[1, 2].item() == -100.0
+        and out[1, 6].item() == 1.5
+    )
+
+
+def test_min_tokens_masks_the_stop_ids_of_the_row_only():
+    logits = torch.zeros((2, V))
+    plan = _plan([0, 1], min_rows=torch.tensor([1]), min_ids=_ids([[0, 7]]))
+    out = apply_logits_processors(logits, plan, V)
+    assert out[1, 0].item() == float("-inf") and out[1, 7].item() == float("-inf")
+    assert out[1, 3].item() == 0.0
+    assert torch.isfinite(out[0]).all()
+
+
+def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability():
+    logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]]))
+    plan = _plan(
+        [0], min_p=torch.tensor([[0.15]])
+    )  # threshold 0.075, between 0.1 and 0.05
+    out = apply_logits_processors(logits, plan, V)
+    kept = torch.isfinite(out[0])
+    assert kept.tolist() == [True, True, True, False, False, False, False, False]
+    # a row with min_p 0 is left alone
+    plan0 = _plan([0], min_p=torch.tensor([[0.0]]))
+    assert torch.isfinite(apply_logits_processors(logits, plan0, V)[0]).all()
+
+
+def _req(
+    prompt: list[int], generated: list[int], max_tokens: int, **sp
+) -> SimpleNamespace:
+    ids = torch.tensor(prompt + generated, dtype=torch.int32)
+    return SimpleNamespace(
+        input_ids=ids,
+        output_len=max_tokens,
+        max_device_len=len(prompt) + max_tokens,
+        sampling_params=SamplingParams(**sp),
+    )
+
+
+def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
+    sampler = Sampler(CPU, V)
+    plain = SimpleNamespace(reqs=[_req([1, 2], [3], 4)])
+    assert sampler.prepare(plain).plan is None
+
+    # all greedy: the argmax path runs on the CPU, the sampling kernels need a GPU
+    batch = SimpleNamespace(
+        reqs=[
+            _req([1, 2], [3], 4),
+            _req(
+                [1, 2],
+                [3, 3],
+                4,
+                presence_penalty=1.0,
+                repetition_penalty=1.5,
+                logit_bias=[[5, 2.0]],
+            ),
+            _req([4], [], 4, min_tokens=2, min_tokens_stop_ids=[0, 7]),
+            _req(
+                [4], [6, 6], 4, min_tokens=2, min_tokens_stop_ids=[0]
+            ),  # already past min_tokens
+        ]
+    )
+    args = sampler.prepare(batch)
+    plan = args.plan
+    assert plan is not None
+    assert plan.rows.tolist() == [1, 2, 3]
+    assert plan.hist_ids[0].tolist()[:2] == [3, 3]  # row 1's generated tokens
+    assert plan.prompt_ids[0].tolist()[:2] == [1, 2]
+    assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5]
+    assert plan.min_rows.tolist() == [1]  # local row of the third request only
+    assert plan.min_ids[0].tolist() == [0, 7]
+
+    logits = torch.zeros((4, V))
+    out = apply_logits_processors(logits, plan, V)
+    assert (
+        out[1, 3].item() == -1.0
+    )  # presence on the repeated 3 (repetition on 0 is a no-op)
+    assert out[1, 5].item() == 2.0
+    assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf")
+    assert torch.isfinite(out[3]).all()
+    # greedy path: argmax over the processed logits
+    assert sampler.sample(logits, args)[1].item() == 5
diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py
new file mode 100644
index 0000000000..115112b9d0
--- /dev/null
+++ b/tests/server/test_openai_extras.py
@@ -0,0 +1,383 @@
+"""OpenAI protocol coverage: n > 1 fan-out, echo and suffix (FIM) on completions, the
+sampling extras reaching SamplingParams, request_id, usage details, and the tokenize /
+detokenize routes."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from types import SimpleNamespace
+
+import pytest
+from freetoken.message import TokenizeMsg, UserReply
+from freetoken.server.api_models import (
+    ChatCompletionRequest,
+    CompletionRequest,
+    DetokenizeRequest,
+    TokenizeRequest,
+)
+from freetoken.server.openai_api import (
+    _usage,
+    chat_request_to_genspec,
+    handle_chat_completion,
+    handle_completion,
+    handle_detokenize,
+    handle_tokenize,
+)
+from freetoken.server.stats import prometheus_text
+
+
+def run(coro):
+    return asyncio.run(coro)
+
+
+class FakeState:
+    """Distinct uids per generation; every uid gets the same two-ack reply."""
+
+    def __init__(self, text=("hel", "lo"), reasoning_tokens=0, tokenizer=None):
+        self.config = SimpleNamespace(
+            model_path="/models/unit-model",
+            served_model_name="unit-model",
+            tool_call_parser="llama3",
+            reasoning_parser=None,
+            max_seq_len=4096,
+        )
+        self.text = text
+        self.reasoning_tokens = reasoning_tokens
+        self._next = 100
+        self.sent: list[TokenizeMsg] = []
+        self.aborted: list[int] = []
+        if tokenizer is not None:
+            manager = SimpleNamespace(
+                tokenizer=tokenizer,
+                render_prompt=lambda msg: "RENDERED:" + json.dumps(msg.text),
+            )
+            self.frontend_tokenizer = lambda: manager
+
+    def new_user(self) -> int:
+        self._next += 1
+        return self._next
+
+    async def send_one(self, msg):
+        self.sent.append(msg)
+
+    async def abort_user(self, uid):
+        self.aborted.append(uid)
+
+    async def wait_for_ack(self, uid: int):
+        yield UserReply(
+            uid=uid, incremental_output="", finished=False, prompt_tokens_delta=5
+        )
+        for i, piece in enumerate(self.text):
+            last = i == len(self.text) - 1
+            yield UserReply(
+                uid=uid,
+                incremental_output=piece,
+                finished=last,
+                completion_tokens_delta=1,
+                finish_reason="stop" if last else None,
+                reasoning_tokens=self.reasoning_tokens if last else 0,
+            )
+
+
+def chat(**kwargs) -> ChatCompletionRequest:
+    body = {"model": "unit-model", "messages": [{"role": "user", "content": "hi"}]}
+    body.update(kwargs)
+    return ChatCompletionRequest(**body)
+
+
+def parse_sse(chunks: list[bytes]) -> list:
+    out = []
+    for chunk in chunks:
+        for line in chunk.decode().split("\n"):
+            if not line.startswith("data: "):
+                continue
+            payload = line[len("data: ") :]
+            out.append("[DONE]" if payload == "[DONE]" else json.loads(payload))
+    return out
+
+
+async def _collect(generator):
+    return [item async for item in generator]
+
+
+# ---- n > 1 ------------------------------------------------------------------------------
+
+
+def test_chat_n_returns_one_choice_per_generation_and_counts_the_prompt_once():
+    state = FakeState(reasoning_tokens=2)
+    body = run(handle_chat_completion(chat(n=2), None, state, {}))
+    assert [c["index"] for c in body["choices"]] == [0, 1]
+    assert all(c["message"]["content"] == "hello" for c in body["choices"])
+    assert len(state.sent) == 2 and len({m.uid for m in state.sent}) == 2
+    usage = body["usage"]
+    assert (
+        usage["prompt_tokens"] == 5
+        and usage["completion_tokens"] == 4
+        and usage["total_tokens"] == 9
+    )
+    assert usage["completion_tokens_details"] == {"reasoning_tokens": 4}
+    assert body["system_fingerprint"] is None
+
+
+def test_chat_n_stream_interleaves_choice_indices_and_ends_once():
+    state = FakeState()
+    response = run(
+        handle_chat_completion(
+            chat(n=2, stream=True, stream_options={"include_usage": True}),
+            None,
+            state,
+            {},
+        )
+    )
+    events = parse_sse(run(_collect(response.body_iterator)))
+    assert events.count("[DONE]") == 1 and events[-1] == "[DONE]"
+    indices = {c["index"] for ev in events[:-1] for c in ev.get("choices", [])}
+    assert indices == {0, 1}
+    usage_chunks = [ev for ev in events[:-1] if ev.get("usage")]
+    assert len(usage_chunks) == 1
+    assert usage_chunks[0]["usage"]["prompt_tokens"] == 5
+    assert usage_chunks[0]["usage"]["completion_tokens"] == 4
+    finals = [
+        c for ev in events[:-1] for c in ev.get("choices", []) if c.get("finish_reason")
+    ]
+    assert len(finals) == 2
+
+
+def test_n_out_of_range_is_400():
+    for n in (0, 17):
+        body = run(handle_chat_completion(chat(n=n), None, FakeState(), {}))
+        assert body.status_code == 400
+        assert json.loads(body.body)["error"]["param"] == "n"
+
+
+# ---- completions: echo, suffix, n ---------------------------------------------------------
+
+
+def test_completion_echo_prepends_the_prompt():
+    state = FakeState()
+    req = CompletionRequest(model="unit-model", prompt="Once upon", echo=True)
+    body = run(handle_completion(req, None, state, {}))
+    assert body["choices"][0]["text"] == "Once uponhello"
+
+
+def test_completion_echo_streams_the_prompt_first():
+    state = FakeState()
+    req = CompletionRequest(
+        model="unit-model", prompt="Once upon", echo=True, stream=True
+    )
+    response = run(handle_completion(req, None, state, {}))
+    events = parse_sse(run(_collect(response.body_iterator)))
+    assert events[0]["choices"][0]["text"] == "Once upon"
+    assert events[-1] == "[DONE]"
+
+
+class _FimTokenizer:
+    unk_token_id = -1
+
+    def convert_tokens_to_ids(self, name):
+        return {"<|fim_prefix|>": 10, "<|fim_suffix|>": 11, "<|fim_middle|>": 12}.get(
+            name, -1
+        )
+
+
+class _NoFimTokenizer:
+    unk_token_id = 3
+
+    def convert_tokens_to_ids(self, name):
+        return 3
+
+
+def test_completion_suffix_builds_a_fill_in_the_middle_prompt():
+    state = FakeState(tokenizer=_FimTokenizer())
+    req = CompletionRequest(model="unit-model", prompt="def f(", suffix="    return 1")
+    run(handle_completion(req, None, state, {}))
+    assert (
+        state.sent[0].text
+        == "<|fim_prefix|>def f(<|fim_suffix|>    return 1<|fim_middle|>"
+    )
+
+
+def test_completion_suffix_needs_fim_tokens():
+    state = FakeState(tokenizer=_NoFimTokenizer())
+    req = CompletionRequest(model="unit-model", prompt="def f(", suffix="x")
+    body = run(handle_completion(req, None, state, {}))
+    assert (
+        body.status_code == 400 and json.loads(body.body)["error"]["param"] == "suffix"
+    )
+
+
+def test_completion_n_orders_choices_prompt_major():
+    state = FakeState()
+    req = CompletionRequest(model="unit-model", prompt=["a", "b"], n=2)
+    body = run(handle_completion(req, None, state, {}))
+    assert [c["index"] for c in body["choices"]] == [0, 1, 2, 3]
+    assert [m.text for m in state.sent] == ["a", "a", "b", "b"]
+    assert body["usage"]["prompt_tokens"] == 10  # two prompts, each counted once
+    assert body["usage"]["completion_tokens"] == 8
+
+
+# ---- sampling extras --------------------------------------------------------------------
+
+
+def test_sampling_extras_reach_sampling_params():
+    spec = chat_request_to_genspec(
+        chat(
+            min_p=0.1,
+            repetition_penalty=1.2,
+            presence_penalty=0.5,
+            frequency_penalty=0.3,
+            logit_bias={"5": 150, "7": -3},
+            min_tokens=3,
+            stop_token_ids=[9],
+            include_stop_str_in_output=True,
+            skip_special_tokens=True,
+            seed=7,
+            user="u1",
+        ),
+        {},
+    )
+    sp = spec.sampling_params
+    assert sp.min_p == 0.1 and sp.repetition_penalty == 1.2
+    assert sp.presence_penalty == 0.5 and sp.frequency_penalty == 0.3
+    assert sp.logit_bias == [[5, 100.0], [7, -3.0]]  # clamped to the OpenAI range
+    assert sp.min_tokens == 3 and sp.stop_token_ids == [9]
+    assert sp.include_stop_str_in_output is True and sp.skip_special_tokens is True
+    assert sp.needs_logits_processing
+
+
+@pytest.mark.parametrize(
+    "field",
+    [
+        {"min_p": 1.5},
+        {"presence_penalty": 3.0},
+        {"frequency_penalty": -2.5},
+        {"repetition_penalty": 0.0},
+        {"logit_bias": {"x": 1.0}},
+        {"min_tokens": -1},
+        {"stop_token_ids": [-1]},
+    ],
+)
+def test_bad_sampling_extras_are_400(field):
+    body = run(handle_chat_completion(chat(**field), None, FakeState(), {}))
+    assert body.status_code == 400, field
+
+
+def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template():
+    spec = chat_request_to_genspec(
+        chat(
+            messages=[
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "Sure,"},
+            ],
+            continue_final_message=True,
+        ),
+        {},
+    )
+    assert spec.chat_template_kwargs["continue_final_message"] is True
+    with pytest.raises(ValueError):
+        chat(continue_final_message=True)
+
+
+# ---- ids and usage ------------------------------------------------------------------------
+
+
+def test_request_id_becomes_the_response_id():
+    state = FakeState()
+    body = run(handle_chat_completion(chat(request_id="abc-1"), None, state, {}))
+    assert body["id"] == "chatcmpl-abc-1"
+    response = run(
+        handle_chat_completion(chat(request_id="abc-2", stream=True), None, state, {})
+    )
+    events = parse_sse(run(_collect(response.body_iterator)))
+    assert events[0]["id"] == "chatcmpl-abc-2"
+
+
+def test_usage_details_appear_only_when_nonzero():
+    assert "completion_tokens_details" not in _usage(10, 5)
+    assert _usage(10, 5, reasoning_tokens=3)["completion_tokens_details"] == {
+        "reasoning_tokens": 3
+    }
+    assert _usage(10, 5, cached_tokens=4)["prompt_tokens_details"] == {
+        "cached_tokens": 4
+    }
+
+
+# ---- tokenize / detokenize / metrics -----------------------------------------------------
+
+
+class _Tokenizer:
+    def encode(self, text, add_special_tokens=True):
+        ids = [ord(c) for c in text]
+        return ([1] + ids) if add_special_tokens else ids
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(chr(i) for i in ids if not (skip_special_tokens and i == 1))
+
+    def convert_ids_to_tokens(self, ids):
+        return [chr(i) for i in ids]
+
+
+def test_tokenize_prompt_and_messages_and_detokenize():
+    state = FakeState(tokenizer=_Tokenizer())
+    body = run(handle_tokenize(TokenizeRequest(prompt="ab"), state))
+    assert body == {"count": 3, "max_model_len": 4096, "tokens": [1, 97, 98]}
+    body = run(
+        handle_tokenize(
+            TokenizeRequest(
+                prompt="ab", add_special_tokens=False, return_token_strs=True
+            ),
+            state,
+        )
+    )
+    assert body["tokens"] == [97, 98] and body["token_strs"] == ["a", "b"]
+    body = run(
+        handle_tokenize(
+            TokenizeRequest(messages=[{"role": "user", "content": "x"}]), state
+        )
+    )
+    assert body["count"] == len(
+        "RENDERED:" + json.dumps([{"role": "user", "content": "x"}])
+    )
+    bad = run(
+        handle_tokenize(
+            TokenizeRequest(prompt="a", messages=[{"role": "user", "content": "x"}]),
+            state,
+        )
+    )
+    assert bad.status_code == 400
+    body = run(
+        handle_detokenize(
+            DetokenizeRequest(tokens=[1, 104, 105], skip_special_tokens=True), state
+        )
+    )
+    assert body == {"prompt": "hi"}
+
+
+def test_metrics_exposition_lists_the_stats_document():
+    doc = {
+        "instance_id": "abc",
+        "model": {"id": "unit-model"},
+        "uptime_s": 12,
+        "kv": {"used_pages": 3, "total_pages": 10, "page_size": 64},
+        "mamba": {"used_slots": 1, "total_slots": 4},
+        "swa": None,
+        "vram_bytes": 123,
+        "throughput": {"decode_tps": 99.5, "prefill_tps": 1000.0},
+        "requests": {
+            "active": 1,
+            "completed": 7,
+            "p95_ms": 50,
+            "ttft_mean_ms": 20,
+            "prompt_tokens_total": 100,
+            "completion_tokens_total": 40,
+        },
+    }
+    text = prometheus_text(doc)
+    assert 'freetoken_info{instance_id="abc",model="unit-model"} 1' in text
+    assert "freetoken_requests_active 1" in text
+    assert "freetoken_requests_completed_total 7" in text
+    assert "freetoken_kv_pages_total 10" in text
+    assert "freetoken_mamba_slots_used 1" in text
+    assert "freetoken_decode_tokens_per_second 99.5" in text
+    assert "# TYPE freetoken_prompt_tokens_total counter" in text
diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py
new file mode 100644
index 0000000000..5617a10537
--- /dev/null
+++ b/tests/tokenizer/test_detokenize_extras.py
@@ -0,0 +1,81 @@
+"""Per-request detokenizer options: include_stop_str_in_output, skip_special_tokens, and
+the reasoning token count reported on the finished reply."""
+
+from __future__ import annotations
+
+from freetoken.message import DetokenizeMsg
+from freetoken.tokenizer.detokenize import DetokenizeManager
+
+EOS, SPECIAL, THINK_END = 0, 99, 7
+
+
+class _Tok:
+    """Token ids decode to letters; 99 is a special marker, 7 the reasoning end tag."""
+
+    eos_token_id = EOS
+    unk_token_id = -1
+
+    def decode(self, ids, skip_special_tokens=False):
+        out = []
+        for t in ids:
+            if t == SPECIAL:
+                if not skip_special_tokens:
+                    out.append("<s>")
+            elif t == THINK_END:
+                out.append("</think>")
+            else:
+                out.append(chr(ord("a") + t))
+        return "".join(out)
+
+    def batch_decode(self, ids_list, skip_special_tokens=False):
+        return [self.decode(ids, skip_special_tokens) for ids in ids_list]
+
+
+def _run(manager, uid, tokens, *, finished_kwargs=None, **msg_kwargs):
+    text = ""
+    meta = 0
+    for i, t in enumerate(tokens):
+        last = i == len(tokens) - 1
+        kwargs = dict(msg_kwargs)
+        if last and finished_kwargs:
+            kwargs.update(finished_kwargs)
+        strs, counts = manager.detokenize_with_meta(
+            [DetokenizeMsg(uid=uid, next_token=t, finished=last, **kwargs)]
+        )
+        text += strs[0]
+        meta = counts[0]
+    return text, meta
+
+
+def test_stop_string_is_trimmed_by_default_and_kept_on_request():
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    trimmed, _ = _run(
+        manager, 1, [1, 2, 3], stop_strs=["cd"], finished_kwargs={"matched_stop": "cd"}
+    )
+    assert trimmed == "b"
+    kept, _ = _run(
+        manager,
+        2,
+        [1, 2, 3],
+        stop_strs=["cd"],
+        keep_stop_str=True,
+        finished_kwargs={"matched_stop": "cd"},
+    )
+    assert kept == "bcd"
+
+
+def test_skip_special_tokens_is_per_request():
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    shown, _ = _run(manager, 1, [1, SPECIAL, 2, EOS])
+    hidden, _ = _run(manager, 2, [1, SPECIAL, 2, EOS], skip_special_tokens=True)
+    assert shown == "b<s>c"
+    assert hidden == "bc"
+
+
+def test_reasoning_tokens_count_up_to_the_end_tag():
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END)
+    text, reasoning = _run(manager, 1, [1, 2, THINK_END, 3, 4, EOS])
+    assert text == "bc</think>de"
+    assert reasoning == 3  # b, c and the tag itself
+    _, none = _run(manager, 2, [1, 2, EOS])
+    assert none == 0

From 77009dd9b7a9232fafa7de1dda95516c2ff0e407 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 17 Sep 2026 22:53:22 -0400
Subject: [PATCH 145/249] fix(models): skip the ftw layout check for a bare ftw
 without an hf config

Fork #19 asks the model family to check the stored fp8 split against the env flags; a
synthetic FTW written straight from tensors (upstream's test_ftw_weights) carries no
config.json, so there is no family to ask and the read proceeds as before.

Assisted-by: Claude Fable 5.1
(cherry picked from commit cab9be35b8e56fa3cf08a109017183691e2e69c5)
(cherry picked from commit e59e5441e44dac96d530f4e2cac34d2c6273411f)
---
 python/freetoken/models/weight.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py
index 514d73705b..19268ac619 100644
--- a/python/freetoken/models/weight.py
+++ b/python/freetoken/models/weight.py
@@ -224,9 +224,13 @@ def load_weight(
     # a text-only engine never built the tower, so its tensors are not even read
     keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES))
     if is_ftw_checkpoint(model_path):
-        spec = get_model_spec(cached_load_hf_config(model_path).architectures[0])
-        # a family whose stored tensors depend on load-time settings checks them against the index before the read
-        check = _model_override(spec, "check_ftw_weights")
+        # a family whose stored tensors depend on load-time settings checks them against the index before the
+        # read; a bare FTW without a readable HF config (synthetic test checkpoints) has no family to ask
+        try:
+            spec = get_model_spec(cached_load_hf_config(model_path).architectures[0])
+        except (OSError, ValueError):
+            spec = None
+        check = _model_override(spec, "check_ftw_weights") if spec is not None else None
         if check is not None:
             reader = FTWReader(model_path)
             names = {e["name"] for e in reader.entries("weight")}

From db1ffafc74270056e8ba8046d20c3f89cf2d35e7 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 17 Sep 2026 22:53:25 -0400
Subject: [PATCH 146/249] test(scheduler): give the rejection test's engine
 fake the pool size the admission clamp reads

Assisted-by: Claude Fable 5.1
(cherry picked from commit 4584957f3d63449a153097f380ffdac035a245dd)
(cherry picked from commit 0ab1764343228845db270505185dd2be7090086d)
---
 tests/scheduler/test_cost_accounting_core.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py
index 045010c6b4..116d4b88c7 100644
--- a/tests/scheduler/test_cost_accounting_core.py
+++ b/tests/scheduler/test_cost_accounting_core.py
@@ -132,7 +132,8 @@ def fail_prepare(_batch):
 
 def test_scheduler_rejection_emits_error_but_no_admission():
     scheduler = Scheduler.__new__(Scheduler)
-    scheduler.engine = SimpleNamespace(max_seq_len=4)
+    scheduler.engine = SimpleNamespace(max_seq_len=4, num_pages=4)
+    scheduler.config = SimpleNamespace(page_size=1)
     added = []
     scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append)
     sent = []

From bbd13370db7867d1ccc98afcc6c5e304bbc679b6 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 17 Sep 2026 22:41:19 -0400
Subject: [PATCH 147/249] test(server): abort every sample of an n>1
 non-streaming request on disconnect

Guards the #222 resolution against #393's fan-out: the disconnect watcher takes the
whole uid list, not the first sample's uid.

Assisted-by: Claude Fable 5.1
(cherry picked from commit 93ff96349235ea3ca1504733a2abbd6e6bcccf94)
(cherry picked from commit 69a5efa3300e93bd331b4394e54eea8cef82ad3f)
(cherry picked from commit 40e994121fb79a21ba0d7ffff70e4394a4498899)
---
 tests/server/test_stream_cancellation.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index e38fa56bfb..ba7da7c17b 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -294,3 +294,25 @@ def test_non_stream_connected_client_gets_result_without_abort(monkeypatch):
 
     assert result["choices"][0]["message"]["content"] == "Hi"
     assert state.aborted == []
+
+
+def test_chat_non_stream_disconnect_aborts_every_sample(monkeypatch):
+    # n > 1 fans out into one engine request per sample; a disconnect must stop all of them
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+
+    class _State(_ApiState):
+        def __init__(self):
+            super().__init__(acks=None)
+            self._next = 7
+
+        def new_user(self):
+            self._next += 1
+            return self._next
+
+    state = _State()
+    resp = asyncio.run(
+        openai_api.handle_chat_completion(_chat_req(n=2), _Request(disconnected=True), state, {})
+    )
+
+    assert resp.status_code == 499
+    assert sorted(state.aborted) == [8, 9]

From e9717bbee807c8b3aab3f0cce0b44444c8ab18c2 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 17 Sep 2026 23:04:45 -0400
Subject: [PATCH 148/249] fix(server): record requests with a pure asgi
 middleware so handlers can see a client disconnect

Starlette's @app.middleware("http") (BaseHTTPMiddleware) wraps the handler's receive channel
and never yields the client's http.disconnect while a non-streaming request is still
computing, so request.is_disconnected() stays False and the disconnect watcher from
FlashML-org/FreeToken#222 could not stop an abandoned request. Streaming requests only
worked because the socket write failed. The recorder is the one such middleware; as pure
ASGI it records the same row at response start.

Assisted-by: Claude Fable 5.1
(cherry picked from commit 8adde913c94e2f76b3962b8d10f10ded4b1cccc6)
(cherry picked from commit 202d7fec30b29a9a4320370b0c62426422422339)
---
 python/freetoken/server/api_server.py         | 79 ++++++++++++-------
 .../test_request_recorder_disconnect.py       | 75 ++++++++++++++++++
 2 files changed, 124 insertions(+), 30 deletions(-)
 create mode 100644 tests/server/test_request_recorder_disconnect.py

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 9d53df1bf0..0184e51230 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -506,38 +506,57 @@ def _served_model_name() -> str | None:
     return getattr(cfg, "served_model_name", None)
 
 
-@app.middleware("http")
-async def _record_request_middleware(request: Request, call_next):
+class _RecordRequestMiddleware:
     """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single-
     model server, so model = served_model_name; stream is inferred from the response media
-    type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None."""
-    path = request.url.path
-    if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith(
-        _TRACKED_REQUEST_PREFIXES
-    ):
-        return await call_next(request)
-    import time as _time
-
-    start = _time.monotonic()
-    response = await call_next(request)
-    duration_ms = int((_time.monotonic() - start) * 1000)
-    ctype = response.headers.get("content-type", "")
-    request_ring.record_request(
-        request_ring.RequestRecord(
-            ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()),
-            method=request.method,
-            path=path,
-            status=response.status_code,
-            model=_served_model_name(),
-            duration_ms=duration_ms,
-            ttft_ms=None,
-            prompt_tokens=None,
-            completion_tokens=None,
-            stream=ctype.startswith("text/event-stream"),
-            error=None,
-        )
-    )
-    return response
+    type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.
+
+    Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the
+    handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a
+    non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an
+    abandoned request runs to max_tokens. The recorder is the only such middleware here."""
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        path = scope.get("path", "") if scope["type"] == "http" else ""
+        if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES):
+            await self.app(scope, receive, send)
+            return
+        import time as _time
+
+        start = _time.monotonic()
+
+        async def send_recording(message):
+            # the base-class version recorded once the handler returned its response object,
+            # i.e. at response start, before a streaming body -- keep that timing
+            if message["type"] == "http.response.start":
+                ctype = ""
+                for name, value in message.get("headers", ()):
+                    if name.lower() == b"content-type":
+                        ctype = value.decode("latin-1")
+                request_ring.record_request(
+                    request_ring.RequestRecord(
+                        ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()),
+                        method=scope["method"],
+                        path=path,
+                        status=message["status"],
+                        model=_served_model_name(),
+                        duration_ms=int((_time.monotonic() - start) * 1000),
+                        ttft_ms=None,
+                        prompt_tokens=None,
+                        completion_tokens=None,
+                        stream=ctype.startswith("text/event-stream"),
+                        error=None,
+                    )
+                )
+            await send(message)
+
+        await self.app(scope, receive, send_recording)
+
+
+app.add_middleware(_RecordRequestMiddleware)
 
 
 class CacheRebuildRequest(BaseModel):
diff --git a/tests/server/test_request_recorder_disconnect.py b/tests/server/test_request_recorder_disconnect.py
new file mode 100644
index 0000000000..06ff17f99b
--- /dev/null
+++ b/tests/server/test_request_recorder_disconnect.py
@@ -0,0 +1,75 @@
+"""The request recorder must not hide a client disconnect from a non-streaming handler.
+
+Real sockets on purpose: the failure only shows through uvicorn's receive channel (a
+BaseHTTPMiddleware wrapper never yields ``http.disconnect`` while the handler computes), which
+TestClient does not model.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import socket
+import threading
+import time
+
+from fastapi import FastAPI, Request
+
+from freetoken.server.api_server import _RecordRequestMiddleware
+
+
+def _free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _serve(app: FastAPI, port: int):
+    import uvicorn
+
+    server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error"))
+    thread = threading.Thread(target=server.run, daemon=True)
+    thread.start()
+    deadline = time.monotonic() + 10
+    while not server.started:
+        assert time.monotonic() < deadline, "uvicorn did not start"
+        time.sleep(0.02)
+    return server
+
+
+def _post_and_hang_up(port: int) -> None:
+    body = b'{"prompt": "x"}'
+    with socket.create_connection(("127.0.0.1", port)) as sock:
+        sock.sendall(
+            b"POST /v1/completions HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n"
+            b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body
+        )
+        time.sleep(0.2)  # the handler is now polling
+
+
+def test_non_streaming_handler_sees_the_client_leave():
+    app = FastAPI()
+    app.add_middleware(_RecordRequestMiddleware)
+    seen: dict[str, float | None] = {}
+
+    @app.post("/v1/completions")
+    async def slow(request: Request):
+        t0 = time.monotonic()
+        for _ in range(60):
+            await asyncio.sleep(0.05)
+            if await request.is_disconnected():
+                seen["after"] = time.monotonic() - t0
+                return {"detected": True}
+        seen["after"] = None
+        return {"detected": False}
+
+    port = _free_port()
+    server = _serve(app, port)
+    try:
+        _post_and_hang_up(port)
+        deadline = time.monotonic() + 5
+        while "after" not in seen and time.monotonic() < deadline:
+            time.sleep(0.05)
+    finally:
+        server.should_exit = True
+    assert seen.get("after") is not None, "handler never saw the disconnect"
+    assert seen["after"] < 2.0

From 78a7b87fd13eb79c9380a0c30da17082524bd244 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sat, 19 Sep 2026 17:21:12 -0400
Subject: [PATCH 149/249] fix(kvcache): refuse the host kv tiers on an nvfp4
 pool whose block scales they cannot mirror

Upstream #499/#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4
pool (#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor
the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without
scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only
tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable.

Assisted-by: Claude Fable 5.1
(cherry picked from commit 4bb35fc34b0cb65b41584c65d7390b6c5f4f854a)
---
 python/freetoken/kvcache/host_prefix_tier.py |  3 ++
 python/freetoken/kvcache/kv_host_offload.py  | 14 +++++++
 tests/kvcache/test_host_tier_nvfp4_guard.py  | 44 ++++++++++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 tests/kvcache/test_host_tier_nvfp4_guard.py

diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 0265b126ef..3b2d2d5dc3 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -79,6 +79,9 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
         self._kv_scales: HostBank | None = None
         self._kv_free: list[int] = []
         if kv_pool is not None and kv_budget_bytes > 0:
+            from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool
+
+            _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST")
             codes = kv_pool._kv_buffer            # [2, L, P, page, H, D]
             scales = kv_pool._scale_buffer        # [2, L, P*page, H] | None (fp8)
             _, layers, _, page_size, heads, dim = codes.shape
diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py
index dbb49f4814..1f8c358426 100644
--- a/python/freetoken/kvcache/kv_host_offload.py
+++ b/python/freetoken/kvcache/kv_host_offload.py
@@ -33,10 +33,24 @@
 logger = init_logger(__name__)
 
 
+def _reject_block_scaled_pool(pool, knob: str) -> None:
+    """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool
+    (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror
+    copies, so a rehydrated page would be codes without scales. Refuse instead of decoding
+    garbage."""
+    if getattr(pool, "_block_scale_buffer", None) is not None:
+        raise ValueError(
+            f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps "
+            "block scales the host mirror does not carry -- leave the tier off or serve with "
+            "--kv-cache-dtype fp8"
+        )
+
+
 class KVHostOffloader:
     """GPU-slot LRU cache over the pinned host mirror of the logical KV page space."""
 
     def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None:
+        _reject_block_scaled_pool(pool, "--kv-host-pages")
         buf = pool._kv_buffer  # [2, L, P_gpu + 1, page_size, kv_heads, head_dim]
         assert buf.shape[0] == 2 and buf.is_contiguous()
         _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape
diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py
new file mode 100644
index 0000000000..a793fb294a
--- /dev/null
+++ b/tests/kvcache/test_host_tier_nvfp4_guard.py
@@ -0,0 +1,44 @@
+"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales;
+an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+from freetoken.kvcache.kv_host_offload import KVHostOffloader
+
+
+def _kv_pool(block_scaled: bool):
+    return SimpleNamespace(
+        _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8),  # [2, L, P+1, page, H, D]
+        _scale_buffer=None,
+        _block_scale_buffer=torch.zeros(1) if block_scaled else None,
+    )
+
+
+def _state_pool():
+    return SimpleNamespace(
+        conv_states=torch.zeros(2, 2, 4, 3),
+        recurrent_states=torch.zeros(2, 2, 2, 4),
+        slot_states={},
+    )
+
+
+def test_kv_host_offloader_refuses_an_nvfp4_pool():
+    with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"):
+        KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu"))
+
+
+def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool():
+    with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"):
+        HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2,
+                       kv_budget_bytes=1 << 16, device="cpu")
+
+
+def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool():
+    # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not
+    tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2,
+                          kv_budget_bytes=0, device="cpu")
+    assert tier.num_slots == 2 and not tier.has_kv

From f5e4f40f11a0026f53a5f10bc317517e7c6a8f28 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:57:03 -0400
Subject: [PATCH 150/249] test(scheduler): give the scheduler stubs the
 prefill-start map the metrics timing reads

Upstream #504 times the prefill span in Scheduler._prefill_start; the stubs from #222 and
#505's tests build the scheduler without it.

Assisted-by: Claude Opus 5.5
---
 tests/scheduler/test_abort_inflight_prefill.py | 1 +
 tests/scheduler/test_hybrid_cache_manager.py   | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py
index 85c217859f..f2cbd6e0d2 100644
--- a/tests/scheduler/test_abort_inflight_prefill.py
+++ b/tests/scheduler/test_abort_inflight_prefill.py
@@ -63,6 +63,7 @@ def _setup():
         config=SimpleNamespace(page_size=1),
         status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None),
         send_result=sent.extend,
+        _prefill_start={},
         _kv_usage_pages=cm.page_usage,
         _mamba_slot_usage=lambda: None,
         _swa_token_usage=lambda: None,
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index 282ab2ca43..07385d7046 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -213,7 +213,7 @@ def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, ta
         assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected])
 
     # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path.
-    stub = SimpleNamespace(cache_manager=cm, table_manager=tm)
+    stub = SimpleNamespace(cache_manager=cm, table_manager=tm, _prefill_start={})
     Scheduler._free_req_resources(stub, final)
     Scheduler._free_req_resources(stub, final)
     match = cm.match_req(_pend(prompt.tolist() + [999]))

From d5a94a879beaeeeeb74365ecc90821b6001076ba Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 19:59:17 -0400
Subject: [PATCH 151/249] test(engine): give the logprobs sampler fakes the
 logits-processing flag the sampler plan reads

Fork #47's sampler plan asks each request whether it needs logits processing; the
upstream #224 test fakes predate it.

Assisted-by: Claude Opus 5.5
---
 tests/engine/test_sample_logprobs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
index aabbdc5fe7..6bd204604d 100644
--- a/tests/engine/test_sample_logprobs.py
+++ b/tests/engine/test_sample_logprobs.py
@@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc
 
     batch = SimpleNamespace(
         reqs=[
-            SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)),
-            SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)),
+            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)),
+            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)),
         ]
     )
     args = sampler.prepare(batch)

From ec25773ec46ba9cf07a87c3cd35df84082895975 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:48:29 -0400
Subject: [PATCH 152/249] merge fork review fixes: 11 bugs found re-reviewing
 the 5080deploy stack

Assisted-by: Claude Opus 5.5

# Conflicts:
#	python/freetoken/server/openai_api.py
#	tests/server/test_stream_cancellation.py
---
 python/freetoken/engine/sample.py             |  4 +-
 python/freetoken/kvcache/host_prefix_tier.py  | 53 +++++++++--------
 python/freetoken/models/qwen4_exp/ple.py      | 13 +++--
 python/freetoken/scheduler/cache.py           |  7 ++-
 python/freetoken/server/anthropic_api.py      |  8 ++-
 python/freetoken/server/generation.py         | 13 ++++-
 python/freetoken/server/openai_api.py         | 15 +++--
 python/freetoken/server/responses_api.py      |  8 ++-
 python/freetoken/tokenizer/tokenize.py        | 10 +++-
 python/freetoken/utils/mp.py                  |  4 +-
 tests/engine/test_logits_processors.py        | 10 ++++
 tests/kvcache/test_host_tier_live_pool.py     | 53 +++++++++++++++++
 .../kvcache/test_host_tier_rehydrate_exact.py | 39 +++++++++++++
 tests/models/qwen4_exp/test_ple.py            | 13 +++++
 tests/server/test_openai_api.py               | 20 +++++++
 tests/server/test_sampling_validation.py      | 23 ++++++++
 tests/server/test_stream_cancellation.py      | 58 +++++++++++++++++++
 tests/tokenizer/test_tokenize.py              | 18 ++++++
 tests/utils/test_mp_large_frame.py            | 13 +++++
 19 files changed, 338 insertions(+), 44 deletions(-)
 create mode 100644 tests/kvcache/test_host_tier_live_pool.py
 create mode 100644 tests/kvcache/test_host_tier_rehydrate_exact.py
 create mode 100644 tests/server/test_sampling_validation.py
 create mode 100644 tests/utils/test_mp_large_frame.py

diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index f1984f7576..bca85c56fa 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -187,7 +187,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
                 hist.append(ids[prompt_len:].tolist() if penalized else [])
                 if p.repetition_penalty != 1.0:
                     want_prompt = True
-                    prompts.append(ids[:prompt_len].tolist())
+                    # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE);
+                    # send them to the scratch column, not past the scatter target
+                    prompts.append(ids[:prompt_len].clamp(max=pad).tolist())
                 else:
                     prompts.append([])
             plan.hist_ids = _padded_rows(hist, pad, self.device)
diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 3b2d2d5dc3..39c96dfd6d 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -59,13 +59,15 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
         tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states}
         tensors.update(state_pool.slot_states)
         self._order = list(tensors)
-        self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {}
+        # banks only: the pool tensors are re-read on every copy, since
+        # LinearStatePool.rebuild replaces them (a runtime cache rebuild)
+        self._banks: dict[str, HostBank] = {}
         for name, t in tensors.items():
             # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest]
             bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype)
             if self.device.type == "cuda" and torch.cuda.is_available():
                 bank.pin()
-            self._banks[name] = (t, bank)
+            self._banks[name] = bank
         self.num_slots = gdn_slots_host
         self._free = list(range(gdn_slots_host))
         self._in_use: set[int] = set()
@@ -102,40 +104,41 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
             torch.cuda.Stream(self.device) if self.device.type == "cuda" else None)
 
     # ---------------- GDN snapshot storage ----------------
+    def _pairs(self):
+        """(live pool tensor [layers, slots, ...], host bank) per state tensor."""
+        pool = self.pool
+        live = {"conv": pool.conv_states, "rec": pool.recurrent_states, **pool.slot_states}
+        return [(live[name], self._banks[name]) for name in self._order]
+
+    def _copy(self, copies) -> None:
+        if self.copy_stream is None:
+            for dst, src in copies:
+                dst.copy_(src)
+            return
+        # Order the copies after the caller's stream: the scheduler stream (itself made to wait
+        # on the engine stream) or the engine stream. Overlap launches the next batch before
+        # the previous one is drained, so a just-freed slot may still be written by it.
+        self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
+        with torch.cuda.stream(self.copy_stream):
+            for dst, src in copies:
+                dst.copy_(src, non_blocking=True)
+        self.copy_stream.synchronize()
+
     def write_gdn(self, gpu_slot: int) -> int:
         """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot.
-        Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do
-        engine garante leitura antes de qualquer reescrita do slot (ver spec)."""
+        Deve ser enfileirado ANTES de state_pool.free(gpu_slot)."""
         if not self._free:
             raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)")
         hs = self._free.pop()
         self._in_use.add(hs)
-        if self.copy_stream is not None:
-            with torch.cuda.stream(self.copy_stream):
-                for name in self._order:
-                    t, bank = self._banks[name]
-                    bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True)
-            self.copy_stream.synchronize()
-        else:
-            for name in self._order:
-                t, bank = self._banks[name]
-                bank.tensor[hs].copy_(t[:, gpu_slot])
+        self._copy([(bank.tensor[hs], t[:, gpu_slot]) for t, bank in self._pairs()])
         return hs
 
     def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None:
         """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller)."""
         if host_slot not in self._in_use:
             raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso")
-        if self.copy_stream is not None:
-            with torch.cuda.stream(self.copy_stream):
-                for name in self._order:
-                    t, bank = self._banks[name]
-                    t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True)
-            self.copy_stream.synchronize()
-        else:
-            for name in self._order:
-                t, bank = self._banks[name]
-                t[:, dst_gpu_slot].copy_(bank.tensor[host_slot])
+        self._copy([(t[:, dst_gpu_slot], bank.tensor[host_slot]) for t, bank in self._pairs()])
 
     def free_gdn(self, host_slot: int) -> None:
         self._in_use.discard(host_slot)
@@ -198,6 +201,7 @@ def _go():
                     del staged_s
 
         if self.copy_stream is not None:
+            self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
             with torch.cuda.stream(self.copy_stream):
                 _go()
             self.copy_stream.synchronize()
@@ -235,6 +239,7 @@ def _go():
                     del staged_s
 
         if self.copy_stream is not None:
+            self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
             with torch.cuda.stream(self.copy_stream):
                 _go()
             self.copy_stream.synchronize()
diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py
index 284e023bdd..44973bc0a6 100644
--- a/python/freetoken/models/qwen4_exp/ple.py
+++ b/python/freetoken/models/qwen4_exp/ple.py
@@ -49,6 +49,10 @@
 _FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH"
 
 
+def _capturing(device: torch.device) -> bool:
+    return device.type == "cuda" and torch.cuda.is_current_stream_capturing()
+
+
 def _fused_row_ids_enabled() -> bool:
     """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops."""
     return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False")
@@ -482,12 +486,13 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]:
 
         The fused kernel addresses the hash window through these instead of materializing the
         ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on)
-        so a captured replay reads a stable address instead of re-running the build; a build
-        that happens DURING capture is not cached, since its buffers live in the graph pool.
+        for eager forwards only: a capture always builds its own, which the graph recomputes
+        into its pool on every replay. A cached pair (the eager warmup forward right before
+        capture leaves one) would be baked in by address, then evicted and reused.
         """
         device = meta.input_ids.device
         num_tokens = meta.input_ids.numel()
-        capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing()
+        capturing = _capturing(device)
         # is_decode is part of the key, not just the shape: a decode of B requests and a
         # prefill of ONE B-token request are the same [T] and mean opposite things (one token
         # per request at offset 0 vs B offsets inside one request).
@@ -496,7 +501,7 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]:
             (num_tokens,) if meta.is_decode else tuple(meta.seq_lens),
             str(device),
         )
-        cached = self._token_index_cache.get(key)
+        cached = None if capturing else self._token_index_cache.get(key)
         if cached is not None:
             return cached
         if meta.is_decode:  # one token per request, each at offset 0
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 7e4aa8e695..4e6b91df88 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -624,8 +624,11 @@ def _rehydrate_snapshot(self, node) -> int | None:
         prefix, copy the snapshot back into a fresh pool slot and return the slot; None =
         keep walking (the match truncates as before)."""
         tier = self.host_tier
-        hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size)
-        if hit is None or hit.gdn_slot is None:
+        key = self.prefix_cache._collect_key(node)
+        hit = tier.lookup(key, self.page_size)
+        # lookup is longest-prefix; a shorter entry's state would resume this node's full
+        # depth with the tokens in between missing from the recurrence
+        if hit is None or hit.gdn_slot is None or hit.host_len != len(key):
             return None
         pool = self.linear_state_pool
         if pool.num_free_slots < 1:
diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 385c87f3f1..4cb4b34822 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -55,6 +55,7 @@
     submit_generation,
     with_keepalive,
 )
+from .openai_api import _await_watching_disconnect
 from .request_logger import log_request
 
 # Emit a protocol-native `ping` event after this many seconds of stream silence,
@@ -141,9 +142,14 @@ async def handle_anthropic_messages(
         return StreamingResponse(events, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/messages")
+        # an abandoned request must not keep decoding to max_tokens (#222's watcher)
+        result = await _await_watching_disconnect(
+            generate_full(uid, spec, state, source="/v1/messages"), request, state, [uid]
+        )
     except GenerationError as exc:
         return _anthropic_error_response(400, "invalid_request_error", str(exc))
+    if result is None:
+        return _anthropic_error_response(499, "api_error", "client disconnected before the response was ready")
     response = anthropic_full_response(
         result, req.model, uid, cache_report=cache_report, metrics=metrics
     )
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index e7d2b970a7..3db3913909 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import math
 import json
 import time
 from collections.abc import AsyncIterator
@@ -231,8 +232,9 @@ def pick(value, key, framework):
     for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)):
         if not -2.0 <= float(value) <= 2.0:
             raise ValueError(f"{name} must be in [-2, 2], got {value}")
-    if resolved_rep <= 0.0:
-        raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}")
+    # not `<= 0`: NaN compares False and would poison every seen token's logit
+    if not (math.isfinite(resolved_rep) and resolved_rep > 0.0):
+        raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}")
     if min_tokens < 0:
         raise ValueError(f"min_tokens must be >= 0, got {min_tokens}")
     bias: list[list[float]] | None = None
@@ -487,7 +489,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None:
         # The qwen3 chat template opens an implicit <think> (thinking on) unless
         # enable_thinking is explicitly false, so the model emits only the closing
         # </think>. Mirror that default here, else the chain-of-thought leaks into content.
-        force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False
+        # continue_final_message is the exception: the template renders the final assistant
+        # turn as <think>...</think> + content, so the continuation starts in content.
+        ctk = spec.chat_template_kwargs or {}
+        force_reasoning = ctk.get("enable_thinking") is not False and not ctk.get(
+            "continue_final_message"
+        )
     elif parser_name == "glm":
         # GLM's template honors enable_thinking (default on) even with tools; the
         # generic fallback would force thinking and mislabel disabled output as reasoning.
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b67825bfaf..9ce63060ab 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -401,14 +401,17 @@ async def handle_chat_completion(
             chunks = state.stream_with_cancellation(chunks, request, uids)
         return StreamingResponse(chunks, media_type="text/event-stream")
 
+    tasks = [
+        asyncio.ensure_future(generate_full(u, spec, state, source="/v1/chat/completions"))
+        for u in uids
+    ]
     try:
-        results = await _await_watching_disconnect(
-            asyncio.gather(*(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids)),
-            request,
-            state,
-            uids,
-        )
+        results = await _await_watching_disconnect(asyncio.gather(*tasks), request, state, uids)
     except GenerationError as exc:
+        # gather does not cancel the other choices; the aborts below drop the events they
+        # wait on, so without the cancel they would stay pending forever
+        for task in tasks:
+            task.cancel()
         for u in uids:
             await state.abort_user(u)
         return create_error_response(str(exc), code=exc.code)
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index eea41722bb..9fd063a992 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -79,6 +79,7 @@
     submit_generation,
     with_keepalive,
 )
+from .openai_api import _await_watching_disconnect
 from .request_logger import log_request
 
 # Seconds of event silence before a keep-alive frame is emitted on the stream.
@@ -177,9 +178,14 @@ async def handle_responses(
         return StreamingResponse(events, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/responses")
+        # an abandoned request must not keep decoding to max_tokens (#222's watcher)
+        result = await _await_watching_disconnect(
+            generate_full(uid, spec, state, source="/v1/responses"), request, state, [uid]
+        )
     except GenerationError as exc:
         return _error_response(400, str(exc), exc.code)
+    if result is None:
+        return _error_response(499, "client disconnected before the response was ready")
     response = build_responses_response(
         result, req, response_id, created, cache_report=cache_report, metrics=metrics
     )
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index f512930dee..9648974dab 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -67,7 +67,15 @@ def _map_developer_role(
         m.get("role") == "developer" for m in messages
     ):
         return messages
-    return [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages]
+    mapped = [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages]
+    # The frontend hoisted and merged system messages before this mapping ran (templates such
+    # as Qwen's raise on a system message that is not first), so do it again for the new ones.
+    system = [m for m in mapped if m.get("role") == "system"]
+    if len(system) == 1 and mapped[0] is system[0]:
+        return mapped
+    rest = [m for m in mapped if m.get("role") != "system"]
+    content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content"))
+    return [{"role": "system", "content": content}] + rest
 
 
 class TokenizeManager:
diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py
index f91498fbb2..ee35c057c1 100644
--- a/python/freetoken/utils/mp.py
+++ b/python/freetoken/utils/mp.py
@@ -22,7 +22,9 @@ class _CoalescedUnpacker:
     """
 
     def __init__(self) -> None:
-        self._unpacker = msgpack.Unpacker(raw=False)
+        # 0 = effectively uncapped (2**31-1), as unpackb had; the Unpacker default of 100 MiB would
+        # kill the worker on a large message (e.g. an image's pixel values)
+        self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0)
         self._pending: list[Any] = []
 
     def feed(self, frame: bytes) -> None:
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index f051dec51f..a9645f4210 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -153,3 +153,13 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
     assert torch.isfinite(out[3]).all()
     # greedy path: argmax over the processed logits
     assert sampler.sample(logits, args)[1].item() == 5
+
+
+def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt():
+    # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the
+    # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA)
+    sampler = Sampler(CPU, V)
+    batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)])
+    args = sampler.prepare(batch)
+    out = apply_logits_processors(torch.ones(1, V), args.plan, V)
+    assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0
diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py
new file mode 100644
index 0000000000..bd12836893
--- /dev/null
+++ b/tests/kvcache/test_host_tier_live_pool.py
@@ -0,0 +1,53 @@
+"""The GDN snapshot tier (upstream #525) must copy against the pool's CURRENT tensors, which
+LinearStatePool.rebuild replaces, and must order its copies after the engine stream."""
+
+import contextlib
+from types import SimpleNamespace
+
+import torch
+
+import freetoken.kvcache.host_prefix_tier as tier_module
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+
+
+def _pool(slots: int):
+    return SimpleNamespace(
+        conv_states=torch.zeros(2, slots, 4, 3),
+        recurrent_states=torch.zeros(2, slots, 2, 4),
+        slot_states={},
+    )
+
+
+def test_read_after_a_pool_rebuild_lands_in_the_new_tensors():
+    pool = _pool(3)
+    tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu")
+    pool.conv_states[:, 1] = 5.0
+    pool.recurrent_states[:, 1] = 7.0
+    host = tier.write_gdn(1)
+    # a runtime rebuild swaps in fresh (larger) state tensors
+    pool.conv_states = torch.zeros(2, 6, 4, 3)
+    pool.recurrent_states = torch.zeros(2, 6, 2, 4)
+    tier.read_gdn(host, 5)
+    assert torch.all(pool.conv_states[:, 5] == 5.0)
+    assert torch.all(pool.recurrent_states[:, 5] == 7.0)
+
+
+def test_copies_wait_for_the_engine_stream(monkeypatch):
+    events = []
+    engine = object()
+
+    class FakeStream:
+        def wait_stream(self, other):
+            events.append(("wait", other))
+
+        def synchronize(self):
+            events.append(("sync",))
+
+    monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine)
+    monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext())
+    pool = _pool(3)
+    tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu")
+    tier.copy_stream = FakeStream()
+    host = tier.write_gdn(1)
+    tier.read_gdn(host, 2)
+    assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)]
diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py
new file mode 100644
index 0000000000..f67402041f
--- /dev/null
+++ b/tests/kvcache/test_host_tier_rehydrate_exact.py
@@ -0,0 +1,39 @@
+"""The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of
+exactly that node's prefix: the host index lookup is longest-prefix, and a shorter
+snapshot would resume the recurrence with the tokens in between silently missing."""
+
+from types import SimpleNamespace
+
+import torch
+
+from freetoken.kvcache.host_prefix_tier import HostMatch
+from freetoken.scheduler.cache import CacheManager
+
+
+def _manager(host_len: int, key_len: int = 128):
+    reads = []
+    tier = SimpleNamespace(
+        lookup=lambda key, page_size: HostMatch(host_len, None, 7),
+        read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)),
+    )
+    pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3])
+    manager = SimpleNamespace(
+        host_tier=tier,
+        prefix_cache=SimpleNamespace(_collect_key=lambda node: torch.arange(key_len)),
+        page_size=64,
+        linear_state_pool=pool,
+        ensure_mamba_slots=lambda n: None,
+    )
+    return manager, reads
+
+
+def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node():
+    manager, reads = _manager(host_len=64)
+    assert CacheManager._rehydrate_snapshot(manager, node=object()) is None
+    assert reads == []
+
+
+def test_an_exact_host_snapshot_revives_the_node():
+    manager, reads = _manager(host_len=128)
+    assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3
+    assert reads == [(7, 3)]
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index 74d4fd614e..02abaee1f4 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -201,6 +201,19 @@ def test_token_index_cache_is_keyed_by_shape():
     assert embedding._token_index(other)[0] is not first[0]
 
 
+def test_token_index_capture_never_reuses_an_evictable_cache_entry(monkeypatch):
+    """The eager warmup forward caches the decode pair right before capture; a capture that
+    reused it would bake an address the FIFO later evicts while the graph still reads it."""
+    embedding = _make_layer(_config()).ple_embedding
+    decode = _meta([[5]], [[EOS, EOS]], decode=True)
+    eager = embedding._token_index(decode)  # the warmup forward
+    monkeypatch.setattr(ple_module, "_capturing", lambda device: True)
+    captured = embedding._token_index(decode)
+    assert captured[0] is not eager[0] and captured[1] is not eager[1]
+    assert all(v[0] is not captured[0] for v in embedding._token_index_cache.values())
+    assert captured[0].tolist() == [0] and captured[1].tolist() == [0]
+
+
 def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill():
     """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request."""
     layer = _make_layer(_config())
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index ce675a2508..7d5b3cd114 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -223,6 +223,26 @@ def test_glm_reasoning_parser_honors_disabled_thinking_with_tools():
     assert parser is not None and parser.detector.force_reasoning is True
 
 
+def test_qwen3_continuation_of_a_final_assistant_message_is_content_not_reasoning():
+    # The Qwen template closes the think block before the assistant's content, so a
+    # continue_final_message continuation never emits </think>; starting the parser inside
+    # reasoning filed the whole continuation under reasoning_content.
+    from freetoken.server.generation import _make_reasoning_parser
+
+    state = FakeState([], reasoning_parser="qwen3")
+    req = chat_request(
+        messages=[
+            {"role": "user", "content": "Count to five."},
+            {"role": "assistant", "content": "One, two,"},
+        ],
+        continue_final_message=True,
+    )
+    parser = _make_reasoning_parser(chat_request_to_genspec(req, {}), state)
+    assert parser is not None and parser.detector.force_reasoning is False
+    plain = _make_reasoning_parser(chat_request_to_genspec(chat_request(), {}), state)
+    assert plain.detector.force_reasoning is True
+
+
 def test_non_stream_chat_completion_returns_openai_tool_calls_and_sends_tools():
     output = '[TOOL_CALLS] [{"name":"get_weather","arguments":{"city":"Paris"}}]'
     state = FakeState(
diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py
new file mode 100644
index 0000000000..3b2e1d6473
--- /dev/null
+++ b/tests/server/test_sampling_validation.py
@@ -0,0 +1,23 @@
+"""resolve_sampling rejects non-finite penalties: NaN passes a `<= 0` check and would
+turn every already-seen token's logit into NaN."""
+
+import pytest
+
+from freetoken.server.generation import resolve_sampling
+
+
+def _resolve(**extra):
+    return resolve_sampling(
+        temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False,
+        model_sampling={}, **extra,
+    )
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0])
+def test_repetition_penalty_must_be_positive_and_finite(value):
+    with pytest.raises(ValueError, match="repetition_penalty"):
+        _resolve(repetition_penalty=value)
+
+
+def test_a_normal_repetition_penalty_is_kept():
+    assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1)
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index ba7da7c17b..71d5ac8b85 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -316,3 +316,61 @@ def new_user(self):
 
     assert resp.status_code == 499
     assert sorted(state.aborted) == [8, 9]
+
+
+def test_a_failed_choice_cancels_the_other_choices(monkeypatch):
+    # n > 1: one choice errors while another is still generating. gather passes the error up
+    # without cancelling the sibling, and the aborts drop the event it waits on.
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+
+    class _State(_ApiState):
+        def __init__(self):
+            super().__init__(acks=None)
+            self._next = 7
+
+        def new_user(self):
+            self._next += 1
+            return self._next
+
+        async def wait_for_ack(self, uid):
+            if uid == 8:
+                yield SimpleNamespace(**{**vars(_ack("")), "error": "boom"})
+                return
+            await asyncio.sleep(3600)
+            yield _ack("never")
+
+    async def scenario():
+        state = _State()
+        resp = await openai_api.handle_chat_completion(
+            _chat_req(n=2), _Request(disconnected=False), state, {}
+        )
+        await asyncio.sleep(0)
+        pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+        return resp, state, pending
+
+    resp, state, pending = asyncio.run(scenario())
+    assert resp.status_code == 400
+    assert sorted(state.aborted) == [8, 9]
+    assert pending == []
+
+
+def test_messages_and_responses_non_stream_disconnect_delivers_abort(monkeypatch):
+    # #222 only wrapped the OpenAI handlers; the Anthropic and Responses endpoints kept
+    # decoding an abandoned non-streaming request to max_tokens.
+    from freetoken.server import anthropic_api, responses_api
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+    from freetoken.server.responses_api import ResponsesRequest
+
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+
+    state = _ApiState(acks=None)
+    req = AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "hi"}])
+    handler = anthropic_api.handle_anthropic_messages(req, _Request(disconnected=True), state, {})
+    resp = asyncio.run(asyncio.wait_for(handler, timeout=5))  # unwatched: hangs, then times out
+    assert resp.status_code == 499 and state.aborted == [7]
+
+    state = _ApiState(acks=None)
+    req = ResponsesRequest(model="m", input="hi", max_output_tokens=8)
+    handler = responses_api.handle_responses(req, _Request(disconnected=True), state, {})
+    resp = asyncio.run(asyncio.wait_for(handler, timeout=5))
+    assert resp.status_code == 499 and state.aborted == [7]
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index c82a76861c..390fcd4cce 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -827,3 +827,21 @@ def test_developer_role_maps_to_system_unless_the_template_knows_it():
     assert mapped[0]["content"] == "be brief" and msgs[0]["role"] == "developer"  # input untouched
     assert _map_developer_role(msgs, "{% if message.role == 'developer' %}") is msgs
     assert _map_developer_role(msgs[1:], None) == msgs[1:]  # nothing to map
+
+
+def test_mapped_developer_messages_are_merged_into_one_leading_system_message():
+    """The frontend hoists system messages before the tokenizer maps developer -> system; a
+    template that rejects a non-first system message must still get exactly one, first."""
+    from freetoken.tokenizer.tokenize import _map_developer_role
+
+    template = "{% if message.role == 'system' %}"
+    msgs = [
+        {"role": "system", "content": "be brief"},
+        {"role": "developer", "content": "answer in German"},
+        {"role": "user", "content": "hi"},
+    ]
+    mapped = _map_developer_role(msgs, template)
+    assert [m["role"] for m in mapped] == ["system", "user"]
+    assert mapped[0]["content"] == "be brief\n\nanswer in German"
+    late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}]
+    assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"]
diff --git a/tests/utils/test_mp_large_frame.py b/tests/utils/test_mp_large_frame.py
new file mode 100644
index 0000000000..f91ee5604d
--- /dev/null
+++ b/tests/utils/test_mp_large_frame.py
@@ -0,0 +1,13 @@
+"""The coalesced-frame decoder (#466) must accept a message larger than msgpack's
+Unpacker default buffer (100 MiB), which unpackb never limited."""
+
+import msgpack
+
+from freetoken.utils.mp import _CoalescedUnpacker
+
+
+def test_a_message_over_100_mib_decodes():
+    payload = {"pixels": b"x" * (101 * 1024 * 1024)}
+    unpacker = _CoalescedUnpacker()
+    unpacker.feed(msgpack.packb(payload))
+    assert len(unpacker) == 1 and len(unpacker.take()["pixels"]) == len(payload["pixels"])

From 78000e8d8dbb3c9f042a28d7a6c902acf9e66076 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 23:52:44 -0400
Subject: [PATCH 153/249] merge fork fix (issue #65): run the PLE graph fill on
 its own thread, submitted before the launch

# Conflicts:
#	python/freetoken/models/qwen4_exp/ple_disk.py
---
 python/freetoken/models/qwen4_exp/ple_disk.py | 126 ++++++++++++----
 .../qwen4_exp/test_ple_disk_fill_order.py     | 136 ++++++++++++++++++
 2 files changed, 237 insertions(+), 25 deletions(-)
 create mode 100644 tests/models/qwen4_exp/test_ple_disk_fill_order.py

diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index 536853be3b..a8e6d550c7 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -6,6 +6,9 @@
 from __future__ import annotations
 
 import os
+import queue
+import threading
+from concurrent.futures import Future
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Sequence
@@ -129,6 +132,52 @@ def resolve_row_source(model_path: str) -> PleRowSource:
     return source_from_safetensors(download_hf_weight(model_path))
 
 
+class _Filler:
+    """The one thread that drives the PleStore (engine-thread-only C++, no locks), FIFO.
+
+    Graph fills run here, submitted before the launch: the fill releases the captured decode's
+    WAIT, and cuGraphLaunch may block until the GPU drains, so it must not wait for the launch
+    to return (a large graph deadlocked that way; fork issue #65)."""
+
+    def __init__(self) -> None:
+        self._jobs: queue.SimpleQueue = queue.SimpleQueue()
+        self._thread = threading.Thread(target=self._run, name="ple-filler", daemon=True)
+        self._thread.start()
+
+    def submit(self, fn) -> Future:
+        fut: Future = Future()
+        self._jobs.put((fn, fut))
+        return fut
+
+    def _run(self) -> None:
+        while (item := self._jobs.get()) is not None:
+            fn, fut = item
+            try:
+                fut.set_result(fn())
+            except BaseException as exc:  # noqa: BLE001 -- surfaced to the engine thread
+                fut.set_exception(exc)
+
+    def close(self) -> None:
+        self._jobs.put(None)
+        self._thread.join(timeout=5)
+
+
+class _PendingFill:
+    """A graph fill submitted before its launch; ``cancel`` after a failed launch."""
+
+    def __init__(self, table: "DiskRowTable", future: Future, cancelled: threading.Event) -> None:
+        self._table, self._future, self._cancelled = table, future, cancelled
+
+    def cancel(self) -> None:
+        # no WAIT consumed the signal: clear it, or the next step's WAIT passes on stale rows
+        self._cancelled.set()
+        try:
+            self._future.result()
+        except BaseException:  # noqa: BLE001 -- the launch error is the one to surface
+            pass
+        self._table._flag.zero_()
+
+
 class DiskRowTable:
     """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk)."""
 
@@ -184,9 +233,6 @@ def __init__(
         self._graph_dev = torch.empty(
             max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device
         )
-        # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the
-        # dispatch context after the step's fill: a record inside the capture would not be observable)
-        self._graph_consumed = torch.cuda.Event()
         eager_bytes = max_extend_tokens * self._token_bytes
         # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the
         # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy
@@ -211,10 +257,36 @@ def __init__(
 
     # ---------------- host side (engine thread, before the forward launches) ----------------
 
+    def _filler(self) -> _Filler:
+        filler = self.__dict__.get("_filler_thread")
+        if filler is None:
+            filler = self._filler_thread = _Filler()
+        return filler
+
+    def close(self) -> None:
+        filler = self.__dict__.pop("_filler_thread", None)
+        if filler is not None:
+            filler.close()
+
+    def _current_stream(self):
+        return torch.cuda.current_stream(self._device)
+
+    def _raise_failed_fill(self) -> None:
+        pending = self.__dict__.pop("_pending_fill", None)
+        if pending is not None and pending.done() and pending.exception() is not None:
+            raise pending.exception()
+        if pending is not None and not pending.done():
+            self._pending_fill = pending
+
+    def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
+        """A fill that completes before its launch, on the store's thread (after any queued graph fill)."""
+        self._filler().submit(lambda: self.fill(runs, graph=graph)).result()
+
     def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
         """Stage per-request token runs (two context ids, then the new tokens) in batch order."""
         if graph:
-            self._graph_consumed.synchronize()
+            # no launch fence: the job already waited on the readback recorded after the last graph,
+            # and a fence recorded after this step's launch would wait on the WAIT this fill releases
             pinned = self._graph_pinned
         else:
             self._eager_slot ^= 1
@@ -236,22 +308,27 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]:
             return _context(ids, position, self.eos_token_id)
         return [self.image_token_id if t >= MM_PAD_SHIFT_VALUE else t for t in _context(ids, position, self.eos_token_id)]
 
-    def host_fill_batch(self, batch: Batch, use_graph: bool):
-        """Stage this batch's rows; returns the post-dispatch fill callable under flag-sync, else None."""
-        eos = self.eos_token_id
+    def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None:
+        """Stage this batch's rows before the dispatch. Under flag-sync a graph fill is only submitted:
+        it completes on the filler thread, independently of the launch; returns its handle, else None."""
+        self._raise_failed_fill()
         if batch.is_decode:
             reqs = list(batch.reqs)
+            # the context comes from the request as it is NOW: the engine mutates it once the launch returns
+            contexts = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs]
             if use_graph and self._wait_sync:
                 bs = batch.padded_size
                 self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
-                self._readback_event.record(torch.cuda.current_stream(self._device))
+                self._readback_event.record(self._current_stream())
+                cancelled = threading.Event()
 
-                def _complete() -> None:
+                def _job() -> None:
                     try:
                         self._readback_event.synchronize()
                         tokens = self._token_readback[:bs].to(torch.int64).tolist()
-                        runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64)
-                                for r, t in zip(reqs, tokens)]
+                        if cancelled.is_set():
+                            return
+                        runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)]
                         self.fill(runs, graph=True)
                     except BaseException:
                         from freetoken.kernel.row_store import signal
@@ -260,12 +337,12 @@ def _complete() -> None:
                         signal(self._flag)
                         raise
 
-                return _complete
+                self._pending_fill = self._filler().submit(_job)
+                return _PendingFill(self, self._pending_fill, cancelled)
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
             tokens = batch.input_ids.to("cpu").to(torch.int64).tolist()
-            runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64)
-                    for r, t in zip(reqs, tokens)]
-            self.fill(runs, graph=use_graph)
+            runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)]
+            self._fill_now(runs, graph=use_graph)
             return None
         runs = [
             torch.cat((
@@ -274,20 +351,19 @@ def _complete() -> None:
             ))
             for req in batch.padded_reqs
         ]
-        self.fill(runs, graph=False)
+        self._fill_now(runs, graph=False)
         return None
 
     @contextmanager
     def forward_host_ctx(self, batch: Batch, use_graph: bool):
-        """Around one dispatch: stage on enter, run the deferred fill+signal on exit."""
-        deferred = self.host_fill_batch(batch, use_graph)
-        yield
-        # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run
-        if deferred is not None:
-            deferred()
-        if use_graph:
-            # only after this step's fill: the next fill waits on this launch, not the fill on it
-            self._graph_consumed.record(torch.cuda.current_stream(self._device))
+        """Around one dispatch: stage (or submit the graph fill) on enter; cancel a submitted fill if the launch fails."""
+        pending = self.host_fill_batch(batch, use_graph)
+        try:
+            yield
+        except BaseException:
+            if pending is not None:
+                pending.cancel()
+            raise
 
     # ---------------- device side (PLETableBackend protocol) ----------------
 
diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
new file mode 100644
index 0000000000..e7fe58f7e0
--- /dev/null
+++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
@@ -0,0 +1,136 @@
+"""The PLE disk backend's wait-sync fill must not depend on the graph launch returning.
+
+The captured decode graph WAITs on ``_flag``; ``cuGraphLaunch`` may block until the GPU drains
+(a large graph fills the launch queue), and the GPU is parked on that WAIT. So the fill that sets
+the flag has to run on its own thread, started before the launch (fork issue #65).
+These tests stand in for the GPU with host fakes: the "launch" blocks until the flag is set.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import freetoken.models.qwen4_exp.ple_disk as ple_disk
+from freetoken.models.qwen4_exp.ple_disk import DiskRowTable
+
+EOS = 9
+
+
+class _Store:
+    """Stands in for the C++ PleStore: records calls and the thread that made them."""
+
+    def __init__(self, flag: torch.Tensor):
+        self.flag = flag
+        self.calls: list[tuple[str, list[int] | None, str]] = []
+
+    def stage(self, tokens_addr, n, staging_addr):
+        tokens = ctypes_read(tokens_addr, n + 2)
+        self.calls.append(("stage", tokens, threading.current_thread().name))
+
+    def flush(self, signal_addr):
+        self.calls.append(("flush", None, threading.current_thread().name))
+        if signal_addr:
+            self.flag[0] = 1
+
+
+def ctypes_read(addr: int, n: int) -> list[int]:
+    import ctypes
+
+    return list((ctypes.c_int64 * n).from_address(addr))
+
+
+class _Event:
+    def record(self, stream=None):
+        pass
+
+    def synchronize(self):
+        pass
+
+
+def _table(monkeypatch) -> DiskRowTable:
+    # host-only: no CUDA stream behind the (faked) readback event
+    monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None)
+    t = DiskRowTable.__new__(DiskRowTable)
+    t._device = None
+    t._wait_sync = True
+    t._flag = torch.zeros(1, dtype=torch.int64)
+    t._store = _Store(t._flag)
+    t._token_readback = torch.zeros(8, dtype=torch.int32)
+    t._readback_event = _Event()
+    t._graph_pinned = torch.zeros(64, dtype=torch.uint8)
+    t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)]
+    t._eager_read = [_Event(), _Event()]
+    t._eager_slot = 0
+    t._token_bytes = 1
+    t.eos_token_id = EOS
+    t.image_token_id = None
+    return t
+
+
+def _close(t: DiskRowTable) -> None:
+    close = getattr(t, "close", None)
+    if close is not None:
+        close()
+
+
+def _decode_batch(prev_tokens: list[int], new_token: int):
+    req = SimpleNamespace(input_ids=torch.tensor(prev_tokens + [new_token]), device_len=len(prev_tokens) + 1)
+    return SimpleNamespace(is_decode=True, reqs=[req], padded_size=1,
+                           input_ids=torch.tensor([new_token], dtype=torch.int32))
+
+
+def _blocking_launch(flag: torch.Tensor, timeout: float = 5.0) -> bool:
+    """A cuGraphLaunch that cannot return until the GPU passes the graph's WAIT(flag)."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if int(flag[0]) == 1:
+            flag[0] = 0  # the graph's WAIT-then-RESET
+            return True
+        time.sleep(0.001)
+    return False
+
+
+def test_graph_fill_does_not_wait_for_the_launch_to_return(monkeypatch):
+    t = _table(monkeypatch)
+    with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True):
+        assert _blocking_launch(t._flag), "deadlock: the fill only runs after the launch returns"
+    _close(t)
+
+
+def test_fills_run_in_step_order_on_one_thread_with_contexts_taken_before_the_launch(monkeypatch):
+    t = _table(monkeypatch)
+    for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8), ([1, 2, 3, 7, 8], 4)):
+        batch = _decode_batch(prev, new)
+        with t.forward_host_ctx(batch, use_graph=True):
+            # the engine mutates the request as soon as the launch returns (overlap scheduling)
+            batch.reqs[0].input_ids = torch.tensor([0, 0, 0, 0, 0, 0])
+            batch.reqs[0].device_len += 1
+            assert _blocking_launch(t._flag)
+    _close(t)
+    staged = [c for c in t._store.calls if c[0] == "stage"]
+    assert [c[1] for c in staged] == [[2, 3, 7], [3, 7, 8], [7, 8, 4]]
+    assert len({c[2] for c in t._store.calls}) == 1  # the store stays single-threaded
+
+
+def test_eager_fill_completes_before_the_launch(monkeypatch):
+    t = _table(monkeypatch)
+    req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4)
+    batch = SimpleNamespace(is_decode=False, padded_reqs=[req])
+    with t.forward_host_ctx(batch, use_graph=False):
+        assert [c[0] for c in t._store.calls] == ["stage", "flush"]
+    _close(t)
+
+
+def test_a_failed_launch_leaves_no_stale_signal(monkeypatch):
+    t = _table(monkeypatch)
+    with pytest.raises(RuntimeError, match="launch failed"):
+        with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True):
+            raise RuntimeError("launch failed")
+    # no WAIT consumed it, so a leftover 1 would let the next step read stale rows
+    assert int(t._flag[0]) == 0
+    _close(t)

From abd60a17f578dcdf8bdd9c59a3ee8e710b6e2dc3 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 23:52:44 -0400
Subject: [PATCH 154/249] merge fork perf: load only the routed experts for
 short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off by default)

---
 python/freetoken/layers/moe.py            | 19 +++++++
 tests/moe/test_small_prefill_on_demand.py | 69 +++++++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 tests/moe/test_small_prefill_on_demand.py

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 598bd1193f..a8ae742e07 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -24,6 +24,9 @@
 # default. Set FREETOKEN_HYBRID_OVERLAP=0 to force the serial path (CPU sync before the
 # GPU work) -- a measurement-only escape hatch to A/B the overlap benefit.
 _HYBRID_OVERLAP = os.getenv("FREETOKEN_HYBRID_OVERLAP", "1") != "0"
+# Prefills of at most this many tokens load only their routed experts (the decode path's
+# on-demand LRU) instead of streaming every expert of every layer; 0 = off.
+_SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0"))
 
 
 class MoELayer(BaseOP):
@@ -354,6 +357,22 @@ def _prefill_routed(
         pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
+        if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"):
+            # A short extension (an agent turn over a cached prefix) touches a fraction of each
+            # layer's experts; fetching those beats streaming all of them (~68 GB here). The same
+            # token count holds in every layer, so the overlap double buffer is never begun.
+            cache.ensure_experts(self.layer_id, topk_ids)
+            cache.copy_missing()
+            return self._expert_gemm(
+                cache,
+                hidden_states,
+                topk_weights,
+                topk_ids,
+                views=cache.bank_views(),
+                n=None,
+                alphas=cache.alphas_for_slots(self.layer_id),
+                is_prefill=False,
+            )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
new file mode 100644
index 0000000000..5d984a0dec
--- /dev/null
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -0,0 +1,69 @@
+"""Short prefills load only their routed experts (FREETOKEN_MOE_SMALL_PREFILL_TOKENS) instead of
+streaming whole layers; above the threshold, or with it off, the whole-layer path is unchanged."""
+
+import pytest
+import torch
+
+import freetoken.layers.moe as moe_mod
+from freetoken.distributed import set_tp_info, try_get_tp_info
+
+
+def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32):
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(),
+                            prefix="model.layers.0.mlp.experts")
+    cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device)
+    g = torch.Generator().manual_seed(0)
+    cache.set_bank_sources({
+        "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1],
+        "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1],
+    })
+    layer.offload_cache = cache
+    return layer, cache
+
+
+@pytest.mark.parametrize("tokens, threshold, on_demand", [
+    (3, 8, True), (8, 8, True), (9, 8, False), (3, 0, False),
+])
+def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand):
+    layer, cache = _layer_and_cache(torch.device("cpu"))
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+    calls = []
+    monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure"))
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize"))
+    monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy"))
+    monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False)
+    got = {}
+
+    def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill):
+        got.update(n=n, is_prefill=is_prefill)
+        return hs
+
+    monkeypatch.setattr(layer, "_expert_gemm", fake_gemm)
+    hs = torch.randn(tokens, 16)
+    layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32))
+    if on_demand:
+        assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False}
+    else:
+        assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_small_prefill_matches_the_whole_layer_path(monkeypatch):
+    from freetoken.core import get_global_ctx  # noqa: F401  (layer reads the ctx only in forward())
+
+    dev = torch.device("cuda")
+    hs = torch.randn(5, 16, dtype=torch.bfloat16, device=dev) * 0.5
+    ids = torch.tensor([[0, 3], [1, 3], [7, 2], [3, 4], [5, 0]], dtype=torch.int32, device=dev)
+    w = torch.tensor([[0.6, 0.4], [0.5, 0.5], [0.9, 0.1], [0.3, 0.7], [0.2, 0.8]], device=dev)
+    outs = {}
+    for threshold in (0, 8):
+        monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+        layer, cache = _layer_and_cache(dev)
+        outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float()
+    torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2)

From db2cee587a145c3f96d2bfec47a6d001a8e07c86 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 24 Sep 2026 00:27:13 -0400
Subject: [PATCH 155/249] merge fork fix: small-prefill on-demand path hands
 the LRU kernel unique ids under Triton's element limit

---
 python/freetoken/layers/moe.py            | 49 +++++++++++++++++------
 tests/moe/test_small_prefill_on_demand.py | 37 +++++++++++++++--
 2 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index a8ae742e07..6ab51f8a87 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -27,6 +27,11 @@
 # Prefills of at most this many tokens load only their routed experts (the decode path's
 # on-demand LRU) instead of streaming every expert of every layer; 0 = off.
 _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0"))
+_TRITON_MAX_NUMEL = 1 << 20  # Triton's per-tensor element limit (the LRU kernel's block)
+
+
+def triton_next_power_of_2(n: int) -> int:
+    return 1 << max(0, int(n) - 1).bit_length()
 
 
 class MoELayer(BaseOP):
@@ -361,18 +366,18 @@ def _prefill_routed(
             # A short extension (an agent turn over a cached prefix) touches a fraction of each
             # layer's experts; fetching those beats streaming all of them (~68 GB here). The same
             # token count holds in every layer, so the overlap double buffer is never begun.
-            cache.ensure_experts(self.layer_id, topk_ids)
-            cache.copy_missing()
-            return self._expert_gemm(
-                cache,
-                hidden_states,
-                topk_weights,
-                topk_ids,
-                views=cache.bank_views(),
-                n=None,
-                alphas=cache.alphas_for_slots(self.layer_id),
-                is_prefill=False,
-            )
+            slots = self._ensure_unique(cache, topk_ids)
+            if slots is not None:
+                return self._expert_gemm(
+                    cache,
+                    hidden_states,
+                    topk_weights,
+                    slots,
+                    views=cache.bank_views(),
+                    n=None,
+                    alphas=cache.alphas_for_slots(self.layer_id),
+                    is_prefill=False,
+                )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
@@ -402,6 +407,26 @@ def _prefill_routed(
         cache.release_prefill_layer(self.layer_id)
         return out
 
+    def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None:
+        """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or
+        None (take the whole-layer path) when they cannot all be resident at once.
+
+        The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a
+        prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids
+        instead (at most num_experts), chunked under that limit; each chunk's misses are copied
+        before the next ensure reuses the copy plan."""
+        uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True)
+        n = uniq.numel()
+        if n > cache.cache_size // 2:  # leave room so a chunk never evicts the one before it
+            return None
+        uniq = uniq.to(torch.int32)
+        per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size))
+        for start in range(0, n, per_call):
+            part = uniq[start : start + per_call]  # a contiguous view: ensure rewrites it in place
+            cache.ensure_experts(self.layer_id, part)
+            cache.copy_missing()
+        return uniq[inverse].view_as(topk_ids)
+
     def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]:
         """Double-buffer choreography for this layer's overlap prefill: kick off the
         next layer's full-layer H2D copy, then return this layer's bank views (in
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index 5d984a0dec..e7fb70721a 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -8,7 +8,7 @@
 from freetoken.distributed import set_tp_info, try_get_tp_info
 
 
-def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32):
+def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None):
     from freetoken.layers.moe import OffloadMoELayer
     from freetoken.layers.quantization import NoQuantConfig
     from freetoken.moe.offload_cache import OffloadMoeCache
@@ -17,7 +17,7 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32):
         set_tp_info(rank=0, size=1)
     layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(),
                             prefix="model.layers.0.mlp.experts")
-    cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device)
+    cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device)
     g = torch.Generator().manual_seed(0)
     cache.set_bank_sources({
         "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1],
@@ -64,6 +64,35 @@ def test_small_prefill_matches_the_whole_layer_path(monkeypatch):
     outs = {}
     for threshold in (0, 8):
         monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
-        layer, cache = _layer_and_cache(dev)
-        outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float()
+        layer, cache = _layer_and_cache(dev, cache_size=32)
+        outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
     torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2)
+
+
+def _run_and_check_path(layer, hs, w, ids, *, on_demand):
+    taken = []
+    real = getattr(layer, "_ensure_unique", None)
+    if real is not None:
+        layer._ensure_unique = lambda cache, t: taken.append(r := real(cache, t)) or r
+    out = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float()
+    assert (bool(taken) and taken[0] is not None) == on_demand
+    return out
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch):
+    # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the
+    # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's
+    # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead.
+    dev = torch.device("cuda")
+    tokens, experts, top_k = 200, 512, 10
+    g = torch.Generator(device=dev).manual_seed(1)
+    hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 0.5
+    ids = torch.randint(0, experts, (tokens, top_k), dtype=torch.int32, device=dev, generator=g)
+    w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1)
+    outs = {}
+    for threshold in (0, 1024):
+        monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+        layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100)
+        outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
+    torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2)

From 908546cc0aae72a04ccd420e7543fed08e9ac97e Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 24 Sep 2026 00:50:11 -0400
Subject: [PATCH 156/249] merge fork fix: small-prefill ensures its experts in
 decode-width chunks

---
 python/freetoken/layers/moe.py            | 19 ++++++++-----------
 tests/moe/test_small_prefill_on_demand.py |  6 +++++-
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 6ab51f8a87..2f02a0d2fe 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -27,11 +27,9 @@
 # Prefills of at most this many tokens load only their routed experts (the decode path's
 # on-demand LRU) instead of streaming every expert of every layer; 0 = off.
 _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0"))
-_TRITON_MAX_NUMEL = 1 << 20  # Triton's per-tensor element limit (the LRU kernel's block)
-
-
-def triton_next_power_of_2(n: int) -> int:
-    return 1 << max(0, int(n) - 1).bit_length()
+# ids per LRU ensure in that path: the kernel's block is next_pow2(ids) x next_pow2(slots), and a
+# block much wider than decode's spills to local memory whose launch reservation can OOM
+_ENSURE_CHUNK_IDS = 32
 
 
 class MoELayer(BaseOP):
@@ -412,17 +410,16 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc
         None (take the whole-layer path) when they cannot all be resident at once.
 
         The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a
-        prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids
-        instead (at most num_experts), chunked under that limit; each chunk's misses are copied
-        before the next ensure reuses the copy plan."""
+        prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it
+        can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in
+        decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan."""
         uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True)
         n = uniq.numel()
         if n > cache.cache_size // 2:  # leave room so a chunk never evicts the one before it
             return None
         uniq = uniq.to(torch.int32)
-        per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size))
-        for start in range(0, n, per_call):
-            part = uniq[start : start + per_call]  # a contiguous view: ensure rewrites it in place
+        for start in range(0, n, _ENSURE_CHUNK_IDS):
+            part = uniq[start : start + _ENSURE_CHUNK_IDS]  # a contiguous view: ensure rewrites it in place
             cache.ensure_experts(self.layer_id, part)
             cache.copy_missing()
         return uniq[inverse].view_as(topk_ids)
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index e7fb70721a..eea5842d23 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -83,7 +83,8 @@ def _run_and_check_path(layer, hs, w, ids, *, on_demand):
 def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch):
     # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the
     # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's
-    # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead.
+    # 1M-element limit, and a 512-id block failed to launch on a full GPU (CUDA OOM). The path
+    # must hand it the unique ids in decode-width chunks.
     dev = torch.device("cuda")
     tokens, experts, top_k = 200, 512, 10
     g = torch.Generator(device=dev).manual_seed(1)
@@ -94,5 +95,8 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch):
     for threshold in (0, 1024):
         monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
         layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100)
+        widths, ensure = [], cache.ensure_experts
+        monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1])
         outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
+    assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's"
     torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2)

From 072b29a2c9f29cc54e22209367fe96a9c3b0f7c4 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sat, 26 Sep 2026 16:14:49 -0400
Subject: [PATCH 157/249] merge fork perf: scored expert eviction
 (--moe-cache-policy), shared expert beside the miss copy
 (FREETOKEN_MOE_COPY_OVERLAP), opt-in slim copy

---
 benchmarks/README.md                          |  30 +
 benchmarks/bench_expert_copy_corun.py         | 518 +++++++++++++++
 benchmarks/bench_moe_copy_overlap.py          | 233 +++++++
 benchmarks/bench_scored_ensure.py             | 158 +++++
 python/freetoken/env.py                       |   2 +
 python/freetoken/kernel/aot.py                |  19 +
 .../kernel/csrc/jit/fast_index_copy.cuh       | 109 +++
 python/freetoken/kernel/fast_index_copy.py    |  42 ++
 python/freetoken/layers/moe.py                |  21 +-
 python/freetoken/models/deepseek_v4/moe.py    |   2 +-
 python/freetoken/models/qwen4_exp/moe.py      |  12 +-
 python/freetoken/moe/offload_cache.py         | 170 ++++-
 python/freetoken/moe/offload_kernels.py       |  45 ++
 python/freetoken/moe/scored_ensure.py         | 395 +++++++++++
 python/freetoken/server/args.py               |   9 +-
 tests/README.md                               |   1 +
 tests/models/qwen4_exp/test_skeleton.py       | 124 ++++
 tests/moe/ref_scored_cache.py                 | 159 +++++
 tests/moe/test_fused_copy.py                  | 173 ++++-
 tests/moe/test_offload.py                     | 117 +++-
 tests/moe/test_scored_ensure.py               | 622 ++++++++++++++++++
 tests/moe/test_small_prefill_on_demand.py     |  43 +-
 22 files changed, 2971 insertions(+), 33 deletions(-)
 create mode 100644 benchmarks/bench_expert_copy_corun.py
 create mode 100644 benchmarks/bench_moe_copy_overlap.py
 create mode 100644 benchmarks/bench_scored_ensure.py
 create mode 100644 python/freetoken/moe/scored_ensure.py
 create mode 100644 tests/moe/ref_scored_cache.py
 create mode 100644 tests/moe/test_scored_ensure.py

diff --git a/benchmarks/README.md b/benchmarks/README.md
index 5cbb0b34e7..3f0df72f81 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -27,6 +27,36 @@ batch size x miss rate.
 python benchmarks/bench_offload_cache_copy.py
 ```
 
+**`bench_scored_ensure.py`** — synthetic (no checkpoint): us per decode `ensure` call at a
+fixed miss count, flashlib `lru_ensure` vs the `--moe-cache-policy` kernels (lru, kd, kdfb,
+rule), from a CUDA graph of two decode steps on a full cache. `--stats` times it with
+`--moe-collect-stats` accumulation on.
+
+```bash
+python benchmarks/bench_scored_ensure.py --slots 1650 --k 10,20
+```
+
+**`bench_moe_copy_overlap.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_COPY_OVERLAP`
+off vs on for Qwen4ExpMoE decode layers at the Qwen3.8-Flash-Next expert geometry, as CUDA
+graphs. Checks the outputs are bitwise identical, times replays ABAB, and profiles the copy's
+start after the ensure and how much of the side-stream shared expert runs inside it.
+
+```bash
+python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160
+```
+
+**`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel
+alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps
+`fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim`
+(blocks x threads x unroll) over 1..10-expert plans (GB/s, "hot" back-to-back and "cold"
+spacer-separated copies), then reports the stretch sigma of bf16/fp8 GEMVs, a small matmul and
+the NVFP4 routed decode GEMV while a copy loop runs on another stream. `--json` keeps the raw
+per-trial times.
+
+```bash
+python benchmarks/bench_expert_copy_corun.py --trials 5 --json copy_corun.json
+```
+
 For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw`
 instead — it writes the JSON profile the engine reads.
 
diff --git a/benchmarks/bench_expert_copy_corun.py b/benchmarks/bench_expert_copy_corun.py
new file mode 100644
index 0000000000..b76b82ee6f
--- /dev/null
+++ b/benchmarks/bench_expert_copy_corun.py
@@ -0,0 +1,518 @@
+"""Expert-copy kernel sweep and co-run stretch on real-size NVFP4 expert banks.
+
+Two measurements, no checkpoint needed:
+
+* sweep: GB/s of one fused expert copy (all 6 banks, 1..10 experts per plan) for each
+  ``fast_index_copy_multi`` (num_threads, blocks_per_bank) config, and for each
+  ``fast_index_copy_multi_slim`` (blocks, threads, unroll) config when that kernel exists;
+* co-run: the stretch ``sigma = t_corun / t_alone - 1`` of decode compute kernels (bf16 and
+  fp8 GEMV [1,H]x[H,N], a small bf16 matmul, the NVFP4 routed decode GEMV) while a copy loop
+  runs on another stream, and the copy's GB/s while each compute loop runs.
+
+The host banks are registered anonymous mmaps (``HostBank``, like the serving banks) with
+the qwen4_exp triton-NVFP4 row sizes. Source rows are not reused within a loop, the GEMV
+weights are cycled over copies that together exceed L2, and every loop is a CUDA graph.
+Alone and co-run trials are interleaved; each cell reports min and median over trials.
+
+Copy loops run in two modes. "hot": copy nodes back to back. "cold": a one-element spacer
+kernel after every copy, as the ensure kernel precedes every decode copy. On the RTX 5080, in
+graph replay, a copy node that follows another kernel's node runs ~8 us longer than one that
+follows a copy node, for every variant (CUPTI puts it in the copy's own duration; eager
+launches do not show it; cause not identified). Cold GB/s subtracts the spacer's alone time.
+
+Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_expert_copy_corun.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import statistics
+import tempfile
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg
+from freetoken.kernel import fast_index_copy as fic
+from freetoken.kernel.pinned import device_ptr
+from freetoken.moe.host_banks import HostBank
+from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT
+
+L2_BYTES = 64 << 20
+SLEEP_CYCLES = 2_000_000  # ~1 ms head start so the host enqueues every graph before the GPU starts
+LMAX = 16  # plan buffer length; the kernels read the live count from num_indices
+
+
+def nvfp4_bank_bytes(h: int, i: int) -> list[int]:
+    # triton NVFP4 layout() order: gate_up, gate_up_scale, gate_up_global, down, down_scale, down_global
+    return [2 * i * (h // 2), 2 * i * (h // 16), 2 * i * 2, h * (i // 2), h * (i // 16), h * 2]
+
+
+@dataclass(frozen=True)
+class CopyVariant:
+    name: str
+    launch: Callable[..., None]
+
+
+def multi_variant(threads: int, blocks_per_bank: int) -> CopyVariant:
+    def launch(dst, src, feat, di, si, num):
+        fic.fast_index_copy_multi_jit(
+            dst, src, feat, di, si, num, num_threads=threads, blocks_per_bank=blocks_per_bank
+        )
+
+    return CopyVariant(f"multi {threads}x{blocks_per_bank}", launch)
+
+
+def slim_variant(blocks: int, threads: int, unroll: int) -> CopyVariant:
+    def launch(dst, src, feat, di, si, num):
+        fic.fast_index_copy_multi_slim_jit(
+            dst, src, feat, di, si, num, num_blocks=blocks, num_threads=threads, unroll=unroll
+        )
+
+    return CopyVariant(f"slim {blocks}x{threads}u{unroll}", launch)
+
+
+class CopyRig:
+    """Registered host banks, GPU slot caches and the fused-copy descriptor tensors."""
+
+    def __init__(self, feats: list[int], n_src: int, n_slots: int, device: torch.device, seed: int):
+        self.feats = feats
+        self.expert_bytes = sum(feats)
+        self.n_src = n_src
+        self.n_slots = n_slots
+        self.device = device
+        gen = torch.Generator().manual_seed(seed)
+        self.host = []
+        for feat in feats:
+            bank = HostBank((n_src, feat), torch.uint8)
+            bank.tensor.copy_(torch.randint(0, 256, (n_src, feat), dtype=torch.uint8, generator=gen))
+            bank.pin()
+            self.host.append(bank)
+        self.slots = [torch.empty((n_slots, feat), dtype=torch.uint8, device=device) for feat in feats]
+        self.dst_ptrs = torch.tensor([c.data_ptr() for c in self.slots], dtype=torch.int64, device=device)
+        self.src_ptrs = torch.tensor([device_ptr(b.tensor) for b in self.host], dtype=torch.int64, device=device)
+        self.feat_bytes = torch.tensor(feats, dtype=torch.int64, device=device)
+        self.spacer_buf = torch.zeros(1, dtype=torch.int32, device=device)
+        self.gen = gen
+        self._cursor = 0
+        self._perm = torch.randperm(n_src, generator=gen)
+
+    def plans(self, n_launch: int, experts: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """[n_launch, LMAX] dst/src plans: source rows walk a permutation of the host rows
+        (no reuse until all are used), destination slots are distinct within a plan."""
+        src = torch.zeros((n_launch, LMAX), dtype=torch.int32)
+        dst = torch.zeros((n_launch, LMAX), dtype=torch.int32)
+        for r in range(n_launch):
+            for j in range(experts):
+                if self._cursor == self.n_src:
+                    self._perm = torch.randperm(self.n_src, generator=self.gen)
+                    self._cursor = 0
+                src[r, j] = self._perm[self._cursor]
+                self._cursor += 1
+            dst[r, :experts] = torch.randperm(self.n_slots, generator=self.gen)[:experts]
+        num = torch.full((1,), experts, dtype=torch.int64, device=self.device)
+        return dst.to(self.device), src.to(self.device), num
+
+    def copy_loop(self, variant: CopyVariant, n_launch: int, experts: int, cold: bool) -> Callable[[], None]:
+        dst, src, num = self.plans(n_launch, experts)
+
+        def run():
+            for r in range(n_launch):
+                variant.launch(self.dst_ptrs, self.src_ptrs, self.feat_bytes, dst[r], src[r], num)
+                if cold:
+                    self.spacer_buf.add_(1)
+
+        return run
+
+    def spacer_us(self, stream: torch.cuda.Stream, trials: int) -> float:
+        n = 200
+
+        def run():
+            for _ in range(n):
+                self.spacer_buf.add_(1)
+
+        g = capture(run, stream)
+        return statistics.median(time_alone(g, n, stream) for _ in range(trials))
+
+
+@dataclass
+class ComputeOp:
+    name: str
+    make: Callable[[int], Callable[[], None]]  # n_launch -> loop body
+    single: Callable[[], None]
+
+
+def _weight_copies(nbytes: int) -> int:
+    return max(2, math.ceil(2 * L2_BYTES / nbytes))
+
+
+def gemv_bf16(h: int, n: int, device) -> ComputeOp:
+    x = torch.randn(1, h, dtype=torch.bfloat16, device=device)
+    ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))]
+
+    def make(n_launch):
+        def run():
+            for r in range(n_launch):
+                F.linear(x, ws[r % len(ws)])
+
+        return run
+
+    return ComputeOp(f"gemv_bf16 1x{h}x{n}", make, lambda: F.linear(x, ws[0]))
+
+
+def gemv_fp8(h: int, n: int, device) -> ComputeOp:
+    fp8 = torch.float8_e4m3fn
+    x = torch.randn(1, h, device=device).to(fp8)
+    ws = [torch.randn(n, h, device=device).to(fp8) for _ in range(_weight_copies(n * h))]
+    one = torch.ones((), dtype=torch.float32, device=device)
+
+    def mm(w):
+        torch._scaled_mm(x, w.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16)
+
+    def make(n_launch):
+        def run():
+            for r in range(n_launch):
+                mm(ws[r % len(ws)])
+
+        return run
+
+    return ComputeOp(f"gemv_fp8 1x{h}x{n}", make, lambda: mm(ws[0]))
+
+
+def matmul_bf16(m: int, h: int, n: int, device) -> ComputeOp:
+    a = torch.randn(m, h, dtype=torch.bfloat16, device=device)
+    ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))]
+
+    def make(n_launch):
+        def run():
+            for r in range(n_launch):
+                F.linear(a, ws[r % len(ws)])
+
+        return run
+
+    return ComputeOp(f"mm_bf16 {m}x{h}x{n}", make, lambda: F.linear(a, ws[0]))
+
+
+def nvfp4_routed(h: int, i: int, top_k: int, n_slots: int, device) -> ComputeOp:
+    """The decode routed-expert GEMV pair (gate_up -> act -> down) over its own slot banks."""
+    from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin
+
+    fp8 = torch.float8_e4m3fn
+    gu = torch.randint(0, 256, (n_slots, 2 * i, h // 2), dtype=torch.uint8, device=device)
+    gus = torch.full((n_slots, 2 * i, h // 16), 1.0, device=device).to(fp8)
+    gug = torch.full((n_slots, 2 * i), 1.0, dtype=torch.float16, device=device)
+    dn = torch.randint(0, 256, (n_slots, h, i // 2), dtype=torch.uint8, device=device)
+    dns = torch.full((n_slots, h, i // 16), 1.0, device=device).to(fp8)
+    dng = torch.full((n_slots, h), 1.0, dtype=torch.float16, device=device)
+    x = torch.randn(1, h, dtype=torch.bfloat16, device=device) * 0.01
+    w = torch.full((1, top_k), 1.0 / top_k, dtype=torch.float32, device=device)
+    groups = max(1, n_slots // top_k)
+    ids = [
+        (torch.arange(top_k, dtype=torch.int32, device=device) + g * top_k).view(1, top_k) % n_slots
+        for g in range(groups)
+    ]
+
+    def call(t):
+        fused_experts_decode_nvfp4_marlin(x, gu, gus, gug, dn, dns, dng, w, t, "silu", False)
+
+    def make(n_launch):
+        def run():
+            for r in range(n_launch):
+                call(ids[r % groups])
+
+        return run
+
+    return ComputeOp(f"nvfp4_routed k={top_k}", make, lambda: call(ids[0]))
+
+
+def capture(fn: Callable[[], None], stream: torch.cuda.Stream) -> torch.cuda.CUDAGraph:
+    with torch.cuda.stream(stream):
+        fn()  # warm-up (JIT / Triton compile, cuBLAS heuristics) outside capture
+    stream.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g, stream=stream):
+        fn()
+    torch.cuda.synchronize()
+    g.keepalive = fn  # the closure owns the plan tensors the graph reads
+    return g
+
+
+def _event() -> torch.cuda.Event:
+    return torch.cuda.Event(enable_timing=True)
+
+
+def time_alone(g: torch.cuda.CUDAGraph, n: int, stream: torch.cuda.Stream) -> float:
+    """us per launch of a graph holding ``n`` launches, replayed alone."""
+    start, end = _event(), _event()
+    with torch.cuda.stream(stream):
+        torch.cuda._sleep(SLEEP_CYCLES)
+        start.record(stream)
+        g.replay()
+        end.record(stream)
+    end.synchronize()
+    return start.elapsed_time(end) * 1e3 / n
+
+
+def time_corun(
+    g_fg: torch.cuda.CUDAGraph, n_fg: int, s_fg: torch.cuda.Stream,
+    g_bg: torch.cuda.CUDAGraph, s_bg: torch.cuda.Stream,
+) -> tuple[float, bool]:
+    """us per launch of the foreground graph while the background graph runs on another
+    stream, and whether the background outlasted the foreground (a valid co-run)."""
+    ev0, bg_end, fg_start, fg_end = _event(), _event(), _event(), _event()
+    s0 = torch.cuda.current_stream()
+    torch.cuda._sleep(SLEEP_CYCLES)
+    ev0.record(s0)
+    s_bg.wait_event(ev0)
+    s_fg.wait_event(ev0)
+    with torch.cuda.stream(s_bg):
+        g_bg.replay()
+        bg_end.record(s_bg)
+    with torch.cuda.stream(s_fg):
+        torch.cuda._sleep(SLEEP_CYCLES // 20)  # let the background reach steady state
+        fg_start.record(s_fg)
+        g_fg.replay()
+        fg_end.record(s_fg)
+    torch.cuda.synchronize()
+    covered = ev0.elapsed_time(bg_end) >= ev0.elapsed_time(fg_end)
+    return fg_start.elapsed_time(fg_end) * 1e3 / n_fg, covered
+
+
+def kernel_info(fn: Callable[[], None]) -> list[dict]:
+    """(name, grid, block, regs/thread, smem) of every kernel one call launches."""
+    from torch.profiler import ProfilerActivity, profile
+
+    fn()
+    torch.cuda.synchronize()
+    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+        fn()
+        torch.cuda.synchronize()
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "trace.json")
+        prof.export_chrome_trace(path)
+        with open(path) as f:
+            events = json.load(f)["traceEvents"]
+    out = []
+    for e in events:
+        if e.get("cat") != "kernel":
+            continue
+        a = e.get("args", {})
+        out.append({
+            "name": e["name"][:80],
+            "grid": a.get("grid"),
+            "block": a.get("block"),
+            "regs": a.get("registers per thread"),
+            "smem": a.get("shared memory"),
+        })
+    return out
+
+
+def _fmt_info(info: list[dict]) -> str:
+    return "; ".join(
+        f"{k['name'][:48]} grid={k['grid']} block={k['block']} regs={k['regs']} smem={k['smem']}" for k in info
+    )
+
+
+def _stats(xs: list[float]) -> tuple[float, float]:
+    return min(xs), statistics.median(xs)
+
+
+def parse_pairs(text: str, n: int) -> list[tuple[int, ...]]:
+    out = []
+    for item in text.split(","):
+        item = item.strip()
+        if not item:
+            continue
+        parts = tuple(int(p) for p in item.lower().replace("u", "x").split("x"))
+        if len(parts) != n:
+            raise argparse.ArgumentTypeError(f"{item!r}: expected {n} 'x'-separated ints")
+        out.append(parts)
+    return out
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--gpu", type=single_gpu_arg, default=None,
+                   help="GPU UUID or nvidia-smi index (default: the first visible GPU)")
+    p.add_argument("--hidden", type=int, default=2560)
+    p.add_argument("--inter", type=int, default=640)
+    p.add_argument("--src-experts", type=int, default=256, help="host rows per bank (pinned)")
+    p.add_argument("--slots", type=int, default=64, help="GPU slot-cache rows per bank")
+    p.add_argument("--multi", default="1024x8,1024x4,512x8,512x4,256x16,256x8,128x32",
+                   help="fast_index_copy_multi configs, THREADSxBLOCKS_PER_BANK")
+    p.add_argument("--slim", default="8x256x4,16x256x2,12x256x4,16x256x4,8x256x2,4x256x4",
+                   help="fast_index_copy_multi_slim configs, BLOCKSxTHREADSxUNROLL "
+                        "(skipped if the kernel is absent)")
+    p.add_argument("--experts", default="1,2,3,4,5,6,7,8,9,10", help="plan sizes for the sweep")
+    p.add_argument("--corun-experts", default="4", help="plan sizes for the co-run copy loop")
+    p.add_argument("--link", choices=["hot", "cold", "both"], default="both", help="sweep link mode(s)")
+    p.add_argument("--corun-link", choices=["hot", "cold"], default="cold", help="co-run copy loop link mode")
+    p.add_argument("--trials", type=int, default=5)
+    p.add_argument("--copy-ms", type=float, default=6.0, help="target length of one timed copy loop")
+    p.add_argument("--compute-ms", type=float, default=3.0, help="target length of one timed compute loop")
+    p.add_argument("--cover", type=float, default=3.0,
+                   help="background loop length as a multiple of the foreground's alone time")
+    p.add_argument("--skip-sweep", action="store_true")
+    p.add_argument("--skip-corun", action="store_true")
+    p.add_argument("--ops", default="all",
+                   help="comma list of co-run ops by prefix: gemv_bf16,gemv_fp8,mm_bf16,nvfp4_routed")
+    p.add_argument("--json", default=None, help="write all results to this JSON path")
+    p.add_argument("--seed", type=int, default=0)
+    return p.parse_args()
+
+
+def build_variants(args) -> list[CopyVariant]:
+    variants = [multi_variant(t, b) for t, b in parse_pairs(args.multi, 2)]
+    if hasattr(fic, "fast_index_copy_multi_slim_jit"):
+        variants += [slim_variant(g, t, u) for g, t, u in parse_pairs(args.slim, 3)]
+    else:
+        print("fast_index_copy_multi_slim_jit not available: slim configs skipped", flush=True)
+    return variants
+
+
+def run_sweep(rig: CopyRig, variants, experts_list, args, stream, spacer_us, results) -> None:
+    modes = ["hot", "cold"] if args.link == "both" else [args.link]
+    for mode in modes:
+        cold = mode == "cold"
+        sub = spacer_us if cold else 0.0
+        print(f"\n== sweep ({mode} link{f', spacer {sub:.2f} us subtracted' if cold else ''}): copy alone, "
+              f"expert={rig.expert_bytes} B, trials={args.trials}; cells are GB/s at min-time / "
+              f"median-time (us/copy median)", flush=True)
+        print("variant".ljust(18) + "".join(f"{e:>20d}" for e in experts_list), flush=True)
+        for v in variants:
+            row = v.name.ljust(18)
+            for e in experts_list:
+                n_launch = max(8, math.ceil(args.copy_ms * 1e3 / (e * 55.0)))
+                g = capture(rig.copy_loop(v, n_launch, e, cold), stream)
+                ts = [time_alone(g, n_launch, stream) - sub for _ in range(args.trials)]
+                t_min, t_med = _stats(ts)
+                nbytes = e * rig.expert_bytes
+                results["sweep"].append({
+                    "link": mode, "variant": v.name, "experts": e, "us_min": t_min, "us_median": t_med,
+                    "gbps_best": nbytes / (t_min * 1e3), "gbps_median": nbytes / (t_med * 1e3), "us_all": ts,
+                })
+                row += f"{nbytes / (t_min * 1e3):>7.1f}/{nbytes / (t_med * 1e3):5.1f} ({t_med:5.0f})"
+                del g
+            print(row, flush=True)
+
+
+def run_corun(rig: CopyRig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) -> None:
+    cold = args.corun_link == "cold"
+    sub = spacer_us if cold else 0.0
+    print(f"\n== co-run ({args.corun_link} link copy loop): compute on one stream, copy loop on another; "
+          f"trials={args.trials} (interleaved alone/co-run); sigma = t_corun / t_alone - 1", flush=True)
+    for op in ops:
+        n_fg_probe = 20
+        g_probe = capture(op.make(n_fg_probe), s_fg)
+        t_probe = statistics.median(time_alone(g_probe, n_fg_probe, s_fg) for _ in range(3))
+        del g_probe
+        n_op = max(20, math.ceil(args.compute_ms * 1e3 / t_probe))
+        g_op = capture(op.make(n_op), s_fg)
+        print(f"\n-- {op.name}: ~{t_probe:.1f} us/launch alone, loop of {n_op}", flush=True)
+        results["kernels"][op.name] = kernel_info(op.single)
+        print(f"   kernels: {_fmt_info(results['kernels'][op.name])}", flush=True)
+        print("   " + "copy variant".ljust(18) + "E".rjust(3)
+              + "  op_alone_us(min/med)  op_corun_us(min/med)  sigma(min/med)"
+              + "  copy_GBps alone(med) -> under op(med)  valid  (! = alone median >10% over the probe:"
+              + " other GPU load, trust the min column)", flush=True)
+        for v in variants:
+            for e in corun_experts:
+                copy_launch_us = e * 55.0
+                n_copy_long = max(4, math.ceil(args.cover * n_op * t_probe / copy_launch_us))
+                g_copy_long = capture(rig.copy_loop(v, n_copy_long, e, cold), s_bg)
+                n_copy = max(8, math.ceil(args.copy_ms * 1e3 / copy_launch_us))
+                g_copy = capture(rig.copy_loop(v, n_copy, e, cold), s_bg)
+                n_op_long = max(n_op, math.ceil(args.cover * n_copy * copy_launch_us / t_probe))
+                g_op_long = capture(op.make(n_op_long), s_fg) if n_op_long > n_op else g_op
+                op_alone, op_co, cp_alone, cp_co = [], [], [], []
+                valid_op = valid_cp = 0
+                for _ in range(args.trials):
+                    op_alone.append(time_alone(g_op, n_op, s_fg))
+                    t, ok = time_corun(g_op, n_op, s_fg, g_copy_long, s_bg)
+                    op_co.append(t)
+                    valid_op += ok
+                    cp_alone.append(time_alone(g_copy, n_copy, s_bg) - sub)
+                    t, ok = time_corun(g_copy, n_copy, s_bg, g_op_long, s_fg)
+                    cp_co.append(t - sub)
+                    valid_cp += ok
+                a_min, a_med = _stats(op_alone)
+                c_min, c_med = _stats(op_co)
+                sig_min, sig_med = c_min / a_min - 1, c_med / a_med - 1
+                nbytes = e * rig.expert_bytes
+                gb_alone = nbytes / (statistics.median(cp_alone) * 1e3)
+                gb_co = nbytes / (statistics.median(cp_co) * 1e3)
+                suspect = a_med > 1.1 * t_probe
+                results["corun"].append({
+                    "link": args.corun_link, "op": op.name, "variant": v.name, "experts": e, "suspect": suspect,
+                    "op_alone_us": op_alone, "op_corun_us": op_co,
+                    "sigma_min": sig_min, "sigma_median": sig_med,
+                    "copy_alone_us": cp_alone, "copy_corun_us": cp_co,
+                    "copy_gbps_alone_median": gb_alone, "copy_gbps_corun_median": gb_co,
+                    "valid_op_trials": valid_op, "valid_copy_trials": valid_cp,
+                })
+                print(f"   {v.name.ljust(18)}{e:>3d}  {a_min:8.2f}/{a_med:8.2f}    {c_min:8.2f}/{c_med:8.2f}"
+                      f"    {sig_min:+6.3f}/{sig_med:+6.3f}   {gb_alone:6.1f} -> {gb_co:6.1f}"
+                      f"          {valid_op}/{valid_cp}{'  !' if suspect else ''}", flush=True)
+                del g_copy_long, g_copy, g_op_long
+        del g_op
+
+
+def main() -> None:
+    args = parse_args()
+    assert torch.cuda.is_available(), "CUDA is required"
+    try:
+        assign_gpu(args.gpu)
+        device = bind_assigned_gpu()
+    except (ValueError, RuntimeError) as e:
+        raise SystemExit(f"error: {e}") from e
+
+    feats = nvfp4_bank_bytes(args.hidden, args.inter)
+    assert sum(feats) == _BANK_BYTES_PER_EXPERT["nvfp4"](args.hidden, args.inter)
+    if (args.hidden, args.inter) == (2560, 640):
+        assert sum(feats) == 2_772_480
+    print("gpu", torch.cuda.get_device_name(device), "| banks", feats, "| expert", sum(feats), "B", flush=True)
+    rig = CopyRig(feats, args.src_experts, args.slots, device, args.seed)
+    variants = build_variants(args)
+    results = {"feats": feats, "sweep": [], "corun": [], "kernels": {}}
+
+    s_fg = torch.cuda.Stream(device=device)
+    s_bg = torch.cuda.Stream(device=device)
+    spacer_us = rig.spacer_us(s_bg, args.trials)
+    results["spacer_us"] = spacer_us
+    print(f"spacer kernel alone: {spacer_us:.2f} us", flush=True)
+    for v in variants:
+        dst, src, num = rig.plans(1, 4)
+        info = kernel_info(lambda: v.launch(rig.dst_ptrs, rig.src_ptrs, rig.feat_bytes, dst[0], src[0], num))
+        results["kernels"][v.name] = info
+        print(f"{v.name.ljust(18)} {_fmt_info(info)}", flush=True)
+
+    if not args.skip_sweep:
+        experts_list = [int(x) for x in args.experts.split(",") if x]
+        run_sweep(rig, variants, experts_list, args, s_bg, spacer_us, results)
+
+    if not args.skip_corun:
+        h, i = args.hidden, args.inter
+        all_ops = [
+            gemv_bf16(h, 5120, device), gemv_bf16(h, 12288, device),
+            gemv_fp8(h, 5120, device), gemv_fp8(h, 12288, device),
+            matmul_bf16(256, h, h, device),
+            nvfp4_routed(h, i, 10, args.slots, device),
+        ]
+        wanted = [w.strip() for w in args.ops.split(",") if w.strip()]
+        ops = all_ops if wanted == ["all"] else [o for o in all_ops if any(o.name.startswith(w) for w in wanted)]
+        corun_experts = [int(x) for x in args.corun_experts.split(",") if x]
+        run_corun(rig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results)
+
+    if args.json:
+        with open(args.json, "w") as f:
+            json.dump(results, f, indent=1)
+        print(f"\nwrote {args.json}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/bench_moe_copy_overlap.py b/benchmarks/bench_moe_copy_overlap.py
new file mode 100644
index 0000000000..535f7b1bd2
--- /dev/null
+++ b/benchmarks/bench_moe_copy_overlap.py
@@ -0,0 +1,233 @@
+"""Decode MoE copy/compute overlap bench (FREETOKEN_MOE_COPY_OVERLAP).
+
+Qwen4ExpMoE layers at the Qwen3.8-Flash-Next expert geometry (hidden 2560, moe and shared
+intermediate 640, top-10, NVFP4 triton banks: 2,772,480 B per expert) over one offload cache,
+with fewer experts per layer than the model so host and GPU memory stay small. Each arm
+(overlap off / on) captures one decode step through every layer as a CUDA graph, then:
+  1. bitwise check of the two arms over the same fresh inputs,
+  2. replay time per step with CUDA events, arms interleaved ABAB, fresh inputs every round,
+  3. a torch.profiler (CUPTI) trace of graph replays: per layer, the copy kernel's start slip
+     after the ensure, the shared expert's span (side stream) and whether it ends inside the copy, the join
+     gap, and the copy's duration against the off arm.
+No attention or hyper-connection work runs, so the percentages overstate the model's.
+
+Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_copy_overlap.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+from types import SimpleNamespace
+
+import torch
+
+H, I, K = 2560, 640, 10
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--layers", type=int, default=4)
+    parser.add_argument("--experts", type=int, default=64, help="experts per layer")
+    parser.add_argument("--slots", type=int, default=160, help="cache slots (160 of 4 x 64 gives ~4.4 misses/layer)")
+    parser.add_argument("--bs", type=int, default=1)
+    parser.add_argument("--steps", type=int, default=40, help="decode steps per timing round")
+    parser.add_argument("--rounds", type=int, default=6)
+    return parser.parse_args()
+
+
+def nvfp4_quant():
+    from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend
+
+    set_quant_backend(QuantBackend.parse("moe.nvfp4=triton"))
+    return QuantConfig.from_hf(
+        {"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}
+    )
+
+
+def build_arm(args, overlap: bool, banks) -> SimpleNamespace:
+    import freetoken.core as core
+    from freetoken.core import Context, set_global_ctx
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE
+    from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.utils.torch_utils import torch_dtype
+
+    E, L = args.experts, args.layers
+    config = SimpleNamespace(
+        hidden_size=H, num_experts=E, num_experts_per_tok=K, moe_intermediate_size=I,
+        shared_expert_intermediate_size=I, norm_topk_prob=True, quant=None,
+        moe_strategy="offload", decode_target="gpu",
+    )
+    device = torch.device("cuda")
+    with torch.device(device), torch_dtype(torch.bfloat16):
+        moes = [Qwen4ExpMoE(config, l, prefix=f"model.layers.{l}.mlp") for l in range(L)]
+    gen = torch.Generator(device=device).manual_seed(11)
+    for moe in moes:
+        for t in moe.state_dict().values():
+            if t.is_floating_point():
+                t.normal_(0.0, 0.02, generator=gen)
+    quant = nvfp4_quant()
+    experts = [
+        OffloadMoELayer(l, E, K, H, I, quant_config=quant, prefix=f"model.layers.{l}.mlp.experts")
+        for l in range(L)
+    ]
+    cache = OffloadMoeCache(
+        num_layers=L, num_experts=E, cache_size=args.slots, device=device,
+        quant_format=banks.quant_format, layout=banks.layout,
+        max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap,
+    )
+    cache.set_bank_sources(banks.sources)
+    cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
+    cache.collect_stats = True
+    for moe, ex in zip(moes, experts):
+        ex.offload_cache = cache
+        moe.experts = ex
+    cache.reset()
+    core._GLOBAL_CTX = None
+    ctx = Context(page_size=1)
+    set_global_ctx(ctx)
+    ctx._batch = SimpleNamespace(is_prefill=False)
+
+    static = torch.randn(args.bs, H, device=device, dtype=torch.bfloat16)
+    out = torch.empty(L, args.bs, H, device=device, dtype=torch.bfloat16)
+
+    def step():
+        for l, moe in enumerate(moes):
+            out[l].copy_(moe.forward(static.clone()))
+
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        step()  # the eager warm-up the graph runner also does
+    torch.cuda.current_stream().wait_stream(side)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        step()
+    cache.reset()
+    cache.reset_stats()
+    # the graph holds raw pointers only: the modules must outlive it or replays read freed weights
+    return SimpleNamespace(graph=graph, static=static, out=out, cache=cache, moes=moes)
+
+
+def fresh_inputs(args, n: int, seed: int) -> list[torch.Tensor]:
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    return [torch.randn(args.bs, H, device="cuda", dtype=torch.bfloat16, generator=gen) for _ in range(n)]
+
+
+def replay(arm, inputs, keep: bool = False) -> list[torch.Tensor]:
+    outs = []
+    for x in inputs:
+        arm.static.copy_(x)
+        arm.graph.replay()
+        if keep:
+            outs.append(arm.out.clone())
+    return outs
+
+
+def time_arm(arm, inputs) -> float:
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    start.record()
+    replay(arm, inputs)
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / len(inputs)
+
+
+def trace(arm, inputs) -> list[tuple[float, float, str]]:
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+        replay(arm, inputs)
+        torch.cuda.synchronize()
+    return sorted(
+        (e.time_range.start, e.time_range.end, e.name) for e in prof.events() if e.device_type.name == "CUDA"
+    )
+
+
+def layer_calls(ev, overlap: bool) -> list[dict]:
+    """Per MoE layer call, located by kernel name (graph replay spreads branches over streams)."""
+    rows = []
+    for i, (_, ens_end, name) in enumerate(ev):
+        if "_lru_ensure_kernel" not in name:
+            continue
+        copy = next((k for k in ev[i + 1:] if "fast_index_copy" in k[2]), None)
+        gemm = next((k for k in ev[i + 1:] if "_decode_nvfp4" in k[2]), None)
+        if copy is None or gemm is None:
+            continue
+        # the shared expert chain starts at the gate kernel before the ensure; with the overlap it runs on
+        # the side stream beside the routed path, without it serially ahead of that path's topk
+        j = max(j for j in range(i) if "_gate_sigmoid_kernel" in ev[j][2])
+        routed_path = ("router", "topk", "ensure", "fast_index_copy")
+        chain = [k for k in (ev[j:] if overlap else ev[j:i]) if k[1] <= gemm[0]
+                 and not any(m in k[2] for m in routed_path)]
+        rows.append({
+            "copy_us": copy[1] - copy[0],
+            "slip_us": copy[0] - ens_end,
+            "join_gap_us": gemm[0] - copy[1],
+            "shared_span_us": max(k[1] for k in chain) - chain[0][0],
+            "shared_ends_in_copy": max(k[1] for k in chain) <= copy[1],
+        })
+    return rows
+
+
+def pct(xs, q: float) -> float:
+    xs = sorted(xs)
+    return xs[min(len(xs) - 1, int(q * len(xs)))]
+
+
+def main() -> None:
+    args = parse_args()
+    from freetoken.distributed import set_tp_info
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.moe.expert_banks import build_expert_banks
+
+    set_tp_info(rank=0, size=1)
+    torch.manual_seed(0)
+    method = OffloadMoELayer(0, args.experts, K, H, I, quant_config=nvfp4_quant(), prefix="model.layers.0.mlp.experts").quant_method
+    banks = build_expert_banks(method, args.layers, None, device=torch.device("cuda"), dummy=True)
+    expert_bytes = sum(layers[0][0].numel() * layers[0].element_size() for layers in banks.sources.values())
+    print(f"expert {expert_bytes} B; {args.layers} layers x {args.experts} experts, {args.slots} slots, bs {args.bs}")
+    arms = {False: build_arm(args, False, banks), True: build_arm(args, True, banks)}
+
+    seed = 100
+    inputs = fresh_inputs(args, args.steps, seed)
+    outs = {k: replay(arm, inputs, keep=True) for k, arm in arms.items()}
+    same = all(torch.equal(a, b) for a, b in zip(outs[True], outs[False]))
+    active, missed, calls = arms[False].cache.lru_stats.sum(0).tolist()
+    print(f"bitwise identical on/off over {args.steps} steps: {same}; misses per layer call {missed / calls:.2f}")
+
+    times = {False: [], True: []}
+    for r in range(args.rounds):
+        seed += 1
+        inputs = fresh_inputs(args, args.steps, seed)
+        for k in (False, True) if r % 2 == 0 else (True, False):
+            times[k].append(time_arm(arms[k], inputs))
+    for k in (False, True):
+        t = times[k]
+        print(f"overlap {'on ' if k else 'off'}: {statistics.median(t):.4f} ms/step median (min {min(t):.4f}, max {max(t):.4f}, n={len(t)})")
+    per_layer = [1000 * (a - b) / args.layers for a, b in zip(times[False], times[True])]
+    print(f"saving per layer, paired rounds (us): {[round(x, 1) for x in per_layer]}, median {statistics.median(per_layer):.1f}")
+
+    inputs = fresh_inputs(args, 12, seed + 1)
+    rows = {k: layer_calls(trace(arms[k], inputs), k) for k in (False, True)}
+    for k in (False, True):
+        r = [x for x in rows[k] if x["copy_us"] > 5.0]  # layer calls with at least one miss
+        if not r:
+            continue
+        slips = [x["slip_us"] for x in r]
+        print(
+            f"overlap {'on ' if k else 'off'}: {len(r)} layer calls with misses; copy {statistics.median(x['copy_us'] for x in r):.1f} us; "
+            f"copy start - ensure end median {statistics.median(slips):.2f} us (p90 {pct(slips, 0.9):.2f}, <=5 us in {100 * sum(s <= 5 for s in slips) / len(slips):.0f}%); "
+            f"GEMM start - copy end median {statistics.median(x['join_gap_us'] for x in r):.2f} us; "
+            f"shared expert span median {statistics.median(x['shared_span_us'] for x in r):.1f} us"
+            + (f", ends inside the copy in {100 * sum(x['shared_ends_in_copy'] for x in r) / len(r):.0f}%" if k else "")
+        )
+    pairs = [(a["copy_us"], b["copy_us"]) for a, b in zip(rows[False], rows[True]) if a["copy_us"] > 5.0]
+    if pairs:
+        ratio = [b / a for a, b in pairs]
+        print(f"copy duration on/off: median {statistics.median(ratio):.3f} (p10 {pct(ratio, 0.1):.3f}, p90 {pct(ratio, 0.9):.3f})")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/bench_scored_ensure.py b/benchmarks/bench_scored_ensure.py
new file mode 100644
index 0000000000..5aa852716e
--- /dev/null
+++ b/benchmarks/bench_scored_ensure.py
@@ -0,0 +1,158 @@
+"""Decode ensure cost per call: flashlib lru_ensure vs the scored --moe-cache-policy kernels.
+
+One ensure call per MoE layer per decode step, on a full cache (S slots over L x E experts).
+Each timed call routes K ids of which exactly m miss: the hits are hot, recently used experts
+and the victims are cold filler, under every policy. A CUDA graph of 2 decode steps (2 * L
+calls; 1 step when that many fresh ids would not fit) is replayed from a restored snapshot,
+so the number is the graph-node cost the decode graph pays. The miss count is verified
+through the kernels' stats before timing.
+
+Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_scored_ensure.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+
+import numpy as np
+import torch
+from flashlib.kernels.slot_cache import N_STATS, Stat, lru_ensure
+
+from freetoken.moe import scored_ensure as se
+
+L, E, TOP_K = 48, 512, 10
+VARIANTS = ("flashlib-lru", "lru", "kd", "kdfb", "rule")
+
+
+def build(variant: str, slots: int, k: int, m: int, dev: torch.device):
+    """Full-cache state plus 1-2 steps of L queries of k ids, the first k - m hot, the rest fresh misses."""
+    policy = 0 if variant.endswith("lru") else se.POLICY_IDS[variant]
+    hot = k - m
+    calls = 2 * L if hot * L + 2 * L * m <= slots else L
+    step0, tok0 = 1_000_000, 1000
+    slot_of_id = np.full(L * E, -1, np.int32)
+    id_of_slot = np.full(slots, -1, np.int32)
+    usage = np.zeros(slots, np.int64)
+    last_tok = np.full(L * E, se.LAST_TOK_NEVER, np.int64)
+    lc = np.full(L * E, se.LC_NEVER, np.int32)
+    ct = np.zeros(L * E, np.int64)
+    ids = [layer * E + e for layer in range(L) for e in range(hot)]
+    assert len(ids) + calls * m <= slots, "fresh ids would evict hot ones"
+    filler = (layer * E + e for e in range(k + m, E) for layer in range(L))
+    ids += [next(filler) for _ in range(slots - len(ids))]
+    for s, key in enumerate(ids):
+        id_of_slot[s], slot_of_id[key] = key, s
+        is_hot = key % E < hot
+        usage[s] = step0 - 1 if is_hot else 1 + s
+        last_tok[key] = tok0 if is_hot else tok0 - 500 - s % 7
+        if is_hot:
+            lc[key], ct[key] = 5 << se.Q, tok0 * L + key // E
+    queries = np.empty((calls, k), np.int32)
+    for i in range(calls):
+        visit = i // L
+        queries[i] = np.r_[np.arange(hot), hot + visit * m + np.arange(m)]
+    rng = np.random.default_rng(0)
+    logits = rng.normal(size=(calls, k // TOP_K, E)).astype(np.float32)
+    rows = queries.reshape(calls, k // TOP_K, -1).astype(np.int64)
+    np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, -1) + 3.0, -1)
+    held = id_of_slot.astype(np.int64)
+    t = lambda a, dtype=None: torch.from_numpy(a).to(dev, dtype)  # noqa: E731
+    state = dict(
+        owner=t(id_of_slot.copy()), slot_last=t(last_tok[held]), slot_lc=t(lc[held]), slot_ct=t(ct[held]),
+        slot_of_id=t(slot_of_id), id_of_slot=t(id_of_slot), usage=t(usage),
+        step=torch.tensor(step0, dtype=torch.int64, device=dev),
+        tok=torch.tensor(tok0, dtype=torch.int64, device=dev),
+        last_tok=t(last_tok), lc=t(lc), ct=t(ct), queries=t(queries),
+        src=torch.zeros(slots, dtype=torch.int32, device=dev),
+        dst=torch.zeros(slots, dtype=torch.int32, device=dev),
+        num=torch.zeros((), dtype=torch.int64, device=dev),
+        stats=torch.zeros(L, N_STATS, dtype=torch.int64, device=dev),
+    )
+    return policy, state, t(logits, torch.bfloat16), se.softplus2_table(dev)
+
+
+def step_fn(variant, policy, st, logits, g_table, stats: bool):
+    def run():
+        for i in range(st["queries"].shape[0]):
+            layer = i % L
+            q = st["queries"][i]
+            s = st["stats"][layer] if stats else None
+            if variant == "flashlib-lru":
+                lru_ensure(q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q,
+                           st["src"], st["dst"], st["num"], stats=s, id_base=layer * E)
+                continue
+            se.scored_ensure(
+                q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q,
+                st["src"], st["dst"], st["num"], policy=policy, num_layers=L, num_experts=E,
+                tok=st["tok"], last_tok=st["last_tok"], lc=st["lc"], ct=st["ct"], g_table=g_table,
+                slot_owner=st["owner"], slot_last_tok=st["slot_last"], slot_lc=st["slot_lc"], slot_ct=st["slot_ct"],
+                router_logits=logits[i] if policy == 3 else None, near_miss_thr=0.25,
+                stats=s, id_base=layer * E, bump_tok=layer == 0,
+            )
+    return run
+
+
+def bench(variant, slots, k, m, dev, reps, stats):
+    policy, st, logits, g_table = build(variant, slots, k, m, dev)
+    snap = {name: v.clone() for name, v in st.items()}
+
+    def restore():
+        for name, v in snap.items():
+            st[name].copy_(v)
+
+    # verify the miss count (a stats build of the same calls), then time the requested build
+    restore()
+    step_fn(variant, policy, st, logits, g_table, True)()
+    got = int(st["stats"][:, Stat.MISS].sum())
+    calls = st["queries"].shape[0]
+    assert got == calls * m, f"{variant} K={k} m={m}: {got} misses, expected {calls * m}"
+    restore()
+    run = step_fn(variant, policy, st, logits, g_table, stats)
+    run()  # compile outside capture
+    restore()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    times = []
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    for rep in range(reps + 3):
+        restore()
+        torch.cuda.synchronize()
+        start.record()
+        graph.replay()
+        end.record()
+        end.synchronize()
+        if rep >= 3:
+            times.append(start.elapsed_time(end) * 1000.0 / calls)
+    return statistics.median(times)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--slots", type=int, default=1650)
+    ap.add_argument("--k", default="10,20", help="routed ids per call (10 = bs 1, 20 = bs 2)")
+    ap.add_argument("--misses", default="0-10", help="list or range; capped at K per call")
+    ap.add_argument("--variants", default=",".join(VARIANTS))
+    ap.add_argument("--reps", type=int, default=30)
+    ap.add_argument("--stats", action="store_true", help="time with --moe-collect-stats accumulation on")
+    args = ap.parse_args()
+    dev = torch.device("cuda")
+    variants = args.variants.split(",")
+    misses = []
+    for part in args.misses.split(","):
+        lo, _, hi = part.partition("-")
+        misses += range(int(lo), int(hi or lo) + 1)
+    print(f"# {torch.cuda.get_device_name(dev)}, S={args.slots}, us per ensure call "
+          f"(median of {args.reps} graph replays), stats={'on' if args.stats else 'off'}")
+    print(f"{'K':>3} {'m':>3} " + " ".join(f"{v:>13}" for v in variants) + "  max delta vs flashlib-lru")
+    for k in (int(x) for x in args.k.split(",")):
+        for m in (x for x in misses if x <= k):
+            us = {v: bench(v, args.slots, k, m, dev, args.reps, args.stats) for v in variants}
+            base = us.get("flashlib-lru")
+            delta = "" if base is None else f"  {max(us[v] - base for v in variants):+.2f}"
+            print(f"{k:>3} {m:>3} " + " ".join(f"{us[v]:>13.2f}" for v in variants) + delta, flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/freetoken/env.py b/python/freetoken/env.py
index c8eb1e7513..0e56633f48 100644
--- a/python/freetoken/env.py
+++ b/python/freetoken/env.py
@@ -77,6 +77,8 @@ class EnvClassSingleton:
     MAMBA_SSM_DTYPE = EnvStr("float32")
     # cap the longest prefill warmup forward; 0 = the engine's own chunk cap
     WARMUP_MAX_LEN = EnvInt(0)
+    # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert
+    MOE_COPY_OVERLAP = EnvBool(False)
 
     def __new__(cls):
         # single instance
diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py
index 76250af7b6..c9e378d5b6 100644
--- a/python/freetoken/kernel/aot.py
+++ b/python/freetoken/kernel/aot.py
@@ -124,6 +124,23 @@ def build(build_directory: pathlib.Path) -> object:
     return KernelSpec(name=_make_name("fast_index_copy_multi", *args), build=build)
 
 
+def _fast_index_copy_multi_slim_spec() -> KernelSpec:
+    from .fast_index_copy import SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL
+
+    args = make_cpp_args(SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL)
+
+    def build(build_directory: pathlib.Path) -> object:
+        return load_jit(
+            "fast_index_copy_multi_slim",
+            *args,
+            cuda_files=["fast_index_copy.cuh"],
+            cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")],
+            build_directory=str(build_directory),
+        )
+
+    return KernelSpec(name=_make_name("fast_index_copy_multi_slim", *args), build=build)
+
+
 def _batch_memcpy_spec() -> KernelSpec:
     def build(build_directory: pathlib.Path) -> object:
         return load_jit(
@@ -165,6 +182,8 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]:
     specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32))
     # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding.
     specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64))
+    # opt-in decode copy (FREETOKEN_MOE_SLIM_COPY=1)
+    specs.append(_fast_index_copy_multi_slim_spec())
     specs.append(_batch_memcpy_spec())
     specs.append(_radix_spec())
     specs.append(_ggml_spec())
diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
index 43f3fc1ba5..e37e3d48ad 100644
--- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
+++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
@@ -595,3 +595,112 @@ struct MultiIndexCopyKernel {
                      device.unwrap())(kernel, params);
     }
 };
+
+
+// ---------------------------------------------------------------------------
+// Slim multi-bank copy: the fused copy from a small grid, so it can share the GPU with compute. The
+// grid strides each (row, bank) segment in turn, carrying the phase across, so no bank runs alone at the tail.
+template <std::size_t kUnroll>
+__always_inline __device__ void slim_copy_segment(
+    uint4* __restrict__ dst, const uint4* __restrict__ src, int32_t count, int32_t i, int32_t stride
+) {
+    constexpr int32_t kU = static_cast<int32_t>(kUnroll);
+    // kUnroll independent loads in flight per thread: host-link latency, not issue rate, bounds a small grid
+    for (; i + (kU - 1) * stride < count; i += kU * stride) {
+        uint4 v[kUnroll];
+#pragma unroll
+        for (int32_t k = 0; k < kU; ++k) {
+            v[k] = details::load_nc(src + i + k * stride);
+        }
+#pragma unroll
+        for (int32_t k = 0; k < kU; ++k) {
+            dst[i + k * stride] = v[k];
+        }
+    }
+    uint4 v[kUnroll];
+#pragma unroll
+    for (int32_t k = 0; k < kU; ++k) {
+        if (i + k * stride < count) {
+            v[k] = details::load_nc(src + i + k * stride);
+        }
+    }
+#pragma unroll
+    for (int32_t k = 0; k < kU; ++k) {
+        if (i + k * stride < count) {
+            dst[i + k * stride] = v[k];
+        }
+    }
+}
+
+template <typename IdType, std::size_t kNumThreads, std::size_t kUnroll>
+__global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi_slim(
+    const __grid_constant__ MultiIndexCopyParams p
+) {
+    const int64_t n = p.valid_length ? p.valid_length[0] : p.length;
+    const auto stride = static_cast<int32_t>(gridDim.x * kNumThreads);
+    const auto tid = static_cast<int32_t>(blockIdx.x * kNumThreads + threadIdx.x);
+    const auto* di = static_cast<const IdType*>(p.dst_indices);
+    const auto* si = static_cast<const IdType*>(p.src_indices);
+    int32_t phase = 0;  // grid thread that takes the next segment's first chunk
+    for (int64_t row = 0; row < n; ++row) {
+        const auto ps = static_cast<int64_t>(si[row]);
+        const auto pd = static_cast<int64_t>(di[row]);
+        for (int b = 0; b < p.num_banks; ++b) {
+            const int64_t feat = p.feat_bytes[b];
+            const auto units = static_cast<int32_t>(feat >> 4);  // feat % 16 == 0; int32 chunk index: rows < 32 GiB
+            const auto* src = reinterpret_cast<const uint4*>(reinterpret_cast<const uint8_t*>(p.src_ptrs[b]) + ps * feat);
+            auto* dst = reinterpret_cast<uint4*>(reinterpret_cast<uint8_t*>(p.dst_ptrs[b]) + pd * feat);
+            const int32_t first = tid >= phase ? tid - phase : tid - phase + stride;
+            slim_copy_segment<kUnroll>(dst, src, units, first, stride);
+            phase = static_cast<int32_t>((static_cast<int64_t>(phase) + units) % stride);
+        }
+    }
+}
+
+template <std::size_t kNumBlocks, std::size_t kNumThreads, std::size_t kUnroll>
+struct MultiIndexCopySlimKernel {
+    static void run(
+        tvm::ffi::TensorView dst_ptrs,
+        tvm::ffi::TensorView src_ptrs,
+        tvm::ffi::TensorView feat_bytes,
+        tvm::ffi::TensorView dst_indices,
+        tvm::ffi::TensorView src_indices,
+        tvm::ffi::Optional<tvm::ffi::TensorView> num_indices
+    ) {
+        using namespace host;
+        auto device = SymbolicDevice{};
+        auto B = SymbolicSize{"num_banks"};
+        auto L = SymbolicSize{"indices length"};
+        auto ptr_dtype = SymbolicDType{};
+        auto indices_dtype = SymbolicDType{};
+        auto num_indices_dtype = SymbolicDType{};
+
+        TensorMatcher({B}).with_dtype<int64_t>(ptr_dtype).with_device<kDLCUDA>(device)
+            .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes);
+        TensorMatcher({L}).with_dtype<int32_t, int64_t>(indices_dtype).with_device<kDLCUDA>(device)
+            .verify(dst_indices).verify(src_indices);
+
+        const int64_t* valid_length = nullptr;
+        if (num_indices.has_value()) {
+            TensorMatcher({1}).with_dtype<int64_t>(num_indices_dtype).with_device<kDLCUDA>(device)
+                .verify(num_indices.value());
+            valid_length = static_cast<const int64_t*>(num_indices.value().data_ptr());
+        }
+
+        const auto params = MultiIndexCopyParams{
+            static_cast<const int64_t*>(dst_ptrs.data_ptr()),
+            static_cast<const int64_t*>(src_ptrs.data_ptr()),
+            static_cast<const int64_t*>(feat_bytes.data_ptr()),
+            dst_indices.data_ptr(),
+            src_indices.data_ptr(),
+            valid_length,
+            static_cast<int64_t>(L.unwrap()),
+            static_cast<int>(B.unwrap()),
+        };
+        const auto use_int32 = indices_dtype.unwrap().bits == 32;
+        const auto kernel = use_int32
+            ? fast_index_copy_multi_slim<int32_t, kNumThreads, kUnroll>
+            : fast_index_copy_multi_slim<int64_t, kNumThreads, kUnroll>;
+        LaunchKernel(kNumBlocks, kNumThreads, device.unwrap())(kernel, params);
+    }
+};
diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py
index b2030620ee..3047947c78 100644
--- a/python/freetoken/kernel/fast_index_copy.py
+++ b/python/freetoken/kernel/fast_index_copy.py
@@ -14,6 +14,11 @@
 
 
 DEFAULT_NUM_BLOCKS = 4
+# fast_index_copy_multi_slim grid (8K loads in flight, the least that holds ~51 GB/s on an
+# RTX 5080 at 4 experts); the AOT spec (kernel/aot.py) builds exactly this config
+SLIM_COPY_BLOCKS = 8
+SLIM_COPY_THREADS = 256
+SLIM_COPY_UNROLL = 4
 SKIP_FAST_INDEX_COPY_ENV = "FREETOKEN_SKIP_FAST_INDEX_COPY"
 _TRUE_VALUES = {"1", "true", "yes", "on"}
 
@@ -202,6 +207,43 @@ def fast_index_copy_multi_jit(
     module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices)
 
 
+@lru_cache(maxsize=None)
+def _jit_fast_index_copy_multi_slim_module(*, num_blocks: int, num_threads: int, unroll: int) -> Module:
+    args = make_cpp_args(num_blocks, num_threads, unroll)
+    return load_jit(
+        "fast_index_copy_multi_slim",
+        *args,
+        cuda_files=["fast_index_copy.cuh"],
+        cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")],
+    )
+
+
+def fast_index_copy_multi_slim_jit(
+    dst_ptrs: torch.Tensor,
+    src_ptrs: torch.Tensor,
+    feat_bytes: torch.Tensor,
+    dst_indices: torch.Tensor,
+    src_indices: torch.Tensor,
+    num_indices: torch.Tensor | None = None,
+    *,
+    num_blocks: int = SLIM_COPY_BLOCKS,
+    num_threads: int = SLIM_COPY_THREADS,
+    unroll: int = SLIM_COPY_UNROLL,
+) -> None:
+    """:func:`fast_index_copy_multi_jit` from a small fixed grid, for copies that run beside compute.
+
+    ``num_blocks`` CTAs split the concatenated (row, bank) 16-byte chunk range evenly and each
+    thread keeps ``unroll`` loads in flight, so a few SMs still saturate the host link while
+    the rest stay free for concurrent kernels. Same arguments and alignment contract.
+    """
+    if _skip_fast_index_copy_enabled():
+        return
+    module = _jit_fast_index_copy_multi_slim_module(
+        num_blocks=num_blocks, num_threads=num_threads, unroll=unroll
+    )
+    module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices)
+
+
 def update_copy_flag_jit(sync_flag: torch.Tensor, delta: int) -> None:
     assert sync_flag.is_cuda
     assert sync_flag.numel() == 1
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 2f02a0d2fe..7200cb2f4e 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -240,7 +240,19 @@ def decode_forward(
             topk=self.top_k,
             renormalize=self.renormalize,
         )
-        return self._decode_routed(hidden_states, topk_weights, topk_ids)
+        return self._decode_routed(hidden_states, topk_weights, topk_ids, router_logits=router_logits)
+
+    def decode_side_applies(self) -> bool:
+        """Whether FREETOKEN_MOE_COPY_OVERLAP applies to this layer's next forward: a GPU decode on a
+        cache with a side stream (not prefill, CPU/hybrid decode or TP > 1); see decode_side."""
+        cache = self.offload_cache
+        return not (
+            cache is None
+            or cache.decode_copy_stream is None
+            or self.tp_size > 1
+            or get_global_ctx().batch.is_prefill
+            or cache.is_cpu_layer(self.layer_id)
+        )
 
     def prefill_forward(
         self,
@@ -268,6 +280,7 @@ def _decode_routed(
         hidden_states: torch.Tensor,
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
+        router_logits: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot
         ids in place (loading missing experts), then the GEMM reads the full slot
@@ -287,7 +300,7 @@ def _decode_routed(
             return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids)
         if cache.decode_target == "hybrid":
             return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids)
-        cache.ensure_experts(self.layer_id, topk_ids)
+        cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits)
         cache.copy_missing()
         return self._expert_gemm(
             cache,
@@ -418,9 +431,11 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc
         if n > cache.cache_size // 2:  # leave room so a chunk never evicts the one before it
             return None
         uniq = uniq.to(torch.int32)
+        # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly
+        pin_since = cache.step.clone() if cache.cache_policy_id else None
         for start in range(0, n, _ENSURE_CHUNK_IDS):
             part = uniq[start : start + _ENSURE_CHUNK_IDS]  # a contiguous view: ensure rewrites it in place
-            cache.ensure_experts(self.layer_id, part)
+            cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since)  # not a decode step
             cache.copy_missing()
         return uniq[inverse].view_as(topk_ids)
 
diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py
index aadcdc0481..65de7f85fe 100644
--- a/python/freetoken/models/deepseek_v4/moe.py
+++ b/python/freetoken/models/deepseek_v4/moe.py
@@ -126,7 +126,7 @@ def _prefill_routed(
             or cache.is_unpinned_layer(self.layer_id)
         ):
             return super()._prefill_routed(hidden_states, topk_weights, topk_ids)
-        cache.ensure_experts(self.layer_id, topk_ids)  # in-place expert-id -> slot
+        cache.ensure_experts(self.layer_id, topk_ids, update_state=False)  # in-place expert-id -> slot; not a decode step
         cache.copy_missing()
         if cache.collect_stats:
             cache.record_decode_stats(self.layer_id)
diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py
index c10407861f..6f729231f1 100644
--- a/python/freetoken/models/qwen4_exp/moe.py
+++ b/python/freetoken/models/qwen4_exp/moe.py
@@ -35,9 +35,19 @@ def __init__(
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
+        se, ex = self.shared_expert, self.experts
+        if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies():
+            # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router,
+            # ensure and miss copy stay back to back on this one; joined before the combine
+            cache = ex.offload_cache
+            with cache.decode_side(ex.layer_id):
+                gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1))
+                shared = se.forward(hidden_states)
+            routed = ex.forward(hidden_states=hidden_states, router_logits=self.gate.forward(hidden_states))
+            cache.join_decode_side(ex.layer_id)
+            return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim)
         router_logits = self.gate.forward(hidden_states)
         gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1))
-        se, ex = self.shared_expert, self.experts
         if (
             self._tp_size > 1
             and isinstance(se.down_proj, LinearRowParallel)
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 50f06e1073..956e34f2ea 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -2,6 +2,7 @@
 
 import math
 import os
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Iterator
 
@@ -13,6 +14,9 @@
 # (kept for A/B profiling). Falls back to per-bank automatically if a bank's row bytes or
 # base address are not 16-byte aligned.
 _FUSED_COPY = os.getenv("FREETOKEN_FUSED_COPY", "1").strip().lower() not in {"0", "false", "no", "off"}
+# Opt-in: run the fused copy_missing on the slim kernel (a few small CTAs, leaves most SMs to
+# concurrent compute) instead of the 8-CTAs-per-bank one. Off by default.
+_SLIM_COPY = os.getenv("FREETOKEN_MOE_SLIM_COPY", "0").strip().lower() in {"1", "true", "yes", "on"}
 
 # cudaMemcpyBatchAsync silently degrades to a SYNCHRONOUS copy when a batch mixes
 # large entries with sub-~256KB entries on registered host memory (H100 + CUDA 13.0,
@@ -26,6 +30,7 @@
 _SMALL_BANK_FEAT_BYTES = 256 * 1024
 
 from freetoken.kernel.fast_index_copy import default_blocks_per_bank
+from freetoken.env import ENV
 from freetoken.utils import init_logger
 
 if TYPE_CHECKING:
@@ -167,11 +172,20 @@ class OffloadMoeCache:
     # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers.
     pool_of_layer: list[int] | None = None
     pool_sizes: list[int] | None = None
+    # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the
+    # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP.
+    decode_copy_overlap: bool | None = None
 
     def __post_init__(self) -> None:
-        policy_ids = {"lru": 0}
-        assert self.cache_policy in policy_ids
+        from freetoken.moe.scored_ensure import POLICY_IDS
+
+        assert self.cache_policy in POLICY_IDS, self.cache_policy
         assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target
+        if self.cache_policy != "lru" and self.decode_target != "gpu":
+            # the hybrid and CPU decode paths run their own kernels and would silently stay LRU
+            raise ValueError(
+                f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}"
+            )
         if self.layer_layouts is not None:
             assert self.layout is None, "pass layout or layer_layouts, not both"
             assert len(self.layer_layouts) == self.num_layers, (len(self.layer_layouts), self.num_layers)
@@ -208,7 +222,11 @@ def __post_init__(self) -> None:
             "cache, so cache_size must be at least 2 * num_experts "
             "(raise moe_cache_size or disable moe_prefill_overlap)"
         )
-        self.cache_policy_id = policy_ids[self.cache_policy]
+        self.cache_policy_id = POLICY_IDS[self.cache_policy]
+        # rule: an unrouted expert within this many logits of the top-k-th is a near miss
+        self.near_miss_thr = (
+            float(os.getenv("FREETOKEN_MOE_NEAR_MISS_THR", "0.25")) if self.cache_policy == "rule" else None
+        )
         self.slot_for_id = torch.full(
             (self.num_layers, self.num_experts),
             -1,
@@ -242,6 +260,7 @@ def __post_init__(self) -> None:
         self.expert_recency = torch.full(
             (self.num_layers, self.num_experts), -1, dtype=torch.int64, device=self.device
         )
+        self._init_evict_state()
         # Host source banks (one [num_experts, ...] tensor per layer, so layers can
         # carry independent host attributes -- see layer_residency) and their GPU
         # slot caches, keyed by the format's bank schema (attached by
@@ -341,6 +360,14 @@ def __post_init__(self) -> None:
         self.prefill_hit_rows = 0
         self.prefill_total_rows = 0
         self._slice_pools()
+        # Decode copy overlap: one side stream plus a (forked, done) event pair per layer.
+        # Neither is cache_size-shaped, so rebuild() keeps them.
+        if self.decode_copy_overlap is None:
+            self.decode_copy_overlap = bool(ENV.MOE_COPY_OVERLAP)
+        self.decode_copy_stream: torch.cuda.Stream | None = None
+        self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = []
+        if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu":
+            self._init_decode_copy_overlap()
 
     def _pool_layout(self, pool: int) -> int:
         return self.layer_layout[self.pool_of_layer.index(pool)]
@@ -449,6 +476,78 @@ def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]:
         """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids."""
         assert self.banks, "set_bank_sources must register the banks first"
         return self._pool_bank_views[self.pool_of_layer[layer_id]]
+    def _init_decode_copy_overlap(self) -> None:
+        # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine
+        # or capture stream, and the side work would then run serially (correct, but no overlap)
+        from cuda.bindings import runtime as cudart
+
+        with torch.cuda.device(self.device):
+            err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking)
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise RuntimeError(f"cudaStreamCreateWithFlags failed: {err}")
+        self.decode_copy_stream = torch.cuda.ExternalStream(int(handle), device=self.device)
+        self._decode_copy_events = [
+            (torch.cuda.Event(), torch.cuda.Event()) for _ in range(self.num_layers)
+        ]
+        # torch creates the CUDA event on its first record: do that now, never inside a capture
+        for pair in self._decode_copy_events:
+            for event in pair:
+                event.record(self.decode_copy_stream)
+        logger.info("MoE decode copy overlap on: expert misses copy on a side stream")
+
+    def _init_evict_state(self) -> None:
+        """Scored policies' per-(layer, expert) state (moe/scored_ensure.py); None under LRU.
+
+        It is not slot-shaped: it survives eviction and prefill invalidation, and a rebuild
+        only resets it. ``evict_tok`` is the decode-step clock the first ensure of each step bumps.
+        """
+        from freetoken.moe import scored_ensure as se
+
+        self.evict_tok = self.evict_last_tok = self.evict_lc = self.evict_ct = None
+        self.evict_g_table = None
+        n, dev = self.num_layers * self.num_experts, self.device
+        if self.cache_policy_id >= 1:
+            self.evict_tok = torch.zeros((), dtype=torch.int64, device=dev)
+            self.evict_last_tok = torch.full((n,), se.LAST_TOK_NEVER, dtype=torch.int64, device=dev)
+        if self.cache_policy_id >= 2:
+            self.evict_lc = torch.full((n,), se.LC_NEVER, dtype=torch.int32, device=dev)
+            self.evict_ct = torch.zeros((n,), dtype=torch.int64, device=dev)
+            self.evict_g_table = se.softplus2_table(dev)
+        self._alloc_evict_mirrors()
+
+    def _alloc_evict_mirrors(self) -> None:
+        """Slot-shaped copies of each slot owner's state, read coalesced by the victim scan.
+
+        A mirror is valid while ``evict_slot_owner`` equals ``id_of_slot``; the kernel keeps it
+        so, and anything else that installs an id into a slot must set the owner to -1.
+        """
+        self.evict_slot_owner = self.evict_slot_last_tok = self.evict_slot_lc = self.evict_slot_ct = None
+        if self.evict_tok is None:
+            return
+        size, dev = self.cache_size, self.device
+        self.evict_slot_owner = torch.full((size,), -1, dtype=torch.int32, device=dev)
+        self.evict_slot_last_tok = torch.zeros((size,), dtype=torch.int64, device=dev)
+        if self.evict_lc is not None:
+            self.evict_slot_lc = torch.zeros((size,), dtype=torch.int32, device=dev)
+            self.evict_slot_ct = torch.zeros((size,), dtype=torch.int64, device=dev)
+
+    def reset_evict_state(self) -> None:
+        if self.evict_tok is None:
+            return
+        from freetoken.moe import scored_ensure as se
+
+        self.evict_tok.zero_()
+        self.evict_last_tok.fill_(se.LAST_TOK_NEVER)
+        self.evict_slot_owner.fill_(-1)
+        if self.evict_lc is not None:
+            self.evict_lc.fill_(se.LC_NEVER)
+            self.evict_ct.zero_()
+
+    def evict_clock_layer(self) -> int:
+        """The first layer that ensures on the GPU each decode step; its call bumps ``evict_tok``."""
+        if not self.cpu_layer_ids:
+            return 0
+        return min(set(range(self.num_layers)) - self.cpu_layer_ids, default=0)
 
     def set_bank_sources(
         self,
@@ -727,6 +826,8 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         self.num_indices.zero_()
         self.num_missing_full.zero_()
         self.expert_recency.fill_(-1)
+        self._alloc_evict_mirrors()
+        self.reset_evict_state()
         self.stat_missing.zero_()
         self.stat_active.zero_()
         self.stat_calls.zero_()
@@ -1076,8 +1177,21 @@ def release_prefill_layer(self, layer_id: int) -> None:
             self._prefill_buffer_has_release_event[buffer_id] = True
         self._prefill_buffer_released[buffer_id] = True
 
-    def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
-        from freetoken.moe.offload_kernels import ensure_experts
+    def ensure_experts(
+        self,
+        layer_id: int,
+        expert_ids: torch.Tensor,
+        *,
+        update_state: bool = True,
+        router_logits: torch.Tensor | None = None,
+        pin_since: torch.Tensor | None = None,
+    ) -> None:
+        """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and
+        decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule``
+        policy's near misses; ``pin_since`` (a ``step`` value) keeps every slot touched after it
+        (a small prefill's earlier chunks). All are ignored under LRU, whose recency order
+        already keeps those slots."""
+        from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored
 
         if self.collect_decode_freq:
             # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to
@@ -1086,7 +1200,18 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        ensure_experts(self, layer_id, expert_ids)
+        if self.cache_policy_id == 0:
+            ensure_experts(self, layer_id, expert_ids)
+            return
+        ensure_experts_scored(
+            self,
+            layer_id,
+            expert_ids,
+            bump_tok=update_state and layer_id == self.evict_clock_layer(),
+            update_state=update_state,
+            router_logits=router_logits,
+            pin_since=pin_since,
+        )
 
     def ensure_experts_hybrid(
         self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False
@@ -1117,6 +1242,8 @@ def materialize_layer(self, layer_id: int) -> None:
         self._pending_src_layer = layer_id
         self._pending_whole_layer = True
         materialize_layer(self, layer_id)
+        if self.evict_slot_owner is not None:
+            self.evict_slot_owner[: self.num_experts].fill_(-1)  # it installs the layer into these slots
 
     def reset(self) -> None:
         from freetoken.moe.offload_kernels import reset_cache
@@ -1125,6 +1252,7 @@ def reset(self) -> None:
         # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe
         # it here so a new sequence starts with cold hybrid fetch priorities.
         self.expert_recency.fill_(-1)
+        self.reset_evict_state()
 
     def reset_stats(self) -> None:
         self.prefill_hit_rows = 0
@@ -1231,9 +1359,12 @@ def decode_routing_stats(self) -> dict:
             "norm_entropy": norm_ent,
         }
 
-    def copy_missing(self) -> None:
+    def copy_missing(self, layer_id: int | None = None) -> None:
+        """Copy the misses the last ensure_experts/materialize_layer staged, from ``layer_id``'s
+        host banks (default: the layer that staged them)."""
         assert self.banks, "set_bank_sources must register the banks first"
-        layer_id = self._pending_src_layer
+        if layer_id is None:
+            layer_id = self._pending_src_layer
         assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)"
         if layer_id in self._unpinned_layers:
             if not self._pending_whole_layer:
@@ -1250,13 +1381,14 @@ def copy_missing(self) -> None:
         # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool
         whole = self._pending_whole_layer
         if self._copy_fused_ok:
-            from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
+            from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit
 
             # One launch copies the missing rows for every bank (instead of one launch per
             # bank). evict_slots/src_indices/num_indices are shared across banks;
             # src_indices holds layer-local expert rows, resolved against this layer's
             # source pointers (layer_id is a static int per captured graph node).
-            fast_index_copy_multi_jit(
+            copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit
+            copy(
                 self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]],
                 self._copy_src_ptrs[layer_id],
                 self._copy_feat_bytes[self.layer_layout[layer_id]],
@@ -1281,6 +1413,24 @@ def copy_missing(self) -> None:
                 self.num_indices,
             )
 
+    @contextmanager
+    def decode_side(self, layer_id: int) -> Iterator[None]:
+        """Run the body on ``decode_copy_stream``, ordered after the work already on the current
+        stream; ``join_decode_side`` must follow before its results are read.
+
+        The copy stays on the current stream right after its ensure, so it never waits on a
+        cross-stream edge; the side work (slack-rich) takes that latency and any SM contention.
+        Under capture the event pair becomes a fork and a join edge of the graph."""
+        forked, done = self._decode_copy_events[layer_id]
+        forked.record(torch.cuda.current_stream(self.device))
+        self.decode_copy_stream.wait_event(forked)
+        with torch.cuda.stream(self.decode_copy_stream):
+            yield
+            done.record(self.decode_copy_stream)
+
+    def join_decode_side(self, layer_id: int) -> None:
+        torch.cuda.current_stream(self.device).wait_event(self._decode_copy_events[layer_id][1])
+
 
 def iter_offload_moe_layers(model) -> Iterator:
     from freetoken.layers import OffloadMoELayer, iter_moe_layers
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index 8c4ad8f393..3523a9850b 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -50,6 +50,51 @@ class HybridRoutes(NamedTuple):
     cpu_ids: torch.Tensor  # expert ids, -1 for a GPU route
 
 
+def ensure_experts_scored(
+    cache,
+    layer_id: int,
+    expert_ids: torch.Tensor,
+    *,
+    bump_tok: bool,
+    update_state: bool,
+    router_logits: torch.Tensor | None = None,
+    pin_since: torch.Tensor | None = None,
+) -> None:
+    """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule)."""
+    from freetoken.moe.scored_ensure import scored_ensure
+
+    scored_ensure(
+        expert_ids,
+        cache.slot_for_id.view(-1),
+        cache.id_of_slot,
+        cache.usage,
+        cache.step,
+        expert_ids,
+        cache.src_indices,
+        cache.evict_slots,
+        cache.num_indices,
+        policy=cache.cache_policy_id,
+        num_layers=cache.num_layers,
+        num_experts=cache.num_experts,
+        tok=cache.evict_tok,
+        last_tok=cache.evict_last_tok,
+        lc=cache.evict_lc,
+        ct=cache.evict_ct,
+        g_table=cache.evict_g_table,
+        slot_owner=cache.evict_slot_owner,
+        slot_last_tok=cache.evict_slot_last_tok,
+        slot_lc=cache.evict_slot_lc,
+        slot_ct=cache.evict_slot_ct,
+        router_logits=router_logits if cache.near_miss_thr is not None else None,
+        near_miss_thr=cache.near_miss_thr or 0.0,
+        stats=cache.lru_stats[layer_id] if cache.collect_stats else None,
+        id_base=layer_id * cache.num_experts,
+        bump_tok=bump_tok,
+        update_state=update_state,
+        pin_since=pin_since,
+    )
+
+
 def ensure_experts_hybrid(
     cache,
     layer_id: int,
diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py
new file mode 100644
index 0000000000..123086ed51
--- /dev/null
+++ b/python/freetoken/moe/scored_ensure.py
@@ -0,0 +1,395 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright 2026 The FlashLib Authors
+# Adapted from flashlib 0.3.0 (kernels/slot_cache/triton/lru_ensure.py: _phase1, _stats,
+# _lru_ensure_kernel, _num_warps_for), https://github.com/FlashML-org/flashlib
+"""Scored-victim slot-cache admission for the MoE expert cache (``--moe-cache-policy``).
+
+flashlib's sequential-argmin ``lru_ensure`` with the victim key made pluggable. POLICY 0 is
+LRU and matches flashlib bit for bit; the others evict the lowest packed score key:
+
+- 1 ``kd``: ``-(k + beta * d / L)``. ``k`` is the tokens the expert has been passed over
+  since its last use, ``d`` the layer-steps until its layer comes up again.
+- 2 ``kdfb``: kd plus ``w * log2(decayed use count)`` with a ``halflife``-token decay.
+- 3 ``rule``: kdfb plus the near-miss refresh: this layer's experts whose router logit is
+  within a margin of the top-k-th count as used one token ago (``last_tok``), unpinned.
+
+The per-(layer, expert) state survives eviction; slot-indexed mirrors of it (tagged with the
+owning id) let the victim scan read it coalesced instead of gathering by id. Scores and
+counts are int32 Q16 fixed point and the count update ``log2(2^x + 1)`` reads a host-built
+table, so a CPU reference can reproduce every victim exactly (``tl.exp2``/``tl.log2`` are
+approximate).
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+POLICY_IDS = {"lru": 0, "kd": 1, "kdfb": 2, "rule": 3}
+
+# Replay-tuned on sep5_s01 at 1650 slots; the plateau is flat for beta 1-8, w 2-4, halflife 32-128.
+BETA = 4.0
+W = 3.0
+HALFLIFE = 64
+
+Q = 16  # fraction bits of scores and counts
+LAST_TOK_NEVER = -1_000_000
+LC_NEVER = -(1 << 30)
+LC_FLOOR = -(30 << Q)
+K_MAX = 1 << 14  # tokens; with the caps on beta and w below, a score fits int32
+MAX_BETA, MAX_W = 64.0, 16.0
+DECAY_SHIFT = 8
+SCORE_BIAS = 1 << 31
+MAX_SLOT_BITS = 30
+# g(u) = log2(1 + 2^-u) sampled every 2^-G_STEP_BITS for u in [0, G_RANGE]; zero beyond
+G_STEP_BITS = 6
+G_RANGE = 24
+G_LAST = G_RANGE << G_STEP_BITS
+# jit code may only read globals that are constexpr
+_Q = tl.constexpr(Q)
+_LC_NEVER = tl.constexpr(LC_NEVER)
+_LC_FLOOR = tl.constexpr(LC_FLOOR)
+_K_MAX = tl.constexpr(K_MAX)
+_DECAY_SHIFT = tl.constexpr(DECAY_SHIFT)
+_SCORE_BIAS = tl.constexpr(SCORE_BIAS)
+_G_SHIFT = tl.constexpr(Q - G_STEP_BITS)
+_G_LAST = tl.constexpr(G_LAST)
+
+
+def softplus2_table(device=None) -> torch.Tensor:
+    """Q16 ``log2(1 + 2^-u)`` at ``u = i / 64``, two entries past ``G_LAST`` for the interpolation."""
+    step = 1 << G_STEP_BITS
+    vals = [round(math.log2(1.0 + 2.0 ** (-i / step)) * (1 << Q)) for i in range(G_LAST + 2)]
+    return torch.tensor(vals, dtype=torch.int32, device=device)
+
+
+def score_params(beta: float, w: float, halflife: int, num_layers: int) -> tuple[int, int, int, int]:
+    """``(BETA_STEP, W_Q4, DECAY_MUL, DT_MAX)`` shared by the kernel and the CPU reference.
+
+    ``beta * d / L`` is ``d * BETA_STEP`` in Q16; the decay over ``dt`` layer-steps is
+    ``(dt * DECAY_MUL) >> DECAY_SHIFT``, with ``dt`` capped at 64 half-lives (past the -30 floor).
+    """
+    assert 0 <= beta <= MAX_BETA and 0 <= w <= MAX_W and halflife >= 1
+    span = num_layers * int(halflife)
+    decay_mul = round((1 << (Q + DECAY_SHIFT)) / span)
+    dt_max = 64 * span
+    assert dt_max * decay_mul < (1 << 31)
+    return round(beta * (1 << Q) / num_layers), round(w * 16), decay_mul, dt_max
+
+
+@triton.jit
+def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K,
+            BLOCK_K: tl.constexpr, id_base):
+    """Vendored unchanged: dedup the query, split hit/miss, rank the misses, bump the hits."""
+    k = tl.arange(0, BLOCK_K)
+    kmask = k < K
+    q = tl.load(query_ptr + k, mask=kmask, other=-1) + id_base
+    s = tl.load(slot_of_id_ptr + q, mask=kmask, other=-1)
+    hit = kmask & (s >= 0)
+    miss = kmask & (s == -1)
+    same = (
+        (q[:, None] == q[None, :])
+        & (k[:, None] > k[None, :])
+        & kmask[:, None]
+        & kmask[None, :]
+    )
+    first = kmask & (tl.sum(same.to(tl.int32), axis=1) == 0)
+    first_miss = miss & first
+    smaller = (q[None, :] < q[:, None]) & first_miss[None, :]
+    rank = tl.sum(smaller.to(tl.int32), axis=1)
+    num_missing = tl.sum(first_miss.to(tl.int32))
+    tl.store(num_copy_ptr, num_missing.to(tl.int64))
+    # Duplicated hits write the same value to the same slot -- idempotent.
+    tl.store(lru_usage_ptr + s, step, mask=hit)
+    return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1)
+
+
+@triton.jit
+def _stats(stats_ptr, first, num_missing):
+    # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA.
+    si = tl.arange(0, 4)
+    v = tl.where(si == 0, tl.sum(first.to(tl.int32)), tl.where(si == 1, num_missing, 1))
+    # relaxed: only later kernels read the counters; flashlib's acq_rel waits on this call's stores (~0.3 us)
+    tl.atomic_add(stats_ptr + si, v.to(tl.int64), mask=si < 3, sem="relaxed")
+
+
+@triton.jit
+def _decayed(lc, ct, now, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr):
+    """int32 Q16 log2 count decayed from ``ct`` to ``now`` (int64 layer-steps)."""
+    dt = tl.minimum(tl.maximum(now - ct, 0), DT_MAX).to(tl.int32)
+    return lc - ((dt * DECAY_MUL) >> _DECAY_SHIFT)
+
+
+@triton.jit
+def _softplus2(x, g_ptr, mask):
+    """Q16 ``log2(2^x + 1)`` = ``max(x, 0) + g(|x|)``, g linearly interpolated from the table."""
+    u = tl.abs(x)
+    i = tl.minimum(u >> _G_SHIFT, _G_LAST)
+    frac = u & ((1 << _G_SHIFT) - 1)
+    g0 = tl.load(g_ptr + i, mask=mask, other=0)
+    g1 = tl.load(g_ptr + i + 1, mask=mask, other=0)
+    g = tl.where(i < _G_LAST, g0 + (((g1 - g0) * frac) >> _G_SHIFT), 0)
+    return tl.maximum(x, 0) + g
+
+
+@triton.jit
+def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other):
+    """The slot owner's state: its coalesced mirror, or a gather by id where the mirror is stale."""
+    return tl.where(
+        mirrored,
+        tl.load(mirror_ptr + c, mask=mirrored, other=other),
+        tl.load(state_ptr + oid, mask=stale, other=other),
+    )
+
+
+@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"])
+def _scored_ensure_kernel(
+    query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr,
+    out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr,
+    tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr,
+    owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr,
+    K, num_cached, id_base, nm_topk, nm_stride, nm_thr,
+    BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr,
+    NM_ROWS: tl.constexpr,
+    USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr,
+    POLICY: tl.constexpr, BUMP_TOK: tl.constexpr, UPDATE_STATE: tl.constexpr, PIN_SINCE: tl.constexpr,
+    NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr,
+    BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr,
+    SLOT_BITS: tl.constexpr,
+):
+    step = tl.load(lru_step_ptr) + 1
+    tl.store(lru_step_ptr, step)
+    if POLICY != 0:
+        tok = tl.load(tok_ptr)
+        if BUMP_TOK:
+            tok = tok + 1
+            tl.store(tok_ptr, tok)
+        layer = id_base // NUM_EXPERTS
+        now = tok * NUM_LAYERS + layer
+        if POLICY == 3:
+            if UPDATE_STATE:
+                # Routed ids qualify too; they are pinned now and set to tok below, so that is harmless.
+                ex = tl.arange(0, BLOCK_E)
+                j = tl.arange(0, BLOCK_TOPK)
+                prev = tl.load(last_tok_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=0)
+                held_at = tl.load(slot_of_id_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=-1)
+                near = ex < 0
+                for r in tl.static_range(NM_ROWS):
+                    routed = tl.load(query_ptr + r * nm_topk + j, mask=j < nm_topk, other=-1)
+                    row_ptr = logits_ptr + r * nm_stride
+                    kth = tl.min(tl.load(row_ptr + routed, mask=routed >= 0, other=float("inf")).to(tl.float32), axis=0)
+                    row = tl.load(row_ptr + ex, mask=ex < NUM_EXPERTS, other=float("-inf")).to(tl.float32)
+                    near = near | (row >= kth - nm_thr)
+                refreshed = tl.maximum(prev, tok - 1)
+                tl.store(last_tok_ptr + id_base + ex, refreshed, mask=near)
+                tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0))
+                # the victim scan and the routed update below read and rewrite what this wrote
+                tl.debug_barrier()
+    q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1(
+        query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K,
+        id_base)
+
+    if num_missing > 0:
+        # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted.
+        tl.debug_barrier()
+        c = tl.arange(0, BLOCK_C)
+        cmask = c < num_cached
+        u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX)
+        if POLICY == 0:
+            umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype)
+            u = tl.where((u == step) | (~cmask), umax, u)
+        else:
+            oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1)
+            held = cmask & (oid >= 0)
+            owner = tl.load(owner_ptr + c, mask=cmask, other=-1)
+            mirrored = held & (owner == oid)
+            stale = held & (owner != oid)
+            lk = oid // NUM_EXPERTS
+            ahead = lk > layer
+            d = tl.where(ahead, lk - layer, NUM_LAYERS - layer + lk)
+            last = _owner_state(slot_last_ptr, last_tok_ptr, c, oid, mirrored, stale, 0)
+            k = tl.minimum(tl.maximum(tl.where(ahead, tok - 1, tok) - last, 0), _K_MAX).to(tl.int32)
+            # int32 throughout: |score| < 2^31 by the caps on k, beta and w
+            score = -((k << _Q) + d * BETA_STEP)
+            if POLICY >= 2:
+                lc = _owner_state(slot_lc_ptr, lc_ptr, c, oid, mirrored, stale, _LC_NEVER)
+                ct = _owner_state(slot_ct_ptr, ct_ptr, c, oid, mirrored, stale, 0)
+                lcv = _decayed(lc, ct, now, DECAY_MUL, DT_MAX)
+                lcv = tl.where(lc == _LC_NEVER, _LC_FLOOR, tl.maximum(lcv, _LC_FLOOR))
+                score += (W_Q4 * lcv) >> 4
+            # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one.
+            key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64))
+            evictable = cmask & (u != step)
+            if PIN_SINCE:
+                # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too
+                evictable = evictable & (u <= tl.load(pin_ptr))
+            # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin
+            pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1)
+            key = tl.where(evictable, key, pinned_key)
+        for i in tl.range(num_missing):
+            if POLICY == 0:
+                victim = tl.argmin(u, axis=0).to(tl.int32)
+            else:
+                victim = (tl.min(key, axis=0) & ((1 << SLOT_BITS) - 1)).to(tl.int32)
+            # Scalar load: victims are distinct, so no earlier iteration wrote this slot.
+            old = tl.load(id_of_slot_ptr + victim)
+            if old >= 0:
+                tl.store(slot_of_id_ptr + old, -1)
+            e = tl.sum(tl.where((rank == i) & first_miss, q, 0))
+            tl.store(id_of_slot_ptr + victim, e)
+            tl.store(slot_of_id_ptr + e, victim)
+            tl.store(lru_usage_ptr + victim, step)
+            tl.store(dst_ptr + i, victim)
+            tl.store(src_ptr + i, e - id_base)  # back to the caller's id space
+            out = tl.where((rank == i) & miss, victim, out)
+            if POLICY == 0:
+                u = tl.where(c == victim, umax, u)  # claim in-register
+            else:
+                key = tl.where(c == victim, pinned_key, key)
+
+    # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr.
+    tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask)
+    if POLICY != 0:
+        if UPDATE_STATE:
+            # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot
+            tl.store(last_tok_ptr + q, tok, mask=first)
+            tl.store(owner_ptr + out, q, mask=first)
+            tl.store(slot_last_ptr + out, tok, mask=first)
+            if POLICY >= 2:
+                lc = tl.load(lc_ptr + q, mask=first, other=_LC_NEVER)
+                x = _decayed(lc, tl.load(ct_ptr + q, mask=first, other=0), now, DECAY_MUL, DT_MAX)
+                lc_new = tl.where(lc == _LC_NEVER, 0, _softplus2(x, g_ptr, first))
+                tl.store(lc_ptr + q, lc_new, mask=first)
+                tl.store(ct_ptr + q, now, mask=first)
+                tl.store(slot_lc_ptr + out, lc_new, mask=first)
+                tl.store(slot_ct_ptr + out, now, mask=first)
+        else:
+            # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it
+            tl.store(owner_ptr + out, -1, mask=first_miss)
+    if COLLECT_STATS:
+        _stats(stats_ptr, first, num_missing)
+
+
+def _num_warps_for(block_c: int) -> int:
+    # flashlib's rule: 8 warps from 2048 slots, more only past ~32 elements per thread.
+    warps = 8 if block_c >= 2048 else 4
+    return max(warps, min(triton.next_power_of_2(max(block_c // 1024, 1)), 32))
+
+
+def scored_ensure(
+    query: torch.Tensor,
+    slot_of_id: torch.Tensor,
+    id_of_slot: torch.Tensor,
+    lru_usage: torch.Tensor,
+    lru_step: torch.Tensor,
+    out_indices: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    num_copy: torch.Tensor,
+    *,
+    policy: int,
+    num_layers: int,
+    num_experts: int,
+    tok: torch.Tensor | None = None,
+    last_tok: torch.Tensor | None = None,
+    lc: torch.Tensor | None = None,
+    ct: torch.Tensor | None = None,
+    g_table: torch.Tensor | None = None,
+    slot_owner: torch.Tensor | None = None,
+    slot_last_tok: torch.Tensor | None = None,
+    slot_lc: torch.Tensor | None = None,
+    slot_ct: torch.Tensor | None = None,
+    router_logits: torch.Tensor | None = None,
+    near_miss_thr: float = 0.0,
+    stats: torch.Tensor | None = None,
+    id_base: int = 0,
+    bump_tok: bool = False,
+    update_state: bool = True,
+    pin_since: torch.Tensor | None = None,
+    beta: float = BETA,
+    w: float = W,
+    halflife: int = HALFLIFE,
+) -> None:
+    """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``.
+
+    Same contract as ``lru_ensure``; ``policy`` > 0 also reads and updates the per-(layer,
+    expert) state ``last_tok``/``lc``/``ct`` (flat ``layer * num_experts + expert``) against
+    the device token counter ``tok``, and keeps its ``[num_cached]`` mirrors ``slot_last_tok``/
+    ``slot_lc``/``slot_ct`` valid where ``slot_owner`` equals ``id_of_slot`` (-1 marks a stale
+    mirror). ``bump_tok`` advances the counter first (once per decode step);
+    ``update_state=False`` leaves the state untouched (small-prefill ensures).
+    ``pin_since`` (a ``lru_step`` value) also pins every slot touched after it, so a small
+    prefill's chunked calls cannot evict each other's experts; policy > 0 only.
+    Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of
+    ``query``) are within ``near_miss_thr`` of their row's lowest routed logit.
+    """
+    k = query.numel()
+    num_cached = id_of_slot.numel()
+    assert query.dtype == torch.int32 and query.is_contiguous()
+    assert slot_of_id.dtype == torch.int32 and id_of_slot.dtype == torch.int32
+    assert out_indices.dtype == torch.int32 and out_indices.numel() == k
+    assert lru_usage.dtype == lru_step.dtype and lru_usage.numel() == num_cached
+    plan = min(k, num_cached)
+    assert src_indices.numel() >= plan and dst_indices.numel() >= plan
+    block_c = triton.next_power_of_2(num_cached)
+    slot_bits = max(block_c.bit_length() - 1, 1)
+    beta_step, w_q4, decay_mul, dt_max = score_params(beta, w, halflife, num_layers)
+    if policy != 0:
+        assert slot_bits <= MAX_SLOT_BITS, f"{num_cached} slots overflow the packed score key"
+        assert tok is not None and last_tok is not None
+        assert last_tok.dtype == torch.int64 and last_tok.numel() == num_layers * num_experts
+        assert slot_owner is not None and slot_last_tok is not None
+        assert slot_owner.dtype == torch.int32 and slot_owner.numel() == num_cached
+        assert slot_last_tok.dtype == torch.int64 and slot_last_tok.numel() == num_cached
+    if pin_since is not None:
+        assert policy != 0 and pin_since.dtype == lru_step.dtype and pin_since.numel() == 1
+    if policy >= 2:
+        assert lc is not None and ct is not None and g_table is not None
+        assert lc.dtype == torch.int32 and ct.dtype == torch.int64
+        assert slot_lc is not None and slot_lc.dtype == torch.int32 and slot_lc.numel() == num_cached
+        assert slot_ct is not None and slot_ct.dtype == torch.int64 and slot_ct.numel() == num_cached
+    nm_rows = nm_topk = nm_stride = 0
+    if policy == 3 and router_logits is not None:
+        assert router_logits.ndim == 2 and router_logits.shape[1] == num_experts
+        assert router_logits.stride(1) == 1 and k % router_logits.shape[0] == 0
+        nm_rows, nm_stride = router_logits.shape[0], router_logits.stride(0)
+        nm_topk = k // nm_rows
+    # unused pointers still need a tensor argument
+    dummy = lru_step
+    _scored_ensure_kernel[(1,)](
+        query, slot_of_id, id_of_slot, lru_usage, lru_step,
+        out_indices, src_indices, dst_indices, num_copy, stats,
+        dummy if tok is None else tok,
+        dummy if last_tok is None else last_tok,
+        dummy if lc is None else lc,
+        dummy if ct is None else ct,
+        dummy if g_table is None else g_table,
+        dummy if nm_rows == 0 else router_logits,
+        dummy if slot_owner is None else slot_owner,
+        dummy if slot_last_tok is None else slot_last_tok,
+        dummy if slot_lc is None else slot_lc,
+        dummy if slot_ct is None else slot_ct,
+        dummy if pin_since is None else pin_since,
+        k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr),
+        BLOCK_K=triton.next_power_of_2(k),
+        BLOCK_C=block_c,
+        BLOCK_E=triton.next_power_of_2(num_experts),
+        BLOCK_TOPK=triton.next_power_of_2(max(nm_topk, 1)),
+        NM_ROWS=nm_rows,
+        USAGE_MAX=torch.iinfo(lru_usage.dtype).max,
+        COLLECT_STATS=stats is not None,
+        POLICY=policy,
+        BUMP_TOK=bool(bump_tok) and policy != 0,
+        UPDATE_STATE=bool(update_state) and policy != 0,
+        PIN_SINCE=pin_since is not None,
+        NUM_LAYERS=num_layers,
+        NUM_EXPERTS=num_experts,
+        BETA_STEP=beta_step,
+        W_Q4=w_q4,
+        DECAY_MUL=decay_mul,
+        DT_MAX=dt_max,
+        SLOT_BITS=slot_bits,
+        num_warps=_num_warps_for(block_c),
+    )
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index a5940f8201..46db8d8aa4 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -844,8 +844,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
     parser.add_argument(
         "--moe-cache-policy",
         default=ServerArgs.moe_cache_policy,
-        choices=["lru"],
-        help="The unified MoE cache eviction policy.",
+        choices=["lru", "kd", "kdfb", "rule"],
+        help=(
+            "The unified MoE cache eviction policy. lru (default); kd: evict the expert passed "
+            "over for the most tokens, then the one whose layer comes up last; kdfb: kd plus a "
+            "decayed per-expert use count; rule: kdfb plus router near misses counted as recent "
+            "(margin FREETOKEN_MOE_NEAR_MISS_THR logits). GPU decode only."
+        ),
     )
 
     parser.add_argument(
diff --git a/tests/README.md b/tests/README.md
index 85865673e2..1d82e9fa73 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -65,6 +65,7 @@ checkpoint is set:
 | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) |
 | `FREETOKEN_DSV4_VL_MODEL` | `tokenizer/test_mm_tokenize.py` - local DeepSeek-V4 vision checkpoint |
 | `FREETOKEN_DSV4_VISION_CKPT` | `models/test_dsv4_vision.py` - the same checkpoint, for the parity tests against its `inference/` reference code |
+| `FREETOKEN_MOE_ROUTING_TRACE` | `moe/test_scored_ensure.py` — a compact decode routing trace directory (`experts.npy`, `logit_idx.npy`, `logit_val.npy`); its rows are replayed through the scored eviction kernel and the CPU reference |
 
 `test_aime.py` takes its sampling protocol from the checkpoint's own
 `generation_config.json` (pass@N at the recommended temperature, or a single greedy
diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py
index 8fb9112df2..a2412f8346 100644
--- a/tests/models/qwen4_exp/test_skeleton.py
+++ b/tests/models/qwen4_exp/test_skeleton.py
@@ -432,6 +432,130 @@ def test_shared_gate_kernels_match_torch(num_tokens, hidden, dtype):
     assert (fused.float() - ref).abs().max() <= (eager.float() - ref).abs().max() + 1e-6
 
 
+def _offload_experts(kind: str, layer_id: int, num_experts: int, top_k: int, hidden: int, inter: int):
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig, QuantBackend, QuantConfig, set_quant_backend
+
+    if kind == "bf16":
+        quant = NoQuantConfig()
+    else:
+        set_quant_backend(QuantBackend.parse("moe.nvfp4=triton"))  # the sm_120 decode kernel
+        quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}})
+    return OffloadMoELayer(
+        layer_id, num_experts, top_k, hidden, inter, quant_config=quant,
+        prefix=f"model.layers.{layer_id}.mlp.experts",
+    )
+
+
+def _copy_overlap_run(kind, banks, overlap, bs, graph, inputs, monkeypatch):
+    """Two Qwen4ExpMoE layers sharing one offload cache as small as a layer, so every step evicts
+    and copies: warm up eagerly (and capture), reset the cache, then decode ``inputs``."""
+    import freetoken.kernel.fast_index_copy as fic
+    from flashlib.kernels.slot_cache import Stat
+    from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE
+    from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.utils.torch_utils import torch_dtype
+
+    num_experts, top_k, hidden, inter = 8, 2, 256, 128
+    config = parse_config(toy_hf_config(
+        hidden_size=hidden, num_experts=num_experts, num_experts_per_tok=top_k,
+        moe_intermediate_size=inter, shared_expert_intermediate_size=inter,
+    ))
+    device = torch.device("cuda")
+    with torch.device(device), torch_dtype(torch.bfloat16):
+        moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(2)]
+    gen = torch.Generator(device=device).manual_seed(41)
+    for moe in moes:
+        _fill(moe, gen, scale=0.2)
+    experts = [_offload_experts(kind, i, num_experts, top_k, hidden, inter) for i in range(2)]
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=num_experts, cache_size=num_experts, device=device,
+        quant_format=banks.quant_format, layout=banks.layout,
+        max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap,
+    )
+    cache.set_bank_sources(banks.sources)
+    cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
+    cache.collect_stats = True
+    for moe, ex in zip(moes, experts):
+        ex.offload_cache = cache
+        moe.experts = ex
+    cache.reset()
+
+    import freetoken.models.qwen4_exp.moe as qmoe
+
+    copy_streams, gate_streams = [], []
+    real_gate = qmoe.shared_gate_sigmoid
+    for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"):  # FREETOKEN_MOE_SLIM_COPY
+        real_copy = getattr(fic, name)
+        monkeypatch.setattr(
+            fic, name,
+            lambda *a, _real=real_copy, **k: (copy_streams.append(torch.cuda.current_stream().cuda_stream), _real(*a, **k))[1],
+        )
+    monkeypatch.setattr(
+        qmoe, "shared_gate_sigmoid",
+        lambda *a, **k: (gate_streams.append(torch.cuda.current_stream().cuda_stream), real_gate(*a, **k))[1],
+    )
+
+    def step(x0, x1):
+        return moes[0].forward(x0.clone()), moes[1].forward(x1.clone())
+
+    _fresh_ctx(_batch=SimpleNamespace(is_prefill=False))
+    static = [torch.randn(bs, hidden, device=device, dtype=torch.bfloat16) for _ in range(2)]
+    step(*static)  # the eager warm-up the graph runner also does
+    if graph:
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            outs = step(*static)
+    cache.reset()
+    cache.reset_stats()
+    results = []
+    for x0, x1 in inputs:
+        if graph:
+            static[0].copy_(x0)
+            static[1].copy_(x1)
+            g.replay()
+        else:
+            outs = step(x0, x1)
+        results.append(tuple(o.clone() for o in outs))
+    torch.cuda.synchronize()
+    state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())]
+    misses = int(cache.lru_stats[:, Stat.MISS].sum())
+    side = getattr(cache.decode_copy_stream, "cuda_stream", None)
+    return results, state, misses, [s == side for s in copy_streams], [s == side for s in gate_streams]
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch):
+    """FREETOKEN_MOE_COPY_OVERLAP moves the shared expert onto a side stream beside the decode miss
+    copy; the MoE outputs, the slot map and the slot contents must not change by a single bit."""
+    from freetoken.moe.expert_banks import build_expert_banks
+
+    torch.manual_seed(7)
+    method = _offload_experts(kind, 0, 8, 2, 256, 128).quant_method
+    banks = build_expert_banks(method, 2, None, device=torch.device("cuda"), dummy=True)
+    gen = torch.Generator(device="cuda").manual_seed(bs)
+    inputs = [
+        tuple(torch.randn(bs, 256, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(2))
+        for _ in range(6)
+    ]
+    off = _copy_overlap_run(kind, banks, False, bs, graph, inputs, monkeypatch)
+    on = _copy_overlap_run(kind, banks, True, bs, graph, inputs, monkeypatch)
+
+    assert off[2] > 0 and on[2] == off[2], "the sequence must miss, and copy, in both runs"
+    assert off[3] and not any(off[3]) and not any(off[4]), "flag off: everything stays on the compute stream"
+    assert on[3] and not any(on[3]), "flag on: the miss copy stays right behind its ensure"
+    assert on[4] and all(on[4]), "flag on: the shared expert runs on the cache's side stream"
+    for got, want in zip(on[0], off[0]):
+        for a, b in zip(got, want):
+            assert torch.isfinite(b.float()).all()
+            assert torch.equal(a, b)
+    for a, b in zip(on[1], off[1]):
+        assert torch.equal(a, b)
+
+
 @requires_cuda
 def test_decoder_stack_prefill_and_decode(monkeypatch):
     """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights."""
diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py
new file mode 100644
index 0000000000..d7d8b0f2fb
--- /dev/null
+++ b/tests/moe/ref_scored_cache.py
@@ -0,0 +1,159 @@
+"""CPU reference for the scored slot-cache ensure (freetoken.moe.scored_ensure).
+
+A port of the adversarial-synthesis replay (global slot cache with lru_ensure semantics: the
+call's ids pinned, empty slots first, ties to the lowest slot, missing ids ascending paired with
+ascending victims) onto the kernel's Q16 integer arithmetic, so victims match the GPU exactly.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from freetoken.moe import scored_ensure as se
+
+_KEY_MAX = np.iinfo(np.int64).max
+
+
+class RefScoredCache:
+    def __init__(self, num_layers, num_experts, cache_size, policy,
+                 beta=se.BETA, w=se.W, halflife=se.HALFLIFE):
+        self.L, self.E, self.S, self.policy = num_layers, num_experts, cache_size, policy
+        self.beta_step, self.w_q4, self.decay_mul, self.dt_max = se.score_params(beta, w, halflife, num_layers)
+        block_c = 1 << max(cache_size - 1, 1).bit_length()
+        self.slot_bits = max(block_c.bit_length() - 1, 1)
+        self.g = se.softplus2_table().numpy().astype(np.int64)
+        self.reset()
+
+    def reset(self):
+        n = self.L * self.E
+        self.slot_of_id = np.full(n, -1, np.int64)
+        self.id_of_slot = np.full(self.S, -1, np.int64)
+        self.usage = np.zeros(self.S, np.int64)
+        self.step = 0
+        self.tok = 0
+        self.last_tok = np.full(n, se.LAST_TOK_NEVER, np.int64)
+        self.lc = np.full(n, se.LC_NEVER, np.int64)
+        self.ct = np.zeros(n, np.int64)
+
+    def _decayed(self, ids, now):
+        dt = np.minimum(np.maximum(now - self.ct[ids], 0), self.dt_max)
+        return self.lc[ids] - ((dt * self.decay_mul) >> se.DECAY_SHIFT)
+
+    def _softplus2(self, x):
+        shift = se.Q - se.G_STEP_BITS
+        u = np.abs(x)
+        i = np.minimum(u >> shift, se.G_LAST)
+        frac = u & ((1 << shift) - 1)
+        g0, g1 = self.g[i], self.g[i + 1]
+        g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0)
+        return np.maximum(x, 0) + g
+
+    def _keys(self, layer, tok, now, pin_since=None):
+        c = np.arange(self.S, dtype=np.int64)
+        oid = self.id_of_slot
+        held = oid >= 0
+        if self.policy == 0:
+            key = (self.usage << self.slot_bits) | c
+        else:
+            ids = np.where(held, oid, 0)
+            lk = ids // self.E
+            ahead = lk > layer
+            d = np.where(ahead, lk - layer, self.L - layer + lk)
+            k = np.where(ahead, tok - 1, tok) - self.last_tok[ids]
+            k = np.minimum(np.maximum(k, 0), se.K_MAX)
+            score = -((k << se.Q) + d * self.beta_step)
+            if self.policy >= 2:
+                lc = self.lc[ids]
+                lcv = self._decayed(ids, now)
+                lcv = np.where(lc == se.LC_NEVER, se.LC_FLOOR, np.maximum(lcv, se.LC_FLOOR))
+                score = score + ((self.w_q4 * lcv) >> 4)
+            assert np.abs(score).max() < 2**31  # the kernel computes it in int32
+            key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c)
+        evictable = self.usage != self.step
+        if pin_since is not None:
+            evictable &= self.usage <= pin_since
+        return np.where(evictable, key, _KEY_MAX)
+
+    def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None,
+               pin_since=None):
+        """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan.
+
+        ``pin_since`` (a ``step`` value) also pins every slot touched after it.
+
+        Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their
+        row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask.
+        """
+        self.step += 1
+        base = layer * self.E
+        tok = now = 0
+        if self.policy != 0:
+            if bump_tok:
+                self.tok += 1
+            tok = self.tok
+            now = tok * self.L + layer
+            if self.policy == 3 and update_state and logits is not None:
+                lg = np.asarray(logits, dtype=np.float32).reshape(-1, self.E)
+                routed = np.asarray(ids, dtype=np.int64).reshape(lg.shape[0], -1)
+                kth = lg[np.arange(lg.shape[0])[:, None], routed].min(axis=1)
+                near_miss = lg >= (kth - np.float32(thr))[:, None]
+            if self.policy == 3 and update_state and near_miss is not None:
+                near = base + np.flatnonzero(np.asarray(near_miss).reshape(-1, self.E).any(0))
+                self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1)
+        q = np.asarray(ids, dtype=np.int64).reshape(-1) + base
+        slots = self.slot_of_id[q]
+        self.usage[slots[slots >= 0]] = self.step
+        uniq = np.unique(q)
+        missing = uniq[self.slot_of_id[uniq] < 0]  # ascending, as the kernel ranks them
+        src = missing - base
+        dst = np.empty(missing.size, np.int64)
+        if missing.size:
+            keys = self._keys(layer, tok, now, pin_since)
+            part = np.argpartition(keys, missing.size - 1)[: missing.size]
+            victims = part[np.argsort(keys[part])]
+            for i, (e, v) in enumerate(zip(missing, victims)):
+                old = self.id_of_slot[v]
+                if old >= 0:
+                    self.slot_of_id[old] = -1
+                self.id_of_slot[v] = e
+                self.slot_of_id[e] = v
+                self.usage[v] = self.step
+                dst[i] = v
+        out = self.slot_of_id[q]
+        if self.policy != 0 and update_state:
+            self.last_tok[uniq] = tok
+            if self.policy >= 2:
+                lc = self.lc[uniq]
+                x = self._decayed(uniq, now)
+                self.lc[uniq] = np.where(lc == se.LC_NEVER, 0, self._softplus2(x))
+                self.ct[uniq] = now
+        return out, src, dst
+
+
+    def materialize(self, layer):
+        """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id."""
+        base = layer * self.E
+        same = (self.id_of_slot >= base) & (self.id_of_slot < base + self.E)
+        self.id_of_slot[same] = -1
+        self.usage[same] = 0
+        old = self.id_of_slot[: self.E].copy()
+        self.slot_of_id[old[old >= 0]] = -1
+        self.step += 1
+        self.id_of_slot[: self.E] = base + np.arange(self.E)
+        self.slot_of_id[base : base + self.E] = np.arange(self.E)
+        self.usage[: self.E] = self.step
+
+
+def replay(cache: RefScoredCache, rows, logits=None, thr=0.0, near_miss=None):
+    """Decode ``rows`` ``[N, L, K]`` (one step per row, all layers); per-row miss counts.
+
+    ``rows`` may be ``[N, L, B, K]`` for a batch of B rows per step, one K*B call per layer.
+    Policy 3 takes the matching ``[N, L, B, E]`` router ``logits`` (or a ``near_miss`` mask).
+    """
+    rows = np.asarray(rows)
+    misses = np.zeros(rows.shape[0], np.int64)
+    for r in range(rows.shape[0]):
+        for layer in range(rows.shape[1]):
+            lg = None if logits is None else logits[r, layer]
+            nm = None if near_miss is None else near_miss[r, layer]
+            _, src, _ = cache.ensure(layer, rows[r, layer], bump_tok=layer == 0, logits=lg, thr=thr, near_miss=nm)
+            misses[r] += src.size
+    return misses
diff --git a/tests/moe/test_fused_copy.py b/tests/moe/test_fused_copy.py
index 7a0c086444..c5336ea117 100644
--- a/tests/moe/test_fused_copy.py
+++ b/tests/moe/test_fused_copy.py
@@ -1,6 +1,8 @@
 """The fused multi-bank ``copy_missing`` path must move exactly the same bytes as the
 legacy per-bank ``fast_index_copy_jit`` loop, for every miss count (including the
-zero-copy case), across banks of differing per-row sizes.
+zero-copy case), across banks of differing per-row sizes. The opt-in slim kernel
+(``FREETOKEN_MOE_SLIM_COPY``) must move exactly the same bytes as the fused one, also when
+replayed from a CUDA graph with a device-side plan.
 """
 
 from __future__ import annotations
@@ -8,6 +10,7 @@
 import pytest
 import torch
 
+from freetoken.moe import offload_cache
 from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache
 
 CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
@@ -17,22 +20,30 @@
 # kernel/aot_models.py must be listed in its TEST_FEATURE_SIZES so the per-bank kernels
 # stay prebuilt under FREETOKEN_DISABLE_JIT=1.
 FEATS = [8192, 512, 256, 4096, 512, 256]
+# qwen4_exp (Qwen3.8-Flash-Next) triton-NVFP4 rows in _BANK_SCHEMAS["nvfp4"] order: gate_up
+# packed/scale/global, down packed/scale/global -- 2,772,480 B per expert.
+NVFP4_FEATS = [1638400, 204800, 2560, 819200, 102400, 5120]
+# 16 * odd, several rows shorter than one grid pass, so segments end mid-round and the
+# stride phase wraps across banks and rows.
+ODD_FEATS = [16 * 1031, 16 * 7, 16 * 250, 16, 16 * 4099, 16 * 33]
 
 
-def _build_cache(num_layers, num_experts, cache_size):
+def _build_cache(num_layers, num_experts, cache_size, quant_format="mxfp4_triton", feats=FEATS, pinned=False):
     dev = torch.device("cuda")
     cache = OffloadMoeCache(
         num_layers=num_layers, num_experts=num_experts, cache_size=cache_size,
-        device=dev, cache_policy="lru", prefill_overlap=False, quant_format="mxfp4_triton",
+        device=dev, cache_policy="lru", prefill_overlap=False, quant_format=quant_format,
     )
-    schema = _BANK_SCHEMAS["mxfp4_triton"]
+    schema = _BANK_SCHEMAS[quant_format]
     # Views into one flat tensor: only per-layer addressing matters here, not
     # independent allocations.
-    sources = {
-        name: list(torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev)
-                   .split(num_experts))
-        for name, feat in zip(schema, FEATS)
-    }
+    sources = {}
+    for name, feat in zip(schema, feats):
+        if pinned:
+            flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8).pin_memory()
+        else:
+            flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev)
+        sources[name] = list(flat.split(num_experts))
     cache.set_bank_sources(sources)  # also builds the fused-copy descriptor
     return cache
 
@@ -74,3 +85,147 @@ def test_fused_copy_matches_per_bank(num_indices):
 
     for b, (r, (_, c)) in enumerate(zip(ref, cache.banks)):
         assert torch.equal(r, c), f"bank {b} (feat={FEATS[b]}) fused != per-bank at num_indices={num_indices}"
+
+
+def _host_banks(feats, n_src, seed):
+    from freetoken.kernel.pinned import device_ptr
+
+    gen = torch.Generator().manual_seed(seed)
+    src = [torch.randint(0, 256, (n_src, f), dtype=torch.uint8, generator=gen).pin_memory() for f in feats]
+    dev = torch.device("cuda")
+    src_ptrs = torch.tensor([device_ptr(t) for t in src], dtype=torch.int64, device=dev)
+    feat_bytes = torch.tensor(feats, dtype=torch.int64, device=dev)
+    return src, src_ptrs, feat_bytes
+
+
+def _slot_banks(feats, n_slots):
+    dst = [torch.full((n_slots, f), 0xAB, dtype=torch.uint8, device="cuda") for f in feats]
+    return dst, torch.tensor([t.data_ptr() for t in dst], dtype=torch.int64, device="cuda")
+
+
+def _expected(src, feats, n_slots, di, si, num):
+    # untouched slots keep the 0xAB fill; plan rows land at their dst slots
+    out = []
+    for b, f in enumerate(feats):
+        e = torch.full((n_slots, f), 0xAB, dtype=torch.uint8)
+        e[di[:num].long().cpu()] = src[b][si[:num].long().cpu()]
+        out.append(e.cuda())
+    return out
+
+
+def _random_plan(gen, n_src, n_slots, plan, dtype):
+    di = torch.randperm(n_slots, generator=gen)[:plan].to(dtype).cuda()  # distinct: no write races
+    si = torch.randint(0, n_src, (plan,), generator=gen).to(dtype).cuda()
+    return di, si
+
+
+def _check_slim_matches_multi(feats, idx_dtype, **grid):
+    from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit
+
+    n_src, n_slots, plan = 24, 16, 16
+    src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=0)
+    gen = torch.Generator().manual_seed(1)
+    for num in [0, 1, 3, 10, plan, None]:  # None: no num_indices, the whole index length
+        di, si = _random_plan(gen, n_src, n_slots, plan, idx_dtype)
+        count = None if num is None else torch.tensor([num], dtype=torch.int64, device="cuda")
+        ref, ref_ptrs = _slot_banks(feats, n_slots)
+        fast_index_copy_multi_jit(ref_ptrs, src_ptrs, feat_bytes, di, si, count)
+        out, out_ptrs = _slot_banks(feats, n_slots)
+        fast_index_copy_multi_slim_jit(out_ptrs, src_ptrs, feat_bytes, di, si, count, **grid)
+        torch.cuda.synchronize()
+        exp = _expected(src, feats, n_slots, di, si, plan if num is None else num)
+        for b, f in enumerate(feats):
+            assert torch.equal(ref[b], exp[b]), f"multi reference wrong: bank {b} (feat={f}) num={num}"
+            assert torch.equal(out[b], ref[b]), f"slim != multi: bank {b} (feat={f}) num={num} grid={grid}"
+
+
+@CUDA
+@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"])
+def test_slim_copy_matches_multi(feats, idx_dtype):
+    _check_slim_matches_multi(feats, idx_dtype)
+
+
+@CUDA
+@pytest.mark.slow
+@pytest.mark.parametrize("grid", [(3, 128, 2), (16, 256, 2), (5, 64, 8)])
+@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"])
+def test_slim_copy_matches_multi_other_grids(feats, grid):
+    blocks, threads, unroll = grid
+    _check_slim_matches_multi(feats, torch.int32, num_blocks=blocks, num_threads=threads, unroll=unroll)
+
+
+@CUDA
+def test_slim_copy_graph_replay_follows_device_plan():
+    from freetoken.kernel.fast_index_copy import fast_index_copy_multi_slim_jit
+
+    feats, n_src, n_slots, plan = NVFP4_FEATS, 24, 16, 16
+    src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=2)
+    dst, dst_ptrs = _slot_banks(feats, n_slots)
+    di = torch.zeros(plan, dtype=torch.int32, device="cuda")
+    si = torch.zeros(plan, dtype=torch.int32, device="cuda")
+    count = torch.zeros(1, dtype=torch.int64, device="cuda")
+    fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count)  # JIT outside capture
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count)
+
+    gen = torch.Generator().manual_seed(3)
+    for num in [4, 0, plan, 1, 10]:
+        new_di, new_si = _random_plan(gen, n_src, n_slots, plan, torch.int32)
+        di.copy_(new_di)
+        si.copy_(new_si)
+        count.fill_(num)
+        for d in dst:
+            d.fill_(0xAB)
+        graph.replay()
+        torch.cuda.synchronize()
+        exp = _expected(src, feats, n_slots, di, si, num)
+        for b, f in enumerate(feats):
+            assert torch.equal(dst[b], exp[b]), f"graph replay: bank {b} (feat={f}) num={num}"
+
+
+def test_slim_copy_flag_defaults_off(monkeypatch):
+    import importlib.util
+
+    monkeypatch.delenv("FREETOKEN_MOE_SLIM_COPY", raising=False)
+    spec = importlib.util.find_spec("freetoken.moe.offload_cache")
+    fresh = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(fresh)  # a private copy: the imported module and its classes stay untouched
+    assert fresh._SLIM_COPY is False
+    monkeypatch.setenv("FREETOKEN_MOE_SLIM_COPY", "1")
+    spec.loader.exec_module(fresh)
+    assert fresh._SLIM_COPY is True
+
+
+@CUDA
+@pytest.mark.parametrize("slim", [False, True], ids=["multi", "slim"])
+def test_copy_missing_slim_switch(monkeypatch, slim):
+    from freetoken.kernel import fast_index_copy as fic
+
+    calls = []
+    for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"):
+        real = getattr(fic, name)
+        monkeypatch.setattr(fic, name, lambda *a, _n=name, _r=real, **k: (calls.append(_n), _r(*a, **k)))
+    monkeypatch.setattr(offload_cache, "_SLIM_COPY", slim)
+
+    num_layers, num_experts, cache_size, layer_id, num = 2, 8, 16, 1, 5
+    cache = _build_cache(num_layers, num_experts, cache_size, "nvfp4", NVFP4_FEATS, pinned=True)
+    assert cache._copy_fused_ok
+    cache._pending_src_layer = layer_id
+    dev = torch.device("cuda")
+    cache.num_indices.fill_(num)
+    cache.evict_slots[:num] = torch.tensor([9, 2, 15, 0, 7], dtype=torch.int32, device=dev)
+    cache.src_indices[:num] = torch.tensor([3, 3, 0, 7, 5], dtype=torch.int32, device=dev)
+    for _, c in cache.banks:
+        c.fill_(0xAB)
+    cache.copy_missing()
+    torch.cuda.synchronize()
+
+    assert calls == ["fast_index_copy_multi_slim_jit" if slim else "fast_index_copy_multi_jit"]
+    slots, rows = cache.evict_slots[:num].long().cpu(), cache.src_indices[:num].long().cpu()
+    for b, (per_layer, c) in enumerate(cache.banks):
+        exp = torch.full(c.shape, 0xAB, dtype=torch.uint8)
+        exp[slots] = per_layer[layer_id][rows]
+        assert torch.equal(c.cpu(), exp), f"bank {b} (feat={NVFP4_FEATS[b]})"
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 1f8245ace9..3e558c73af 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -387,7 +387,7 @@ def test_offload_moe_layer_decode_forward_uses_remapped_slot_ids(monkeypatch):
         lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids),
     )
 
-    def fake_ensure(layer_id, expert_ids):
+    def fake_ensure(layer_id, expert_ids, **kw):
         calls["ensure_layer_id"] = layer_id
         calls["ensure_expert_ids"] = expert_ids.clone()
         expert_ids.copy_(torch.tensor([[5, 0]], dtype=torch.int32))
@@ -427,6 +427,121 @@ def fake_fused_decode(
     assert calls["topk_ids"].tolist() == [[5, 0]]
 
 
+def test_decode_copy_overlap_follows_the_flag_and_the_decode_target(monkeypatch):
+    from freetoken.env import ENV
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    def cache(device, **kw):
+        return OffloadMoeCache(num_layers=3, num_experts=4, cache_size=8, device=torch.device(device), **kw)
+
+    monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", True)
+    assert cache("cpu").decode_copy_overlap is True  # unset -> FREETOKEN_MOE_COPY_OVERLAP
+    assert cache("cpu").decode_copy_stream is None
+    monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", False)
+    assert cache("cpu").decode_copy_overlap is False
+    if not torch.cuda.is_available():
+        return
+    assert cache("cuda").decode_copy_stream is None
+    for target in ("cpu", "hybrid"):
+        assert cache("cuda", decode_copy_overlap=True, decode_target=target).decode_copy_stream is None
+    on = cache("cuda", decode_copy_overlap=True)
+    assert on.decode_copy_stream is not None and len(on._decode_copy_events) == 3
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_copy_missing_takes_an_explicit_layer_id():
+    # a second copy path (the overlap fork, later prefetch) names its source layer instead of trusting whichever call staged last
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    num_layers, num_experts = 2, 4
+    fingerprint = torch.arange(num_layers * num_experts, dtype=torch.float32).view(num_layers, num_experts, 1, 1)
+    sources = {
+        "gate_up": [fingerprint[l].expand(num_experts, 32, 8).contiguous().pin_memory() for l in range(num_layers)],
+        "down": [fingerprint[l].expand(num_experts, 8, 16).contiguous().pin_memory() for l in range(num_layers)],
+    }
+    cache = OffloadMoeCache(
+        num_layers=num_layers, num_experts=num_experts, cache_size=num_experts, device=torch.device("cuda")
+    )
+    cache.set_bank_sources(sources)
+    cache.reset()
+    ids = torch.tensor([2, 3], dtype=torch.int32, device="cuda")
+    cache.ensure_experts(1, ids)
+    cache._pending_src_layer = 0
+    cache.copy_missing(1)
+    torch.cuda.synchronize()
+    for bank in cache.bank_views():
+        assert [int(bank[s].flatten()[0]) for s in ids.tolist()] == [num_experts + 2, num_experts + 3]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_decode_side_applies_only_to_gpu_decode(monkeypatch):
+    # prefill, CPU-routed layers, TP > 1 and the flag off keep the shared expert on the compute stream
+    from types import SimpleNamespace
+
+    import freetoken.core as core
+    from freetoken.core import Context
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    layer = _bf16_offload_layer(0, 4, 2, 8, 16)
+    ctx = Context(page_size=1)
+    monkeypatch.setattr(core, "_GLOBAL_CTX", ctx)
+
+    def applies(is_prefill=False, cpu_layers=(), **cache_kw):
+        layer.offload_cache = OffloadMoeCache(
+            num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), **cache_kw
+        )
+        layer.offload_cache.cpu_layer_ids = frozenset(cpu_layers)
+        ctx._batch = SimpleNamespace(is_prefill=is_prefill)
+        return layer.decode_side_applies()
+
+    assert not applies(decode_copy_overlap=False)
+    assert not applies(decode_copy_overlap=True, decode_target="hybrid")
+    assert not applies(is_prefill=True, decode_copy_overlap=True)
+    assert not applies(cpu_layers=(0,), decode_copy_overlap=True)
+    monkeypatch.setattr(layer, "tp_size", 2)
+    assert not applies(decode_copy_overlap=True)
+    monkeypatch.setattr(layer, "tp_size", 1)
+    assert applies(decode_copy_overlap=True)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_decode_side_stream_never_aliases_a_pooled_stream():
+    # torch.cuda.Stream() hands out 32 pooled streams per priority round robin; a pooled side stream
+    # can come back as the engine or capture stream, and then the side work silently runs serially
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(
+        num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True
+    )
+    side = cache.decode_copy_stream.cuda_stream
+    pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)}
+    assert side not in pooled and side != torch.cuda.current_stream().cuda_stream
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_decode_side_survives_a_rebuild():
+    # the side stream and events are not cache_size-shaped, so a runtime rebuild keeps them working
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True
+    )
+    cache.set_bank_sources({
+        "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+        "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+    })
+    stream = cache.decode_copy_stream
+    cache.rebuild(6)
+    assert cache.decode_copy_stream is stream is not None
+    x = torch.arange(1024, dtype=torch.float32, device="cuda")
+    for layer_id in range(2):
+        with cache.decode_side(layer_id):
+            assert torch.cuda.current_stream() == stream
+            y = x * 2
+        cache.join_decode_side(layer_id)
+        assert torch.equal(y + 0, x * 2)
+
 
 def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch():
     import pytest
diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py
new file mode 100644
index 0000000000..db734180b6
--- /dev/null
+++ b/tests/moe/test_scored_ensure.py
@@ -0,0 +1,622 @@
+"""--moe-cache-policy kd / kdfb / rule: the vendored scored-victim ensure kernel.
+
+The kernel must reproduce flashlib's LRU bit for bit under POLICY=lru, and match the CPU
+reference (ref_scored_cache.py, a port of the eviction replay) victim for victim under the
+scored policies. Set FREETOKEN_MOE_ROUTING_TRACE to a compact routing trace directory
+(experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows.
+"""
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+L, E, S, TOP_K = 48, 512, 1650, 10
+NEAR_MISS_THR = 0.25
+_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+_TRACE = os.environ.get("FREETOKEN_MOE_ROUTING_TRACE", "")
+
+
+def _trace():
+    if not _TRACE or not (Path(_TRACE) / "experts.npy").exists():
+        pytest.skip("FREETOKEN_MOE_ROUTING_TRACE is not set to a routing trace directory")
+    load = lambda name: np.load(Path(_TRACE) / name, mmap_mode="r")  # noqa: E731
+    return load("experts.npy"), load("logit_idx.npy"), load("logit_val.npy")
+
+
+def _trace_logits(li, lv):
+    """Router logits rebuilt from the trace's recorded top-32 (the rest far below any margin)."""
+    li, lv = np.asarray(li).astype(np.int64), np.asarray(lv).astype(np.float32)
+    logits = np.full(li.shape[:-1] + (E,), -1e4, np.float32)
+    np.put_along_axis(logits, li, lv, axis=-1)
+    return logits
+
+
+def _zipf_rows(n, rows_per_step=1, seed=0, num_layers=L, num_experts=E, top_k=TOP_K, skew=0.9):
+    """Skewed decode routing ``[n, L, rows_per_step, top_k]``: distinct ids within a row, duplicates across rows."""
+    rng = np.random.default_rng(seed)
+    p = 1.0 / np.arange(1, num_experts + 1) ** skew
+    perm = [rng.permutation(num_experts) for _ in range(num_layers)]
+    out = np.empty((n, num_layers, rows_per_step, top_k), np.int64)
+    for r in range(n):
+        for layer in range(num_layers):
+            for b in range(rows_per_step):
+                pick = rng.choice(num_experts, top_k, replace=False, p=p / p.sum())
+                out[r, layer, b] = perm[layer][pick]
+    return out
+
+
+def _random_logits(rows, seed=1, num_experts=E):
+    """bf16-exact router logits shaped like ``rows`` with the routed ids on top (~10% near misses at 0.25)."""
+    rng = np.random.default_rng(seed)
+    logits = rng.normal(size=rows.shape[:-1] + (num_experts,)).astype(np.float32)
+    np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, axis=-1) + 3.0, axis=-1)
+    return torch.from_numpy(logits).bfloat16().float().numpy()
+
+
+def _cache(policy, cache_size=S, num_layers=L, num_experts=E, device="cuda", **kw):
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    return OffloadMoeCache(
+        num_layers=num_layers, num_experts=num_experts, cache_size=cache_size,
+        device=torch.device(device), cache_policy=policy, **kw,
+    )
+
+
+def _ref(policy, cache_size=S, num_layers=L, num_experts=E):
+    from freetoken.moe.scored_ensure import POLICY_IDS
+
+    from .ref_scored_cache import RefScoredCache
+
+    return RefScoredCache(num_layers, num_experts, cache_size, POLICY_IDS[policy])
+
+
+def _run_gpu(cache, rows, logits=None, dtype=torch.bfloat16):
+    """Decode ``rows`` ``[N, L, B, K]`` through ``cache.ensure_experts``; per-row miss counts."""
+    dev = cache.device
+    ids = torch.from_numpy(np.ascontiguousarray(rows, dtype=np.int32)).to(dev)
+    lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits)).to(dev, dtype)
+    cache.collect_stats = True
+    cache.lru_stats.zero_()
+    marks = torch.zeros(rows.shape[0], dtype=torch.int64, device=dev)
+    for r in range(rows.shape[0]):
+        for layer in range(rows.shape[1]):
+            cache.ensure_experts(layer, ids[r, layer], router_logits=None if lg is None else lg[r, layer])
+        marks[r] = cache.lru_stats[:, 1].sum()
+    return np.diff(marks.cpu().numpy(), prepend=0)
+
+
+def _run_ref(ref, rows, logits=None):
+    from .ref_scored_cache import replay
+
+    return replay(ref, rows, logits, NEAR_MISS_THR)
+
+
+def _assert_same_tables(cache, ref):
+    got = lambda t: t.detach().cpu().numpy().astype(np.int64)  # noqa: E731
+    np.testing.assert_array_equal(got(cache.slot_for_id.view(-1)), ref.slot_of_id)
+    np.testing.assert_array_equal(got(cache.id_of_slot), ref.id_of_slot)
+    np.testing.assert_array_equal(got(cache.usage), ref.usage)
+    assert int(cache.step) == ref.step
+    if ref.policy:
+        assert int(cache.evict_tok) == ref.tok
+        np.testing.assert_array_equal(got(cache.evict_last_tok), ref.last_tok)
+    if ref.policy >= 2:
+        np.testing.assert_array_equal(got(cache.evict_lc), ref.lc)
+        np.testing.assert_array_equal(got(cache.evict_ct), ref.ct)
+    if ref.policy:
+        _assert_mirrors_hold(cache)
+
+
+def _assert_mirrors_hold(cache):
+    """Every slot whose owner tag matches its id mirrors that id's state."""
+    got = lambda t: t.detach().cpu().numpy().astype(np.int64)  # noqa: E731
+    ids, owner = got(cache.id_of_slot), got(cache.evict_slot_owner)
+    tagged = (ids >= 0) & (owner == ids)
+    np.testing.assert_array_equal(got(cache.evict_slot_last_tok)[tagged], got(cache.evict_last_tok)[ids[tagged]])
+    if cache.evict_lc is not None:
+        np.testing.assert_array_equal(got(cache.evict_slot_lc)[tagged], got(cache.evict_lc)[ids[tagged]])
+        np.testing.assert_array_equal(got(cache.evict_slot_ct)[tagged], got(cache.evict_ct)[ids[tagged]])
+    return tagged
+
+
+def _check_against_ref(policy, rows, logits=None, dtype=torch.bfloat16):
+    cache, ref = _cache(policy), _ref(policy)
+    assert policy != "rule" or cache.near_miss_thr == NEAR_MISS_THR
+    got = _run_gpu(cache, rows, logits, dtype)
+    want = _run_ref(ref, rows, logits)
+    np.testing.assert_array_equal(got, want)
+    _assert_same_tables(cache, ref)
+    return got
+
+
+# (a) POLICY=lru is flashlib's kernel -------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_vendored_lru_is_bit_identical_to_flashlib(rows_per_step):
+    from flashlib.kernels.slot_cache import N_STATS, lru_ensure
+
+    from freetoken.moe.scored_ensure import scored_ensure
+
+    dev = torch.device("cuda")
+    k = rows_per_step * TOP_K
+    # skewed enough that some calls hit every id
+    stream = _zipf_rows(63, rows_per_step, seed=rows_per_step, skew=1.6).reshape(-1, k)[:3000]
+    layers = np.arange(stream.shape[0]) % L
+
+    def state():
+        return dict(
+            slot_of_id=torch.full((L * E,), -1, dtype=torch.int32, device=dev),
+            id_of_slot=torch.full((S,), -1, dtype=torch.int32, device=dev),
+            usage=torch.zeros(S, dtype=torch.int64, device=dev),
+            step=torch.zeros((), dtype=torch.int64, device=dev),
+            src=torch.zeros(S, dtype=torch.int32, device=dev),
+            dst=torch.zeros(S, dtype=torch.int32, device=dev),
+            num=torch.zeros((), dtype=torch.int64, device=dev),
+            stats=torch.zeros(N_STATS, dtype=torch.int64, device=dev),
+        )
+
+    a, b = state(), state()
+    queries = torch.from_numpy(stream.astype(np.int32)).to(dev)
+    qa, qb = queries.clone(), queries.clone()
+    plans = ([], [])
+    for i in range(stream.shape[0]):
+        base = int(layers[i]) * E
+        lru_ensure(qa[i], a["slot_of_id"], a["id_of_slot"], a["usage"], a["step"], qa[i],
+                   a["src"], a["dst"], a["num"], stats=a["stats"], id_base=base)
+        scored_ensure(qb[i], b["slot_of_id"], b["id_of_slot"], b["usage"], b["step"], qb[i],
+                      b["src"], b["dst"], b["num"], policy=0, num_layers=L, num_experts=E,
+                      stats=b["stats"], id_base=base)
+        for st, plan in zip((a, b), plans):
+            plan.append(torch.cat([st["num"].view(1), st["src"][:k].long(), st["dst"][:k].long()]))
+    assert torch.equal(qa, qb), "out slots differ"
+    pa, pb = (torch.stack(plan) for plan in plans)
+    n = pa[:, 0]
+    assert int(n.sum()) > 1000 and int((n == 0).sum()) > 0  # the stream exercises misses and all-hit calls
+    for i in range(pa.shape[0]):
+        m = int(n[i])
+        assert torch.equal(pa[i, : 1 + m], pb[i, : 1 + m]) and torch.equal(pa[i, 1 + k : 1 + k + m], pb[i, 1 + k : 1 + k + m])
+    for key in ("slot_of_id", "id_of_slot", "usage", "step", "stats"):
+        assert torch.equal(a[key], b[key]), key
+
+
+def test_lru_policy_stays_on_flashlib(monkeypatch):
+    import freetoken.moe.offload_kernels as ok
+
+    cache = _cache("lru", device="cpu")
+    calls = []
+    monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib"))
+    monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append("scored"))
+    cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32))
+    cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32), update_state=False)
+    assert calls == ["flashlib", "flashlib"]
+    assert cache.evict_tok is None and cache.evict_last_tok is None and cache.near_miss_thr is None
+
+
+# (b) scored policies == CPU reference ------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"])
+def test_scored_kernel_matches_cpu_reference_on_synthetic_routing(policy):
+    rows = _zipf_rows(40, seed=3)
+    got = _check_against_ref(policy, rows, _random_logits(rows) if policy == "rule" else None)
+    assert got.sum() > 0
+
+
+@_cuda
+@pytest.mark.slow
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_scored_kernel_matches_cpu_reference_on_trace_rows(policy):
+    experts, li, lv = _trace()
+    rows = np.asarray(experts[:2000]).astype(np.int64)[:, :, None, :]
+    logits = _trace_logits(li[:2000], lv[:2000])[:, :, None, :] if policy == "rule" else None
+    got = _check_against_ref(policy, rows, logits, torch.float32)
+    lru = _run_ref(_ref("lru"), rows)
+    assert got.sum() < lru.sum()  # the point of the policy
+
+
+# (c) victims: never the call's own ids, empty slots first ----------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_victims_skip_the_call_and_take_empty_slots_first(policy, rows_per_step):
+    from freetoken.moe import scored_ensure as se
+
+    num_layers, num_experts, cache_size = 4, 64, 100
+    cache = _cache(policy, cache_size, num_layers, num_experts)
+    ref = _ref(policy, cache_size, num_layers, num_experts)
+    warm = _zipf_rows(6, rows_per_step, seed=5, num_layers=num_layers, num_experts=num_experts)
+    _run_gpu(cache, warm)
+    _run_ref(ref, warm)
+    assert int((cache.id_of_slot >= 0).sum()) == cache_size
+
+    layer = 2
+    held = cache.id_of_slot.cpu().numpy()
+    resident = sorted(int(i) - layer * num_experts for i in held if layer * num_experts <= i < (layer + 1) * num_experts)
+    absent = sorted(set(range(num_experts)) - set(resident))
+    k = rows_per_step * TOP_K
+    hits = resident[: k // 2]
+    ids = np.array((hits + absent)[:k], np.int64)
+    if rows_per_step == 2:
+        ids[-1] = ids[0]  # a duplicate hit
+    # make this call's hits the stalest ids in the cache, so only the pin keeps them
+    hit_keys = [layer * num_experts + e for e in hits]
+    for key in hit_keys:
+        cache.evict_last_tok[key] = se.LAST_TOK_NEVER
+        ref.last_tok[key] = se.LAST_TOK_NEVER
+        if cache.evict_lc is not None:
+            cache.evict_lc[key] = se.LC_NEVER
+            ref.lc[key] = se.LC_NEVER
+    hit_slots = {int(cache.slot_for_id[layer, e]) for e in hits}
+    # invalidate three slots (as the prefill overlap does): they must be the first victims
+    empty = sorted([s for s in (7, 40, 91, 13, 55, 77) if s not in hit_slots][:3])
+    for slot in empty:
+        for t, v in ((cache.id_of_slot, held[slot]), (ref.id_of_slot, held[slot])):
+            t[slot] = -1
+        cache.slot_for_id.view(-1)[int(held[slot])] = -1
+        ref.slot_of_id[int(held[slot])] = -1
+        cache.usage[slot] = 0
+        ref.usage[slot] = 0
+
+    q = torch.tensor(ids, dtype=torch.int32, device="cuda")
+    lg = _random_logits(ids.reshape(rows_per_step, TOP_K), seed=9, num_experts=num_experts)
+    cache.ensure_experts(layer, q, router_logits=torch.from_numpy(lg).cuda())
+    out, src, dst = ref.ensure(layer, ids, logits=lg, thr=NEAR_MISS_THR)
+
+    n = int(cache.num_indices)
+    got_dst = cache.evict_slots[:n].tolist()
+    assert n == len(set(ids.tolist()) - set(hits)) == src.size
+    assert got_dst[: len(empty)] == empty
+    assert not hit_slots & set(got_dst)
+    assert q.tolist() == out.tolist()
+    assert all(int(cache.slot_for_id[layer, e]) >= 0 for e in ids.tolist())
+    assert got_dst == dst.tolist()
+    _assert_same_tables(cache, ref)
+
+
+# (d) the decode-step clock -----------------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_token_clock_moves_once_per_decode_step(rows_per_step):
+    cache = _cache("kdfb")
+    rows = _zipf_rows(5, rows_per_step, seed=7)
+    ids = torch.from_numpy(rows.astype(np.int32)).cuda()
+    for r in range(rows.shape[0]):
+        for layer in range(L):
+            cache.ensure_experts(layer, ids[r, layer])
+            assert int(cache.evict_tok) == r + 1
+    # a CPU-decoded first layer moves the clock to the first GPU layer
+    cache.reset()
+    cache.cpu_layer_ids = frozenset({0, 1})
+    ids = torch.from_numpy(rows.astype(np.int32)).cuda()
+    for layer in range(2, L):
+        cache.ensure_experts(layer, ids[0, layer])
+    assert int(cache.evict_tok) == 1
+
+
+@_cuda
+def test_small_prefill_ensures_leave_the_state_alone(monkeypatch):
+    import freetoken.layers.moe as moe_mod
+    from freetoken.distributed import set_tp_info, try_get_tp_info
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    dev = torch.device("cuda")
+    experts, top_k, hidden, inter = 64, 4, 16, 32
+    layer = OffloadMoELayer(0, experts, top_k, hidden, inter, quant_config=NoQuantConfig(),
+                            prefix="model.layers.0.mlp.experts")
+    cache = _cache("kdfb", 96, 1, experts)
+    g = torch.Generator().manual_seed(0)
+    cache.set_bank_sources({
+        "gate_up": [torch.randn(experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1],
+        "down": [torch.randn(experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1],
+    })
+    layer.offload_cache = cache
+    decode = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=dev)
+    cache.ensure_experts(0, decode)
+    before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)]
+
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 64)
+    tokens = 40
+    ids = torch.randint(0, experts // 2, (tokens, top_k), dtype=torch.int32, device=dev)
+    w = torch.full((tokens, top_k), 1.0 / top_k, device=dev)
+    layer._prefill_routed(torch.randn(tokens, hidden, dtype=torch.bfloat16, device=dev), w, ids.clone())
+    assert bool((cache.slot_for_id[0, ids.long().unique()] >= 0).all())  # it did make them resident
+    assert int((cache.slot_for_id[0] >= 0).sum()) < experts  # through the on-demand path, not the whole layer
+    after = (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)
+    for b, a in zip(before, after):
+        assert torch.equal(b, a)
+
+
+@_cuda
+def test_dsv4_on_demand_prefill_leaves_the_state_alone():
+    from types import SimpleNamespace
+
+    from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer
+
+    experts, top_k = 64, 4
+    cache = _cache("kdfb", 96, 1, experts)
+    cache.set_bank_sources({"gate_up": [torch.zeros(experts, 2, 4)], "down": [torch.zeros(experts, 4, 2)]})
+    cache.ensure_experts(0, torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="cuda"))
+    before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)]
+    # few enough routes (3 x 4 < 64) for the slot path rather than whole-layer streaming
+    layer = SimpleNamespace(offload_cache=cache, num_experts=experts, top_k=top_k, layer_id=0,
+                            _expert_gemm=lambda *a, **kw: None)
+    ids = torch.tensor([[5, 6, 7, 8], [9, 10, 11, 12], [5, 13, 14, 15]], dtype=torch.int32, device="cuda")
+    DSV4OffloadMoELayer._prefill_routed(layer, torch.zeros(3, 8, device="cuda"), torch.full((3, top_k), 0.25), ids)
+    assert bool((cache.slot_for_id[0, 5:16] >= 0).all())
+    for b, a in zip(before, (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)):
+        assert torch.equal(b, a)
+
+
+# (e) reset() and rebuild() -----------------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_reset_and_rebuild_clear_the_state(policy):
+    from freetoken.moe import scored_ensure as se
+
+    cache = _cache(policy, 1480)
+    cache.set_bank_sources({"gate_up": [torch.zeros(E, 2, 4)] * L, "down": [torch.zeros(E, 4, 2)] * L})
+    rows = _zipf_rows(3, seed=11)
+    fresh = _ref(policy)
+
+    def assert_cold():
+        assert int(cache.evict_tok) == 0
+        assert bool((cache.evict_last_tok == se.LAST_TOK_NEVER).all())
+        assert bool((cache.evict_slot_owner == -1).all()) and cache.evict_slot_owner.numel() == cache.cache_size
+        if cache.evict_lc is not None:
+            assert bool((cache.evict_lc == se.LC_NEVER).all()) and bool((cache.evict_ct == 0).all())
+
+    for clear in (cache.reset, lambda: cache.rebuild(S)):
+        _run_gpu(cache, rows)
+        assert int(cache.evict_tok) == rows.shape[0]
+        clear()
+        assert_cold()
+    assert cache.cache_size == S
+    # a rebuilt cache replays exactly like a fresh one
+    np.testing.assert_array_equal(_run_gpu(cache, rows), _run_ref(fresh, rows))
+    _assert_same_tables(cache, fresh)
+
+
+# (f) CUDA graph replay == eager ------------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kdfb", "rule"])
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_graph_replay_matches_eager(policy, rows_per_step):
+    steps = 50
+    rows = _zipf_rows(steps + 1, rows_per_step, seed=13)
+    ids = torch.from_numpy(rows.astype(np.int32)).cuda()
+    logits = torch.from_numpy(_random_logits(rows, seed=14)).cuda().bfloat16()
+
+    eager = _cache(policy)
+    for r in range(1, steps + 1):
+        for layer in range(L):
+            eager.ensure_experts(layer, ids[r, layer].clone(), router_logits=logits[r, layer])
+
+    graphed = _cache(policy)
+    static_ids = ids[0].clone()
+    static_logits = logits[0].clone()
+
+    def step():
+        for layer in range(L):
+            graphed.ensure_experts(layer, static_ids[layer], router_logits=static_logits[layer])
+
+    step()  # warm-up compiles the kernel variants outside capture
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        step()
+    graphed.reset()  # as graph.py does after capture
+    for r in range(1, steps + 1):
+        static_ids.copy_(ids[r])
+        static_logits.copy_(logits[r])
+        graph.replay()
+    torch.cuda.synchronize()
+
+    assert int(graphed.evict_tok) == steps
+    for name in ("slot_for_id", "id_of_slot", "usage", "step", "evict_tok", "evict_last_tok", "evict_lc", "evict_ct",
+                 "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"):
+        assert torch.equal(getattr(eager, name), getattr(graphed, name)), name
+
+
+# (g) bs=2: two rows per call ---------------------------------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_two_rows_per_call_match_cpu_reference(policy):
+    rows = _zipf_rows(30, 2, seed=17)
+    assert any(len(set(rows[r, layer].ravel())) < 2 * TOP_K for r in range(30) for layer in range(L))
+    _check_against_ref(policy, rows, _random_logits(rows, seed=18) if policy == "rule" else None)
+
+
+@_cuda
+@pytest.mark.slow
+@pytest.mark.parametrize("policy", ["kdfb", "rule"])
+def test_two_trace_streams_per_call_match_cpu_reference(policy):
+    experts, li, lv = _trace()
+    a, b = slice(0, 600), slice(45000, 45600)  # two different requests decoding together
+    rows = np.stack([np.asarray(experts[a]), np.asarray(experts[b])], axis=2).astype(np.int64)
+    logits = None
+    if policy == "rule":
+        logits = np.stack([_trace_logits(li[a], lv[a]), _trace_logits(li[b], lv[b])], axis=2)
+    _check_against_ref(policy, rows, logits, torch.float32)
+
+
+# the slot mirrors: foreign installs and state-free ensures fall back to the per-id state -----
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_materialized_prefill_between_decodes_matches_cpu_reference(policy):
+    cache, ref = _cache(policy), _ref(policy)
+    rows = _zipf_rows(24, seed=21)
+    logits = _random_logits(rows, seed=22) if policy == "rule" else None
+    # layer 5 comes back to slots it lost to layer 9: its tags there would otherwise look current
+    for part, layer in ((slice(0, 6), 5), (slice(6, 12), 9), (slice(12, 18), 5), (slice(18, 24), None)):
+        lg = None if logits is None else logits[part]
+        np.testing.assert_array_equal(_run_gpu(cache, rows[part], lg), _run_ref(ref, rows[part], lg))
+        _assert_same_tables(cache, ref)
+        if layer is not None:
+            # the whole-layer prefill path re-seats this layer's experts at slots [0, E)
+            cache.materialize_layer(layer)
+            ref.materialize(layer)
+            assert not bool((cache.evict_slot_owner[:E] >= 0).any())
+            _assert_same_tables(cache, ref)
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb"])
+def test_materialize_untags_the_slots_it_fills(policy):
+    # expert 3 of layer 0 leaves slot 3, is used elsewhere, then materialize puts it back at slot 3
+    cache = _cache(policy, 24, 2, 8)
+    dev = torch.device("cuda")
+    ids = lambda *e: torch.tensor(e, dtype=torch.int32, device=dev)  # noqa: E731
+    cache.materialize_layer(0)
+    cache.ensure_experts(0, ids(3, 4))  # a hit at slot 3 tags it with expert 3
+    assert int(cache.evict_slot_owner[3]) == 3
+    cache.materialize_layer(1)  # slot 3 now holds layer 1's expert 3
+    cache.ensure_experts(1, ids(0, 1))
+    cache.ensure_experts(0, ids(3, 5))  # expert 3 comes back elsewhere, with new state
+    assert int(cache.slot_for_id[0, 3]) != 3
+    cache.materialize_layer(0)
+    assert int(cache.slot_for_id[0, 3]) == 3
+    _assert_mirrors_hold(cache)
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+@pytest.mark.parametrize("chunks", [1, 2])
+def test_state_free_ensures_between_decodes_match_cpu_reference(policy, chunks):
+    cache, ref = _cache(policy), _ref(policy)
+    rows = _zipf_rows(12, seed=23)
+    extra = _zipf_rows(4, 3, seed=24, skew=0.3)
+    logits = _random_logits(rows, seed=25) if policy == "rule" else None
+    for r in range(rows.shape[0]):
+        lg = None if logits is None else logits[r : r + 1]
+        np.testing.assert_array_equal(_run_gpu(cache, rows[r : r + 1], lg), _run_ref(ref, rows[r : r + 1], lg))
+        if r % 3 == 2:
+            for layer in range(0, L, 7):  # a small prefill's chunked ensures
+                ids = np.unique(extra[r // 3, layer])
+                # two chunks pin each other's slots as _ensure_unique does; one chunk keeps the plain pin
+                pin = cache.step.clone() if chunks > 1 else None
+                ref_pin = ref.step if chunks > 1 else None
+                for part in np.array_split(ids, chunks):
+                    q = torch.from_numpy(part.astype(np.int32)).cuda()
+                    cache.ensure_experts(layer, q, update_state=False, pin_since=pin)
+                    out, _, dst = ref.ensure(layer, part, update_state=False, pin_since=ref_pin)
+                    assert q.tolist() == out.tolist()
+                    assert bool((cache.evict_slot_owner[torch.from_numpy(dst).cuda()] == -1).all())
+                assert bool((cache.slot_for_id[layer, torch.from_numpy(ids).cuda()] >= 0).all())
+        _assert_same_tables(cache, ref)
+    assert not _assert_mirrors_hold(cache).all()  # some slots still read the per-id state
+
+
+# rule: the near misses come from the router logits ----------------------------------------
+
+
+@_cuda
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_rule_refreshes_the_router_near_misses(rows_per_step):
+    from freetoken.moe import scored_ensure as se
+
+    cache = _cache("rule")
+    layer = 5
+    g = torch.Generator(device="cuda").manual_seed(rows_per_step)
+    logits = (torch.randn(rows_per_step, E, device="cuda", generator=g) * 1.5).bfloat16()
+    ids = torch.topk(logits.float(), TOP_K, dim=-1).indices.to(torch.int32)
+    cache.ensure_experts(0, torch.arange(TOP_K, dtype=torch.int32, device="cuda"))  # step 1 starts
+    cache.ensure_experts(layer, ids.clone(), router_logits=logits)
+    kth = logits.float().gather(1, ids.long()).min(dim=1, keepdim=True).values
+    near = (logits.float() >= kth - NEAR_MISS_THR).any(0)
+    routed = torch.zeros(E, dtype=torch.bool, device="cuda")
+    routed[ids.long().view(-1)] = True
+    last = cache.evict_last_tok[layer * E : (layer + 1) * E]
+    assert bool((last[routed] == 1).all())
+    assert bool((last[near & ~routed] == 0).all())
+    assert bool((last[~near & ~routed] == se.LAST_TOK_NEVER).all())
+    assert 0 < int((near & ~routed).sum()) < E - TOP_K * rows_per_step
+
+
+@pytest.mark.parametrize("policy", ["lru", "kdfb", "rule"])
+def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy):
+    from freetoken.distributed import set_tp_info, try_get_tp_info
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    layer = OffloadMoELayer(0, 8, 2, 8, 16, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts")
+    layer.offload_cache = cache = _cache(policy, 12, 1, 8, device="cpu")
+    seen = {}
+    monkeypatch.setattr(
+        "freetoken.layers.moe.fused_topk",
+        lambda **kw: (torch.full((2, 2), 0.5), torch.zeros((2, 2), dtype=torch.int32)),
+    )
+    monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw))
+    monkeypatch.setattr(cache, "copy_missing", lambda: None)
+    monkeypatch.setattr(cache, "bank_views", lambda: ())
+    monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None)
+    router_logits = torch.randn(2, 8)
+    layer.decode_forward(torch.randn(2, 8), router_logits)
+    assert seen["router_logits"] is router_logits
+    assert cache.near_miss_thr == (NEAR_MISS_THR if policy == "rule" else None)
+
+
+def test_policy_flag_and_gpu_only():
+    import contextlib
+    import io
+
+    from freetoken.server.args import parse_args
+
+    buf = io.StringIO()
+    with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit):
+        parse_args(["--help"])
+    assert "{lru,kd,kdfb,rule}" in buf.getvalue()
+    for policy in ("kd", "kdfb", "rule"):
+        _cache(policy, 12, 1, 8, device="cpu")
+        with pytest.raises(ValueError, match="GPU decode"):
+            _cache(policy, 12, 1, 8, device="cpu", decode_target="hybrid")
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy):
+    # 5 slots pad to an 8-lane block; with every slot pinned the sentinel key must not decode to lane 7.
+    # The cache enforces cache_size >= num_experts, so the kernel gets 5-slot views of a 16-slot cache.
+    from freetoken.moe.offload_kernels import ensure_experts_scored
+
+    s, pad = 5, 3
+    cache = _cache(policy, cache_size=16, num_layers=1, num_experts=16)
+    guarded = {}
+    for name in ("id_of_slot", "usage", "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"):
+        t = getattr(cache, name)
+        if t is None:  # kd keeps no lc/ct mirrors
+            continue
+        buf = torch.full((s + pad,), 7, dtype=t.dtype, device=t.device)
+        buf[:s] = t[:s]
+        setattr(cache, name, buf[:s])
+        guarded[name] = buf
+    dev = cache.id_of_slot.device
+    ensure_experts_scored(cache, 0, torch.arange(s, dtype=torch.int32, device=dev), bump_tok=True, update_state=True)
+    # 6 distinct ids against 5 slots breaks the caller contract; the kernel must still stay in bounds
+    q = torch.arange(s + 1, dtype=torch.int32, device=dev)
+    ensure_experts_scored(cache, 0, q, bump_tok=True, update_state=True)
+    torch.cuda.synchronize()
+    for name, buf in guarded.items():
+        assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}"
+    assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index eea5842d23..70cec00019 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -8,7 +8,7 @@
 from freetoken.distributed import set_tp_info, try_get_tp_info
 
 
-def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None):
+def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None, policy="lru", num_layers=1):
     from freetoken.layers.moe import OffloadMoELayer
     from freetoken.layers.quantization import NoQuantConfig
     from freetoken.moe.offload_cache import OffloadMoeCache
@@ -17,11 +17,14 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_
         set_tp_info(rank=0, size=1)
     layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(),
                             prefix="model.layers.0.mlp.experts")
-    cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device)
+    cache = OffloadMoeCache(num_layers=num_layers, num_experts=num_experts, cache_size=cache_size or num_experts,
+                            device=device, cache_policy=policy)
     g = torch.Generator().manual_seed(0)
     cache.set_bank_sources({
-        "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1],
-        "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1],
+        "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1
+                    for _ in range(num_layers)],
+        "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1
+                 for _ in range(num_layers)],
     })
     layer.offload_cache = cache
     return layer, cache
@@ -34,7 +37,7 @@ def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand):
     layer, cache = _layer_and_cache(torch.device("cpu"))
     monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
     calls = []
-    monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure"))
+    monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids, **kw: calls.append("ensure"))
     monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize"))
     monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy"))
     monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False)
@@ -96,7 +99,35 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch):
         monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
         layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100)
         widths, ensure = [], cache.ensure_experts
-        monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1])
+        monkeypatch.setattr(cache, "ensure_experts", lambda lid, q, **kw: (widths.append(q.numel()), ensure(lid, q, **kw))[1])
         outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
     assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's"
     torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"])
+def test_a_later_chunk_keeps_the_experts_an_earlier_chunk_loaded(monkeypatch, policy):
+    # Decode fills every slot, then 40 never-decoded experts take two ensure chunks. A scored
+    # policy ranks the first chunk's fresh installs as its coldest slots, so they need a pin.
+    dev = torch.device("cuda")
+    experts, top_k, slots, tokens = 64, 4, 80, 10
+    g = torch.Generator(device=dev).manual_seed(2)
+    hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 4  # a wrong expert is ~1.0 off
+    ids = torch.arange(16, 16 + tokens * top_k, dtype=torch.int32, device=dev).view(tokens, top_k)
+    w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1)
+    outs = {}
+    for threshold in (0, 64):
+        monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+        layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=slots, policy=policy,
+                                        num_layers=2)
+        for layer_id, first, last in ((1, 0, experts), (0, 0, 16)):
+            for start in range(first, last, 8):
+                cache.ensure_experts(layer_id, torch.arange(start, start + 8, dtype=torch.int32, device=dev))
+                cache.copy_missing()
+        assert bool((cache.id_of_slot >= 0).all())
+        outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
+        if threshold:
+            slot_of = cache.slot_for_id[0, ids.long()]
+            assert bool((slot_of >= 0).all()), f"{int((slot_of < 0).sum())} routed experts were evicted"
+    torch.testing.assert_close(outs[64], outs[0], rtol=2e-2, atol=0.1)

From f62f8edf72d30e4533c1865f6d94a1466b34c0ce Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sun, 27 Sep 2026 03:33:34 -0400
Subject: [PATCH 158/249] merge fork perf: expert prefetch into low-priority
 slots with single-writer slot maps (FREETOKEN_MOE_PREFETCH), plus decode
 verify and slot audit debug modes

---
 benchmarks/README.md                       |   23 +
 benchmarks/bench_moe_prefetch.py           |  580 ++++++
 python/freetoken/engine/engine.py          |   72 +-
 python/freetoken/env.py                    |   14 +
 python/freetoken/layers/moe.py             |   48 +-
 python/freetoken/models/qwen4_exp/model.py |    3 +-
 python/freetoken/models/qwen4_exp/moe.py   |   17 +-
 python/freetoken/moe/offload_cache.py      |  233 ++-
 python/freetoken/moe/offload_kernels.py    |   44 +-
 python/freetoken/moe/prefetch.py           |  438 +++++
 python/freetoken/moe/scored_ensure.py      |  168 +-
 python/freetoken/moe/slot_audit.py         |  688 +++++++
 python/freetoken/moe/verify.py             |  516 ++++++
 tests/models/qwen4_exp/conftest.py         |   13 +
 tests/models/qwen4_exp/test_build_cpu.py   |   29 +
 tests/models/qwen4_exp/test_skeleton.py    | 1911 ++++++++++++++++++++
 tests/moe/ref_prefetch.py                  |   22 +
 tests/moe/ref_scored_cache.py              |   44 +-
 tests/moe/test_prefetch.py                 |  424 +++++
 tests/moe/test_scored_ensure.py            |  319 ++++
 20 files changed, 5541 insertions(+), 65 deletions(-)
 create mode 100644 benchmarks/bench_moe_prefetch.py
 create mode 100644 python/freetoken/moe/prefetch.py
 create mode 100644 python/freetoken/moe/slot_audit.py
 create mode 100644 python/freetoken/moe/verify.py
 create mode 100644 tests/moe/ref_prefetch.py
 create mode 100644 tests/moe/test_prefetch.py

diff --git a/benchmarks/README.md b/benchmarks/README.md
index 3f0df72f81..817f99da18 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -45,6 +45,29 @@ start after the ensure and how much of the side-stream shared expert runs inside
 python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160
 ```
 
+**`bench_moe_prefetch.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_PREFETCH=measure`
+cost. The router-lookahead predictor alone per layer (GEMV `[bs, 2560] x [2560, 512]`, top-K
+select, and the count kernel), then prefetch off vs measure on a real-geometry Qwen4ExpMoE decode
+stack (512 experts, rule eviction, copy overlap on) with forced routing that misses exactly m
+experts per layer, as CUDA graphs. `--on` times whole 48-layer decode steps with prefetch off vs
+on instead: a GEMV per layer stands in for attention (`--filler-us`), each layer prefetches 3 or
+4 candidates of which `--useful` are its misses, and `--delay-us` holds every prefetch copy back
+to price a late one. It reports ms/step, demand misses/step and the late fraction. `--steady
+--trace DIR` runs steady-state decode on a recorded routing trace instead (experts.npy,
+logit_idx.npy, logit_val.npy, req_index.npy): no cache restore, so wrong prefetches and evictions
+carry over, and a noisy copy of each layer's traced logits (`--sigma`) stands in for the
+lookahead. Arms such as `off,measure,on,on:4/4` (budgets before GDN / QSA layers) replay step by
+step in turn, each on its own copy of the cache state; it adds pollution (non-resident demands
+per layer against off's misses) and the late fraction per layer kind. Slots and sigma were
+calibrated to production's measure line (misses/layer ~4.3, precision ~0.56) at 24 layers.
+
+```bash
+python benchmarks/bench_moe_prefetch.py --bs 1,2 --misses 0,4
+python benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 --delay-us 0,50,150,300
+python benchmarks/bench_moe_prefetch.py --steady --trace DIR --bs 1,2 --steady-slots 740 --sigma 0.55 \
+    --arms off,measure,on,on:empty,on:3/3,on:4/4,on:4/5
+```
+
 **`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel
 alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps
 `fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim`
diff --git a/benchmarks/bench_moe_prefetch.py b/benchmarks/bench_moe_prefetch.py
new file mode 100644
index 0000000000..be5d38ebd7
--- /dev/null
+++ b/benchmarks/bench_moe_prefetch.py
@@ -0,0 +1,580 @@
+"""FREETOKEN_MOE_PREFETCH cost and gain: the router-lookahead predictor alone, its effect on the
+decode compute stream in measure mode, and (--on) whole decode steps with prefetch off vs on.
+
+1. Predictor alone, per layer: the lookahead GEMV ([bs, 2560] x [2560, 512] bf16, a distinct router
+   weight per layer as in the model), the top-K + select kernel, and both; one stream, CUDA graph.
+2. The count kernel alone (it runs on the compute stream in measure mode).
+3. A real-geometry decode stack (Qwen3.8-Flash-Next: hidden 2560, 512 experts, top-10, moe and
+   shared intermediate 640, NVFP4 triton banks) of N Qwen4ExpMoE blocks: router GEMV, shared
+   expert (on the FREETOKEN_MOE_COPY_OVERLAP side stream by default), top-k, rule ensure, miss
+   copy of exactly m experts from pinned host, grouped GEMM. The routing is forced (10 - m hot
+   experts + m cold ones per row) and the cache restored before every replay, so both arms copy
+   the same bytes. One CUDA graph per step, prefetch off vs measure. No attention runs, so the
+   percentages overstate the model's.
+4. --on: the same stack at the model's depth (48 layers), with a bf16 GEMV "filler" before every
+   block standing in for attention and the rest of the layer (--filler-us, ~225 us between one
+   demand copy and the next in production). Prefetch on installs a fixed candidate list per layer:
+   3 before GDN / 4 before attention layers, of which --useful are the layer's cold experts and the
+   rest never-routed ids (production: ~3.24 issued, ~1.75 useful, ~4.2 misses per layer at bs 1).
+   Reports ms/step, demand misses/step and the late fraction; --delay-us adds a sleep ahead of every
+   prefetch copy to price the compute stream's wait on a late copy.
+5. --steady --trace DIR: steady-state decode on a recorded routing trace (no restore, so wrong
+   prefetches and evictions carry over), --steady-layers layers of it with --steady-slots slots,
+   a GEMV filler of --filler-gdn-us / --filler-qsa-us before GDN / QSA layers, and a lookahead
+   that is each layer's traced top-32 logits plus N(0, --sigma) noise. Every arm (off, measure,
+   on[:gdn/qsa budgets | :empty]) keeps its own cache state and all arms replay step by step in
+   turn, so drift hits them alike. Reports ms/step (mean, p50, p90), demand misses per layer,
+   issued/useful/precision, non-resident demands per layer (misses + prefetch hits: above off's
+   misses is pollution) and the late fraction before GDN and QSA layers.
+
+Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_prefetch.py
+     ... benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4
+     ... benchmarks/bench_moe_prefetch.py --steady --trace DIR --steady-slots 740 --sigma 0.55
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import statistics
+
+import torch
+import torch.nn.functional as F
+
+H, E, TOP_K, INTER = 2560, 512, 10, 640
+
+
+def time_graph(run, reps: int, before=None) -> float:
+    """Median ms of one replay of ``run`` captured as a CUDA graph (after an eager warm-up)."""
+    if before:
+        before()
+    run()
+    torch.cuda.synchronize()
+    if before:
+        before()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    times = []
+    for rep in range(reps + 3):
+        if before:
+            before()
+        torch.cuda.synchronize()
+        start.record()
+        graph.replay()
+        end.record()
+        end.synchronize()
+        if rep >= 3:
+            times.append(start.elapsed_time(end))
+    return statistics.median(times)
+
+
+def bench_predictor(bs: int, layers: int, k: int, budget: int, reps: int) -> dict:
+    from freetoken.moe.prefetch import lookahead_select, prefetch_count
+
+    dev = torch.device("cuda")
+    gen = torch.Generator(device=dev).manual_seed(bs)
+    weights = [torch.randn(E, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.02 for _ in range(layers)]
+    x = torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen)
+    resident = torch.where(torch.rand(E, device=dev, generator=gen) < 0.4, 1, -1).to(torch.int32)
+    sel = torch.full((layers, 2 * k), -1, dtype=torch.int32, device=dev)
+    res = sel.clone()
+    logits = [F.linear(x, w) for w in weights]
+
+    def gemv():
+        for w in weights:
+            F.linear(x, w)
+
+    def select():
+        for layer in range(layers):
+            lookahead_select(logits[layer], resident, sel[layer], res[layer], k=k, budget=budget)
+
+    def both():
+        for layer, w in enumerate(weights):
+            lookahead_select(F.linear(x, w), resident, sel[layer], res[layer], k=k, budget=budget)
+
+    slots = torch.randint(0, 1024, (bs * TOP_K,), device=dev, dtype=torch.int32, generator=gen)
+    id_of_slot = torch.randint(0, E, (1024,), device=dev, dtype=torch.int32, generator=gen)
+    misses = torch.ones(1, dtype=torch.int64, device=dev)
+    counters = torch.zeros((layers, 6), dtype=torch.int64, device=dev)
+
+    def count():
+        for layer in range(layers):
+            prefetch_count(slots, id_of_slot, misses, sel[layer], res[layer], counters[layer], id_base=0, rows=bs)
+
+    per = 1000.0 / layers
+    return {name: time_graph(fn, reps) * per for name, fn in
+            (("gemv", gemv), ("select", select), ("predictor", both), ("count", count))}
+
+
+class ForcedGate:
+    """A layer's real router GEMV whose logits are then replaced, so every replay routes (and
+    misses) the same experts; the next layer's lookahead runs through it too."""
+
+    def __init__(self, gate, logits: torch.Tensor):
+        self.gate, self.logits = gate, logits
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.gate.forward(x)
+        return self.logits
+
+
+class DecodeStack:
+    """N Qwen4ExpMoE blocks (router, shared expert, nvfp4 offload experts) over one shared cache,
+    wired like qwen4_exp for the lookahead."""
+
+    def __init__(self, layers: int, slots: int, bs: int, misses: int, mode: str, overlap: bool, filler=None,
+                 routing=None, budgets=(3, 4)):
+        import freetoken.core as core
+        from types import SimpleNamespace
+
+        from freetoken.core import Context, set_global_ctx
+        from freetoken.distributed import set_tp_info, try_get_tp_info
+        from freetoken.layers.moe import OffloadMoELayer
+        from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend
+        from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE
+        from freetoken.moe.expert_banks import build_expert_banks
+        from freetoken.moe.offload_cache import OffloadMoeCache
+        from freetoken.utils.torch_utils import torch_dtype
+
+        if try_get_tp_info() is None:
+            set_tp_info(rank=0, size=1)
+        set_quant_backend(QuantBackend.parse("moe.nvfp4=triton"))
+        quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}})
+        dev = torch.device("cuda")
+        config = SimpleNamespace(
+            hidden_size=H, num_experts=E, num_experts_per_tok=TOP_K, moe_intermediate_size=INTER,
+            shared_expert_intermediate_size=INTER, norm_topk_prob=True, quant=None,
+            moe_strategy="offload", decode_target="gpu",
+        )
+        with torch.device(dev), torch_dtype(torch.bfloat16):
+            self.moes = [Qwen4ExpMoE(config, i, prefix=f"model.layers.{i}.mlp") for i in range(layers)]
+        gen = torch.Generator(device=dev).manual_seed(0)
+        for moe in self.moes:
+            for t in moe.state_dict().values():
+                if t.is_floating_point():
+                    t.normal_(0.0, 0.02, generator=gen)
+        experts = [
+            OffloadMoELayer(i, E, TOP_K, H, INTER, quant_config=quant, prefix=f"model.layers.{i}.mlp.experts")
+            for i in range(layers)
+        ]
+        method = experts[0].quant_method
+        if not hasattr(DecodeStack, "_banks"):
+            DecodeStack._banks = build_expert_banks(method, 1, None, device=dev, dummy=True)  # one layer, shared
+        banks = DecodeStack._banks
+        self.cache = cache = OffloadMoeCache(
+            num_layers=layers, num_experts=E, cache_size=slots, device=dev, cache_policy="rule",
+            quant_format=banks.quant_format, layout=banks.layout, max_slots=method.slot_limit(),
+            decode_copy_overlap=overlap, prefetch_mode=mode,
+        )
+        cache.set_bank_sources({role: per_layer * layers for role, per_layer in banks.sources.items()})
+        cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
+        core._GLOBAL_CTX = None
+        ctx = Context(page_size=1)
+        set_global_ctx(ctx)
+        ctx._batch = SimpleNamespace(is_prefill=False)
+        self.x = [torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(layers)]
+        self.filler = filler
+        if routing is not None:
+            # --steady: every layer routes this step's row of route_buf and layer L's lookahead (from
+            # layer L-1) returns look_buf[L]; both run the real router GEMV first
+            route_buf, look_buf = routing
+            gates = [moe.gate for moe in self.moes]
+            for i, (moe, ex) in enumerate(zip(self.moes, experts)):
+                ex.offload_cache = cache
+                moe.experts = ex
+                moe.gate = ForcedGate(gates[i], route_buf[i])
+            for i in range(1, layers):
+                experts[i - 1].set_lookahead(ForcedGate(gates[i], look_buf[i]), i, budgets[(i + 1) % 4 == 0])
+            self.bs = bs
+            cache.reset()
+            torch.cuda.synchronize()
+            return
+        # forced routing: row r routes hot experts [r*hot, (r+1)*hot) and m fresh ones past all rows' hot ones
+        hot = TOP_K - misses
+        for moe, ex in zip(self.moes, experts):
+            ex.offload_cache = cache
+            moe.experts = ex
+            lg = torch.zeros(bs, E, device=dev, dtype=torch.bfloat16)
+            for r in range(bs):
+                lg[r, list(range(r * hot, (r + 1) * hot))] = 10.0
+                lg[r, list(range(bs * hot + r * misses, bs * hot + (r + 1) * misses))] = 10.0
+            moe.gate = ForcedGate(moe.gate, lg)
+        for i in range(1, layers):
+            experts[i - 1].set_lookahead(self.moes[i].gate, i, budgets[(i + 1) % 4 == 0])
+        self.bs, self.hot, self.misses = bs, hot, misses
+        # warm the hot experts in, then drop the fresh ones: every replay then misses exactly those
+        cache.reset()
+        self.step()
+        cache.slot_for_id[:, bs * hot:].fill_(-1)
+        fresh = (cache.id_of_slot >= 0) & (cache.id_of_slot % E >= bs * hot)
+        cache.id_of_slot[fresh] = -1
+        cache.usage[fresh] = 0
+        torch.cuda.synchronize()
+        self.state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step]
+        for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner",
+                     "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"):
+            if getattr(cache, name) is not None:
+                self.state.append(getattr(cache, name))
+        self.snap = [t.clone() for t in self.state]
+
+    def restore(self):
+        for dst, src in zip(self.state, self.snap):
+            dst.copy_(src)
+
+    def step(self):
+        for i, (moe, x) in enumerate(zip(self.moes, self.x)):
+            filler = self.filler[i] if isinstance(self.filler, list) else self.filler
+            if filler is not None:
+                F.linear(x, filler)
+            moe.forward(x.clone())
+
+    def cold(self, layer: int) -> list[int]:
+        """The experts ``layer`` misses every replay, merged rank by rank across rows like the select."""
+        per_row = [list(range(self.bs * self.hot + r * self.misses, self.bs * self.hot + (r + 1) * self.misses))
+                   for r in range(self.bs)]
+        return [row[i] for i in range(self.misses) for row in per_row]
+
+
+def bench_stack(layers, slots, bs, misses, reps, overlap) -> dict:
+    from flashlib.kernels.slot_cache import Stat
+
+    out = {}
+    for mode in ("off", "measure"):
+        stack = DecodeStack(layers, slots, bs, misses, mode, overlap)
+        stack.cache.collect_stats = True
+        stack.restore()
+        stack.cache.lru_stats.zero_()
+        stack.step()
+        got = stack.cache.lru_stats[:, Stat.MISS].tolist()
+        assert got == [bs * misses] * layers, f"expected {bs * misses} misses per layer, got {got}"
+        stack.cache.collect_stats = False
+        out[mode] = time_graph(stack.step, reps, before=stack.restore) * 1000.0 / layers
+        if stack.cache.prefetch is not None:
+            stack.cache.prefetch.counters.zero_()
+            stack.restore()
+            stack.step()
+            c = stack.cache.prefetch.counters.sum(0).tolist()
+            out["issued/layer"] = c[0] / max(c[3], 1)
+        del stack
+        torch.cuda.empty_cache()
+    return out
+
+
+def budget_of(layer: int) -> int:
+    return 3 if (layer + 1) % 4 else 4
+
+
+def bench_on(layers, slots, bs, misses, useful, reps, filler_us, delays_us) -> list[dict]:
+    """ms per decode step, prefetch off vs on (and on with every prefetch copy delayed)."""
+    from flashlib.kernels.slot_cache import Stat
+
+    from freetoken.moe.prefetch import COPIED, ISSUED, LATE, USEFUL
+
+    dev = torch.device("cuda")
+    x = torch.randn(1, H, device=dev, dtype=torch.bfloat16)
+    probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16)
+    per_row_us = time_graph(lambda: F.linear(x, probe), reps) * 1000.0 / 32768
+    del probe
+    filler = torch.randn(max(int(filler_us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01
+    cycles_per_us = 1_000_000 / (time_graph(lambda: torch.cuda._sleep(1_000_000), 5) * 1000.0)
+    out = []
+    for arm in ["off", "on", *[f"on+{d}us" for d in delays_us if d]]:
+        stack = DecodeStack(layers, slots, bs, misses, "off" if arm == "off" else "on", True, filler=filler)
+        cache = stack.cache
+        want = [bs * misses] * layers
+        if arm != "off":
+            pf = cache.prefetch
+            pf.sel_override = torch.full_like(pf.sel, -1)
+            for layer in range(1, layers):
+                u = min(useful[layer % len(useful)], bs * misses)
+                wrong = list(range(E - 1, E - 1 - (budget_of(layer) - u), -1))
+                ids = stack.cold(layer)[:u] + wrong
+                pf.sel_override[layer, : len(ids)] = torch.tensor(ids, dtype=torch.int32)
+                want[layer] -= u
+            if arm.startswith("on+"):
+                pf.delay_copy_cycles = int(float(arm[3:-2]) * cycles_per_us)
+        cache.collect_stats = True
+        stack.restore()
+        cache.lru_stats.zero_()
+        stack.step()
+        got = cache.lru_stats[:, Stat.MISS].tolist()
+        assert got == want, f"{arm}: expected demand misses {want}, got {got}"
+        cache.collect_stats = False
+        stack.restore()
+        stack.step()
+        torch.cuda.synchronize()
+        stack.restore()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            stack.step()
+        if cache.prefetch is not None:
+            cache.prefetch.stats.zero_()
+        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        times = []
+        for rep in range(reps + 3):
+            stack.restore()
+            torch.cuda.synchronize()
+            start.record()
+            graph.replay()
+            end.record()
+            end.synchronize()
+            if rep >= 3:
+                times.append(start.elapsed_time(end))
+        row = {"arm": arm, "ms": statistics.median(times), "misses": sum(want)}
+        if cache.prefetch is not None:
+            st = cache.prefetch.stats.sum(0).tolist()
+            calls = max(reps + 3, 1) * (layers - 1)
+            row.update(issued=st[ISSUED] / calls, useful=st[USEFUL] / calls, late=st[LATE] / max(st[COPIED], 1))
+        out.append(row)
+        del graph, stack, cache
+        gc.collect()  # the module graph holds cycles: free this arm's slot banks before the next one
+        torch.cuda.empty_cache()
+    return out
+
+
+def load_trace(path: str, layers: int, offset: int, bs: int, steps: int, skip: int, sigma: float, seed: int):
+    """Per-step routing and lookahead logits ``[steps, layers, bs, E]`` from a routing trace.
+
+    ``path`` holds experts.npy ``[N, 48, 10]`` (routed ids), logit_idx.npy / logit_val.npy
+    ``[N, 48, 32]`` (each layer's own top-32 router logits) and req_index.npy ``[N]``. Row r of the
+    batch follows the r-th longest request from its ``skip``-th token on. The routing logits pick
+    exactly the traced top-10; the lookahead of layer L is layer L's traced top-32 logits plus
+    N(0, sigma) noise (everything else far below), a stand-in for W_gate[L] . x_{L-1}."""
+    import os
+    from collections import Counter
+
+    import numpy as np
+
+    ex = np.load(os.path.join(path, "experts.npy"), mmap_mode="r")
+    li = np.load(os.path.join(path, "logit_idx.npy"), mmap_mode="r")
+    lv = np.load(os.path.join(path, "logit_val.npy"), mmap_mode="r")
+    req = np.load(os.path.join(path, "req_index.npy"))
+    reqs = [r for r, n in Counter(req.tolist()).most_common(bs)]
+    rng = np.random.default_rng(seed)
+    route = np.full((steps, layers, bs, E), -10.0, np.float32)
+    look = np.full((steps, layers, bs, E), -30.0, np.float32)
+    lsl = slice(offset, offset + layers)
+    for r, rid in enumerate(reqs):
+        rows = np.flatnonzero(req == rid)[skip: skip + steps]
+        assert len(rows) == steps, f"request {rid} has {len(rows) + skip} tokens, fewer than skip + steps"
+        routed = np.asarray(ex[rows, lsl]).astype(np.int64)  # [steps, layers, 10]
+        np.put_along_axis(route[:, :, r], routed, 10.0 - 0.1 * np.arange(TOP_K), -1)
+        idx = np.asarray(li[rows, lsl]).astype(np.int64)
+        val = np.asarray(lv[rows, lsl]).astype(np.float32) + rng.normal(0.0, sigma, idx.shape).astype(np.float32)
+        np.put_along_axis(look[:, :, r], idx, val, -1)
+    dev = torch.device("cuda")
+    return (torch.from_numpy(route).to(dev, torch.bfloat16), torch.from_numpy(look).to(dev, torch.bfloat16), reqs)
+
+
+def bench_steady(args, bs: int, arms: list[str]) -> list[dict]:
+    """ms per decode step in steady state on traced routing: no restore, the cache evolves.
+
+    One stack, one graph per arm, each arm with its own copy of the cache's slot maps, eviction
+    state and counters: every step replays all arms back to back (rotating their order), so
+    clock or link drift hits every arm alike. Slot bytes are shared (timing does not read them)."""
+    from flashlib.kernels.slot_cache import Stat
+
+    from freetoken.moe.prefetch import COPIED, ISSUED, LATE, MISSES, USEFUL
+
+    dev = torch.device("cuda")
+    steps = args.warm + args.timed
+    route_all, look_all, reqs = load_trace(args.trace, args.steady_layers, args.layer_offset, bs, steps,
+                                           args.skip, args.sigma, args.seed)
+    layers = args.steady_layers
+    x = torch.randn(1, H, device=dev, dtype=torch.bfloat16)
+    probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16)
+    per_row_us = time_graph(lambda: F.linear(x, probe), args.reps) * 1000.0 / 32768
+    del probe
+    fill = {us: torch.randn(max(int(us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01
+            for us in (args.filler_gdn_us, args.filler_qsa_us)}
+    # the filler of layer i stands in for its attention: QSA layers are 3, 7, ... (i % 4 == 3)
+    fillers = [fill[args.filler_qsa_us if i % 4 == 3 else args.filler_gdn_us] for i in range(layers)]
+    qsa = torch.tensor([i % 4 == 3 for i in range(layers)])
+    route_buf = torch.empty_like(route_all[0])
+    look_buf = torch.empty_like(look_all[0])
+    stack = DecodeStack(layers, args.steady_slots, bs, 0, "on", True, filler=fillers, routing=(route_buf, look_buf))
+    cache, pf = stack.cache, stack.cache.prefetch
+    cache.collect_stats = True
+    experts = [moe.experts for moe in stack.moes]
+    state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step, cache.lru_stats, pf.stats]
+    for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner",
+                 "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"):
+        if getattr(cache, name) is not None:
+            state.append(getattr(cache, name))
+    init = [t.clone() for t in state]
+
+    def load(snap):
+        for dst, src in zip(state, snap):
+            dst.copy_(src)
+
+    def save(snap):
+        for dst, src in zip(snap, state):
+            dst.copy_(src)
+
+    route_buf.copy_(route_all[0])
+    look_buf.copy_(look_all[0])
+    graphs, snaps, keep = {}, {}, []
+    for arm in dict.fromkeys(arms):
+        mode, _, spec = arm.partition(":")
+        budgets = tuple(int(b) for b in spec.split("/")) if spec and spec != "empty" else (3, 4)
+        for i in range(1, layers):
+            experts[i - 1].set_lookahead(experts[i - 1]._lookahead[0], i, budgets[(i + 1) % 4 == 0])
+        # every mode switch is read at enqueue time, so each arm's graph bakes its own path in
+        pf.mode = "on" if mode == "off" else mode
+        pf.sel_override = torch.full_like(pf.sel, -1) if spec == "empty" else None
+        keep.append(pf.sel_override)
+        cache.prefetch = None if mode == "off" else pf
+        load(init)
+        stack.step()  # compile outside capture
+        torch.cuda.synchronize()
+        load(init)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            stack.step()
+        graphs[arm], snaps[arm] = graph, [t.clone() for t in init]
+    cache.prefetch, pf.mode, pf.sel_override = pf, "on", None
+    names = list(graphs)
+    events = {arm: [] for arm in names}
+    for t in range(steps):
+        if t == args.warm:
+            for arm in names:
+                snaps[arm][4].zero_()
+                snaps[arm][5].zero_()
+        route_buf.copy_(route_all[t])
+        look_buf.copy_(look_all[t])
+        for j in range(len(names)):
+            arm = names[(t + j) % len(names)]
+            load(snaps[arm])
+            ev = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True))
+            ev[0].record()
+            graphs[arm].replay()
+            ev[1].record()
+            save(snaps[arm])
+            if t >= args.warm:
+                events[arm].append(ev)
+        if t % 50 == 49:
+            torch.cuda.synchronize()  # bound the event backlog
+    torch.cuda.synchronize()
+    out = []
+    n = args.timed
+    for arm in names:
+        mode = arm.partition(":")[0]
+        times = sorted(a.elapsed_time(b) for a, b in events[arm])
+        miss = snaps[arm][4][:, Stat.MISS].double().cpu()
+        row = {"arm": arm, "mean": statistics.fmean(times), "p50": times[len(times) // 2],
+               "p90": times[int(len(times) * 0.9)], "miss": miss.sum().item() / (n * layers),
+               "miss1": miss[1:].sum().item() / (n * (layers - 1))}
+        if mode != "off":
+            st = snaps[arm][5].double().cpu()
+            calls = n * (layers - 1)
+            row.update(issued=st[:, ISSUED].sum().item() / calls, useful=st[:, USEFUL].sum().item() / calls)
+            row["precision"] = row["useful"] / row["issued"] if row["issued"] else 0.0
+            if mode == "measure":
+                row["coverage"] = st[:, USEFUL].sum().item() / max(st[:, MISSES].sum().item(), 1)
+            if mode == "on":
+                # every routed id that was not a normal resident: residual misses plus prefetch hits
+                row["nonres1"] = row["miss1"] + row["useful"]
+                for name, mask in (("late_gdn", ~qsa), ("late_qsa", qsa)):
+                    m = mask.clone()
+                    m[0] = False
+                    row[name] = st[m, LATE].sum().item() / max(st[m, COPIED].sum().item(), 1)
+                row["late"] = st[:, LATE].sum().item() / max(st[:, COPIED].sum().item(), 1)
+        out.append(row)
+    print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB, requests {reqs}", flush=True)
+    del graphs, stack, cache, pf, events
+    gc.collect()
+    torch.cuda.empty_cache()
+    return out
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--bs", default="1,2")
+    ap.add_argument("--k", type=int, default=16)
+    ap.add_argument("--budget", type=int, default=4)
+    ap.add_argument("--layers", type=int, default=8, help="MoE layers in the decode stack")
+    ap.add_argument("--slots", type=int, default=512, help="expert slots (2.77 MB each)")
+    ap.add_argument("--misses", default="0,4", help="demand misses per layer per row")
+    ap.add_argument("--reps", type=int, default=30)
+    ap.add_argument("--overlap", type=int, default=1, help="FREETOKEN_MOE_COPY_OVERLAP in the stack (production: 1)")
+    ap.add_argument("--on", action="store_true", help="time whole decode steps with prefetch off vs on instead")
+    ap.add_argument("--on-layers", type=int, default=48)
+    ap.add_argument("--on-slots", type=int, default=0, help="0 = every layer's routed experts + 64 (the warm-up must not evict)")
+    ap.add_argument("--filler-us", type=float, default=160.0, help="per-layer compute standing in for attention")
+    ap.add_argument("--useful", default="2,1,2,2", help="useful prefetches per layer, cycled over layers")
+    ap.add_argument("--delay-us", default="0,150,300", help="sleep ahead of every prefetch copy (0 = none)")
+    ap.add_argument("--steady", action="store_true", help="steady-state decode on a routing trace instead")
+    ap.add_argument("--trace", default=None, help="--steady: dir with experts/logit_idx/logit_val/req_index .npy")
+    ap.add_argument("--steady-layers", type=int, default=24, help="--steady: layers (VRAM: slots * 2.77 MB)")
+    ap.add_argument("--layer-offset", type=int, default=0, help="--steady: first traced layer")
+    ap.add_argument("--steady-slots", type=int, default=825, help="--steady: slots (1650 * 24 / 48 by default)")
+    ap.add_argument("--warm", type=int, default=300, help="--steady: untimed steps before timing")
+    ap.add_argument("--timed", type=int, default=500)
+    ap.add_argument("--skip", type=int, default=200, help="--steady: tokens skipped at each request's start")
+    ap.add_argument("--sigma", type=float, default=0.3, help="--steady: lookahead noise on the traced logits")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--filler-gdn-us", type=float, default=171.0, help="--steady: compute before a GDN layer's MoE")
+    ap.add_argument("--filler-qsa-us", type=float, default=212.0, help="--steady: compute before a QSA layer's MoE")
+    ap.add_argument("--arms", default="off,measure,on,on:empty,on:2/3,on:3/3,on:4/4",
+                    help="--steady: mode[:gdn/qsa budgets | :empty (on, nothing selected)]")
+    args = ap.parse_args()
+    from freetoken.env import ENV
+
+    ENV.MOE_PREFETCH_K.value, ENV.MOE_PREFETCH_BUDGET.value = args.k, 0 if args.steady else args.budget
+    dev = torch.device("cuda")
+    if args.steady:
+        print(f"# {torch.cuda.get_device_name(dev)}, steady state on {args.trace}: {args.steady_layers} layers from "
+              f"{args.layer_offset}, {args.steady_slots} slots, rule, copy overlap on, filler {args.filler_gdn_us:.0f}/"
+              f"{args.filler_qsa_us:.0f} us (GDN/QSA), sigma {args.sigma}, {args.warm} warm + {args.timed} timed steps")
+        print(f"{'bs':>3} {'arm':>10} {'mean ms':>8} {'vs off':>7} {'p50':>7} {'p90':>7} {'miss/L':>7} {'miss/L1+':>8} "
+              f"{'iss/L':>6} {'use/L':>6} {'prec':>6} {'nonres/L1+':>10} {'late':>6} {'lateGDN':>7} {'lateQSA':>7}")
+        for bs in (int(b) for b in args.bs.split(",")):
+            rows = bench_steady(args, bs, args.arms.split(","))
+            base = rows[0]["mean"]
+            for r in rows:
+                g = r.get
+                print(f"{bs:>3} {r['arm']:>10} {r['mean']:>8.3f} {r['mean'] / base - 1:>+7.1%} {r['p50']:>7.3f} "
+                      f"{r['p90']:>7.3f} {r['miss']:>7.3f} {r['miss1']:>8.3f} {g('issued', 0):>6.2f} {g('useful', 0):>6.2f} "
+                      f"{g('precision', 0):>6.3f} {g('nonres1', 0):>10.3f} {g('late', 0):>6.3f} {g('late_gdn', 0):>7.3f} "
+                      f"{g('late_qsa', 0):>7.3f}", flush=True)
+        return
+    if args.on:
+        useful = [int(u) for u in args.useful.split(",")]
+        delays = [float(d) for d in args.delay_us.split(",")]
+        print(f"# {torch.cuda.get_device_name(dev)}, {args.on_layers} Qwen4ExpMoE nvfp4 layers, "
+              f"{args.on_slots or 'layers * bs * top-k + 64'} slots, "
+              f"rule, copy overlap on, filler {args.filler_us:.0f} us/layer, useful {useful}, median of {args.reps} replays")
+        print(f"{'bs':>3} {'m':>3} {'arm':>10} {'ms/step':>9} {'vs off':>8} {'misses/step':>12} "
+              f"{'issued/layer':>13} {'useful/layer':>13} {'late':>6}")
+        for bs in (int(b) for b in args.bs.split(",")):
+            for m in (int(x) for x in args.misses.split(",")):
+                slots = args.on_slots or args.on_layers * bs * TOP_K + 64
+                rows = bench_on(args.on_layers, slots, bs, m, useful, args.reps, args.filler_us, delays)
+                base = rows[0]["ms"]
+                for r in rows:
+                    extra = (f"{r['issued']:>13.2f} {r['useful']:>13.2f} {r['late']:>6.3f}" if "late" in r
+                             else f"{'':>13} {'':>13} {'':>6}")
+                    print(f"{bs:>3} {m:>3} {r['arm']:>10} {r['ms']:>9.3f} {r['ms'] / base - 1:>+8.1%} "
+                          f"{r['misses']:>12d} {extra}", flush=True)
+        print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB")
+        return
+    print(f"# {torch.cuda.get_device_name(dev)}, K={args.k}, budget={args.budget}, median of {args.reps} graph replays")
+    print("# predictor alone, us per layer (47 layers, one stream)")
+    print(f"{'bs':>3} {'gemv':>8} {'select':>8} {'predictor':>10} {'count':>8}")
+    for bs in (int(b) for b in args.bs.split(",")):
+        r = bench_predictor(bs, 47, args.k, args.budget, args.reps)
+        print(f"{bs:>3} {r['gemv']:>8.2f} {r['select']:>8.2f} {r['predictor']:>10.2f} {r['count']:>8.2f}", flush=True)
+    print(f"# decode stack ({args.layers} Qwen4ExpMoE nvfp4 layers, {args.slots} slots, rule, overlap={args.overlap}), "
+          "us per layer, prefetch off vs measure")
+    print(f"{'bs':>3} {'m':>3} {'off':>9} {'measure':>9} {'delta':>8} {'issued/layer':>13}")
+    for bs in (int(b) for b in args.bs.split(",")):
+        for m in (int(x) for x in args.misses.split(",")):
+            r = bench_stack(args.layers, args.slots, bs, m, args.reps, bool(args.overlap))
+            print(f"{bs:>3} {m:>3} {r['off']:>9.2f} {r['measure']:>9.2f} {r['measure'] - r['off']:>+8.2f} "
+                  f"{r.get('issued/layer', 0.0):>13.2f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 5eb7ce3c53..518c55c1fe 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1320,6 +1320,9 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         use_graph = self.graph_runner.can_use_cuda_graph(batch)
         with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph):
             logits = self.graph_runner.replay(batch) if use_graph else self.model.forward()
+        audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None
+        if audit is not None and batch.is_decode:
+            audit.end_decode_step(self.moe_offload_cache)  # every side stream has joined the compute stream
         if self.cpu_moe_executor is not None:
             # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator
             # -> stale expert outputs) as a loud error instead of silent corruption.
@@ -1334,11 +1337,21 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args)
         copy_done_event = torch.cuda.Event()
         copy_done_event.record(self.stream)
-        if self.moe_offload_cache is not None and self.moe_offload_cache.collect_stats and batch.is_decode:
+        cache = self.moe_offload_cache
+        if cache is not None and batch.is_decode and (
+            cache.collect_stats or cache.prefetch is not None or cache.verify is not None or cache.audit is not None
+        ):
             self._moe_stats_step += 1
             if self._moe_stats_step >= MOE_STATS_INTERVAL:
                 self._moe_stats_step = 0
-                self._emit_moe_stats()
+                if cache.prefetch is not None:
+                    self._emit_prefetch_stats()
+                if cache.verify is not None:
+                    self._emit_verify_stats()
+                if cache.audit is not None:
+                    self._emit_audit_stats()
+                if cache.collect_stats:
+                    self._emit_moe_stats()
         if logprobs_out is None:
             return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
         chosen_logprobs, top_ids, top_logprobs = logprobs_out
@@ -1381,6 +1394,40 @@ def _warmup_prefill_lens(self) -> list[int]:
         if mc.has_linear_attention and cap >= 8193:
             lens.add(8193)  # chunk_delta_h NT bucket 2 (NT > 128)
         return sorted(n for n in lens if 2 <= n <= cap)
+
+    def _emit_prefetch_stats(self) -> None:
+        """Report one window of the FREETOKEN_MOE_PREFETCH counters (one host sync), then restart them."""
+        from freetoken.moe.prefetch import format_per_layer, format_summary, summarize
+
+        prefetch = self.moe_offload_cache.prefetch
+        window = prefetch.take_window()
+        stats = summarize(window, prefetch.mode)
+        if not stats["layer_calls"]:
+            return
+        logger.info_rank0(f"MoE prefetch {prefetch.mode} ({MOE_STATS_INTERVAL} decode steps): {format_summary(stats)}")
+        if ENV.MOE_PREFETCH_DEBUG:
+            logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(window)}")
+
+    def _emit_verify_stats(self) -> None:
+        """Report one window of the FREETOKEN_MOE_PREFETCH_VERIFY counts and the new bad records (host syncs)."""
+        _log_lines(self.moe_offload_cache.verify.report_window(MOE_STATS_INTERVAL))
+
+    def _emit_audit_stats(self) -> None:
+        """Report one window of FREETOKEN_MOE_SLOT_AUDIT totals and the new bad slots (host syncs)."""
+        _log_lines(self.moe_offload_cache.audit.report_window(MOE_STATS_INTERVAL))
+
+    def _log_prefetch_totals(self) -> None:
+        from freetoken.moe.prefetch import format_per_layer, format_summary, summarize
+
+        prefetch = self.moe_offload_cache.prefetch if self.moe_offload_cache is not None else None
+        if prefetch is None:
+            return
+        totals = prefetch.totals + prefetch.stats.cpu()
+        stats = summarize(totals, prefetch.mode)
+        if stats["layer_calls"]:
+            logger.info_rank0(f"MoE prefetch {prefetch.mode} (session, {stats['tokens']} tokens): {format_summary(stats)}")
+            logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(totals)}")
+
     def _emit_moe_stats(self) -> None:
         """Report one window of expert-cache behaviour, then reset the miss counters.
 
@@ -1503,11 +1550,32 @@ def _warmup_prefill(self) -> None:
         )
 
     def shutdown(self) -> None:
+        try:
+            self._log_prefetch_totals()
+        except Exception as exc:  # noqa: BLE001 -- a diagnostic must never block shutdown
+            logger.warning(f"MoE prefetch totals unavailable at shutdown: {exc}")
+        try:
+            verify = self.moe_offload_cache.verify if self.moe_offload_cache is not None else None
+            if verify is not None:
+                _log_lines(verify.report_session())
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(f"MoE verify totals unavailable at shutdown: {exc}")
+        try:
+            audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None
+            if audit is not None:
+                _log_lines(audit.report_session())
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(f"MoE slot audit totals unavailable at shutdown: {exc}")
         self.graph_runner.destroy_cuda_graphs()
         torch.distributed.destroy_process_group()
         destroy_distributed()
 
 
+def _log_lines(lines: list[tuple[str, str]]) -> None:
+    for level, line in lines:
+        (logger.warning_rank0 if level == "warning" else logger.info_rank0)(line)
+
+
 def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]:
     """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA."""
     if not torch.cuda.is_available():
diff --git a/python/freetoken/env.py b/python/freetoken/env.py
index 0e56633f48..9a06b1c029 100644
--- a/python/freetoken/env.py
+++ b/python/freetoken/env.py
@@ -79,6 +79,20 @@ class EnvClassSingleton:
     WARMUP_MAX_LEN = EnvInt(0)
     # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert
     MOE_COPY_OVERLAP = EnvBool(False)
+    # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy)
+    MOE_PREFETCH = EnvStr("off")
+    # router-lookahead candidates per layer, and the non-resident ones kept per target layer
+    # (0 = 3 before GDN layers, 4 before full-attention layers)
+    MOE_PREFETCH_K = EnvInt(16)
+    MOE_PREFETCH_BUDGET = EnvInt(0)
+    # measure mode: also log the per-layer breakdown with every summary
+    MOE_PREFETCH_DEBUG = EnvBool(False)
+    # debug: check every GPU decode layer call's slot maps (meta) and slot bytes before/after the GEMM (1)
+    MOE_PREFETCH_VERIFY = EnvStr("0")
+    # routed experts a call byte-checks (scratch rows, ~2.8 MB each on Qwen3.8-Flash)
+    MOE_PREFETCH_VERIFY_ROWS = EnvInt(20)
+    # debug: record every expert slot writer and audit every held slot each N decode steps (0 = off)
+    MOE_SLOT_AUDIT = EnvInt(0)
 
     def __new__(cls):
         # single instance
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 7200cb2f4e..19e6d8e77c 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -189,6 +189,13 @@ def __init__(
             prefix=prefix,
         )
         self.offload_cache: OffloadMoeCache | None = None
+        # FREETOKEN_MOE_PREFETCH: (next MoE layer's router op, its layer id, its budget); the
+        # underscore keeps the borrowed router out of this layer's state dict
+        self._lookahead: tuple[BaseOP, int, int] | None = None
+
+    def set_lookahead(self, gate: BaseOP, target_layer: int, budget: int) -> None:
+        """Let this layer's GPU decode predict ``target_layer``'s experts with that layer's router."""
+        self._lookahead = (gate, target_layer, budget)
 
     @property
     def skips_inactive_slots(self) -> bool:
@@ -300,9 +307,25 @@ def _decode_routed(
             return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids)
         if cache.decode_target == "hybrid":
             return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids)
+        prefetch, verify = cache.prefetch, cache.verify
+        if verify is not None:
+            verify.begin(cache, self.layer_id, topk_ids, prefetch and prefetch.plan_of(self.layer_id))
         cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits)
+        target = None
+        if prefetch is not None:
+            # on: install the next layer's prediction (forked at this block's start) on this stream, the slot
+            # maps' only writer; before this layer's copy join, so a late prefetch copy never delays it
+            target = prefetch.install(self._lookahead and self._lookahead[1])
+            prefetch.join_copy(self.layer_id)
         cache.copy_missing()
-        return self._expert_gemm(
+        if prefetch is not None:
+            prefetch.join_and_count(self.layer_id, topk_ids, cache.id_of_slot, cache.num_indices)
+            if target is not None:
+                prefetch.issue_copy(target)
+        if verify is not None:
+            # after issue_copy: the next layer's prefetch copy starts when it would without the checks
+            verify.before_gemm(cache, topk_ids, None if target is None else prefetch.plan_of(target))
+        out = self._expert_gemm(
             cache,
             hidden_states,
             topk_weights,
@@ -312,6 +335,29 @@ def _decode_routed(
             alphas=cache.alphas_for_slots(self.layer_id),
             is_prefill=False,
         )
+        if verify is not None:
+            verify.after_gemm(cache)
+        return out
+
+    def fork_lookahead(self, hidden_states: torch.Tensor) -> None:
+        """FREETOKEN_MOE_PREFETCH: at the start of this layer's MoE block, predict the next layer's
+        experts from this block's router input on the predictor stream, beside this layer's router
+        and ensure. Only on the GPU decode path that ``_decode_routed`` takes and not under TP, and
+        never toward a CPU-decoded layer, whose ensure would not use it."""
+        cache = self.offload_cache
+        if (
+            self._lookahead is None
+            or cache is None
+            or cache.prefetch is None
+            or self.tp_size > 1
+            or get_global_ctx().batch.is_prefill
+            or cache.is_cpu_layer(self.layer_id)
+            or cache.decode_target == "hybrid"
+        ):
+            return
+        gate, target, budget = self._lookahead
+        if not cache.is_cpu_layer(target):
+            cache.prefetch.fork(target, hidden_states, gate, cache.slot_for_id[target], budget)
 
     def _decode_hybrid(
         self,
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index c4fe8da1bb..fdef4298ba 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -26,7 +26,7 @@
 
 from .attention import Qwen4ExpAttention
 from .hc import GatedResidual
-from .moe import Qwen4ExpMoE
+from .moe import Qwen4ExpMoE, wire_router_lookahead
 from .ple import PLELayer
 from freetoken.models.blocks import embed_input_ids
 from freetoken.models.qwen3_vl.vision import Qwen3VLVisionModel, QwenVLVisionMixin
@@ -104,6 +104,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None:
                 for layer_id in range(config.num_layers)
             ]
         )
+        wire_router_lookahead([layer.mlp for layer in self.layers.op_list], config)
         self.hyper_connection_mixer = GatedResidual(config, use_combine=False, prefix=f"{prefix}.hyper_connection_mixer")
         # plain tuple (not an OP child), so it never shows up in the state dict
         self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None)
diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py
index 6f729231f1..4a1907d2ce 100644
--- a/python/freetoken/models/qwen4_exp/moe.py
+++ b/python/freetoken/models/qwen4_exp/moe.py
@@ -36,6 +36,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         se, ex = self.shared_expert, self.experts
+        if isinstance(ex, OffloadMoELayer):
+            # FREETOKEN_MOE_PREFETCH: the next layer's prediction needs only this block's input, so it
+            # starts here and runs beside this layer's router and ensure
+            ex.fork_lookahead(hidden_states)
         if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies():
             # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router,
             # ensure and miss copy stay back to back on this one; joined before the combine
@@ -65,4 +69,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim)
 
 
-__all__ = ["Qwen4ExpMoE"]
+def wire_router_lookahead(moes: list[Qwen4ExpMoE], config: ModelConfig) -> None:
+    """Give each offloaded MoE block the next layer's router, for FREETOKEN_MOE_PREFETCH's
+    lookahead (inert when it is off). ``moes[i]`` is decoder layer ``i``'s block."""
+    from freetoken.moe.prefetch import default_budget
+
+    for target in range(1, len(moes)):
+        experts = moes[target - 1].experts
+        if isinstance(experts, OffloadMoELayer):
+            experts.set_lookahead(moes[target].gate, target, default_budget(config.is_linear_layer(target)))
+
+
+__all__ = ["Qwen4ExpMoE", "wire_router_lookahead"]
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 956e34f2ea..510cd9aae9 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -2,6 +2,7 @@
 
 import math
 import os
+import weakref
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Iterator
@@ -175,6 +176,13 @@ class OffloadMoeCache:
     # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the
     # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP.
     decode_copy_overlap: bool | None = None
+    # GPU decode only: cross-layer expert prefetch (moe/prefetch.py). None = FREETOKEN_MOE_PREFETCH.
+    prefetch_mode: str | None = None
+    # GPU decode only: debug slot checks (moe/verify.py). None = FREETOKEN_MOE_PREFETCH_VERIFY.
+    verify_mode: str | None = None
+    # GPU decode only: debug slot writer history + audit every N decode steps (moe/slot_audit.py).
+    # None = FREETOKEN_MOE_SLOT_AUDIT, 0 = off.
+    slot_audit: int | None = None
 
     def __post_init__(self) -> None:
         from freetoken.moe.scored_ensure import POLICY_IDS
@@ -340,6 +348,9 @@ def __post_init__(self) -> None:
         # _stage_tails[g][b]: the pad after buffer b's rows in layout g
         self._stage_tails: list[list[list[torch.Tensor]]] = []
         self._inval: list[slice | torch.Tensor] = []
+        # the slot audit's label for the staged copy, and the most entries its plan can have
+        self._pending_kind = 0
+        self._pending_width: int | None = None
         self.prefill_copy_stream: torch.cuda.Stream | None = None
         self.prefill_begin_event: torch.cuda.Event | None = None
         self.prefill_ready_events: list[torch.cuda.Event] = []
@@ -368,6 +379,70 @@ def __post_init__(self) -> None:
         self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = []
         if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu":
             self._init_decode_copy_overlap()
+        self.prefetch = self._init_prefetch()
+        self.verify = self._init_verify()
+        self.audit = self._init_audit()
+
+    def _init_audit(self):
+        from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled
+        from freetoken.moe.slot_audit import SlotAuditor, resolve_interval
+
+        from_env = self.slot_audit is None
+        self.slot_audit = resolve_interval(self.slot_audit)
+        if not self.slot_audit or self.device.type != "cuda":
+            return None
+        if self.decode_target != "gpu":
+            if not from_env:
+                raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} needs GPU decode, not {self.decode_target!r}")
+            logger.warning(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} ignored: decode target is {self.decode_target!r}")
+            self.slot_audit = 0
+            return None
+        if _skip_fast_index_copy_enabled():
+            logger.warning("FREETOKEN_MOE_SLOT_AUDIT ignored: FREETOKEN_SKIP_FAST_INDEX_COPY makes every copy a no-op")
+            self.slot_audit = 0
+            return None
+        return SlotAuditor(self.num_layers, self.num_experts, self.device, interval=self.slot_audit)
+
+    def _init_verify(self):
+        from freetoken.moe.verify import ExpertVerifier, resolve_mode
+
+        from_env = self.verify_mode is None
+        self.verify_mode = resolve_mode(self.verify_mode)
+        if self.verify_mode == "off" or self.device.type != "cuda":
+            return None
+        if self.decode_target != "gpu":
+            if not from_env:
+                raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} needs GPU decode, not {self.decode_target!r}")
+            logger.warning(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} ignored: decode target is {self.decode_target!r}")
+            self.verify_mode = "off"
+            return None
+        return ExpertVerifier(self.num_layers, self.num_experts, self.device, mode=self.verify_mode)
+
+    def _init_prefetch(self):
+        from freetoken.moe.prefetch import ExpertPrefetcher, resolve_mode
+
+        from_env = self.prefetch_mode is None
+        self.prefetch_mode = resolve_mode(self.prefetch_mode)
+        if self.prefetch_mode == "off":
+            return None
+        if self.decode_target != "gpu":
+            if not from_env:
+                raise ValueError(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} needs GPU decode, not {self.decode_target!r}")
+            # the process-wide flag skips hybrid/CPU-decoded caches, like FREETOKEN_MOE_COPY_OVERLAP
+            logger.warning(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} ignored: decode target is {self.decode_target!r}")
+            self.prefetch_mode = "off"
+            return None
+        if self.device.type != "cuda":
+            return None
+        prefetch = ExpertPrefetcher(self.num_layers, self.num_experts, self.device, mode=self.prefetch_mode)
+        # a proxy, not a cycle: dropping the cache must free its slot banks without waiting for gc
+        prefetch.cache = weakref.proxy(self)
+        return prefetch
+
+    @property
+    def prefetch_on(self) -> bool:
+        """FREETOKEN_MOE_PREFETCH=on: every ensure keys prefetched slots low priority."""
+        return self.prefetch is not None and self.prefetch.mode == "on"
 
     def _pool_layout(self, pool: int) -> int:
         return self.layer_layout[self.pool_of_layer.index(pool)]
@@ -651,6 +726,10 @@ def set_bank_sources(
             self.prefill_overlap = False
         self._slice_pools()
         self._build_copy_plan()
+        if self.verify is not None:
+            self.verify.bind(self)
+        if self.audit is not None:
+            self.audit.bind(self)
         if self.prefill_overlap:
             self._init_prefill_overlap_buffers()
 
@@ -818,6 +897,8 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device)
         self._slice_pools()
         self._build_copy_plan()  # slot caches were reallocated -> refresh fused-copy addrs
+        if self.verify is not None:
+            self.verify.bind(self)
         plan_slots = max(self.num_experts, cache_size)
         self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
         self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
@@ -839,6 +920,10 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         self.stat_fetched_layer.zero_()
         self.stat_steps_layer.zero_()
         self.decode_freq.zero_()
+        if self.prefetch is not None:
+            self.prefetch.reset()
+        if self.audit is not None:
+            self.audit.bind(self)  # new slot arrays: a new history
         self.prefill_hit_rows = 0
         self.prefill_total_rows = 0
         self._hit_d2d_fallback_logged = False  # geometry changed; re-log if still unusable
@@ -943,6 +1028,10 @@ def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int) -> None:
         """Take buffer ``buffer_id``'s slots for staging ``layer_id``."""
         slots = self._inval[buffer_id]
         if isinstance(slots, slice):
+            if self.audit is not None:
+                from freetoken.moe.slot_audit import INVALIDATE
+
+                self.audit.record_clear(self, INVALIDATE, layer_id, slots.start, slots.stop - slots.start)
             # One fixed-shape launch. The boolean-mask index
             # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every
             # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48
@@ -1009,6 +1098,10 @@ def copy() -> None:
             self._invalidate_prefill_buffer(buffer_id, layer_id)
             for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)):
                 view.copy_(per_layer[layer_id], non_blocking=True)
+            if self.audit is not None:
+                from freetoken.moe.slot_audit import PREFILL_BUFFER
+
+                self.audit.record_range(self, PREFILL_BUFFER, layer_id, buffer_id * self.num_experts, self.num_experts, meta=False)
 
         if self._prefill_hit_d2d_active:
             self._prefetch_split(layer_id, buffer_id)
@@ -1110,6 +1203,11 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
                 self._prefill_hit_num,
                 blocks_per_bank=64,
             )
+            if self.audit is not None:
+                from freetoken.moe.slot_audit import PREFILL_HIT_D2D
+
+                self.audit.record_plan(self, PREFILL_HIT_D2D, layer_id, self._prefill_hit_dst, None, self._prefill_hit_num,
+                                       meta=False, aux=self._prefill_hit_src, row_base=buffer_id * E)
         miss = np.nonzero(~hit_mask)[0]
         with torch.cuda.stream(self.prefill_copy_stream):
             if self._prefill_buffer_has_release_event[buffer_id]:
@@ -1142,6 +1240,11 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None:
                     torch.tensor(nbytes, dtype=torch.int64),
                     torch.cuda.current_stream(self.device).cuda_stream,
                 )
+                if self.audit is not None:
+                    from freetoken.moe.slot_audit import PREFILL_SPLIT_H2D
+
+                    # every row gets at least its small banks from the host, the misses all of them
+                    self.audit.record_range(self, PREFILL_SPLIT_H2D, layer_id, buffer_id * E, E, meta=False)
             self.prefill_ready_events[buffer_id].record(self.prefill_copy_stream)
 
     def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]:
@@ -1200,18 +1303,49 @@ def ensure_experts(
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        if self.cache_policy_id == 0:
+        audit = self.audit
+        if audit is not None and update_state and layer_id == self.evict_clock_layer():
+            audit.begin_step()
+        lowpri = self.prefetch_on
+        width = expert_ids.numel()
+        if self.cache_policy_id == 0 and not lowpri:
             ensure_experts(self, layer_id, expert_ids)
-            return
-        ensure_experts_scored(
-            self,
-            layer_id,
-            expert_ids,
-            bump_tok=update_state and layer_id == self.evict_clock_layer(),
-            update_state=update_state,
-            router_logits=router_logits,
-            pin_since=pin_since,
-        )
+        else:
+            # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots)
+            ensure_experts_scored(
+                self,
+                layer_id,
+                expert_ids,
+                bump_tok=update_state and layer_id == self.evict_clock_layer(),
+                update_state=update_state,
+                router_logits=router_logits,
+                pin_since=pin_since,
+                lowpri=lowpri,
+                pf_count=self.prefetch.count_args(layer_id) if lowpri and update_state else None,
+            )
+        if audit is not None:
+            from freetoken.moe import slot_audit as sa
+
+            kind = sa.DEMAND_INSTALL if update_state else sa.PREFILL_INSTALL
+            self._pending_kind = sa.DEMAND_COPY if update_state else sa.PREFILL_COPY
+            self._pending_width = width
+            audit.record_plan(self, kind, layer_id, self.evict_slots, self.src_indices, self.num_indices, meta=True, width=width)
+
+    def prefetch_ensure(
+        self, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor,
+        num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None,
+    ) -> None:
+        """Install the first ``budget`` of ``query``'s non-resident experts of ``layer_id`` as
+        low-priority slots and plan their copy into ``dst_slots``/``src_rows``/``num`` (never the
+        demand plan, which copy(L-1) still reads). Runs on the compute stream like every other slot
+        map writer."""
+        from freetoken.moe.offload_kernels import prefetch_ensure_experts
+
+        prefetch_ensure_experts(self, layer_id, query, dst_slots, src_rows, num, stats_row, ready, budget=budget)
+        if self.audit is not None:
+            from freetoken.moe.slot_audit import PREFETCH_INSTALL
+
+            self.audit.record_plan(self, PREFETCH_INSTALL, layer_id, dst_slots, src_rows, num, meta=True)
 
     def ensure_experts_hybrid(
         self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False
@@ -1241,18 +1375,36 @@ def materialize_layer(self, layer_id: int) -> None:
 
         self._pending_src_layer = layer_id
         self._pending_whole_layer = True
+        audit = self.audit
+        if audit is not None:
+            from freetoken.moe import slot_audit as sa
+
+            # the kernel also frees this layer's slots above E; record them while they still name it
+            audit.record_clear(self, sa.MATERIALIZE_CLEAR, layer_id, 0, self.cache_size, layer_above_e=True)
         materialize_layer(self, layer_id)
         if self.evict_slot_owner is not None:
             self.evict_slot_owner[: self.num_experts].fill_(-1)  # it installs the layer into these slots
+        if audit is not None:
+            self._pending_kind = sa.MATERIALIZE_COPY
+            self._pending_width = self.num_experts
+            audit.record_plan(self, sa.MATERIALIZE_INSTALL, layer_id, self.evict_slots, self.src_indices, self.num_indices,
+                              meta=True, width=self.num_experts)
 
     def reset(self) -> None:
         from freetoken.moe.offload_kernels import reset_cache
 
+        if self.audit is not None:
+            from freetoken.moe.slot_audit import RESET
+
+            self.audit.record_clear(self, RESET, -1, 0, self.cache_size)
         reset_cache(self)
         # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe
         # it here so a new sequence starts with cold hybrid fetch priorities.
         self.expert_recency.fill_(-1)
         self.reset_evict_state()
+        if self.prefetch is not None:
+            # the graph runner resets after each capture: drop the capture-time counts too
+            self.prefetch.reset_counters()
 
     def reset_stats(self) -> None:
         self.prefill_hit_rows = 0
@@ -1377,26 +1529,49 @@ def copy_missing(self, layer_id: int | None = None) -> None:
             # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor)
             for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)):
                 view.copy_(per_layer[layer_id])
+            if self.audit is not None:
+                from freetoken.moe.slot_audit import PAGEABLE_COPY
+
+                self.audit.record_range(self, PAGEABLE_COPY, layer_id, 0, self.num_experts, meta=False)
             return
         # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool
-        whole = self._pending_whole_layer
+        self.copy_rows(layer_id, self.evict_slots, self.src_indices, self.num_indices, slim=_SLIM_COPY,
+                       kind=self._pending_kind, width=self._pending_width, whole=self._pending_whole_layer)
+
+    def copy_rows(
+        self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool,
+        kind: int = 0, width: int | None = None, whole: bool = False,
+    ) -> None:
+        """Copy host rows ``src_rows[:num]`` of ``layer_id``'s banks into slots ``dst_slots[:num]``
+        on the current stream (the demand plan, or a prefetch plan on the prefetch copy stream);
+        ``kind`` labels the copy for the slot audit and ``width`` bounds ``num`` for it. ``dst_slots``
+        are local to the layer's pool, or rows of staging buffer 0 when ``whole`` (materialize)."""
+        self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim, whole=whole)
+        if self.audit is not None:
+            from freetoken.moe.slot_audit import OTHER_COPY
+
+            self.audit.record_plan(self, kind or OTHER_COPY, layer_id, dst_slots, src_rows, num, meta=False, width=width)
+
+    def _copy_rows(
+        self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool,
+        whole: bool = False,
+    ) -> None:
         if self._copy_fused_ok:
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit
 
             # One launch copies the missing rows for every bank (instead of one launch per
-            # bank). evict_slots/src_indices/num_indices are shared across banks;
-            # src_indices holds layer-local expert rows, resolved against this layer's
-            # source pointers (layer_id is a static int per captured graph node).
-            copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit
-            copy(
-                self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]],
-                self._copy_src_ptrs[layer_id],
-                self._copy_feat_bytes[self.layer_layout[layer_id]],
-                self.evict_slots,
-                self.src_indices,
-                self.num_indices,
-                blocks_per_bank=self._copy_blocks_per_bank,
-            )
+            # bank). The slot/row plan is shared across banks; src_rows holds layer-local
+            # expert rows, resolved against this layer's source pointers (layer_id is a
+            # static int per captured graph node).
+            dst_ptrs = self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]]
+            feat_bytes = self._copy_feat_bytes[self.layer_layout[layer_id]]
+            if slim:
+                fast_index_copy_multi_slim_jit(dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num)
+            else:
+                fast_index_copy_multi_jit(
+                    dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num,
+                    blocks_per_bank=self._copy_blocks_per_bank,
+                )
             return
 
         from freetoken.kernel import fast_index_copy_jit
@@ -1405,13 +1580,7 @@ def copy_missing(self, layer_id: int | None = None) -> None:
         for (per_layer, _), cache in zip(self.banks, views):
             if cache.numel() == 0:
                 continue
-            fast_index_copy_jit(
-                cache,
-                self.evict_slots,
-                per_layer[layer_id],
-                self.src_indices,
-                self.num_indices,
-            )
+            fast_index_copy_jit(cache, dst_slots, per_layer[layer_id], src_rows, num)
 
     @contextmanager
     def decode_side(self, layer_id: int) -> Iterator[None]:
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index 3523a9850b..93a7909eb4 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -59,10 +59,15 @@ def ensure_experts_scored(
     update_state: bool,
     router_logits: torch.Tensor | None = None,
     pin_since: torch.Tensor | None = None,
+    lowpri: bool = False,
+    pf_count: tuple | None = None,
 ) -> None:
-    """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule)."""
+    """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule), and for
+    every policy under FREETOKEN_MOE_PREFETCH=on (``lowpri``; ``pf_count`` is the prefetcher's
+    ``(stats_row, ready, num, rows)`` when this layer was prefetched this step)."""
     from freetoken.moe.scored_ensure import scored_ensure
 
+    stats_row, ready, num, rows = pf_count if pf_count is not None else (None, None, None, 0)
     scored_ensure(
         expert_ids,
         cache.slot_for_id.view(-1),
@@ -73,6 +78,23 @@ def ensure_experts_scored(
         cache.src_indices,
         cache.evict_slots,
         cache.num_indices,
+        router_logits=router_logits if cache.near_miss_thr is not None else None,
+        near_miss_thr=cache.near_miss_thr or 0.0,
+        stats=cache.lru_stats[layer_id] if cache.collect_stats else None,
+        bump_tok=bump_tok,
+        update_state=update_state,
+        pin_since=pin_since,
+        lowpri=lowpri,
+        pf_stats=stats_row,
+        pf_ready=ready,
+        pf_num=num,
+        pf_rows=rows,
+        **_evict_state(cache, layer_id),
+    )
+
+
+def _evict_state(cache, layer_id: int) -> dict:
+    return dict(
         policy=cache.cache_policy_id,
         num_layers=cache.num_layers,
         num_experts=cache.num_experts,
@@ -85,13 +107,21 @@ def ensure_experts_scored(
         slot_last_tok=cache.evict_slot_last_tok,
         slot_lc=cache.evict_slot_lc,
         slot_ct=cache.evict_slot_ct,
-        router_logits=router_logits if cache.near_miss_thr is not None else None,
-        near_miss_thr=cache.near_miss_thr or 0.0,
-        stats=cache.lru_stats[layer_id] if cache.collect_stats else None,
         id_base=layer_id * cache.num_experts,
-        bump_tok=bump_tok,
-        update_state=update_state,
-        pin_since=pin_since,
+    )
+
+
+def prefetch_ensure_experts(
+    cache, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor,
+    num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None,
+) -> None:
+    """FREETOKEN_MOE_PREFETCH=on: install the first ``budget`` of ``query``'s non-resident experts
+    of ``layer_id`` as low-priority slots and plan their copy (scored_ensure.prefetch_ensure)."""
+    from freetoken.moe.scored_ensure import prefetch_ensure
+
+    prefetch_ensure(
+        query, cache.slot_for_id.view(-1), cache.id_of_slot, cache.usage, cache.step,
+        dst_slots, src_rows, num, stats_row, ready, budget=budget, **_evict_state(cache, layer_id),
     )
 
 
diff --git a/python/freetoken/moe/prefetch.py b/python/freetoken/moe/prefetch.py
new file mode 100644
index 0000000000..01129d3997
--- /dev/null
+++ b/python/freetoken/moe/prefetch.py
@@ -0,0 +1,438 @@
+"""Cross-layer expert prefetch for GPU decode (``FREETOKEN_MOE_PREFETCH``).
+
+A training-free router lookahead predicts layer L's experts from the tensor layer L-1's router
+consumed (``W_gate[L] . x_{L-1}``, top-K) and keeps the first BUDGET candidates that are not
+resident. At bs 2 both rows' top-K lists merge rank by rank (row 0's first, then row 1's
+first, ...), duplicates dropped, under one budget.
+
+``measure``: after layer L's ensure, count how many of them layer L actually routed. Nothing is
+copied and the cache is never touched. The predictor reads ``slot_for_id[L]`` without ordering
+against the ensures; a late read only skews the counters, never the model.
+
+The prediction forks at the start of layer L-1's MoE block, where x_{L-1} first exists, so it runs
+beside layer L-1's own router and ensure:
+
+    compute: -fork(L)-> router_{L-1} ensure(L-1) copy(L-1) GEMM(L-1) ... ensure(L) copy(L) -join-> count(L) GEMM(L)
+    predict:       \\-> router_L(x_{L-1}) -> top-K -> select(L) -rec sel(L) -/
+
+``on``: the predictor only merges the candidates (it reads no cache state); the compute stream
+installs the first BUDGET of them that are not resident as low-priority slots (prefetch_ensure:
+held, usage 0, keyed between the empty slots and every resident) and a third stream copies them
+from layer L's host banks with the slim copy kernel:
+
+    compute: -fork(L)-> ... ensure(L-1) [wait sel(L)] prefetch_ensure(L) [wait pfcopy(L-1)] copy(L-1) -rec dcopy(L)-> GEMM(L-1) ... ensure(L) [wait pfcopy(L)] copy(L) GEMM(L)
+    predict:       \\-> router_L -> top-K -> merge(L) -rec sel(L)
+    copy:                                                                                   [wait dcopy(L)] slim copy(L) -> ready(L) -rec pfcopy(L)
+
+Every writer of the slot maps and their mirrors runs on the compute stream, so the maps have one
+writer in stream order and no cross-stream edge guards them, and the installs follow layer L's
+residency right after ensure(L-1) however late the predictor runs. dcopy(L) follows
+prefetch_ensure(L) and the demand copy(L-1) on the compute stream, so the copy stream reads a
+finished plan and owns the host link second. prefetch_ensure(L) keeps usage == step, i.e. layer
+L-1's routed slots, pinned while GEMM(L-1) may read them; the compute stream joins the prefetch copy
+before its own copy(L), so ensure(L) may reuse a prefetched slot as a demand victim and GEMM(L) only
+ever reads landed bytes. A demand hit on a prefetched slot sets its usage (a normal resident) and
+counts as useful. Layer 0 is never predicted.
+"""
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.env import ENV
+from freetoken.moe.slot_audit import PREFETCH_COPY
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+MODES = ("off", "measure", "on")
+# non-resident candidates per target layer: GDN layers leave a shorter copy window than QSA ones
+GDN_BUDGET = 3
+ATTN_BUDGET = 4
+# counters[layer, col]; on mode adds LATE and COPIED (calls whose plan had a copy) in a wider stats row
+ISSUED, USEFUL, RESIDENT_HITS, CALLS, ROWS, MISSES, LATE, COPIED = range(8)
+NUM_COLS = 6
+NUM_STAT_COLS = 8
+# the select kernel's dedup tile is (rows * k)^2; larger (eager-only) decode batches skip the prediction
+MAX_ROWS = 8
+
+
+def resolve_mode(mode: str | None = None) -> str:
+    mode = (ENV.MOE_PREFETCH.value if mode is None else mode).strip().lower()
+    if mode not in MODES:
+        raise ValueError(f"FREETOKEN_MOE_PREFETCH={mode!r}: expected one of {', '.join(MODES)}")
+    return mode
+
+
+def default_budget(before_linear: bool) -> int:
+    """Non-resident candidates kept for a target layer, by the target's attention kind."""
+    return GDN_BUDGET if before_linear else ATTN_BUDGET
+
+
+def dedicated_stream(device: torch.device, *, highest_priority: bool = False) -> torch.cuda.ExternalStream:
+    # not one of torch's 32 pooled streams: a pooled one can alias the engine or capture stream
+    from cuda.bindings import runtime as cudart
+
+    with torch.cuda.device(device):
+        if highest_priority:
+            err, _, greatest = cudart.cudaDeviceGetStreamPriorityRange()
+            if err == cudart.cudaError_t.cudaSuccess:
+                err, handle = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, greatest)
+        else:
+            err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking)
+    if err != cudart.cudaError_t.cudaSuccess:
+        raise RuntimeError(f"creating a dedicated CUDA stream failed: {err}")
+    return torch.cuda.ExternalStream(int(handle), device=device)
+
+
+@triton.jit(do_not_specialize=["stride", "budget"])
+def _lookahead_select_kernel(
+    logits_ptr, stride, resident_ptr, sel_ptr, res_ptr, budget,
+    ROWS: tl.constexpr, ROWS_P: tl.constexpr, E: tl.constexpr, E_P: tl.constexpr,
+    K: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr,
+):
+    b = tl.arange(0, ROWS_P)
+    e = tl.arange(0, E_P)
+    live = (b < ROWS)[:, None] & (e < E)[None, :]
+    v = tl.load(logits_ptr + b[:, None] * stride + e[None, :], mask=live, other=float("-inf")).to(tl.float32)
+    # candidates in merged order: position r * ROWS + b holds row b's rank-r id
+    j = tl.arange(0, N_P)
+    cand = tl.full([N_P], -1, tl.int32)
+    for r in tl.static_range(K):
+        top = tl.argmax(v, axis=1).to(tl.int32)  # ties go to the lowest id
+        at = (j[:, None] == r * ROWS + b[None, :]) & (b < ROWS)[None, :]
+        cand = tl.where(tl.sum(at.to(tl.int32), axis=1) > 0, tl.sum(tl.where(at, top[None, :], 0), axis=1), cand)
+        v = tl.where(e[None, :] == top[:, None], float("-inf"), v)
+    valid = j < ROWS * K
+    earlier = (cand[:, None] == cand[None, :]) & (j[None, :] < j[:, None]) & valid[None, :]
+    first = valid & (tl.sum(earlier.to(tl.int32), axis=1) == 0)
+    held = tl.load(resident_ptr + cand, mask=first, other=-1) >= 0
+    pick = first & (~held)
+    rank = tl.cumsum(pick.to(tl.int32), axis=0) - 1
+    kept = first & held & (tl.cumsum(first.to(tl.int32), axis=0) <= budget)
+    krank = tl.cumsum(kept.to(tl.int32), axis=0) - 1
+    p = tl.arange(0, W_P)
+    m_sel = pick[None, :] & (rank[None, :] == p[:, None]) & (p < budget)[:, None]
+    m_res = kept[None, :] & (krank[None, :] == p[:, None])
+    # at most one candidate matches each output position
+    sel = tl.sum(tl.where(m_sel, cand[None, :] + 1, 0), axis=1) - 1
+    res = tl.sum(tl.where(m_res, cand[None, :] + 1, 0), axis=1) - 1
+    tl.store(sel_ptr + p, sel, mask=p < W)
+    tl.store(res_ptr + p, res, mask=p < W)
+
+
+@triton.jit(do_not_specialize=["id_base"])
+def _prefetch_count_kernel(
+    slots_ptr, id_of_slot_ptr, misses_ptr, sel_ptr, res_ptr, counters_ptr, id_base,
+    N: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, ROWS: tl.constexpr,
+):
+    i = tl.arange(0, N_P)
+    slot = tl.load(slots_ptr + i, mask=i < N, other=-1)
+    # ensure rewrote the routed ids to slots, all pinned to this layer's ids until the next ensure
+    routed = tl.where(slot >= 0, tl.load(id_of_slot_ptr + slot, mask=slot >= 0, other=0) - id_base, -1)
+    p = tl.arange(0, W_P)
+    sel = tl.load(sel_ptr + p, mask=p < W, other=-1)
+    res = tl.load(res_ptr + p, mask=p < W, other=-1)
+    sel_hit = tl.sum(((sel[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0
+    res_hit = tl.sum(((res[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0
+    issued = tl.sum((sel >= 0).to(tl.int32))
+    useful = tl.sum((sel_hit & (sel >= 0)).to(tl.int32))
+    res_hits = tl.sum((res_hit & (res >= 0)).to(tl.int32))
+    misses = tl.load(misses_ptr)
+    c = tl.arange(0, 8)
+    add = tl.where(c == 0, issued, tl.where(c == 1, useful, tl.where(c == 2, res_hits, tl.where(c == 3, 1, ROWS))))
+    add = tl.where(c == 5, misses, add.to(tl.int64))
+    # one writer per layer row, stream-ordered: no atomics needed
+    old = tl.load(counters_ptr + c, mask=c < 6, other=0)
+    tl.store(counters_ptr + c, old + add, mask=c < 6)
+
+
+def lookahead_select(
+    logits: torch.Tensor, resident: torch.Tensor, sel: torch.Tensor, res: torch.Tensor, *, k: int, budget: int
+) -> None:
+    """Top-``k`` of each ``logits`` row, merged rank by rank; ``sel`` gets the first ``budget``
+    distinct ids with ``resident[id] < 0`` and ``res`` the resident ids among the first ``budget``
+    distinct ones, both in merge order and padded with -1. One CTA, no host sync."""
+    rows, num_experts = logits.shape
+    width = sel.numel()
+    assert logits.stride(1) == 1 and resident.numel() == num_experts and resident.dtype == torch.int32
+    assert sel.dtype == res.dtype == torch.int32 and res.numel() == width and 1 <= k <= num_experts
+    _lookahead_select_kernel[(1,)](
+        logits, logits.stride(0), resident, sel, res, min(budget, width),
+        ROWS=rows, ROWS_P=triton.next_power_of_2(rows),
+        E=num_experts, E_P=triton.next_power_of_2(num_experts),
+        K=k, N_P=triton.next_power_of_2(rows * k),
+        W=width, W_P=triton.next_power_of_2(width),
+        num_warps=4,
+    )
+
+
+def prefetch_count(
+    slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor,
+    sel: torch.Tensor, res: torch.Tensor, counters: torch.Tensor, *, id_base: int, rows: int,
+) -> None:
+    """Add one call's issued / useful / resident-hit / call / row / miss counts to ``counters``."""
+    assert slots.dtype == torch.int32 and slots.is_contiguous() and counters.dtype == torch.int64
+    n, width = slots.numel(), sel.numel()
+    _prefetch_count_kernel[(1,)](
+        slots, id_of_slot, misses, sel, res, counters, id_base,
+        N=n, N_P=triton.next_power_of_2(n), W=width, W_P=triton.next_power_of_2(width), ROWS=rows,
+        num_warps=1,
+    )
+
+
+@triton.jit
+def _mark_ready_kernel(ready_ptr):
+    tl.store(ready_ptr, 1)
+
+
+class ExpertPrefetcher:
+    """Predictor stream, per-layer buffers and device counters for the router lookahead.
+
+    ``fork`` runs layer L's prediction on the dedicated stream at the start of layer L-1's MoE
+    block. ``measure``: ``join_and_count`` waits for it after layer L's copy is enqueued and counts
+    on the compute stream. ``on``: ``install`` (right after layer L-1's ensure) waits for it and
+    runs prefetch_ensure(L) on the compute stream; ``issue_copy`` puts its copy on a second
+    dedicated stream after layer L-1's demand copy; the compute stream joins the copy before its
+    copy(L) (``join_copy``) and ensure(L) counts the call (``count_args``). The events are created
+    up front, the shapes are fixed and nothing syncs the host, so every side captures into the
+    decode graph. Nothing here is cache_size-shaped: a rebuild keeps it and only resets the counters.
+    """
+
+    def __init__(
+        self, num_layers: int, num_experts: int, device: torch.device, *,
+        mode: str = "measure", k: int | None = None, budget: int | None = None,
+    ) -> None:
+        self.mode = mode
+        self.num_layers = num_layers
+        self.num_experts = num_experts
+        self.device = device
+        k = ENV.MOE_PREFETCH_K.value if k is None else int(k)
+        if k < 1:
+            raise ValueError(f"FREETOKEN_MOE_PREFETCH_K={k} must be >= 1")
+        self.k = min(k, num_experts)
+        # 0 = the per-layer default the model wires in
+        self.budget_override = ENV.MOE_PREFETCH_BUDGET.value if budget is None else int(budget)
+        if self.budget_override < 0:
+            raise ValueError(f"FREETOKEN_MOE_PREFETCH_BUDGET={self.budget_override} must be >= 0")
+        self.stream = dedicated_stream(device)
+        # (forked, selected): the predictor starts after forked and records selected after the select
+        self._events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)]
+        # torch creates the CUDA event on its first record: do that now, never inside a capture
+        for pair in self._events:
+            for event in pair:
+                event.record(self.stream)
+        # on mode's select sees nothing resident: prefetch_ensure applies residency and the budget
+        self._none_resident = torch.full((num_experts,), -1, dtype=torch.int32, device=device)
+        self._budget = [0] * num_layers
+        width = 2 * self.k  # room for the whole bs-2 union
+        self.sel = torch.full((num_layers, width), -1, dtype=torch.int32, device=device)
+        self.res = torch.full((num_layers, width), -1, dtype=torch.int32, device=device)
+        self.stats = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64, device=device)
+        # measure mode's six columns, a view: its count kernel and readers keep the stage-A layout
+        self.counters = self.stats[:, :NUM_COLS]
+        self.totals = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64)
+        # the router input of each forked prediction, alive until the compute stream has waited for
+        # the select so the allocator cannot hand its block to later work the predictor still reads
+        self._inflight: list[torch.Tensor | None] = [None] * num_layers
+        # on mode: whether install ran prefetch_ensure for the forked prediction (its copy is then due)
+        self._installed = [False] * num_layers
+        # on mode: the cache whose slot maps prefetch_ensure writes and whose banks the copy reads
+        self.cache = None
+        self.copy_stream = None
+        if mode == "on":
+            # the slim copy is an SM kernel: at default priority a late copy queues for SMs behind the next
+            # compute grid while the compute stream waits on it (bench --on: a 50 us slip cost 14% per step)
+            self.copy_stream = dedicated_stream(device, highest_priority=True)
+            # (dcopy: after prefetch_ensure(L) and layer L-1's demand copy on the compute stream,
+            # pfcopy: after layer L's prefetch copy)
+            self._copy_events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)]
+            for pair in self._copy_events:
+                for event in pair:
+                    event.record(self.copy_stream)
+            # one plan per target layer: prefetch_ensure(L+1) must not overwrite the plan copy(L) reads
+            self.pf_slots = torch.full((num_layers, width), -1, dtype=torch.int32, device=device)
+            self.pf_src = torch.full((num_layers, width), -1, dtype=torch.int32, device=device)
+            self.pf_num = torch.zeros((num_layers, 1), dtype=torch.int64, device=device)
+            self.pf_ready = torch.zeros((num_layers, 1), dtype=torch.int32, device=device)
+        # test hooks, read at enqueue (capture) time: a [num_layers, width] device buffer replacing
+        # sel (on mode still installs at most the budget), and sleeps put on the predictor / copy
+        # stream ahead of their work
+        self.sel_override: torch.Tensor | None = None
+        self.delay_predict_cycles = 0
+        self.delay_copy_cycles = 0
+        logger.info(
+            f"MoE prefetch {mode}: router lookahead top-{self.k}, budget "
+            + (str(self.budget_override) if self.budget_override else f"{GDN_BUDGET} before GDN / {ATTN_BUDGET} before attention layers")
+            + (", slim copies on a dedicated high-priority stream" if mode == "on" else "")
+        )
+
+    def budget_for(self, wired: int) -> int:
+        return min(self.budget_override or wired, self.sel.shape[1])
+
+    def fork(self, target: int, x: torch.Tensor, gate, resident: torch.Tensor, budget: int) -> bool:
+        """Predict ``target``'s experts from ``x`` (the previous layer's router input) on the
+        predictor stream, ordered after the work already on the current stream. Measure mode skips
+        the ids ``resident`` holds (a racy read of the cache); on mode reads no cache state and
+        leaves residency and the budget to ``install``. Returns whether it forked."""
+        if x.shape[0] > MAX_ROWS:
+            return False
+        budget = self.budget_for(budget)
+        if self.mode == "on":
+            self._budget[target] = budget
+            resident, budget = self._none_resident, self.sel.shape[1]
+        forked, selected = self._events[target]
+        forked.record(torch.cuda.current_stream(self.device))
+        self.stream.wait_event(forked)
+        with torch.cuda.stream(self.stream):
+            if self.delay_predict_cycles:
+                torch.cuda._sleep(self.delay_predict_cycles)
+            logits = gate.forward(x)
+            lookahead_select(logits, resident, self.sel[target], self.res[target], k=self.k, budget=budget)
+            if self.sel_override is not None:
+                self.sel[target].copy_(self.sel_override[target])
+            selected.record(self.stream)
+        self._inflight[target] = x
+        self._installed[target] = False
+        return True
+
+    def install(self, target: int | None) -> int | None:
+        """On mode, right after the previous layer's ensure: wait for ``target``'s prediction and
+        install its first budget non-resident candidates with prefetch_ensure on the current
+        (compute) stream. Returns ``target`` when its copy is due (``issue_copy``), else None."""
+        if target is None or self.mode != "on" or self._inflight[target] is None:
+            return None
+        torch.cuda.current_stream(self.device).wait_event(self._events[target][1])
+        self.cache.prefetch_ensure(
+            target, self.sel[target], self.pf_slots[target], self.pf_src[target],
+            self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target],
+        )
+        self._installed[target] = True
+        return target
+
+    def _pending(self, layer_id: int) -> bool:
+        """On mode: ``layer_id``'s prediction was installed and its copy is not joined yet."""
+        return self.mode == "on" and self._installed[layer_id] and self._inflight[layer_id] is not None
+
+    def count_args(self, layer_id: int) -> tuple | None:
+        """On mode: what ``layer_id``'s demand ensure counts into, or None when it was not prefetched."""
+        if not self._pending(layer_id):
+            return None
+        return self.stats[layer_id], self.pf_ready[layer_id], self.pf_num[layer_id], self._inflight[layer_id].shape[0]
+
+    def plan_of(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None:
+        """On mode: ``layer_id``'s prefetch copy plan (slots, count) while its prefetch is in flight."""
+        if not self._pending(layer_id):
+            return None
+        return self.pf_slots[layer_id], self.pf_num[layer_id]
+
+    def join_copy(self, layer_id: int) -> None:
+        """On mode: the compute stream waits for ``layer_id``'s prefetch copy before its own copy, so
+        a prefetched slot ensure(L) picked as a demand victim is never written by two copies."""
+        if self.mode != "on" or self._inflight[layer_id] is None:
+            return
+        # a prediction that was never installed still has to rejoin (a capture needs every fork joined)
+        done = self._copy_events[layer_id][1] if self._pending(layer_id) else self._events[layer_id][1]
+        torch.cuda.current_stream(self.device).wait_event(done)
+        self._inflight[layer_id] = None
+        self._installed[layer_id] = False
+
+    def issue_copy(self, target: int) -> None:
+        """On mode: copy ``target``'s prefetch plan on the copy stream after the work already on the
+        current stream: its prefetch_ensure and the demand copy (which owns the host link first)."""
+        if not self._pending(target):
+            return
+        dcopy, pfcopy = self._copy_events[target]
+        dcopy.record(torch.cuda.current_stream(self.device))
+        self.copy_stream.wait_event(dcopy)
+        with torch.cuda.stream(self.copy_stream):
+            if self.delay_copy_cycles:
+                torch.cuda._sleep(self.delay_copy_cycles)
+            self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True,
+                                 kind=PREFETCH_COPY)
+            _mark_ready_kernel[(1,)](self.pf_ready[target])
+            pfcopy.record(self.copy_stream)
+
+    def join_and_count(self, layer_id: int, slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor) -> None:
+        """Measure mode: join ``layer_id``'s prediction (if one was forked this step) and count it
+        against the routed slots ``slots`` of this layer's ensure."""
+        x = self._inflight[layer_id] if self.mode == "measure" else None
+        if x is None:
+            return
+        torch.cuda.current_stream(self.device).wait_event(self._events[layer_id][1])
+        prefetch_count(
+            slots.view(-1), id_of_slot, misses, self.sel[layer_id], self.res[layer_id], self.counters[layer_id],
+            id_base=layer_id * self.num_experts, rows=x.shape[0],
+        )
+        self._inflight[layer_id] = None
+
+    def reset_counters(self) -> None:
+        self.stats.zero_()
+
+    def reset(self) -> None:
+        """Counters and in-flight state; the streams, events and buffers stay."""
+        self.stats.zero_()
+        self.totals.zero_()
+        self._inflight = [None] * self.num_layers
+        self._installed = [False] * self.num_layers
+
+    def take_window(self) -> torch.Tensor:
+        """This window's ``[num_layers, NUM_STAT_COLS]`` stats on the host (one sync), added to
+        ``totals``; the device ones restart."""
+        window = self.stats.cpu()
+        self.stats.zero_()
+        self.totals += window
+        return window
+
+
+def summarize(counters: torch.Tensor, mode: str = "measure") -> dict:
+    """Aggregate ``[num_layers, NUM_COLS or NUM_STAT_COLS]`` counters into per-layer-call and
+    per-token rates. In on mode the misses are what was left after the prefetch, so coverage is
+    useful / (useful + misses), and late is the share of copied calls whose copy had not landed
+    when the compute stream reached it."""
+    total = counters.sum(0).tolist()
+    calls, issued, useful = total[CALLS], total[ISSUED], total[USEFUL]
+    tokens = int(counters[:, ROWS].max()) if counters.numel() else 0
+    wanted = useful + total[MISSES] if mode == "on" else total[MISSES]
+    stats = {
+        "layer_calls": calls,
+        "tokens": tokens,
+        "issued_per_layer": issued / calls if calls else 0.0,
+        "useful_per_layer": useful / calls if calls else 0.0,
+        "resident_hits_per_layer": total[RESIDENT_HITS] / calls if calls else 0.0,
+        "misses_per_layer": total[MISSES] / calls if calls else 0.0,
+        "precision": useful / issued if issued else 0.0,
+        "coverage": useful / wanted if wanted else 0.0,
+        "useful_per_token": useful / tokens if tokens else 0.0,
+    }
+    if mode == "on":
+        copied = total[COPIED] if len(total) > COPIED else 0
+        stats["copied_calls"] = copied
+        stats["late_fraction"] = total[LATE] / copied if copied else 0.0
+    return stats
+
+
+def format_summary(stats: dict) -> str:
+    if "late_fraction" in stats:
+        return (
+            f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, "
+            f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, "
+            f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, "
+            f"late={stats['late_fraction']:.3f} of {stats['copied_calls']} copies"
+        )
+    return (
+        f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, "
+        f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, "
+        f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, "
+        f"resident_hits/layer={stats['resident_hits_per_layer']:.2f}"
+    )
+
+
+def format_per_layer(counters: torch.Tensor) -> str:
+    parts = []
+    for layer, row in enumerate(counters.tolist()):
+        if row[CALLS]:
+            c = row[CALLS]
+            parts.append(f"L{layer}={row[USEFUL] / c:.2f}/{row[ISSUED] / c:.2f}/{row[MISSES] / c:.2f}")
+    return ", ".join(parts)
diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py
index 123086ed51..ba570b5169 100644
--- a/python/freetoken/moe/scored_ensure.py
+++ b/python/freetoken/moe/scored_ensure.py
@@ -106,6 +106,22 @@ def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K,
     return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1)
 
 
+@triton.jit
+def _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K: tl.constexpr, id_base):
+    """_phase1 for a prefetch: -1 entries are padding, hits are left alone (a prefetch is not an
+    access), and the misses rank in query order so the select's best candidates install first."""
+    k = tl.arange(0, BLOCK_K)
+    raw = tl.load(query_ptr + k, mask=k < K, other=-1)
+    kmask = raw >= 0
+    q = tl.where(kmask, raw + id_base, 0)
+    s = tl.load(slot_of_id_ptr + q, mask=kmask, other=0)
+    miss = kmask & (s == -1)
+    same = (q[:, None] == q[None, :]) & (k[:, None] > k[None, :]) & kmask[:, None] & kmask[None, :]
+    first_miss = miss & (tl.sum(same.to(tl.int32), axis=1) == 0)
+    rank = tl.cumsum(first_miss.to(tl.int32), axis=0) - 1
+    return q, first_miss, rank, tl.sum(first_miss.to(tl.int32))
+
+
 @triton.jit
 def _stats(stats_ptr, first, num_missing):
     # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA.
@@ -144,13 +160,14 @@ def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other):
     )
 
 
-@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"])
+@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride", "pf_rows", "pf_budget"])
 def _scored_ensure_kernel(
     query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr,
     out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr,
     tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr,
     owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr,
-    K, num_cached, id_base, nm_topk, nm_stride, nm_thr,
+    pf_stats_ptr, pf_ready_ptr, pf_num_ptr,
+    K, num_cached, id_base, nm_topk, nm_stride, nm_thr, pf_rows, pf_budget,
     BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr,
     NM_ROWS: tl.constexpr,
     USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr,
@@ -158,15 +175,27 @@ def _scored_ensure_kernel(
     NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr,
     BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr,
     SLOT_BITS: tl.constexpr,
+    LOWPRI: tl.constexpr, PF_COUNT: tl.constexpr, PREFETCH: tl.constexpr,
 ):
-    step = tl.load(lru_step_ptr) + 1
-    tl.store(lru_step_ptr, step)
+    if PREFETCH:
+        # not an access: the clock stays, so usage == step still pins the last demand call's slots
+        # (at step 0, right after a reset, every slot reads pinned: nothing installs before a demand call)
+        step = tl.load(lru_step_ptr)
+        tl.store(num_copy_ptr, tl.zeros([], tl.int64))
+        tl.store(pf_ready_ptr, 0)
+    else:
+        step = tl.load(lru_step_ptr) + 1
+        tl.store(lru_step_ptr, step)
     if POLICY != 0:
         tok = tl.load(tok_ptr)
         if BUMP_TOK:
             tok = tok + 1
             tl.store(tok_ptr, tok)
         layer = id_base // NUM_EXPERTS
+        if PREFETCH:
+            # score from where the pass is, right after layer - 1's demand call: the target layer's
+            # residents are next up, not a whole pass away (else it evicts what the layer is about to route)
+            layer = layer - 1
         now = tok * NUM_LAYERS + layer
         if POLICY == 3:
             if UPDATE_STATE:
@@ -187,9 +216,18 @@ def _scored_ensure_kernel(
                 tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0))
                 # the victim scan and the routed update below read and rewrite what this wrote
                 tl.debug_barrier()
-    q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1(
-        query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K,
-        id_base)
+    if PREFETCH:
+        q, first_miss, rank, num_missing = _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K, id_base)
+    else:
+        if PF_COUNT:
+            kk = tl.arange(0, BLOCK_K)
+            sk = tl.load(slot_of_id_ptr + tl.load(query_ptr + kk, mask=kk < K, other=0) + id_base, mask=kk < K, other=-1)
+            # a hit on a held slot with usage 0 consumes a prefetch; read it before _phase1 bumps it
+            was_pf = (sk >= 0) & (tl.load(lru_usage_ptr + sk, mask=sk >= 0, other=1) == 0)
+            tl.debug_barrier()
+        q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1(
+            query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K,
+            id_base)
 
     if num_missing > 0:
         # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted.
@@ -199,7 +237,13 @@ def _scored_ensure_kernel(
         u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX)
         if POLICY == 0:
             umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype)
-            u = tl.where((u == step) | (~cmask), umax, u)
+            if LOWPRI:
+                held0 = cmask & (tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) >= 0)
+                # empty slots, then unconsumed prefetches (held, usage 0), then residents by recency
+                u = tl.where((u == step) | (~cmask), umax, tl.where(held0, u + 1, 0))
+            else:
+                u = tl.where((u == step) | (~cmask), umax, u)
+            n_free = tl.sum((u != umax).to(tl.int32))
         else:
             oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1)
             held = cmask & (oid >= 0)
@@ -221,6 +265,10 @@ def _scored_ensure_kernel(
                 score += (W_Q4 * lcv) >> 4
             # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one.
             key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64))
+            if LOWPRI:
+                # an unconsumed prefetch (held, usage 0) keys above every empty slot and below every
+                # resident: the caps keep score + 2^31 > 2^30 - 2^25 - 2^22, far above this band's 1
+                key = tl.where(held & (u == 0), (1 << SLOT_BITS) | c.to(tl.int64), key)
             evictable = cmask & (u != step)
             if PIN_SINCE:
                 # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too
@@ -228,7 +276,16 @@ def _scored_ensure_kernel(
             # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin
             pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1)
             key = tl.where(evictable, key, pinned_key)
-        for i in tl.range(num_missing):
+            n_free = tl.sum(evictable.to(tl.int32))
+        if PREFETCH:
+            # never the all-pinned fallback: that slot may be under the running GEMM; misses rank in
+            # query order, so the budget keeps the first non-resident candidates
+            n_iter = tl.minimum(tl.minimum(num_missing, n_free), pf_budget)
+            tl.store(num_copy_ptr, n_iter.to(tl.int64))
+            tl.store(pf_stats_ptr, tl.load(pf_stats_ptr) + n_iter)
+        else:
+            n_iter = num_missing
+        for i in tl.range(n_iter):
             if POLICY == 0:
                 victim = tl.argmin(u, axis=0).to(tl.int32)
             else:
@@ -240,17 +297,38 @@ def _scored_ensure_kernel(
             e = tl.sum(tl.where((rank == i) & first_miss, q, 0))
             tl.store(id_of_slot_ptr + victim, e)
             tl.store(slot_of_id_ptr + e, victim)
-            tl.store(lru_usage_ptr + victim, step)
+            if PREFETCH:
+                # usage 0 on a held slot marks it low priority until a demand call touches it
+                tl.store(lru_usage_ptr + victim, 0)
+                if POLICY != 0:
+                    tl.store(owner_ptr + victim, -1)
+            else:
+                tl.store(lru_usage_ptr + victim, step)
             tl.store(dst_ptr + i, victim)
             tl.store(src_ptr + i, e - id_base)  # back to the caller's id space
-            out = tl.where((rank == i) & miss, victim, out)
+            if not PREFETCH:
+                out = tl.where((rank == i) & miss, victim, out)
             if POLICY == 0:
                 u = tl.where(c == victim, umax, u)  # claim in-register
             else:
                 key = tl.where(c == victim, pinned_key, key)
 
-    # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr.
-    tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask)
+    if not PREFETCH:
+        # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr.
+        tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask)
+    if PF_COUNT:
+        # one writer per column: pf_ensure adds ISSUED (column 0), earlier on this stream
+        pn = tl.load(pf_num_ptr)
+        late = (pn > 0) & (tl.load(pf_ready_ptr, volatile=True) == 0)
+        col = tl.arange(0, 8)
+        add = tl.where(col == 1, tl.sum((first & was_pf).to(tl.int32)).to(tl.int64), 0)
+        add = tl.where(col == 3, 1, add)
+        add = tl.where(col == 4, pf_rows, add)
+        add = tl.where(col == 5, num_missing.to(tl.int64), add)
+        add = tl.where(col == 6, late.to(tl.int64), add)
+        add = tl.where(col == 7, (pn > 0).to(tl.int64), add)
+        cols = (col == 1) | (col >= 3)
+        tl.store(pf_stats_ptr + col, tl.load(pf_stats_ptr + col, mask=cols, other=0) + add, mask=cols)
     if POLICY != 0:
         if UPDATE_STATE:
             # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot
@@ -265,7 +343,7 @@ def _scored_ensure_kernel(
                 tl.store(ct_ptr + q, now, mask=first)
                 tl.store(slot_lc_ptr + out, lc_new, mask=first)
                 tl.store(slot_ct_ptr + out, now, mask=first)
-        else:
+        elif not PREFETCH:
             # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it
             tl.store(owner_ptr + out, -1, mask=first_miss)
     if COLLECT_STATS:
@@ -311,6 +389,13 @@ def scored_ensure(
     beta: float = BETA,
     w: float = W,
     halflife: int = HALFLIFE,
+    lowpri: bool = False,
+    pf_stats: torch.Tensor | None = None,
+    pf_ready: torch.Tensor | None = None,
+    pf_num: torch.Tensor | None = None,
+    pf_rows: int = 0,
+    pf_budget: int | None = None,
+    prefetch: bool = False,
 ) -> None:
     """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``.
 
@@ -324,7 +409,21 @@ def scored_ensure(
     prefill's chunked calls cannot evict each other's experts; policy > 0 only.
     Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of
     ``query``) are within ``near_miss_thr`` of their row's lowest routed logit.
+
+    FREETOKEN_MOE_PREFETCH=on: ``lowpri`` keys held slots with usage 0 (unconsumed prefetches)
+    between the empty slots and every resident. ``pf_stats`` (a ``[8]`` int64 prefetch stats row)
+    also counts this demand call: hits on such slots, the call, ``pf_rows``, the misses, and whether
+    the prefetch plan ``pf_num`` had rows whose copy had not set ``pf_ready`` yet.
+    ``prefetch=True`` is :func:`prefetch_ensure`'s install mode instead, installing at most
+    ``pf_budget`` ids (default: all).
     """
+    if prefetch:
+        assert lowpri and pf_stats is not None and pf_ready is not None and stats is None
+        assert not update_state and not bump_tok and pin_since is None and router_logits is None
+    if pf_stats is not None:
+        assert lowpri and pf_stats.dtype == torch.int64 and pf_stats.numel() == 8 and pf_stats.is_contiguous()
+        assert pf_ready is not None and pf_ready.dtype == torch.int32 and pf_ready.numel() == 1
+        assert prefetch or (pf_num is not None and pf_num.dtype == torch.int64 and pf_num.numel() == 1)
     k = query.numel()
     num_cached = id_of_slot.numel()
     assert query.dtype == torch.int32 and query.is_contiguous()
@@ -372,7 +471,11 @@ def scored_ensure(
         dummy if slot_lc is None else slot_lc,
         dummy if slot_ct is None else slot_ct,
         dummy if pin_since is None else pin_since,
-        k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr),
+        dummy if pf_stats is None else pf_stats,
+        dummy if pf_ready is None else pf_ready,
+        dummy if pf_num is None else pf_num,
+        k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), int(pf_rows),
+        k if pf_budget is None else int(pf_budget),
         BLOCK_K=triton.next_power_of_2(k),
         BLOCK_C=block_c,
         BLOCK_E=triton.next_power_of_2(num_experts),
@@ -391,5 +494,40 @@ def scored_ensure(
         DECAY_MUL=decay_mul,
         DT_MAX=dt_max,
         SLOT_BITS=slot_bits,
+        LOWPRI=bool(lowpri),
+        PF_COUNT=pf_stats is not None and not prefetch,
+        PREFETCH=bool(prefetch),
         num_warps=_num_warps_for(block_c),
     )
+
+
+def prefetch_ensure(
+    query: torch.Tensor,
+    slot_of_id: torch.Tensor,
+    id_of_slot: torch.Tensor,
+    lru_usage: torch.Tensor,
+    lru_step: torch.Tensor,
+    dst_slots: torch.Tensor,
+    src_rows: torch.Tensor,
+    num_copy: torch.Tensor,
+    pf_stats: torch.Tensor,
+    pf_ready: torch.Tensor,
+    budget: int | None = None,
+    **state,
+) -> None:
+    """Install the ids of ``query`` (-1 = padding) that are not resident as low-priority slots
+    (held, usage 0), at most the first ``budget`` of them in query order (default: all), and
+    write their copy plan to ``dst_slots``/``src_rows``/``num_copy``.
+
+    Unlike a demand call it bumps no clock, touches no hit and no per-id state, and so keeps
+    every slot of the last demand call (usage == step) pinned; victims follow the demand key
+    (empty, then earlier unconsumed prefetches, then the coldest residents), scored as of the
+    previous layer's call, the position the prefetch runs at. With fewer
+    evictable slots than misses it installs only the first ones, in query order. Adds the
+    installs to ``pf_stats[0]`` and clears ``pf_ready``.
+    """
+    scored_ensure(
+        query, slot_of_id, id_of_slot, lru_usage, lru_step, query, src_rows, dst_slots, num_copy,
+        update_state=False, lowpri=True, pf_stats=pf_stats, pf_ready=pf_ready, pf_budget=budget, prefetch=True,
+        **state,
+    )
diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py
new file mode 100644
index 0000000000..981239758f
--- /dev/null
+++ b/python/freetoken/moe/slot_audit.py
@@ -0,0 +1,688 @@
+"""Expert slot audit (``FREETOKEN_MOE_SLOT_AUDIT=N``): a debug detector for PERSISTENT slot corruption.
+
+Two parts, both on the device and off by default:
+
+- writer history: every writer of the slot maps (demand and prefetch ensures, prefill ensures, the
+  whole-layer materialize, prefill buffer invalidation, reset) and every writer of the slot banks
+  (demand, prefetch and prefill copies, the prefill double buffers) is followed on its writer's
+  stream by a tiny kernel keyed on the same plan (or slot range) that appends ``(seq, kind, decode
+  step, lru step, layer, expert row, ...)`` to a per-slot ring of the last ``DEPTH`` metadata events
+  and one of the last ``DEPTH`` byte events. The map writers all run on the compute stream except the
+  prefill buffer invalidation (on the prefill copy stream, fenced behind it); the byte writers also
+  run on the prefetch and prefill copy streams. ``seq`` is one global atomic counter, so the two rings
+  merge into one cross-stream order. Clears are recorded just before their writer, so they carry the
+  old owner.
+  Fixed shapes and no host sync: the recorders capture into the decode graph.
+- audit: every N decode steps, after the forward on the compute stream, every held slot is checked:
+  its bytes in every bank against the host row of the (layer, expert) its ``id_of_slot`` names
+  (host rows gathered with the fused copy into a scratch of fewer than 64 MiB, one layer chunk at a
+  time), ``slot_for_id[id_of_slot[s]] == s``, and, over every id, ``id_of_slot[slot_for_id[id]] == id``
+  (a stale entry routes that id to another expert's bytes). The first ``RING`` bad (slot, owner)
+  pairs are kept with both history rings; the engine logs them and the totals every
+  MOE_STATS_INTERVAL decode steps and at shutdown, with a diagnosis read from the history.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.env import ENV
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+# events kept per slot and ring (metadata / bytes)
+DEPTH = 4
+# event fields
+(E_SEQ, E_KIND, E_DSTEP, E_LSTEP, E_LAYER, E_ROW, E_AUX, E_OWNER, E_USAGE, E_POS) = range(10)
+EVENT_FIELDS = 10
+# metadata kinds (installs name the new owner; clears the old one, recorded before the writer)
+DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL = 1, 2, 3, 4
+MATERIALIZE_CLEAR, INVALIDATE, RESET = 5, 6, 7
+# byte kinds
+DEMAND_COPY, PREFETCH_COPY, PREFILL_COPY, MATERIALIZE_COPY, PAGEABLE_COPY = 11, 12, 13, 14, 15
+PREFILL_BUFFER, PREFILL_SPLIT_H2D, PREFILL_HIT_D2D, OTHER_COPY = 16, 17, 18, 19
+KIND_NAMES = {
+    DEMAND_INSTALL: "demand install", PREFETCH_INSTALL: "prefetch install", PREFILL_INSTALL: "prefill ensure install",
+    MATERIALIZE_INSTALL: "materialize install", MATERIALIZE_CLEAR: "materialize clear",
+    INVALIDATE: "prefill buffer invalidate", RESET: "reset",
+    DEMAND_COPY: "demand copy", PREFETCH_COPY: "prefetch copy", PREFILL_COPY: "prefill ensure copy",
+    MATERIALIZE_COPY: "materialize copy", PAGEABLE_COPY: "pageable materialize copy",
+    PREFILL_BUFFER: "prefill buffer copy", PREFILL_SPLIT_H2D: "prefill buffer H2D (misses and small banks)",
+    PREFILL_HIT_D2D: "prefill buffer D2D", OTHER_COPY: "copy",
+}
+INSTALLS = (DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL)
+# the copy each install expects next
+COPY_FOR = {
+    DEMAND_INSTALL: DEMAND_COPY, PREFETCH_INSTALL: PREFETCH_COPY, PREFILL_INSTALL: PREFILL_COPY,
+    MATERIALIZE_INSTALL: MATERIALIZE_COPY,
+}
+# recorder row sources: the plan's src rows, slot - row_base, the range index, the old owner
+ROW_PLAN, ROW_FROM_SLOT, ROW_INDEX, ROW_OLD_OWNER = range(4)
+# recorder predicates: every slot, held slots, slots >= E held by the layer
+PRED_ALL, PRED_HELD, PRED_LAYER_ABOVE_E = range(3)
+# totals: audits, held slots seen, slots byte-checked, held slots not byte-checked, bad slots found
+# (summed over audits), of them with bad bytes / a forward map miss / named by a stale slot_for_id /
+# an owner out of range, slot_for_id entries past the cache, new records, repeat findings, events
+# recorded, plan entries naming no slot
+(T_AUDITS, T_HELD, T_BYTES_CHECKED, T_UNCHECKED, T_BAD, T_BYTES, T_FWD, T_STALE, T_OWNER_RANGE, T_MAP_RANGE,
+ T_NEW, T_STILL, T_EVENTS, T_BAD_PLAN) = range(14)
+NUM_TOTALS = 16
+# per-slot audit flags
+F_BYTES, F_FWD, F_STALE, F_OWNER_RANGE = 1, 2, 4, 8
+# record header, then the metadata ring and the byte ring as stored ([DEPTH, EVENT_FIELDS] each)
+(H_AUDIT, H_DSTEP, H_LSTEP, H_SLOT, H_OWNER, H_MAPPED, H_USAGE, H_FLAGS, H_BANKS, H_FIRST_BANK, H_FIRST_BYTE,
+ H_BAD_WORDS, H_STALE_ID, H_STALE_N, H_META_N, H_BYTES_N) = range(16)
+HEADER = 16
+RING_EVENTS = DEPTH * EVENT_FIELDS
+RECORD_FIELDS = HEADER + 2 * RING_EVENTS
+RING = 32
+SCRATCH_BYTES = 64 << 20
+MAX_CHUNK_ROWS = 64
+COMPARE_WORDS = 4096
+_NO_BAD = tl.constexpr(0x7FFFFFFF)
+_D = tl.constexpr(DEPTH)
+_F = tl.constexpr(EVENT_FIELDS)
+
+
+def resolve_interval(interval: int | None = None) -> int:
+    value = ENV.MOE_SLOT_AUDIT.value if interval is None else interval
+    value = int(value or 0)
+    if value < 0:
+        raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={value}: expected 0 (off) or the decode steps between audits")
+    return value
+
+
+@triton.jit
+def _put(vals, f, idx: tl.constexpr, v):
+    return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals)
+
+
+@triton.jit
+def _bump_step_kernel(dstep_ptr):
+    tl.store(dstep_ptr, tl.load(dstep_ptr) + 1)
+
+
+@triton.jit(do_not_specialize=["kind", "layer", "start", "row_base", "width", "num_slots"])
+def _record_kernel(
+    hist_ptr, cur_ptr, seq_ptr, dstep_ptr, lstep_ptr, id_of_slot_ptr, usage_ptr, totals_ptr,
+    dst_ptr, src_ptr, aux_ptr, num_ptr,
+    kind, layer, start, row_base, width, num_slots,
+    PLAN: tl.constexpr, ROW: tl.constexpr, HAS_AUX: tl.constexpr, HAS_NUM: tl.constexpr, PRED: tl.constexpr,
+    E: tl.constexpr, BLOCK: tl.constexpr,
+):
+    """Append one event per written slot to that slot's ring: plan entries ``i < min(*num, width)``
+    (``PLAN``) or the slots ``start + i`` for ``i < width``."""
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    n = width
+    if HAS_NUM:
+        n = tl.minimum(tl.load(num_ptr).to(tl.int32), width)
+    valid = i < n
+    if PLAN:
+        slot = tl.load(dst_ptr + i, mask=valid, other=-1)
+    else:
+        slot = start + i
+    ok = valid & (slot >= 0) & (slot < num_slots)
+    bad_plan = valid & ~ok
+    owner = tl.load(id_of_slot_ptr + slot, mask=ok, other=-3)
+    usage = tl.load(usage_ptr + slot, mask=ok, other=-3)
+    if PRED == 1:
+        ok = ok & (owner >= 0)
+    if PRED == 2:
+        ok = ok & (owner >= layer * E) & (owner < layer * E + E) & (slot >= E)
+    lay = tl.zeros([BLOCK], tl.int32) + layer
+    aux = tl.zeros([BLOCK], tl.int32) - 1
+    if ROW == 0:
+        row = tl.load(src_ptr + i, mask=ok, other=-1)
+    elif ROW == 1:
+        row = slot - row_base
+    elif ROW == 2:
+        row = i
+    else:
+        row = tl.where(owner >= 0, owner % E, -1)
+        lay = tl.where(owner >= 0, owner // E, -1)
+        aux = aux + 1 + layer
+    if HAS_AUX:
+        aux = tl.load(aux_ptr + i, mask=ok, other=-1)
+    cnt = tl.sum(ok.to(tl.int64), axis=0)
+    seq = tl.atomic_add(seq_ptr, cnt) + tl.cumsum(ok.to(tl.int64), axis=0) - 1
+    pos = tl.atomic_add(cur_ptr + slot, 1, mask=ok)
+    f = tl.arange(0, 16)
+    z = tl.zeros([BLOCK], tl.int64)
+    vals = tl.zeros([BLOCK, 16], tl.int64) - 1
+    vals = _put(vals, f, 0, seq)
+    vals = _put(vals, f, 1, z + kind)
+    vals = _put(vals, f, 2, z + tl.load(dstep_ptr))
+    vals = _put(vals, f, 3, z + tl.load(lstep_ptr))
+    vals = _put(vals, f, 4, lay)
+    vals = _put(vals, f, 5, row)
+    vals = _put(vals, f, 6, aux)
+    vals = _put(vals, f, 7, owner)
+    vals = _put(vals, f, 8, usage)
+    vals = _put(vals, f, 9, i)
+    at = (slot.to(tl.int64) * _D + (pos % _D)) * _F
+    tl.store(hist_ptr + at[:, None] + f[None, :], vals, mask=ok[:, None] & (f[None, :] < _F))
+    t = tl.arange(0, 2)
+    add = tl.where(t == 0, cnt, tl.sum(bad_plan.to(tl.int64), axis=0))
+    tl.atomic_add(totals_ptr + 12 + t, add, sem="relaxed")
+
+
+@triton.jit(do_not_specialize=["num_slots", "num_ids"])
+def _audit_meta_kernel(
+    id_of_slot_ptr, slot_for_id_ptr, layer_ok_ptr, count_ptr, rows_ptr, slots_ptr,
+    flags_ptr, stale_id_ptr, stale_n_ptr, totals_ptr, num_slots, num_ids,
+    E: tl.constexpr, CAP: tl.constexpr, BLOCK: tl.constexpr,
+):
+    """Both map checks, and each byte-checkable held slot appended to its layer's list."""
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    ms = i < num_slots
+    o = tl.load(id_of_slot_ptr + i, mask=ms, other=-1)
+    held = ms & (o >= 0)
+    in_range = held & (o < num_ids)
+    mapped = tl.load(slot_for_id_ptr + o, mask=in_range, other=-2)
+    fwd = in_range & (mapped != i)
+    out_of_range = held & ~in_range
+    fl = tl.where(fwd, 2, 0) | tl.where(out_of_range, 8, 0)
+    tl.atomic_or(flags_ptr + i, fl, mask=ms & (fl != 0))
+    layer = tl.where(in_range, o // E, 0)
+    ok = in_range & (tl.load(layer_ok_ptr + layer, mask=in_range, other=0) != 0)
+    p = tl.atomic_add(count_ptr + layer, 1, mask=ok)
+    take = ok & (p < CAP)
+    tl.store(rows_ptr + layer * CAP + p, o - layer * E, mask=take)
+    tl.store(slots_ptr + layer * CAP + p, i, mask=take)
+    mi = i < num_ids
+    s = tl.load(slot_for_id_ptr + i, mask=mi, other=-1)
+    live = mi & (s >= 0)
+    s_ok = live & (s < num_slots)
+    owner = tl.load(id_of_slot_ptr + s, mask=s_ok, other=-1)
+    stale = s_ok & (owner != i)
+    tl.atomic_or(flags_ptr + s, 4, mask=stale)
+    tl.store(stale_id_ptr + s, i, mask=stale)
+    tl.atomic_add(stale_n_ptr + s, 1, mask=stale)
+    t = tl.arange(0, 16)
+    add = tl.where(t == 1, tl.sum(held.to(tl.int64), axis=0), 0)
+    add = tl.where(t == 2, tl.sum(take.to(tl.int64), axis=0), add)
+    add = tl.where(t == 3, tl.sum((held & ~take).to(tl.int64), axis=0), add)
+    add = tl.where(t == 6, tl.sum(fwd.to(tl.int64), axis=0), add)
+    add = tl.where(t == 7, tl.sum(stale.to(tl.int64), axis=0), add)
+    add = tl.where(t == 8, tl.sum(out_of_range.to(tl.int64), axis=0), add)
+    add = tl.where(t == 9, tl.sum((live & ~s_ok).to(tl.int64), axis=0), add)
+    add = tl.where((t == 0) & (tl.program_id(0) == 0), 1, add)
+    tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed")
+
+
+@triton.jit
+def _audit_nums_kernel(count_ptr, nums_ptr, L: tl.constexpr, C: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr):
+    j = tl.arange(0, BLOCK)
+    m = j < L * C
+    cnt = tl.load(count_ptr + j // C, mask=m, other=0)
+    tl.store(nums_ptr + j, tl.minimum(tl.maximum(cnt - (j % C) * R, 0), R).to(tl.int64), mask=m)
+
+
+@triton.jit
+def _audit_compare_kernel(cache_ptrs, ref_ptrs, words_ptr, slots_ptr, num_ptr, bad_cnt_ptr, bad_first_ptr,
+                          NB: tl.constexpr, BLOCK: tl.constexpr):
+    """Program (row, bank, chunk): one chunk of a listed slot's bank row against its gathered host row."""
+    r = tl.program_id(0)
+    if r < tl.load(num_ptr):
+        b = tl.program_id(1)
+        words = tl.load(words_ptr + b)
+        off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        m = off < words
+        slot = tl.load(slots_ptr + r).to(tl.int64)
+        got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0)
+        want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0)
+        bad = m & (got != want)
+        nbad = tl.sum(bad.to(tl.int32), axis=0)
+        if nbad > 0:
+            tl.atomic_add(bad_cnt_ptr + slot * NB + b, nbad)
+            tl.atomic_min(bad_first_ptr + slot * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32))
+
+
+@triton.jit(do_not_specialize=["num_slots", "num_ids"])
+def _audit_report_kernel(
+    id_of_slot_ptr, slot_for_id_ptr, usage_ptr, lstep_ptr, dstep_ptr,
+    flags_ptr, stale_id_ptr, stale_n_ptr, bad_cnt_ptr, bad_first_ptr, reported_ptr,
+    meta_hist_ptr, meta_cur_ptr, bytes_hist_ptr, bytes_cur_ptr,
+    ring_ptr, cursor_ptr, totals_ptr, count_ptr,
+    num_slots, num_ids,
+    NB: tl.constexpr, NB_P: tl.constexpr, L: tl.constexpr, L_P: tl.constexpr, RING_N: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Keep each new bad (slot, owner) with its history, count the findings and clear the work arrays."""
+    s = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    ms = s < num_slots
+    b = tl.arange(0, NB_P)
+    cell = ms[:, None] & (b < NB)[None, :]
+    cnt = tl.load(bad_cnt_ptr + s[:, None] * NB + b[None, :], mask=cell, other=0)
+    first = tl.load(bad_first_ptr + s[:, None] * NB + b[None, :], mask=cell, other=_NO_BAD)
+    bank_bad = cnt > 0
+    bytes_bad = tl.sum(bank_bad.to(tl.int32), axis=1) > 0
+    bank_mask = tl.sum(tl.where(bank_bad, 1 << b[None, :], 0), axis=1)
+    first_bank = tl.min(tl.where(bank_bad, b[None, :], NB_P), axis=1)
+    first_word = tl.sum(tl.where(b[None, :] == first_bank[:, None], first, 0), axis=1)
+    fl = tl.load(flags_ptr + s, mask=ms, other=0) | tl.where(bytes_bad, 1, 0)
+    bad = ms & (fl != 0)
+    owner = tl.load(id_of_slot_ptr + s, mask=ms, other=-1)
+    sid = tl.load(stale_id_ptr + s, mask=ms, other=-1)
+    key = (owner.to(tl.int64) + 1) * (num_ids + 2) + sid + 1
+    new = bad & (tl.load(reported_ptr + s, mask=ms, other=-1) != key)
+    tl.store(reported_ptr + s, tl.where(bad, key, -1), mask=ms)
+    pos = tl.atomic_add(cursor_ptr + s * 0, 1, mask=new)
+    keep = new & (pos < RING_N)
+    in_range = ms & (owner >= 0) & (owner < num_ids)
+    f = tl.arange(0, 16)
+    z = tl.zeros([BLOCK], tl.int64)
+    vals = tl.zeros([BLOCK, 16], tl.int64) - 1
+    vals = _put(vals, f, 0, z + tl.load(totals_ptr) - 1)
+    vals = _put(vals, f, 1, z + tl.load(dstep_ptr))
+    vals = _put(vals, f, 2, z + tl.load(lstep_ptr))
+    vals = _put(vals, f, 3, s)
+    vals = _put(vals, f, 4, owner)
+    vals = _put(vals, f, 5, tl.load(slot_for_id_ptr + owner, mask=in_range, other=-2))
+    vals = _put(vals, f, 6, tl.load(usage_ptr + s, mask=ms, other=-1))
+    vals = _put(vals, f, 7, fl)
+    vals = _put(vals, f, 8, bank_mask)
+    vals = _put(vals, f, 9, tl.where(bytes_bad, first_bank, -1))
+    vals = _put(vals, f, 10, tl.where(bytes_bad, first_word.to(tl.int64) * 4, -1))
+    vals = _put(vals, f, 11, tl.sum(cnt, axis=1))
+    vals = _put(vals, f, 12, sid)
+    vals = _put(vals, f, 13, tl.load(stale_n_ptr + s, mask=ms, other=0))
+    vals = _put(vals, f, 14, tl.load(meta_cur_ptr + s, mask=ms, other=0))
+    vals = _put(vals, f, 15, tl.load(bytes_cur_ptr + s, mask=ms, other=0))
+    rec = ring_ptr + pos.to(tl.int64) * (16 + 2 * _D * _F)
+    tl.store(rec[:, None] + f[None, :], vals, mask=keep[:, None])
+    j = tl.arange(0, 64)
+    jm = keep[:, None] & (j < _D * _F)[None, :]
+    src = s.to(tl.int64)[:, None] * (_D * _F) + j[None, :]
+    tl.store(rec[:, None] + 16 + j[None, :], tl.load(meta_hist_ptr + src, mask=jm, other=-1), mask=jm)
+    tl.store(rec[:, None] + 16 + _D * _F + j[None, :], tl.load(bytes_hist_ptr + src, mask=jm, other=-1), mask=jm)
+    t = tl.arange(0, 16)
+    add = tl.where(t == 4, tl.sum(bad.to(tl.int64), axis=0), 0)
+    add = tl.where(t == 5, tl.sum((ms & bytes_bad).to(tl.int64), axis=0), add)
+    add = tl.where(t == 10, tl.sum(new.to(tl.int64), axis=0), add)
+    add = tl.where(t == 11, tl.sum((bad & ~new).to(tl.int64), axis=0), add)
+    tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed")
+    tl.store(flags_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms)
+    tl.store(stale_id_ptr + s, tl.zeros([BLOCK], tl.int32) - 1, mask=ms)
+    tl.store(stale_n_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms)
+    tl.store(bad_cnt_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32), mask=cell)
+    tl.store(bad_first_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32) + _NO_BAD, mask=cell)
+    if tl.program_id(0) == 0:
+        ll = tl.arange(0, L_P)
+        tl.store(count_ptr + ll, tl.zeros([L_P], tl.int32), mask=ll < L)
+
+
+class SlotAuditor:
+    """Per-slot writer history (recorded by ``record_*`` on the writer's stream) and the periodic
+    audit (``end_decode_step`` -> ``scan`` on the compute stream). ``bind`` sizes everything for a
+    cache's slots and banks; a rebuild rebinds and starts a new history, the totals and records stay.
+    Holds no reference to its cache: every call takes it."""
+
+    def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, interval: int,
+                 ring: int | None = None, scratch_bytes: int = SCRATCH_BYTES) -> None:
+        self.num_layers = num_layers
+        self.num_experts = num_experts
+        self.device = device
+        self.interval = interval
+        self.ring_n = RING if ring is None else ring
+        self.scratch_bytes = scratch_bytes
+        ring = self.ring_n
+        self.dstep = torch.zeros((1,), dtype=torch.int64, device=device)
+        self.seq = torch.zeros((1,), dtype=torch.int64, device=device)
+        self.totals = torch.zeros((NUM_TOTALS,), dtype=torch.int64, device=device)
+        self.ring = torch.full((ring, RECORD_FIELDS), -1, dtype=torch.int64, device=device)
+        self.cursor = torch.zeros((1,), dtype=torch.int64, device=device)
+        self.count = torch.zeros((num_layers,), dtype=torch.int32, device=device)
+        self.num_slots = 0
+        self.bank_names: list[str] = []
+        self._feats: list[int] = []
+        self._steps = 0
+        self._printed = 0
+        self._last = torch.zeros((NUM_TOTALS,), dtype=torch.int64)
+        self._dummy = torch.zeros((1,), dtype=torch.int32, device=device)
+        logger.info(f"MoE slot audit: recording every slot writer, auditing every held slot each {interval} decode steps")
+
+    # ------------------------------------------------------------------ setup
+
+    def bind(self, cache) -> None:
+        """Size the history and audit arrays for ``cache``'s slots and banks (a rebuild starts a new
+        history) and point the byte checks at its banks."""
+        caches = [bank for _, bank in cache.banks]
+        feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches]
+        assert all(f % 4 == 0 for f in feats), feats
+        if self.bank_names:
+            assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed"
+        else:
+            self.bank_names = list(cache.bank_schema)
+            self._feats = feats
+            self._alloc_scratch(cache, feats)
+        self.cache_ptrs = torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64, device=self.device)
+        S, dev = cache.cache_size, self.device
+        self.num_slots = S
+        self.meta_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev)
+        self.bytes_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev)
+        self.meta_cur = torch.zeros((S,), dtype=torch.int32, device=dev)
+        self.bytes_cur = torch.zeros((S,), dtype=torch.int32, device=dev)
+        self.flags = torch.zeros((S,), dtype=torch.int32, device=dev)
+        self.stale_id = torch.full((S,), -1, dtype=torch.int32, device=dev)
+        self.stale_n = torch.zeros((S,), dtype=torch.int32, device=dev)
+        self.reported = torch.full((S,), -1, dtype=torch.int64, device=dev)
+        nb = len(feats)
+        self.bad_cnt = torch.zeros((S * nb,), dtype=torch.int32, device=dev)
+        self.bad_first = torch.full((S * nb,), 0x7FFFFFFF, dtype=torch.int32, device=dev)
+        # a correct cache holds at most num_experts slots of a layer; more are counted unchecked
+        per_layer = min(self.num_experts, S)
+        self.chunks = triton.cdiv(per_layer, self.rows)
+        self.cap = self.chunks * self.rows
+        self.list_rows = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev)
+        self.list_slots = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev)
+        self.nums = torch.zeros((self.num_layers * self.chunks,), dtype=torch.int64, device=dev)
+        # layers without a device alias (LOCKED/PAGEABLE host banks) cannot be gathered
+        self.layer_ok_host = [layer not in cache._unpinned_layers for layer in range(self.num_layers)]
+        self.layer_ok = torch.tensor(self.layer_ok_host, dtype=torch.int32, device=dev)
+        self.count.zero_()
+
+    def _alloc_scratch(self, cache, feats: list[int]) -> None:
+        row = sum(feats)
+        self.rows = max(1, min(MAX_CHUNK_ROWS, (self.scratch_bytes - 1) // row))
+        rows = self.rows
+        self.scratch = torch.empty((rows * row,), dtype=torch.uint8, device=self.device)
+        offsets = [rows * sum(feats[:b]) for b in range(len(feats))]
+        self.scratch_views = [
+            self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:])
+            for o, f, (_, bank) in zip(offsets, feats, cache.banks)
+        ]
+        self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device)
+        self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device)
+        self.iota = torch.arange(rows, dtype=torch.int32, device=self.device)
+        self._word_chunks = triton.cdiv(max(feats) // 4, COMPARE_WORDS)
+        logger.info(f"MoE slot audit scratch: {self.scratch.numel() / 2**20:.1f} MiB, {rows} experts a chunk")
+
+    # ------------------------------------------------------------------ recording (graph-capturable)
+
+    def begin_step(self) -> None:
+        """Advance the decode-step clock (the first GPU layer's decode ensure calls this)."""
+        _bump_step_kernel[(1,)](self.dstep)
+
+    def _launch(self, cache, kind: int, layer: int, *, meta: bool, width: int, dst=None, src=None, aux=None,
+                num=None, start: int = 0, row: int = ROW_PLAN, row_base: int = 0, pred: int = PRED_ALL) -> None:
+        if width <= 0 or not self.num_slots:
+            return
+        block = min(triton.next_power_of_2(width), 128)
+        d = self._dummy
+        _record_kernel[(triton.cdiv(width, block),)](
+            self.meta_hist if meta else self.bytes_hist, self.meta_cur if meta else self.bytes_cur,
+            self.seq, self.dstep, cache.step, cache.id_of_slot, cache.usage, self.totals,
+            d if dst is None else dst, d if src is None else src, d if aux is None else aux, d if num is None else num,
+            kind, layer, start, row_base, width, self.num_slots,
+            PLAN=dst is not None, ROW=row, HAS_AUX=aux is not None, HAS_NUM=num is not None, PRED=pred,
+            E=self.num_experts, BLOCK=block, num_warps=4,
+        )
+
+    def record_plan(self, cache, kind: int, layer: int, dst: torch.Tensor, src: torch.Tensor | None, num: torch.Tensor,
+                    *, meta: bool, width: int | None = None, aux: torch.Tensor | None = None, row_base: int = 0) -> None:
+        """The writer just wrote slots ``dst[:num]`` (rows ``src[:num]``, else ``dst - row_base``)."""
+        width = dst.numel() if width is None else min(width, dst.numel())
+        self._launch(cache, kind, layer, meta=meta, width=width, dst=dst, src=src, aux=aux, num=num,
+                     row=ROW_PLAN if src is not None else ROW_FROM_SLOT, row_base=row_base)
+
+    def record_range(self, cache, kind: int, layer: int, start: int, n: int, *, meta: bool) -> None:
+        """The writer just wrote expert ``i`` of ``layer`` into slot ``start + i`` for ``i < n``."""
+        self._launch(cache, kind, layer, meta=meta, width=n, start=start, row=ROW_INDEX)
+
+    def record_clear(self, cache, kind: int, context_layer: int, start: int, n: int, *, layer_above_e: bool = False) -> None:
+        """Before a writer clears held slots in ``[start, start + n)`` (with ``layer_above_e``: only
+        those >= E held by ``context_layer``); the event names the old owner."""
+        self._launch(cache, kind, context_layer, meta=True, width=n, start=start, row=ROW_OLD_OWNER,
+                     pred=PRED_LAYER_ABOVE_E if layer_above_e else PRED_HELD)
+
+    # ------------------------------------------------------------------ audit
+
+    def end_decode_step(self, cache) -> bool:
+        """After a decode forward, on the compute stream: audit every ``interval`` steps."""
+        self._steps += 1
+        if self._steps % self.interval:
+            return False
+        self.scan(cache)
+        return True
+
+    def scan(self, cache) -> None:
+        """Audit every held slot now (enqueued on the current stream; no host sync)."""
+        S, num_ids = self.num_slots, self.num_layers * self.num_experts
+        assert S == cache.cache_size and cache.id_of_slot.numel() == S, "bind after a rebuild"
+        block = 1024
+        _audit_meta_kernel[(triton.cdiv(max(S, num_ids), block),)](
+            cache.id_of_slot, cache.slot_for_id.view(-1), self.layer_ok, self.count, self.list_rows, self.list_slots,
+            self.flags, self.stale_id, self.stale_n, self.totals, S, num_ids,
+            E=self.num_experts, CAP=self.cap, BLOCK=block, num_warps=4,
+        )
+        n_lists = self.num_layers * self.chunks
+        _audit_nums_kernel[(1,)](self.count, self.nums, L=self.num_layers, C=self.chunks, R=self.rows,
+                                 BLOCK=triton.next_power_of_2(n_lists))
+        nb = len(self.bank_names)
+        R = self.rows
+        for layer in range(self.num_layers):
+            if not self.layer_ok_host[layer]:
+                continue
+            for c in range(self.chunks):
+                rows = self.list_rows[layer, c * R : (c + 1) * R]
+                slots = self.list_slots[layer, c * R : (c + 1) * R]
+                num = self.nums[layer * self.chunks + c : layer * self.chunks + c + 1]
+                self._gather(cache, layer, rows, num)
+                _audit_compare_kernel[(R, nb, self._word_chunks)](
+                    self.cache_ptrs, self.scratch_ptrs, self.words, slots, num, self.bad_cnt, self.bad_first,
+                    NB=nb, BLOCK=COMPARE_WORDS, num_warps=8,
+                )
+        rblock = 256
+        _audit_report_kernel[(triton.cdiv(S, rblock),)](
+            cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, self.dstep,
+            self.flags, self.stale_id, self.stale_n, self.bad_cnt, self.bad_first, self.reported,
+            self.meta_hist, self.meta_cur, self.bytes_hist, self.bytes_cur,
+            self.ring, self.cursor, self.totals, self.count, S, num_ids,
+            NB=nb, NB_P=triton.next_power_of_2(nb), L=self.num_layers, L_P=triton.next_power_of_2(self.num_layers),
+            RING_N=self.ring_n, BLOCK=rblock, num_warps=4,
+        )
+
+    def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> None:
+        # the plain fused copy, not the slim one prefetch copies use: a slim-kernel bug cannot hide itself
+        if cache._copy_fused_ok:
+            from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
+
+            fast_index_copy_multi_jit(
+                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, rows, num,
+            )
+            return
+        from freetoken.kernel import fast_index_copy_jit
+
+        for (per_layer, _), view in zip(cache.banks, self.scratch_views):
+            fast_index_copy_jit(view, self.iota, per_layer[layer], rows, num)
+
+    # ------------------------------------------------------------------ reporting (host syncs)
+
+    def take_window(self) -> tuple[list[int], list[int], list[list[int]], int]:
+        """(totals now, totals of this window, records kept since the last call, records in all)."""
+        totals = self.totals.cpu()
+        window = (totals - self._last).tolist()
+        self._last = totals
+        cursor = int(self.cursor.item())
+        upto = min(cursor, self.ring_n)
+        records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else []
+        self._printed = max(self._printed, upto)
+        return totals.tolist(), window, records, cursor
+
+    def report_window(self, steps: int) -> list[tuple[str, str]]:
+        """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no audit ran."""
+        totals, window, records, cursor = self.take_window()
+        if not window[T_AUDITS] and not records:
+            return []
+        bad = bool(window[T_BAD] or window[T_MAP_RANGE] or window[T_BAD_PLAN] or records)
+        lines = [("warning" if bad else "info",
+                  f"MoE slot audit ({steps} decode steps, every {self.interval}): {self.format_counts(window)}; "
+                  f"session: {totals[T_AUDITS]} audits, {totals[T_NEW]} bad (slot, owner) pairs")]
+        lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records]
+        if records and cursor > self.ring_n:
+            lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept"))
+        return lines
+
+    def report_session(self) -> list[tuple[str, str]]:
+        """Session totals and every kept record (shutdown)."""
+        totals = self.totals.cpu().tolist()
+        cursor = int(self.cursor.item())
+        if not totals[T_AUDITS]:
+            return []
+        bad = bool(totals[T_BAD] or totals[T_MAP_RANGE] or totals[T_BAD_PLAN])
+        lines = [("warning" if bad else "info", f"MoE slot audit (session, every {self.interval} decode steps): "
+                                                  f"{self.format_counts(totals)}")]
+        records = self.ring[: min(cursor, self.ring_n)].cpu().tolist()
+        lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records]
+        if cursor > self.ring_n:
+            lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept"))
+        return lines
+
+    @staticmethod
+    def format_counts(t: list[int]) -> str:
+        audits = t[T_AUDITS]
+        per = (lambda v: f"{v / audits:.0f}") if audits else str
+        text = (
+            f"{audits} audits, {per(t[T_HELD])} held slots an audit ({per(t[T_BYTES_CHECKED])} byte-checked), "
+            f"bad slots={t[T_BAD]} (bytes={t[T_BYTES]}, map={t[T_FWD]}, stale slot_for_id={t[T_STALE]}, "
+            f"owner out of range={t[T_OWNER_RANGE]}), new={t[T_NEW]}, repeat={t[T_STILL]}, events={t[T_EVENTS]}"
+        )
+        if t[T_UNCHECKED]:
+            text += f", held slots not byte-checked={t[T_UNCHECKED]}"
+        if t[T_MAP_RANGE]:
+            text += f", slot_for_id entries past the cache={t[T_MAP_RANGE]}"
+        if t[T_BAD_PLAN]:
+            text += f", plan entries naming no slot={t[T_BAD_PLAN]}"
+        return text
+
+    # ------------------------------------------------------------------ record decoding
+
+    def _name(self, flat: int) -> str:
+        e = self.num_experts
+        return f"L{flat // e}/e{flat % e}" if flat >= 0 else str(flat)
+
+    @staticmethod
+    def events(rec: list[int]) -> list[dict]:
+        """Both rings of a record as event dicts in ``seq`` order (``ring`` 'meta' or 'bytes')."""
+        out = []
+        for ring, base in (("meta", HEADER), ("bytes", HEADER + RING_EVENTS)):
+            for d in range(DEPTH):
+                ev = rec[base + d * EVENT_FIELDS : base + (d + 1) * EVENT_FIELDS]
+                if ev[E_SEQ] >= 0:
+                    out.append({
+                        "ring": ring, "seq": ev[E_SEQ], "kind": ev[E_KIND], "dstep": ev[E_DSTEP], "lstep": ev[E_LSTEP],
+                        "layer": ev[E_LAYER], "row": ev[E_ROW], "aux": ev[E_AUX], "owner": ev[E_OWNER],
+                        "usage": ev[E_USAGE], "pos": ev[E_POS],
+                    })
+        return sorted(out, key=lambda ev: ev["seq"])
+
+    def _flat(self, ev: dict) -> int:
+        return ev["layer"] * self.num_experts + ev["row"] if ev["layer"] >= 0 and ev["row"] >= 0 else -1
+
+    def format_event(self, ev: dict) -> str:
+        e = self.num_experts
+        text = f"#{ev['seq']} {KIND_NAMES.get(ev['kind'], ev['kind'])} L{ev['layer']}/e{ev['row']}"
+        if ev["kind"] == PREFILL_HIT_D2D:
+            text += f" from slot {ev['aux']}"
+        elif ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE) and ev["aux"] >= 0:
+            text += f" (for layer {ev['aux']})"
+        seen = "slot held" if ev["ring"] == "bytes" else "then held"
+        if ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE, RESET):
+            seen = "held"
+        owner = f"L{ev['owner'] // e}/e{ev['owner'] % e}" if ev["owner"] >= 0 else "nothing"
+        return f"{text} @step {ev['dstep']}/lru {ev['lstep']} ({seen} {owner}, usage {ev['usage']})"
+
+    def _kind(self, ev: dict) -> str:
+        return KIND_NAMES.get(ev["kind"], str(ev["kind"]))
+
+    def _pending_install(self, copy: dict, meta: list[dict], byte: list[dict], before: dict) -> dict | None:
+        """The earlier install (before ``before``) that ``copy`` completes: an install of the id it
+        copied with no copy of that id between them."""
+        flat = self._flat(copy)
+        for inst in reversed([ev for ev in meta if ev["kind"] in INSTALLS and ev["seq"] < before["seq"]]):
+            if self._flat(inst) == flat:
+                done = any(self._flat(ev) == flat and inst["seq"] < ev["seq"] < copy["seq"] for ev in byte)
+                return None if done else inst
+        return None
+
+    def _bytes_story(self, owner: int, meta: list[dict], byte: list[dict]) -> str:
+        insts = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == owner]
+        if not insts:
+            return (f"no recorded install of {self._name(owner)} in the kept history "
+                    f"(an unrecorded map writer, or it is older than the last {DEPTH} events)")
+        inst = insts[-1]
+        name = self._kind(inst)
+        after = [ev for ev in byte if ev["seq"] > inst["seq"]]
+        own = [ev for ev in after if self._flat(ev) == owner]
+        if own:
+            last = after[-1]
+            if self._flat(last) == owner:
+                return (f"{name} and {self._kind(last)} of {self._name(owner)} (#{last['seq']}), then the bytes changed "
+                        f"with no recorded write: overwritten after the last recorded write, or that copy did not land")
+            text = (f"{name} and {self._kind(own[-1])} of {self._name(owner)} (#{own[-1]['seq']}), then "
+                    f"{self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}) overwrote it")
+            prior = self._pending_install(last, meta, byte, inst)
+            if prior is not None:
+                text += f": the late copy of the earlier {self._kind(prior)} of {self._name(self._flat(prior))} (#{prior['seq']})"
+            return text
+        text = f"{name}, no {KIND_NAMES[COPY_FOR[inst['kind']]]}"
+        if not after:
+            return text
+        last = after[-1]
+        prior = self._pending_install(last, meta, byte, inst)
+        if prior is not None:
+            return (f"{text}; the last write, {self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}), "
+                    f"is the copy of the earlier {self._kind(prior)} (#{prior['seq']}) and landed after this install")
+        where = (f"from layer {last['layer']}'s source" if last["row"] == owner % self.num_experts
+                 else f"of {self._name(self._flat(last))}")
+        return f"{name} of {self._name(owner)}, then {self._kind(last)} {where} ({self._name(self._flat(last))}, #{last['seq']})"
+
+    def diagnose(self, rec: list[int]) -> str:
+        """What the history says about a bad slot, in words."""
+        owner, fl = rec[H_OWNER], rec[H_FLAGS]
+        evs = self.events(rec)
+        meta = [ev for ev in evs if ev["ring"] == "meta"]
+        byte = [ev for ev in evs if ev["ring"] == "bytes"]
+        parts = []
+        if fl & F_OWNER_RANGE:
+            parts.append(f"id_of_slot={owner} names no expert")
+        if fl & F_BYTES and owner >= 0:
+            parts.append(self._bytes_story(owner, meta, byte))
+        if fl & F_FWD:
+            parts.append(f"slot_for_id[{self._name(owner)}]={rec[H_MAPPED]}, not this slot")
+        if fl & F_STALE:
+            sid = rec[H_STALE_ID]
+            text = (f"stale slot_for_id: {self._name(sid)} -> this slot, which holds {self._name(owner)}"
+                    + (f" ({rec[H_STALE_N]} such ids)" if rec[H_STALE_N] > 1 else ""))
+            gone = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == sid]
+            if gone:
+                later = [ev for ev in meta if ev["seq"] > gone[-1]["seq"]]
+                text += (f"; {self._name(sid)} was installed here by {self._kind(gone[-1])} #{gone[-1]['seq']}"
+                         + (f" and replaced by {self._kind(later[0])} #{later[0]['seq']} without clearing "
+                            f"its map entry" if later else ""))
+            parts.append(text)
+        # an install's recorder reads the slot right after the install on the same stream
+        for ev in meta:
+            if ev["kind"] in INSTALLS and ev["owner"] != self._flat(ev):
+                parts.append(f"right after the {self._kind(ev)} of {self._name(self._flat(ev))} (#{ev['seq']}) the slot "
+                             f"named {self._name(ev['owner'])}: another map writer ran concurrently")
+        return "; ".join(parts) or "clean"
+
+    def format_record(self, rec: list[int]) -> str:
+        fl = rec[H_FLAGS]
+        text = (f"audit={rec[H_AUDIT]} step={rec[H_DSTEP]} lru={rec[H_LSTEP]} slot={rec[H_SLOT]} "
+                f"holds {self._name(rec[H_OWNER])} (slot_for_id={rec[H_MAPPED]}, usage={rec[H_USAGE]})")
+        if fl & F_BYTES:
+            banks = [self.bank_names[b] if b < len(self.bank_names) else str(b)
+                     for b in range(max(len(self.bank_names), 1)) if rec[H_BANKS] >> b & 1]
+            first = self.bank_names[rec[H_FIRST_BANK]] if 0 <= rec[H_FIRST_BANK] < len(self.bank_names) else rec[H_FIRST_BANK]
+            text += (f" | bytes differ in {', '.join(banks)} (first at {first}+{rec[H_FIRST_BYTE]}, "
+                     f"{rec[H_BAD_WORDS]} bad words)")
+        text += f" | diagnosis: {self.diagnose(rec)}"
+        history = "; ".join(self.format_event(ev) for ev in self.events(rec))
+        return text + f" | history ({rec[H_META_N]} map / {rec[H_BYTES_N]} byte events in all): {history or 'none'}"
diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py
new file mode 100644
index 0000000000..1f096c9c3a
--- /dev/null
+++ b/python/freetoken/moe/verify.py
@@ -0,0 +1,516 @@
+"""Decode expert-cache verifier (``FREETOKEN_MOE_PREFETCH_VERIFY``): a debug detector for slot races.
+
+For every GPU decode layer call it checks, on the compute stream and in its order, what the
+grouped GEMM is about to read and has just read:
+
+- meta (before and after the GEMM): each routed entry's slot maps to the expert the router chose,
+  in both slot maps (``id_of_slot[s] == layer * E + id`` and ``slot_for_id[layer, id] == s``); the
+  router's ids are copied before ensure rewrites them in place;
+- bytes before the GEMM (``full``): every routed slot's row in every bank equals the host bank row
+  of that expert for that layer, gathered into a scratch buffer with the fused copy;
+- bytes after the GEMM (``full``): the same slots against the same scratch, so a slot overwritten
+  while the GEMM ran shows up even when a later copy restores it.
+
+Nothing here records or waits on an event: the prefetch streams keep exactly the edges they have
+without it. The checks do delay the compute stream (``full`` gathers every routed expert over
+PCIe), which moves where the other streams land relative to it. ``meta`` skips the byte checks and
+costs a few one-CTA launches per layer. Counts accumulate per layer on the device and the first
+``RING`` offending records are kept with the slot state seen at check time; the engine reports
+them every MOE_STATS_INTERVAL decode steps and at shutdown.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+from freetoken.env import ENV
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
+MODES = ("off", "meta", "full")
+# counters[layer, col]: checked calls, offending records per check, routed experts left unchecked
+CHECKS, META_BAD, PRE_BAD, POST_BAD, POST_META_BAD, UNCHECKED = range(6)
+NUM_COLS = 6
+COL_NAMES = ("checks", "meta_bad", "pre_bad", "post_bad", "post_meta_bad", "unchecked")
+META_PRE, BYTES_PRE, META_POST, BYTES_POST = range(4)
+KINDS = ("meta_pre", "bytes_pre", "meta_post", "bytes_post")
+# record fields; owner / mapped / usage / step are read when the check runs
+(R_CALL, R_LAYER, R_ROW, R_EXPERT, R_SLOT, R_BANK, R_KIND, R_OWNER, R_MAPPED, R_USAGE, R_STEP,
+ R_SOURCE, R_FIRST_BYTE, R_BAD_WORDS) = range(14)
+NUM_FIELDS = 14
+# R_SOURCE bits: the slot is in this call's demand copy plan, this layer's prefetch copy plan, or
+# the next layer's prefetch plan (installed earlier on the compute stream)
+SRC_DEMAND, SRC_PREFETCH, SRC_NEXT_PREFETCH = 1, 2, 4
+RING = 64
+# routed entries a call may have (bs * top_k); wider eager batches are counted unchecked
+MAX_ENTRIES = 64
+CHUNK_WORDS = 4096
+_NF = tl.constexpr(NUM_FIELDS)
+_NC = tl.constexpr(NUM_COLS)
+_NO_BAD = tl.constexpr(0x7FFFFFFF)
+
+
+def resolve_mode(mode: str | None = None) -> str:
+    raw = (ENV.MOE_PREFETCH_VERIFY.value if mode is None else mode).strip().lower()
+    if raw in ("", "0", "off", "false", "no"):
+        return "off"
+    if raw in ("1", "on", "true", "yes", "full"):
+        return "full"
+    if raw == "meta":
+        return "meta"
+    raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={raw!r}: expected 0, 1 (full) or meta")
+
+
+@triton.jit
+def _put(vals, f, idx: tl.constexpr, v):
+    return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals)
+
+
+@triton.jit
+def _in_plan(s, plan_ptr, num_ptr, W: tl.constexpr, W_P: tl.constexpr):
+    j = tl.arange(0, W_P)
+    plan = tl.load(plan_ptr + j, mask=j < tl.minimum(tl.load(num_ptr), W), other=-3)
+    return tl.sum((s[:, None] == plan[None, :]).to(tl.int32), axis=1) > 0
+
+
+@triton.jit
+def _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+            DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr,
+            HAS_PF: tl.constexpr, HAS_NX: tl.constexpr):
+    src = tl.where(_in_plan(s, dem_ptr, dem_num_ptr, DEM_W, DEM_P), 1, 0)
+    if HAS_PF:
+        src += tl.where(_in_plan(s, pf_ptr, pf_num_ptr, PF_W, PF_P), 2, 0)
+    if HAS_NX:
+        src += tl.where(_in_plan(s, nx_ptr, nx_num_ptr, PF_W, PF_P), 4, 0)
+    return src
+
+
+@triton.jit
+def _write(ring_ptr, cursor, bad, vals, RING_N: tl.constexpr):
+    """Append the ``bad`` rows of ``vals`` after ``cursor`` while the ring has room; the new cursor."""
+    f = tl.arange(0, 16)
+    pos = cursor + tl.cumsum(bad.to(tl.int64), axis=0) - 1
+    keep = bad & (pos < RING_N)
+    tl.store(ring_ptr + pos[:, None] * _NF + f[None, :], vals, mask=keep[:, None] & (f[None, :] < _NF))
+    return cursor + tl.sum(bad.to(tl.int64), axis=0)
+
+
+@triton.jit
+def _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr,
+          dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+          ring_ptr, cursor, counters_ptr, col: tl.constexpr, kind: tl.constexpr, call, layer, num_slots,
+          N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr,
+          DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr,
+          HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, RING_N: tl.constexpr):
+    """Check every routed entry's slot against both maps; returns the new ring cursor."""
+    i = tl.arange(0, N_P)
+    m = i < N
+    ids = tl.load(ids_ptr + i, mask=m, other=0)
+    s = tl.load(slots_ptr + i, mask=m, other=-1)
+    ok_s = m & (s >= 0) & (s < num_slots)
+    ok_id = m & (ids >= 0) & (ids < E)
+    fid = layer * E + ids
+    owner = tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2)
+    mapped = tl.load(slot_for_id_ptr + fid, mask=ok_id, other=-2)
+    bad = m & ~(ok_s & ok_id & (owner == fid) & (mapped == s))
+    usage = tl.load(usage_ptr + s, mask=ok_s, other=-1)
+    src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX)
+    f = tl.arange(0, 16)
+    z = tl.zeros([N_P], tl.int64)
+    vals = tl.zeros([N_P, 16], tl.int64) - 1
+    vals = _put(vals, f, 0, z + call)
+    vals = _put(vals, f, 1, z + layer)
+    vals = _put(vals, f, 2, i)
+    vals = _put(vals, f, 3, ids)
+    vals = _put(vals, f, 4, s)
+    vals = _put(vals, f, 6, z + kind)
+    vals = _put(vals, f, 7, owner)
+    vals = _put(vals, f, 8, mapped)
+    vals = _put(vals, f, 9, usage)
+    vals = _put(vals, f, 10, z + tl.load(step_ptr))
+    vals = _put(vals, f, 11, src)
+    nbad = tl.sum(bad.to(tl.int64), axis=0)
+    tl.store(counters_ptr + col, tl.load(counters_ptr + col) + nbad)
+    return _write(ring_ptr, cursor, bad, vals, RING_N)
+
+
+@triton.jit(do_not_specialize=["layer", "num_slots"])
+def _plan_kernel(
+    ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr,
+    dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+    ring_ptr, cursor_ptr, counters_ptr, calls_ptr,
+    u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, slots_copy_ptr,
+    layer, num_slots,
+    N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr,
+    DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr,
+    HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, RING_N: tl.constexpr,
+):
+    """Before the GEMM: count the call, check the maps, keep the slots for the after-GEMM checks and
+    list the distinct (expert, slot) pairs to byte-check (first CAP of them)."""
+    call = tl.load(calls_ptr + layer)
+    tl.store(calls_ptr + layer, call + 1)
+    row = counters_ptr + layer * _NC
+    tl.store(row + 0, tl.load(row + 0) + 1)
+    cursor = tl.load(cursor_ptr)
+    cursor = _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr,
+                   dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+                   ring_ptr, cursor, row, 1, 0, call, layer, num_slots,
+                   N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N)
+    tl.store(cursor_ptr, cursor)
+    i = tl.arange(0, N_P)
+    m = i < N
+    ids = tl.load(ids_ptr + i, mask=m, other=-1)
+    s = tl.load(slots_ptr + i, mask=m, other=-1)
+    tl.store(slots_copy_ptr + i, s, mask=m)
+    if BYTES:
+        # an unreadable slot or id is already a meta record; a duplicate pair shares its first row
+        ok = m & (s >= 0) & (s < num_slots) & (ids >= 0) & (ids < E)
+        same = (ids[:, None] == ids[None, :]) & (s[:, None] == s[None, :]) & (i[None, :] < i[:, None]) & ok[None, :]
+        first = ok & (tl.sum(same.to(tl.int32), axis=1) == 0)
+        rank = tl.cumsum(first.to(tl.int32), axis=0) - 1
+        take = first & (rank < CAP)
+        tl.store(u_ids_ptr + rank, ids, mask=take)
+        tl.store(u_slots_ptr + rank, s, mask=take)
+        tl.store(u_rows_ptr + rank, i, mask=take)
+        total = tl.sum(first.to(tl.int64), axis=0)
+        kept = tl.minimum(total, CAP)
+        tl.store(u_num_ptr, kept)
+        tl.store(row + 5, tl.load(row + 5) + total - kept)
+
+
+@triton.jit
+def _compare_kernel(cache_ptrs, ref_ptrs, words_ptr, u_slots_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr,
+                    NB: tl.constexpr, BLOCK: tl.constexpr):
+    """Program (row, bank, chunk): compare one chunk of a routed slot's row with its host copy."""
+    r = tl.program_id(0)
+    b = tl.program_id(1)
+    words = tl.load(words_ptr + b)
+    off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    m = (r < tl.load(u_num_ptr)) & (off < words)
+    slot = tl.load(u_slots_ptr + r).to(tl.int64)
+    got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0)
+    want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0)
+    bad = m & (got != want)
+    nbad = tl.sum(bad.to(tl.int32), axis=0)
+    if nbad > 0:
+        tl.atomic_add(bad_cnt_ptr + r * NB + b, nbad)
+        tl.atomic_min(bad_first_ptr + r * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32))
+
+
+@triton.jit(do_not_specialize=["layer", "num_slots"])
+def _tally_kernel(
+    ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr,
+    dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+    ring_ptr, cursor_ptr, counters_ptr, calls_ptr,
+    u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr,
+    layer, num_slots,
+    N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, CAP_P: tl.constexpr,
+    NB: tl.constexpr, NB_P: tl.constexpr,
+    DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr,
+    HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, POST: tl.constexpr,
+    BYTES_KIND: tl.constexpr, BYTES_COL: tl.constexpr, RING_N: tl.constexpr,
+):
+    """Turn the compare flags into counts and records (and clear them); after the GEMM also recheck
+    the maps of the slots it read."""
+    call = tl.load(calls_ptr + layer) - 1
+    row = counters_ptr + layer * _NC
+    cursor = tl.load(cursor_ptr)
+    if POST:
+        cursor = _meta(ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr,
+                       dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr,
+                       ring_ptr, cursor, row, 4, 2, call, layer, num_slots,
+                       N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N)
+    if BYTES:
+        j = tl.arange(0, CAP_P * NB_P)
+        r = j // NB_P
+        b = j % NB_P
+        cell = (r < CAP) & (b < NB)
+        live = cell & (r < tl.load(u_num_ptr))
+        cnt = tl.load(bad_cnt_ptr + r * NB + b, mask=live, other=0)
+        first = tl.load(bad_first_ptr + r * NB + b, mask=live, other=_NO_BAD)
+        tl.store(bad_cnt_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32), mask=cell)
+        tl.store(bad_first_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32) + _NO_BAD, mask=cell)
+        bad = live & (cnt > 0)
+        ids = tl.load(u_ids_ptr + r, mask=live, other=0)
+        s = tl.load(u_slots_ptr + r, mask=live, other=-1)
+        ok_s = live & (s >= 0) & (s < num_slots)
+        src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX)
+        f = tl.arange(0, 16)
+        z = tl.zeros([CAP_P * NB_P], tl.int64)
+        vals = tl.zeros([CAP_P * NB_P, 16], tl.int64) - 1
+        vals = _put(vals, f, 0, z + call)
+        vals = _put(vals, f, 1, z + layer)
+        vals = _put(vals, f, 2, tl.load(u_rows_ptr + r, mask=live, other=-1))
+        vals = _put(vals, f, 3, ids)
+        vals = _put(vals, f, 4, s)
+        vals = _put(vals, f, 5, b)
+        vals = _put(vals, f, 6, z + BYTES_KIND)
+        vals = _put(vals, f, 7, tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2))
+        vals = _put(vals, f, 8, tl.load(slot_for_id_ptr + layer * E + ids, mask=live, other=-2))
+        vals = _put(vals, f, 9, tl.load(usage_ptr + s, mask=ok_s, other=-1))
+        vals = _put(vals, f, 10, z + tl.load(step_ptr))
+        vals = _put(vals, f, 11, src)
+        vals = _put(vals, f, 12, first.to(tl.int64) * 4)
+        vals = _put(vals, f, 13, cnt)
+        tl.store(row + BYTES_COL, tl.load(row + BYTES_COL) + tl.sum(bad.to(tl.int64), axis=0))
+        cursor = _write(ring_ptr, cursor, bad, vals, RING_N)
+    tl.store(cursor_ptr, cursor)
+
+
+class ExpertVerifier:
+    """Device counters, the record ring and (``full``) the scratch the byte checks compare against.
+
+    ``begin`` copies the router's ids before ensure rewrites them, ``before_gemm`` runs after the
+    joins and the demand copy, ``after_gemm`` right after the GEMM; all three enqueue on the current
+    (compute) stream with fixed shapes and no host sync, so they capture into the decode graph.
+    ``bind`` sizes the scratch from the banks once and refreshes the slot bank addresses in place
+    (a rebuild reallocates the banks); counters survive rebuilds and resets.
+    """
+
+    def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, mode: str = "full",
+                 rows: int | None = None) -> None:
+        from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled
+
+        if mode == "full" and _skip_fast_index_copy_enabled():
+            logger.warning(
+                "FREETOKEN_MOE_PREFETCH_VERIFY=1 with FREETOKEN_SKIP_FAST_INDEX_COPY: the host rows cannot be "
+                "gathered, so only the slot maps are checked (meta)"
+            )
+            mode = "meta"
+        self.mode = mode
+        self.num_layers = num_layers
+        self.num_experts = num_experts
+        self.device = device
+        self.rows = ENV.MOE_PREFETCH_VERIFY_ROWS.value if rows is None else int(rows)
+        if self.rows < 1:
+            raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY_ROWS={self.rows} must be >= 1")
+        self.counters = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64, device=device)
+        self.totals = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64)
+        self.calls = torch.zeros((num_layers,), dtype=torch.int64, device=device)
+        self.ring = torch.full((RING, NUM_FIELDS), -1, dtype=torch.int64, device=device)
+        self.cursor = torch.zeros((1,), dtype=torch.int64, device=device)
+        self.ids = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device)
+        self.slots = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device)
+        self.u_ids = torch.zeros((self.rows,), dtype=torch.int32, device=device)
+        self.u_slots = torch.zeros((self.rows,), dtype=torch.int32, device=device)
+        self.u_rows = torch.zeros((self.rows,), dtype=torch.int32, device=device)
+        self.u_num = torch.zeros((1,), dtype=torch.int64, device=device)
+        self.iota = torch.arange(self.rows, dtype=torch.int32, device=device)
+        self.bank_names: list[str] = []
+        self.scratch: torch.Tensor | None = None
+        self._printed = 0
+        self._top_k = 1
+        self._call = None  # (layer, entries, this layer's prefetch plan or None) between begin and after_gemm
+        self._next = None
+        logger.info(
+            f"MoE verify {self.mode}: checking every GPU decode layer call's slot maps"
+            + (f" and slot bytes before and after the GEMM (up to {self.rows} experts a call)" if self.mode == "full" else "")
+        )
+
+    def bind(self, cache) -> None:
+        """Size the scratch for ``cache``'s banks (first call) or point the checks at its new banks."""
+        caches = [bank for _, bank in cache.banks]
+        feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches]
+        if self.bank_names:
+            assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed"
+        else:
+            assert all(f % 4 == 0 for f in feats), feats
+            self.bank_names = list(cache.bank_schema)
+            self._feats = feats
+            self.cache_ptrs = torch.zeros((len(feats),), dtype=torch.int64, device=self.device)
+            if self.mode == "full":
+                self._alloc_scratch(cache, feats)
+        self.cache_ptrs.copy_(torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64))
+
+    def _alloc_scratch(self, cache, feats: list[int]) -> None:
+        rows = self.rows
+        self.scratch = torch.empty((rows * sum(feats),), dtype=torch.uint8, device=self.device)
+        offsets = [rows * sum(feats[:b]) for b in range(len(feats))]
+        self.scratch_views = [
+            self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:])
+            for o, f, (_, bank) in zip(offsets, feats, cache.banks)
+        ]
+        self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device)
+        self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device)
+        self.bad_cnt = torch.zeros((rows * len(feats),), dtype=torch.int32, device=self.device)
+        self.bad_first = torch.full((rows * len(feats),), 0x7FFFFFFF, dtype=torch.int32, device=self.device)
+        self._chunks = triton.cdiv(max(feats) // 4, CHUNK_WORDS)
+        logger.info(f"MoE verify scratch: {self.scratch.numel() / 2**20:.1f} MiB for {rows} experts")
+
+    def begin(self, cache, layer_id: int, topk_ids: torch.Tensor, plan: tuple | None) -> None:
+        """Copy the router's ids of ``layer_id`` (``plan``: this layer's prefetch copy plan, if forked)."""
+        assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous()
+        n = topk_ids.numel()
+        self._top_k = topk_ids.shape[-1]
+        if n > MAX_ENTRIES:
+            self.counters[layer_id, UNCHECKED] += n
+            self._call = None
+            return
+        self.ids[:n].copy_(topk_ids.view(-1))
+        self._call = (layer_id, n, plan)
+
+    def _plans(self, cache) -> tuple:
+        _, n, plan = self._call
+        pf, nx = plan, self._next
+        dummy = (self.u_slots, self.u_num)
+        width = (pf or nx or dummy)[0].numel()
+        return (
+            cache.evict_slots, cache.num_indices, *(pf or dummy), *(nx or dummy),
+            dict(DEM_W=n, DEM_P=triton.next_power_of_2(n), PF_W=width, PF_P=triton.next_power_of_2(width),
+                 HAS_PF=pf is not None, HAS_NX=nx is not None),
+        )
+
+    def before_gemm(self, cache, slots: torch.Tensor, next_plan: tuple | None) -> None:
+        """Check the maps and (``full``) the bytes of the slots ``slots`` the GEMM is about to read
+        (``next_plan``: the next layer's prefetch plan, if one was issued)."""
+        if self._call is None:
+            return
+        layer, n, _ = self._call
+        self._next = next_plan
+        *plans, plan_kw = self._plans(cache)
+        full = self.mode == "full"
+        _plan_kernel[(1,)](
+            self.ids, slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans,
+            self.ring, self.cursor, self.counters, self.calls,
+            self.u_ids, self.u_slots, self.u_rows, self.u_num, self.slots,
+            layer, cache.cache_size,
+            N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows,
+            BYTES=full, RING_N=RING, num_warps=4, **plan_kw,
+        )
+        if not full:
+            return
+        self._gather(cache, layer)
+        self._compare_and_tally(cache, post=False)
+
+    def after_gemm(self, cache) -> None:
+        """Recheck the maps and (``full``) the bytes of the slots the GEMM just read."""
+        if self._call is None:
+            return
+        if self.mode == "full":
+            self._compare_and_tally(cache, post=True)
+        else:
+            self._tally(cache, post=True, bytes_=False)
+        self._call = self._next = None
+
+    def _gather(self, cache, layer: int) -> None:
+        if cache._copy_fused_ok:
+            from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
+
+            fast_index_copy_multi_jit(
+                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, self.u_ids, self.u_num,
+            )
+            return
+        from freetoken.kernel import fast_index_copy_jit
+
+        for (per_layer, _), view in zip(cache.banks, self.scratch_views):
+            fast_index_copy_jit(view, self.iota, per_layer[layer], self.u_ids, self.u_num)
+
+    def _compare_and_tally(self, cache, *, post: bool) -> None:
+        nb = len(self.bank_names)
+        _compare_kernel[(self.rows, nb, self._chunks)](
+            self.cache_ptrs, self.scratch_ptrs, self.words, self.u_slots, self.u_num, self.bad_cnt, self.bad_first,
+            NB=nb, BLOCK=CHUNK_WORDS, num_warps=8,
+        )
+        self._tally(cache, post=post, bytes_=True)
+
+    def _tally(self, cache, *, post: bool, bytes_: bool) -> None:
+        layer, n, _ = self._call
+        *plans, plan_kw = self._plans(cache)
+        nb = max(len(self.bank_names), 1)
+        _tally_kernel[(1,)](
+            self.ids, self.slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans,
+            self.ring, self.cursor, self.counters, self.calls,
+            self.u_ids, self.u_slots, self.u_rows, self.u_num,
+            self.bad_cnt if bytes_ else self.u_rows, self.bad_first if bytes_ else self.u_rows,
+            layer, cache.cache_size,
+            N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows,
+            CAP_P=triton.next_power_of_2(self.rows), NB=nb, NB_P=triton.next_power_of_2(nb),
+            BYTES=bytes_, POST=post, BYTES_KIND=BYTES_POST if post else BYTES_PRE, BYTES_COL=POST_BAD if post else PRE_BAD,
+            RING_N=RING, num_warps=4, **plan_kw,
+        )
+
+    # ------------------------------------------------------------------ reporting (host syncs)
+
+    def take_window(self) -> tuple[torch.Tensor, list[list[int]], int]:
+        """This window's ``[num_layers, NUM_COLS]`` counts (added to ``totals``; the device ones
+        restart), the records kept since the last call, and how many records there were in all."""
+        window = self.counters.cpu()
+        self.counters.zero_()
+        self.totals += window
+        cursor = int(self.cursor.item())
+        upto = min(cursor, RING)
+        records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else []
+        self._printed = max(self._printed, upto)
+        return window, records, cursor
+
+    def report_window(self, steps: int) -> list[tuple[str, str]]:
+        """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no call was checked."""
+        window, records, cursor = self.take_window()
+        if not int(window[:, CHECKS].sum()) and not records:
+            return []
+        bad = self._bad(window) or bool(records)
+        session = self._bad(self.totals)
+        lines = [("warning" if bad else "info",
+                  f"MoE verify {self.mode} ({steps} decode steps): {self.format_counts(window)}; session bad={session}")]
+        lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records]
+        if bad and cursor > RING:
+            lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept"))
+        return lines
+
+    def report_session(self) -> list[tuple[str, str]]:
+        """Session counts and every kept record (shutdown)."""
+        totals = self.totals + self.counters.cpu()
+        cursor = int(self.cursor.item())
+        if not int(totals[:, CHECKS].sum()):
+            return []
+        bad = self._bad(totals)
+        lines = [("warning" if bad else "info", f"MoE verify {self.mode} (session): {self.format_counts(totals)}")]
+        records = self.ring[: min(cursor, RING)].cpu().tolist()
+        lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records]
+        if cursor > RING:
+            lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept"))
+        return lines
+
+    @staticmethod
+    def _bad(counts: torch.Tensor) -> int:
+        return int(counts[:, [META_BAD, PRE_BAD, POST_BAD, POST_META_BAD]].sum())
+
+    def format_counts(self, counts: torch.Tensor) -> str:
+        total = counts.sum(0).tolist()
+        text = f"{total[CHECKS]} layer calls checked, " + ", ".join(
+            f"{COL_NAMES[c]}={total[c]}" for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)
+        )
+        if total[UNCHECKED]:
+            text += f", unchecked experts={total[UNCHECKED]}"
+        per_layer = [
+            f"L{layer}=" + "/".join(str(row[c]) for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD))
+            for layer, row in enumerate(counts.tolist())
+            if any(row[c] for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD))
+        ]
+        if per_layer:
+            text += " (meta/pre/post/post_meta by layer: " + ", ".join(per_layer) + ")"
+        return text
+
+    def format_record(self, rec: list[int]) -> str:
+        e = self.num_experts
+        kind = KINDS[rec[R_KIND]] if 0 <= rec[R_KIND] < len(KINDS) else str(rec[R_KIND])
+        owner = rec[R_OWNER]
+        owner_text = f"L{owner // e}/e{owner % e}" if owner >= 0 else str(owner)
+        src = [name for bit, name in ((SRC_DEMAND, "demand copy"), (SRC_PREFETCH, "prefetch copy"),
+                                       (SRC_NEXT_PREFETCH, "next layer's prefetch plan")) if rec[R_SOURCE] & bit]
+        row = rec[R_ROW]
+        text = (
+            f"{kind} call={rec[R_CALL]} layer={rec[R_LAYER]} row={row} (token {row // self._top_k}, rank {row % self._top_k}) "
+            f"expert={rec[R_EXPERT]} slot={rec[R_SLOT]}"
+        )
+        if rec[R_BANK] >= 0:
+            bank = self.bank_names[rec[R_BANK]] if rec[R_BANK] < len(self.bank_names) else rec[R_BANK]
+            text += f" bank={bank} first_bad_byte={rec[R_FIRST_BYTE]} bad_words={rec[R_BAD_WORDS]}"
+        return text + (
+            f" | slot now holds {owner_text}, slot_for_id={rec[R_MAPPED]}, usage={rec[R_USAGE]}, lru_step={rec[R_STEP]}, "
+            f"slot in: {', '.join(src) or 'no copy plan (a resident hit)'}"
+        )
diff --git a/tests/models/qwen4_exp/conftest.py b/tests/models/qwen4_exp/conftest.py
index 2b9e80214f..4e2bb67bb8 100644
--- a/tests/models/qwen4_exp/conftest.py
+++ b/tests/models/qwen4_exp/conftest.py
@@ -13,3 +13,16 @@ def _runtime():
     core._GLOBAL_CTX = None
     yield
     core._GLOBAL_CTX = None
+    import torch
+
+    if torch.cuda.is_available() and torch.cuda.is_initialized():
+        import gc
+
+        # torch keeps a cuBLAS workspace per stream for good: each offload cache's dedicated
+        # streams would otherwise pin ~16 MiB for the rest of the session
+        gc.collect()
+        torch.cuda.synchronize()
+        clear = getattr(torch._C, "_cuda_clearCublasWorkspaces", None)  # private torch API
+        if clear is not None:
+            clear()
+        torch.cuda.empty_cache()
diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py
index 58ad0c4dc3..85366638a5 100644
--- a/tests/models/qwen4_exp/test_build_cpu.py
+++ b/tests/models/qwen4_exp/test_build_cpu.py
@@ -47,3 +47,32 @@ def test_every_layer_gets_its_own_prefixed_weights():
     mlp_keys = {k for k in sd if ".mlp." in k}
     layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k}
     assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}"
+
+
+def test_offloaded_layers_borrow_the_next_router_for_the_lookahead():
+    # FREETOKEN_MOE_PREFETCH: layer L-1's experts hold layer L's router, outside the state dict
+    # and without the engine's offload-layer walk finding extra layers through it
+    import dataclasses
+
+    from freetoken.moe.offload_cache import iter_offload_moe_layers
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    previous = rotary._ROPE_DEVICE
+    set_rope_device(torch.device("cpu"))
+    try:
+        config = dataclasses.replace(parse_config(toy_hf_config()), moe_strategy="offload")
+        with torch.device("meta"):
+            model = Qwen4ExpForCausalLM(config)
+    finally:
+        rotary.get_rope.cache_clear()
+        set_rope_device(previous)
+    mlps = [layer.mlp for layer in model.model.layers.op_list]
+    assert len(list(iter_offload_moe_layers(model))) == len(mlps) == 4
+    wired = [mlp.experts._lookahead for mlp in mlps]
+    assert wired[-1] is None
+    assert [(gate is mlps[t].gate, t, budget) for gate, t, budget in wired[:-1]] == [
+        (True, 1, 3), (True, 2, 3), (True, 3, 4),  # layer 3 is the full-attention one
+    ]
+    assert not any("lookahead" in k for k in model.state_dict())
+    assert sum(k.endswith(".mlp.gate.weight") for k in model.state_dict()) == 4
diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py
index a2412f8346..2b47748f53 100644
--- a/tests/models/qwen4_exp/test_skeleton.py
+++ b/tests/models/qwen4_exp/test_skeleton.py
@@ -556,6 +556,1917 @@ def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch):
         assert torch.equal(a, b)
 
 
+# FREETOKEN_MOE_PREFETCH stack: 4 layers (full attention on layer 3), 16 experts, top-2, and one
+# shared cache as small as a layer so every step evicts
+_PF_LAYERS, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER = 4, 16, 2, 256, 128
+# ~50-100 us between layers: the predictor must finish before the next ensure even at toy size
+_PF_SLEEP_CYCLES = 200_000
+
+
+def _prefetch_banks(kind):
+    from freetoken.moe.expert_banks import build_expert_banks
+
+    torch.manual_seed(7)
+    method = _offload_experts(kind, 0, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER).quant_method
+    return build_expert_banks(method, _PF_LAYERS, None, device=torch.device("cuda"), dummy=True)
+
+
+def _prefetch_stack(kind, banks, mode, monkeypatch, *, k=6, budget=0, wire=True, overlap=False, policy="rule",
+                    cache_size=_PF_EXPERTS, **cache_kw):
+    from freetoken.env import ENV
+    from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE, wire_router_lookahead
+    from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.utils.torch_utils import torch_dtype
+
+    monkeypatch.setattr(ENV.MOE_PREFETCH_K, "value", k)
+    monkeypatch.setattr(ENV.MOE_PREFETCH_BUDGET, "value", budget)
+    config = parse_config(toy_hf_config(
+        _PF_LAYERS, hidden_size=_PF_HIDDEN, num_experts=_PF_EXPERTS, num_experts_per_tok=_PF_TOPK,
+        moe_intermediate_size=_PF_INTER, shared_expert_intermediate_size=_PF_INTER,
+    ))
+    device = torch.device("cuda")
+    with torch.device(device), torch_dtype(torch.bfloat16):
+        moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(_PF_LAYERS)]
+    gen = torch.Generator(device=device).manual_seed(43)
+    for moe in moes:
+        _fill(moe, gen, scale=0.2)
+    experts = [_offload_experts(kind, i, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER) for i in range(_PF_LAYERS)]
+    cache = OffloadMoeCache(
+        num_layers=_PF_LAYERS, num_experts=_PF_EXPERTS, cache_size=cache_size, device=device,
+        cache_policy=policy, quant_format=banks.quant_format, layout=banks.layout,
+        max_slots=experts[0].quant_method.slot_limit(), prefetch_mode=mode, decode_copy_overlap=overlap,
+        **cache_kw,
+    )
+    cache.set_bank_sources(banks.sources)
+    cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
+    cache.collect_stats = True
+    for moe, ex in zip(moes, experts):
+        ex.offload_cache = cache
+        moe.experts = ex
+    if wire:
+        wire_router_lookahead(moes, config)
+    cache.reset()
+    return moes, cache
+
+
+def _prefetch_run(moes, cache, bs, graph, inputs, before_decode=None, sleep=_PF_SLEEP_CYCLES, before_step=None):
+    """Warm up eagerly (and capture), reset the cache, then decode ``inputs`` (per step, one
+    ``[bs, hidden]`` input per layer); ``sleep`` cycles between layers, ``before_step(i)`` before step i."""
+    from flashlib.kernels.slot_cache import Stat
+
+    _fresh_ctx(_batch=SimpleNamespace(is_prefill=False))
+
+    def step(xs):
+        outs = []
+        for layer, (moe, x) in enumerate(zip(moes, xs)):
+            if layer and sleep:
+                torch.cuda._sleep(sleep)
+            outs.append(moe.forward(x.clone()))
+        return outs
+
+    static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes]
+    step(static)  # the eager warm-up the graph runner also does
+    if graph:
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            outs = step(static)
+    cache.reset()
+    cache.reset_stats()
+    if before_decode is not None:
+        before_decode()
+    results = []
+    for i, xs in enumerate(inputs):
+        if before_step is not None:
+            before_step(i)
+        if graph:
+            for s, x in zip(static, xs):
+                s.copy_(x)
+            g.replay()
+        else:
+            outs = step(xs)
+        if cache.audit is not None:
+            cache.audit.end_decode_step(cache)  # where the engine calls it: after the forward
+        results.append([o.clone() for o in outs])
+    torch.cuda.synchronize()
+    state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())]
+    counters = None if cache.prefetch is None else cache.prefetch.counters.cpu()
+    return results, state, counters, cache.lru_stats[:, Stat.MISS].cpu()
+
+
+def _prefetch_inputs(bs, steps=8, chain=False, seed=0):
+    gen = torch.Generator(device="cuda").manual_seed(seed * 10 + bs)
+
+    def x():
+        return torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+
+    if chain:  # every layer sees the same input, so the lookahead equals the next router
+        return [[x()] * _PF_LAYERS for _ in range(steps)]
+    return [[x() for _ in range(_PF_LAYERS)] for _ in range(steps)]
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("overlap,policy", [(False, "rule"), (True, "rule"), (True, "lru")], ids=["serial", "overlap", "overlap-lru"])
+def test_moe_prefetch_measure_is_bitwise_identical(kind, graph, bs, overlap, policy, monkeypatch):
+    """FREETOKEN_MOE_PREFETCH=measure predicts and counts beside the decode (with and without
+    FREETOKEN_MOE_COPY_OVERLAP's side stream); the MoE outputs, the slot map, the usage clock and
+    the slot contents must not change by a single bit."""
+    banks = _prefetch_banks(kind)
+    inputs = _prefetch_inputs(bs)
+    off = _prefetch_run(*_prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs)
+    on = _prefetch_run(*_prefetch_stack(kind, banks, "measure", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs)
+
+    assert off[2] is None and int(off[3].sum()) > 0 and torch.equal(on[3], off[3])
+    for got, want in zip(on[0], off[0]):
+        for a, b in zip(got, want):
+            assert torch.isfinite(b.float()).all()
+            assert torch.equal(a, b)
+    for a, b in zip(on[1], off[1]):
+        assert torch.equal(a, b)
+    from freetoken.moe.prefetch import CALLS, ISSUED, ROWS
+
+    counters = on[2]
+    assert counters[0].tolist() == [0] * counters.shape[1], "layer 0 has no predecessor"
+    assert counters[1:, CALLS].tolist() == [len(inputs)] * (_PF_LAYERS - 1)
+    assert counters[1:, ROWS].tolist() == [len(inputs) * bs] * (_PF_LAYERS - 1)
+    assert int(counters[1:, ISSUED].sum()) > 0
+
+
+def _reference_counters(selects, ensures, budget_override, k):
+    """CPU counters from the spied predictor logits and the residency its select read, and each
+    target layer's routing and residency at its ensure."""
+    from freetoken.moe.prefetch import NUM_COLS, default_budget
+    from tests.moe.ref_prefetch import ref_count, ref_select
+
+    config = parse_config(toy_hf_config(_PF_LAYERS))
+    counters = torch.zeros((_PF_LAYERS, NUM_COLS), dtype=torch.int64)
+    for layer in range(1, _PF_LAYERS):
+        assert len(selects[layer]) == len(ensures[layer]) > 0
+        want_budget = budget_override or default_budget(config.is_linear_layer(layer))
+        for (logits, budget, seen), (raw, resident) in zip(selects[layer], ensures[layer]):
+            assert budget == want_budget
+            sel, res = ref_select(logits.tolist(), seen.tolist(), k, budget)
+            routed = raw.view(-1).tolist()
+            misses = len({e for e in routed if resident[e] < 0})
+            counters[layer] += torch.tensor(ref_count(sel, res, routed, misses, logits.shape[0]))
+    return counters
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("budget", [0, 5])
+def test_moe_prefetch_counters_match_reference(kind, bs, budget, monkeypatch):
+    """Eager counters equal a CPU reference built from the predictor's own logits and the residency
+    its select read, and each layer's routing and pre-ensure residency; the captured graph counts
+    the same. The predictor runs on its dedicated stream, the count on the compute stream."""
+    import freetoken.moe.prefetch as pf_mod
+
+    banks = _prefetch_banks(kind)
+    inputs = _prefetch_inputs(bs, seed=1)
+    k = 6
+    moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget)
+    selects = {layer: [] for layer in range(_PF_LAYERS)}
+    ensures = {layer: [] for layer in range(_PF_LAYERS)}
+    streams = {"select": set(), "count": set()}
+    real_select, real_count, real_ensure = pf_mod.lookahead_select, pf_mod.prefetch_count, cache.ensure_experts
+
+    def spy_select(logits, resident, sel, res, *, k, budget):
+        # the select reads residency beside layer L-1's ensure: syncing on both sides pins what it saw
+        torch.cuda.synchronize()
+        layer = (sel.data_ptr() - cache.prefetch.sel.data_ptr()) // cache.prefetch.sel[0].nbytes
+        selects[layer].append((logits.float().cpu(), budget, resident.cpu()))
+        streams["select"].add(torch.cuda.current_stream().cuda_stream)
+        real_select(logits, resident, sel, res, k=k, budget=budget)
+        torch.cuda.synchronize()
+
+    def spy_count(*args, **kwargs):
+        streams["count"].add(torch.cuda.current_stream().cuda_stream)
+        return real_count(*args, **kwargs)
+
+    def spy_ensure(layer_id, expert_ids, **kwargs):
+        torch.cuda.synchronize()  # the prediction forked after the previous ensure has landed
+        ensures[layer_id].append((expert_ids.cpu(), cache.slot_for_id[layer_id].cpu()))
+        return real_ensure(layer_id, expert_ids, **kwargs)
+
+    def install():
+        monkeypatch.setattr(pf_mod, "lookahead_select", spy_select)
+        monkeypatch.setattr(pf_mod, "prefetch_count", spy_count)
+        monkeypatch.setattr(cache, "ensure_experts", spy_ensure)
+
+    eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install)
+    want = _reference_counters(selects, ensures, budget, k)
+    assert torch.equal(eager[2], want), (eager[2], want)
+    assert int(want[:, 1].sum()) > 0, "the sequence must have useful predictions"
+    assert streams["select"] == {cache.prefetch.stream.cuda_stream}
+    assert streams["count"] == {torch.cuda.current_stream().cuda_stream}
+
+    monkeypatch.undo()
+    moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget)
+    graph = _prefetch_run(moes, cache, bs, True, inputs)
+    assert torch.equal(graph[2], eager[2])
+    for got, want_out in zip(graph[0], eager[0]):
+        for a, b in zip(got, want_out):
+            assert torch.equal(a, b)
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+def test_moe_prefetch_perfect_lookahead_counts_every_miss(graph, bs, monkeypatch):
+    """Known answer: when layer L-1's router input equals layer L's, the lookahead is layer L's own
+    router, so with the budget covering every candidate each demand miss is a useful prediction."""
+    from freetoken.moe.prefetch import ISSUED, MISSES, RESIDENT_HITS, USEFUL
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, "measure", monkeypatch, k=4, budget=8)
+    _, _, counters, lru_misses = _prefetch_run(moes, cache, bs, graph, _prefetch_inputs(bs, chain=True, seed=2))
+    assert int(counters[1:, MISSES].sum()) > 0
+    assert counters[1:, USEFUL].tolist() == counters[1:, MISSES].tolist() == lru_misses[1:].tolist()
+    assert (counters[1:, ISSUED] >= counters[1:, USEFUL]).all()
+    assert int(counters[1:, RESIDENT_HITS].sum()) > 0, "routed hits among the top candidates"
+
+
+@requires_cuda
+def test_moe_prefetch_off_launches_nothing_new(monkeypatch):
+    """With the flag off the wired lookahead is inert: no prefetcher, and the decode launches
+    exactly the kernels of an unwired stack (measure adds some, so the check can see them)."""
+    from torch.profiler import ProfilerActivity, profile
+
+    banks = _prefetch_banks("bf16")
+    inputs = _prefetch_inputs(1, steps=2)
+
+    def kernels(mode, wire):
+        moes, cache = _prefetch_stack("bf16", banks, mode, monkeypatch, wire=wire)
+        _prefetch_run(moes, cache, 1, False, inputs[:1])  # warm up and compile outside the trace
+        with profile(activities=[ProfilerActivity.CUDA]) as prof:
+            _prefetch_run(moes, cache, 1, False, inputs)
+        events = [e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA]
+        return cache, [e.name for e in sorted(events, key=lambda e: e.time_range.start)]
+
+    cache, wired_off = kernels("off", True)
+    assert cache.prefetch is None
+    _, unwired_off = kernels("off", False)
+    _, measured = kernels("measure", True)
+    assert wired_off == unwired_off
+    assert len(measured) > len(wired_off)
+
+
+@requires_cuda
+def test_moe_prefetch_survives_a_rebuild(monkeypatch):
+    """A runtime rebuild resets the counters and keeps the predictor stream; a recaptured graph
+    counts like an eager run on the rebuilt cache."""
+    banks = _prefetch_banks("bf16")
+    inputs = _prefetch_inputs(1, seed=3)
+    moes, cache = _prefetch_stack("bf16", banks, "measure", monkeypatch)
+    before = _prefetch_run(moes, cache, 1, True, inputs)[2]
+    stream = cache.prefetch.stream
+    assert int(before.sum()) > 0
+    cache.rebuild(_PF_EXPERTS + 4)
+    assert int(cache.prefetch.counters.abs().sum()) == 0 and cache.prefetch.stream is stream
+    graph = _prefetch_run(moes, cache, 1, True, inputs)
+    eager = _prefetch_run(moes, cache, 1, False, inputs)
+    assert int(graph[2].sum()) > 0 and torch.equal(graph[2], eager[2])
+    for a, b in zip(graph[1], eager[1]):
+        assert torch.equal(a, b)
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("bs", [1, 2])
+def test_moe_prefetch_predicts_layer_l_from_layer_l_minus_1_input(kind, bs, monkeypatch):
+    """With a distinct input per layer, the prediction for layer L reads exactly layer L-1's router
+    input through layer L's router, and the select sees those logits and layer L's residency."""
+    import freetoken.moe.prefetch as pf_mod
+
+    banks = _prefetch_banks(kind)
+    inputs = _prefetch_inputs(bs, steps=3, seed=4)
+    moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch)
+    forks, selects = [], []
+    real_fork, real_select = cache.prefetch.fork, pf_mod.lookahead_select
+
+    def spy_fork(target, x, gate, resident, budget):
+        forks.append((target, x.clone(), gate, resident.data_ptr()))
+        return real_fork(target, x, gate, resident, budget)
+
+    def spy_select(logits, resident, sel, res, *, k, budget):
+        selects.append(logits.clone())
+        return real_select(logits, resident, sel, res, k=k, budget=budget)
+
+    def install():
+        monkeypatch.setattr(cache.prefetch, "fork", spy_fork)
+        monkeypatch.setattr(pf_mod, "lookahead_select", spy_select)
+
+    _prefetch_run(moes, cache, bs, False, inputs, before_decode=install)
+    assert [f[0] for f in forks] == list(range(1, _PF_LAYERS)) * len(inputs) and len(selects) == len(forks)
+    for call, ((target, x, gate, resident), logits) in enumerate(zip(forks, selects)):
+        assert torch.equal(x, inputs[call // (_PF_LAYERS - 1)][target - 1])
+        assert gate is moes[target].gate and resident == cache.slot_for_id[target].data_ptr()
+        assert torch.equal(logits, F.linear(x, moes[target].gate.weight))
+
+
+# FREETOKEN_MOE_PREFETCH=on: the same stack, with copies. No sleep between layers, so the prefetch
+# streams race the compute stream at toy size and only the joins keep the GEMMs on landed bytes.
+_PF_ON_STEPS = 200
+
+
+def _assert_slots_hold_their_experts(cache):
+    """Lossless: every held slot's bytes equal the host rows of the expert its id names."""
+    held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu()
+    ids = cache.id_of_slot[held.cuda()].long().cpu()
+    assert held.numel() > 0
+    for (per_layer, slots), view in zip(cache.banks, cache.bank_views()):
+        host = torch.stack(per_layer).flatten(0, 1)
+        want = host[ids].contiguous().view(torch.uint8)
+        got = view[held.cuda()].cpu().contiguous().view(torch.uint8)
+        assert torch.equal(got, want)
+
+
+def _routed_ids(moes, inputs):
+    """[step][layer] -> the experts each layer routes (the router alone decides; no cache input)."""
+    from freetoken.moe.fused import fused_topk
+
+    out = []
+    for xs in inputs:
+        row = []
+        for moe, x in zip(moes, xs):
+            _, ids = fused_topk(hidden_states=x, gating_output=moe.gate.forward(x), topk=_PF_TOPK, renormalize=True)
+            row.append(ids.view(-1).tolist())
+        out.append(row)
+    return out
+
+
+def _garbage(routed, width, kind, seed):
+    """Per step a [layers, width] candidate override: 'random' ids (and -1), 'adversarial' ones
+    that mix layer L's routed ids with ids it does not route, duplicates and padding, or 'flood':
+    every expert of layer L, so a prefetch wants more slots than the unpinned ones."""
+    gen = torch.Generator().manual_seed(seed)
+    steps = []
+    for row in routed:
+        sel = torch.full((_PF_LAYERS, width), -1, dtype=torch.int32)
+        for layer in range(1, _PF_LAYERS):
+            if kind == "random":
+                ids = torch.randint(-1, _PF_EXPERTS, (width,), generator=gen)
+            elif kind == "flood":
+                ids = torch.randperm(_PF_EXPERTS, generator=gen)[:width]
+            else:
+                hit = list(dict.fromkeys(row[layer]))
+                miss = [e for e in torch.randperm(_PF_EXPERTS, generator=gen).tolist() if e not in hit][:3]
+                ids = torch.tensor((hit[:3] + miss + hit[:1] + [-1] + miss[:1]) * 2)[:width]
+                ids = ids[torch.randperm(ids.numel(), generator=gen)]
+            sel[layer, : ids.numel()] = ids.to(torch.int32)
+        steps.append(sel)
+    return steps
+
+
+def _on_stack(kind, banks, monkeypatch, *, policy, overlap, k=6, budget=0):
+    return _prefetch_stack(kind, banks, "on", monkeypatch, k=k, budget=budget, overlap=overlap, policy=policy)
+
+
+def _with_override(cache, steps, **delays):
+    """before_step hook that loads step i's candidate override into the prefetcher's static buffer."""
+    pf = cache.prefetch
+    pf.sel_override = torch.full_like(pf.sel, -1)
+    for name, cycles in delays.items():
+        setattr(pf, name, cycles)
+
+    def before_step(i):
+        pf.sel_override.copy_(steps[i][:, : pf.sel.shape[1]])
+
+    return before_step
+
+
+def _assert_same_outputs(got, want):
+    for a_step, b_step in zip(got, want):
+        for a, b in zip(a_step, b_step):
+            assert torch.isfinite(b.float()).all()
+            assert torch.equal(a, b)
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+@pytest.mark.parametrize("overlap", [False, True], ids=["serial", "overlap"])
+def test_moe_prefetch_on_is_bitwise_identical(kind, graph, bs, policy, overlap, monkeypatch):
+    """FREETOKEN_MOE_PREFETCH=on copies predicted experts into low-priority slots beside the decode:
+    over 200 steps of a cache that evicts every step the MoE outputs are those of prefetch off to
+    the bit, every held slot holds its expert's bytes, and the prefetches were used."""
+    from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, USEFUL
+
+    banks = _prefetch_banks(kind)
+    inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=5)
+    moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    _assert_slots_hold_their_experts(cache)
+    moes, cache = _on_stack(kind, banks, monkeypatch, policy=policy, overlap=overlap)
+    on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    _assert_same_outputs(on[0], off[0])
+    _assert_slots_hold_their_experts(cache)
+    stats = cache.prefetch.stats.cpu()
+    assert stats[0].tolist() == [0] * stats.shape[1], "layer 0 is never prefetched"
+    assert stats[1:, CALLS].tolist() == [_PF_ON_STEPS] * (_PF_LAYERS - 1)
+    assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 and int(stats[:, COPIED].sum()) > 0
+    assert int(on[3].sum()) < int(off[3].sum()), "useful prefetches remove demand misses"
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+@pytest.mark.parametrize("garbage", ["random", "adversarial"])
+def test_moe_prefetch_on_survives_a_garbage_predictor(graph, bs, policy, garbage, monkeypatch):
+    """Whatever the candidates (random ids and padding, or the next layer's routed ids mixed with ids
+    it will not route, duplicates and -1), prefetch on changes no output bit and no slot's bytes."""
+    from freetoken.moe.prefetch import ISSUED, USEFUL
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=6)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8)
+    steps = _garbage(_routed_ids(moes, inputs), 12, garbage, seed=bs)
+    on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps))
+    _assert_same_outputs(on[0], off[0])
+    _assert_slots_hold_their_experts(cache)
+    stats = cache.prefetch.stats.cpu()
+    assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0
+    assert int(stats[:, USEFUL].sum()) < int(stats[:, ISSUED].sum()), "some prefetches must be wrong"
+
+
+# ~0.5 ms at the RTX 5080's clock: far longer than a toy layer, so every delayed stage is late
+_PF_DELAY_CYCLES = 1_000_000
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+@pytest.mark.parametrize("delay", ["copy", "predict", "both"])
+def test_moe_prefetch_on_joins_hold_when_the_prefetch_streams_run_late(graph, bs, policy, delay, monkeypatch):
+    """Race test: a sleep on the prefetch copy stream (every copy lands after its layer's ensure) or on
+    the predictor stream (prefetch_ensure runs long after the fork) must not change a bit. The
+    candidates include each layer's routed ids, so GEMMs do read prefetched slots. The joins also
+    make the slot maps timing-free: the delayed run ends with the undelayed run's maps and counts."""
+    from freetoken.moe.prefetch import COPIED, LATE, NUM_STAT_COLS, USEFUL
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=7)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    runs = []
+    for cycles in (0, _PF_DELAY_CYCLES):
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8)
+        steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 10)
+        delays = {"delay_copy_cycles": cycles * (delay in ("copy", "both")),
+                  "delay_predict_cycles": cycles * (delay in ("predict", "both"))}
+        on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps, **delays))
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+        runs.append((on[1][:3], cache.prefetch.stats.cpu()))
+    (maps, stats), (delayed_maps, delayed_stats) = runs
+    assert int(stats[:, USEFUL].sum()) > 0
+    for a, b in zip(maps, delayed_maps):
+        assert torch.equal(a, b)
+    cols = [c for c in range(NUM_STAT_COLS) if c != LATE]
+    assert torch.equal(stats[:, cols], delayed_stats[:, cols])
+    if delay == "copy":
+        # the late counter sees what the join waited for: the copies were still sleeping at ensure(L)
+        # (all of them under replay; in eager a loaded host can take longer than the sleep per layer)
+        late, copied = int(delayed_stats[:, LATE].sum()), int(delayed_stats[:, COPIED].sum())
+        assert copied > 0 and (late == copied if graph else late > 0)
+
+
+@requires_cuda
+def test_moe_prefetch_race_test_catches_a_missing_copy_join(monkeypatch):
+    """Negative control for the race test: without the compute stream's wait on the prefetch copy,
+    GEMMs read prefetched slots before their bytes land and the outputs change."""
+    from freetoken.moe.prefetch import ExpertPrefetcher
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=40, seed=7)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True)
+    off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8)
+
+    def no_wait(self, layer_id):
+        self._inflight[layer_id] = None
+
+    monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait)
+    steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11)
+    on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0,
+                       before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES))
+    assert any(not torch.equal(a, b) for got, want in zip(on[0], off[0]) for a, b in zip(got, want))
+
+
+def _replay_on_reference(policy, calls, rows):
+    """Replay the spied prefetch_ensure / ensure sequence on the CPU reference; per-layer stats."""
+    from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, MISSES, ROWS, USEFUL
+    from freetoken.moe.scored_ensure import POLICY_IDS
+    from tests.moe.ref_scored_cache import RefScoredCache
+
+    ref = RefScoredCache(_PF_LAYERS, _PF_EXPERTS, _PF_EXPERTS, POLICY_IDS[policy])
+    stats = torch.zeros((_PF_LAYERS, 8), dtype=torch.int64)
+    planned = {}
+    for call in calls:
+        if call[0] == "prefetch":
+            _, layer, sel, budget = call
+            src, _ = ref.prefetch(layer, sel.tolist(), budget)
+            stats[layer, ISSUED] += src.size
+            planned[layer] = src.size
+        else:
+            _, layer, ids, logits = call
+            _, src, _ = ref.ensure(layer, ids.view(-1).numpy(), bump_tok=layer == 0, lowpri=True,
+                                   logits=None if logits is None else logits.float().numpy(), thr=0.25)
+            if layer in planned:
+                stats[layer, [USEFUL, CALLS, ROWS, MISSES, COPIED]] += torch.tensor(
+                    [ref.useful, 1, rows, src.size, int(planned.pop(layer) > 0)])
+    return ref, stats
+
+
+@requires_cuda
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_prefetch_on_counters_match_reference(bs, policy, monkeypatch):
+    """issued / useful / calls / rows / misses / copied equal a CPU replay of the spied candidates
+    and routings, the slot maps end identical, late is 0 when each ensure waits for the copies
+    (known answer), and with the candidates fixed a captured graph counts the same as eager.
+    prefetch_ensure and the demand copy run on the compute stream, the slim prefetch copy on the
+    copy stream."""
+    import freetoken.kernel.fast_index_copy as fic
+    from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, RESIDENT_HITS
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=60, seed=8)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True)
+    calls, streams = [], {"pf_ensure": set(), "slim": set(), "demand": set()}
+    real_pf, real_ensure = cache.prefetch_ensure, cache.ensure_experts
+    real_slim, real_multi = fic.fast_index_copy_multi_slim_jit, fic.fast_index_copy_multi_jit
+
+    def spy_pf(layer, query, *args, **kw):
+        streams["pf_ensure"].add(torch.cuda.current_stream().cuda_stream)
+        calls.append(("prefetch", layer, query.clone(), kw.get("budget")))  # on the compute stream after its wait on the select
+        return real_pf(layer, query, *args, **kw)
+
+    def spy_ensure(layer, ids, **kw):
+        torch.cuda.synchronize()  # the copies forked before this ensure have landed: late must be 0
+        logits = kw.get("router_logits")
+        calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone()))
+        return real_ensure(layer, ids, **kw)
+
+    def spy_slim(*args):
+        streams["slim"].add(torch.cuda.current_stream().cuda_stream)
+        return real_slim(*args)
+
+    def spy_multi(*args):
+        streams["demand"].add(torch.cuda.current_stream().cuda_stream)
+        return real_multi(*args)
+
+    def install():
+        monkeypatch.setattr(cache, "prefetch_ensure", spy_pf)
+        monkeypatch.setattr(cache, "ensure_experts", spy_ensure)
+        monkeypatch.setattr(fic, "fast_index_copy_multi_slim_jit", spy_slim)
+        monkeypatch.setattr(fic, "fast_index_copy_multi_jit", spy_multi)
+
+    eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install, sleep=0)
+    torch.cuda.synchronize()
+    calls = [c if c[0] == "prefetch" else (c[0], c[1], c[2].cpu(), None if c[3] is None else c[3].cpu()) for c in calls]
+    calls = [(c[0], c[1], c[2].cpu(), c[3]) if c[0] == "prefetch" else c for c in calls]
+    ref, want = _replay_on_reference(policy, calls, bs)
+    got = cache.prefetch.stats.cpu()
+    assert got.shape == (_PF_LAYERS, NUM_STAT_COLS)
+    assert int(got[:, LATE].sum()) == 0 and int(got[:, RESIDENT_HITS].sum()) == 0
+    assert torch.equal(got, want), (got, want)
+    assert int(want[:, 1].sum()) > 0, "the sequence must have useful prefetches"
+    assert torch.equal(cache.slot_for_id.view(-1).cpu().long(), torch.from_numpy(ref.slot_of_id))
+    assert torch.equal(cache.id_of_slot.cpu().long(), torch.from_numpy(ref.id_of_slot))
+    assert torch.equal(cache.usage.cpu(), torch.from_numpy(ref.usage))
+    assert streams["pf_ensure"] == {torch.cuda.current_stream().cuda_stream}
+    assert streams["slim"] == {cache.prefetch.copy_stream.cuda_stream}
+    assert streams["demand"] == {torch.cuda.current_stream().cuda_stream}
+
+    # the select reads residency beside layer L-1's ensure, so two runs may pick different candidates:
+    # with the candidates fixed, a captured graph must install, count and end like eager
+    monkeypatch.undo()
+    runs = []
+    for graph in (False, True):
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True)
+        steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 8)
+        run = _prefetch_run(moes, cache, bs, graph, inputs, sleep=_PF_SLEEP_CYCLES, before_step=_with_override(cache, steps))
+        _assert_same_outputs(run[0], eager[0])
+        runs.append((run[1][:3], cache.prefetch.stats.cpu()))
+    cols = [c for c in range(NUM_STAT_COLS) if c != LATE]
+    (maps, stats), (graph_maps, graph_stats) = runs
+    for a, b in zip(maps, graph_maps):
+        assert torch.equal(a, b)
+    assert torch.equal(graph_stats[:, cols], stats[:, cols]) and int(stats[:, 1].sum()) > 0
+
+
+@requires_cuda
+def test_moe_prefetch_on_rebuild_and_reset_clear_the_prefetch_state(monkeypatch):
+    """A rebuild keeps both prefetch streams, the events and plan buffers, and drops the counters,
+    the in-flight state and every low-priority slot; a recaptured graph then matches eager and off.
+    A reset (as after a capture) also drops the counters and low-priority slots."""
+    banks = _prefetch_banks("bf16")
+    inputs = _prefetch_inputs(1, steps=40, seed=9)
+    moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True)
+    _prefetch_run(moes, cache, 1, True, inputs, sleep=0)
+    pf = cache.prefetch
+    streams, plan = (pf.stream, pf.copy_stream), (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready)
+    assert int(pf.stats.sum()) > 0 and bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any())
+    pf._inflight[2] = torch.zeros(1)
+    cache.rebuild(_PF_EXPERTS + 4)
+    assert (pf.stream, pf.copy_stream) == streams and (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) == plan
+    assert int(pf.stats.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 and pf._inflight == [None] * _PF_LAYERS
+    assert not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any())
+    graph = _prefetch_run(moes, cache, 1, True, inputs, sleep=0)
+    eager = _prefetch_run(moes, cache, 1, False, inputs, sleep=0)
+    _assert_same_outputs(graph[0], eager[0])
+    _assert_slots_hold_their_experts(cache)
+    assert int(pf.stats.sum()) > 0
+    cache.reset()
+    assert int(pf.stats.abs().sum()) == 0 and not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any())
+    moes_off, cache_off = _prefetch_stack("bf16", banks, "off", monkeypatch, overlap=True)
+    cache_off.rebuild(_PF_EXPERTS + 4)
+    off = _prefetch_run(moes_off, cache_off, 1, False, inputs, sleep=0)
+    _assert_same_outputs(eager[0], off[0])
+
+
+@requires_cuda
+def test_moe_prefetch_on_skips_batches_past_the_select_tile(monkeypatch):
+    """An eager decode wider than MAX_ROWS forks nothing: no join, no count, the same outputs."""
+    from freetoken.moe.prefetch import MAX_ROWS
+
+    banks = _prefetch_banks("bf16")
+    bs = MAX_ROWS + 1
+    inputs = _prefetch_inputs(bs, steps=6, seed=10)
+    off = _prefetch_run(*_prefetch_stack("bf16", banks, "off", monkeypatch), bs, False, inputs, sleep=0)
+    moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True)
+    on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0)
+    _assert_same_outputs(on[0], off[0])
+    assert int(cache.prefetch.stats.abs().sum()) == 0 and cache.prefetch._inflight == [None] * _PF_LAYERS
+
+
+def _stall_hooks(monkeypatch):
+    """Sleeps (cycles, read at enqueue time) on the compute stream right before each decode GEMM, i.e.
+    after issue_copy(L+1), and on the stream that runs the shared expert gate (the side stream)."""
+    import freetoken.models.qwen4_exp.moe as q4moe
+    from freetoken.layers.moe import OffloadMoELayer
+
+    stall = {"gemm": 0, "shared": 0}
+    real_gemm, real_gate = OffloadMoELayer._expert_gemm, q4moe.shared_gate_sigmoid
+
+    def gemm(self, cache, *args, is_prefill, **kw):
+        if stall["gemm"] and not is_prefill:
+            torch.cuda._sleep(stall["gemm"])
+        return real_gemm(self, cache, *args, is_prefill=is_prefill, **kw)
+
+    def gate(*args, **kw):
+        if stall["shared"]:
+            torch.cuda._sleep(stall["shared"])
+        return real_gate(*args, **kw)
+
+    monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm)
+    monkeypatch.setattr(q4moe, "shared_gate_sigmoid", gate)
+    return stall
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+@pytest.mark.parametrize("where", ["gemm", "shared"])
+@pytest.mark.parametrize("cands", ["adversarial", "flood"])
+def test_moe_prefetch_on_holds_when_the_compute_side_stalls(graph, bs, policy, where, cands, monkeypatch):
+    """Race test from the other side: a sleep on the compute stream between issue_copy(L+1) and
+    GEMM(L) lets copy(L+1) land before GEMM(L) reads, so a prefetch into any of layer L's routed
+    slots changes outputs deterministically; a sleep on the shared expert's side stream must not
+    matter either. 'flood' asks for all 16 experts a layer, more than the unpinned slots, so even
+    LRU's recency order would reach the pinned ones. Bit-identical to off, and the same final maps
+    as the unstalled run."""
+    k, width = (8, 16) if cands == "flood" else (6, 12)
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=80, seed=21)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    stall = _stall_hooks(monkeypatch)
+    runs = []
+    for cycles in (0, _PF_DELAY_CYCLES):
+        stall[where] = cycles
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, k=k, budget=width)
+        steps = _garbage(_routed_ids(moes, inputs), width, cands, seed=bs + 20)
+        on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps))
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+        runs.append(on[1][:3])
+    for a, b in zip(*runs):
+        assert torch.equal(a, b)
+
+
+def _assert_held_slots_hold_their_experts(cache):
+    """_assert_slots_hold_their_experts, but a prefill may leave no slot held."""
+    if bool((cache.id_of_slot >= 0).any()):
+        _assert_slots_hold_their_experts(cache)
+
+
+def _mixed_run(moes, cache, bs, graph, schedule, before_step=None, late=None, settle=True):
+    """Warm up (and capture) the decode like _prefetch_run, reset, then run ``schedule``: ("decode", xs)
+    is one decode step (a replay under ``graph``), ("prefill", xs) an eager prefill of xs's tokens.
+    ``late=(cycles, steps)`` sleeps that long on the predictor stream in those decode steps (a second
+    graph captured with the sleep, in the same pool, replays them); ``settle`` syncs and checks the
+    slots after each prefill, else the next decode is enqueued right behind it."""
+    ctx = _fresh_ctx(_batch=SimpleNamespace(is_prefill=False))
+
+    def step(xs):
+        return [moe.forward(x.clone()) for moe, x in zip(moes, xs)]
+
+    static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes]
+    step(static)
+    if graph:
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            graphs = {False: (g, step(static))}
+        if late is not None:
+            cache.prefetch.delay_predict_cycles = late[0]
+            g_late = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g_late, pool=g.pool()):
+                graphs[True] = (g_late, step(static))
+            cache.prefetch.delay_predict_cycles = 0
+    cache.reset()
+    results = []
+    for i, (phase, xs) in enumerate(schedule):
+        if before_step is not None:
+            before_step(i)
+        ctx._batch.is_prefill = phase == "prefill"
+        slow = late is not None and i in late[1]
+        if phase == "decode" and graph:
+            for s, x in zip(static, xs):
+                s.copy_(x)
+            g, outs = graphs[slow]
+            g.replay()
+        else:
+            if late is not None:
+                cache.prefetch.delay_predict_cycles = late[0] if slow else 0
+            outs = step(xs)
+        if phase == "decode" and cache.audit is not None:
+            cache.audit.end_decode_step(cache)
+        results.append([o.clone() for o in outs])
+        if phase == "prefill" and settle:
+            torch.cuda.synchronize()
+            _assert_held_slots_hold_their_experts(cache)
+    torch.cuda.synchronize()
+    _assert_held_slots_hold_their_experts(cache)
+    return results
+
+
+_PREFILL_PATHS = {
+    # (cache_size, prefill tokens, cache kwargs, FREETOKEN_MOE_SMALL_PREFILL_TOKENS)
+    "materialize": (_PF_EXPERTS + 4, 24, {}, 0),
+    "overlap": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True}, 0),
+    "overlap_d2d": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True, "prefill_hit_d2d": True}, 0),
+    "small": (_PF_EXPERTS, 3, {}, 8),
+}
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+@pytest.mark.parametrize("path", list(_PREFILL_PATHS))
+def test_moe_prefetch_on_survives_prefills_between_decodes(graph, bs, policy, path, monkeypatch):
+    """Prefills between decode steps rewrite slots behind the prefetch's back: the whole-layer
+    materialize installs slots [0, E), the overlap double buffers invalidate and stream [0, 2E)
+    (with the hit-D2D gather reading resident slots), and small prefills ensure their experts
+    under the low-priority keys. Decode, prefill, decode, prefill, decode: every output of prefetch
+    on equals off to the bit, and after each prefill every held slot holds its expert."""
+    import freetoken.layers.moe as layers_moe
+
+    cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path]
+    monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small)
+    banks = _prefetch_banks("nvfp4")
+    decode = _prefetch_inputs(bs, steps=18, seed=31)
+    gen = torch.Generator(device="cuda").manual_seed(32)
+    prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+                 for _ in range(_PF_LAYERS)] for _ in range(2)]
+    schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])]
+                + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])]
+                + [("decode", xs) for xs in decode[12:]])
+    kw = dict(overlap=True, policy=policy, cache_size=cache_size, **cache_kw)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw)
+    off = _mixed_run(moes, cache, bs, graph, schedule)
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, **kw)
+    steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 30)
+    on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps))
+    _assert_same_outputs(on, off)
+    if path == "overlap_d2d":
+        assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots"
+    assert int(cache.prefetch.stats[:, 0].sum()) > 0
+
+
+def _record_compute_streams(monkeypatch, moes):
+    """The streams the MoE blocks are entered on (the default stream, and the capture stream under a
+    graph): the compute streams of the decode."""
+    compute = set()
+    for moe in moes:
+        real = moe.forward
+
+        def forward(x, real=real):
+            compute.add(torch.cuda.current_stream().cuda_stream)
+            return real(x)
+
+        monkeypatch.setattr(moe, "forward", forward)
+    return compute
+
+
+def _spy_map_writers(monkeypatch):
+    """(writer, stream) of every launch that writes the slot maps or their mirrors, at enqueue time."""
+    import freetoken.kernel.triton.moe as triton_moe
+    import freetoken.moe.offload_kernels as ok
+    import freetoken.moe.scored_ensure as se
+
+    seen = []
+
+    def spy(name, real):
+        def launch(*args, **kwargs):
+            seen.append((name(kwargs) if callable(name) else name, torch.cuda.current_stream().cuda_stream))
+            return real(*args, **kwargs)
+
+        return launch
+
+    # scored_ensure is the kernel of both the demand ensures and prefetch_ensure (prefetch=True)
+    monkeypatch.setattr(se, "scored_ensure", spy(lambda kw: "prefetch_ensure" if kw.get("prefetch") else "ensure",
+                                                 se.scored_ensure))
+    monkeypatch.setattr(ok, "lru_ensure", spy("ensure", ok.lru_ensure))
+    monkeypatch.setattr(ok, "materialize_layer", spy("materialize", ok.materialize_layer))
+    monkeypatch.setattr(ok, "reset_cache", spy("reset", ok.reset_cache))
+    monkeypatch.setattr(triton_moe, "invalidate_prefill_slots", spy("invalidate", triton_moe.invalidate_prefill_slots))
+    return seen
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("path", list(_PREFILL_PATHS))
+def test_moe_prefetch_on_writes_the_slot_maps_only_on_the_compute_stream(graph, bs, path, monkeypatch):
+    """Single writer: every launch that writes the slot maps (the demand and prefetch ensures, prefill
+    ensures, the whole-layer materialize, reset) runs on the stream the MoE blocks run on (the capture
+    stream under a graph), never on the predictor, prefetch copy or shared-expert side stream. The only
+    other writer is the overlap prefill's buffer invalidation, on the prefill copy stream it fences.
+    The predictor's select reads no cache state."""
+    import freetoken.layers.moe as layers_moe
+    import freetoken.moe.prefetch as pf_mod
+    from freetoken.moe.prefetch import ISSUED
+
+    cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path]
+    monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small)
+    banks = _prefetch_banks("nvfp4")
+    decode = _prefetch_inputs(bs, steps=6, seed=95)
+    gen = torch.Generator(device="cuda").manual_seed(96)
+    prefill = [torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+               for _ in range(_PF_LAYERS)]
+    schedule = [("decode", xs) for xs in decode[:3]] + [("prefill", prefill)] + [("decode", xs) for xs in decode[3:]]
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, overlap=True, cache_size=cache_size,
+                                  **cache_kw)
+    steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 97)
+    seen = _spy_map_writers(monkeypatch)
+    compute = _record_compute_streams(monkeypatch, moes)
+    selects, real_select = [], pf_mod.lookahead_select
+
+    def spy_select(logits, resident, *args, **kwargs):
+        selects.append((resident.data_ptr(), torch.cuda.current_stream().cuda_stream))
+        return real_select(logits, resident, *args, **kwargs)
+
+    monkeypatch.setattr(pf_mod, "lookahead_select", spy_select)
+    _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps))
+    pf = cache.prefetch
+    assert not compute & {pf.stream.cuda_stream, pf.copy_stream.cuda_stream, cache.decode_copy_stream.cuda_stream}
+    assert len(compute) == 1 + graph
+    writers = {name for name, _ in seen}
+    want = {"ensure", "prefetch_ensure", "reset"} | {"materialize": {"materialize"}, "overlap": {"invalidate"},
+                                                     "overlap_d2d": {"invalidate"}, "small": set()}[path]
+    assert writers == want, writers
+    off_compute = [(name, stream) for name, stream in seen if stream not in compute]
+    fence = {cache.prefill_copy_stream.cuda_stream} if cache.prefill_copy_stream is not None else set()
+    assert all(name == "invalidate" and stream in fence for name, stream in off_compute), off_compute
+    assert {stream for name, stream in seen if name == "prefetch_ensure"} == compute  # eager, and the capture
+    assert int(pf.stats[:, ISSUED].sum()) > 0
+    # after the writer checks, which are what a predictor-side install fails
+    assert selects and set(selects) == {(pf._none_resident.data_ptr(), pf.stream.cuda_stream)}
+
+
+# ~4 ms at the RTX 5080's clock: a select far later than its whole layer
+_PF_LATE_CYCLES = 8 * _PF_DELAY_CYCLES
+
+
+def _step_starts(times):
+    """before_step hook recording a timed event on the compute stream at the start of each step."""
+    def before_step(i):
+        times.append(torch.cuda.Event(enable_timing=True))
+        times[-1].record()
+
+    return before_step
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_prefetch_on_waits_for_a_late_select(graph, bs, policy, monkeypatch):
+    """The compute stream waits for the select before it installs, however late the predictor runs:
+    with every select ~4 ms late (far past its whole layer) a step lasts at least its three late
+    selects, and the outputs equal off's and the maps and counts (all but LATE) those of the run
+    without the delay."""
+    from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, USEFUL
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=24, seed=100)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    start.record()
+    torch.cuda._sleep(_PF_LATE_CYCLES)
+    end.record()
+    end.synchronize()
+    late_ms = start.elapsed_time(end)
+    runs = []
+    for cycles in (0, _PF_LATE_CYCLES):
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8)
+        steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 101)
+        load, times = _with_override(cache, steps, delay_predict_cycles=cycles), []
+        mark = _step_starts(times)
+
+        def before_step(i, load=load, mark=mark):
+            mark(i)
+            load(i)
+
+        on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step)
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+        step_ms = sorted(a.elapsed_time(b) for a, b in zip(times, times[1:]))
+        runs.append((on[1][:3], cache.prefetch.stats.cpu(), step_ms[len(step_ms) // 2]))
+    (maps, stats, fast_ms), (late_maps, late_stats, slow_ms) = runs
+    assert int(stats[:, USEFUL].sum()) > 0
+    for a, b in zip(maps, late_maps):
+        assert torch.equal(a, b)
+    cols = [c for c in range(NUM_STAT_COLS) if c != LATE]
+    assert torch.equal(stats[:, cols], late_stats[:, cols])
+    # layers 1-3 are predicted; each prediction forks after the previous one was waited for
+    assert slow_ms >= 0.9 * (_PF_LAYERS - 1) * late_ms > 4 * fast_ms, (slow_ms, late_ms, fast_ms)
+
+
+@requires_cuda
+def test_moe_prefetch_on_late_select_without_the_wait_changes_the_maps_but_no_output(monkeypatch):
+    """Negative control for the test above (eager: a capture needs the predictor joined): without the
+    compute stream's wait, prefetch_ensure installs stale candidates, so the maps differ from the
+    undelayed run's, yet every output still equals off's: the slot maps have one writer, so a wrong
+    prediction only costs a copy."""
+    from freetoken.moe.prefetch import ExpertPrefetcher
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=24, seed=100)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True)
+    off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0)
+
+    def no_wait(self, target):
+        """install without the compute stream's wait on the select."""
+        if target is None or self._inflight[target] is None:
+            return None
+        self.cache.prefetch_ensure(
+            target, self.sel[target], self.pf_slots[target], self.pf_src[target],
+            self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target],
+        )
+        self._installed[target] = True
+        return target
+
+    runs = []
+    for skip in (False, True):
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8)
+        if skip:
+            monkeypatch.setattr(ExpertPrefetcher, "install", no_wait)
+        steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=102)
+        on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0,
+                           before_step=_with_override(cache, steps, delay_predict_cycles=_PF_LATE_CYCLES * skip))
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+        runs.append(on[1][:3])
+    assert any(not torch.equal(a, b) for a, b in zip(*runs))
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("path", list(_PREFILL_PATHS))
+def test_moe_prefetch_on_late_select_right_after_a_prefill(graph, bs, path, monkeypatch):
+    """The suspected production trigger: a decode step enqueued right behind a prefill (which
+    rewrote slots behind the prefetch's back, the overlap buffers' [0, 2E) included) with the
+    predictor ~4 ms late. Every output equals off's to the bit, the maps equal those of the same run
+    with no late step, every held slot holds its expert and the decode verifier finds nothing."""
+    import freetoken.layers.moe as layers_moe
+    from freetoken.moe.verify import CHECKS
+
+    cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path]
+    monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small)
+    banks = _prefetch_banks("nvfp4")
+    decode = _prefetch_inputs(bs, steps=15, seed=105)
+    gen = torch.Generator(device="cuda").manual_seed(106)
+    prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+                 for _ in range(_PF_LAYERS)] for _ in range(2)]
+    schedule = ([("decode", xs) for xs in decode[:5]] + [("prefill", prefills[0])]
+                + [("decode", xs) for xs in decode[5:10]] + [("prefill", prefills[1])]
+                + [("decode", xs) for xs in decode[10:]])
+    after_prefill = {i + 1 for i, (phase, _) in enumerate(schedule) if phase == "prefill"}
+    kw = dict(overlap=True, cache_size=cache_size, **cache_kw)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw)
+    off = _mixed_run(moes, cache, bs, graph, schedule)
+    runs = []
+    for late in (None, (_PF_LATE_CYCLES, after_prefill)):
+        moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, verify_mode="meta", **kw)
+        steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 107)
+        on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps), late=late,
+                        settle=False)
+        _assert_same_outputs(on, off)
+        counts, records = _verify_state(cache)
+        assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records
+        runs.append([t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage)])
+    for a, b in zip(*runs):
+        assert torch.equal(a, b)
+    if path == "overlap_d2d":
+        assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots"
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_prefetch_on_install_may_evict_a_slot_whose_prefetch_copy_is_in_flight(graph, bs, policy, monkeypatch):
+    """prefetch_ensure(L) runs on the compute stream before it joins copy(L-1), so it may evict a slot
+    that plan L-1 installed and whose slim copy is still writing it. The next writer of that slot's
+    bytes (copy(L) on the same copy stream, or a demand copy after join_copy) is ordered after copy(L-1):
+    with every copy late and flood candidates the eviction happens (counted on the device, so it holds
+    under replay too), and still no output bit and no slot's bytes change."""
+    from freetoken.moe.prefetch import ExpertPrefetcher
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=40, seed=110)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8)
+    pf = cache.prefetch
+    evicted_in_flight = torch.zeros((), dtype=torch.int64, device="cuda")
+    col = torch.arange(pf.pf_slots.shape[1], device="cuda")
+    real_install = ExpertPrefetcher.install
+
+    def install(self, target):
+        got = real_install(self, target)
+        if got is not None and got >= 2:
+            # on the compute stream right after prefetch_ensure(got), before join_copy(got - 1)
+            new = torch.where(col < self.pf_num[got], self.pf_slots[got], -2)
+            old = torch.where(col < self.pf_num[got - 1], self.pf_slots[got - 1], -3)
+            shared = (new[:, None] == old[None, :]).any()
+            evicted_in_flight.add_((shared & (self.pf_ready[got - 1][0] == 0)).long())
+        return got
+
+    monkeypatch.setattr(ExpertPrefetcher, "install", install)
+    steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=bs + 111)
+    on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0,
+                       before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES))
+    _assert_same_outputs(on[0], off[0])
+    _assert_slots_hold_their_experts(cache)
+    assert int(evicted_in_flight) > 0, "no install evicted a slot of the previous layer's in-flight copy"
+
+
+@requires_cuda
+def test_moe_prefetch_in_flight_eviction_test_catches_an_unordered_next_copy(monkeypatch):
+    """Negative control for the test above: order each prefetch copy only after its own install (not
+    after copy(L-1)'s join) and put odd and even layers' copies on two streams, only the even ones
+    late. A late copy then lands on a slot the next install took over, and the outputs or the held
+    slots' bytes go wrong."""
+    from freetoken.moe.prefetch import ExpertPrefetcher, _mark_ready_kernel, dedicated_stream
+    from freetoken.moe.slot_audit import PREFETCH_COPY
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=40, seed=110)
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True)
+    off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8)
+    pf = cache.prefetch
+    streams = [pf.copy_stream, dedicated_stream(torch.device("cuda"), highest_priority=True)]
+    installed = [torch.cuda.Event() for _ in range(_PF_LAYERS)]
+    real_install = ExpertPrefetcher.install
+
+    def install(self, target):
+        got = real_install(self, target)
+        if got is not None:
+            installed[got].record(torch.cuda.current_stream())
+        return got
+
+    def issue_copy(self, target):
+        if not self._pending(target):
+            return
+        stream = streams[target % 2]
+        stream.wait_event(installed[target])
+        with torch.cuda.stream(stream):
+            if target % 2 == 0:
+                torch.cuda._sleep(_PF_DELAY_CYCLES)
+            self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True,
+                                 kind=PREFETCH_COPY)
+            _mark_ready_kernel[(1,)](self.pf_ready[target])
+            self._copy_events[target][1].record(stream)
+
+    monkeypatch.setattr(ExpertPrefetcher, "install", install)
+    monkeypatch.setattr(ExpertPrefetcher, "issue_copy", issue_copy)
+    steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=112)
+    on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, before_step=_with_override(cache, steps))
+    try:
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+    except AssertionError:
+        return
+    pytest.fail("a prefetch copy landing on a slot the next install took over went unnoticed")
+
+
+@requires_cuda
+@pytest.mark.parametrize("kind", ["bf16", "nvfp4"])
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+def test_moe_prefetch_on_per_bank_copy_fallback_is_bitwise_identical(kind, graph, bs, monkeypatch):
+    """Without the fused copy plan (FREETOKEN_FUSED_COPY=0 or misaligned banks) the prefetch copy
+    takes copy_rows' per-bank launches on the copy stream; late copies still change no bit."""
+    import freetoken.moe.offload_cache as offload_cache
+
+    monkeypatch.setattr(offload_cache, "_FUSED_COPY", False)
+    banks = _prefetch_banks(kind)
+    inputs = _prefetch_inputs(bs, steps=60, seed=5)
+    moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=True)
+    off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0)
+    moes, cache = _on_stack(kind, banks, monkeypatch, policy="rule", overlap=True, budget=8)
+    assert not cache._copy_fused_ok
+    steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs)
+    on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0,
+                       before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2))
+    _assert_same_outputs(on[0], off[0])
+    _assert_slots_hold_their_experts(cache)
+    assert int(cache.prefetch.stats[:, 0].sum()) > 0
+
+
+def _multi_graph_run(moes, cache, schedule, before_step=None):
+    """The graph runner's sequence: per graph size (largest first) an eager warm-up, a capture into
+    one shared pool and a reset; then ``schedule`` of (bs, xs, replay) steps, replays and eager
+    steps interleaved."""
+    _fresh_ctx(_batch=SimpleNamespace(is_prefill=False))
+
+    def step(xs):
+        return [moe.forward(x.clone()) for moe, x in zip(moes, xs)]
+
+    graphs, pool = {}, None
+    for bs in (2, 1):
+        static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes]
+        step(static)
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, pool=pool):
+            outs = step(static)
+        cache.reset()
+        pool = g.pool() if pool is None else pool
+        graphs[bs] = (g, static, outs)
+    results = []
+    for i, (bs, xs, replay) in enumerate(schedule):
+        if before_step is not None:
+            before_step(i)
+        if replay:
+            g, static, outs = graphs[bs]
+            for s, x in zip(static, xs):
+                s.copy_(x)
+            g.replay()
+        else:
+            outs = step(xs)
+        if cache.audit is not None:
+            cache.audit.end_decode_step(cache)
+        results.append([o.clone() for o in outs])
+    torch.cuda.synchronize()
+    return results
+
+
+@requires_cuda
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_prefetch_on_interleaves_graph_sizes_and_eager_steps(policy, monkeypatch):
+    """bs 1 and bs 2 graphs captured into one pool (as the graph runner does), replayed in an
+    arbitrary order with eager steps between them, including an eager bs 9 step that forks nothing:
+    every output equals off's to the bit and every held slot holds its expert."""
+    from freetoken.moe.prefetch import MAX_ROWS
+
+    banks = _prefetch_banks("nvfp4")
+    gen = torch.Generator().manual_seed(41)
+    order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (60,), generator=gen), torch.randint(0, 2, (60,), generator=gen))]
+    sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1}
+    xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=60, seed=42 + bs)) for bs in sizes.values()}
+    schedule = []
+    for b, r in order:
+        bs = sizes[b]
+        schedule.append((bs, next(xs_by_bs[bs]), r and bs <= 2))
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _multi_graph_run(moes, cache, schedule)
+    moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8)
+    steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=43)
+    on = _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps))
+    _assert_same_outputs(on, off)
+    _assert_slots_hold_their_experts(cache)
+    assert cache.prefetch._inflight == [None] * _PF_LAYERS
+    assert int(cache.prefetch.stats[:, 0].sum()) > 0 and int(cache.prefetch.stats[:, 1].sum()) > 0
+
+
+@requires_cuda
+@pytest.mark.parametrize("seed", range(4))
+@pytest.mark.parametrize("rows", ["bs1", "bs2", "bs2_same"])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_prefetch_on_fuzzes_per_layer_stream_delays(seed, rows, policy, monkeypatch):
+    """Eager fuzz: before every layer a random choice of sleeps on the predictor stream, the prefetch
+    copy stream, the compute stream before the GEMM and the shared expert's side stream, with up to
+    12 candidates a layer (most of the 16 slots re-installed every call) and, for bs2_same, two
+    identical rows. Outputs equal off's, held slots hold their experts, and the final maps and
+    stats (all but LATE) equal those of the same run without delays."""
+    from freetoken.moe.prefetch import LATE, NUM_STAT_COLS
+
+    bs = 1 if rows == "bs1" else 2
+    inputs = _prefetch_inputs(bs, steps=60, seed=50 + seed)
+    if rows == "bs2_same":
+        inputs = [[x[:1].expand(2, -1).contiguous() for x in xs] for xs in inputs]
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy)
+    off = _prefetch_run(moes, cache, bs, False, inputs, sleep=0)
+    stall = _stall_hooks(monkeypatch)
+    gen = torch.Generator().manual_seed(seed)
+    choices = torch.randint(0, 16, (len(inputs), _PF_LAYERS), generator=gen).tolist()
+    runs = []
+    for delayed in (False, True):
+        moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=12)
+        steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=60 + seed)
+        pf = cache.prefetch
+        load = _with_override(cache, steps)
+        at = {"step": -1}
+
+        def before_step(i, load=load):
+            at["step"] = i
+            load(i)
+
+        for layer, moe in enumerate(moes):
+            real = moe.forward
+
+            def forward(x, layer=layer, real=real, pf=pf):
+                bits = choices[at["step"]][layer] if delayed and at["step"] >= 0 else 0
+                # layer L's fork runs layer L+1's prediction and copy: these delays hit them
+                pf.delay_predict_cycles = _PF_DELAY_CYCLES // 2 if bits & 1 else 0
+                pf.delay_copy_cycles = _PF_DELAY_CYCLES // 2 if bits & 2 else 0
+                stall["gemm"] = _PF_DELAY_CYCLES // 2 if bits & 4 else 0
+                stall["shared"] = _PF_DELAY_CYCLES // 2 if bits & 8 else 0
+                return real(x)
+
+            monkeypatch.setattr(moe, "forward", forward)
+        on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0, before_step=before_step)
+        _assert_same_outputs(on[0], off[0])
+        _assert_slots_hold_their_experts(cache)
+        runs.append((on[1][:3], pf.stats.cpu()))
+    (maps, stats), (delayed_maps, delayed_stats) = runs
+    for a, b in zip(maps, delayed_maps):
+        assert torch.equal(a, b)
+    cols = [c for c in range(NUM_STAT_COLS) if c != LATE]
+    assert torch.equal(stats[:, cols], delayed_stats[:, cols])
+
+
+# FREETOKEN_MOE_PREFETCH_VERIFY: the debug slot checks on the same stacks, quiet on a correct decode
+# and loud on each kind of fault they exist for
+_VERIFY_STEPS = 60
+
+
+def _verify_state(cache):
+    """(counters summed over layers, kept records) of the stack's verifier."""
+    from freetoken.moe.verify import RING
+
+    v = cache.verify
+    torch.cuda.synchronize()
+    return v.counters.cpu().sum(0).tolist(), v.ring[: min(int(v.cursor), RING)].cpu().tolist()
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+@pytest.mark.parametrize("verify", ["full", "meta"])
+def test_moe_prefetch_verify_is_quiet_on_a_correct_decode(graph, bs, prefetch, verify, monkeypatch):
+    """FREETOKEN_MOE_PREFETCH_VERIFY checks every decode layer call and finds nothing on a correct
+    decode: prefetch off, and prefetch on with candidates the GEMMs read, late copies and a late
+    predictor. The checks only read, so the outputs equal those of verify off to the bit."""
+    from freetoken.moe.prefetch import USEFUL
+    from freetoken.moe.verify import CHECKS
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(bs, steps=_VERIFY_STEPS, seed=70)
+    runs = []
+    for mode in ("off", verify):
+        moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, verify_mode=mode)
+        before_step = None
+        if prefetch == "on":
+            steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 70)
+            before_step = _with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2,
+                                         delay_predict_cycles=_PF_DELAY_CYCLES // 4)
+        runs.append(_prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step))
+    _assert_same_outputs(runs[1][0], runs[0][0])
+    counts, records = _verify_state(cache)
+    assert counts[CHECKS] == (_VERIFY_STEPS + 1) * _PF_LAYERS  # the eager warm-up is checked too
+    assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records
+    assert (cache.verify.scratch is None) == (verify == "meta")
+    if prefetch == "on":
+        assert int(cache.prefetch.stats[:, USEFUL].sum()) > 0, "the GEMMs must have read prefetched slots"
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+@pytest.mark.parametrize("path", list(_PREFILL_PATHS))
+def test_moe_prefetch_verify_is_quiet_across_prefills(graph, prefetch, path, monkeypatch):
+    """Every request prefills before it decodes: the whole-layer materialize, the overlap double
+    buffers (with the hit-D2D gather) and small prefills rewrite slots between decode steps, and the
+    decode checks after each of them still find nothing."""
+    import freetoken.layers.moe as layers_moe
+    from freetoken.moe.verify import CHECKS
+
+    cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path]
+    monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small)
+    banks = _prefetch_banks("nvfp4")
+    decode = _prefetch_inputs(2, steps=18, seed=75)
+    gen = torch.Generator(device="cuda").manual_seed(76)
+    prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+                 for _ in range(_PF_LAYERS)] for _ in range(2)]
+    schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])]
+                + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])]
+                + [("decode", xs) for xs in decode[12:]])
+    moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, cache_size=cache_size,
+                                  verify_mode="full", **cache_kw)
+    before_step = None
+    if prefetch == "on":
+        steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=77)
+        before_step = _with_override(cache, steps)
+    _mixed_run(moes, cache, 2, graph, schedule, before_step=before_step)
+    counts, records = _verify_state(cache)
+    if path == "overlap_d2d":
+        assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots"
+    assert counts[CHECKS] == (len(decode) + 1) * _PF_LAYERS
+    assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+def test_moe_prefetch_verify_catches_a_missing_copy_join(graph, monkeypatch):
+    """Negative control: without the compute stream's wait on the prefetch copy (and with that copy
+    late), GEMMs read prefetched slots before their bytes land. The byte check before the GEMM flags
+    them, and the records place the slot in the layer's prefetch copy plan."""
+    from freetoken.moe.prefetch import ExpertPrefetcher
+    from freetoken.moe.verify import BYTES_PRE, META_BAD, PRE_BAD, R_KIND, R_SOURCE, SRC_PREFETCH
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=40, seed=7)
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, verify_mode="full")
+
+    def no_wait(self, layer_id):
+        self._inflight[layer_id] = None
+
+    monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait)
+    if graph:
+        # a capture needs the copy stream joined back: after the last layer, past every GEMM it races
+        last, copy_stream = moes[-1].forward, cache.prefetch.copy_stream
+
+        def forward(x):
+            out = last(x)
+            torch.cuda.current_stream().wait_stream(copy_stream)
+            return out
+
+        monkeypatch.setattr(moes[-1], "forward", forward)
+    steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11)
+    # ~2.5 ms a copy: the copy stream falls ever further behind even an eager host's slow enqueue
+    _prefetch_run(moes, cache, 1, graph, inputs, sleep=0,
+                  before_step=_with_override(cache, steps, delay_copy_cycles=5 * _PF_DELAY_CYCLES))
+    counts, records = _verify_state(cache)
+    assert counts[PRE_BAD] > 0 and counts[META_BAD] == 0
+    assert any(r[R_KIND] == BYTES_PRE and r[R_SOURCE] & SRC_PREFETCH for r in records)
+
+
+def _on_call(layer, call, action):
+    """Wrap ExpertVerifier.before_gemm so ``action(verifier, cache, slots, real)`` replaces it on the
+    ``call``-th decode call of ``layer`` (call 0 is the eager warm-up)."""
+    from freetoken.moe.verify import ExpertVerifier
+
+    real, seen = ExpertVerifier.before_gemm, {}
+
+    def before_gemm(self, cache, slots, next_plan):
+        at = self._call[0]
+        seen[at] = seen.get(at, -1) + 1
+        if (at, seen[at]) == (layer, call):
+            return action(self, cache, slots, lambda: real(self, cache, slots, next_plan))
+        return real(self, cache, slots, next_plan)
+
+    return before_gemm
+
+
+@requires_cuda
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+def test_moe_prefetch_verify_catches_a_corrupt_slot(prefetch, monkeypatch):
+    """Negative control: bytes flipped in one bank of one routed slot right before the checks are
+    flagged before the GEMM, with the call, layer, slot, expert, bank and byte offset of the flip."""
+    from freetoken.moe.verify import (BYTES_PRE, META_BAD, PRE_BAD, R_BANK, R_CALL, R_EXPERT, R_FIRST_BYTE,
+                                      R_KIND, R_LAYER, R_SLOT, ExpertVerifier)
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full")
+    hit = {}
+
+    def corrupt(verifier, cache, slots, real):
+        torch.cuda.synchronize()
+        hit.update(slot=int(slots.view(-1)[0]), expert=int(verifier.ids[0]))
+        cache.bank_views()[3][hit["slot"]].view(torch.uint8).view(-1)[100:108].bitwise_not_()
+        real()
+
+    monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(2, 10, corrupt))
+    _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=20, seed=71), sleep=0)
+    counts, records = _verify_state(cache)
+    assert counts[PRE_BAD] >= 1 and counts[META_BAD] == 0
+    first = records[0]
+    assert [first[f] for f in (R_KIND, R_CALL, R_LAYER, R_SLOT, R_EXPERT, R_BANK, R_FIRST_BYTE)] == [
+        BYTES_PRE, 10, 2, hit["slot"], hit["expert"], 3, 100]
+
+
+@requires_cuda
+@pytest.mark.parametrize("verify", ["full", "meta"])
+def test_moe_prefetch_verify_catches_a_wrong_slot_owner(verify, monkeypatch):
+    """Negative control: a routed slot's id_of_slot names another layer's expert while the checks
+    before the GEMM run (restored right after them): exactly that entry is flagged, with the owner it
+    saw, and the checks after the GEMM find the maps whole again."""
+    from freetoken.moe.verify import (CHECKS, META_BAD, META_PRE, R_CALL, R_EXPERT, R_KIND, R_LAYER, R_MAPPED,
+                                      R_OWNER, R_ROW, R_SLOT, ExpertVerifier)
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, verify_mode=verify)
+    hit = {}
+
+    def corrupt(verifier, cache, slots, real):
+        torch.cuda.synchronize()
+        s, e = int(slots.view(-1)[1]), int(verifier.ids[1])
+        hit.update(slot=s, expert=e, saved=int(cache.id_of_slot[s]))
+        cache.id_of_slot[s] = 3 * _PF_EXPERTS + e  # in range: an ensure reading it stays in bounds
+        real()
+        cache.id_of_slot[s] = hit["saved"]
+
+    monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(1, 7, corrupt))
+    _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=72), sleep=0)
+    counts, records = _verify_state(cache)
+    assert hit["saved"] == _PF_EXPERTS + hit["expert"]
+    assert counts[META_BAD] == 1 and sum(counts[CHECKS + 1 :]) == 1 and len(records) == 1
+    rec = records[0]
+    assert [rec[f] for f in (R_KIND, R_CALL, R_LAYER, R_ROW, R_SLOT, R_EXPERT, R_OWNER, R_MAPPED)] == [
+        META_PRE, 7, 1, 1, hit["slot"], hit["expert"], 3 * _PF_EXPERTS + hit["expert"], hit["slot"]]
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+def test_moe_prefetch_verify_catches_a_write_during_the_gemm(graph, prefetch, monkeypatch):
+    """Negative control: another stream rewrites a routed slot while the GEMM runs, after the checks
+    before it (a sleep on the compute stream holds the GEMM until the write lands). The byte check
+    after the GEMM flags it; the check before it saw nothing."""
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.moe.verify import (BYTES_POST, BYTES_PRE, POST_BAD, R_BANK, R_CALL, R_FIRST_BYTE, R_KIND, R_LAYER,
+                                      R_SLOT, ExpertVerifier)
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full")
+    side, real, armed = torch.cuda.Stream(), OffloadMoELayer._expert_gemm, {"on": False}
+    slot = torch.zeros((1,), dtype=torch.int64, device="cuda")
+    offset = torch.arange(16, device="cuda")
+
+    def gemm(self, cache, hidden, weights, topk_ids, *, is_prefill, **kw):
+        if armed["on"] and self.layer_id == 3 and not is_prefill:
+            slot.copy_(topk_ids.view(-1)[:1])  # device side, so a captured graph rewrites this step's slot
+            side.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(side):
+                torch.cuda._sleep(_PF_DELAY_CYCLES // 20)
+                row = cache.bank_views()[0].view(cache.cache_size, -1)
+                row.index_put_((slot, offset), row[slot, offset].bitwise_not())
+            torch.cuda._sleep(_PF_DELAY_CYCLES)
+            out = real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw)
+            torch.cuda.current_stream().wait_stream(side)  # a capture must join the side stream back
+            return out
+        return real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw)
+
+    monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm)
+    real_run = ExpertVerifier.before_gemm
+
+    def arm_at_call_5(self, cache, slots, next_plan):
+        # graph: the arm must be set while capturing, so the rewrite replays every step; eager: once at call 5
+        armed["on"] = self._call[0] == 3 and (graph or int(self.calls[3]) == 5)
+        return real_run(self, cache, slots, next_plan)
+
+    monkeypatch.setattr(ExpertVerifier, "before_gemm", arm_at_call_5)
+    _prefetch_run(moes, cache, 1, graph, _prefetch_inputs(1, steps=8, seed=73), sleep=0)
+    torch.cuda.synchronize()
+    counts, records = _verify_state(cache)
+    assert counts[POST_BAD] >= 1
+    first = records[0]
+    assert [first[f] for f in (R_KIND, R_LAYER, R_BANK, R_FIRST_BYTE)] == [BYTES_POST, 3, 0, 0]
+    assert not any(r[R_KIND] == BYTES_PRE and (r[R_CALL], r[R_SLOT]) == (first[R_CALL], first[R_SLOT]) for r in records)
+
+
+@requires_cuda
+def test_moe_prefetch_verify_off_builds_and_launches_nothing(monkeypatch):
+    """With the flag off there is no verifier and a decode launches only what it launches without
+    one; meta and full only add launches (their kernels, the id copy and the host-row gather)."""
+    from collections import Counter
+
+    from torch.profiler import ProfilerActivity, profile
+
+    from freetoken.env import ENV
+
+    monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0")
+    banks = _prefetch_banks("bf16")
+    inputs = _prefetch_inputs(1, steps=2)
+
+    def kernels(verify):
+        moes, cache = _prefetch_stack("bf16", banks, "off", monkeypatch, verify_mode=verify)
+        _prefetch_run(moes, cache, 1, False, inputs[:1])  # warm up and compile outside the trace
+        with profile(activities=[ProfilerActivity.CUDA]) as prof:
+            _prefetch_run(moes, cache, 1, False, inputs)
+        return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA)
+
+    cache, off = kernels(None)
+    assert cache.verify is None and cache.verify_mode == "off"
+    ours = ("_plan_kernel", "_compare_kernel", "_tally_kernel")
+    assert not any(k in name for name in off for k in ours)
+    for mode in ("meta", "full"):
+        _, on = kernels(mode)
+        assert not off - on, f"{mode} dropped a launch"
+        added = on - off
+        assert all(any(k in name for k in (*ours, "Memcpy", "fast_index_copy_multi")) for name in added), added
+        assert any("_compare_kernel" in name for name in added) == (mode == "full")
+
+
+@requires_cuda
+def test_moe_prefetch_verify_follows_a_rebuild(monkeypatch):
+    """A rebuild reallocates the slot banks (production rebuilds to its final slot count after start):
+    the verifier keeps its scratch and counts, points its byte checks at the new banks, and a
+    recaptured graph checks clean."""
+    from freetoken.moe.verify import CHECKS
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(2, steps=30, seed=74)
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, verify_mode="full")
+    v = cache.verify
+    scratch = v.scratch
+    _prefetch_run(moes, cache, 2, True, inputs, sleep=0)
+    cache.rebuild(_PF_EXPERTS + 4)
+    assert cache.verify is v and v.scratch is scratch
+    assert v.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()]
+    _prefetch_run(moes, cache, 2, True, inputs, sleep=0)
+    counts, records = _verify_state(cache)
+    assert counts[CHECKS] == 2 * (30 + 1) * _PF_LAYERS
+    assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == []
+
+
+# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): quiet and output-neutral on a correct decode with
+# prefills between the steps, and each injected persistent fault found with the history that explains it
+
+
+def _audit_records(cache):
+    """(totals, kept records) of the stack's slot audit."""
+    a = cache.audit
+    torch.cuda.synchronize()
+    return a.totals.cpu().tolist(), a.ring[: min(int(a.cursor), a.ring_n)].cpu().tolist()
+
+
+def _bad_slots(cache):
+    """Ground truth: the (slot, owner) pairs whose bytes differ from the owner's host rows."""
+    torch.cuda.synchronize()
+    held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu()
+    ids = cache.id_of_slot[held.cuda()].long().cpu()
+    bad = torch.zeros(held.numel(), dtype=torch.bool)
+    for (per_layer, _), view in zip(cache.banks, cache.bank_views()):
+        want = torch.stack(per_layer).flatten(0, 1)[ids].contiguous().view(torch.uint8).view(held.numel(), -1)
+        got = view[held.cuda()].cpu().contiguous().view(torch.uint8).view(held.numel(), -1)
+        bad |= (got != want).any(dim=1)
+    return {(int(s), int(o)) for s, o, b in zip(held, ids, bad) if b}
+
+
+def _spy_audit_kinds(monkeypatch):
+    """The event kinds the recorders were asked for, at enqueue time, each with the streams its
+    recorders ran on (a recorder follows its writer on the writer's stream)."""
+    from freetoken.moe.slot_audit import SlotAuditor
+
+    kinds, real = {}, SlotAuditor._launch
+
+    def launch(self, cache, kind, *args, **kwargs):
+        kinds.setdefault(kind, set()).add(torch.cuda.current_stream().cuda_stream)
+        return real(self, cache, kind, *args, **kwargs)
+
+    monkeypatch.setattr(SlotAuditor, "_launch", launch)
+    return kinds
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("bs", [1, 2])
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+@pytest.mark.parametrize("path", list(_PREFILL_PATHS))
+def test_moe_slot_audit_is_quiet_and_changes_no_bit(graph, bs, prefetch, path, monkeypatch):
+    """FREETOKEN_MOE_SLOT_AUDIT=1 records every slot writer (the decode and prefetch ensures and copies,
+    and between the decode steps the whole-layer materialize, the overlap double buffers with the
+    hit-D2D gather, or small prefills) and audits every held slot after each decode step: it finds
+    nothing, and the outputs, slot maps and slot bytes equal those of the audit off to the bit. Every
+    map writer it records ran on the compute stream but the prefill buffer invalidation (on the
+    prefill copy stream), and the prefetch copy on the prefetch copy stream."""
+    import freetoken.layers.moe as layers_moe
+    import freetoken.moe.offload_cache as offload_cache
+    from freetoken.moe import slot_audit as sa
+
+    cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path]
+    monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small)
+    if path == "overlap_d2d":
+        # every toy bank is under the 256 KiB small-bank floor, which would leave the hit gather idle
+        monkeypatch.setattr(offload_cache, "_SMALL_BANK_FEAT_BYTES", 8192)
+    banks = _prefetch_banks("nvfp4")
+    decode = _prefetch_inputs(bs, steps=18, seed=81)
+    gen = torch.Generator(device="cuda").manual_seed(82)
+    prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5
+                 for _ in range(_PF_LAYERS)] for _ in range(2)]
+    schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])]
+                + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])]
+                + [("decode", xs) for xs in decode[12:]])
+    kinds = _spy_audit_kinds(monkeypatch)
+    runs = []
+    for interval in (0, 1):
+        moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True,
+                                      cache_size=cache_size, slot_audit=interval, **cache_kw)
+        assert (cache.audit is None) == (interval == 0)
+        before_step = None
+        if prefetch == "on":
+            steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 83)
+            before_step = _with_override(cache, steps)
+        compute = _record_compute_streams(monkeypatch, moes)
+        out = _mixed_run(moes, cache, bs, graph, schedule, before_step=before_step)
+        runs.append((out, [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())]))
+    _assert_same_outputs(runs[1][0], runs[0][0])
+    for a, b in zip(runs[1][1], runs[0][1]):
+        assert torch.equal(a, b)
+    totals, records = _audit_records(cache)
+    assert records == [] and totals[sa.T_BAD] == 0, [cache.audit.format_record(r) for r in records]
+    assert totals[sa.T_MAP_RANGE] == totals[sa.T_BAD_PLAN] == totals[sa.T_UNCHECKED] == 0
+    assert totals[sa.T_AUDITS] == len(decode)
+    assert totals[sa.T_HELD] > 0 and totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] and totals[sa.T_EVENTS] > 0
+    expected = {sa.DEMAND_INSTALL, sa.DEMAND_COPY, sa.RESET}
+    if prefetch == "on":
+        expected |= {sa.PREFETCH_INSTALL, sa.PREFETCH_COPY}
+    expected |= {
+        "materialize": {sa.MATERIALIZE_CLEAR, sa.MATERIALIZE_INSTALL, sa.MATERIALIZE_COPY},
+        "overlap": {sa.INVALIDATE, sa.PREFILL_BUFFER},
+        "overlap_d2d": {sa.INVALIDATE, sa.PREFILL_SPLIT_H2D, sa.PREFILL_HIT_D2D},
+        "small": {sa.PREFILL_INSTALL, sa.PREFILL_COPY},
+    }[path]
+    assert set(kinds) == expected, (set(kinds) - expected, expected - set(kinds))
+    fence = set() if cache.prefill_copy_stream is None else {cache.prefill_copy_stream.cuda_stream}
+    for kind in (sa.DEMAND_INSTALL, sa.PREFETCH_INSTALL, sa.PREFILL_INSTALL, sa.MATERIALIZE_INSTALL,
+                 sa.MATERIALIZE_CLEAR, sa.RESET, sa.INVALIDATE):
+        assert kinds.get(kind, set()) <= (fence if kind == sa.INVALIDATE else compute), kind
+    if prefetch == "on":
+        assert kinds[sa.PREFETCH_COPY] == {cache.prefetch.copy_stream.cuda_stream}
+
+
+@requires_cuda
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_slot_audit_is_quiet_across_graph_sizes_and_eager_steps(policy, monkeypatch):
+    """The graph runner's shape: bs 1 and bs 2 graphs in one pool, replayed in any order with eager
+    steps between them (a bs 9 one forks no prefetch); audited after every step, with the decode
+    verifier on beside it, neither finds anything."""
+    from freetoken.moe import slot_audit as sa
+    from freetoken.moe.prefetch import MAX_ROWS
+    from freetoken.moe.verify import CHECKS
+
+    banks = _prefetch_banks("nvfp4")
+    gen = torch.Generator().manual_seed(90)
+    order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (40,), generator=gen), torch.randint(0, 2, (40,), generator=gen))]
+    sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1}
+    xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=40, seed=91 + bs)) for bs in sizes.values()}
+    schedule = [(sizes[b], next(xs_by_bs[sizes[b]]), r and sizes[b] <= 2) for b, r in order]
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1,
+                                  verify_mode="meta")
+    steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=92)
+    _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps))
+    totals, records = _audit_records(cache)
+    assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == len(schedule)
+    assert totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] > 0
+    counts, verify_records = _verify_state(cache)
+    assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and verify_records == []
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("policy", ["rule", "lru"])
+def test_moe_slot_audit_explains_a_skipped_prefetch_copy(graph, policy, monkeypatch):
+    """Fault injection: layer 2's prefetch copies copy nothing. Every slot the audit finds bad holds a
+    layer-2 expert whose history reads 'prefetch install, no prefetch copy', and every bad slot still
+    held at the end was found."""
+    from freetoken.moe import slot_audit as sa
+
+    monkeypatch.setattr(sa, "RING", 1024)
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=40, seed=84)
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1)
+    real = cache.copy_rows
+
+    def copy_rows(layer_id, dst, src, num, *, slim, **kw):
+        if kw.get("kind") == sa.PREFETCH_COPY and layer_id == 2:
+            return None  # the fault: no bytes move (and so nothing is recorded)
+        return real(layer_id, dst, src, num, slim=slim, **kw)
+
+    monkeypatch.setattr(cache, "copy_rows", copy_rows)
+    steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=85)
+    _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, before_step=_with_override(cache, steps))
+    truth = _bad_slots(cache)
+    totals, records = _audit_records(cache)
+    assert records and totals[sa.T_BAD] > 0 and totals[sa.T_NEW] == len(records)
+    a = cache.audit
+    for rec in records:
+        assert rec[sa.H_FLAGS] == sa.F_BYTES and rec[sa.H_OWNER] // _PF_EXPERTS == 2, a.format_record(rec)
+        # a layer-1 prefetch copy still in flight may land after the install and is named as such
+        assert a.diagnose(rec).startswith("prefetch install, no prefetch copy"), a.format_record(rec)
+    assert truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records}
+
+
+@requires_cuda
+@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+def test_moe_slot_audit_explains_a_demand_copy_from_the_wrong_layer(graph, prefetch, monkeypatch):
+    """Fault injection: layer 2's demand copies read layer 1's host banks (the right rows). The bad
+    slots are layer-2 demand installs whose last copy came from layer 1's source."""
+    from freetoken.moe import slot_audit as sa
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    monkeypatch.setattr(sa, "RING", 1024)
+    real = OffloadMoeCache.copy_missing
+
+    def copy_missing(self, layer_id=None):
+        if self._pending_src_layer == 2 and not self._pending_whole_layer:
+            return self.copy_rows(1, self.evict_slots, self.src_indices, self.num_indices, slim=False, kind=self._pending_kind)
+        return real(self, layer_id)
+
+    monkeypatch.setattr(OffloadMoeCache, "copy_missing", copy_missing)
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(1, steps=30, seed=86)
+    moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1)
+    _prefetch_run(moes, cache, 1, graph, inputs, sleep=0)
+    truth = _bad_slots(cache)
+    totals, records = _audit_records(cache)
+    assert records and truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records}
+    for rec in records:
+        owner = rec[sa.H_OWNER]
+        assert owner // _PF_EXPERTS == 2 and rec[sa.H_FLAGS] == sa.F_BYTES
+        e = owner % _PF_EXPERTS
+        assert cache.audit.diagnose(rec).startswith(
+            f"demand install of L2/e{e}, then demand copy from layer 1's source (L1/e{e}, #"), cache.audit.format_record(rec)
+
+
+@requires_cuda
+@pytest.mark.parametrize("prefetch", ["off", "on"])
+def test_moe_slot_audit_explains_bytes_overwritten_after_install(prefetch, monkeypatch):
+    """Fault injection: bytes of a held slot flipped behind every recorder's back. The next audit keeps
+    one record (bank, byte offset, bad words) whose history ends in the owner's install and copy, a
+    later audit counts it again without a second record, and a clean slot is forgotten."""
+    from freetoken.moe import slot_audit as sa
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1000)
+    _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=87), sleep=0)
+    a = cache.audit
+    a.scan(cache)
+    assert _audit_records(cache)[1] == []
+    slot = int((cache.id_of_slot >= 0).nonzero()[0])
+    owner = int(cache.id_of_slot[slot])
+    row = cache.bank_views()[3][slot].view(torch.uint8).view(-1)
+    row[100:108].bitwise_not_()
+    a.scan(cache)
+    a.scan(cache)
+    totals, records = _audit_records(cache)
+    assert len(records) == 1 and [totals[t] for t in (sa.T_AUDITS, sa.T_BAD, sa.T_NEW, sa.T_STILL)] == [3, 2, 1, 1]
+    rec = records[0]
+    assert [rec[f] for f in (sa.H_AUDIT, sa.H_SLOT, sa.H_OWNER, sa.H_FLAGS, sa.H_BANKS, sa.H_FIRST_BANK,
+                             sa.H_FIRST_BYTE, sa.H_BAD_WORDS)] == [1, slot, owner, sa.F_BYTES, 1 << 3, 3, 100, 2]
+    events = a.events(rec)
+    install = [ev for ev in events if ev["kind"] in sa.INSTALLS][-1]
+    copy = [ev for ev in events if ev["ring"] == "bytes"][-1]
+    assert a._flat(install) == owner == a._flat(copy) and copy["seq"] > install["seq"]
+    assert "then the bytes changed with no recorded write" in a.diagnose(rec)
+    row[100:108].bitwise_not_()
+    a.scan(cache)
+    row[100:108].bitwise_not_()
+    a.scan(cache)
+    totals, records = _audit_records(cache)
+    assert len(records) == 2 and records[1][sa.H_AUDIT] == 4 and totals[sa.T_BAD] == 3
+
+
+@requires_cuda
+def test_moe_slot_audit_explains_stale_and_missing_map_entries(monkeypatch):
+    """Fault injection on the maps: slot_for_id still names a slot for its previous owner (an eviction
+    that forgot to clear it: routing that expert would read another's bytes), and a held slot's owner
+    has lost its entry. Both are found and the stale one is traced to the install it outlived."""
+    from freetoken.moe import slot_audit as sa
+
+    banks = _prefetch_banks("nvfp4")
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, slot_audit=1000)
+    _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=88), sleep=0)
+    a = cache.audit
+    flat = cache.slot_for_id.view(-1)
+    held = (cache.id_of_slot >= 0).nonzero().view(-1).tolist()
+    found = None
+    for s in held:
+        for ev in sorted(a.meta_hist[s].tolist()):
+            fid = ev[sa.E_LAYER] * _PF_EXPERTS + ev[sa.E_ROW]
+            if ev[sa.E_KIND] in sa.INSTALLS and int(flat[fid]) == -1:
+                found = (s, fid)
+        if found:
+            break
+    assert found, "some slot must have outlived an earlier owner"
+    stale_slot, stale_id = found
+    fwd_slot = next(t for t in held if t != stale_slot)
+    fwd_owner = int(cache.id_of_slot[fwd_slot])
+    flat[stale_id] = stale_slot
+    flat[fwd_owner] = -1
+    a.scan(cache)
+    totals, records = _audit_records(cache)
+    assert [totals[t] for t in (sa.T_BAD, sa.T_STALE, sa.T_FWD, sa.T_BYTES)] == [2, 1, 1, 0] and len(records) == 2
+    by_slot = {r[sa.H_SLOT]: r for r in records}
+    rec = by_slot[stale_slot]
+    assert rec[sa.H_FLAGS] == sa.F_STALE and rec[sa.H_STALE_ID] == stale_id
+    holder = a._name(int(cache.id_of_slot[stale_slot]))
+    text = a.diagnose(rec)
+    assert text.startswith(f"stale slot_for_id: {a._name(stale_id)} -> this slot, which holds {holder}; "
+                           f"{a._name(stale_id)} was installed here by "), text
+    assert text.endswith("without clearing its map entry"), text
+    rec = by_slot[fwd_slot]
+    assert rec[sa.H_FLAGS] == sa.F_FWD and rec[sa.H_MAPPED] == -1
+    assert a.diagnose(rec) == f"slot_for_id[{a._name(fwd_owner)}]=-1, not this slot"
+
+
+@requires_cuda
+def test_moe_slot_audit_off_builds_and_launches_nothing(monkeypatch):
+    """With the flag off there is no auditor and a decode launches only what it launches without one;
+    on, a decode only adds the recorders (and the audit its own kernels and host-row gathers)."""
+    from collections import Counter
+
+    from torch.profiler import ProfilerActivity, profile
+
+    from freetoken.env import ENV
+
+    monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0)
+    banks = _prefetch_banks("bf16")
+    inputs = _prefetch_inputs(1, steps=2)
+
+    def kernels(interval):
+        moes, cache = _prefetch_stack("bf16", banks, "on", monkeypatch, slot_audit=interval)
+        _prefetch_run(moes, cache, 1, False, inputs[:1])  # warm up and compile outside the trace
+        with profile(activities=[ProfilerActivity.CUDA]) as prof:
+            _prefetch_run(moes, cache, 1, False, inputs)
+        return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA)
+
+    cache, off = kernels(None)
+    assert cache.audit is None and cache.slot_audit == 0
+    ours = ("_record_kernel", "_bump_step_kernel", "_audit_")
+    assert not any(k in name for name in off for k in ours)
+    _, on = kernels(1)
+    assert not off - on, "the audit dropped a launch"
+    added = on - off
+    assert all(any(k in name for k in (*ours, "fast_index_copy_multi")) for name in added), added
+    assert any("_audit_report_kernel" in name for name in added)
+
+
+@requires_cuda
+def test_moe_slot_audit_follows_a_rebuild(monkeypatch):
+    """A rebuild reallocates the slot banks and maps: the auditor keeps its scratch and totals, starts a
+    new history sized for the new slot count, points its byte checks at the new banks, and a recaptured
+    graph audits clean."""
+    from freetoken.moe import slot_audit as sa
+
+    banks = _prefetch_banks("nvfp4")
+    inputs = _prefetch_inputs(2, steps=20, seed=89)
+    moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, slot_audit=1)
+    a = cache.audit
+    scratch = a.scratch
+    _prefetch_run(moes, cache, 2, True, inputs, sleep=0)
+    cache.rebuild(_PF_EXPERTS + 4)
+    assert cache.audit is a and a.scratch is scratch and a.num_slots == _PF_EXPERTS + 4
+    assert a.meta_hist.shape[0] == a.bytes_cur.shape[0] == _PF_EXPERTS + 4 and int(a.meta_cur.sum()) == 0
+    assert a.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()]
+    _prefetch_run(moes, cache, 2, True, inputs, sleep=0)
+    totals, records = _audit_records(cache)
+    assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == 2 * len(inputs)
+
+
 @requires_cuda
 def test_decoder_stack_prefill_and_decode(monkeypatch):
     """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights."""
diff --git a/tests/moe/ref_prefetch.py b/tests/moe/ref_prefetch.py
new file mode 100644
index 0000000000..c40d9351a8
--- /dev/null
+++ b/tests/moe/ref_prefetch.py
@@ -0,0 +1,22 @@
+"""CPU reference for the router-lookahead select and count (freetoken.moe.prefetch)."""
+from __future__ import annotations
+
+
+def ref_select(logits, resident, k: int, budget: int):
+    """Each row's top-k (ties to the lowest id), merged rank by rank, duplicates dropped; the first
+    ``budget`` non-resident ids, and the resident ids among the first ``budget`` distinct ones."""
+    rows = [sorted(range(len(row)), key=lambda e: (-row[e], e))[:k] for row in logits]
+    merged = []
+    for r in range(k):
+        for row in rows:
+            if row[r] not in merged:
+                merged.append(row[r])
+    sel = [e for e in merged if resident[e] < 0][:budget]
+    res = [e for e in merged[:budget] if resident[e] >= 0]
+    return sel, res
+
+
+def ref_count(sel, res, routed, misses: int, rows: int):
+    """``[issued, useful, resident_hits, calls, rows, misses]`` for one layer call."""
+    routed = set(routed)
+    return [len(sel), sum(e in routed for e in sel), sum(e in routed for e in res), 1, rows, misses]
diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py
index d7d8b0f2fb..15153c86d5 100644
--- a/tests/moe/ref_scored_cache.py
+++ b/tests/moe/ref_scored_cache.py
@@ -47,12 +47,13 @@ def _softplus2(self, x):
         g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0)
         return np.maximum(x, 0) + g
 
-    def _keys(self, layer, tok, now, pin_since=None):
+    def _keys(self, layer, tok, now, pin_since=None, lowpri=False):
         c = np.arange(self.S, dtype=np.int64)
         oid = self.id_of_slot
         held = oid >= 0
         if self.policy == 0:
-            key = (self.usage << self.slot_bits) | c
+            # lowpri: empty (0) < unconsumed prefetch (held, usage 0 -> 1) < resident (usage + 1)
+            key = (np.where(held, self.usage + 1, 0) << self.slot_bits) | c if lowpri else (self.usage << self.slot_bits) | c
         else:
             ids = np.where(held, oid, 0)
             lk = ids // self.E
@@ -68,16 +69,20 @@ def _keys(self, layer, tok, now, pin_since=None):
                 score = score + ((self.w_q4 * lcv) >> 4)
             assert np.abs(score).max() < 2**31  # the kernel computes it in int32
             key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c)
+            if lowpri:
+                key = np.where(held & (self.usage == 0), (1 << self.slot_bits) | c, key)
         evictable = self.usage != self.step
         if pin_since is not None:
             evictable &= self.usage <= pin_since
         return np.where(evictable, key, _KEY_MAX)
 
     def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None,
-               pin_since=None):
+               pin_since=None, lowpri=False):
         """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan.
 
-        ``pin_since`` (a ``step`` value) also pins every slot touched after it.
+        ``pin_since`` (a ``step`` value) also pins every slot touched after it. ``lowpri``
+        (FREETOKEN_MOE_PREFETCH=on) keys unconsumed prefetches between the empty slots and the
+        residents; ``self.useful`` is then the call's distinct hits on such slots.
 
         Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their
         row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask.
@@ -100,13 +105,15 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None,
                 self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1)
         q = np.asarray(ids, dtype=np.int64).reshape(-1) + base
         slots = self.slot_of_id[q]
-        self.usage[slots[slots >= 0]] = self.step
         uniq = np.unique(q)
+        held_at = self.slot_of_id[uniq]
+        self.useful = int(((held_at >= 0) & (self.usage[np.maximum(held_at, 0)] == 0)).sum())
+        self.usage[slots[slots >= 0]] = self.step
         missing = uniq[self.slot_of_id[uniq] < 0]  # ascending, as the kernel ranks them
         src = missing - base
         dst = np.empty(missing.size, np.int64)
         if missing.size:
-            keys = self._keys(layer, tok, now, pin_since)
+            keys = self._keys(layer, tok, now, pin_since, lowpri)
             part = np.argpartition(keys, missing.size - 1)[: missing.size]
             victims = part[np.argsort(keys[part])]
             for i, (e, v) in enumerate(zip(missing, victims)):
@@ -127,6 +134,31 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None,
                 self.ct[uniq] = now
         return out, src, dst
 
+    def prefetch(self, layer, sel, budget=None):
+        """``prefetch_ensure``: install ``sel``'s non-resident ids (-1 = padding, query order), at most
+        the first ``budget``, as held slots with usage 0, without bumping a clock or touching a hit;
+        the last demand call's slots (usage == step) stay pinned. Returns ``(src_ids, dst_slots)``,
+        the copy plan."""
+        base = layer * self.E
+        tok = self.tok if self.policy else 0
+        pos = layer - 1  # the kernel scores from the previous layer's call, where the prefetch runs
+        now = tok * self.L + pos
+        order = []
+        for e in sel:
+            if e >= 0 and e + base not in order:
+                order.append(int(e) + base)
+        missing = [i for i in order if self.slot_of_id[i] < 0]
+        keys = self._keys(pos, tok, now, lowpri=True)
+        n = min(len(missing), int((keys != _KEY_MAX).sum()), len(missing) if budget is None else budget)
+        victims = np.argsort(keys, kind="stable")[:n]
+        for e, v in zip(missing[:n], victims):
+            old = self.id_of_slot[v]
+            if old >= 0:
+                self.slot_of_id[old] = -1
+            self.id_of_slot[v] = e
+            self.slot_of_id[e] = v
+            self.usage[v] = 0
+        return np.asarray(missing[:n], np.int64) - base, victims.astype(np.int64)
 
     def materialize(self, layer):
         """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id."""
diff --git a/tests/moe/test_prefetch.py b/tests/moe/test_prefetch.py
new file mode 100644
index 0000000000..47c26066a4
--- /dev/null
+++ b/tests/moe/test_prefetch.py
@@ -0,0 +1,424 @@
+"""FREETOKEN_MOE_PREFETCH stage A: the router-lookahead select / count kernels against a Python
+reference, the mode switch, and the predictor stream's lifetime across resets and rebuilds.
+The qwen4_exp stack tests (bit-identical outputs, counters over a real decode) live in
+tests/models/qwen4_exp/test_skeleton.py."""
+
+from __future__ import annotations
+
+import weakref
+
+import pytest
+import torch
+
+from .ref_prefetch import ref_count, ref_select
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def _padded(ids, width):
+    return ids + [-1] * (width - len(ids))
+
+
+def _select(logits, resident, k, budget, width=16):
+    from freetoken.moe.prefetch import lookahead_select
+
+    sel = torch.full((width,), -7, dtype=torch.int32, device="cuda")
+    res = torch.full((width,), -7, dtype=torch.int32, device="cuda")
+    lookahead_select(
+        torch.tensor(logits, dtype=torch.bfloat16, device="cuda"),
+        torch.tensor(resident, dtype=torch.int32, device="cuda"),
+        sel, res, k=k, budget=budget,
+    )
+    return sel.tolist(), res.tolist()
+
+
+def test_mode_switch():
+    from freetoken.moe.prefetch import default_budget, resolve_mode
+
+    assert resolve_mode("off") == "off" and resolve_mode(" Measure ") == "measure" and resolve_mode("ON") == "on"
+    with pytest.raises(ValueError):
+        resolve_mode("lookahead")
+    assert (default_budget(True), default_budget(False)) == (3, 4)
+
+
+@requires_cuda
+def test_cache_builds_the_prefetcher_only_when_asked():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    def cache(**kw):
+        return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw)
+
+    assert cache(prefetch_mode="off").prefetch is None
+    assert cache(prefetch_mode="measure").prefetch.copy_stream is None
+    on = cache(prefetch_mode="on")
+    assert on.prefetch.mode == "on" and on.prefetch.cache.prefetch is on.prefetch and on.prefetch_on
+    # no reference cycle: dropping the cache frees its tensors at once
+    ref = weakref.ref(on)
+    del on
+    assert ref() is None
+    for mode in ("measure", "on"):
+        with pytest.raises(ValueError):
+            cache(prefetch_mode=mode, decode_target="hybrid")
+
+
+@requires_cuda
+def test_the_prefetch_env_flag_skips_caches_that_do_not_decode_on_the_gpu(monkeypatch):
+    # a process-wide flag must not break hybrid or CPU-decoded caches, as FREETOKEN_MOE_COPY_OVERLAP does not
+    from freetoken.env import ENV
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    monkeypatch.setattr(ENV.MOE_PREFETCH, "value", "on")
+    for target in ("hybrid", "cpu"):
+        c = OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_target=target)
+        assert c.prefetch is None and not c.prefetch_on
+    assert OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda")).prefetch_on
+
+
+@requires_cuda
+def test_select_known_answers():
+    # order by logit: 1, 4, 6, 2, 3, 7, 0, 5; experts 4 and 2 are resident
+    logits = [[0.1, 0.9, 0.5, 0.3, 0.8, 0.0, 0.7, 0.2]]
+    resident = [-1, -1, 7, -1, 3, -1, -1, -1]
+    assert _select(logits, resident, 5, 2) == (_padded([1, 6], 16), _padded([4], 16))
+    assert _select(logits, resident, 5, 4) == (_padded([1, 6, 3], 16), _padded([4, 2], 16))
+    # the top-k cap: 3 and 7 are out of a top-4 list
+    assert _select(logits, resident, 4, 4) == (_padded([1, 6], 16), _padded([4, 2], 16))
+    # all-equal logits rank by id
+    assert _select([[0.5] * 8], [-1] * 8, 8, 3) == (_padded([0, 1, 2], 16), _padded([], 16))
+    # bs 2: row 0 ranks 1, 4, 6 and row 1 ranks 4, 0, 5 -> merged 1, 4, 0, 6, 5; expert 0 resident
+    rows = [[0.0, 0.9, 0.1, 0.1, 0.8, 0.1, 0.7, 0.1], [0.8, 0.1, 0.1, 0.1, 0.9, 0.7, 0.1, 0.1]]
+    resident = [2, -1, -1, -1, -1, -1, -1, -1]
+    assert _select(rows, resident, 3, 3) == (_padded([1, 4, 6], 16), _padded([0], 16))
+    assert _select(rows, resident, 3, 16) == (_padded([1, 4, 6, 5], 16), _padded([0], 16))
+
+
+@requires_cuda
+@pytest.mark.parametrize("rows", [1, 2])
+@pytest.mark.parametrize("k,budget", [(3, 3), (16, 4), (16, 3), (8, 16)])
+def test_select_matches_reference(rows, k, budget):
+    gen = torch.Generator().manual_seed(rows * 100 + k + budget)
+    for trial in range(20):
+        # coarse values make bf16 ties common, so the tie rule is exercised
+        logits = (torch.randint(-24, 24, (rows, 512), generator=gen).float() / 8).tolist()
+        resident = torch.where(torch.rand(512, generator=gen) < 0.4, 1, -1).tolist()
+        want = ref_select(logits, resident, k, budget)
+        width = 2 * k
+        got = _select(logits, resident, k, budget, width=width)
+        assert got == (_padded(want[0], width), _padded(want[1], width)), trial
+
+
+@requires_cuda
+def test_count_known_answers():
+    from freetoken.moe.prefetch import prefetch_count
+
+    num_experts, layer = 8, 2
+    base = layer * num_experts
+    # slots 0, 1, 3 hold this layer's experts 3, 1, 6; slot 4 another layer's expert
+    id_of_slot = torch.tensor([base + 3, base + 1, -1, base + 6, 5, -1], dtype=torch.int32, device="cuda")
+    slots = torch.tensor([0, 1, 3], dtype=torch.int32, device="cuda")
+    sel = torch.tensor(_padded([3, 5], 6), dtype=torch.int32, device="cuda")
+    res = torch.tensor(_padded([1, 4], 6), dtype=torch.int32, device="cuda")
+    misses = torch.tensor([2], dtype=torch.int64, device="cuda")
+    counters = torch.zeros(6, dtype=torch.int64, device="cuda")
+    for _ in range(2):
+        prefetch_count(slots, id_of_slot, misses, sel, res, counters, id_base=base, rows=1)
+    want = ref_count([3, 5], [1, 4], [3, 1, 6], 2, 1)
+    assert counters.tolist() == [2 * v for v in want] == [4, 2, 2, 2, 2, 4]
+    # bs 2 with a duplicate routed expert counts it once
+    counters.zero_()
+    slots2 = torch.tensor([0, 1, 0, 3], dtype=torch.int32, device="cuda")
+    prefetch_count(slots2, id_of_slot, misses, sel, res, counters, id_base=base, rows=2)
+    assert counters.tolist() == ref_count([3, 5], [1, 4], [3, 1, 3, 6], 2, 2)
+
+
+@requires_cuda
+def test_prefetch_stream_never_aliases_a_pooled_stream():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"),
+        decode_copy_overlap=True, prefetch_mode="measure",
+    )
+    pred = cache.prefetch.stream.cuda_stream
+    pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)}
+    assert pred not in pooled and pred != torch.cuda.current_stream().cuda_stream
+    assert pred != cache.decode_copy_stream.cuda_stream
+
+
+@requires_cuda
+def test_prefetch_copy_stream_is_a_third_dedicated_stream():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"),
+        decode_copy_overlap=True, prefetch_mode="on",
+    )
+    pf = cache.prefetch
+    copy = pf.copy_stream.cuda_stream
+    pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)}
+    assert copy not in pooled and pf.stream.cuda_stream not in pooled
+    assert len({copy, pf.stream.cuda_stream, cache.decode_copy_stream.cuda_stream, torch.cuda.current_stream().cuda_stream}) == 4
+    # its CTAs go ahead of pending compute ones, so a late copy does not also wait for SMs
+    from cuda.bindings import runtime as cudart
+
+    _, _, greatest = cudart.cudaDeviceGetStreamPriorityRange()
+    assert cudart.cudaStreamGetPriority(copy)[1] == greatest < 0
+    assert cudart.cudaStreamGetPriority(pf.stream.cuda_stream)[1] == 0
+
+
+@requires_cuda
+def test_rebuild_and_reset_clear_the_counters():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), prefetch_mode="measure"
+    )
+    cache.set_bank_sources({
+        "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+        "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)],
+    })
+    pf = cache.prefetch
+    stream, sel = pf.stream, pf.sel
+    pf.counters.fill_(3)
+    cache.reset()
+    assert int(pf.counters.abs().sum()) == 0
+    pf.counters.fill_(5)
+    pf.take_window()
+    assert int(pf.totals.sum()) > 0
+    pf.counters.fill_(2)
+    pf._inflight[1] = torch.zeros(1)
+    cache.rebuild(6)
+    assert cache.prefetch is pf and pf.stream is stream and pf.sel is sel
+    assert int(pf.counters.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0
+    assert pf._inflight == [None, None]
+
+
+def _emit_prefetch(window, monkeypatch, *, debug=False, mode="measure"):
+    from types import SimpleNamespace
+
+    from freetoken.engine import engine as engine_mod
+    from freetoken.env import ENV
+
+    taken, lines = [], []
+    prefetch = SimpleNamespace(mode=mode, take_window=lambda: taken.append(1) or window)
+    engine = SimpleNamespace(moe_offload_cache=SimpleNamespace(prefetch=prefetch))
+    monkeypatch.setattr(ENV.MOE_PREFETCH_DEBUG, "value", debug)
+    # the freetoken loggers do not propagate, so caplog only sees them in some import orders
+    monkeypatch.setattr(engine_mod.logger, "info_rank0", lambda msg, *a, **k: lines.append(msg))
+    engine_mod.Engine._emit_prefetch_stats(engine)
+    return "\n".join(lines), taken
+
+
+def test_engine_reports_one_prefetch_window(monkeypatch):
+    from freetoken.engine.engine import MOE_STATS_INTERVAL
+
+    window = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]])
+    out, taken = _emit_prefetch(window, monkeypatch)
+    assert taken == [1]
+    assert f"MoE prefetch measure ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out
+    assert "per layer" not in out
+    out, _ = _emit_prefetch(window, monkeypatch, debug=True)
+    assert "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" in out and "L0" not in out
+    assert _emit_prefetch(torch.zeros((3, 6), dtype=torch.int64), monkeypatch)[0] == ""
+
+
+def test_engine_reports_the_on_mode_window(monkeypatch):
+    from freetoken.engine.engine import MOE_STATS_INTERVAL
+
+    # issued, useful, resident_hits, calls, rows, misses, late, copied
+    window = torch.tensor([[0] * 8, [30, 12, 0, 10, 10, 18, 2, 10], [40, 20, 0, 10, 10, 20, 0, 6]])
+    out, _ = _emit_prefetch(window, monkeypatch, mode="on")
+    assert f"MoE prefetch on ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out
+    assert "misses/layer=1.90, coverage=0.457, late=0.125 of 16 copies" in out and "resident_hits" not in out
+
+
+def test_summary_rates():
+    from freetoken.moe.prefetch import format_per_layer, format_summary, summarize
+
+    # layer 0 never predicts; layers 1 and 2 over 10 bs-1 steps
+    counters = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]])
+    stats = summarize(counters)
+    assert stats["layer_calls"] == 20 and stats["tokens"] == 10
+    assert stats["issued_per_layer"] == 3.5 and stats["useful_per_layer"] == 1.6
+    assert stats["precision"] == 32 / 70 and stats["coverage"] == 32 / 70
+    assert stats["useful_per_token"] == 3.2 and stats["resident_hits_per_layer"] == 0.5
+    assert "useful/token=3.2" in format_summary(stats)
+    assert format_per_layer(counters) == "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00"
+
+
+# FREETOKEN_MOE_PREFETCH_VERIFY (moe/verify.py): the mode switch, where the verifier is built, and
+# the engine's report; the checks themselves run on the qwen4_exp stacks in test_skeleton.py
+
+
+def test_verify_mode_switch():
+    from freetoken.moe.verify import resolve_mode
+
+    assert [resolve_mode(m) for m in ("0", "", "off", "1", " True ", "full", "META")] == [
+        "off", "off", "off", "full", "full", "full", "meta"]
+    with pytest.raises(ValueError):
+        resolve_mode("bytes")
+
+
+@requires_cuda
+def test_cache_builds_the_verifier_only_for_gpu_decode(monkeypatch):
+    from freetoken.env import ENV
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    def cache(**kw):
+        return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw)
+
+    monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0")
+    assert cache().verify is None and cache(verify_mode="off").verify is None
+    full = cache(verify_mode="full", prefetch_mode="on")
+    assert full.verify.mode == "full" and cache(verify_mode="meta").verify.mode == "meta"
+    # the verifier holds no reference to its cache
+    ref = weakref.ref(full)
+    del full
+    assert ref() is None
+    with pytest.raises(ValueError):
+        cache(verify_mode="full", decode_target="hybrid")
+    monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "1")
+    for target in ("hybrid", "cpu"):
+        assert cache(decode_target=target).verify is None
+    assert cache().verify.mode == "full"
+
+
+@requires_cuda
+def test_engine_reports_the_verify_window_and_session(monkeypatch):
+    from types import SimpleNamespace
+
+    from freetoken.engine import engine as engine_mod
+    from freetoken.engine.engine import MOE_STATS_INTERVAL
+    from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.moe.verify import BYTES_PRE, CHECKS, NUM_FIELDS, PRE_BAD
+
+    cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=4, device=torch.device("cuda"), verify_mode="full")
+    cache.set_bank_sources({
+        "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)],
+        "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)],
+    })
+    v = cache.verify
+    lines = []
+    for level in ("info_rank0", "warning_rank0"):
+        monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg)))
+    engine = SimpleNamespace(moe_offload_cache=cache)
+
+    v.counters[:, CHECKS] = 256
+    engine_mod.Engine._emit_verify_stats(engine)
+    assert lines == [("info_rank0", f"MoE verify full ({MOE_STATS_INTERVAL} decode steps): 768 layer calls checked, "
+                                    "meta_bad=0, pre_bad=0, post_bad=0, post_meta_bad=0; session bad=0")]
+    lines.clear()
+    v.counters[:, CHECKS] = 256
+    v.counters[1, PRE_BAD] = 1
+    # call 7, layer 1, row 3 of a top-2 call, expert 2, slot 1, bank down, before the GEMM; slot 1 holds
+    # layer 1's expert 2, usage 9 at lru step 9, demand-copied; bytes 64.. differ in 4 words
+    record = [7, 1, 3, 2, 1, 1, BYTES_PRE, 6, 1, 9, 9, 1, 64, 4]
+    assert len(record) == NUM_FIELDS
+    v.ring[0] = torch.tensor(record)
+    v.cursor.fill_(1)
+    v._top_k = 2
+    engine_mod.Engine._emit_verify_stats(engine)
+    assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"]
+    assert "pre_bad=1" in lines[0][1] and "(meta/pre/post/post_meta by layer: L1=0/1/0/0); session bad=1" in lines[0][1]
+    assert lines[1][1] == (
+        "MoE verify record: bytes_pre call=7 layer=1 row=3 (token 1, rank 1) expert=2 slot=1 bank=down "
+        "first_bad_byte=64 bad_words=4 | slot now holds L1/e2, slot_for_id=1, usage=9, lru_step=9, slot in: demand copy")
+    lines.clear()
+    engine_mod.Engine._emit_verify_stats(engine)  # nothing checked and no new record: silent
+    assert lines == []
+    session = v.report_session()
+    assert session[0] == ("warning", "MoE verify full (session): 1536 layer calls checked, meta_bad=0, pre_bad=1, "
+                                     "post_bad=0, post_meta_bad=0 (meta/pre/post/post_meta by layer: L1=0/1/0/0)")
+    assert session[1:] == [("warning", "MoE verify record: " + v.format_record(record))]
+
+
+# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): the switch, where the auditor is built, and the
+# engine's report; the recorders and the audit run on the qwen4_exp stacks in test_skeleton.py
+
+
+def test_slot_audit_interval_switch(monkeypatch):
+    from freetoken.env import ENV
+    from freetoken.moe.slot_audit import resolve_interval
+
+    monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 64)
+    assert [resolve_interval(v) for v in (None, 0, 1, 8)] == [64, 0, 1, 8]
+    with pytest.raises(ValueError):
+        resolve_interval(-1)
+
+
+@requires_cuda
+def test_cache_builds_the_auditor_only_for_gpu_decode(monkeypatch):
+    from freetoken.env import ENV
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    def cache(**kw):
+        return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw)
+
+    monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0)
+    assert cache().audit is None and cache(slot_audit=0).audit is None
+    on = cache(slot_audit=8, prefetch_mode="on")
+    assert on.audit.interval == 8 and on.slot_audit == 8
+    ref = weakref.ref(on)  # the auditor holds no reference to its cache
+    del on
+    assert ref() is None
+    with pytest.raises(ValueError):
+        cache(slot_audit=8, decode_target="hybrid")
+    monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 16)
+    for target in ("hybrid", "cpu"):
+        assert cache(decode_target=target).audit is None
+    assert cache().audit.interval == 16
+    monkeypatch.setenv("FREETOKEN_SKIP_FAST_INDEX_COPY", "1")
+    assert cache().audit is None, "copies that move nothing cannot be audited"
+
+
+@requires_cuda
+def test_engine_reports_the_audit_window_and_session(monkeypatch):
+    from types import SimpleNamespace
+
+    from freetoken.engine import engine as engine_mod
+    from freetoken.engine.engine import MOE_STATS_INTERVAL
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=6, device=torch.device("cuda"), slot_audit=2)
+    cache.set_bank_sources({
+        "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)],
+        "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)],
+    })
+    cache.reset()
+    lines = []
+    for level in ("info_rank0", "warning_rank0"):
+        monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg)))
+    engine = SimpleNamespace(moe_offload_cache=cache)
+    engine_mod.Engine._emit_audit_stats(engine)
+    assert lines == [], "no audit ran: nothing to say"
+
+    # one decode step of layer 1 routing experts 2 and 0: two demand installs and their copy
+    ids = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda")
+    cache.ensure_experts(1, ids)
+    cache.copy_missing()
+    audit = cache.audit
+    assert not audit.end_decode_step(cache) and audit.end_decode_step(cache), "every 2 decode steps"
+    engine_mod.Engine._emit_audit_stats(engine)
+    assert lines == [("info_rank0", f"MoE slot audit ({MOE_STATS_INTERVAL} decode steps, every 2): 1 audits, 2 held slots "
+                                    "an audit (2 byte-checked), bad slots=0 (bytes=0, map=0, stale slot_for_id=0, owner out "
+                                    "of range=0), new=0, repeat=0, events=4; session: 1 audits, 0 bad (slot, owner) pairs")]
+    lines.clear()
+    slot = int(ids[0, 0])
+    cache.bank_views()[1][slot].view(torch.uint8).view(-1)[8:12].bitwise_not_()
+    audit.scan(cache)
+    engine_mod.Engine._emit_audit_stats(engine)
+    assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"]
+    assert "1 audits, 2 held slots an audit (2 byte-checked), bad slots=1 (bytes=1," in lines[0][1]
+    assert lines[1][1].startswith(
+        f"MoE slot audit record: audit=1 step=0 lru=1 slot={slot} holds L1/e2 (slot_for_id={slot}, usage=1) | bytes differ in "
+        "down (first at down+8, 1 bad words) | diagnosis: demand install and demand copy of L1/e2 (#"), lines[1][1]
+    assert "then the bytes changed with no recorded write" in lines[1][1]
+    # the kernel ranks misses by id: e0 installs first (#0), e2 second (#1); the copy records after both
+    assert lines[1][1].endswith("history (1 map / 1 byte events in all): #1 demand install L1/e2 @step 0/lru 1 (then held "
+                                "L1/e2, usage 1); #3 demand copy L1/e2 @step 0/lru 1 (slot held L1/e2, usage 1)"), lines[1][1]
+    lines.clear()
+    engine_mod.Engine._emit_audit_stats(engine)
+    assert lines == []
+    session = audit.report_session()
+    assert session[0][0] == "warning" and session[0][1].startswith("MoE slot audit (session, every 2 decode steps): 2 audits")
+    assert session[1:] == [("warning", "MoE slot audit record: " + audit.format_record(audit.ring[0].tolist()))]
diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py
index db734180b6..6bfa84a5ed 100644
--- a/tests/moe/test_scored_ensure.py
+++ b/tests/moe/test_scored_ensure.py
@@ -6,6 +6,7 @@
 (experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows.
 """
 import os
+from types import SimpleNamespace
 from pathlib import Path
 
 import numpy as np
@@ -620,3 +621,321 @@ def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy):
     for name, buf in guarded.items():
         assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}"
     assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s
+
+
+# (h) FREETOKEN_MOE_PREFETCH=on: prefetch installs and low-priority keys ----------------------
+
+_PF_L, _PF_E, _PF_S, _PF_K, _PF_W = 4, 64, 100, 6, 16
+
+
+def _pf_cache(policy, cache_size=_PF_S, num_layers=_PF_L, num_experts=_PF_E):
+    return _cache(policy, cache_size, num_layers, num_experts, prefetch_mode="on")
+
+
+def _gpu_prefetch(cache, layer, sel, stats_row, width=_PF_W, budget=None):
+    """prefetch_ensure on the current stream; returns (src_ids, dst_slots), the plan it wrote."""
+    from freetoken.moe.offload_kernels import prefetch_ensure_experts
+
+    dev = cache.id_of_slot.device
+    query = torch.tensor(list(sel) + [-1] * (width - len(sel)), dtype=torch.int32, device=dev)
+    dst = torch.full((width,), -7, dtype=torch.int32, device=dev)
+    src = torch.full((width,), -7, dtype=torch.int32, device=dev)
+    num = torch.full((1,), 99, dtype=torch.int64, device=dev)
+    ready = torch.ones(1, dtype=torch.int32, device=dev)
+    prefetch_ensure_experts(cache, layer, query, dst, src, num, stats_row, ready, budget=budget)
+    n = int(num)
+    assert int(ready) == 0, "prefetch_ensure clears the copy's ready flag"
+    assert dst[n:].tolist() == [-7] * (width - n)
+    return src[:n].cpu().numpy().astype(np.int64), dst[:n].cpu().numpy().astype(np.int64), num
+
+
+def _lowpri_slots(cache):
+    return set(((cache.id_of_slot >= 0) & (cache.usage == 0)).nonzero().view(-1).tolist())
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"])
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_prefetch_mode_matches_cpu_reference(policy, rows_per_step):
+    """prefetch_ensure and the low-priority demand ensure against the CPU reference victim for victim,
+    with adversarial candidates (next routed ids, resident ids, duplicates, padding, junk) and a
+    materialized prefill in between; the per-layer prefetch stats match the reference's counts."""
+    from freetoken.moe.offload_kernels import ensure_experts_scored
+    from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, LATE, MISSES, RESIDENT_HITS, ROWS, USEFUL
+
+    rng = np.random.default_rng(31 + rows_per_step)
+    rows = _zipf_rows(40, rows_per_step, seed=33, num_layers=_PF_L, num_experts=_PF_E, top_k=_PF_K)
+    logits = _random_logits(rows, seed=34, num_experts=_PF_E) if policy == "rule" else None
+    cache, ref = _pf_cache(policy), _ref(policy, _PF_S, _PF_L, _PF_E)
+    dev = cache.id_of_slot.device
+    stats = torch.zeros((_PF_L, 8), dtype=torch.int64, device=dev)
+    want = np.zeros((_PF_L, 8), np.int64)
+    ready = torch.zeros(1, dtype=torch.int32, device=dev)
+    for r in range(rows.shape[0]):
+        if r == 20:
+            cache.materialize_layer(2)
+            ref.materialize(2)
+            assert not {s for s in _lowpri_slots(cache) if s < _PF_E}, "materialize installs at usage step"
+        for layer in range(_PF_L):
+            ids = rows[r, layer]
+            plan = None
+            if layer and rng.random() < 0.9:
+                routed = list(dict.fromkeys(ids.ravel().tolist()))
+                held = [e for e in range(_PF_E) if ref.slot_of_id[layer * _PF_E + e] >= 0]
+                sel = rng.choice(routed, size=min(3, len(routed)), replace=False).tolist()
+                sel += rng.choice(_PF_E, size=4).tolist() + held[:2] + sel[:1] + [-1]
+                rng.shuffle(sel)
+                pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist())
+                budget = [None, 0, 1, 2, 4][int(rng.integers(5))]
+                src, dst, num = _gpu_prefetch(cache, layer, sel, stats[layer], budget=budget)
+                want_src, want_dst = ref.prefetch(layer, sel, budget)
+                np.testing.assert_array_equal(src, want_src)
+                np.testing.assert_array_equal(dst, want_dst)
+                assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call"
+                assert set(dst.tolist()) <= _lowpri_slots(cache)
+                want[layer, ISSUED] += len(want_src)
+                _assert_same_tables(cache, ref)
+                plan = num
+                ready.fill_(int(rng.random() < 0.5))  # a copy that has (1) or has not (0) landed
+            q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1)
+            lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16)
+            lowpri_before = _lowpri_slots(cache)
+            ensure_experts_scored(
+                cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True,
+                pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step),
+            )
+            out, src, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(),
+                                     thr=NEAR_MISS_THR, lowpri=True)
+            np.testing.assert_array_equal(q.cpu().numpy(), out)
+            _assert_same_tables(cache, ref)
+            # a demand hit makes a prefetched slot an ordinary resident, tagged with its owner's state
+            converted = lowpri_before & set(out.tolist())
+            assert all(int(cache.usage[s]) == int(cache.step) for s in converted)
+            if policy != "lru":
+                assert all(int(cache.evict_slot_owner[s]) == int(cache.id_of_slot[s]) for s in converted)
+            if plan is not None:
+                n = int(plan)
+                want[layer, [USEFUL, CALLS, ROWS, MISSES]] += [ref.useful, 1, rows_per_step, src.size]
+                want[layer, LATE] += int(n > 0 and int(ready) == 0)
+                want[layer, COPIED] += int(n > 0)
+    got = stats.cpu().numpy()
+    np.testing.assert_array_equal(got, want)
+    assert got[:, USEFUL].sum() > 0 and got[:, LATE].sum() > 0 and got[:, RESIDENT_HITS].sum() == 0
+    assert got[:, COPIED].sum() > got[:, LATE].sum()
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "rule"])
+@pytest.mark.parametrize("rows_per_step", [1, 2, 8])
+def test_prefetch_mode_matches_cpu_reference_at_production_geometry(policy, rows_per_step):
+    """The same victim-for-victim check at the deployment's shape: 48 layers x 512 experts, 1480
+    slots (a 2048-wide, 8-warp victim scan), top-10, up to 8 rows (every graph size that forks)
+    and 32-wide prefetch queries, with small-prefill ensures (pinned since their first chunk) and a
+    materialized layer in between. No prefetch ever evicts a slot of the last demand call."""
+    from freetoken.moe.offload_kernels import ensure_experts_scored
+
+    num_layers, num_experts, size, width = L, E, 1480, 32
+    rng = np.random.default_rng(71 + rows_per_step)
+    rows = _zipf_rows(4, rows_per_step, seed=72, num_layers=num_layers, num_experts=num_experts, top_k=TOP_K)
+    logits = _random_logits(rows, seed=73, num_experts=num_experts) if policy == "rule" else None
+    cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts)
+    dev = cache.id_of_slot.device
+    stats = torch.zeros((num_layers, 8), dtype=torch.int64, device=dev)
+    ready = torch.zeros(1, dtype=torch.int32, device=dev)
+    installs = 0
+    for r in range(rows.shape[0]):
+        if r == 2:
+            cache.materialize_layer(5)
+            ref.materialize(5)
+        for layer in range(num_layers):
+            ids = rows[r, layer]
+            plan = None
+            if layer:
+                routed = list(dict.fromkeys(ids.ravel().tolist()))
+                held = [e for e in range(num_experts) if ref.slot_of_id[layer * num_experts + e] >= 0]
+                sel = rng.choice(routed, size=min(6, len(routed)), replace=False).tolist()
+                sel += rng.choice(num_experts, size=16).tolist() + held[:4] + sel[:2] + [-1, -1]
+                sel = sel[:width]
+                rng.shuffle(sel)
+                pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist())
+                src, dst, plan = _gpu_prefetch(cache, layer, sel, stats[layer], width=width)
+                want_src, want_dst = ref.prefetch(layer, sel)
+                np.testing.assert_array_equal(src, want_src)
+                np.testing.assert_array_equal(dst, want_dst)
+                assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call"
+                installs += len(dst)
+                _assert_same_tables(cache, ref)
+            q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1)
+            lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16)
+            ensure_experts_scored(
+                cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True,
+                pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step),
+            )
+            out, _, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(),
+                                   thr=NEAR_MISS_THR, lowpri=True)
+            np.testing.assert_array_equal(q.cpu().numpy(), out)
+            _assert_same_tables(cache, ref)
+        # a short prefill of one layer between decode steps: chunked, pinned since its first chunk
+        layer = int(rng.integers(num_layers))
+        uniq = np.unique(rng.choice(num_experts, size=60, replace=False))
+        pin = cache.step.clone() if policy != "lru" else None
+        ref_pin = ref.step if policy != "lru" else None
+        for start in range(0, uniq.size, 32):
+            part = torch.from_numpy(uniq[start : start + 32].astype(np.int32)).to(dev)
+            cache.ensure_experts(layer, part, update_state=False, pin_since=pin)
+            out, _, _ = ref.ensure(layer, uniq[start : start + 32], update_state=False, pin_since=ref_pin, lowpri=True)
+            np.testing.assert_array_equal(part.cpu().numpy(), out)
+            _assert_same_tables(cache, ref)
+    assert installs > 0 and int((ref.id_of_slot >= 0).sum()) == size, "the cache must be full and evicting"
+
+
+def _plant(cache, ref, slot, flat_id, usage):
+    for t_id, t_slot in ((cache.id_of_slot, cache.slot_for_id.view(-1)), (ref.id_of_slot, ref.slot_of_id)):
+        t_id[slot] = flat_id
+        t_slot[flat_id] = slot
+    cache.usage[slot] = usage
+    ref.usage[slot] = usage
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"])
+def test_lowpri_band_sits_between_empty_slots_and_every_resident(policy):
+    """An unconsumed prefetch is evicted after every empty slot and before any resident, even the
+    coldest resident the score caps allow, whatever the slot order."""
+    from freetoken.moe import scored_ensure as se
+    from freetoken.moe.offload_kernels import ensure_experts_scored
+
+    num_layers, num_experts, size = 2, 16, 16
+    cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts)
+    dev = cache.id_of_slot.device
+    cache.step.fill_(5)
+    ref.step = 5
+    # slots 0-12: layer-0 residents at the oldest usage 1 with never-seen state (the lowest possible
+    # score); slot 13 empty; slots 14, 15 unconsumed prefetches
+    for slot in range(13):
+        _plant(cache, ref, slot, slot, 1)
+    _plant(cache, ref, 14, 13, 0)
+    _plant(cache, ref, 15, 14, 0)
+    if policy != "lru":
+        cache.evict_slot_owner.fill_(-1)
+    victims = []
+    for e in (15, 16 + 0, 16 + 1, 16 + 2):  # four misses, one per call
+        layer, expert = divmod(e, num_experts)
+        q = torch.tensor([expert], dtype=torch.int32, device=dev)
+        ensure_experts_scored(cache, layer, q, bump_tok=False, update_state=True, lowpri=True)
+        ref.ensure(layer, [expert], lowpri=True)
+        victims.append(int(cache.evict_slots[0]))
+        _assert_same_tables(cache, ref)
+    assert victims[:3] == [13, 14, 15] and victims[3] < 13
+    # the band's keys are below 2 << SLOT_BITS; the caps' worst resident score keeps score + 2^31 >= 2
+    beta_step, w_q4, _, _ = se.score_params(se.MAX_BETA, se.MAX_W, 1, 1)
+    worst = -((se.K_MAX << se.Q) + 1 * beta_step) + ((w_q4 * se.LC_FLOOR) >> 4)
+    assert worst + se.SCORE_BIAS >= 2
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "rule"])
+def test_prefetch_budget_keeps_the_first_non_resident_candidates(policy):
+    """The install budget counts only non-resident candidates, in query order: resident ids, padding
+    and duplicates ahead of them take none of it, and ISSUED counts the installs."""
+    num_experts = 16
+    cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts)
+    stats = torch.zeros(8, dtype=torch.int64, device="cuda")
+    cache.ensure_experts(0, torch.tensor([0, 1], dtype=torch.int32, device="cuda"))  # step 0 pins every slot
+    ref.ensure(0, np.arange(2), bump_tok=True, lowpri=True)
+    src, _, _ = _gpu_prefetch(cache, 1, [4, 6], stats)
+    ref.prefetch(1, [4, 6])
+    assert src.tolist() == [4, 6]
+    sel = [4, -1, 6, 9, 9, 3, 7, 1]
+    src, dst, _ = _gpu_prefetch(cache, 1, sel, stats, budget=2)
+    assert src.tolist() == [9, 3] and int(stats[0]) == 4
+    np.testing.assert_array_equal(np.stack(ref.prefetch(1, sel, 2)), np.stack([src, dst]))
+    _assert_same_tables(cache, ref)
+    src, _, _ = _gpu_prefetch(cache, 1, [5, 8], stats, budget=0)
+    assert src.size == 0 and int(stats[0]) == 4
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["lru", "rule"])
+def test_prefetch_installs_only_into_evictable_slots(policy):
+    """With fewer evictable slots than candidates a prefetch installs the first candidates only, and
+    with every slot pinned by the last demand call it installs nothing (no all-pinned fallback)."""
+    num_experts = 16
+    cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts)
+    stats = torch.zeros(8, dtype=torch.int64, device="cuda")
+    ids = np.arange(12)
+    cache.ensure_experts(0, torch.from_numpy(ids.astype(np.int32)).cuda())
+    ref.ensure(0, ids, bump_tok=True, lowpri=True)
+    pinned = cache.id_of_slot.clone()
+    src, dst, _ = _gpu_prefetch(cache, 1, [9, 3, 7, 1, 5, 11, 2], stats)
+    assert src.tolist() == [9, 3, 7, 1] and sorted(dst.tolist()) == [12, 13, 14, 15]
+    np.testing.assert_array_equal(np.stack(ref.prefetch(1, [9, 3, 7, 1, 5, 11, 2])), np.stack([src, dst]))
+    assert torch.equal(cache.id_of_slot[:12], pinned[:12]) and int(stats[0]) == 4
+    all16 = np.arange(16)
+    cache.ensure_experts(0, torch.from_numpy(all16.astype(np.int32)).cuda())
+    before = [t.clone() for t in (cache.id_of_slot, cache.slot_for_id, cache.usage)]
+    src, dst, num = _gpu_prefetch(cache, 1, [0, 1, 2], stats)
+    assert int(num) == 0 and src.size == 0 and int(stats[0]) == 4
+    for a, b in zip(before, (cache.id_of_slot, cache.slot_for_id, cache.usage)):
+        assert torch.equal(a, b)
+
+
+@_cuda
+@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"])
+def test_prefetch_scores_victims_from_the_previous_layers_position(policy):
+    """prefetch_ensure(L) runs right after layer L-1's demand call and scores victims from there:
+    layer L's residents are next up (d = 1), so one used last token outlives an older resident of
+    a later layer instead of looking a whole pass away."""
+    num_layers, num_experts, size = 8, 16, 16
+    cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts)
+    cache.step.fill_(7)
+    ref.step = 7
+    cache.evict_tok.fill_(10)
+    ref.tok = 10
+    cache.evict_slot_owner.fill_(-1)
+    # slots 0-13: layer 1's routed experts (pinned), slot 14: a layer-2 expert used last token,
+    # slot 15: a layer-5 expert last used three tokens ago; all with the same (never-counted) use count
+    planted = [(slot, 16 + slot, 7, 10) for slot in range(14)] + [(14, 2 * 16 + 3, 6, 9), (15, 5 * 16 + 4, 5, 7)]
+    for slot, flat, usage, last in planted:
+        _plant(cache, ref, slot, flat, usage)
+        cache.evict_last_tok[flat] = last
+        ref.last_tok[flat] = last
+    stats = torch.zeros(8, dtype=torch.int64, device="cuda")
+    src, dst, _ = _gpu_prefetch(cache, 2, [9], stats)
+    assert src.tolist() == [9] and dst.tolist() == [15], "the prefetch evicted the target layer's fresh resident"
+    np.testing.assert_array_equal(np.stack(ref.prefetch(2, [9])), np.stack([src, dst]))
+    _assert_same_tables(cache, ref)
+
+
+@_cuda
+@pytest.mark.parametrize("rows_per_step", [1, 2])
+def test_lowpri_lru_without_prefetches_is_flashlib(rows_per_step):
+    """Prefetch on routes LRU through the vendored kernel with the low-priority key; while no slot is
+    a prefetch it evicts exactly what flashlib's lru_ensure does."""
+    num_layers, num_experts, size = 8, 64, 200
+    on = _pf_cache("lru", size, num_layers, num_experts)
+    off = _cache("lru", size, num_layers, num_experts)
+    rows = _zipf_rows(30, rows_per_step, seed=41, num_layers=num_layers, num_experts=num_experts)
+    ids = torch.from_numpy(rows.astype(np.int32)).cuda()
+    for r in range(rows.shape[0]):
+        for layer in range(num_layers):
+            a, b = ids[r, layer].clone(), ids[r, layer].clone()
+            on.ensure_experts(layer, a)
+            off.ensure_experts(layer, b)
+            assert torch.equal(a, b)
+            assert torch.equal(on.evict_slots[: int(on.num_indices)], off.evict_slots[: int(off.num_indices)])
+    for name in ("slot_for_id", "id_of_slot", "usage", "step"):
+        assert torch.equal(getattr(on, name), getattr(off, name)), name
+
+
+def test_lru_policy_leaves_flashlib_only_when_prefetch_is_on(monkeypatch):
+    import freetoken.moe.offload_kernels as ok
+
+    calls = []
+    monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib"))
+    monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append(("scored", kw["lowpri"])))
+    for mode in ("off", "measure", "on"):
+        cache = _cache("lru", 12, 2, 8, device="cpu", prefetch_mode=mode)
+        cache.prefetch = SimpleNamespace(mode=mode, count_args=lambda layer: None) if mode != "off" else None
+        cache.ensure_experts(1, torch.zeros(2, dtype=torch.int32))
+    assert calls == ["flashlib", "flashlib", ("scored", True)]

From 6ccca4255b15fe6c5d312e7ed85d2001005d9300 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:18:22 -0400
Subject: [PATCH 159/249] test(engine): give the kv quant config fake the rope
 table size #573's max_seq_len reads

Assisted-by: Claude Opus 5.5
---
 tests/engine/test_kv_quant_config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py
index ba4196be09..1f460166d1 100644
--- a/tests/engine/test_kv_quant_config.py
+++ b/tests/engine/test_kv_quant_config.py
@@ -39,7 +39,7 @@ def _model_config(kind):
         has_swa_attention=False,
         has_linear_attention=False,
         num_layers=4,
-        rotary_config=SimpleNamespace(max_position=1024),
+        rotary_config=SimpleNamespace(max_position=1024, table_positions=1024),
     )
     specs = {
         "full": (_spec("full", AttnType.FULL),),

From 9af2aaca6125c18fdb4f2c07f4d39c33cbb1445b Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:55:23 -0400
Subject: [PATCH 160/249] fix(engine): hand the rendezvous listen fd to
 TCPStore instead of keeping its socket

TCPStore owns master_listen_fd and closes it with the store; #461 also kept the Python
socket on the engine, whose finalizer then closed whatever file reused that fd number
after destroy_process_group at shutdown.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit b54cced6af48bffcce88c280124692d54c69f792)
---
 python/freetoken/engine/engine.py      |  8 ++---
 tests/engine/test_distributed_store.py | 49 +++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 9407b76ae8..3ac453581c 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -674,12 +674,12 @@ def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Sto
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         sock.bind(("127.0.0.1", config.distributed_port))
         sock.listen(1024)
-        # TCPStore uses this fd directly rather than duplicating it, so it must outlive the
-        # store; kept alive on self, closed only when the process (and the store) exits.
-        self._distributed_listen_socket = sock
+        self._distributed_listen_addr = sock.getsockname()
+        # TCPStore takes ownership of the fd and closes it with the store; detach so the
+        # Python socket never closes that number again after it has been reused.
         return torch.distributed.TCPStore(
             "127.0.0.1", config.distributed_port, config.tp_info.size,
-            is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.fileno(),
+            is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(),
         )
 
     def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py
index f70be12668..4fc2473762 100644
--- a/tests/engine/test_distributed_store.py
+++ b/tests/engine/test_distributed_store.py
@@ -8,7 +8,9 @@
 
 from __future__ import annotations
 
+import gc
 import multiprocessing as mp
+import os
 from datetime import timedelta
 from types import SimpleNamespace
 
@@ -38,7 +40,7 @@ def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None:
     )
     total = torch.tensor([float(rank)])
     dist.all_reduce(total)
-    bind_host = engine._distributed_listen_socket.getsockname()[0] if rank == 0 else None
+    bind_host = engine._distributed_listen_addr[0] if rank == 0 else None
     result_q.put((rank, total.item(), bind_host))
     dist.destroy_process_group()
 
@@ -64,3 +66,48 @@ def test_two_ranks_rendezvous_with_loopback_only_store():
     # the master's listening socket was pre-bound to loopback, not the wildcard
     # address the C10d TCPStore server binds to on its own
     assert results[0][1] == "127.0.0.1"
+
+
+def _teardown_rank(port: int, result_q: mp.Queue) -> None:
+    listen_fds = []
+    real_store = dist.TCPStore
+
+    def recording_store(*args, **kwargs):
+        listen_fds.append(kwargs.get("master_listen_fd"))
+        return real_store(*args, **kwargs)
+
+    dist.TCPStore = recording_store
+    engine = Engine.__new__(Engine)
+    config = SimpleNamespace(
+        tp_info=DistributedInfo(0, 1), distributed_timeout=TIMEOUT.total_seconds(),
+        distributed_port=port,
+    )
+    store = engine._make_distributed_store(config)
+    dist.init_process_group(backend="gloo", rank=0, world_size=1, timeout=TIMEOUT, store=store)
+    dist.destroy_process_group()
+    del store
+    gc.collect()
+    # the listening fd number is reused by an unrelated file, as a later open() would
+    fd = listen_fds[0]
+    r, w = os.pipe()
+    os.dup2(r, fd)
+    del engine
+    gc.collect()
+    try:
+        os.fstat(fd)
+        result_q.put("open")
+    except OSError as exc:
+        result_q.put(f"closed: {exc}")
+
+
+def test_the_engine_does_not_close_the_listen_fd_the_store_owns():
+    # TCPStore takes ownership of master_listen_fd and closes it itself; an engine that also
+    # kept the Python socket closed whatever file reused that number when it was freed
+    ctx = mp.get_context("spawn")
+    result_q = ctx.Queue()
+    proc = ctx.Process(target=_teardown_rank, args=(29513, result_q))
+    proc.start()
+    outcome = result_q.get(timeout=30)
+    proc.join(timeout=30)
+    assert proc.exitcode == 0
+    assert outcome == "open"

From 49b3f8f85d858271db8e82011576adf41abfd127 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:52:14 -0400
Subject: [PATCH 161/249] fix(scheduler): trim the stop-match EOS only at the
 delivered output limit

#464 used `not req.can_decode`, but since #546 overlap runs device_len a step ahead of
the host, so a non-terminal EOS was dropped from the matched text one step early.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 1814a3acc1848366a59d16e40d25c3d66a3be301)
---
 python/freetoken/scheduler/scheduler.py |  4 +-
 tests/scheduler/test_match_stop_str.py  | 71 +++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 tests/scheduler/test_match_stop_str.py

diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 01f1678761..895f012c75 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -478,9 +478,11 @@ def _match_stop_str(self, req: Req) -> str | None:
         prompt_len = req.max_device_len - req.output_len
         end = len(req.input_ids)
         # The frontend omits a terminal EOS, including at the output limit with ignore_eos.
+        # The limit is the delivered length, as in hit_length: under overlap device_len (and
+        # so can_decode) runs a step ahead of the host.
         if (
             end > prompt_len
-            and not req.can_decode
+            and end >= req.max_device_len
             and int(req.input_ids[-1]) in self.eos_token_ids
         ):
             end -= 1
diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py
new file mode 100644
index 0000000000..351509b759
--- /dev/null
+++ b/tests/scheduler/test_match_stop_str.py
@@ -0,0 +1,71 @@
+"""Scheduler-side stop-string matching (Scheduler._match_stop_str) against the same
+incrementally decoded text the frontend streams."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+from freetoken.core import Req, SamplingParams
+from freetoken.scheduler.scheduler import Scheduler
+
+EOS = 0
+
+
+class _ByteTok:
+    """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS."""
+
+    def decode(self, ids, skip_special_tokens=False):
+        return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode(
+            "utf-8", errors="replace"
+        )
+
+
+def _req(prompt: list[int], output_len: int, **sp) -> Req:
+    return Req(
+        input_ids=torch.tensor(prompt, dtype=torch.int32),
+        table_idx=0,
+        cached_len=0,
+        output_len=output_len,
+        uid=1,
+        sampling_params=SamplingParams(**sp),
+        cache_handle=None,
+    )
+
+
+def _match(req: Req) -> str | None:
+    stub = SimpleNamespace(tokenizer=_ByteTok(), eos_token_ids={EOS})
+    return Scheduler._match_stop_str(stub, req)
+
+
+def _feed(req: Req, ids: list[int]) -> list[str | None]:
+    seen = []
+    for t in ids:
+        req.append_host(torch.tensor([t], dtype=torch.int32))
+        req.complete_one()
+        seen.append(_match(req))
+    return seen
+
+
+def test_a_stop_spanning_more_tokens_than_characters_is_matched():
+    dna = list("\N{DNA DOUBLE HELIX}".encode())  # one character, four byte tokens
+    req = _req([65, 66], 32, stop_strs=["\N{DNA DOUBLE HELIX}"])
+    text = list(b"Here is DNA ") + dna + list(b" and more")
+    seen = _feed(req, text)
+    first = next(i for i, m in enumerate(seen) if m is not None)
+    assert first == len(b"Here is DNA ") + len(dna) - 1
+
+
+def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text():
+    # Overlap scheduling advances device_len ahead of the host: the budget is spent on the
+    # device (can_decode is False) while the host holds one token less than the limit. That
+    # EOS is not terminal, so the frontend decodes it and the scheduler must too.
+    req = _req([65], 3, stop_strs=["zz"], ignore_eos=True)
+    for t in (65, EOS):
+        req.append_host(torch.tensor([t], dtype=torch.int32))
+    while req.can_decode:
+        req.complete_one()
+    assert req.input_ids.numel() < req.max_device_len
+    assert _match(req) is None
+    assert req.stop_decode_status.decoded_ids == [65, EOS]

From 0c62979ee9109ffe88e34629b47102bc688693d8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:52:27 -0400
Subject: [PATCH 162/249] test(tokenizer): cover a detokenize batch that
 carries one uid more than once

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 577427a2cecf047ac64616854a51aa338197336b)
---
 tests/tokenizer/test_detokenize_extras.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py
index 5617a10537..4717617b9a 100644
--- a/tests/tokenizer/test_detokenize_extras.py
+++ b/tests/tokenizer/test_detokenize_extras.py
@@ -79,3 +79,19 @@ def test_reasoning_tokens_count_up_to_the_end_tag():
     assert reasoning == 3  # b, c and the tag itself
     _, none = _run(manager, 2, [1, 2, EOS])
     assert none == 0
+
+
+def test_a_batch_with_one_uid_twice_advances_its_state_in_order():
+    # each message must see the decode state the previous one left, and the reasoning
+    # counts must stay aligned with the texts
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END)
+    msgs = [
+        DetokenizeMsg(uid=1, next_token=1, finished=False),
+        DetokenizeMsg(uid=1, next_token=THINK_END, finished=False),
+        DetokenizeMsg(uid=2, next_token=3, finished=False),
+        DetokenizeMsg(uid=1, next_token=2, finished=True),
+    ]
+    texts, counts = manager.detokenize_with_meta(msgs)
+    assert texts == ["b", "</think>", "d", "c"]
+    assert counts == [0, 0, 0, 2]
+    assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"]

From 435bd69d42926a32f0b05158ca904e6fb3ee7f29 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:53:13 -0400
Subject: [PATCH 163/249] fix(scheduler): match stop strings in the text the
 request's skip_special_tokens decodes

#464's scheduler-side DecodeStatus decoded with special tokens while #393 lets a request
skip them on the frontend, so a stop spanning a special token matched different text.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit fc2392b5367e841b329454bddf33514f808417e1)
---
 python/freetoken/scheduler/scheduler.py  |  7 ++++++-
 python/freetoken/tokenizer/detokenize.py |  7 +++++--
 tests/scheduler/test_match_stop_str.py   | 25 +++++++++++++++++++-----
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 895f012c75..cb6085d344 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -490,7 +490,12 @@ def _match_stop_str(self, req: Req) -> str | None:
             return None
         if req.stop_decode_status is None:
             req.stop_decode_status = DecodeStatus(
-                decoded_ids=[], decoded_str="", read_offset=0, surr_offset=0, sent_offset=0
+                decoded_ids=[],
+                decoded_str="",
+                read_offset=0,
+                surr_offset=0,
+                sent_offset=0,
+                skip_special=req.sampling_params.skip_special_tokens,
             )
         state = req.stop_decode_status
         state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist())
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index b4c0b04733..2eb3029df0 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -84,9 +84,12 @@ def update(self, read_str: str, surr_str: str) -> str:
         return self.decoded_str + find_printable_text(new_text)
 
     def decode(self, tokenizer: PreTrainedTokenizerBase) -> str:
+        skip = self.skip_special
         return self.update(
-            tokenizer.decode(self.decoded_ids[self.surr_offset :]),
-            tokenizer.decode(self.decoded_ids[self.surr_offset : self.read_offset]),
+            tokenizer.decode(self.decoded_ids[self.surr_offset :], skip_special_tokens=skip),
+            tokenizer.decode(
+                self.decoded_ids[self.surr_offset : self.read_offset], skip_special_tokens=skip
+            ),
         )
 
 
diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py
index 351509b759..e9a5c07da5 100644
--- a/tests/scheduler/test_match_stop_str.py
+++ b/tests/scheduler/test_match_stop_str.py
@@ -10,16 +10,22 @@
 from freetoken.core import Req, SamplingParams
 from freetoken.scheduler.scheduler import Scheduler
 
-EOS = 0
+EOS, SPECIAL = 0, 1
 
 
 class _ByteTok:
-    """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS."""
+    """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS,
+    1 a special marker that skip_special_tokens drops."""
 
     def decode(self, ids, skip_special_tokens=False):
-        return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode(
-            "utf-8", errors="replace"
-        )
+        def piece(t):
+            if t == EOS:
+                return b""
+            if t == SPECIAL:
+                return b"" if skip_special_tokens else b"<s>"
+            return bytes([t])
+
+        return b"".join(piece(t) for t in ids).decode("utf-8", errors="replace")
 
 
 def _req(prompt: list[int], output_len: int, **sp) -> Req:
@@ -69,3 +75,12 @@ def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text():
     assert req.input_ids.numel() < req.max_device_len
     assert _match(req) is None
     assert req.stop_decode_status.decoded_ids == [65, EOS]
+
+
+def test_stops_match_the_text_decoded_with_the_requests_skip_special_tokens():
+    # the frontend streams "ab" for a request that skips special tokens; the scheduler has
+    # to match the stop in that same text, not in "a<s>b"
+    req = _req([65], 8, stop_strs=["ab"], skip_special_tokens=True)
+    assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, "ab"]
+    req = _req([65], 8, stop_strs=["ab"])
+    assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, None]

From 07b8be03fcb9fbd0931f8c4c8b974bddb838f8cd Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:43:14 -0400
Subject: [PATCH 164/249] perf(ple): skip the device n-gram hash for the disk
 table, which reads only its shape

DiskRowTable hashes on the host in its fill; its lookup and prefetch read only the
row_ids shape, so every prefill chunk and captured decode step ran the hash for nothing.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 52a18015ffff759a9f2b7733954bd9998b900f32)
---
 python/freetoken/models/qwen4_exp/ple.py      | 17 ++++++--
 python/freetoken/models/qwen4_exp/ple_disk.py |  3 ++
 tests/models/qwen4_exp/test_ple.py            | 40 +++++++++++++++++++
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py
index 44973bc0a6..22bda9e6c8 100644
--- a/python/freetoken/models/qwen4_exp/ple.py
+++ b/python/freetoken/models/qwen4_exp/ple.py
@@ -581,8 +581,19 @@ def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor:
             blocks.append(head_ids + self.ngram_heads_offsets[start:end])
         return torch.cat(blocks, dim=-1)
 
+    def lookup_row_ids(self, meta: PLEMetadata) -> torch.Tensor:
+        """``row_ids`` as the attached table's ``prefetch``/``lookup`` take them. A table that
+        hashes on the host (``reads_row_ids = False``) reads only their shape: skip the hash."""
+        if not getattr(self.table, "reads_row_ids", True):
+            return torch.empty(
+                (meta.input_ids.numel(), self.num_heads),
+                dtype=torch.int64,
+                device=meta.input_ids.device,
+            )
+        return self.row_ids(meta)
+
     def forward(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor:
-        return self.table.lookup(self.row_ids(meta), out)
+        return self.table.lookup(self.lookup_row_ids(meta), out)
 
 
 class _DepthwiseConv1d(BaseOP):
@@ -680,7 +691,7 @@ def start_prefetch(self, batch: Batch, meta: PLEMetadata | None = None) -> None:
         """Hash this forward's n-grams and start the table gather on the side stream."""
         if meta is None:
             meta = build_ple_metadata(batch, self.args, batch.input_ids.device)
-        row_ids = self.ple_embedding.row_ids(meta)
+        row_ids = self.ple_embedding.lookup_row_ids(meta)
         self._pending = (meta, row_ids)
         self.ple_embedding.table.prefetch(row_ids)
 
@@ -701,7 +712,7 @@ def forward(
         elif pending is not None and pending[0] is meta:
             row_ids = pending[1]
         if row_ids is None:
-            row_ids = self.ple_embedding.row_ids(meta)
+            row_ids = self.ple_embedding.lookup_row_ids(meta)
 
         embeddings = self.ple_embedding.table.lookup(row_ids).to(R.dtype)
         key = self.norm_key.forward(self.key_proj.forward(embeddings))
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index a8e6d550c7..e317d0af89 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -181,6 +181,9 @@ def cancel(self) -> None:
 class DiskRowTable:
     """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk)."""
 
+    # the fill hashes on the host; lookup/prefetch read only the row_ids shape
+    reads_row_ids = False
+
     def __init__(
         self,
         source: PleRowSource,
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index 02abaee1f4..954d898945 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -339,6 +339,46 @@ def test_the_fused_hash_refuses_a_geometry_it_cannot_address(ctx_len, heads_per_
         )
 
 
+
+def test_a_host_hashing_table_skips_the_device_hash(monkeypatch):
+    """The disk table hashes on the host and reads only the row_ids shape, so neither the
+    prefetch nor the lookup path runs the device hash for it."""
+    from freetoken.models.qwen4_exp.ple import ZeroTable
+    from freetoken.models.qwen4_exp.ple_disk import DiskRowTable
+
+    assert DiskRowTable.reads_row_ids is False
+    config = _config()
+    args = config.qwen4_args
+    seen = []
+
+    class HostHashedTable(ZeroTable):
+        reads_row_ids = False
+
+        def lookup(self, row_ids, out=None):
+            seen.append(("lookup", tuple(row_ids.shape), row_ids.dtype))
+            return super().lookup(row_ids, out)
+
+        def prefetch(self, row_ids):
+            seen.append(("prefetch", tuple(row_ids.shape), row_ids.dtype))
+
+    layer = _make_layer(config, table=HostHashedTable(_padded_vocab(args), args.ngram_head_dim))
+    embedding = layer.ple_embedding
+
+    def no_hash(*a, **k):
+        raise AssertionError("the device hash ran for a table that never reads it")
+
+    monkeypatch.setattr(embedding, "row_ids", no_hash)
+    sequences, contexts = [[3, 4, EOS, 5], [2, 7]], [[EOS, EOS], [21, 22]]
+    meta = _meta(sequences, contexts)
+    total = sum(len(s) for s in sequences)
+    states = torch.zeros(len(sequences), args.ple_state_width, args.ple_conv_state_len)
+
+    layer.start_prefetch(None, meta)
+    _forward(layer, torch.randn(total, layer.hc_count * layer.hidden_size), meta, states)
+    embedding.forward(meta)
+    rows = (total, embedding.num_heads)
+    assert seen == [("prefetch", rows, torch.int64), ("lookup", rows, torch.int64), ("lookup", rows, torch.int64)]
+
 # --------------------------------------------------------------------------------------
 # table backends
 # --------------------------------------------------------------------------------------

From 3c7be6e2981a4ca72c401a2d6b273faf59b6e9fd Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:24:20 -0400
Subject: [PATCH 165/249] fix(ple): give each graph fill its own readback event
 and keep every pending fill

A shared event re-recorded by the next step before the filler reached
synchronize() made fill N wait on graph N, which waits on fill N; and a fill
that failed after the next step submitted was dropped, so its graph ran on
stale rows without an error.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 79988ccdf596e54b37817971628de84c802f3a58)
---
 python/freetoken/models/qwen4_exp/ple_disk.py | 34 ++++---
 .../qwen4_exp/test_ple_disk_fill_order.py     | 95 ++++++++++++++++++-
 2 files changed, 117 insertions(+), 12 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index e317d0af89..9edd715e4e 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -8,6 +8,7 @@
 import os
 import queue
 import threading
+from collections import deque
 from concurrent.futures import Future
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -175,6 +176,9 @@ def cancel(self) -> None:
             self._future.result()
         except BaseException:  # noqa: BLE001 -- the launch error is the one to surface
             pass
+        pending = self._table._pending_fills()
+        if self._future in pending:
+            pending.remove(self._future)  # its outcome was consumed here
         self._table._flag.zero_()
 
 
@@ -250,11 +254,10 @@ def __init__(
         from freetoken.kernel.row_store import probe_wait_sync
 
         self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device)
-        # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap
+        # one flag for all graphs: each fill's readback event orders it after the previous graph, so signals never overlap
         self._flag = alloc_pinned_tensor(1, dtype=torch.int64)
         self._flag.zero_()
         self._token_readback = alloc_pinned_tensor(max_graph_rows, dtype=torch.int32)
-        self._readback_event = torch.cuda.Event()
         sync = "wait-sync" if self._wait_sync else "launch-gating"
         logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}")
 
@@ -274,12 +277,17 @@ def close(self) -> None:
     def _current_stream(self):
         return torch.cuda.current_stream(self._device)
 
+    def _pending_fills(self) -> deque:
+        return self.__dict__.setdefault("_pending_fill_queue", deque())
+
     def _raise_failed_fill(self) -> None:
-        pending = self.__dict__.pop("_pending_fill", None)
-        if pending is not None and pending.done() and pending.exception() is not None:
-            raise pending.exception()
-        if pending is not None and not pending.done():
-            self._pending_fill = pending
+        """Surface a finished graph fill's error on the engine thread, oldest first; a failed
+        fill signalled its WAIT, so that graph ran on stale rows."""
+        pending = self._pending_fills()
+        while pending and pending[0].done():
+            exc = pending.popleft().exception()
+            if exc is not None:
+                raise exc
 
     def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
         """A fill that completes before its launch, on the store's thread (after any queued graph fill)."""
@@ -322,12 +330,15 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None:
             if use_graph and self._wait_sync:
                 bs = batch.padded_size
                 self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
-                self._readback_event.record(self._current_stream())
+                # one event per fill: a shared one re-recorded by the next step before this job
+                # reached synchronize() would wait on a graph parked on this job's own WAIT
+                readback = torch.cuda.Event()
+                readback.record(self._current_stream())
                 cancelled = threading.Event()
 
                 def _job() -> None:
                     try:
-                        self._readback_event.synchronize()
+                        readback.synchronize()
                         tokens = self._token_readback[:bs].to(torch.int64).tolist()
                         if cancelled.is_set():
                             return
@@ -340,8 +351,9 @@ def _job() -> None:
                         signal(self._flag)
                         raise
 
-                self._pending_fill = self._filler().submit(_job)
-                return _PendingFill(self, self._pending_fill, cancelled)
+                future = self._filler().submit(_job)
+                self._pending_fills().append(future)
+                return _PendingFill(self, future, cancelled)
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
             tokens = batch.input_ids.to("cpu").to(torch.int64).tolist()
             runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)]
diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
index e7fe58f7e0..4d2bbfd71b 100644
--- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py
+++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import queue
 import threading
 import time
 from types import SimpleNamespace
@@ -55,13 +56,13 @@ def synchronize(self):
 def _table(monkeypatch) -> DiskRowTable:
     # host-only: no CUDA stream behind the (faked) readback event
     monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None)
+    monkeypatch.setattr(ple_disk.torch.cuda, "Event", _Event)
     t = DiskRowTable.__new__(DiskRowTable)
     t._device = None
     t._wait_sync = True
     t._flag = torch.zeros(1, dtype=torch.int64)
     t._store = _Store(t._flag)
     t._token_readback = torch.zeros(8, dtype=torch.int32)
-    t._readback_event = _Event()
     t._graph_pinned = torch.zeros(64, dtype=torch.uint8)
     t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)]
     t._eager_read = [_Event(), _Event()]
@@ -134,3 +135,95 @@ def test_a_failed_launch_leaves_no_stale_signal(monkeypatch):
     # no WAIT consumed it, so a leftover 1 would let the next step read stale rows
     assert int(t._flag[0]) == 0
     _close(t)
+
+
+class _Gpu:
+    """One in-order stream: each launched graph parks on WAIT(flag), then RESETs it."""
+
+    def __init__(self, flag: torch.Tensor):
+        self.flag = flag
+        self.enqueued = self.completed = 0
+        self.stalled = False
+        self.cv = threading.Condition()
+        self._launches: queue.SimpleQueue = queue.SimpleQueue()
+        threading.Thread(target=self._run, daemon=True).start()
+
+    def launch(self) -> None:
+        self.enqueued += 1
+        self._launches.put(True)
+
+    def close(self) -> None:
+        self._launches.put(False)
+
+    def _run(self) -> None:
+        while self._launches.get():
+            self.stalled |= not _blocking_launch(self.flag)
+            with self.cv:
+                self.completed += 1
+                self.cv.notify_all()
+
+    def wait_completed(self, n: int, timeout: float = 5.0) -> bool:
+        with self.cv:
+            return self.cv.wait_for(lambda: self.completed >= n, timeout)
+
+
+class _StreamEvent:
+    """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record."""
+
+    def __init__(self, gpu: _Gpu):
+        self.gpu, self.target = gpu, 0
+
+    def record(self, stream=None):
+        self.target = self.gpu.enqueued
+
+    def synchronize(self):
+        if not self.gpu.wait_completed(self.target, timeout=1.0):
+            raise TimeoutError("the readback waits on a graph parked on this fill's own WAIT")
+
+
+def _two_steps_behind_a_slow_filler(t: DiskRowTable, gpu: _Gpu) -> None:
+    """Steps N and N+1 both launch before the filler thread reaches fill N."""
+    gate = threading.Event()
+    t._filler().submit(gate.wait)
+    for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)):
+        with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True):
+            gpu.launch()
+    gate.set()
+    assert gpu.wait_completed(2), "decode deadlocked"
+    t._filler().submit(lambda: None).result()  # every submitted fill has finished
+    gpu.close()
+
+
+def test_each_fill_waits_on_its_own_steps_readback(monkeypatch):
+    """A readback event shared across steps is re-recorded by step N+1 behind graph N, so
+    fill N would wait on the graph that is parked on fill N's own WAIT."""
+    t = _table(monkeypatch)
+    gpu = _Gpu(t._flag)
+    monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu))
+    _two_steps_behind_a_slow_filler(t, gpu)
+    t._raise_failed_fill()
+    _close(t)
+    assert not gpu.stalled
+    # the fake D2H readback is not stream-ordered, so compare only the contexts
+    assert [c[1][:2] for c in t._store.calls if c[0] == "stage"] == [[2, 3], [3, 7]]
+
+
+def test_a_failed_fill_is_raised_after_the_next_step_was_submitted(monkeypatch):
+    """Fill N fails (its graph runs on stale rows) after step N+1 already submitted: the error
+    must still reach the engine thread at the next launch."""
+    t = _table(monkeypatch)
+    gpu = _Gpu(t._flag)
+    monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu))
+    flush = t._store.flush
+
+    def fail_step_n(signal_addr):
+        if t._store.calls[-1][1][:2] == [2, 3]:  # the run staged just before is step N's
+            raise OSError("PLE row read failed")
+        flush(signal_addr)
+
+    t._store.flush = fail_step_n
+    _two_steps_behind_a_slow_filler(t, gpu)
+    req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4)
+    with pytest.raises(OSError, match="PLE row read failed"):
+        t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False)
+    _close(t)

From f15f924380517d4716cdf9e23450ec48db4729fe Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 15:01:03 -0400
Subject: [PATCH 166/249] fix(ple): recycle the fill readback events instead of
 freeing them on the filler

The filler thread dropped each fill's last event reference, and cuEventDestroy there blocks
on the driver lock that a blocked cuGraphLaunch holds (host-func nodes in the decode graph make
the launch wait for the GPU, parked on the next fill's WAIT), so decode deadlocked with
FREETOKEN_CPU_MOE_FLAG_SYNC=0. Events now return to an engine-thread pool once their fill is done.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 198574d266608fd6744f0ad90f01461ecc6b1c2d)
---
 python/freetoken/models/qwen4_exp/ple_disk.py | 21 ++++++++++++++--
 .../qwen4_exp/test_ple_disk_fill_order.py     | 25 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index 9edd715e4e..fa8740897d 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -179,6 +179,7 @@ def cancel(self) -> None:
         pending = self._table._pending_fills()
         if self._future in pending:
             pending.remove(self._future)  # its outcome was consumed here
+        self._table._retire(self._future)
         self._table._flag.zero_()
 
 
@@ -280,12 +281,24 @@ def _current_stream(self):
     def _pending_fills(self) -> deque:
         return self.__dict__.setdefault("_pending_fill_queue", deque())
 
+    def _readback_pool(self) -> list:
+        return self.__dict__.setdefault("_readback_events", [])
+
+    def _retire(self, future: Future) -> None:
+        """Return a finished fill's readback event for reuse (engine thread only)."""
+        event = getattr(future, "readback", None)
+        if event is not None:
+            self._readback_pool().append(event)
+            future.readback = None
+
     def _raise_failed_fill(self) -> None:
         """Surface a finished graph fill's error on the engine thread, oldest first; a failed
         fill signalled its WAIT, so that graph ran on stale rows."""
         pending = self._pending_fills()
         while pending and pending[0].done():
-            exc = pending.popleft().exception()
+            future = pending.popleft()
+            self._retire(future)
+            exc = future.exception()
             if exc is not None:
                 raise exc
 
@@ -332,7 +345,10 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None:
                 self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
                 # one event per fill: a shared one re-recorded by the next step before this job
                 # reached synchronize() would wait on a graph parked on this job's own WAIT
-                readback = torch.cuda.Event()
+                # recycled, never destroyed while serving: cuEventDestroy from the filler blocks on
+                # the driver lock a blocked cuGraphLaunch holds, and the fill it waits for never runs
+                pool = self._readback_pool()
+                readback = pool.pop() if pool else torch.cuda.Event()
                 readback.record(self._current_stream())
                 cancelled = threading.Event()
 
@@ -352,6 +368,7 @@ def _job() -> None:
                         raise
 
                 future = self._filler().submit(_job)
+                future.readback = readback
                 self._pending_fills().append(future)
                 return _PendingFill(self, future, cancelled)
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
index 4d2bbfd71b..9d979b64eb 100644
--- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py
+++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py
@@ -227,3 +227,28 @@ def fail_step_n(signal_addr):
     with pytest.raises(OSError, match="PLE row read failed"):
         t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False)
     _close(t)
+
+
+def test_the_readback_event_is_never_freed_on_the_filler_thread(monkeypatch):
+    """cuEventDestroy on the filler blocks on the driver lock a blocked cuGraphLaunch holds (host
+    nodes in the graph make the launch wait for the GPU, which is parked on the next fill's WAIT)."""
+    import gc
+
+    freed_on: list[str] = []
+
+    class _TrackedEvent(_Event):
+        def __del__(self):
+            freed_on.append(threading.current_thread().name)
+
+    t = _table(monkeypatch)
+    monkeypatch.setattr(ple_disk.torch.cuda, "Event", _TrackedEvent)
+    for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)):
+        with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True):
+            assert _blocking_launch(t._flag)
+    time.sleep(0.05)  # let the filler drop its work items
+    gc.collect()
+    t._raise_failed_fill()
+    gc.collect()
+    assert "ple-filler" not in freed_on, freed_on
+    assert len(t._readback_pool()) >= 1  # finished fills hand their events back for reuse
+    _close(t)

From 45acef309b718474416c9afdfdfe91adb50fcb5d Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:25:06 -0400
Subject: [PATCH 167/249] fix(kvcache): refuse the host kv-page tier on a pool
 whose index slab it cannot restore

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit eb462bd73c2fa5eb96b38abe56fb841d82d68a32)
---
 python/freetoken/kvcache/host_prefix_tier.py | 8 +++++++-
 tests/kvcache/test_host_tier_nvfp4_guard.py  | 7 +++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 39c96dfd6d..5cee24058a 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -84,8 +84,14 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
             from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool
 
             _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST")
+            if getattr(kv_pool, "_cmp_k_buffer", None) is not None:
+                # a QSA pool's indexer slab is keyed by the same pages and is not mirrored here
+                raise ValueError(
+                    "FT_PREFIX_HOST cannot restore a sparse-attention pool's compressed index "
+                    "slab; keep FT_GDN_HOST_TIER alone for this model"
+                )
             codes = kv_pool._kv_buffer            # [2, L, P, page, H, D]
-            scales = kv_pool._scale_buffer        # [2, L, P*page, H] | None (fp8)
+            scales = getattr(kv_pool, "_scale_buffer", None)  # [2, L, P*page, H] | None (fp8)
             _, layers, _, page_size, heads, dim = codes.shape
             page_bytes = 2 * layers * page_size * heads * dim * codes.element_size()
             if scales is not None:
diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py
index a793fb294a..50227267a6 100644
--- a/tests/kvcache/test_host_tier_nvfp4_guard.py
+++ b/tests/kvcache/test_host_tier_nvfp4_guard.py
@@ -42,3 +42,10 @@ def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool():
     tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2,
                           kv_budget_bytes=0, device="cpu")
     assert tier.num_slots == 2 and not tier.has_kv
+
+
+def test_prefix_tier_refuses_kv_pages_on_a_qsa_pool():
+    # the KV-page store mirrors codes and scales, not the QSA indexer's compressed-key slab
+    pool = SimpleNamespace(**vars(_kv_pool(block_scaled=False)), _cmp_k_buffer=torch.zeros(1))
+    with pytest.raises(ValueError, match="FT_PREFIX_HOST .* index slab"):
+        HostPrefixTier(_state_pool(), kv_pool=pool, gdn_slots_host=2, kv_budget_bytes=1 << 16, device="cpu")

From e3f77624f90e9b98ba3126d5451cc8addf1b3d12 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:51:57 -0400
Subject: [PATCH 168/249] perf(kvcache): keep one gpu gdn snapshot slot when
 the host snapshot tier is on

Evicted snapshots spill to RAM and return on a hit, so the floor of four resident ones
only costs VRAM: three 110 MiB Qwen3.8-Flash-Next slots go back to the expert cache.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 6751cf86b302dd539fdf81c1b32aaa8820650a5b)
---
 python/freetoken/kvcache/linear_state_pool.py |  6 +++++-
 tests/scheduler/test_hybrid_cache_manager.py  | 11 +++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py
index ff7539410c..d89556899d 100644
--- a/python/freetoken/kvcache/linear_state_pool.py
+++ b/python/freetoken/kvcache/linear_state_pool.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+import os
 
 import torch
 from freetoken.distributed import get_tp_info
@@ -285,7 +286,10 @@ def _linear_pool_num_slots(config) -> int:
     if config.cache_type != "hybrid_radix":
         return mr + 1  # live + dummy/padding
     ratio = config.linear_state_cache_ratio
-    n_cache = max(4, int(ratio * mr))
+    # With the host snapshot tier (FT_GDN_HOST_TIER=1) evicted snapshots spill to RAM and come
+    # back on a hit, so the GPU cache need not keep four of them resident.
+    floor = 1 if os.environ.get("FT_GDN_HOST_TIER", "0") == "1" else 4
+    n_cache = max(floor, int(ratio * mr))
     return 4 * mr + n_cache + 1  # live + 2 ping-pong + locked committed snapshot + cache + padding
 
 
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index 07385d7046..165689ab3b 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -241,5 +241,16 @@ def test_pool_sizing_covers_4mr_floor():
         assert _linear_pool_num_slots(c) >= 4 * mr + 1, (mr, _linear_pool_num_slots(c))
 
 
+def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch):
+    """With FT_GDN_HOST_TIER evicted snapshots live in RAM, so one GPU snapshot slot is enough."""
+    from types import SimpleNamespace
+    from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots
+    c = SimpleNamespace(max_running_req=1, cache_type="hybrid_radix", linear_state_cache_ratio=0.0)
+    monkeypatch.delenv("FT_GDN_HOST_TIER", raising=False)
+    assert _linear_pool_num_slots(c) == 4 + 4 + 1
+    monkeypatch.setenv("FT_GDN_HOST_TIER", "1")
+    assert _linear_pool_num_slots(c) == 4 + 1 + 1
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))

From 4002afcac2396672e0c160467e97c2297683b992 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:43:02 -0400
Subject: [PATCH 169/249] perf(scheduler): probe the host tier for the exact
 tombstone prefix

Revival only accepts an exact entry, but the longest-prefix lookup re-hashed every
shorter page-aligned prefix on a miss and refreshed the LRU age of a shorter entry it
then rejected. _rehydrate_snapshot on a 262144-token key, CPU, median of 3:
  no host entry          1852 ms -> 4.0 ms
  entry 4096 tokens in   1494 ms -> 2.9 ms
  entry 8192 below        91 ms -> 3.5 ms
  exact hit               1.4 ms -> 1.6 ms

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit b349af0cac8ee9334b49023460bc5d7afa9e92d1)
---
 python/freetoken/kvcache/host_prefix_tier.py  | 10 +++++
 python/freetoken/scheduler/cache.py           |  8 ++--
 .../kvcache/test_host_tier_rehydrate_exact.py | 39 ++++++++++++-------
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 5cee24058a..4295aa8e4b 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -318,6 +318,16 @@ def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | No
             n -= page_size
         return None
 
+    def lookup_exact(self, key_tokens: torch.Tensor) -> HostMatch | None:
+        """The entry for exactly this prefix: one hash, where lookup re-hashes every shorter
+        page-aligned prefix on a miss."""
+        key = self._key_of(key_tokens)
+        hit = self._index.get(key)
+        if hit is None:
+            return None
+        self._index.move_to_end(key)
+        return HostMatch(key_tokens.numel(), hit.kv_slots, hit.gdn_slot)
+
     def evict_host(self, key_tokens: torch.Tensor) -> None:
         """Remove uma entrada do índice e libera os slots host (KV + GDN)."""
         ent = self._index.pop(self._key_of(key_tokens), None)
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 4e6b91df88..715dfdc21a 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -625,10 +625,10 @@ def _rehydrate_snapshot(self, node) -> int | None:
         keep walking (the match truncates as before)."""
         tier = self.host_tier
         key = self.prefix_cache._collect_key(node)
-        hit = tier.lookup(key, self.page_size)
-        # lookup is longest-prefix; a shorter entry's state would resume this node's full
-        # depth with the tokens in between missing from the recurrence
-        if hit is None or hit.gdn_slot is None or hit.host_len != len(key):
+        # exact only: a shorter entry's state would resume this node's full depth with the
+        # tokens in between missing from the recurrence
+        hit = tier.lookup_exact(key)
+        if hit is None or hit.gdn_slot is None:
             return None
         pool = self.linear_state_pool
         if pool.num_free_slots < 1:
diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py
index f67402041f..d6f4d09215 100644
--- a/tests/kvcache/test_host_tier_rehydrate_exact.py
+++ b/tests/kvcache/test_host_tier_rehydrate_exact.py
@@ -1,21 +1,19 @@
 """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of
-exactly that node's prefix: the host index lookup is longest-prefix, and a shorter
-snapshot would resume the recurrence with the tokens in between silently missing."""
+exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens
+in between silently missing."""
 
 from types import SimpleNamespace
 
 import torch
 
-from freetoken.kvcache.host_prefix_tier import HostMatch
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
 from freetoken.scheduler.cache import CacheManager
 
 
-def _manager(host_len: int, key_len: int = 128):
-    reads = []
-    tier = SimpleNamespace(
-        lookup=lambda key, page_size: HostMatch(host_len, None, 7),
-        read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)),
-    )
+def _manager(key_len: int = 128):
+    state = SimpleNamespace(conv_states=torch.zeros(1, 4, 1, 1),
+                            recurrent_states=torch.zeros(1, 4, 1, 1), slot_states={})
+    tier = HostPrefixTier(state, gdn_slots_host=4, device="cpu")
     pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3])
     manager = SimpleNamespace(
         host_tier=tier,
@@ -24,16 +22,27 @@ def _manager(host_len: int, key_len: int = 128):
         linear_state_pool=pool,
         ensure_mamba_slots=lambda n: None,
     )
-    return manager, reads
+    return manager, tier, state
+
+
+def _snapshot(tier, state, key, value):
+    state.conv_states[:, 1] = value
+    tier.put_prefix(key, None, tier.write_gdn(1))
 
 
-def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node():
-    manager, reads = _manager(host_len=64)
+def test_a_shorter_host_snapshot_neither_revives_nor_refreshes_a_deeper_node():
+    manager, tier, state = _manager(key_len=128)
+    _snapshot(tier, state, torch.arange(64), 1.0)
+    _snapshot(tier, state, torch.arange(1000, 1064), 2.0)
     assert CacheManager._rehydrate_snapshot(manager, node=object()) is None
-    assert reads == []
+    # the rejected entry keeps its LRU age
+    assert [len(e.key_tokens) for e in tier._index.values()] == [64, 64]
+    assert next(iter(tier._index.values())).key_tokens[0] == 0
+    assert torch.all(state.conv_states[:, 3] == 0)
 
 
 def test_an_exact_host_snapshot_revives_the_node():
-    manager, reads = _manager(host_len=128)
+    manager, tier, state = _manager(key_len=128)
+    _snapshot(tier, state, torch.arange(128), 5.0)
     assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3
-    assert reads == [(7, 3)]
+    assert torch.all(state.conv_states[:, 3] == 5.0)

From da62c755d58888e1fe1f13bb814446109109c41c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:46:01 -0400
Subject: [PATCH 170/249] fix(scheduler): re-probe the host entry after making
 room to rehydrate it

The evictions that make room for a revived snapshot or a host prefix write their
victims to the host tier, which can drop the hit's own entry and hand its GDN and KV
slots to the victim, so the stale hit restored another prefix's state.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 44f25067541a6fb84161effe4dd2d7ee8e3a0306)
---
 python/freetoken/scheduler/cache.py           | 12 ++-
 .../kvcache/test_host_tier_rehydrate_exact.py | 79 ++++++++++++++++++-
 2 files changed, 87 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 715dfdc21a..ff4467328e 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -632,9 +632,11 @@ def _rehydrate_snapshot(self, node) -> int | None:
             return None
         pool = self.linear_state_pool
         if pool.num_free_slots < 1:
-            # make room: evicting another snapshot writes IT to the host tier first
+            # make room: evicting another snapshot writes IT to the host tier first, which may
+            # drop this entry and hand its host slot to the victim, so probe again
             self.ensure_mamba_slots(1)
-            if pool.num_free_slots < 1:
+            hit = tier.lookup_exact(key)
+            if pool.num_free_slots < 1 or hit is None or hit.gdn_slot is None:
                 return None
         slot = pool.alloc(1)[0]
         tier.read_gdn(hit.gdn_slot, slot)
@@ -651,7 +653,11 @@ def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None:
         pool = self.linear_state_pool
         if hit.gdn_slot is not None:
             self.ensure_mamba_slots(1)
-            if pool.num_free_slots < 1:
+            # _allocate and ensure_mamba_slots write their victims to the tier, which may drop
+            # this entry and hand its host slots to them: use what the key maps to now
+            hit = tier.lookup_exact(ids[: hit.host_len])
+            if (pool.num_free_slots < 1 or hit is None or hit.kv_pages is None
+                    or hit.gdn_slot is None):
                 self._free(kv_new)  # sem slot GDN não dá pra reidratar; devolve as páginas
                 return
             slot = pool.alloc(1)[0]
diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py
index d6f4d09215..77dae0fd08 100644
--- a/tests/kvcache/test_host_tier_rehydrate_exact.py
+++ b/tests/kvcache/test_host_tier_rehydrate_exact.py
@@ -1,12 +1,16 @@
 """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of
 exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens
-in between silently missing."""
+in between silently missing. Making room for the revival evicts other prefixes into the
+tier, which may recycle the entry's own host slots, so it must not read stale ones."""
 
 from types import SimpleNamespace
 
 import torch
 
+import freetoken.core as core
 from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+from freetoken.kvcache.linear_state_pool import LinearStatePool
+from freetoken.models.config import LinearGatedDeltaGroupConfig
 from freetoken.scheduler.cache import CacheManager
 
 
@@ -46,3 +50,76 @@ def test_an_exact_host_snapshot_revives_the_node():
     _snapshot(tier, state, torch.arange(128), 5.0)
     assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3
     assert torch.all(state.conv_states[:, 3] == 5.0)
+
+
+def _cache_manager(monkeypatch, num_pages, page_size, state_slots, **env):
+    for k, v in env.items():
+        monkeypatch.setenv(k, v)
+    g = LinearGatedDeltaGroupConfig(
+        name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4,
+        key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu",
+    )
+    pool = LinearStatePool(group=g, num_slots=state_slots, dtype=torch.bfloat16,
+                           device=torch.device("cpu"), tp_size=1)
+    page_table = torch.zeros(2, 64, dtype=torch.int32)
+    cm = CacheManager(num_pages, page_size, page_table, "hybrid_radix", linear_state_pool=pool)
+    return cm, pool
+
+
+def _put(cm, pool, ids, state, stamp, kv_pool=None, kv_value=None):
+    kv = cm._page_to_token(cm._allocate(len(ids) // cm.page_size))
+    if kv_pool is not None:
+        kv_pool._kv_buffer[:, :, kv[:: cm.page_size] // cm.page_size] = kv_value
+    slot = pool.alloc(1)[0]
+    pool.conv_states[:, slot] = state
+    cm.prefix_cache.insert(ids, kv, slot)
+    cm.prefix_cache.match_prefix(ids).node.timestamp = stamp
+
+
+def _match(cm, ids):
+    tail = torch.tensor([7, 9], dtype=torch.int32)
+    return cm.match_req(SimpleNamespace(input_ids=torch.cat([ids, tail]), input_len=len(ids) + 2))
+
+
+def test_making_room_to_revive_does_not_restore_the_victim_into_the_host_slot(monkeypatch):
+    # one host slot: evicting B for the revived slot drops A's entry and writes B there
+    cm, pool = _cache_manager(monkeypatch, 8, 64, 3, FT_GDN_HOST_TIER="1", FT_GDN_HOST_SLOTS="1")
+    a = torch.arange(128, dtype=torch.int32)
+    b = torch.arange(1000, 1128, dtype=torch.int32)
+    _put(cm, pool, a, 1.0, stamp=1)
+    _put(cm, pool, b, 2.0, stamp=2)
+    cm.ensure_mamba_slots(1)  # A -> host tier, A's node becomes a tombstone
+    pool.conv_states[:, pool.alloc(1)[0]] = 9.0  # someone else takes the freed slot
+
+    mr = _match(cm, a)
+    restored = None if mr.mamba_value is None else pool.conv_states[:, mr.mamba_value].unique().tolist()
+    assert restored in (None, [1.0]), (mr.cuda_handle.cached_len, restored)
+    assert cm.host_tier.lookup_exact(b) is not None
+
+
+def test_making_room_to_rehydrate_a_host_prefix_does_not_restore_the_victim(monkeypatch):
+    # a 4-page KV store: writing the 4-page victim V drops the 2-page hit H and reuses its slots
+    page, n_pages = 4, 6
+    kv_pool = SimpleNamespace(_kv_buffer=torch.zeros(2, 1, n_pages, page, 1, 2),
+                              _scale_buffer=None, _block_scale_buffer=None)
+    monkeypatch.setattr(core, "_GLOBAL_CTX", SimpleNamespace(kv_cache=kv_pool, kv_offloader=None))
+    page_bytes = kv_pool._kv_buffer[:, :, 0].numel() * kv_pool._kv_buffer.element_size()
+    cm, pool = _cache_manager(monkeypatch, n_pages, page, 4, FT_PREFIX_HOST="1",
+                              FT_PREFIX_HOST_GB=repr(4 * page_bytes / (1 << 30)),
+                              FT_GDN_HOST_SLOTS="4")
+    h = torch.arange(2 * page, dtype=torch.int32)
+    v = torch.arange(100, 100 + 4 * page, dtype=torch.int32)
+    _put(cm, pool, h, 1.0, stamp=1, kv_pool=kv_pool, kv_value=1.0)
+    er = cm.prefix_cache.evict_full(len(h))  # H -> host tier (KV + snapshot), out of the tree
+    pool.free(er.mamba_slots)
+    cm._free(er.kv_indices)
+    _put(cm, pool, v, 2.0, stamp=2, kv_pool=kv_pool, kv_value=2.0)
+    cm._allocate(len(cm.free_slots))  # a running request holds the rest of the GPU pages
+
+    mr = _match(cm, h)
+    handle = mr.cuda_handle
+    if handle.cached_len:
+        pages = handle.kv_indices[::page] // page
+        assert pool.conv_states[:, mr.mamba_value].unique().tolist() == [1.0]
+        assert kv_pool._kv_buffer[:, :, pages].unique().tolist() == [1.0]
+    assert cm.host_tier.lookup_exact(v) is not None

From d4de4d0fdaad287e16c758ee6cc2fa0f174d0a01 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:49:19 -0400
Subject: [PATCH 171/249] test(kvcache): cover the engine-stream waits of the
 prefix tier's kv page copies

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit af5ca25838e9c690ba32a6cae20c72692e61477c)
---
 tests/kvcache/test_host_tier_live_pool.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py
index bd12836893..0eea1990e3 100644
--- a/tests/kvcache/test_host_tier_live_pool.py
+++ b/tests/kvcache/test_host_tier_live_pool.py
@@ -32,7 +32,7 @@ def test_read_after_a_pool_rebuild_lands_in_the_new_tensors():
     assert torch.all(pool.recurrent_states[:, 5] == 7.0)
 
 
-def test_copies_wait_for_the_engine_stream(monkeypatch):
+def _fake_copy_stream(monkeypatch):
     events = []
     engine = object()
 
@@ -45,9 +45,28 @@ def synchronize(self):
 
     monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine)
     monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext())
+    return FakeStream(), events, engine
+
+
+def test_copies_wait_for_the_engine_stream(monkeypatch):
+    stream, events, engine = _fake_copy_stream(monkeypatch)
     pool = _pool(3)
     tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu")
-    tier.copy_stream = FakeStream()
+    tier.copy_stream = stream
     host = tier.write_gdn(1)
     tier.read_gdn(host, 2)
     assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)]
+
+
+def test_kv_page_copies_wait_for_the_engine_stream(monkeypatch):
+    stream, events, engine = _fake_copy_stream(monkeypatch)
+    kv_pool = SimpleNamespace(
+        _kv_buffer=torch.zeros(2, 1, 4, 4, 1, 2),  # [2, L, P, page, H, D]
+        _scale_buffer=torch.zeros(2, 1, 16, 1),    # [2, L, P*page, H]
+        _block_scale_buffer=None,
+    )
+    tier = HostPrefixTier(_pool(3), kv_pool, gdn_slots_host=2, kv_budget_bytes=1 << 12, device="cpu")
+    tier.copy_stream = stream
+    slots = tier.write_kv(torch.arange(4, 8, dtype=torch.int32))
+    tier.read_kv(slots, torch.tensor([2]))
+    assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)]

From 3bdd459466b158b93311f05a54793cea851db56a Mon Sep 17 00:00:00 2001
From: fangyuan-3149 <314945674@qq.com>
Date: Sun, 30 Aug 2026 09:35:36 +0800
Subject: [PATCH 172/249] fix(engine): reserve the GDN prefill workspace in the
 cache budget

chunk_gated_delta_rule_fwd_h materializes the per-chunk recurrent states
for the whole extend at once -- [NT, H, V, K] with NT = ceil(extend/64) --
plus a transient v_new buffer. Subtract it at the startup sizing and the
--moe-cache-auto plan for hybrid-GDN models.

(cherry picked from commit 5c3b9141b6b658f4d4ebe59f2a68b2208c6f6771)

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit ec8731a65ff5ca2cce0ff073d273f482a8faa5e6)
---
 python/freetoken/engine/engine.py             |  7 ++++-
 python/freetoken/kvcache/linear_state_pool.py | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 3ac453581c..5e9a93bc60 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -37,7 +37,8 @@
 from freetoken.kvcache.base import CacheRebuildRejected
 from freetoken.kvcache.cache_status import _supports_swa_ratio
 from freetoken.kvcache.linear_state_pool import (
-    _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes,
+    _linear_pool_min_slots, _linear_pool_num_slots, gdn_prefill_workspace_bytes,
+    state_pool_bytes,
 )
 
 logger = init_logger(__name__)
@@ -540,6 +541,9 @@ def __init__(self, config: EngineConfig):
         # off it; the KV pool family owns every geometry-specific formula behind the rest.
         available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free)
         available_memory -= state_pool_bytes(config)
+        # reserve the GDN prefill transient (chunk recurrent states + v_new); without
+        # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools
+        available_memory -= gdn_prefill_workspace_bytes(config)
         num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory)
         if config.kv_host_pages > 0:
             # KV host offload: the scheduler, page table and radix tree run in a LOGICAL page
@@ -795,6 +799,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None,
 
         cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config)
         fixed_cache_size += state_pool_bytes(config)  # sibling GDN state pool, engine-summed
+        fixed_cache_size += gdn_prefill_workspace_bytes(config)  # GDN prefill transient
         num_experts = config.model_config.num_experts
         total_experts = config.model_config.num_moe_layers * num_experts
         return resolve_moe_cache_auto(
diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py
index d89556899d..16d2fb537e 100644
--- a/python/freetoken/kvcache/linear_state_pool.py
+++ b/python/freetoken/kvcache/linear_state_pool.py
@@ -261,6 +261,33 @@ def linear_state_bytes_per_req(
 __all__ = ["LinearStatePool", "linear_state_bytes_per_req"]
 
 
+_GDN_CHUNK_SIZE = 64  # keep in lockstep with kernel/fla/chunk_delta_h.py CHUNK_SIZE
+
+
+def gdn_prefill_workspace_bytes(config) -> int:
+    """Peak transient VRAM of one GDN prefill forward.
+
+    ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for
+    the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` --
+    plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for
+    ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for
+    CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills
+    at aggressive memory ratios. 0 for non-GDN models."""
+    linear_group = config.model_config.linear_attention_group()
+    if linear_group is None:
+        return 0
+    extent = min(config.max_extend_tokens, config.max_seq_len)
+    nt = -(-extent // _GDN_CHUNK_SIZE)
+    states = (
+        nt
+        * linear_group.num_value_heads
+        * linear_group.value_head_dim
+        * linear_group.key_head_dim
+    )
+    v_new = extent * linear_group.num_value_heads * linear_group.value_head_dim
+    return (states + v_new) * config.dtype.itemsize
+
+
 def state_pool_bytes(config, num_slots: int | None = None) -> int:
     """Total GDN state-pool bytes at ``num_slots`` PHYSICAL slots (default: the startup
     slot count). The engine adds this to the KV family's fixed cost when budgeting --

From 097abf7954a7c60b2caf4979e1a6b3f49ee387f3 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:29:45 -0400
Subject: [PATCH 173/249] fix(engine): reserve the gdn prefill workspace in the
 rebuild fit-check too

The cherry-picked reservation covered startup sizing and the auto plan
but not validate_rebuild, so a runtime rebuild could grow the pools back
into it. Also price the scheduler's default extend when the config is a
bare EngineConfig, and document that h + v_new is a lower bound on the
GDN prefill peak.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 9d836ed0b0b6c448975b26b4d86368e087dfe0af)
---
 python/freetoken/engine/engine.py             |  7 ++--
 python/freetoken/kvcache/linear_state_pool.py | 13 +++---
 tests/engine/test_cache_budget.py             | 38 +++++++++++++++++
 tests/kvcache/test_linear_state_pool_alloc.py | 42 +++++++++++++++++++
 4 files changed, 91 insertions(+), 9 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 5e9a93bc60..1a92c53a40 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -541,8 +541,8 @@ def __init__(self, config: EngineConfig):
         # off it; the KV pool family owns every geometry-specific formula behind the rest.
         available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free)
         available_memory -= state_pool_bytes(config)
-        # reserve the GDN prefill transient (chunk recurrent states + v_new); without
-        # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools
+        # Reserve part of the GDN prefill transient (chunk states + v_new); the rest of the
+        # prefill activations still live in the (1 - memory_ratio) headroom.
         available_memory -= gdn_prefill_workspace_bytes(config)
         num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory)
         if config.kv_host_pages > 0:
@@ -1281,7 +1281,8 @@ def rebuild_runtime_cache(
             per_expert_bytes=per_expert_bytes, baseline_free=self._baseline_free,
             weights_bytes=self._weights_bytes, current_num_pages=self.num_pages,
             extra_fixed_bytes=(
-                state_pool_bytes(config, target_mamba) if target_mamba is not None else 0
+                (state_pool_bytes(config, target_mamba) if target_mamba is not None else 0)
+                + gdn_prefill_workspace_bytes(config)
             ),
             extra_note=(
                 f", mamba={target_mamba - 1} slots" if target_mamba is not None else ""
diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py
index 16d2fb537e..df5a3396da 100644
--- a/python/freetoken/kvcache/linear_state_pool.py
+++ b/python/freetoken/kvcache/linear_state_pool.py
@@ -265,18 +265,19 @@ def linear_state_bytes_per_req(
 
 
 def gdn_prefill_workspace_bytes(config) -> int:
-    """Peak transient VRAM of one GDN prefill forward.
+    """Bytes of the GDN chunked-prefill buffers that scale with the extend length.
 
     ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for
     the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` --
-    plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for
-    ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for
-    CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills
-    at aggressive memory ratios. 0 for non-GDN models."""
+    plus the ``v_new`` value buffer (288 MiB at an 8192-token extend for qwen4_exp).
+    This is a lower bound on the GDN prefill peak, not the peak: w/u/A/o and the
+    in_proj/conv activations still come out of the ``(1 - memory_ratio)`` headroom.
+    0 for non-GDN models."""
     linear_group = config.model_config.linear_attention_group()
     if linear_group is None:
         return 0
-    extent = min(config.max_extend_tokens, config.max_seq_len)
+    # EngineConfig has no max_extend_tokens (the served config is a SchedulerConfig subclass)
+    extent = min(getattr(config, "max_extend_tokens", 8192), config.max_seq_len)
     nt = -(-extent // _GDN_CHUNK_SIZE)
     states = (
         nt
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 0fcb968252..6f7184b327 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -398,6 +398,44 @@ def slot_limit(self):
     assert captured["max_slots"] == 5
 
 
+
+def test_rebuild_fit_check_reserves_the_gdn_prefill_workspace():
+    """The runtime-rebuild fit-check prices the GDN prefill workspace like the startup sizing,
+    so a rebuild cannot grow the pools back into the space reserved for it."""
+    from freetoken.engine.engine import Engine
+    from freetoken.kvcache.base import CacheRebuildRejected
+    from freetoken.kvcache.linear_state_pool import gdn_prefill_workspace_bytes, state_pool_bytes
+    from freetoken.models.config import LinearGatedDeltaGroupConfig
+
+    group = LinearGatedDeltaGroupConfig(
+        name="linear", layer_ids=(0, 1), num_key_heads=2, num_value_heads=4,
+        key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu",
+    )
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(linear_attention_group=lambda: group),
+        dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1),
+        max_extend_tokens=1024, max_seq_len=4096,
+    )
+    captured = {}
+
+    def validate_rebuild(config, **kwargs):
+        captured.update(kwargs)
+        raise CacheRebuildRejected("stop before teardown")
+
+    engine = SimpleNamespace(
+        config=config, moe_offload_cache=None, kv_offloader=None,
+        linear_state_pool=SimpleNamespace(num_slots=5),
+        kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild),
+        _baseline_free=1 << 30, _weights_bytes=0, num_pages=64,
+    )
+    engine._target_moe_and_expert_bytes = lambda size: (0, 0)
+    with pytest.raises(CacheRebuildRejected):
+        Engine.rebuild_runtime_cache(engine, num_pages=32)
+
+    workspace = gdn_prefill_workspace_bytes(config)
+    assert workspace > 0
+    assert captured["extra_fixed_bytes"] == state_pool_bytes(config, 5) + workspace
+
 # ---------------------------------------------------------------------------
 # offload-cache sizing guard + auto-resolution (_require_offload_cache_size / _adjust_config),
 # the floor rule compute_cache_floors documents above.
diff --git a/tests/kvcache/test_linear_state_pool_alloc.py b/tests/kvcache/test_linear_state_pool_alloc.py
index e9fe405426..e0edd0438e 100644
--- a/tests/kvcache/test_linear_state_pool_alloc.py
+++ b/tests/kvcache/test_linear_state_pool_alloc.py
@@ -9,6 +9,7 @@
 
 from freetoken.kvcache.linear_state_pool import (
     LinearStatePool,
+    gdn_prefill_workspace_bytes,
     linear_state_bytes_per_req,
     state_pool_bytes,
 )
@@ -189,3 +190,44 @@ def test_slot_state_bytes_for_the_real_geometry():
     delta = linear_state_bytes_per_req(group, 1, torch.bfloat16, (spec,)) - \
         linear_state_bytes_per_req(group, 1, torch.bfloat16)
     assert delta == 4 * 2560 * 9 * 2 == 180 * 1024
+
+
+def _gdn_config(**fields):
+    # qwen4_exp GDN geometry: 48 value heads of 128x128
+    group = LinearGatedDeltaGroupConfig(
+        name="linear", layer_ids=(0, 1),
+        num_key_heads=16, num_value_heads=48,
+        key_head_dim=128, value_head_dim=128, conv_kernel_dim=4, output_gate="silu",
+    )
+    mc = SimpleNamespace(linear_attention_group=lambda: group)
+    return SimpleNamespace(model_config=mc, dtype=torch.bfloat16, max_seq_len=262144, **fields)
+
+
+def test_gdn_prefill_workspace_for_the_real_geometry():
+    # h [8192/64 chunks, 48, 128, 128] bf16 = 192 MiB, plus v_new [8192, 48, 128] = 96 MiB
+    assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=8192)) == 288 * 2**20
+    assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=2048)) == 72 * 2**20
+
+
+def test_gdn_prefill_workspace_without_a_scheduler_extend_budget():
+    # a bare EngineConfig has no max_extend_tokens: price the scheduler's 8192 default
+    assert gdn_prefill_workspace_bytes(_gdn_config()) == 288 * 2**20
+
+
+def test_gdn_prefill_workspace_is_capped_by_the_sequence_length():
+    config = _gdn_config(max_extend_tokens=8192)
+    config.max_seq_len = 100  # 2 chunks, 100 value rows
+    assert gdn_prefill_workspace_bytes(config) == (2 * 48 * 128 * 128 + 100 * 48 * 128) * 2
+
+
+def test_gdn_prefill_workspace_is_zero_without_a_linear_group():
+    config = _gdn_config(max_extend_tokens=8192)
+    config.model_config.linear_attention_group = lambda: None
+    assert gdn_prefill_workspace_bytes(config) == 0
+
+
+def test_gdn_prefill_workspace_chunk_matches_the_kernel():
+    from freetoken.kernel.fla.chunk_delta_h import CHUNK_SIZE
+    from freetoken.kvcache.linear_state_pool import _GDN_CHUNK_SIZE
+
+    assert _GDN_CHUNK_SIZE == CHUNK_SIZE

From 2fe2ebd558b266755643befa7b0abe1e46c7a8b2 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:25:31 -0400
Subject: [PATCH 174/249] fix(engine): count min_tokens progress from the
 device length

Under overlap scheduling the sampler prepares batch N before batch N-1's token
lands in req.input_ids, so the host count lagged by one and kept the stop ids
masked for one token past min_tokens.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 8e4891fd54037f19aa42d43422a50cafc38513d9)
---
 python/freetoken/engine/sample.py      |  4 +++-
 tests/engine/test_logits_processors.py | 26 ++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index bca85c56fa..89d2a6fb88 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -220,7 +220,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
         min_ids: list[list[int]] = []
         for local, (req, p) in enumerate(picked):
             if p.min_tokens > 0 and p.min_tokens_stop_ids:
-                generated = len(req.input_ids) - (req.max_device_len - req.output_len)
+                # device_len, not the host ids: under overlap scheduling the previous token
+                # reaches req.input_ids only after this batch is prepared
+                generated = req.device_len - (req.max_device_len - req.output_len)
                 if generated < p.min_tokens:
                     min_rows.append(local)
                     min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size])
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index a9645f4210..c9f498183f 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -85,6 +85,31 @@ def test_min_tokens_masks_the_stop_ids_of_the_row_only():
     assert torch.isfinite(out[0]).all()
 
 
+@pytest.mark.parametrize("overlap", [True, False])
+def test_min_tokens_releases_the_stop_ids_at_output_index_min_tokens(overlap):
+    """Overlap scheduling prepares batch N before batch N-1's token reaches req.input_ids on
+    the host; the count must not lag by that token and hold EOS one step too long."""
+    from freetoken.core import Req
+
+    sp = SamplingParams(min_tokens=3, max_tokens=10)
+    sp.min_tokens_stop_ids = [0]
+    req = Req(input_ids=torch.tensor([5, 6], dtype=torch.int32), table_idx=0, cached_len=0,
+              output_len=10, uid=1, sampling_params=sp, cache_handle=None)
+    sampler = Sampler(CPU, V)
+    masked, pending = [], None
+    for _ in range(6):
+        plan = sampler.prepare(SimpleNamespace(reqs=[req])).plan
+        masked.append(plan is not None and plan.min_rows is not None)
+        req.complete_one()  # the forward launch advances the device length
+        if overlap:
+            if pending is not None:  # batch N-1 drains after batch N launched
+                req.append_host(pending)
+            pending = torch.tensor([7], dtype=torch.int32)
+        else:
+            req.append_host(torch.tensor([7], dtype=torch.int32))
+    assert [i for i, m in enumerate(masked) if m] == [0, 1, 2]
+
+
 def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability():
     logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]]))
     plan = _plan(
@@ -104,6 +129,7 @@ def _req(
     ids = torch.tensor(prompt + generated, dtype=torch.int32)
     return SimpleNamespace(
         input_ids=ids,
+        device_len=len(ids),
         output_len=max_tokens,
         max_device_len=len(prompt) + max_tokens,
         sampling_params=SamplingParams(**sp),

From 707989978937f0b7dae44a812feb37537df9b5f4 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:23:00 -0400
Subject: [PATCH 175/249] test(engine): import pytest for the parametrized
 min_tokens test

Assisted-by: Claude Opus 5.5
---
 tests/engine/test_logits_processors.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index c9f498183f..abd2e1f939 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -5,6 +5,7 @@
 
 from types import SimpleNamespace
 
+import pytest
 import torch
 from freetoken.core import SamplingParams
 from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors

From d541e96ee1fb6097704d22d35705b4539ed829a0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:01:42 -0400
Subject: [PATCH 176/249] fix(server): let protocol-level thinking requests win
 over --default-thinking-mode

returned early, so reasoning_effort / reasoning.effort / thinking.type were silently turned
off under --default-thinking-mode chat. count_tokens also skipped the default.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 8276ccd49a8fcb9123ea28f113ac39684b725665)
---
 python/freetoken/server/anthropic_api.py   | 16 ++--
 python/freetoken/server/args.py            | 17 ++--
 python/freetoken/server/model_meta.py      | 19 +++++
 python/freetoken/server/openai_api.py      | 33 +-------
 python/freetoken/server/responses_api.py   |  9 +--
 tests/server/test_default_thinking_mode.py | 94 +++++++++++++++++++---
 6 files changed, 127 insertions(+), 61 deletions(-)

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 4cb4b34822..a7c88b141a 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -166,7 +166,9 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state:
     # so it must not fall into the convert/empty-prompt ValueError branch.
     try:
         messages, template_tools, _, ctk = convert_anthropic_prompt(
-            req, reasoning_parser=getattr(state.config, "reasoning_parser", None)
+            req,
+            reasoning_parser=getattr(state.config, "reasoning_parser", None),
+            default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"),
         )
     except ValueError as exc:
         return _anthropic_error_response(400, "invalid_request_error", str(exc))
@@ -198,6 +200,7 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state:
 def convert_anthropic_prompt(
     req: AnthropicMessagesRequest | AnthropicCountTokensRequest,
     reasoning_parser: str | None = None,
+    default_thinking_mode: str | None = None,
 ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None, list[dict[str, Any]] | None, dict[str, Any]]:
     """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt
     side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so
@@ -310,7 +313,7 @@ def convert_anthropic_prompt(
     # Native extended-thinking toggle -> template kwargs, broadcast in every
     # spelling the ecosystem's templates read (a bare enable_thinking bool is
     # inert for templates that read a different knob, e.g. M3's thinking_mode).
-    from .model_meta import thinking_toggle_kwargs
+    from .model_meta import apply_default_thinking_mode, thinking_toggle_kwargs
 
     ctk: dict[str, Any] = {}
     if req.thinking:
@@ -318,6 +321,10 @@ def convert_anthropic_prompt(
             ctk = thinking_toggle_kwargs(True)
         elif req.thinking.get("type") == "disabled":
             ctk = thinking_toggle_kwargs(False)
+    else:
+        # Any thinking block (e.g. "adaptive") is the client's explicit choice; the server
+        # default only covers requests that leave thinking out.
+        ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
 
     return render_messages(messages), template_tools, parser_tools, ctk
 
@@ -330,12 +337,9 @@ def convert_anthropic_to_genspec(
     inline_system_policy: str = "auto",
     default_thinking_mode: str | None = None,
 ) -> GenSpec:
-    from .openai_api import apply_default_thinking_mode
-
     messages, template_tools, parser_tools, ctk = convert_anthropic_prompt(
-        req, reasoning_parser=reasoning_parser
+        req, reasoning_parser=reasoning_parser, default_thinking_mode=default_thinking_mode
     )
-    ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
     return GenSpec(
         messages=messages,
         sampling_params=resolve_sampling(
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 3cf6521112..e25b9e5df0 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -52,10 +52,10 @@ class ServerArgs(SchedulerConfig):
     # responses. None disables it (default for models without a reasoning protocol).
     reasoning_parser: str | None = None
     # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the
-    # current per-request behavior; "chat" forces enable_thinking=False for every request
-    # that does not explicitly set it in chat_template_kwargs (for OpenAI-compatible
-    # clients that never send template kwargs, like Vercel AI SDK or llama-swap);
-    # "thinking" forces thinking on the same way.
+    # current per-request behavior; "chat" turns thinking off for every request that does
+    # not choose itself (template kwargs or a protocol-level effort/thinking field), for
+    # OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK or
+    # llama-swap; "thinking" turns it on the same way.
     default_thinking_mode: str = "auto"
     # "model": fill unspecified request sampling params from generation_config.json
     # (temperature/top_k/top_p), like sglang. "none": use framework defaults only.
@@ -720,10 +720,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         choices=["auto", "chat", "thinking"],
         help=(
             "Server-wide default thinking mode for reasoning-capable models. 'auto' keeps "
-            "the current per-request behavior; 'chat' forces enable_thinking=False for "
-            "every request that does not explicitly set it in chat_template_kwargs (for "
-            "OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK "
-            "or llama-swap); 'thinking' forces thinking on the same way."
+            "the current per-request behavior; 'chat' turns thinking off for every request "
+            "that does not choose itself (chat_template_kwargs, reasoning_effort, "
+            "reasoning.effort, thinking) -- for OpenAI-compatible clients that never send "
+            "template kwargs, like Vercel AI SDK or llama-swap; 'thinking' turns it on the "
+            "same way."
         ),
     )
 
diff --git a/python/freetoken/server/model_meta.py b/python/freetoken/server/model_meta.py
index 8ce8009a05..c9187ad629 100644
--- a/python/freetoken/server/model_meta.py
+++ b/python/freetoken/server/model_meta.py
@@ -119,6 +119,25 @@ def effort_toggle_kwargs(
     return mapped
 
 
+def apply_default_thinking_mode(
+    ctk: dict[str, Any] | None,
+    default_mode: str | None,
+) -> dict[str, Any] | None:
+    """Merge the server's --default-thinking-mode into a request's template kwargs.
+
+    An explicit per-request value always wins; the server default only fills
+    in what the request left unset. "auto" (the default) is a no-op. Adapters
+    apply it after folding their protocol-level toggles (effort_toggle_kwargs),
+    so those count as explicit too.
+    """
+    if default_mode not in ("chat", "thinking"):
+        return ctk
+    ctk = dict(ctk or {})
+    if not any(key in ctk for key in _THINKING_KWARG_KEYS):
+        ctk.update(thinking_toggle_kwargs(default_mode == "thinking"))
+    return ctk
+
+
 def moe_total_experts(config: Any) -> int:
     """Total routed-expert slots the model has: experts per layer x MoE layers. Matches the
     engine's own basis (``Engine._resolve_auto_moe_cache_size``), so a residency rate derived
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b09c480e1d..54c10ac2e6 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -71,34 +71,6 @@ def _thinking_type(req: Any) -> str | None:
 
 
 
-def apply_default_thinking_mode(
-    ctk: dict[str, Any] | None,
-    default_mode: str | None,
-) -> dict[str, Any] | None:
-    """Merge the server's --default-thinking-mode into a request's template kwargs.
-
-    An explicit per-request value always wins; the server default only fills
-    in what the request left unset. "auto" (the default) is a no-op.
-    """
-    if not default_mode or default_mode == "auto":
-        return ctk
-    if ctk is None:
-        ctk = {}
-    else:
-        ctk = dict(ctk)
-    has_explicit = (
-        "enable_thinking" in ctk
-        or "thinking" in ctk
-        or "thinking_mode" in ctk
-    )
-    if not has_explicit:
-        if default_mode == "chat":
-            ctk["enable_thinking"] = False
-        elif default_mode == "thinking":
-            ctk["enable_thinking"] = True
-    return ctk
-
-
 def chat_request_to_genspec(
     req: ChatCompletionRequest,
     model_sampling: dict[str, Any],
@@ -106,12 +78,13 @@ def chat_request_to_genspec(
     default_thinking_mode: str | None = None,
 ) -> GenSpec:
     """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params')."""
-    from .model_meta import effort_toggle_kwargs
+    from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs
 
-    ctk = apply_default_thinking_mode(req.chat_template_kwargs, default_thinking_mode)
+    ctk = req.chat_template_kwargs
     thinking_type = _thinking_type(req)
     if req.reasoning_effort or thinking_type:
         ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type)
+    ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
     if req.continue_final_message:
         ctk = {**ctk, "continue_final_message": True}
     sampling_params = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens)
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index 9fd063a992..a84c68cb6d 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -245,15 +245,12 @@ def convert_responses_to_genspec(
     else:
         template_tools, parser_tools = split_tool_lists(raw_tools, selected)
 
-    from .model_meta import effort_toggle_kwargs
+    from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs
 
-    from .openai_api import apply_default_thinking_mode
-
-    ctk = apply_default_thinking_mode(
-        dict(getattr(req, "chat_template_kwargs", None) or {}), default_thinking_mode
-    )
+    ctk = dict(getattr(req, "chat_template_kwargs", None) or {})
     if req.reasoning:
         ctk = effort_toggle_kwargs(req.reasoning.get("effort"), ctk)
+    ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
 
     return GenSpec(
         messages=render_messages(messages),
diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py
index ac979e625d..7f9938f60c 100644
--- a/tests/server/test_default_thinking_mode.py
+++ b/tests/server/test_default_thinking_mode.py
@@ -1,7 +1,18 @@
 """Unit tests for the --default-thinking-mode server flag merge logic."""
 from __future__ import annotations
 
-from freetoken.server.openai_api import apply_default_thinking_mode
+import asyncio
+from types import SimpleNamespace
+
+from freetoken.server import anthropic_api as A
+from freetoken.server.anthropic_models import AnthropicCountTokensRequest, AnthropicMessagesRequest
+from freetoken.server.api_models import ChatCompletionRequest
+from freetoken.server.model_meta import apply_default_thinking_mode
+from freetoken.server.openai_api import chat_request_to_genspec
+from freetoken.server.responses_api import ResponsesRequest, convert_responses_to_genspec
+
+OFF = {"enable_thinking": False, "thinking_mode": "disabled"}
+ON = {"enable_thinking": True, "thinking_mode": "enabled"}
 
 
 def test_auto_is_noop():
@@ -13,17 +24,14 @@ def test_auto_is_noop():
 
 
 def test_chat_fills_unset_kwargs():
-    assert apply_default_thinking_mode(None, "chat") == {"enable_thinking": False}
-    assert apply_default_thinking_mode({}, "chat") == {"enable_thinking": False}
-    assert apply_default_thinking_mode({"other": 1}, "chat") == {
-        "other": 1,
-        "enable_thinking": False,
-    }
+    assert apply_default_thinking_mode(None, "chat") == OFF
+    assert apply_default_thinking_mode({}, "chat") == OFF
+    assert apply_default_thinking_mode({"other": 1}, "chat") == {"other": 1, **OFF}
 
 
 def test_thinking_fills_unset_kwargs():
-    assert apply_default_thinking_mode(None, "thinking") == {"enable_thinking": True}
-    assert apply_default_thinking_mode({}, "thinking") == {"enable_thinking": True}
+    assert apply_default_thinking_mode(None, "thinking") == ON
+    assert apply_default_thinking_mode({}, "thinking") == ON
 
 
 def test_explicit_request_value_wins():
@@ -38,14 +46,78 @@ def test_explicit_request_value_wins():
     assert apply_default_thinking_mode({"thinking_mode": "thinking"}, "chat") == {
         "thinking_mode": "thinking"
     }
+    # A raw graded-effort request asks for thinking too.
+    assert apply_default_thinking_mode({"reasoning_effort": "high"}, "chat") == {
+        "reasoning_effort": "high"
+    }
 
 
 def test_original_dict_not_mutated():
     original = {"other": 1}
     result = apply_default_thinking_mode(original, "chat")
-    assert result == {"other": 1, "enable_thinking": False}
+    assert result == {"other": 1, **OFF}
     assert original == {"other": 1}
 
 
 def test_unknown_mode_is_noop():
-    assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1}
\ No newline at end of file
+    assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1}
+
+
+def _chat(**extra):
+    body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **extra}
+    return chat_request_to_genspec(
+        ChatCompletionRequest.model_validate(body), {}, default_thinking_mode="chat"
+    ).chat_template_kwargs
+
+
+def test_chat_default_fills_a_request_without_thinking_controls():
+    assert _chat() == OFF
+
+
+def test_chat_default_keeps_protocol_level_thinking_requests():
+    assert _chat(reasoning_effort="high")["enable_thinking"] is True
+    assert _chat(thinking={"type": "enabled"})["enable_thinking"] is True
+    assert _chat(reasoning_effort="none")["enable_thinking"] is False
+
+
+def test_responses_chat_default_keeps_reasoning_effort():
+    def ctk(**extra):
+        req = ResponsesRequest.model_validate({"model": "m", "input": "hi", **extra})
+        return convert_responses_to_genspec(req, {}, default_thinking_mode="chat").chat_template_kwargs
+
+    assert ctk() == OFF
+    on = ctk(reasoning={"effort": "medium"})
+    assert on["enable_thinking"] is True and on["reasoning_effort"] == "medium"
+
+
+def _anthropic(**extra):
+    return AnthropicMessagesRequest.model_validate(
+        {"model": "claude-x", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}], **extra}
+    )
+
+
+def test_anthropic_chat_default_applies_only_without_a_thinking_block():
+    def ctk(**extra):
+        return A.convert_anthropic_to_genspec(
+            _anthropic(**extra), {}, default_thinking_mode="chat"
+        ).chat_template_kwargs
+
+    assert ctk() == OFF
+    assert ctk(thinking={"type": "enabled", "budget_tokens": 1024}) == ON
+    assert ctk(thinking={"type": "adaptive"}) == {}
+
+
+def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch):
+    captured = {}
+
+    async def fake_count(messages, template_tools, ctk, state):
+        captured["ctk"] = ctk
+        return 1
+
+    monkeypatch.setattr(A, "count_prompt_tokens", fake_count)
+    state = SimpleNamespace(config=SimpleNamespace(reasoning_parser=None, default_thinking_mode="chat"))
+    body = {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}]}
+    asyncio.run(A.handle_anthropic_count_tokens(AnthropicCountTokensRequest.model_validate(body), state))
+
+    generation = A.convert_anthropic_to_genspec(_anthropic(), {}, default_thinking_mode="chat")
+    assert captured["ctk"] == generation.chat_template_kwargs == OFF

From 8d51b51fc1854294018af6f21f17de08e92b0807 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:57:18 -0400
Subject: [PATCH 177/249] fix(server): type allOf-wrapped refs and untyped
 enum/const tool parameters

#435 sends any parameter it cannot type to the loose JSON parser, so an untyped
string enum member such as "1" came back as the integer 1, outside the enum.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit ab2693c977c64e425d428468364e2a2743f170e4)
---
 .../freetoken/server/function_call_parser.py  | 38 +++++++++-
 ..._function_call_parser_schema_resolution.py | 70 +++++++++++++++++++
 2 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py
index 08c60a2993..88cc65b367 100644
--- a/python/freetoken/server/function_call_parser.py
+++ b/python/freetoken/server/function_call_parser.py
@@ -382,7 +382,27 @@ def _resolve_local_ref(self, schema: Dict, root_schema: Dict, _seen: Optional[se
             return self._resolve_local_ref(merged, root_schema, _seen)
         return merged
 
-    def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict:
+    @staticmethod
+    def _literal_schema_type(schema: Dict) -> Optional[str]:
+        """Type of an untyped ``enum``/``const`` schema, inferred from its literals."""
+        if "const" in schema:
+            values = [schema["const"]]
+        elif isinstance(schema.get("enum"), list):
+            values = schema["enum"]
+        else:
+            return None
+        values = [v for v in values if v is not None]
+        if not values:
+            return None
+        if all(isinstance(v, str) for v in values):
+            return "string"
+        if all(isinstance(v, bool) for v in values):
+            return "boolean"
+        if all(isinstance(v, int) and not isinstance(v, bool) for v in values):
+            return "integer"
+        return "loose"
+
+    def _normalize_param_schema(self, schema: Dict, root_schema: Dict, _depth: int = 0) -> Dict:
         """Resolve the schema parts needed for tool argument typing."""
         if not isinstance(schema, dict):
             return {"type": "string"}
@@ -395,6 +415,18 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict:
             return result
         if isinstance(param_type, str):
             return schema
+        all_of = schema.get("allOf")
+        if (
+            isinstance(all_of, list)
+            and len(all_of) == 1
+            and isinstance(all_of[0], dict)
+            and _depth < 16  # bounds a self-referential allOf chain
+        ):
+            # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach
+            # siblings such as default/description; type it like the bare $ref.
+            merged = dict(all_of[0])
+            merged.update({k: v for k, v in schema.items() if k != "allOf"})
+            return self._normalize_param_schema(merged, root_schema, _depth + 1)
         for keyword in ("oneOf", "anyOf"):
             branches = schema.get(keyword)
             if not isinstance(branches, list):
@@ -405,6 +437,8 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict:
                     continue
                 branch = self._resolve_local_ref(branch, root_schema)
                 branch_type = branch.get("type")
+                if not isinstance(branch_type, str):
+                    branch_type = self._literal_schema_type(branch)
                 if isinstance(branch_type, str) and branch_type != "null":
                     types.append(branch_type)
             types = list(dict.fromkeys(types))
@@ -420,7 +454,7 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict:
             result["type"] = "array"
             return result
         result = dict(schema)
-        result["type"] = "loose"
+        result["type"] = self._literal_schema_type(schema) or "loose"
         return result
 
     def _get_tool_schema(self, func_name: str, tools: List[Tool]) -> Dict:
diff --git a/tests/server/test_function_call_parser_schema_resolution.py b/tests/server/test_function_call_parser_schema_resolution.py
index fbc0fe270b..468acc3bd2 100644
--- a/tests/server/test_function_call_parser_schema_resolution.py
+++ b/tests/server/test_function_call_parser_schema_resolution.py
@@ -375,6 +375,67 @@ def _parse_args(parser_name: str, prop_schema: dict, root_extra: dict | None, ra
         5,
         id="ambiguous-union-loose-json",
     ),
+
+    # An enum/const with no "type" takes its type from the literals, so a
+    # numeric-looking string member is not coerced out of the enum.
+    pytest.param(
+        {"enum": ["1", "2"]},
+        None,
+        "1",
+        "1",
+        id="untyped-string-enum",
+    ),
+    pytest.param(
+        {"enum": [1, 2]},
+        None,
+        "2",
+        2,
+        id="untyped-integer-enum",
+    ),
+    pytest.param(
+        {"const": "10"},
+        None,
+        "10",
+        "10",
+        id="untyped-string-const",
+    ),
+    pytest.param(
+        {
+            "anyOf": [
+                {"enum": ["1", "2"]},
+                {"type": "null"},
+            ]
+        },
+        None,
+        "1",
+        "1",
+        id="anyof-untyped-string-enum-nullable",
+    ),
+
+    # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach
+    # siblings such as default/description.
+    pytest.param(
+        {"allOf": [{"$ref": "#/$defs/Mode"}], "default": "a"},
+        {
+            "$defs": {
+                "Mode": {"type": "string", "enum": ["a", "true"]},
+            }
+        },
+        "true",
+        "true",
+        id="allof-ref-string-enum",
+    ),
+    pytest.param(
+        {"allOf": [{"$ref": "#/$defs/Limit"}], "description": "max rows"},
+        {
+            "$defs": {
+                "Limit": {"type": "integer"},
+            }
+        },
+        "5",
+        5,
+        id="allof-ref-integer",
+    ),
 ]
 
 
@@ -454,6 +515,15 @@ def test_schema_aware_parsers_resolve_indirect_schema_types(
         },
         id="mutually-recursive-ref",
     ),
+    pytest.param(
+        {"allOf": [{"$ref": "#/$defs/A"}]},
+        {
+            "$defs": {
+                "A": {"allOf": [{"$ref": "#/$defs/A"}]},
+            }
+        },
+        id="self-referential-allof",
+    ),
 ]
 
 

From 538b92707eaa26a1c1d1b11cfc25c03ade0055f4 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:03:44 -0400
Subject: [PATCH 178/249] fix(server): keep /v1/models/ listing the model after
 the retrieve route

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 7ccc433dcc8608b23f33c1a4be6313b5f9307c47)
---
 python/freetoken/server/openai_api.py |  3 +++
 tests/server/test_openai_api.py       | 26 ++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 54c10ac2e6..7f94efc250 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -176,6 +176,9 @@ async def v1_model(model_id: str):
         OpenAI-compatible clients may use this endpoint to validate a model
         name before sending completion requests.
         """
+        if not model_id:
+            # The :path converter also matches "/v1/models/", which used to redirect to the list.
+            return await v1_models()
         state = get_state()
         served_model_id = _served_model_name(state)
 
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index 376cc2a7ba..7e6d4ed30e 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -523,6 +523,32 @@ def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool():
     assert card["max_model_len"] == 40000
 
 
+def test_model_retrieve_route_returns_the_served_card_or_404():
+    state = FakeState([])
+    state.config.max_seq_len = 262144
+    app = FastAPI()
+    register_openai_routes(app, lambda: state, lambda: {})
+    client = TestClient(app)
+
+    card = client.get("/v1/models/unit-model").json()
+    assert card["id"] == "unit-model" and card["context_length"] == 262144
+
+    missing = client.get("/v1/models/other-model")
+    assert missing.status_code == 404
+    assert missing.json()["error"]["code"] == "model_not_found"
+
+
+def test_models_route_with_a_trailing_slash_still_lists():
+    state = FakeState([])
+    app = FastAPI()
+    register_openai_routes(app, lambda: state, lambda: {})
+
+    response = TestClient(app).get("/v1/models/")
+
+    assert response.status_code == 200
+    assert response.json()["data"][0]["id"] == "unit-model"
+
+
 async def _collect(generator):
     return [chunk async for chunk in generator]
 

From 2e1021ea69a85bee5d22d633952318a8c96d1157 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:22:47 -0400
Subject: [PATCH 179/249] fix(server): open the streamed thinking block with
 the signature field

The full response and the closing signature_delta already send signature="", but
exclude_none dropped it from content_block_start, where the SDK types it as required.
Taken from upstream #517 (b0442764b014bd32a68ba29055eb20eda9bfb026).

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 55a8e72f21cb679b585cfe6f7f738aafedeb7369)
---
 python/freetoken/server/anthropic_api.py | 4 +++-
 tests/server/test_anthropic_api.py       | 3 ++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index a7c88b141a..545cf4e7cc 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -480,7 +480,9 @@ def _open_thinking() -> str:
         block_open = "thinking"
         return _event(AnthropicStreamEvent(
             type="content_block_start", index=block_index,
-            content_block=AnthropicContentBlock(type="thinking", thinking=""),
+            # The SDK types signature as a required string on the opening block, before any
+            # signature_delta. Local reasoning is unsigned, like the full-response path.
+            content_block=AnthropicContentBlock(type="thinking", thinking="", signature=""),
         ))
 
     def _stop_block() -> list[str]:
diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py
index 1a09b83a1c..890339522c 100644
--- a/tests/server/test_anthropic_api.py
+++ b/tests/server/test_anthropic_api.py
@@ -431,7 +431,8 @@ def test_stream_reasoning_as_thinking_block():
         "message_stop",
     ]
     blocks = [e[1] for e in collected if e[0] == "content_block_start"]
-    assert blocks[0]["content_block"]["type"] == "thinking"
+    # signature is a required string in the SDK's thinking block, present from the start.
+    assert blocks[0]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""}
     assert blocks[0]["index"] == 0
     assert blocks[1]["content_block"]["type"] == "text"
     assert blocks[1]["index"] == 1

From b64232a8f8a582b17e664be6309421e825cb7225 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:22:47 -0400
Subject: [PATCH 180/249] test(server): exercise detector SDK and lifecycle
 contracts

Upstream #517 tests without its docs and pyproject dev-extra: the SDK tests importorskip
anthropic/openai/httpx so the server suite still collects without them.

(cherry picked from commit b0442764b014bd32a68ba29055eb20eda9bfb026, tests only)

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 50f99d167304858d9eaebfbcd3d00aa3c1f06430)
---
 tests/server/test_accounting.py           | 121 ++++++++++++
 tests/server/test_function_call_parser.py | 129 +++++++++++++
 tests/server/test_protocol_compat.py      | 213 ++++++++++++++++++++++
 tests/server/test_reasoning_parser.py     | 108 +++++++++++
 tests/server/test_supervisor.py           |  37 ++++
 5 files changed, 608 insertions(+)
 create mode 100644 tests/server/test_accounting.py
 create mode 100644 tests/server/test_protocol_compat.py
 create mode 100644 tests/server/test_reasoning_parser.py

diff --git a/tests/server/test_accounting.py b/tests/server/test_accounting.py
new file mode 100644
index 0000000000..fe29d414d8
--- /dev/null
+++ b/tests/server/test_accounting.py
@@ -0,0 +1,121 @@
+"""Failure/retry and concurrency contracts for the terminal accounting barrier."""
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from freetoken.message import UserReply
+from freetoken.server.accounting import AccountingDrainError, prepare_stop_accounting
+from freetoken.server.stats import StatsTracker
+
+
+def state_with_request():
+    stats = StatsTracker()
+    stats.on_new_user(7)
+    return SimpleNamespace(stats=stats, maintenance_state="serving", instance_id="instance",
+                           ready_at=None, config=SimpleNamespace(served_model_name="test"))
+
+
+def terminal(state, tokens=3):
+    state.stats.observe(UserReply(uid=7, incremental_output="", finished=True,
+                                 completion_tokens_delta=tokens, error="aborted"))
+
+
+def test_drain_timeout_then_late_terminal_can_retry_without_losing_tokens():
+    async def run():
+        state = state_with_request()
+        aborted = []
+
+        async def abort(uid):
+            aborted.append(uid)
+            state.stats.on_abort(uid)  # enqueueing abort is NOT a terminal reply
+
+        state.abort_user = abort
+        with pytest.raises(AccountingDrainError, match="abort barrier timed out"):
+            await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0)
+        assert state.maintenance_state == "stopping"
+        assert state.stats.active == 1 and not hasattr(state, "_sealed_accounting")
+        terminal(state, tokens=11)  # last sampled-token delta arrives AFTER the timeout
+        sealed = await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0)
+        assert sealed["completion_tokens_total"] == 11
+        assert sealed["drain_complete"] is True
+        assert aborted == [7]
+        sealed["completion_tokens_total"] = -1
+        assert (await prepare_stop_accounting(state))["completion_tokens_total"] == 11
+    asyncio.run(run())
+
+
+def test_concurrent_prepare_stop_serializes_abort_and_returns_independent_snapshots():
+    async def run():
+        state = state_with_request()
+        entered, release = asyncio.Event(), asyncio.Event()
+        aborted = []
+
+        async def abort(uid):
+            assert state.maintenance_state == "stopping"
+            aborted.append(uid)
+            entered.set()
+            await release.wait()
+            terminal(state)
+
+        state.abort_user = abort
+        first = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0))
+        await entered.wait()
+        second = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0))
+        await asyncio.sleep(0)  # second caller reaches the shared lock, no wall-clock race
+        assert not second.done()
+        release.set()
+        a, b = await asyncio.gather(first, second)
+        assert a == b and a is not b
+        assert aborted == [7]
+        assert a["completion_tokens_total"] == 3
+    asyncio.run(run())
+
+
+def test_abort_transport_failure_preserves_engine_and_closes_gate():
+    state = state_with_request()
+
+    async def abort(uid):
+        raise ConnectionError("scheduler disconnected")
+
+    state.abort_user = abort
+    with pytest.raises(AccountingDrainError, match="failed to abort active requests") as error:
+        asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0))
+    assert isinstance(error.value.__cause__, ConnectionError)
+    assert state.maintenance_state == "stopping" and state.stats.active == 1
+    assert not hasattr(state, "_sealed_accounting")
+
+
+def test_unidentified_active_request_cannot_be_sealed():
+    state = state_with_request()
+    state.stats = SimpleNamespace(active=1, inflight_uids=())
+    with pytest.raises(AccountingDrainError, match="unidentified"):
+        asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0))
+    assert state.maintenance_state == "stopping"
+    assert not hasattr(state, "_sealed_accounting")
+
+
+@pytest.mark.parametrize("maintenance", ["rebuilding", "unknown"])
+def test_invalid_prepare_transition_preserves_previous_state(maintenance):
+    state = state_with_request()
+    state.maintenance_state = maintenance
+    with pytest.raises(AccountingDrainError):
+        asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0))
+    assert state.maintenance_state == maintenance
+    assert not hasattr(state, "_sealed_accounting")
+
+
+def test_cancelled_drain_releases_lock_without_reopening_admission():
+    async def run():
+        state = state_with_request()
+        task = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=10))
+        await asyncio.sleep(0)  # admission closes before the first drain wait
+        assert state.maintenance_state == "stopping"
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert not state._accounting_prepare_lock.locked()
+        assert not hasattr(state, "_sealed_accounting")
+        terminal(state)
+        assert (await prepare_stop_accounting(state))["drain_complete"] is True
+    asyncio.run(run())
diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py
index bc22296760..23cd12272e 100644
--- a/tests/server/test_function_call_parser.py
+++ b/tests/server/test_function_call_parser.py
@@ -371,3 +371,132 @@ def test_streaming_support_flags():
     # test_streaming_model_matrix.py::test_non_streaming_detector_falls_back_to_buffered_parse).
     for name in SUPPORTED_TOOL_CALL_PARSERS:
         assert FunctionCallParser(TOOLS, tool_call_parser=name).supports_streaming() is True
+
+
+# One class per concrete detector makes format coverage visible in pytest output.
+# A shared contract avoids accidentally giving a newly added format weaker checks.
+class _DetectorContract:
+    parser_name: str
+    block: str
+    truncate_before: str
+    buffered_recovery = False
+
+    def parser(self):
+        return FunctionCallParser(OPENCODE_TOOLS, self.parser_name)
+
+    def test_complete_call(self):
+        result = self.parser().parse_non_stream(self.block)
+        assert result.normal_text == ""
+        assert len(result.calls) == 1
+        assert result.calls[0].name == "read"
+        assert json.loads(result.calls[0].parameters) == {"filePath": "/tmp/test_calc.py"}
+
+    def test_every_character_boundary(self):
+        parser = self.parser()
+        texts, calls = _feed(parser, self.block)
+        assert "".join(texts) + parser.finish_stream() == ""
+        assert [call.name for call in calls if call.name] == ["read"]
+        assert json.loads("".join(call.parameters for call in calls)) == {
+            "filePath": "/tmp/test_calc.py"
+        }
+        # A completed call must never be duplicated by the end-of-stream recovery.
+        assert parser.recover_truncated_call() == []
+
+    def test_truncated_mid_call_recovery(self):
+        parser = self.parser()
+        cut = self.block.rindex(self.truncate_before)
+        texts, calls = _feed(parser, self.block[:cut])
+        recovered = parser.recover_truncated_call()
+        assert bool(recovered) is self.buffered_recovery
+        # Incremental detectors have already consumed the opener and emitted the
+        # call identity. Their recovery method MUST NOT invent a duplicate call:
+        # generation.py recovers missing arguments from their partial-parse ledger.
+        # Buffered M3 instead recovers a complete call from the retained markup.
+        all_calls = calls + recovered
+        assert [call.name for call in all_calls if call.name] == ["read"]
+        arguments = "".join(call.parameters for call in all_calls)
+        try:
+            parsed = json.loads(arguments)
+        except json.JSONDecodeError:
+            parsed = json.loads(parser.unstreamed_arguments(0))
+        assert parsed == {"filePath": "/tmp/test_calc.py"}
+        assert parser.recover_truncated_call() == []
+        assert "".join(texts) + parser.finish_stream() == ""
+
+
+class TestQwen25Detector(_DetectorContract):
+    parser_name = 'qwen25'
+    block = '<tool_call>{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}</tool_call>'
+    truncate_before = '</tool_call>'
+
+
+class TestMistralDetector(_DetectorContract):
+    parser_name = 'mistral'
+    block = '[TOOL_CALLS] [{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}]'
+    truncate_before = '}]'
+
+
+class TestLlama32Detector(_DetectorContract):
+    parser_name = 'llama3'
+    block = '<|python_tag|>{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}'
+    truncate_before = '}'
+
+
+class TestGlm47Detector(_DetectorContract):
+    parser_name = 'glm47'
+    block = '<tool_call>read<arg_key>filePath</arg_key><arg_value>/tmp/test_calc.py</arg_value></tool_call>'
+    truncate_before = '</tool_call>'
+
+
+class TestDeepSeekV32Detector(_DetectorContract):
+    parser_name = 'deepseekv32'
+    block = '<｜DSML｜function_calls><｜DSML｜invoke name="read"><｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py</｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>'
+    truncate_before = '</｜DSML｜invoke>'
+
+
+class TestQwen3CoderDetector(_DetectorContract):
+    parser_name = 'qwen3_coder'
+    block = '<tool_call><function=read><parameter=filePath>/tmp/test_calc.py</parameter></function></tool_call>'
+    truncate_before = '</function>'
+
+
+class TestGemma4Detector(_DetectorContract):
+    parser_name = 'gemma4'
+    block = '<|tool_call>call:read{filePath:<|"|>/tmp/test_calc.py<|"|>}<tool_call|>'
+    truncate_before = '<tool_call|>'
+
+
+class TestMiniMaxDetector(_DetectorContract):
+    parser_name = 'minimax'
+    block = '<minimax:tool_call><invoke name="read"><parameter name="filePath">/tmp/test_calc.py</parameter></invoke></minimax:tool_call>'
+    truncate_before = '</invoke>'
+
+
+class TestMiniMaxM3Detector(_DetectorContract):
+    parser_name = 'minimax_m3'
+    block = ']<]minimax[>[<tool_call>\n]<]minimax[>[<invoke name="read">]<]minimax[>[<filePath>/tmp/test_calc.py]<]minimax[>[</filePath>]<]minimax[>[</invoke>\n]<]minimax[>[</tool_call>'
+    truncate_before = ']<]minimax[>[</invoke>'
+    buffered_recovery = True
+
+
+class TestGptOssDetector(_DetectorContract):
+    parser_name = 'gpt_oss'
+    block = '<|start|>assistant<|channel|>commentary to=functions.read <|constrain|>json<|message|>{"filePath": "/tmp/test_calc.py"}<|end|>'
+    truncate_before = '<|end|>'
+
+
+class TestMuseGlimmerDetector(_DetectorContract):
+    parser_name = 'muse_glimmer'
+    block = '<|start|>assistant to=read<|message|><atem:function_calls>\n<atem:invoke name="read">\n<atem:parameter name="filePath">/tmp/test_calc.py</atem:parameter>\n</atem:invoke>\n</atem:function_calls><|eot|>'
+    truncate_before = '</atem:invoke>'
+
+
+def test_contract_classes_cover_every_concrete_detector():
+    from freetoken.server.function_call_parser import BaseFormatDetector
+
+    covered = {
+        type(cls().parser().detector)
+        for cls in _DetectorContract.__subclasses__()
+    }
+    assert covered == set(BaseFormatDetector.__subclasses__())
+    assert covered == set(FunctionCallParser.ToolCallParserEnum.values())
diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py
new file mode 100644
index 0000000000..8dc22908d6
--- /dev/null
+++ b/tests/server/test_protocol_compat.py
@@ -0,0 +1,213 @@
+"""Real SDK model/stream validation over the real HTTP adapters and GenSpec path.
+
+Only the scheduler/tokenizer transport is scripted. In-process ASGI transport
+keeps these deterministic and CPU-only; it does not prove TCP chunk timing,
+disconnect behavior, tokenizer templates, or a real model's output quality.
+"""
+import asyncio
+import json
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI
+
+from freetoken.message import TokenizeMsg, UserReply
+from freetoken.server.anthropic_api import register_anthropic_routes
+from freetoken.server.openai_api import register_openai_routes
+from freetoken.server.responses_api import register_responses_routes
+
+# The SDKs are optional test dependencies: skip what needs a missing one instead of
+# failing collection of the whole server suite. The Anthropic half skips per test.
+httpx = pytest.importorskip("httpx")
+openai = pytest.importorskip("openai")
+
+
+SCHEMA = {"type": "object", "properties": {"city": {"type": "string"}}}
+TOOL = {"type": "function", "function": {"name": "get_weather", "parameters": SCHEMA}}
+CALL = '<tool_call>{"name":"get_weather","arguments":{"city":"Paris"}}</tool_call>'
+
+
+class ScriptedBackend:
+    def __init__(self, text, *, reasoning=False):
+        self.config = SimpleNamespace(
+            model_path="/models/test", served_model_name="test", max_seq_len=4096,
+            tool_call_parser="qwen25", reasoning_parser="qwen3" if reasoning else None,
+        )
+        self.text = text
+        self.sent = []
+        self.maintenance_state = "serving"
+
+    def new_user(self):
+        return 42
+
+    async def send_one(self, message):
+        assert isinstance(message, TokenizeMsg)
+        self.sent.append(message)
+
+    async def wait_for_ack(self, uid):
+        # One character per backend reply exercises adapter stream aggregation
+        # as well as the SDK SSE decoder. Usage deltas remain authoritative.
+        for i, char in enumerate(self.text):
+            yield UserReply(uid=uid, incremental_output=char,
+                            finished=i == len(self.text) - 1,
+                            prompt_tokens_delta=5 if i == 0 else 0,
+                            completion_tokens_delta=1)
+
+    async def stream_with_cancellation(self, events, request, uid):
+        async for event in events:
+            yield event
+
+
+def client(state, vendor="openai"):
+    app = FastAPI()
+    for register in (register_openai_routes, register_anthropic_routes, register_responses_routes):
+        register(app, lambda: state, dict)
+    # Recent Anthropic SDKs use httpx2 and reject httpx client instances.
+    # Select the transport matching its public default client class; older SDKs
+    # in the supported range still use httpx. Neither path bypasses SDK validation.
+    http = httpx
+    if vendor == "anthropic":
+        anthropic = pytest.importorskip("anthropic")
+        if not issubclass(anthropic.DefaultAsyncHttpxClient, httpx.AsyncClient):
+            http = pytest.importorskip("httpx2")
+    transport = http.AsyncClient(transport=http.ASGITransport(app=app))
+    cls = openai.AsyncOpenAI if vendor == "openai" else anthropic.AsyncAnthropic
+    return cls(base_url="http://test/v1" if vendor == "openai" else "http://test",
+               api_key="test", http_client=transport, max_retries=0,
+               _strict_response_validation=True)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("tool", [False, True])
+def test_openai_chat_sdk_parses_reasoning_text_tools_and_usage(stream, tool):
+    async def run():
+        state = ScriptedBackend("<think>consider</think>" + (CALL if tool else "Hello"), reasoning=True)
+        async with client(state) as sdk:
+            response = await sdk.chat.completions.create(
+                model="test", messages=[{"role": "user", "content": "Hi"}],
+                tools=[TOOL], stream=stream, stream_options={"include_usage": True},
+            )
+            if stream:
+                chunks = [chunk async for chunk in response]
+                choices = [choice for chunk in chunks for choice in chunk.choices]
+                assert "".join(getattr(c.delta, "reasoning_content", "") or "" for c in choices) == "consider"
+                assert [c.finish_reason for c in choices if c.finish_reason] == ["tool_calls" if tool else "stop"]
+                if tool:
+                    calls = [tc for c in choices for tc in c.delta.tool_calls or []]
+                    assert [tc.function.name for tc in calls if tc.function.name] == ["get_weather"]
+                    assert json.loads("".join(tc.function.arguments or "" for tc in calls)) == {"city": "Paris"}
+                    assert all(tc.index == 0 for tc in calls)
+                else:
+                    assert "".join(c.delta.content or "" for c in choices) == "Hello"
+                usage = chunks[-1].usage
+            else:
+                assert response.choices[0].message.reasoning_content == "consider"
+                assert response.choices[0].finish_reason == ("tool_calls" if tool else "stop")
+                if tool:
+                    call = response.choices[0].message.tool_calls[0]
+                    assert call.id and call.function.name == "get_weather"
+                    assert json.loads(call.function.arguments) == {"city": "Paris"}
+                else:
+                    assert response.choices[0].message.content == "Hello"
+                usage = response.usage
+            assert usage.prompt_tokens == 5
+            assert usage.completion_tokens == len(state.text)
+            assert usage.total_tokens == 5 + len(state.text)
+            assert len(state.sent) == 1
+    asyncio.run(run())
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("tool", [False, True])
+def test_anthropic_sdk_assembles_typed_message(stream, tool):
+    pytest.importorskip("anthropic")
+
+    async def run():
+        state = ScriptedBackend("<think>consider</think>" + (CALL if tool else "Hello"), reasoning=True)
+        async with client(state, "anthropic") as sdk:
+            kwargs = dict(model="test", max_tokens=256,
+                          messages=[{"role": "user", "content": "Hi"}],
+                          tools=[{"name": "get_weather", "input_schema": SCHEMA}])
+            if stream:
+                async with sdk.messages.stream(**kwargs) as response:
+                    events = [event async for event in response]
+                    result = await response.get_final_message()
+                assert events[0].type == "message_start"
+                assert events[-1].type == "message_stop"
+            else:
+                result = await sdk.messages.create(**kwargs)
+            assert result.role == "assistant" and result.model == "test"
+            assert result.stop_reason == ("tool_use" if tool else "end_turn")
+            assert result.content[0].type == "thinking"
+            assert result.content[0].thinking == "consider"
+            if tool:
+                # The established full-response shape includes an empty text
+                # block before tools; preserving that is part of compatibility.
+                assert all(b.text == "" for b in result.content if b.type == "text")
+                block = next(b for b in result.content if b.type == "tool_use")
+                assert block.id
+                assert block.name == "get_weather" and block.input == {"city": "Paris"}
+            else:
+                block = result.content[1]
+                assert block.type == "text" and block.text == "Hello"
+            assert result.usage.input_tokens == 5
+            assert result.usage.output_tokens == len(state.text)
+            assert len(state.sent) == 1
+    asyncio.run(run())
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("tool", [False, True])
+def test_responses_sdk_parses_output_and_events(stream, tool):
+    async def run():
+        state = ScriptedBackend(CALL if tool else "Hello")
+        async with client(state) as sdk:
+            response = await sdk.responses.create(
+                model="test", input="Hi", stream=stream,
+                tools=[{"type": "function", "name": "get_weather", "parameters": SCHEMA}],
+            )
+            if stream:
+                events = [event async for event in response]
+                assert events[0].type == "response.created"
+                assert events[-1].type == "response.completed"
+                assert [e.sequence_number for e in events] == list(range(len(events)))
+                result = events[-1].response
+                deltas = [e for e in events if e.type == ("response.function_call_arguments.delta" if tool else "response.output_text.delta")]
+                text = "".join(e.delta for e in deltas)
+                assert json.loads(text) == {"city": "Paris"} if tool else text == "Hello"
+            else:
+                result = response
+            assert result.status == "completed"
+            if tool:
+                call = next(item for item in result.output if item.type == "function_call")
+                assert call.call_id and call.name == "get_weather"
+                assert json.loads(call.arguments) == {"city": "Paris"}
+            else:
+                assert result.output_text == "Hello"
+            assert result.usage.input_tokens == 5
+            assert result.usage.output_tokens == len(state.text)
+            assert result.usage.total_tokens == 5 + len(state.text)
+            assert len(state.sent) == 1
+    asyncio.run(run())
+
+
+@pytest.mark.parametrize("protocol", ["chat", "responses", "anthropic"])
+def test_sdk_classifies_invalid_request_as_400(protocol):
+    vendor = "anthropic" if protocol == "anthropic" else "openai"
+    sdk_module = pytest.importorskip("anthropic") if vendor == "anthropic" else openai
+
+    async def run():
+        state = ScriptedBackend("unused")
+        error = sdk_module.BadRequestError
+        async with client(state, vendor) as sdk:
+            with pytest.raises(error) as caught:
+                if protocol == "chat":
+                    await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2)
+                elif protocol == "responses":
+                    await sdk.responses.create(model="test", input="Hi", background=True)
+                else:
+                    await sdk.messages.create(model="test", max_tokens=0, messages=[{"role": "user", "content": "Hi"}])
+            assert caught.value.status_code == 400
+            assert "invalid_request_error" in str(caught.value)
+            assert state.sent == []
+    asyncio.run(run())
diff --git a/tests/server/test_reasoning_parser.py b/tests/server/test_reasoning_parser.py
new file mode 100644
index 0000000000..05c207d0e1
--- /dev/null
+++ b/tests/server/test_reasoning_parser.py
@@ -0,0 +1,108 @@
+"""Character-boundary contracts for every reasoning detector and template mode."""
+from types import SimpleNamespace
+
+import pytest
+
+from freetoken.server.generation import GenSpec, _make_reasoning_parser
+from freetoken.server.reasoning_parser import BaseReasoningParser, ReasoningParser
+
+
+class _ReasoningContract:
+    name: str
+    opener: str
+    closer: str
+    content_opener = ""
+
+    def parser(self):
+        return ReasoningParser(self.name, force_reasoning=False)
+
+    def stream(self, text):
+        parser = self.parser()
+        parts = [parser.parse_stream_chunk(char) for char in text]
+        parts.append(parser.flush())
+        assert parser.flush() == ("", "")
+        return tuple("".join(pair[i] for pair in parts) for i in (0, 1))
+
+    def test_complete_reasoning(self):
+        text = self.opener + "consider" + self.closer + self.content_opener + "answer"
+        assert self.parser().parse_non_stream(text) == ("consider", "answer")
+
+    def test_every_character_boundary(self):
+        text = self.opener + "consider" + self.closer + self.content_opener + "answer"
+        assert self.stream(text) == ("consider", "answer")
+
+    def test_truncated_reasoning_flush(self):
+        # Reasoning detectors expose flush(), not recover_truncated_call(). EOS
+        # inside an open reasoning body must retain it in the reasoning channel.
+        assert self.stream(self.opener + "partial thought") == ("partial thought", "")
+
+
+class TestDeepSeekV32ReasoningParser(_ReasoningContract):
+    name = "deepseekv32"
+    opener = "<think>"
+    closer = "</think>"
+
+
+class TestThinkReasoningParser(_ReasoningContract):
+    name = "qwen3"
+    opener = "<think>"
+    closer = "</think>"
+
+
+class TestMiniMaxM3ReasoningParser(_ReasoningContract):
+    name = "minimax_m3"
+    opener = "<mm:think>"
+    closer = "</mm:think>"
+
+
+class TestGemmaThoughtReasoningParser(_ReasoningContract):
+    name = "gemma4"
+    opener = "<|channel>thought\n"
+    closer = "<channel|>"
+
+
+class TestGptOssHarmonyReasoningParser(_ReasoningContract):
+    name = "gpt_oss"
+    opener = "<|channel|>analysis<|message|>"
+    closer = "<|end|>"
+    content_opener = "<|start|>assistant<|channel|>final<|message|>"
+
+
+class TestMuseGlimmerReasoningParser(_ReasoningContract):
+    name = "muse_glimmer"
+    opener = "<|start|>assistant to=self<|message|>"
+    closer = "<|eom|>"
+    content_opener = "<|start|>assistant to=user<|message|>"
+
+
+def test_contract_classes_cover_every_reasoning_detector():
+    covered = {type(cls().parser().detector) for cls in _ReasoningContract.__subclasses__()}
+    assert covered == set(BaseReasoningParser.__subclasses__())
+    assert covered == set(ReasoningParser.ReasoningParserEnum.values())
+
+
+@pytest.mark.parametrize("tools", [None, [{"type": "function", "function": {"name": "read"}}]])
+@pytest.mark.parametrize(("name", "kwargs", "forced"), [
+    ("qwen3", {}, True),
+    ("qwen3", {"enable_thinking": True}, True),
+    ("qwen3", {"enable_thinking": False}, False),
+    ("glm", {}, True),
+    ("glm", {"enable_thinking": True}, True),
+    ("glm", {"enable_thinking": False}, False),
+    ("gemma4", {}, False),
+    ("gemma4", {"thinking_mode": "thinking"}, True),
+    ("gemma4", {"enable_thinking": True}, True),
+    ("gemma4", {"thinking": True}, True),
+    ("gemma4", {"thinking_mode": "chat", "enable_thinking": False}, False),
+    ("minimax_m3", {}, False),
+    ("minimax_m3", {"thinking_mode": "adaptive"}, False),
+    ("minimax_m3", {"thinking_mode": "disabled"}, False),
+    ("minimax_m3", {"thinking_mode": "enabled"}, True),
+])
+def test_generation_force_reasoning_matches_template(name, kwargs, forced, tools):
+    spec = GenSpec(messages=[], sampling_params=None, chat_template_kwargs=kwargs,
+                   template_tools=tools, parser_tools=tools)
+    parser = _make_reasoning_parser(spec, SimpleNamespace(config=SimpleNamespace(reasoning_parser=name)))
+    # Observe routing, not an internal flag: tools must not turn a disabled
+    # template's ordinary response into invisible reasoning.
+    assert parser.parse_non_stream("probe") == (("probe", "") if forced else ("", "probe"))
diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py
index 3568d90505..c4772603b0 100644
--- a/tests/server/test_supervisor.py
+++ b/tests/server/test_supervisor.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pytest
+
 import queue
 
 from queue import Empty as _Empty
@@ -270,3 +272,38 @@ def test_byte_bar_emits_to_installed_sink_then_stops_after_clear():
     bar.update(100)
     bar.close()
     assert seen == []
+
+
+@pytest.mark.parametrize("shutdown", [False, True])
+def test_shutdown_flag_is_read_at_later_worker_death(monkeypatch, shutdown):
+    """Readiness must not snapshot shutdown state: the flag changes on a later poll.
+
+    Keep TP0 alive and kill a different worker, so the first process cannot mask
+    a dead tokenizer. No real processes or sleep timing enter the assertion.
+    """
+    import queue
+    from types import SimpleNamespace
+    from freetoken.server import supervisor
+
+    alive = {"scheduler": True, "tokenizer": True}
+    processes = [SimpleNamespace(name=name, is_alive=lambda name=name: alive[name])
+                 for name in alive]
+    acks = queue.Queue()
+    acks.put("scheduler ready")
+    acks.put("tokenizer ready")
+    handle = BackendHandle(ack_queue=acks, processes=processes, expected_acks=2)
+    stopping = False
+    seen = []
+
+    def next_poll(_):
+        nonlocal stopping
+        stopping = shutdown
+        alive["tokenizer"] = False
+
+    monkeypatch.setattr(supervisor.time, "sleep", next_poll)
+    supervisor.run_backend_supervisor(
+        handle, LoadProgress(), on_ready=lambda: seen.append("ready"),
+        on_failure=lambda message: seen.append(message),
+        is_shutting_down=lambda: stopping,
+    )
+    assert seen == (["ready"] if shutdown else ["ready", "backend worker tokenizer exited"])

From 9267991b6d1d4ffc3b82aca93350e9ebaa28cc43 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:49:51 -0400
Subject: [PATCH 181/249] fix(server): drop /v1/completions from the
 --enable-metrics-report help, which it does not serve

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit c6b7ce7cc462c99c561a08071ec9cd308fdcc26e)
---
 python/freetoken/server/args.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index e25b9e5df0..69d63c6ce7 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -643,7 +643,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         default=ServerArgs.enable_metrics_report,
         help=(
             "Serve a per-request `metrics` object next to usage on /v1/chat/completions, "
-            "/v1/completions, /v1/messages and /v1/responses: ttft_ms, prefill_time_ms and "
+            "/v1/messages and /v1/responses: ttft_ms, prefill_time_ms and "
             "prefill_tokens_per_second (the prefill span measured by the scheduler itself), "
             "decode_time_ms and decode_tokens_per_second, cached_prompt_tokens and "
             "total_time_ms. Streaming responses carry it on the same final chunk as usage, so "

From 560f88e709a6a9e655de936e5a6b3daa6b634b86 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:36:38 -0400
Subject: [PATCH 182/249] fix(server): keep developer ahead of system and
 reject images when merging system messages

render_messages hoisted system ahead of a leading developer message, which the
tokenizer then merged as "S\n\nD"; it now leaves developer requests to the
tokenizer's per-template mapping. Both merge sites joined str() of an image
part list into the prompt, data URL included; they now raise the template's
400 instead.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit e60b2c39d78626e8ecc5c21c7f5b1d77664698ba)
---
 python/freetoken/server/generation.py  | 13 ++++++-----
 python/freetoken/tokenizer/tokenize.py | 13 +++++++++--
 tests/server/test_mm_input.py          | 32 ++++++++++++++++++++++++++
 tests/tokenizer/test_tokenize.py       |  3 +++
 4 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index b86bcab0f0..0ed933feb8 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -25,7 +25,7 @@
 from freetoken.core import SamplingParams
 from freetoken.message import TokenizeMsg
 from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason
-from freetoken.tokenizer.tokenize import resolve_thinking_mode
+from freetoken.tokenizer.tokenize import join_system_contents, resolve_thinking_mode
 from freetoken.tokenizer.inline_system import InlineSystemError
 
 try:
@@ -525,6 +525,11 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
     hoist system messages to the front and merge multiples into one to satisfy
     that constraint."""
     rendered = [_render_message(m) for m in messages]
+    if any(m.get("role") == "developer" for m in rendered):
+        # The tokenizer maps developer per template and then merges in order
+        # (_map_developer_role); hoisting system alone here would put it ahead of
+        # a developer message that came first.
+        return rendered
     system_msgs = [m for m in rendered if m.get("role") == "system"]
     if system_msgs and rendered[0].get("role") != "system":
         non_system = [m for m in rendered if m.get("role") != "system"]
@@ -533,12 +538,8 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
     # system messages that are not at loop.first, i.e. after the first).
     system_msgs = [m for m in rendered if m.get("role") == "system"]
     if len(system_msgs) > 1:
-        sep = chr(10) + chr(10)
-        merged_content = sep.join(
-            str(m.get("content") or "") for m in system_msgs if m.get("content")
-        )
         non_system = [m for m in rendered if m.get("role") != "system"]
-        rendered = [{"role": "system", "content": merged_content}] + non_system
+        rendered = [{"role": "system", "content": join_system_contents(system_msgs)}] + non_system
     return rendered
 
 
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index 9648974dab..31b74bdac2 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -57,6 +57,16 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An
 _EFFORT_PROBE_MESSAGES = [{"role": "user", "content": "ping"}]
 
 
+def join_system_contents(messages: list[dict[str, Any]]) -> str:
+    """The contents of several system/developer messages as one system prompt, in order.
+    Image content (a part list left by text flattening) cannot be joined into text: raise
+    the error templates give for an image in the system turn."""
+    contents = [m.get("content") for m in messages if m.get("content")]
+    if any(not isinstance(c, str) for c in contents):
+        raise ValueError("System message cannot contain images")
+    return "\n\n".join(contents)
+
+
 def _map_developer_role(
     messages: list[dict[str, Any]], chat_template: str | None
 ) -> list[dict[str, Any]]:
@@ -74,8 +84,7 @@ def _map_developer_role(
     if len(system) == 1 and mapped[0] is system[0]:
         return mapped
     rest = [m for m in mapped if m.get("role") != "system"]
-    content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content"))
-    return [{"role": "system", "content": content}] + rest
+    return [{"role": "system", "content": join_system_contents(system)}] + rest
 
 
 class TokenizeManager:
diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py
index 54f7685afa..1cca773361 100644
--- a/tests/server/test_mm_input.py
+++ b/tests/server/test_mm_input.py
@@ -183,3 +183,35 @@ def test_render_messages_no_system_unchanged():
         {"role": "user", "content": "Bye"},
     ])
     assert [m["role"] for m in msgs] == ["user", "assistant", "user"]
+
+
+def test_render_messages_keeps_a_leading_developer_message_ahead_of_system():
+    """[developer, system, user]: the tokenizer maps developer -> system and merges in order, so
+    render_messages must not hoist system ahead of it first. A template that knows the
+    developer role gets both messages as sent."""
+    from freetoken.tokenizer.tokenize import _map_developer_role
+
+    msgs = [
+        {"role": "developer", "content": "D"},
+        {"role": "system", "content": "S"},
+        {"role": "user", "content": "hi"},
+    ]
+    rendered = render_messages(msgs)
+    assert [m["role"] for m in rendered] == ["developer", "system", "user"]
+    mapped = _map_developer_role(rendered, "{% if message.role == 'system' %}")
+    assert [m["role"] for m in mapped] == ["system", "user"] and mapped[0]["content"] == "D\n\nS"
+    assert _map_developer_role(rendered, "{% if message.role == 'developer' %}") is rendered
+
+
+def test_render_messages_rejects_an_image_in_a_merged_system_message():
+    """Merging system messages joins text; an image part must not be pasted into the prompt as
+    the repr of its part list (with the whole data URL)."""
+    with pytest.raises(ValueError, match="System message cannot contain images"):
+        render_messages([
+            {"role": "system", "content": "A"},
+            {"role": "system", "content": [
+                {"type": "text", "text": "B"},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{PNG}"}},
+            ]},
+            {"role": "user", "content": "hi"},
+        ])
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 390fcd4cce..9e441886b2 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -845,3 +845,6 @@ def test_mapped_developer_messages_are_merged_into_one_leading_system_message():
     assert mapped[0]["content"] == "be brief\n\nanswer in German"
     late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}]
     assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"]
+    image = [{"type": "text", "text": "B"}, {"type": "image", "freetoken_ref": {"kind": "url", "data": "x"}}]
+    with pytest.raises(ValueError, match="System message cannot contain images"):
+        _map_developer_role([{"role": "developer", "content": "A"}, {"role": "system", "content": image}], template)

From ba1bb53077c61eee3599c6cf0c3ad82eb32a0b4c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:24:07 -0400
Subject: [PATCH 183/249] test(server): ask for n=0 in the invalid-request sdk
 test, since n=2 is served here

Assisted-by: Claude Opus 5.5
---
 tests/server/test_protocol_compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py
index 8dc22908d6..d91582e423 100644
--- a/tests/server/test_protocol_compat.py
+++ b/tests/server/test_protocol_compat.py
@@ -202,7 +202,7 @@ async def run():
         async with client(state, vendor) as sdk:
             with pytest.raises(error) as caught:
                 if protocol == "chat":
-                    await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2)
+                    await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=0)
                 elif protocol == "responses":
                     await sdk.responses.create(model="test", input="Hi", background=True)
                 else:

From 943d2d6b8c74ab9fd04c1e84eef30682eedc2e1d Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:40:55 -0400
Subject: [PATCH 184/249] fix(moe): keep unpinned layers off the small-prefill
 path and take it in every layer

An unpinned layer's copy_missing is the whole-layer pageable copy, which cannot serve
the LRU slot remap, so it keeps streaming. The per-layer fallback when a layer routes
to more than half the cache is gone: the engine's num_experts slot floor already keeps
every chunk resident, and a uniform path cannot leave a later layer in an overlap
buffer that was never begun.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 37c2479212b8b7028827218b23d1cee68e2c274c)
---
 python/freetoken/layers/moe.py            | 44 +++++++++---------
 tests/moe/test_small_prefill_on_demand.py | 54 +++++++++++++++++++++++
 2 files changed, 78 insertions(+), 20 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 19e6d8e77c..adce6a49d3 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -419,22 +419,26 @@ def _prefill_routed(
         pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
-        if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"):
+        if (
+            0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS
+            and cache.decode_target in ("gpu", "hybrid")
+            and not cache.is_unpinned_layer(self.layer_id)
+        ):
             # A short extension (an agent turn over a cached prefix) touches a fraction of each
-            # layer's experts; fetching those beats streaming all of them (~68 GB here). The same
-            # token count holds in every layer, so the overlap double buffer is never begun.
-            slots = self._ensure_unique(cache, topk_ids)
-            if slots is not None:
-                return self._expert_gemm(
-                    cache,
-                    hidden_states,
-                    topk_weights,
-                    slots,
-                    views=cache.bank_views(),
-                    n=None,
-                    alphas=cache.alphas_for_slots(self.layer_id),
-                    is_prefill=False,
-                )
+            # layer's experts; fetching those beats streaming all of them (~68 GB here). An
+            # unpinned layer's copy_missing presumes position == expert id, so it streams below;
+            # unpinned layers only exist with the prefill overlap off, so no layer can enter
+            # the overlap double buffer after an earlier layer skipped begin_prefill.
+            return self._expert_gemm(
+                cache,
+                hidden_states,
+                topk_weights,
+                self._ensure_unique(cache, topk_ids),
+                views=cache.bank_views(),
+                n=None,
+                alphas=cache.alphas_for_slots(self.layer_id),
+                is_prefill=False,
+            )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
@@ -464,9 +468,8 @@ def _prefill_routed(
         cache.release_prefill_layer(self.layer_id)
         return out
 
-    def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None:
-        """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or
-        None (take the whole-layer path) when they cannot all be resident at once.
+    def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor:
+        """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots.
 
         The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a
         prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it
@@ -474,8 +477,9 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc
         decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan."""
         uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True)
         n = uniq.numel()
-        if n > cache.cache_size // 2:  # leave room so a chunk never evicts the one before it
-            return None
+        # Each ensure stamps a newer LRU step, so a later chunk evicts the earlier chunks' slots
+        # only once every other slot is gone; the engine's num_experts slot floor rules that out.
+        assert n <= cache.cache_size, f"{n} routed experts exceed the {cache.cache_size}-slot cache"
         uniq = uniq.to(torch.int32)
         # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly
         pin_since = cache.step.clone() if cache.cache_policy_id else None
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index 70cec00019..717f835d9c 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -56,6 +56,60 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill):
         assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts
 
 
+def _record_paths(monkeypatch, layer, cache):
+    """Stub the cache movement and the GEMM; log each call and what the GEMM was handed."""
+    calls, got = [], {}
+
+    def ensure(layer_id, ids):
+        calls.append(("ensure", ids.tolist()))
+        ids.add_(100)  # the LRU kernel rewrites its query to slot ids in place
+
+    monkeypatch.setattr(cache, "ensure_experts", ensure)
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append(("materialize",)))
+    monkeypatch.setattr(cache, "copy_missing", lambda: calls.append(("copy",)))
+
+    def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill):
+        got.update(ids=ids.tolist(), n=n, is_prefill=is_prefill)
+        return hs
+
+    monkeypatch.setattr(layer, "_expert_gemm", fake_gemm)
+    return calls, got
+
+
+def test_small_prefill_leaves_an_unpinned_layer_on_the_whole_layer_path(monkeypatch):
+    # an unpinned (locked) layer's copy_missing is the whole-layer pageable copy, which presumes
+    # position == expert id and cannot serve an LRU slot remap
+    layer, cache = _layer_and_cache(torch.device("cpu"))
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8)
+    cache._unpinned_layers = frozenset({layer.layer_id})
+    calls, got = _record_paths(monkeypatch, layer, cache)
+    layer._prefill_routed(torch.randn(3, 16), torch.full((3, 2), 0.5), torch.zeros(3, 2, dtype=torch.int32))
+    assert calls == [("materialize",), ("copy",)] and got["n"] == layer.num_experts
+
+
+def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_to(monkeypatch):
+    # 6 distinct experts in an 8-slot cache: every layer of a short prefill still takes the
+    # on-demand path, so no layer can fall back into an overlap buffer that was never begun
+    layer, cache = _layer_and_cache(torch.device("cpu"))
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8)
+    monkeypatch.setattr(moe_mod, "_ENSURE_CHUNK_IDS", 4)
+    calls, got = _record_paths(monkeypatch, layer, cache)
+    ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32)
+    layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone())
+    assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)]
+    assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False}
+
+
+def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch):
+    layer, cache = _layer_and_cache(torch.device("cpu"))
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8)
+    _record_paths(monkeypatch, layer, cache)
+    cache.cache_size = 4  # below the engine's num_experts floor
+    ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32)
+    with pytest.raises(AssertionError, match="exceed"):
+        layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 def test_small_prefill_matches_the_whole_layer_path(monkeypatch):
     from freetoken.core import get_global_ctx  # noqa: F401  (layer reads the ctx only in forward())

From 547c5e994a80eac048f99254aac1739247bcddc6 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:24:22 -0400
Subject: [PATCH 185/249] test(moe): let the small-prefill ensure fake take the
 eviction policy kwargs

Assisted-by: Claude Opus 5.5
---
 tests/moe/test_small_prefill_on_demand.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index 717f835d9c..8b7b9d802c 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -60,7 +60,7 @@ def _record_paths(monkeypatch, layer, cache):
     """Stub the cache movement and the GEMM; log each call and what the GEMM was handed."""
     calls, got = [], {}
 
-    def ensure(layer_id, ids):
+    def ensure(layer_id, ids, **_):
         calls.append(("ensure", ids.tolist()))
         ids.add_(100)  # the LRU kernel rewrites its query to slot ids in place
 

From 3608bd81f020e01c4361435e8e0046d63e8923f0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 4 Sep 2026 17:31:34 -0400
Subject: [PATCH 186/249] test(qwen4_exp): compare chunked prefill to one-shot
 at bf16 tolerance

`torch.equal` fails for all three cut points on torch 2.11.0+cu130 /
flashinfer 0.6.18 / triton 3.6.0 (sm_89): cuBLAS picks a different kernel
for M=5000 and M=5000-cut, so the indexer's q and k round differently, a
few rows flip one block at the top-k margin, and the layer output differs
by up to 2.6e-3 (a few bf16 ulps). Both paths are deterministic run to run.
Use the tolerance of the neighbouring oracle tests.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
(cherry picked from commit 3849e3caeebe4f1376708fe677cc0ecb106fb78c)

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit e700e8fb7339eb958507401c39d81b2623f56ca7)
---
 tests/models/qwen4_exp/test_qsa_backend.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py
index 9302fad172..3827b62416 100644
--- a/tests/models/qwen4_exp/test_qsa_backend.py
+++ b/tests/models/qwen4_exp/test_qsa_backend.py
@@ -4,7 +4,8 @@
     tokens every complete block is selected, so QSA IS dense attention: the selection must be
     exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32)
     and a flashinfer dense run over the same pool;
-(b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress);
+(b) chunked prefill at unaligned cut points matches one-shot prefill to bf16 tolerance (the
+    dual-source compress);
 (c) a captured decode replay equals the eager decode step;
 (d) an fp8 KV pool (``--kv-cache-dtype fp8``) keeps block selection bit-identical to the
     16-bit run -- only the selected K/V rows are read back as e4m3 codes -- and the layer
@@ -136,7 +137,11 @@ def test_flashinfer_dense_matches_the_sparse_path():
 @requires_cuda
 @pytest.mark.parametrize("cut", [1001, 4096, 4097], ids=["unaligned", "page-boundary", "boundary+1"])
 def test_chunked_prefill_matches_one_shot(cut: int):
-    """Cut points that are not multiples of index_ratio exercise the dual-source compress."""
+    """Cut points that are not multiples of index_ratio exercise the dual-source compress.
+
+    Not bit for bit: the projections see a different M (5000 vs 5000 - cut rows) and cuBLAS
+    may pick a different kernel per shape, so the indexer's q and k round differently and a
+    few rows flip one block at the top-k margin. Same tolerance as the oracle tests above."""
     config = parsed_config()
     fixture = Fixture(config, num_pages=512)
     attn = fixture.layer(QSA_LAYER)
@@ -148,7 +153,7 @@ def test_chunked_prefill_matches_one_shot(cut: int):
     attn.forward(x[:cut], fixture.batch([head], "prefill"))
     tail = fixture.req(1, cut, length)
     got = attn.forward(x[cut:], fixture.batch([tail], "prefill"))
-    assert torch.equal(got, one_shot[cut:])
+    torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2)
 
 
 @requires_cuda

From afbfea8e02a8597acacce0dc258b8f8218656d20 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 09:23:20 -0400
Subject: [PATCH 187/249] test(qwen4_exp): compare the chunked prefill's
 compressed slab rows to the one-shot run

At bf16 output tolerance a broken dual-source compress only swaps a block in a few rows
and passes; the slab keys are unit-scale, so a wrong straddling-group key stands out.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit cedf179739d2f17e8dc8ee27152f35c5584c24c4)
---
 tests/models/qwen4_exp/test_qsa_backend.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py
index 3827b62416..8e836f5632 100644
--- a/tests/models/qwen4_exp/test_qsa_backend.py
+++ b/tests/models/qwen4_exp/test_qsa_backend.py
@@ -155,6 +155,18 @@ def test_chunked_prefill_matches_one_shot(cut: int):
     got = attn.forward(x[cut:], fixture.batch([tail], "prefill"))
     torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2)
 
+    # A wrong key for the group straddling the cut only swaps one block in a few rows, which
+    # the output tolerance hides; the unit-scale slab rows show it as an O(1) error.
+    ratio = config.qwen4_args.index_ratio
+    slab = fixture.pool.cmp_k_cache(0)
+    group_starts = slice(0, length - length % ratio, ratio)
+    one_shot_rows, chunked_rows = (
+        (fixture.page_table[table_idx, group_starts] // ratio).long() for table_idx in (0, 1)
+    )
+    torch.testing.assert_close(
+        slab[chunked_rows].float(), slab[one_shot_rows].float(), rtol=2e-2, atol=2e-2
+    )
+
 
 @requires_cuda
 def test_decode_graph_replay_matches_eager():

From 7f59703d93c21c45f89850239a2916ce0d97b5b7 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:20:22 -0400
Subject: [PATCH 188/249] feat(engine): keep the token embedding in pinned host
 ram (--embed-weights host)

The embedding is a row lookup, so a single-GPU engine can leave the table in page-locked
host RAM and gather each step's rows over UVA with the PLE row kernel (5 KiB per decode
token, graph-safe): 1.2 GiB of VRAM back for Qwen3.8-Flash-Next.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 9f2b7a2e33424e2c0e06b240f1fb7a408c850212)
---
 python/freetoken/engine/config.py     |  2 +
 python/freetoken/engine/engine.py     | 35 +++++++++++++++++
 python/freetoken/kernel/triton/ple.py |  8 +++-
 python/freetoken/layers/embedding.py  | 26 ++++++++++++-
 python/freetoken/server/args.py       | 10 +++++
 tests/kernels/test_host_embedding.py  | 55 +++++++++++++++++++++++++++
 6 files changed, 133 insertions(+), 3 deletions(-)
 create mode 100644 tests/kernels/test_host_embedding.py

diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index d52f3d808b..2669e63ed7 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -38,6 +38,8 @@ class EngineConfig:
     # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each
     # prompt's last window; "exact" runs them on every token.
     swa_decoder_replay: str = "bounded"
+    # --embed-weights: "host" keeps the token embedding in pinned host RAM, gathered over PCIe.
+    embed_weights: str = "gpu"
     # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered
     # experts in parallel but falls back to serial when free RAM can't cover the banks + the
     # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 1a92c53a40..c2428d43d8 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -108,6 +108,39 @@ def _sgl_flash_attn_available() -> bool:
     return True
 
 
+def _place_embeddings_on_host(model) -> None:
+    """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to."""
+    from freetoken.layers.base import BaseOP
+    from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding
+
+    found: list[VocabParallelEmbedding] = []
+    tied: set[int] = set()
+    stack, seen = [model], set()
+    while stack:
+        op = stack.pop()
+        if id(op) in seen:
+            continue
+        seen.add(id(op))
+        if isinstance(op, ParallelLMHead):
+            if op.tied_embedding is not None:
+                tied.add(id(op.tied_embedding))
+            continue
+        if isinstance(op, VocabParallelEmbedding):
+            found.append(op)
+            continue
+        for value in vars(op).values():
+            items = value if isinstance(value, (list, tuple)) else (value,)
+            stack.extend(v for v in items if isinstance(v, BaseOP))
+    moved = [e for e in found if id(e) not in tied]
+    if not moved:
+        raise ValueError("--embed-weights host: the model has no untied input embedding to move")
+    for emb in moved:
+        emb.place_on_host()
+    torch.cuda.empty_cache()
+    size = sum(e.weight.numel() * e.weight.element_size() for e in moved)
+    logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe")
+
+
 def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int:
     """Bytes available to the KV pool at startup: ratio-scaled pre-load free memory minus
     what the resident model consumed. Kept as a pure function so the composition with the
@@ -491,6 +524,8 @@ def __init__(self, config: EngineConfig):
                 )
             # before the residency snapshot, so streamed blocks are not charged as resident weights
             self.model.place_encoder_weights(config.mm.encoder_weights)
+        if config.embed_weights == "host":
+            _place_embeddings_on_host(self.model)
         post_weights_free = self._sync_get_memory()[0]
         self._weights_bytes = self._baseline_free - post_weights_free
         # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are
diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py
index 2314e94e72..cded101a2e 100644
--- a/python/freetoken/kernel/triton/ple.py
+++ b/python/freetoken/kernel/triton/ple.py
@@ -21,6 +21,7 @@
 
 # Latency-bound over PCIe, so keep the block small and let many of them be in flight.
 _NUM_WARPS = 1
+_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16}
 
 
 @triton.jit
@@ -33,6 +34,7 @@ def _ple_gather_kernel(
     EMB_DIM: tl.constexpr,
     IS_FP8: tl.constexpr,
     BLOCK_D: tl.constexpr,
+    TABLE_DTYPE: tl.constexpr,
 ):
     row = tl.program_id(0)
     idx = tl.load(ids_ptr + row).to(tl.int64)
@@ -50,7 +52,7 @@ def _ple_gather_kernel(
             base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8))
             values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0))
     else:
-        base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16))
+        base = table_ptr.to(tl.int64).to(tl.pointer_type(TABLE_DTYPE))
         values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32)
     values = tl.where(in_range, values * scale, 0.0)
     tl.store(
@@ -68,11 +70,12 @@ def ple_gather_rows(
     out: torch.Tensor,
     scale: float = 1.0,
     is_fp8: bool = True,
+    table_dtype: torch.dtype = torch.bfloat16,
 ) -> torch.Tensor:
     """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``.
 
     ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]``
-    bf16 on the same device. ``table_ptr`` is the address the GPU must dereference
+    on the same device. A non-fp8 table holds ``table_dtype`` elements. ``table_ptr`` is the address the GPU must dereference
     (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``.
     """
     n = row_ids.numel()
@@ -87,6 +90,7 @@ def ple_gather_rows(
             EMB_DIM=embed_dim,
             IS_FP8=is_fp8,
             BLOCK_D=triton.next_power_of_2(embed_dim),
+            TABLE_DTYPE=_TL_TABLE_DTYPES[table_dtype],
             num_warps=_NUM_WARPS,
         )
     return out
diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index 20a6f1241d..528a3404ed 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -47,10 +47,34 @@ def __init__(
         self._embed_scale = embed_scale
         self._embed_scale_t: torch.Tensor | None = None
         self._comm = DistributedCommunicator()
+        self._host_ptr: int | None = None
+
+    def place_on_host(self) -> None:
+        """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA)."""
+        from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr
+        from freetoken.layers.quantization.embedding.unquantized import UnquantizedEmbeddingMethod
+
+        if self.tp_size > 1:
+            raise ValueError("a host-resident embedding serves a single GPU only")
+        if not isinstance(self.quant_method, UnquantizedEmbeddingMethod):
+            raise ValueError("a host-resident embedding needs a dense (unquantized) table")
+        host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype)
+        host.copy_(self.weight)
+        self.weight = host
+        self._host_ptr = device_ptr(host)
 
     @nvtx_annotate("Embedding")
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = self.quant_method.embed(self, x)
+        if self._host_ptr is not None:
+            from freetoken.kernel.triton.ple import ple_gather_rows
+
+            rows, dim = self.weight.shape
+            y = torch.empty(x.numel(), dim, dtype=self.weight.dtype, device=x.device)
+            ple_gather_rows(self._host_ptr, rows, dim, x.reshape(-1), y, is_fp8=False,
+                            table_dtype=self.weight.dtype)
+            y = y.view(*x.shape, dim)
+        else:
+            y = self.quant_method.embed(self, x)
         if self.tp_size > 1:
             y = self._comm.all_reduce(y)
         if self._embed_scale is not None:
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 69d63c6ce7..c493d96102 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -780,6 +780,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--embed-weights",
+        default=ServerArgs.embed_weights,
+        choices=["gpu", "host"],
+        help=(
+            "Where the token embedding table lives. 'host' keeps it in pinned host RAM and gathers "
+            "the rows each step needs over PCIe, freeing its VRAM (single GPU only)."
+        ),
+    )
+
     parser.add_argument(
         "--nvfp4-backend",
         action=_DeprecatedAlias,
diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py
new file mode 100644
index 0000000000..76c8c612c3
--- /dev/null
+++ b/tests/kernels/test_host_embedding.py
@@ -0,0 +1,55 @@
+"""--embed-weights host: the UVA row gather from pinned host RAM must equal the GPU table lookup."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+@pytest.fixture(autouse=True)
+def _tp():
+    from freetoken.distributed import set_tp_info, try_get_tp_info
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+
+
+def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560):
+    from freetoken.layers.embedding import VocabParallelEmbedding
+
+    torch.manual_seed(0)
+    emb = VocabParallelEmbedding(rows, dim)
+    emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype)
+    return emb
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_host_rows_equal_the_gpu_lookup(dtype):
+    emb = _embedding(dtype)
+    ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32)
+    want = emb.forward(ids)
+    emb.place_on_host()
+    assert not emb.weight.is_cuda and emb.weight.is_pinned()
+    got = emb.forward(ids)
+    assert got.is_cuda and got.dtype is dtype
+    assert torch.equal(got, want)
+
+
+def test_host_lookup_replays_in_a_cuda_graph_with_new_ids():
+    emb = _embedding(torch.bfloat16)
+    table = emb.weight.clone()
+    emb.place_on_host()
+    ids = torch.zeros(8, device="cuda", dtype=torch.int64)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        emb.forward(ids)  # compile outside the capture
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = emb.forward(ids)
+    for seed in range(3):
+        ids.copy_(torch.randint(0, 1000, (8,), generator=torch.Generator().manual_seed(seed)))
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(out, table[ids])

From 600f4261711a164f5eba12f6133db5d34ff23018 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:51:34 -0400
Subject: [PATCH 189/249] fix(layers): gather fp32 host embeddings and refuse
 other dtypes at placement

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit ccfb2b8148c8c108897f37b7f337e99cbb5a0134)
---
 python/freetoken/kernel/triton/ple.py |  2 +-
 python/freetoken/layers/embedding.py  |  6 +++---
 tests/kernels/test_host_embedding.py  | 17 +++++++++++++----
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py
index cded101a2e..8fa25acdfd 100644
--- a/python/freetoken/kernel/triton/ple.py
+++ b/python/freetoken/kernel/triton/ple.py
@@ -21,7 +21,7 @@
 
 # Latency-bound over PCIe, so keep the block small and let many of them be in flight.
 _NUM_WARPS = 1
-_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16}
+_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32}
 
 
 @triton.jit
diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index 528a3404ed..eaa3608826 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -52,12 +52,12 @@ def __init__(
     def place_on_host(self) -> None:
         """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA)."""
         from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr
-        from freetoken.layers.quantization.embedding.unquantized import UnquantizedEmbeddingMethod
+        from freetoken.kernel.triton.ple import _TL_TABLE_DTYPES
 
         if self.tp_size > 1:
             raise ValueError("a host-resident embedding serves a single GPU only")
-        if not isinstance(self.quant_method, UnquantizedEmbeddingMethod):
-            raise ValueError("a host-resident embedding needs a dense (unquantized) table")
+        if self.weight.dtype not in _TL_TABLE_DTYPES:
+            raise ValueError(f"a host-resident embedding cannot gather a {self.weight.dtype} table")
         host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype)
         host.copy_(self.weight)
         self.weight = host
diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py
index 76c8c612c3..866687a9f1 100644
--- a/tests/kernels/test_host_embedding.py
+++ b/tests/kernels/test_host_embedding.py
@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 
 
 @pytest.fixture(autouse=True)
@@ -14,16 +14,17 @@ def _tp():
         set_tp_info(rank=0, size=1)
 
 
-def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560):
+def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560, device: str = "cuda"):
     from freetoken.layers.embedding import VocabParallelEmbedding
 
     torch.manual_seed(0)
     emb = VocabParallelEmbedding(rows, dim)
-    emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype)
+    emb.weight = torch.randn(rows, dim, device=device, dtype=dtype)
     return emb
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@needs_cuda
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 def test_host_rows_equal_the_gpu_lookup(dtype):
     emb = _embedding(dtype)
     ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32)
@@ -35,6 +36,14 @@ def test_host_rows_equal_the_gpu_lookup(dtype):
     assert torch.equal(got, want)
 
 
+def test_a_table_the_gather_cannot_read_is_refused_before_pinning():
+    emb = _embedding(torch.float64, rows=4, dim=8, device="cpu")
+    with pytest.raises(ValueError, match="cannot gather a torch.float64 table"):
+        emb.place_on_host()
+    assert emb._host_ptr is None
+
+
+@needs_cuda
 def test_host_lookup_replays_in_a_cuda_graph_with_new_ids():
     emb = _embedding(torch.bfloat16)
     table = emb.weight.clone()

From b1f443358b342071e2bf6687f81db4dd6b0507fc Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:54:43 -0400
Subject: [PATCH 190/249] fix(kernel): compute the host gather output offset in
 int64

No regression test: overflowing it needs an output of 2^31+ elements (4 GiB+ in bf16).

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit cc5af7adc5d99c17c0edc3f7947040ac35a48445)
---
 python/freetoken/kernel/triton/ple.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py
index 8fa25acdfd..466fce11a0 100644
--- a/python/freetoken/kernel/triton/ple.py
+++ b/python/freetoken/kernel/triton/ple.py
@@ -36,7 +36,8 @@ def _ple_gather_kernel(
     BLOCK_D: tl.constexpr,
     TABLE_DTYPE: tl.constexpr,
 ):
-    row = tl.program_id(0)
+    # the output offset row * EMB_DIM can pass 2^31 for a long prefill at a wide hidden size
+    row = tl.program_id(0).to(tl.int64)
     idx = tl.load(ids_ptr + row).to(tl.int64)
     in_range = (idx >= 0) & (idx < num_rows)
     idx = tl.where(in_range, idx, 0)

From 41ec434fbc9a5ede25787e3bcee61679b35ff853 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:56:09 -0400
Subject: [PATCH 191/249] fix(engine): charge the host embedding's pinned bytes
 to the pin budget

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit a8b2f0a6924c9362b014c200dd702115b0c33d89)
---
 python/freetoken/engine/engine.py | 16 +++++++++-------
 tests/engine/test_cache_budget.py | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index c2428d43d8..e2459cb97c 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -108,8 +108,10 @@ def _sgl_flash_attn_available() -> bool:
     return True
 
 
-def _place_embeddings_on_host(model) -> None:
-    """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to."""
+def _place_embeddings_on_host(model) -> int:
+    """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to.
+
+    Returns the pinned bytes."""
     from freetoken.layers.base import BaseOP
     from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding
 
@@ -139,6 +141,7 @@ def _place_embeddings_on_host(model) -> None:
     torch.cuda.empty_cache()
     size = sum(e.weight.numel() * e.weight.element_size() for e in moved)
     logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe")
+    return size
 
 
 def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int:
@@ -524,8 +527,7 @@ def __init__(self, config: EngineConfig):
                 )
             # before the residency snapshot, so streamed blocks are not charged as resident weights
             self.model.place_encoder_weights(config.mm.encoder_weights)
-        if config.embed_weights == "host":
-            _place_embeddings_on_host(self.model)
+        embed_host_bytes = _place_embeddings_on_host(self.model) if config.embed_weights == "host" else 0
         post_weights_free = self._sync_get_memory()[0]
         self._weights_bytes = self._baseline_free - post_weights_free
         # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are
@@ -539,11 +541,11 @@ def __init__(self, config: EngineConfig):
         self.cpu_moe_executor = None
         # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a
         # load failure is not masked, before the MoE offload cache so the bank residency
-        # planning sees the pin quota the table already spent.
-        self._host_tables_bytes = 0
+        # planning sees the pin quota the table (and a host embedding) already spent.
+        self._host_tables_bytes = embed_host_bytes
         if hasattr(self.model, "load_host_tables"):
             with _weight_load_context():
-                self._host_tables_bytes = int(self.model.load_host_tables(config) or 0)
+                self._host_tables_bytes += int(self.model.load_host_tables(config) or 0)
         if is_offload_moe_strategy(config.moe_strategy):
             self._init_offload_moe_cache(config)
         if hasattr(self.model, "prepare_for_runtime"):
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 6f7184b327..29603f2a83 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -671,3 +671,19 @@ def test_moe_pool_flags_and_decode_batch_clamp():
     )
     _clamp_decode_batch_to_pools(config, [12, 16, 9])
     assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2])
+
+
+def test_host_embedding_placement_reports_the_bytes_it_pinned(monkeypatch):
+    # the engine charges these bytes to the pin budget before it plans the expert banks
+    from freetoken.distributed import set_tp_info, try_get_tp_info
+    from freetoken.engine.engine import _place_embeddings_on_host
+    from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    monkeypatch.setattr(VocabParallelEmbedding, "place_on_host", lambda self: None)
+    embed, tied = VocabParallelEmbedding(100, 64), VocabParallelEmbedding(50, 64)
+    embed.weight = torch.empty(100, 64, dtype=torch.bfloat16)
+    head = ParallelLMHead(50, 64, tie_word_embeddings=True, tied_embedding=tied)
+    model = SimpleNamespace(embed_tokens=embed, per_layer_embed=tied, lm_head=head)
+    assert _place_embeddings_on_host(model) == 100 * 64 * 2

From 654176997ed6e16d78b14510f1e47edf1eb7f5a9 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:57:23 -0400
Subject: [PATCH 192/249] fix(engine): refuse --embed-weights host for a tied
 embedding or tp > 1 at config time

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 579e03a7aae077ac209cbf22ed4de6cda5bb390f)
---
 python/freetoken/engine/config.py  | 10 ++++++++
 tests/engine/test_embed_weights.py | 39 ++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 tests/engine/test_embed_weights.py

diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 2669e63ed7..b6d8e40ea0 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -169,8 +169,18 @@ def model_config(self) -> ModelConfig:
         quant = checkpoint_quant_config(self.model_path, hf_config, spec)
         set_quant_config(quant)
         model_config = _load_attr(spec.module, spec.parse_config)(hf_config)
+        self._check_embed_weights(model_config)
         return replace(model_config, quant=quant)
 
+    def _check_embed_weights(self, model_config) -> None:
+        # the parsed flag, not the hf key: gemma4 ties by default and gguf infers it from the tensors
+        if self.embed_weights != "host":
+            return
+        if model_config.tie_word_embeddings:
+            raise ValueError("--embed-weights host needs an untied embedding (a tied lm_head reads the table on the GPU)")
+        if self.tp_info.size > 1:
+            raise ValueError("--embed-weights host serves a single GPU only (--tensor-parallel-size 1)")
+
     @property
     def max_seq_len(self) -> int:
         if self.max_seq_len_override is not None:
diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py
new file mode 100644
index 0000000000..b8cb0fdc26
--- /dev/null
+++ b/tests/engine/test_embed_weights.py
@@ -0,0 +1,39 @@
+"""--embed-weights host: a setup the host gather cannot serve is refused while the config is still a dataclass, not after the weight load."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.distributed import DistributedInfo
+from freetoken.engine.config import EngineConfig
+
+
+def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host"):
+    import freetoken.engine.config as engine_config
+
+    hf = SimpleNamespace(
+        architectures=["Qwen3ForCausalLM"], model_type="qwen3", num_hidden_layers=2, num_attention_heads=4,
+        num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6,
+        max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie,
+    )
+    monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf)
+    monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None)
+    return EngineConfig(
+        model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights
+    )
+
+
+def test_a_tied_embedding_is_refused(monkeypatch):
+    with pytest.raises(ValueError, match="needs an untied embedding"):
+        _config(monkeypatch, tie=True).model_config
+
+
+def test_tensor_parallel_is_refused(monkeypatch):
+    with pytest.raises(ValueError, match="single GPU only"):
+        _config(monkeypatch, tie=False, tp=2).model_config
+
+
+def test_an_untied_single_gpu_model_and_the_gpu_default_pass(monkeypatch):
+    assert not _config(monkeypatch, tie=False).model_config.tie_word_embeddings
+    assert _config(monkeypatch, tie=True, tp=2, embed_weights="gpu").model_config.tie_word_embeddings

From 5cc4fa8e49240af1aadb9c02113aba75d1804ed3 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:26:39 -0400
Subject: [PATCH 193/249] test(engine): fake cached_load_hf_config with the
 overrides argument #573 passes

Assisted-by: Claude Opus 5.5
---
 tests/engine/test_embed_weights.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py
index b8cb0fdc26..d3aa102f94 100644
--- a/tests/engine/test_embed_weights.py
+++ b/tests/engine/test_embed_weights.py
@@ -17,7 +17,7 @@ def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host")
         num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6,
         max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie,
     )
-    monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf)
+    monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, *_: hf)
     monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None)
     return EngineConfig(
         model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights

From 541d528fd7cf94ac67c83dafb36c5ab464deb625 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 1 Oct 2026 23:27:21 -0400
Subject: [PATCH 194/249] test(ple): drive overlapped graph fills against an
 in-order fake stream so a post-launch fence fails as a hang

Upstream #588's _graph_consumed fence is left out of the merge: our graph fill runs on the
filler thread before its own launch passes its WAIT, so that fence deadlocks the second step.

Assisted-by: Claude Opus 5.5
---
 .../qwen4_exp/test_ple_disk_no_graph_fence.py | 136 ++++++++++++++++++
 1 file changed, 136 insertions(+)
 create mode 100644 tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py

diff --git a/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py
new file mode 100644
index 0000000000..cb9bc32356
--- /dev/null
+++ b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py
@@ -0,0 +1,136 @@
+"""The wait-sync graph fill must never wait on an event recorded after its own launch.
+
+Upstream #588 fences ``_graph_pinned`` with ``_graph_consumed``, recorded once the launch returns and
+synchronized by the next graph fill. Upstream runs the fill after the launch returns; here (fork
+#65/#67) the fill runs on the filler thread BEFORE its own graph can pass its WAIT, so that fence
+makes fill k wait on graph k, which waits on fill k: decode hangs on the second overlapped step.
+The previous graph's read of ``_graph_pinned`` is already ordered before the next fill by the
+per-fill readback event. This test drives the overlapped decode loop against an asynchronous,
+in-order fake stream (launch returns at once; events record a stream position) and fails if any
+step stalls, including when the fence is reintroduced.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from types import SimpleNamespace
+
+import torch
+
+import freetoken.models.qwen4_exp.ple_disk as ple_disk
+from freetoken.models.qwen4_exp.ple_disk import DiskRowTable
+
+EOS = 9
+
+
+class _Stream:
+    """One in-order stream: a graph item WAITs on the flag, then RESETs it; launch is async."""
+
+    def __init__(self, flag: torch.Tensor):
+        self.flag = flag
+        self.cv = threading.Condition()
+        self.queue: list[str] = []
+        self.enqueued = self.done = 0
+        self.stop = False
+        threading.Thread(target=self._run, daemon=True).start()
+
+    def enqueue(self, kind: str) -> None:
+        with self.cv:
+            self.queue.append(kind)
+            self.enqueued += 1
+            self.cv.notify_all()
+
+    def _run(self) -> None:
+        while True:
+            with self.cv:
+                self.cv.wait_for(lambda: self.queue or self.stop)
+                if self.stop:
+                    return
+                kind = self.queue.pop(0)
+            if kind == "graph":
+                deadline = time.monotonic() + 10.0
+                while int(self.flag[0]) != 1 and time.monotonic() < deadline:
+                    time.sleep(0.0005)
+                self.flag[0] = 0
+            with self.cv:
+                self.done += 1
+                self.cv.notify_all()
+
+    def wait_done(self, n: int, timeout: float) -> bool:
+        with self.cv:
+            return self.cv.wait_for(lambda: self.done >= n, timeout)
+
+    def close(self) -> None:
+        with self.cv:
+            self.stop = True
+            self.cv.notify_all()
+
+
+class _Event:
+    """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record()."""
+
+    def __init__(self, stream: _Stream):
+        self.stream, self.pos = stream, 0
+
+    def record(self, _stream=None) -> None:
+        self.pos = self.stream.enqueued
+
+    def synchronize(self) -> None:
+        if not self.stream.wait_done(self.pos, 2.0):
+            raise TimeoutError(f"a fill waited on stream position {self.pos}, which never completed")
+
+
+class _Store:
+    def __init__(self, flag: torch.Tensor):
+        self.flag = flag
+
+    def stage(self, tokens_addr, n, staging_addr):
+        pass
+
+    def flush(self, signal_addr):
+        if signal_addr:
+            self.flag[0] = 1
+
+
+def test_overlapped_graph_fills_never_wait_on_their_own_launch(monkeypatch):
+    flag = torch.zeros(1, dtype=torch.int64)
+    stream = _Stream(flag)
+    monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None)
+    monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda *a, **k: _Event(stream))
+    t = DiskRowTable.__new__(DiskRowTable)
+    t._device = None
+    t._wait_sync = True
+    t._flag = flag
+    t._store = _Store(flag)
+    t._token_readback = torch.zeros(8, dtype=torch.int32)
+    t._graph_pinned = torch.zeros(64, dtype=torch.uint8)
+    t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)]
+    t._eager_read = [_Event(stream), _Event(stream)]
+    t._eager_slot = 0
+    # if upstream's post-launch fence is ever reintroduced, model it with real stream semantics
+    # so this test fails by hanging the way production would, not on a missing attribute
+    t._graph_consumed = _Event(stream)
+    t._token_bytes = 1
+    t.eos_token_id = EOS
+    t.image_token_id = None
+
+    tokens = [1, 2, 3]
+    steps = 12
+    try:
+        for step in range(steps):
+            new = 10 + step
+            req = SimpleNamespace(input_ids=torch.tensor(tokens + [new]), device_len=len(tokens) + 1)
+            batch = SimpleNamespace(is_decode=True, reqs=[req], padded_size=1,
+                                    input_ids=torch.tensor([new], dtype=torch.int32))
+            with t.forward_host_ctx(batch, use_graph=True):
+                stream.enqueue("graph")  # cuGraphLaunch returns before the GPU passes the WAIT
+            tokens.append(new)
+            # the overlap loop keeps the engine at most one launch ahead of the GPU
+            assert stream.wait_done(step, 4.0), f"decode hung: graph {step - 1} never completed"
+            t._raise_failed_fill()
+        assert stream.wait_done(steps, 4.0), "decode hung on the last step"
+        t._raise_failed_fill()
+    finally:
+        t.close()
+        stream.close()

From ec5013b2770e5eab9c5d087882a03c352b8daa27 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 7 Oct 2026 17:33:31 -0400
Subject: [PATCH 195/249] fix(kernel): keep one e4m3_f32_to_u8, #408's copy
 duplicates the one #603 added to main

Assisted-by: Claude Opus 5.5
---
 python/freetoken/kernel/triton/e4m3_compat.py | 27 -------------------
 1 file changed, 27 deletions(-)

diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py
index 71ada11edc..b14a7ca68b 100644
--- a/python/freetoken/kernel/triton/e4m3_compat.py
+++ b/python/freetoken/kernel/triton/e4m3_compat.py
@@ -201,33 +201,6 @@ def round_e4m3(x):
     return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub)
 
 
-@jit
-def e4m3_f32_to_u8(x):
-    """Encode an fp32 value that ALREADY lies on the e4m3 grid -- the output of
-    :func:`round_e4m3`, clamped to +-448 -- into its e4m3 byte code. This is the
-    encoder the pre-sm_89 emulated path needs to STORE fp8-sized data (the fp8
-    type itself is unavailable there, so the bytes live in a uint8 buffer that
-    :func:`e4m3_u8_to_f32` decodes back).
-
-    Normal range: read the (unbiased) exponent and the now-zero-padded fp32
-    mantissa back out of the fp32 header. Subnormal range (|x| < 2^-6, grid step
-    2^-9): the value is an exact multiple of 2^-9, so ``|x| * 512`` IS the mantissa
-    field -- the sign bit has to be carried in by hand, since that branch never
-    looks at the header. The same 0.015625 boundary as :func:`round_e4m3` keeps the
-    two consistent: ``e4m3_u8_to_f32(e4m3_f32_to_u8(round_e4m3(x)))`` is x's
-    single-rounded value for every input, and no code it emits is a NaN pattern (the
-    caller's +-448 clamp caps the code at 0x7E). ``-0.0`` encodes as 0x00 after
-    round_e4m3 (which documents returning +0.0 for it).
-    """
-    u = x.to(tl.uint32, bitcast=True)
-    sign = ((u >> 31) & 1).to(tl.int32)
-    exp = ((u >> 23) & 0xFF).to(tl.int32) - 127
-    mant = ((u >> 20) & 7).to(tl.int32)
-    normal = (sign << 7) | ((exp + 7) << 3) | mant
-    sub = (sign << 7) | (tl.abs(x) * 512.0).to(tl.int32)
-    return tl.where(tl.abs(x) >= 0.015625, normal, sub).to(tl.uint8)
-
-
 @jit
 def kv_load_e4m3_tile_f32(ptrs, mask):
     """Load a tile of KV e4m3 codes and widen it to fp32.

From c6c696e8aceee378beed789c62743477436c492f Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 7 Oct 2026 17:38:06 -0400
Subject: [PATCH 196/249] fix(server): leave late Anthropic system messages to
 the inline-system policy instead of hoisting them

#487's hoist ran before #608's fold on /v1/messages and moved a late system reminder into
the system head, so every reminder missed the prefix cache; the OpenAI and Responses
paths keep the hoist the qwen templates need.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/server/anthropic_api.py | 3 ++-
 python/freetoken/server/generation.py    | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 545cf4e7cc..35d902b9d6 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -326,7 +326,8 @@ def convert_anthropic_prompt(
         # default only covers requests that leave thinking out.
         ctk = apply_default_thinking_mode(ctk, default_thinking_mode)
 
-    return render_messages(messages), template_tools, parser_tools, ctk
+    # late system reminders stay in place: the tokenizer folds them instead (prefix reuse)
+    return render_messages(messages, hoist_system=False), template_tools, parser_tools, ctk
 
 
 def convert_anthropic_to_genspec(
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 0ed933feb8..e5a215658b 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -516,16 +516,17 @@ def pick(value, key, framework):
     return params
 
 
-def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+def render_messages(messages: list[dict[str, Any]], *, hoist_system: bool = True) -> list[dict[str, Any]]:
     """Normalize OpenAI-shaped message dicts for the chat template: flatten text
     content parts to a string and decode tool-call arguments from JSON. Raises
     ValueError on a non-text content part (text-only server). Shared by all adapters.
 
     Some chat templates (e.g. Qwen3.6) require the system message at index 0;
     hoist system messages to the front and merge multiples into one to satisfy
-    that constraint."""
+    that constraint. ``hoist_system=False`` leaves late system messages in place
+    for the tokenizer's inline-system policy (the Anthropic path)."""
     rendered = [_render_message(m) for m in messages]
-    if any(m.get("role") == "developer" for m in rendered):
+    if not hoist_system or any(m.get("role") == "developer" for m in rendered):
         # The tokenizer maps developer per template and then merges in order
         # (_map_developer_role); hoisting system alone here would put it ahead of
         # a developer message that came first.

From 9e84b61b603f37de0c4d1830ba1e0169a4f25bb9 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 7 Oct 2026 17:39:01 -0400
Subject: [PATCH 197/249] test(server): give main's DeepSeek-V4.1 detector its
 contract class and the count fake the inline-system kwarg

Assisted-by: Claude Opus 5.5
---
 tests/server/test_default_thinking_mode.py |  2 +-
 tests/server/test_function_call_parser.py  | 13 ++++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py
index 7f9938f60c..56854641f2 100644
--- a/tests/server/test_default_thinking_mode.py
+++ b/tests/server/test_default_thinking_mode.py
@@ -110,7 +110,7 @@ def ctk(**extra):
 def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch):
     captured = {}
 
-    async def fake_count(messages, template_tools, ctk, state):
+    async def fake_count(messages, template_tools, ctk, state, **_):
         captured["ctk"] = ctk
         return 1
 
diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py
index 23cd12272e..8af5bc88b6 100644
--- a/tests/server/test_function_call_parser.py
+++ b/tests/server/test_function_call_parser.py
@@ -454,6 +454,12 @@ class TestDeepSeekV32Detector(_DetectorContract):
     truncate_before = '</｜DSML｜invoke>'
 
 
+class TestDeepSeekV41Detector(_DetectorContract):
+    parser_name = 'deepseekv41'
+    block = '<｜DSML｜ calls><｜DSML｜ invoke name="read"><｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py</｜DSML｜ parameter></｜DSML｜ invoke></｜DSML｜ calls>'
+    truncate_before = '</｜DSML｜ invoke>'
+
+
 class TestQwen3CoderDetector(_DetectorContract):
     parser_name = 'qwen3_coder'
     block = '<tool_call><function=read><parameter=filePath>/tmp/test_calc.py</parameter></function></tool_call>'
@@ -498,5 +504,10 @@ def test_contract_classes_cover_every_concrete_detector():
         type(cls().parser().detector)
         for cls in _DetectorContract.__subclasses__()
     }
-    assert covered == set(BaseFormatDetector.__subclasses__())
+    concrete, todo = set(), list(BaseFormatDetector.__subclasses__())
+    while todo:  # a detector may extend another (DeepSeekV41Detector < DeepSeekV32Detector)
+        cls = todo.pop()
+        concrete.add(cls)
+        todo.extend(cls.__subclasses__())
+    assert covered == concrete
     assert covered == set(FunctionCallParser.ToolCallParserEnum.values())

From 35a8555647278eb97f9653512691c8f9aab79a6e Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 7 Oct 2026 18:46:57 -0400
Subject: [PATCH 198/249] test(moe): let the prefetch counter test's copy spies
 pass #610's blocks_per_bank through

Assisted-by: Claude Opus 5.5
---
 tests/models/qwen4_exp/test_skeleton.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py
index 2b47748f53..0803cbca28 100644
--- a/tests/models/qwen4_exp/test_skeleton.py
+++ b/tests/models/qwen4_exp/test_skeleton.py
@@ -1117,13 +1117,13 @@ def spy_ensure(layer, ids, **kw):
         calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone()))
         return real_ensure(layer, ids, **kw)
 
-    def spy_slim(*args):
+    def spy_slim(*args, **kw):
         streams["slim"].add(torch.cuda.current_stream().cuda_stream)
-        return real_slim(*args)
+        return real_slim(*args, **kw)
 
-    def spy_multi(*args):
+    def spy_multi(*args, **kw):
         streams["demand"].add(torch.cuda.current_stream().cuda_stream)
-        return real_multi(*args)
+        return real_multi(*args, **kw)
 
     def install():
         monkeypatch.setattr(cache, "prefetch_ensure", spy_pf)

From 709cc051df877dddf425136b7841905d594ade83 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 09:43:17 -0400
Subject: [PATCH 199/249] test(engine): give #650's rebuild engine fake the kv
 offloader and gdn group the fit check reads

The carried host kv tier (FlashML-org/FreeToken#499) rejects a pages rebuild through
engine.kv_offloader, and the fork's gdn workspace reserve asks the model config for its
linear attention group; the SimpleNamespace shell had neither.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 tests/engine/test_rebuild_runtime_cache_ordering.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py
index 229431b7a4..aae7d49cf0 100644
--- a/tests/engine/test_rebuild_runtime_cache_ordering.py
+++ b/tests/engine/test_rebuild_runtime_cache_ordering.py
@@ -125,12 +125,14 @@ def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: in
     config = types.SimpleNamespace(
         page_size=16, max_seq_len=4096, max_running_req=4, memory_ratio=1.0,
         cuda_graph_max_bs=4, swa_num_pages_override=None,
-        model_config=types.SimpleNamespace(vocab_size=128, model_is_mrope=False),
+        model_config=types.SimpleNamespace(
+            vocab_size=128, model_is_mrope=False, linear_attention_group=lambda: None
+        ),
     )
     page_table = torch.zeros((config.max_running_req + 1, 256), dtype=torch.int32)
     eng = types.SimpleNamespace(
         config=config, moe_offload_cache=_FakeMoeCache(ledger, moe), kv_cache=_FakeKvPool(ledger, kv),
-        linear_state_pool=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None,
+        linear_state_pool=None, kv_offloader=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None,
         _baseline_free=budget, _weights_bytes=0, page_table=page_table, max_seq_len=4096,
         ctx=types.SimpleNamespace(page_table=page_table),
         dummy_req=types.SimpleNamespace(table_idx=0),

From 2dad6443addfe347a085305da1ba0031f4c31ee8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 09:43:17 -0400
Subject: [PATCH 200/249] test(server): expect out-of-range penalties to fail
 request validation

Main (FlashML-org/FreeToken#527) bounds presence/frequency penalties to [-2, 2] on the
request models, so they never reach the handler's 400 path.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 tests/server/test_openai_extras.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py
index 115112b9d0..19140d8b72 100644
--- a/tests/server/test_openai_extras.py
+++ b/tests/server/test_openai_extras.py
@@ -8,6 +8,7 @@
 import json
 from types import SimpleNamespace
 
+import pydantic
 import pytest
 from freetoken.message import TokenizeMsg, UserReply
 from freetoken.server.api_models import (
@@ -250,8 +251,6 @@ def test_sampling_extras_reach_sampling_params():
     "field",
     [
         {"min_p": 1.5},
-        {"presence_penalty": 3.0},
-        {"frequency_penalty": -2.5},
         {"repetition_penalty": 0.0},
         {"logit_bias": {"x": 1.0}},
         {"min_tokens": -1},
@@ -263,6 +262,13 @@ def test_bad_sampling_extras_are_400(field):
     assert body.status_code == 400, field
 
 
+@pytest.mark.parametrize("field", [{"presence_penalty": 3.0}, {"frequency_penalty": -2.5}])
+def test_out_of_range_penalties_fail_request_validation(field):
+    # main's request models bound both penalties to [-2, 2] (FastAPI answers 422)
+    with pytest.raises(pydantic.ValidationError):
+        chat(**field)
+
+
 def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template():
     spec = chat_request_to_genspec(
         chat(

From 9db9dae350dc3a7de97c4647fabff28d58145238 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 11:28:51 -0400
Subject: [PATCH 201/249] fix(kernel): decode nvfp4 qsa prefill tiles narrow at
 head_dim 256

The wide prefill profile (BLOCK_N 64, 2 warps) that the nvfp4 path inherits from the
bf16 tuning has load_nvfp4 decode a 64 x 256 fp32 tile for K and V on 64 threads. Triton
3.8 compiles it to 40 registers with ~1.4 KB of spills: 439 ms instead of 45 ms at 4096
rows on sm_120, and a ~556 MiB driver local-memory reservation outside the torch
allocator that OOMs this deployment on its first prefill. BLOCK_N 16 at 4 warps does not
spill: 12.3 ms on Triton 3.8 (33 ms on 3.6), same output to bf16 rounding.

Upstream: FlashML-org/FreeToken#675, root cause posted on FlashML-org/FreeToken#408.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/kernel/triton/qsa/attend.py |  7 +++
 tests/kernels/test_qsa_nvfp4.py              | 55 ++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py
index cf74e4e95f..74ebf95c48 100644
--- a/python/freetoken/kernel/triton/qsa/attend.py
+++ b/python/freetoken/kernel/triton/qsa/attend.py
@@ -407,6 +407,13 @@ def qsa_sparse_paged_attention(
     else:
         block_n, target_splits, partial_warps = 64, 1, 2
 
+    if nvfp4 and block_n == 64 and head_dim >= 256:
+        # The wide bf16 profile has load_nvfp4 decode a BLOCK_N x HEAD_DIM fp32 tile for K and V
+        # across 64 threads (256 values per thread at HEAD_DIM 256): Triton 3.8 compiles that to
+        # 40 registers with ~1.4 KB of spills, ~10x slower and a ~550 MiB local-memory reservation
+        # outside the torch allocator. Narrow tiles at 4 warps do not spill and are faster on 3.6 too.
+        block_n, partial_warps = 16, 4
+
     # These tiles were tuned for the GB300's 228KB LDS; AMD parts have less and it varies by arch
     # (e.g. 64KB/CU on RDNA gfx11/gfx12), so k+v staging can overflow. Query the device's actual
     # shared-memory limit rather than assuming a fixed size, and shrink the KV tile only when it
diff --git a/tests/kernels/test_qsa_nvfp4.py b/tests/kernels/test_qsa_nvfp4.py
index 591c6f1d4e..ec96e34563 100644
--- a/tests/kernels/test_qsa_nvfp4.py
+++ b/tests/kernels/test_qsa_nvfp4.py
@@ -139,3 +139,58 @@ def test_qsa_nvfp4_splitk_reads_value_rows():
     ref = _decode(codes, row, block).view(3, PAGE, 1, HEAD_DIM)
     want = qsa_sparse_paged_attention(q, ref, ref, *args)
     assert torch.equal(got, want)
+
+
+def test_qsa_nvfp4_wide_head_prefill_uses_narrow_tiles(monkeypatch):
+    """At HEAD_DIM 256 the prefill profile must not decode 64-token fp32 tiles on 2 warps:
+    Triton 3.8 spills that variant (~10x slower, ~550 MiB of driver local memory)."""
+    from freetoken.kernel.triton.qsa import attend
+
+    head_dim, kv_heads, query_heads, rows, pages = 256, 2, 24, 32, 3
+    seen = []
+    kernel = attend._qsa_sparse_paged_gqa_splitk_kernel
+
+    class Spy:
+        def __getitem__(self, grid):
+            launch = kernel[grid]
+
+            def run(*args, **kwargs):
+                seen.append((kwargs["BLOCK_N"], kwargs["num_warps"]))
+                return launch(*args, **kwargs)
+            return run
+
+    torch.manual_seed(29)
+    slots = pages * PAGE
+    k = torch.randn(slots, kv_heads * head_dim, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k) * 0.25
+    codes = [torch.zeros((slots, kv_heads, head_dim // 2), dtype=torch.uint8, device="cuda") for _ in range(2)]
+    rows_s = [torch.zeros((slots, kv_heads), dtype=torch.float32, device="cuda") for _ in range(2)]
+    blocks = [torch.zeros((slots, kv_heads, head_dim // 16), dtype=torch.uint8, device="cuda") for _ in range(2)]
+    quantize_nvfp4_to_cache(
+        k, v, torch.arange(slots, dtype=torch.int32, device="cuda"), codes[0], codes[1],
+        rows_s[0], rows_s[1], blocks[0], blocks[1],
+    )
+    q = torch.randn(rows, query_heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    indices, block_table, token_to_req = _layout(rows, topk=128)
+    packed = (pages, PAGE, kv_heads, head_dim // 2)
+    monkeypatch.setattr(attend, "_qsa_sparse_paged_gqa_splitk_kernel", Spy())
+    got = qsa_sparse_paged_attention(
+        q, codes[0].view(packed), codes[1].view(packed), indices, block_table, token_to_req,
+        k_scale=rows_s[0], v_scale=rows_s[1], kv_quant="nvfp4",
+        k_block_scale=blocks[0], v_block_scale=blocks[1],
+    )
+    assert seen == [(16, 4)]
+
+    def decode(c, r, b):
+        lo, hi = (c & 15).to(torch.long), (c >> 4).to(torch.long)
+        code = torch.stack((lo, hi), dim=-1).reshape(*c.shape[:-1], head_dim)
+        mag = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=c.device)[code & 7]
+        val = torch.where(code < 8, mag, -mag)
+        blk = b.view(torch.float8_e4m3fn).to(torch.float32).repeat_interleave(16, dim=-1)
+        return (val * blk * r.unsqueeze(-1)).to(torch.bfloat16).view(pages, PAGE, kv_heads, head_dim)
+
+    want = qsa_sparse_paged_attention(
+        q, decode(codes[0], rows_s[0], blocks[0]), decode(codes[1], rows_s[1], blocks[1]),
+        indices, block_table, token_to_req,
+    )
+    torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-2)

From 0448e713a90923fd32dea1b0c09304ad4dfe8250 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 05:16:35 -0400
Subject: [PATCH 202/249] fix(kernel): size triton's autotune flush buffer to
 the device l2

A kernel autotunes on its first launch of each key, which can be the first request after
the cache planner has handed out nearly all VRAM; triton's flat 256 MB flush buffer was an
OOM there on an 11 GB card. Four times the L2 still evicts it between timed runs.

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit ffbf001acec04f4f882d6cccbc1d0d4cbbec633f)
(cherry picked from commit e097977e525302a6f773ace34ca553666e593876)
---
 python/freetoken/engine/engine.py             |  3 +++
 .../freetoken/kernel/triton/autotune_cache.py | 26 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index f5c2e0c183..db06a1e136 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -513,6 +513,9 @@ def __init__(self, config: EngineConfig):
         logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}")
 
         # ======================= Model initialization ========================
+        from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer
+
+        bound_autotune_flush_buffer(self.device.index)
         set_rope_device(self.device)
         with torch.device("meta"), torch_dtype(config.dtype):
             self.model = create_model(config.model_config)
diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py
index ea95768368..2e57fadde9 100644
--- a/python/freetoken/kernel/triton/autotune_cache.py
+++ b/python/freetoken/kernel/triton/autotune_cache.py
@@ -34,4 +34,28 @@
     {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {}
 )
 
-__all__ = ["autotune_cache_kwargs"]
+
+
+def bound_autotune_flush_buffer(device_index: int) -> None:
+    """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB.
+
+    A kernel autotunes on its first launch of each key, which can be the first request, after the
+    cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small
+    card. Four times the L2 still evicts everything between timed runs.
+    """
+    import torch
+
+    try:
+        from triton.backends.nvidia.driver import CudaDriver
+    except ImportError:
+        return
+    l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0)
+    nbytes = min(256 << 20, max(8 << 20, 4 * l2))
+
+    def get_empty_cache_for_benchmark(self):
+        return torch.empty(nbytes // 4, dtype=torch.int, device="cuda")
+
+    CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark
+
+
+__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"]

From 9917f68c48ebd536d4d84e7e22a24ae29baca204 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 10:02:35 -0400
Subject: [PATCH 203/249] test(utils): drive the async coalesced-pull test with
 asyncio.run

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 32007ecb48aee2dc619fc9ceb89ad9b5dd70ef14)
---
 tests/utils/test_mp_coalesced.py | 31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py
index e599db7265..c64ee1c5c3 100644
--- a/tests/utils/test_mp_coalesced.py
+++ b/tests/utils/test_mp_coalesced.py
@@ -92,7 +92,7 @@ def test_sub_queue_coalesced(self):
             sub.stop()
 
 
-async def test_async_pull_coalesced():
+def test_async_pull_coalesced():
     import sys
 
     from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull
@@ -101,16 +101,19 @@ async def test_async_pull_coalesced():
     # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI.
     if sys.platform == "win32":
         pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)")
-    push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o)
-    pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder)
-    try:
-        await push.put({"n": 1})
-        await push.put({"n": 2})
-        await asyncio.sleep(0.1)
-
-        # Even if both messages arrive in a single frame, two gets must succeed.
-        assert await pull.get() == {"n": 1}
-        assert await pull.get() == {"n": 2}
-    finally:
-        push.stop()
-        pull.stop()
\ No newline at end of file
+    async def scenario():
+        push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o)
+        pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder)
+        try:
+            await push.put({"n": 1})
+            await push.put({"n": 2})
+            await asyncio.sleep(0.1)
+
+            # Even if both messages arrive in a single frame, two gets must succeed.
+            assert await pull.get() == {"n": 1}
+            assert await pull.get() == {"n": 2}
+        finally:
+            push.stop()
+            pull.stop()
+
+    asyncio.run(scenario())

From 2019d4f5d68aeee163cad56a828e24b0d3ac5b09 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:35:40 -0400
Subject: [PATCH 204/249] fix(engine): keep multimodal placeholder ids out of
 the repetition-penalty prompt mask

Fork #47's sampler plan scattered the raw prompt ids into a [rows, vocab+1] mask; image
tokens carry pseudo-ids >= 1,000,000, so repetition_penalty with an image indexed out of
bounds -- a device-side assert that takes the server down. They now land in the scratch
column.

Also keeps presence/frequency out of the LogitsPlan: with #527 device counts and
#393 plan both applying them, a request with a penalty plus any plan processor was
penalized twice.

(cherry picked from commit 66053ded)

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit b9883902735cadb2232e3ae4d9dd18e02481a189)
---
 python/freetoken/engine/sample.py      | 23 ++++++-----------------
 tests/engine/test_logits_processors.py | 23 ++++++++++++++++++++---
 tests/engine/test_sample_logprobs.py   |  4 ++--
 3 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 89d2a6fb88..8e05f80694 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -172,35 +172,24 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
         )
         pad = self.vocab_size
 
-        want_hist = any(
-            p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0
-            for _, p in picked
-        )
-        if want_hist:
+        # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties)
+        if any(p.repetition_penalty != 1.0 for _, p in picked):
             hist: list[list[int]] = []
             prompts: list[list[int]] = []
-            want_prompt = False
             for req, p in picked:
                 prompt_len = req.max_device_len - req.output_len
                 ids = req.input_ids
-                penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0
-                hist.append(ids[prompt_len:].tolist() if penalized else [])
                 if p.repetition_penalty != 1.0:
-                    want_prompt = True
+                    hist.append(ids[prompt_len:].tolist())
                     # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE);
                     # send them to the scratch column, not past the scatter target
                     prompts.append(ids[:prompt_len].clamp(max=pad).tolist())
                 else:
+                    hist.append([])
                     prompts.append([])
             plan.hist_ids = _padded_rows(hist, pad, self.device)
-            if want_prompt:
-                plan.prompt_ids = _padded_rows(prompts, pad, self.device)
-            if any(p.presence_penalty != 0.0 for _, p in picked):
-                plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device)
-            if any(p.frequency_penalty != 0.0 for _, p in picked):
-                plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device)
-            if any(p.repetition_penalty != 1.0 for _, p in picked):
-                plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device)
+            plan.prompt_ids = _padded_rows(prompts, pad, self.device)
+            plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device)
 
         bias_rows: list[int] = []
         bias_ids: list[int] = []
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index abd2e1f939..237ce9f602 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -172,9 +172,8 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
 
     logits = torch.zeros((4, V))
     out = apply_logits_processors(logits, plan, V)
-    assert (
-        out[1, 3].item() == -1.0
-    )  # presence on the repeated 3 (repetition on 0 is a no-op)
+    # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op
+    assert plan.presence is None and out[1, 3].item() == 0.0
     assert out[1, 5].item() == 2.0
     assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf")
     assert torch.isfinite(out[3]).all()
@@ -182,6 +181,24 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
     assert sampler.sample(logits, args)[1].item() == 5
 
 
+def test_presence_and_frequency_ride_the_device_counts_not_the_plan():
+    # the plan must not apply them a second time on top of the sampler's own counts
+    sampler = Sampler(CPU, V)
+    only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)])
+    args = sampler.prepare(only)
+    assert args.plan is None and [row for row, *_ in args.penalties] == [0]
+
+    req = _req([1], [], 4, presence_penalty=1.0, logit_bias=[[5, 0.25]])
+    args = sampler.prepare(SimpleNamespace(reqs=[req]))
+    assert args.plan is not None and [row for row, *_ in args.penalties] == [0]
+    req.output_token_counts[3] = 1
+    logits = torch.zeros((1, V))
+    logits[0, 3] = 1.5
+    # once: 3 -> 0.5 beats 5 -> 0.25; twice would drop 3 to -0.5 and pick 5
+    assert sampler.sample(logits, args).item() == 3
+    assert req.output_token_counts[3].item() == 2
+
+
 def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt():
     # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the
     # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA)
diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
index 6bd204604d..fae5b7c25b 100644
--- a/tests/engine/test_sample_logprobs.py
+++ b/tests/engine/test_sample_logprobs.py
@@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc
 
     batch = SimpleNamespace(
         reqs=[
-            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)),
-            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)),
+            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)),
+            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)),
         ]
     )
     args = sampler.prepare(batch)

From b712deed7a284e1b393bac308d531297fc5e628a Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:43:17 -0400
Subject: [PATCH 205/249] fix(engine): apply the repetition penalty from device
 state, not a per-step prompt upload

(up to 262K ids) and a synchronous pageable copy, which stalled the overlap loop and saw
the output one token late. The prompt mask is now built once per request, the output
side reuses #527's device counts, and the order is vLLM's (repetition, then presence and
frequency).

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 343d7421dfebf72244b9e9d16261dba073b48cbc)
(cherry picked from commit 3df811dee9bcd4af287e3cc94abb689d531e58bd)
---
 python/freetoken/core.py               |  11 +-
 python/freetoken/engine/sample.py      | 188 ++++++++++++++-----------
 tests/engine/test_logits_processors.py |  83 +++++++----
 tests/engine/test_sample_logprobs.py   |   5 +-
 4 files changed, 172 insertions(+), 115 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index ab5b81b0af..a7811ed712 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -30,7 +30,8 @@ class SamplingParams:
     # min_p: drop tokens whose probability is below min_p x the top probability.
     min_p: float = 0.0
     # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the
-    # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off).
+    # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler
+    # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask).
     presence_penalty: float = 0.0
     frequency_penalty: float = 0.0
     repetition_penalty: float = 1.0
@@ -60,11 +61,10 @@ def is_greedy(self) -> bool:
 
     @property
     def needs_logits_processing(self) -> bool:
+        # the penalties are not LogitsPlan work: the sampler applies them from its own
+        # device-side state (Sampler.prepare's penalties).
         return bool(
             self.min_p > 0.0
-            or self.presence_penalty != 0.0
-            or self.frequency_penalty != 0.0
-            or self.repetition_penalty != 1.0
             or self.logit_bias
             or self.min_tokens > 0
         )
@@ -106,7 +106,10 @@ class Req:
     # handler must not free resources under an in-flight forward; it sets this flag and
     # _process_last_data frees the request when the batch drains (after copy_done.synchronize).
     aborted: bool = False
+    output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False)
     stop_decode_status: DecodeStatus | None = None
+    # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once.
+    prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         assert self.input_ids.is_cpu
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 8e05f80694..9d12b5f571 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, List
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, NamedTuple
 
 import torch
 from freetoken.utils import is_sm90_supported, nvtx_annotate
@@ -13,41 +13,49 @@
 @dataclass
 class LogitsPlan:
     """Inputs of the logits processors for the rows of one batch that asked for any
-    (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p).
-
-    Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and
-    token histories, applied on the device by ``apply_logits_processors`` before the
-    sampling kernel. ``None`` when no request in the batch uses a processor, so the
-    default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1,
-    in the order of ``rows``); ``rows`` maps them back to the batch.
+    (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies
+    them from per-request device state (``RowPenalty``).
+
+    Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied
+    on the device by ``apply_logits_processors`` before the sampling kernel. ``None``
+    when no request in the batch uses a processor, so the default path stays exactly
+    as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``);
+    ``rows`` maps them back to the batch.
     """
 
     rows: torch.Tensor  # [r] int64 batch rows
     temps: torch.Tensor  # [r, 1] float32 temperatures (min_p works on probabilities)
-    presence: torch.Tensor | None = None  # [r, 1] float32
-    frequency: torch.Tensor | None = None  # [r, 1] float32
-    repetition: torch.Tensor | None = None  # [r, 1] float32 (1.0 = off)
-    # Generated ids per row, right-padded with vocab_size (a scratch column that is
-    # dropped), for the presence / frequency counts and the repetition mask.
-    hist_ids: torch.Tensor | None = None  # [r, L] int64
-    # Prompt ids per row, same padding: repetition_penalty also covers the prompt.
-    prompt_ids: torch.Tensor | None = None  # [r, P] int64
     bias_rows: torch.Tensor | None = None  # [b] int64 local rows
     bias_ids: torch.Tensor | None = None  # [b] int64
     bias_vals: torch.Tensor | None = None  # [b] float32
-    # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must
-    # not sample yet.
+    # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with
+    # vocab_size, a scratch column that is dropped) they must not sample yet.
     min_rows: torch.Tensor | None = None  # [m] int64 local rows
     min_ids: torch.Tensor | None = None  # [m, S] int64
     min_p: torch.Tensor | None = None  # [r, 1] float32 (0 = off for that row)
 
 
+class RowPenalty(NamedTuple):
+    """One batch row's presence / frequency / repetition penalty. ``counts`` are the
+    request's generated-token counts ([V] int32), updated on the device after each draw so
+    overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's
+    ids for the repetition penalty and is None when that is off."""
+
+    row: int
+    counts: torch.Tensor
+    presence: float
+    frequency: float
+    repetition: float = 1.0
+    prompt_mask: torch.Tensor | None = None
+
+
 @dataclass
 class BatchSamplingArgs:
     temperatures: torch.Tensor | None
     top_k: torch.Tensor | None = None
     top_p: torch.Tensor | None = None
     greedy_mask: torch.Tensor | None = None
+    penalties: list[RowPenalty] = field(default_factory=list)
     plan: LogitsPlan | None = None
     logprob_rows: torch.Tensor | None = None
     max_top_logprobs: int = 0
@@ -60,47 +68,39 @@ def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) ->
 
 
 def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor:
-    """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the
-    histories can be long (prompt + output) and a pinned buffer must not be freed under
-    an in-flight non_blocking copy."""
+    """[len(rows), max_len] int64 with ``pad`` on the right."""
     width = max(1, max((len(r) for r in rows), default=1))
-    out = torch.full((len(rows), width), pad, dtype=torch.int64)
-    for i, r in enumerate(rows):
-        if r:
-            out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64)
-    return out.to(device)
+    return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device)
+
+
+def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor:
+    """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order:
+    repetition (over prompt + generated ids), then frequency and presence (generated)."""
+    logits = logits.float().clone()
+    for row, counts, presence, frequency, repetition, prompt_mask in penalties:
+        if repetition != 1.0:
+            x = logits[row]
+            seen = counts > 0
+            if prompt_mask is not None:
+                seen = seen | prompt_mask
+            # HF / vLLM semantics: divide positive logits, multiply negative ones.
+            penalized = torch.where(x > 0, x / repetition, x * repetition)
+            logits[row] = torch.where(seen, penalized, x)
+        if presence or frequency:
+            logits[row] -= frequency * counts + presence * (counts > 0)
+    return logits
 
 
 def apply_logits_processors(
     logits: torch.Tensor, plan: LogitsPlan, vocab_size: int
 ) -> torch.Tensor:
-    """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order:
-    repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then
-    min_p. Rows outside the plan are copied unchanged."""
+    """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the
+    min_tokens mask, then min_p. Rows outside the plan are copied unchanged."""
     out = logits.to(torch.float32, copy=True)
     sub = out[plan.rows]  # [r, V] (advanced indexing copies)
-    r = sub.shape[0]
     device = sub.device
     neg_inf = float("-inf")
 
-    if plan.hist_ids is not None:
-        counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device)
-        counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32))
-        counts = counts[:, :vocab_size]
-        if plan.repetition is not None:
-            seen = counts > 0
-            if plan.prompt_ids is not None:
-                in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device)
-                in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool))
-                seen = seen | in_prompt[:, :vocab_size]
-            # HF / vLLM semantics: divide positive logits, multiply negative ones.
-            penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition)
-            sub = torch.where(seen, penalized, sub)
-        if plan.presence is not None:
-            sub = sub - plan.presence * (counts > 0).to(sub.dtype)
-        if plan.frequency is not None:
-            sub = sub - plan.frequency * counts
-
     if plan.bias_ids is not None:
         sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True)
 
@@ -155,11 +155,7 @@ class Sampler:
     vocab_size: int
 
     def _plan(self, batch: Batch, params) -> LogitsPlan | None:
-        """The LogitsPlan for this batch, or None when no request asks for a processor.
-        Token histories come from the host-side ``req.input_ids``; under overlap
-        scheduling the previous step's token may not be appended yet, so a penalty sees
-        the history one token late. That is the accepted cost of keeping the counts off
-        the device."""
+        """The LogitsPlan for this batch, or None when no request asks for a processor."""
         MIN_T = 1e-6
         rows = [i for i, p in enumerate(params) if p.needs_logits_processing]
         if not rows:
@@ -172,25 +168,6 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
         )
         pad = self.vocab_size
 
-        # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties)
-        if any(p.repetition_penalty != 1.0 for _, p in picked):
-            hist: list[list[int]] = []
-            prompts: list[list[int]] = []
-            for req, p in picked:
-                prompt_len = req.max_device_len - req.output_len
-                ids = req.input_ids
-                if p.repetition_penalty != 1.0:
-                    hist.append(ids[prompt_len:].tolist())
-                    # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE);
-                    # send them to the scratch column, not past the scatter target
-                    prompts.append(ids[:prompt_len].clamp(max=pad).tolist())
-                else:
-                    hist.append([])
-                    prompts.append([])
-            plan.hist_ids = _padded_rows(hist, pad, self.device)
-            plan.prompt_ids = _padded_rows(prompts, pad, self.device)
-            plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device)
-
         bias_rows: list[int] = []
         bias_ids: list[int] = []
         bias_vals: list[float] = []
@@ -223,9 +200,48 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
             plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device)
         return plan
 
+    def _prompt_mask(self, req) -> torch.Tensor:
+        """[V] bool marking the request's prompt ids, uploaded once per request."""
+        prompt_len = req.max_device_len - req.output_len
+        # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE);
+        # send them to the scratch column, not past the scatter target
+        ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size)
+        if self.device.type == "cuda":
+            ids = ids.pin_memory().to(self.device, non_blocking=True)
+        mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device)
+        mask[ids] = True
+        return mask[: self.vocab_size]
+
     def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
         is_greedy = [p.is_greedy for p in params]
+        penalties = []
+        for row, req in enumerate(batch.reqs):
+            p = req.sampling_params
+            repetition = p.repetition_penalty
+            if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0):
+                continue
+            if not req.can_decode:
+                continue
+            if req.output_token_counts is None:
+                req.output_token_counts = torch.zeros(
+                    self.vocab_size, dtype=torch.int32, device=self.device
+                )
+            prompt_mask = None
+            if repetition != 1.0:
+                if req.prompt_token_mask is None:
+                    req.prompt_token_mask = self._prompt_mask(req)
+                prompt_mask = req.prompt_token_mask
+            penalties.append(
+                RowPenalty(
+                    row,
+                    req.output_token_counts,
+                    p.presence_penalty,
+                    p.frequency_penalty,
+                    repetition,
+                    prompt_mask,
+                )
+            )
         plan = self._plan(batch, params)
         want_logprobs = [p.logprobs for p in params]
         logprob_rows = (
@@ -239,6 +255,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
         if all(is_greedy):
             return BatchSamplingArgs(
                 temperatures=None,
+                penalties=penalties,
                 plan=plan,
                 logprob_rows=logprob_rows,
                 max_top_logprobs=max_top_logprobs,
@@ -270,6 +287,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             top_k=top_k,
             top_p=top_p,
             greedy_mask=greedy_mask,
+            penalties=penalties,
             plan=plan,
             logprob_rows=logprob_rows,
             max_top_logprobs=max_top_logprobs,
@@ -278,16 +296,26 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
     @nvtx_annotate("Sampler")
     def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
         with torch.cuda.nvtx.range("Sampler"):
+            if args.penalties:
+                logits = apply_penalties(logits, args.penalties)
+            # after the penalties: min_p must see the final distribution
             if args.plan is not None:
                 logits = apply_logits_processors(logits, args.plan, self.vocab_size)
             if args.temperatures is None:  # greedy sampling
-                return torch.argmax(logits, dim=-1)
-            tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
-            if args.greedy_mask is not None:
-                # Mixed batches still run probability sampling for all rows, but
-                # greedy rows must follow argmax's deterministic tie-breaking.
-                greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
-                tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
+                tokens = torch.argmax(logits, dim=-1)
+            else:
+                tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
+                if args.greedy_mask is not None:
+                    # Mixed batches still run probability sampling for all rows, but
+                    # greedy rows must follow argmax's deterministic tie-breaking.
+                    greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
+                    tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
+            # Update on the sampling stream: overlapped scheduling can prepare the next
+            # batch before the previous token reaches Req.input_ids on the CPU.
+            for row, counts, *_ in args.penalties:
+                counts.scatter_add_(
+                    0, tokens[row : row + 1].long(), counts.new_ones(1)
+                )
             return tokens
 
     def compute_logprobs(
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index 237ce9f602..1829f35f32 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -8,7 +8,13 @@
 import pytest
 import torch
 from freetoken.core import SamplingParams
-from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors
+from freetoken.engine.sample import (
+    LogitsPlan,
+    RowPenalty,
+    Sampler,
+    apply_logits_processors,
+    apply_penalties,
+)
 
 V = 8
 CPU = torch.device("cpu")
@@ -32,33 +38,17 @@ def _ids(rows: list[list[int]]) -> torch.Tensor:
     return out
 
 
-def test_presence_and_frequency_penalties_count_generated_tokens():
-    logits = torch.zeros((2, V))
-    plan = _plan(
-        [1],
-        hist_ids=_ids([[3, 3, 5]]),
-        presence=torch.tensor([[0.5]]),
-        frequency=torch.tensor([[0.25]]),
-    )
-    out = apply_logits_processors(logits, plan, V)
-    assert torch.equal(out[0], logits[0])  # row outside the plan untouched
-    assert out[1, 3].item() == -(0.5 + 2 * 0.25)
-    assert out[1, 5].item() == -(0.5 + 0.25)
-    assert out[1, 0].item() == 0.0
-
-
 def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics():
     logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]])
-    plan = _plan(
-        [0],
-        hist_ids=_ids([[0, 1]]),  # generated: 0 (positive), 1 (negative)
-        prompt_ids=_ids([[2, 3]]),  # prompt: 2 (positive), 3 (negative)
-        repetition=torch.tensor([[2.0]]),
-    )
-    out = apply_logits_processors(logits, plan, V)
+    counts = torch.zeros(V, dtype=torch.int32)
+    counts[[0, 1]] = 1  # generated: 0 (positive), 1 (negative)
+    prompt = torch.zeros(V, dtype=torch.bool)
+    prompt[[2, 3]] = True  # prompt: 2 (positive), 3 (negative)
+    out = apply_penalties(logits, [RowPenalty(0, counts, 0.0, 0.0, 2.0, prompt)])
     assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0  # divided
     assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0  # multiplied
     assert out[0, 4].item() == 1.0  # unseen
+    assert logits[0, 0].item() == 2.0  # the caller's logits stay untouched
 
 
 def test_logit_bias_adds_per_row():
@@ -134,6 +124,9 @@ def _req(
         output_len=max_tokens,
         max_device_len=len(prompt) + max_tokens,
         sampling_params=SamplingParams(**sp),
+        can_decode=True,
+        output_token_counts=None,
+        prompt_token_mask=None,
     )
 
 
@@ -164,16 +157,15 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
     plan = args.plan
     assert plan is not None
     assert plan.rows.tolist() == [1, 2, 3]
-    assert plan.hist_ids[0].tolist()[:2] == [3, 3]  # row 1's generated tokens
-    assert plan.prompt_ids[0].tolist()[:2] == [1, 2]
+    assert [pen.row for pen in args.penalties] == [1]
+    assert batch.reqs[1].prompt_token_mask.nonzero().flatten().tolist() == [1, 2]
     assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5]
     assert plan.min_rows.tolist() == [1]  # local row of the third request only
     assert plan.min_ids[0].tolist() == [0, 7]
 
     logits = torch.zeros((4, V))
     out = apply_logits_processors(logits, plan, V)
-    # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op
-    assert plan.presence is None and out[1, 3].item() == 0.0
+    assert out[1, 3].item() == 0.0  # the penalties are not plan work
     assert out[1, 5].item() == 2.0
     assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf")
     assert torch.isfinite(out[3]).all()
@@ -181,7 +173,7 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
     assert sampler.sample(logits, args)[1].item() == 5
 
 
-def test_presence_and_frequency_ride_the_device_counts_not_the_plan():
+def test_penalties_ride_the_device_state_not_the_plan():
     # the plan must not apply them a second time on top of the sampler's own counts
     sampler = Sampler(CPU, V)
     only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)])
@@ -205,5 +197,38 @@ def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt():
     sampler = Sampler(CPU, V)
     batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)])
     args = sampler.prepare(batch)
-    out = apply_logits_processors(torch.ones(1, V), args.plan, V)
+    out = apply_penalties(torch.ones(1, V), args.penalties)
     assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_min_p_with_temperature_draws_only_the_kept_tokens():
+    # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that
+    # into NaN (which the Pascal fallback then samples as a uniform row)
+    V = 32768
+    sampler = Sampler(torch.device("cuda"), V)
+    req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1)
+    args = sampler.prepare(SimpleNamespace(reqs=[req]))
+    logits = torch.full((1, V), -10.0, device="cuda")
+    kept = [100, 20000, 31000]
+    logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda")
+    for _ in range(32):
+        assert sampler.sample(logits, args).item() in kept
+
+
+def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_device():
+    # the prompt (up to the whole context) is uploaded once per request, not every step,
+    # and the generated side comes from the counts sample() keeps, not the host history
+    sampler = Sampler(CPU, V)
+    req = _req([1, 2], [], 4, repetition_penalty=2.0)
+    batch = SimpleNamespace(reqs=[req])
+    args = sampler.prepare(batch)
+    mask = req.prompt_token_mask
+    logits = torch.tensor([[0.0, 1.0, 1.0, 4.0, 0.5, 0.0, 0.0, 0.0]])
+    assert sampler.sample(logits, args).item() == 3
+    assert req.output_token_counts.tolist() == [0, 0, 0, 1, 0, 0, 0, 0]
+    req.input_ids = torch.tensor([9, 9, 3], dtype=torch.int32)  # a stale host view is not read
+    args = sampler.prepare(batch)
+    assert req.prompt_token_mask is mask
+    out = apply_penalties(logits, args.penalties)
+    assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0]
diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
index fae5b7c25b..88fffcc2ed 100644
--- a/tests/engine/test_sample_logprobs.py
+++ b/tests/engine/test_sample_logprobs.py
@@ -2,6 +2,7 @@
 
 import torch
 
+from freetoken.core import SamplingParams
 from freetoken.engine.sample import BatchSamplingArgs, Sampler
 
 
@@ -137,8 +138,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc
 
     batch = SimpleNamespace(
         reqs=[
-            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)),
-            SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)),
+            SimpleNamespace(sampling_params=SamplingParams(logprobs=True, top_logprobs=17)),
+            SimpleNamespace(sampling_params=SamplingParams(logprobs=False, top_logprobs=0)),
         ]
     )
     args = sampler.prepare(batch)

From 6dc2bb9a738843da55e03263217aaecf8705fc61 Mon Sep 17 00:00:00 2001
From: Siarhei Berdachuk <berdachuk@berdaflex.com>
Date: Mon, 5 Oct 2026 19:11:56 +0300
Subject: [PATCH 206/249] fix(server): clarify unsupported response_format
 error text

After soft-accepting json_object/json_schema, the old "no constrained decoding"
message was misleading for remaining unknown types.

(cherry picked from commit 51e17e6c3380af25d34a301859a8991224f30c2f)
(cherry picked from commit e2df92f2ade5cd0ed4f5f1bd421eef96ded793cb)
---
 python/freetoken/server/openai_api.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 7f94efc250..79dff72963 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -305,8 +305,10 @@ async def handle_chat_completion(
     if logprobs_conflict is not None:
         return create_error_response(logprobs_conflict, param="logprobs")
     if _response_format_unsupported(req.response_format):
+        rtype = (req.response_format or {}).get("type")
         return create_error_response(
-            "response_format json_object/json_schema is not supported (no constrained decoding)",
+            f"response_format.type {rtype!r} is not supported "
+            "(supported: text, json_object, json_schema)",
             param="response_format",
         )
     if not 1 <= req.n <= MAX_N:
@@ -1072,7 +1074,11 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
     if logprobs_error is not None:
         return logprobs_error
     if _response_format_unsupported(req.response_format):
-        return "response_format json_object/json_schema is not supported (no constrained decoding)"
+        rtype = (req.response_format or {}).get("type")
+        return (
+            f"response_format.type {rtype!r} is not supported "
+            "(supported: text, json_object, json_schema)"
+        )
     return None
 
 

From cddbc9dc312e554ee745d4de0ae24fddf182e11c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= <Avicennasis@gmail.com>
Date: Sat, 3 Oct 2026 09:45:33 -0400
Subject: [PATCH 207/249] fix(mm): classify image fetch errors and re-check
 redirect targets

(cherry picked from commit 8952c053603397040c0bdf58c9d3e020b82e319e)
(cherry picked from commit 5eeb4802f3dde07fdc1b4b1c74a9b33d1726e264)
---
 docs/cli.md                   |   2 +-
 python/freetoken/mm/media.py  | 124 ++++++++++++++++++++++++++--------
 tests/server/test_mm_input.py | 121 +++++++++++++++++++++++++++++++++
 3 files changed, 219 insertions(+), 28 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index f3585435b2..ebec8b405f 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -205,7 +205,7 @@ so a client can gate its attachment controls without reading the checkpoint conf
 | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up |
 | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins |
 | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back |
-| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain |
+| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address |
 | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files |
 
 ## ft shell
diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py
index 40c782eed4..90ef790de3 100644
--- a/python/freetoken/mm/media.py
+++ b/python/freetoken/mm/media.py
@@ -1,6 +1,8 @@
 """Client-supplied image handling: ref collection, the accept gate, and byte fetching.
 
 Raises plain ValueError on bad input; the server layer maps it to its wire error type.
+The message names the class of failure only: no local path, allowlist root, errno or
+library text reaches the client (that detail goes to the server log).
 """
 
 from __future__ import annotations
@@ -8,10 +10,15 @@
 import base64
 from typing import TYPE_CHECKING, Any
 
+from freetoken.utils import init_logger
+
 if TYPE_CHECKING:
     from freetoken.server.args import ServerArgs
 
+logger = init_logger(__name__)
+
 _MAX_IMAGE_BYTES = 32 << 20
+_MAX_REDIRECTS = 5
 
 
 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
@@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None:
         )
 
 
+def _check_redirect_target(url: str, config: ServerArgs) -> None:
+    """A redirect hop is the remote server's choice, not the client's URL: it must pass the
+    allowlist like the original, and may never land on a loopback / private / link-local
+    address (an allowed public host 302-ing into the operator's network is the SSRF shape)."""
+    import ipaddress
+    from urllib.parse import urlparse
+
+    if not url.startswith(("http://", "https://")):
+        raise ValueError("the URL redirects to a non-http(s) location")
+    try:
+        _check_media_domain(url, config)
+    except ValueError:
+        raise ValueError("the URL redirects outside the allowed domains") from None
+    host = (urlparse(url).hostname or "").rstrip(".").lower()
+    if host == "localhost" or host.endswith(".localhost"):
+        raise ValueError("the URL redirects to a private or loopback address")
+    try:
+        ip = ipaddress.ip_address(host)
+    except ValueError:
+        return  # a name; the allowlist above is the gate for names (no DNS resolution here)
+    if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None:
+        ip = ip.ipv4_mapped
+    if not ip.is_global or ip.is_multicast:  # covers loopback, private, link-local, CGNAT, unspecified
+        raise ValueError("the URL redirects to a private or loopback address")
+
+
+async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes:
+    """GET an http(s) ref, following redirects one hop at a time so every target is re-checked."""
+    import httpx
+
+    _check_media_domain(url, config)
+    async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client:
+        for _ in range(_MAX_REDIRECTS + 1):
+            async with client.stream("GET", url) as resp:
+                if resp.next_request is not None:  # a 3xx with a Location: httpx built the hop
+                    url = str(resp.next_request.url)
+                    _check_redirect_target(url, config)
+                    continue
+                resp.raise_for_status()
+                declared = resp.headers.get("content-length", "")
+                if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES:
+                    raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes")
+                buf = bytearray()
+                async for chunk in resp.aiter_bytes():
+                    buf += chunk
+                    if len(buf) > _MAX_IMAGE_BYTES:
+                        raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes")
+                return bytes(buf)
+    raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times")
+
+
 def _load_local_media(url: str, config: ServerArgs) -> bytes:
     """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path."""
     from pathlib import Path
@@ -73,41 +131,53 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes:
     # resolve() follows symlinks, so a link inside the root escaping it is rejected too
     resolved = filepath.resolve()
     if Path(root).resolve() not in resolved.parents:
-        raise ValueError(
-            f"the file path {filepath} must be a subpath of "
-            f"--allowed-local-media-path {root}"
-        )
+        raise ValueError("the file path must be a subpath of --allowed-local-media-path")
     return resolved.read_bytes()
 
 
+async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes:
+    """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything
+    a library raises is reduced to its failure class and logged with the detail."""
+    data = ref.get("data") or ""
+    try:
+        if ref.get("kind") == "b64":
+            return base64.b64decode(data)
+        if data.startswith("data:"):
+            return base64.b64decode(data.split(",", 1)[1])
+    except Exception as exc:  # noqa: BLE001 -- binascii.Error / IndexError, input-driven
+        raise ValueError("image data is not valid base64") from exc
+    if data.startswith(("http://", "https://")):
+        import httpx
+
+        try:
+            return await _fetch_remote_media(data, config)
+        except ValueError:
+            raise
+        except httpx.HTTPStatusError as exc:
+            raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc
+        except httpx.TimeoutException as exc:
+            raise ValueError("timed out fetching the image URL") from exc
+        except Exception as exc:  # noqa: BLE001 -- transport / protocol / bad URL
+            logger.warning("image URL fetch failed: %r", exc)
+            raise ValueError("the image URL is unreachable") from exc
+    if data.startswith("file://"):
+        try:
+            return _load_local_media(data, config)
+        except ValueError:
+            raise
+        except OSError as exc:
+            logger.warning("local image read failed: %r", exc)
+            raise ValueError("the local file is missing or unreadable") from exc
+    raise ValueError("unsupported image source (expect http(s)/file url or base64)")
+
+
 async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]:
     """Resolve collected image refs (URLs / base64) to raw bytes, in order."""
     out: list[bytes] = []
     for ref in refs:
-        data = ref.get("data") or ""
         try:
-            if ref.get("kind") == "b64":
-                out.append(base64.b64decode(data))
-            elif data.startswith("data:"):
-                out.append(base64.b64decode(data.split(",", 1)[1]))
-            elif data.startswith(("http://", "https://")):
-                import httpx
-
-                _check_media_domain(data, config)
-                # redirect targets are not re-checked against the allowlist
-                async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
-                    resp = await client.get(data)
-                    resp.raise_for_status()
-                    if len(resp.content) > _MAX_IMAGE_BYTES:
-                        raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes")
-                    out.append(resp.content)
-            elif data.startswith("file://"):
-                out.append(_load_local_media(data, config))
-            else:
-                raise ValueError(
-                    "unsupported image source (expect http(s)/file url or base64)"
-                )
-        except Exception as exc:  # noqa: BLE001 -- input-driven, client-classifiable
+            out.append(await _fetch_one(ref, config))
+        except ValueError as exc:
             raise ValueError(f"could not load image: {exc}") from exc
     return out
 
diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py
index 1cca773361..803ac34c8f 100644
--- a/tests/server/test_mm_input.py
+++ b/tests/server/test_mm_input.py
@@ -10,6 +10,7 @@
 
 import pytest
 
+from freetoken.mm.media import _MAX_IMAGE_BYTES as _MAX
 from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason
 from freetoken.server.generation import GenerationError, render_messages
 from freetoken.server.stats import derive_model_card
@@ -215,3 +216,123 @@ def test_render_messages_rejects_an_image_in_a_merged_system_message():
             ]},
             {"role": "user", "content": "hi"},
         ])
+# ---- error bodies and redirect policy (no network: httpx.MockTransport stands in for the remote) ----
+
+
+def _mock_remote(monkeypatch, routes):
+    """Serve ``routes`` ({url: (status, headers, body) | Exception}) through the client media.py builds;
+    returns the list of URLs the transport actually saw."""
+    import httpx
+
+    seen = []
+
+    def handler(request):
+        seen.append(str(request.url))
+        route = routes[str(request.url)]
+        if isinstance(route, Exception):
+            raise route
+        status, headers, body = route
+        return httpx.Response(status, headers=headers, content=body, request=request)
+
+    real = httpx.AsyncClient
+    monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real(transport=httpx.MockTransport(handler), **kw))
+    return seen
+
+
+def _fetch(url, config):
+    return asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config))
+
+
+def test_local_media_error_bodies_carry_no_path_root_or_errno(tmp_path):
+    root = tmp_path / "media"
+    root.mkdir()
+    config = _config(allowed_local_media_path=str(root))
+    outside = tmp_path / "secret.png"
+    outside.write_bytes(b"x")
+    (root / "dir").mkdir()
+
+    cases = {
+        f"file://{outside}": "subpath",  # escapes the root
+        f"file://{root / 'missing.png'}": "missing or unreadable",  # ENOENT
+        f"file://{root / 'dir'}": "missing or unreadable",  # EISDIR
+    }
+    for url, expect in cases.items():
+        with pytest.raises(ValueError, match=expect) as ei:
+            _fetch(url, config)
+        body = str(ei.value)
+        assert str(tmp_path) not in body and "Errno" not in body, body
+        assert "secret.png" not in body and "missing.png" not in body, body
+
+
+def test_remote_fetch_failures_are_classified_without_library_text(monkeypatch):
+    import httpx
+
+    big = _MAX + 1
+    seen = _mock_remote(monkeypatch, {
+        "https://cdn.example.com/404.png": (404, {}, b""),
+        "https://cdn.example.com/down.png": httpx.ConnectError("[Errno -2] Name or service not known"),
+        "https://cdn.example.com/slow.png": httpx.ReadTimeout("timed out"),
+        "https://cdn.example.com/declared.png": (200, {"content-length": str(big)}, b""),
+        "https://cdn.example.com/streamed.png": (200, {}, b"\0" * big),
+    })
+    config = _config(allowed_media_domains="cdn.example.com")
+    for name, expect in [("404", "HTTP 404"), ("down", "unreachable"), ("slow", "timed out"),
+                         ("declared", "exceeds"), ("streamed", "exceeds")]:
+        with pytest.raises(ValueError, match=expect) as ei:
+            _fetch(f"https://cdn.example.com/{name}.png", config)
+        assert "Errno" not in str(ei.value)
+    assert len(seen) == 5
+
+
+def test_redirects_are_rechecked_against_the_allowlist(monkeypatch):
+    seen = _mock_remote(monkeypatch, {
+        "https://cdn.example.com/a.png": (302, {"location": "https://evil.com/a.png"}, b""),
+        "https://cdn.example.com/b.png": (302, {"location": "/real.png"}, b""),
+        "https://cdn.example.com/real.png": (200, {}, b"fakepng"),
+        "https://evil.com/a.png": (200, {}, b"pwned"),
+    })
+    config = _config(allowed_media_domains="cdn.example.com")
+    with pytest.raises(ValueError, match="redirects outside the allowed domains") as ei:
+        _fetch("https://cdn.example.com/a.png", config)
+    assert "evil.com" not in str(ei.value)
+    # a same-host relative redirect is followed and served
+    assert _fetch("https://cdn.example.com/b.png", config) == [b"fakepng"]
+    # the disallowed hop was never requested
+    assert seen == ["https://cdn.example.com/a.png", "https://cdn.example.com/b.png", "https://cdn.example.com/real.png"]
+
+
+@pytest.mark.parametrize("target", [
+    "http://127.0.0.1/x", "http://10.0.0.1/x", "http://169.254.169.254/latest/meta-data",
+    "http://[::1]/x", "http://[::ffff:10.0.0.1]/x", "http://localhost/x", "http://0.0.0.0/x",
+    "http://100.64.0.1/x",
+])
+def test_redirects_to_private_or_loopback_addresses_are_refused(monkeypatch, target):
+    seen = _mock_remote(monkeypatch, {
+        "https://public.example.com/a.png": (302, {"location": target}, b""),
+        target: (200, {}, b"internal"),
+    })
+    # even with an empty allowlist (any domain admitted) a redirect may not reach inside
+    with pytest.raises(ValueError, match="private or loopback"):
+        _fetch("https://public.example.com/a.png", _config())
+    assert seen == ["https://public.example.com/a.png"]
+
+
+def test_redirect_hop_count_and_scheme_are_bounded(monkeypatch):
+    chain = {f"https://h.example.com/{i}.png": (302, {"location": f"https://h.example.com/{i + 1}.png"}, b"")
+             for i in range(7)}
+    chain["https://h.example.com/7.png"] = (200, {}, b"fakepng")
+    chain["https://h.example.com/file.png"] = (302, {"location": "file:///etc/passwd"}, b"")
+    seen = _mock_remote(monkeypatch, chain)
+    with pytest.raises(ValueError, match="redirects more than"):
+        _fetch("https://h.example.com/0.png", _config())
+    assert len(seen) == 6  # the original plus five hops; the sixth hop is never requested
+    with pytest.raises(ValueError, match="non-http"):
+        _fetch("https://h.example.com/file.png", _config())
+    # a short chain that stays public is fine
+    assert _fetch("https://h.example.com/5.png", _config()) == [b"fakepng"]
+
+
+def test_bad_base64_is_named_without_decoder_text():
+    for ref in ({"kind": "b64", "data": "not base64!!"}, {"kind": "url", "data": "data:image/png;base64"}):
+        with pytest.raises(ValueError, match="not valid base64"):
+            asyncio.run(fetch_image_bytes([ref], _config()))

From 5e8a24b9c00c68b65424eadbbe037596ec8356b8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 06:06:12 -0400
Subject: [PATCH 208/249] test(engine): cover presence/frequency penalties on
 the greedy and sampled paths

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 30b818114ce353da17bf6c3c9933e6c752eb7c11)
(cherry picked from commit dc7cfc13fa5eaef06e105530b55e752529c9d0c6)
---
 tests/engine/test_sample.py | 84 +++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 tests/engine/test_sample.py

diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py
new file mode 100644
index 0000000000..c1cc925f00
--- /dev/null
+++ b/tests/engine/test_sample.py
@@ -0,0 +1,84 @@
+"""Presence/frequency penalties in engine.sample.Sampler (greedy and sampled rows)."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.core import SamplingParams
+from freetoken.engine import sample as S
+
+
+def _batch(*params: SamplingParams, can_decode: bool = True):
+    reqs = [
+        SimpleNamespace(sampling_params=p, can_decode=can_decode, output_token_counts=None)
+        for p in params
+    ]
+    return SimpleNamespace(reqs=reqs)
+
+
+def test_no_penalty_allocates_nothing():
+    sampler = S.Sampler(torch.device("cpu"), 5)
+    batch = _batch(SamplingParams(temperature=0.0))
+    args = sampler.prepare(batch)
+    assert args.penalties == [] and batch.reqs[0].output_token_counts is None
+
+
+def test_frequency_penalty_counts_generated_tokens_on_the_greedy_path():
+    sampler = S.Sampler(torch.device("cpu"), 5)
+    batch = _batch(SamplingParams(temperature=0.0, frequency_penalty=10.0))
+    logits = torch.tensor([[0.0, 5.0, 1.0, 0.0, 0.0]])
+
+    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1]
+    # Token 1 now carries -10, so the runner-up wins; the caller's logits stay untouched.
+    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2]
+    assert batch.reqs[0].output_token_counts.tolist() == [0, 1, 1, 0, 0]
+    assert logits[0, 1] == 5.0
+
+
+def test_intermediate_prefill_chunks_are_not_counted():
+    sampler = S.Sampler(torch.device("cpu"), 5)
+    batch = _batch(SamplingParams(temperature=0.0, presence_penalty=1.0), can_decode=False)
+    assert sampler.prepare(batch).penalties == []
+
+
+def test_penalized_logits_reach_the_probability_sampler(monkeypatch):
+    """Sampled rows (the torch top-k/top-p fallback on Pascal) see the penalized logits."""
+    seen = {}
+
+    def fake_sample_impl(logits, temperatures, top_k, top_p):
+        seen["logits"] = logits.clone()
+        return torch.argmax(logits, dim=-1)
+
+    monkeypatch.setattr(S, "sample_impl", fake_sample_impl)
+    monkeypatch.setattr(S, "make_device_tensor", lambda data, dtype, device: torch.tensor(data, dtype=dtype))
+    sampler = S.Sampler(torch.device("cpu"), 4)
+    batch = _batch(
+        SamplingParams(temperature=1.0, top_k=2, presence_penalty=3.0),
+        SamplingParams(temperature=1.0, top_k=2),
+    )
+    batch.reqs[0].output_token_counts = torch.tensor([0, 2, 0, 0], dtype=torch.int32)
+    logits = torch.tensor([[0.0, 4.0, 2.0, 0.0], [0.0, 4.0, 2.0, 0.0]])
+
+    tokens = sampler.sample(logits, sampler.prepare(batch))
+
+    assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0]  # presence: -3 once, not per count
+    assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0]
+    assert tokens.tolist() == [2, 1]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch):
+    import freetoken.kernel.backend as backend
+    import freetoken.utils.arch as arch
+
+    monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False)
+    monkeypatch.setattr(arch, "is_sm70_supported", lambda: False)
+    sampler = S.Sampler(torch.device("cuda"), 4)
+    batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0))
+    logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda")
+
+    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1]
+    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2]

From faa8116f8317d40ec6bbd28fc14764cdc378dd66 Mon Sep 17 00:00:00 2001
From: "Eva (port)" <eva@local>
Date: Sat, 3 Oct 2026 01:21:44 +0300
Subject: [PATCH 209/249] perf(model): fuse hc combine with the next block's
 RMSNorm via hc_combine_norm

the vendored combine+norm kernel was exported but unwired; mix() takes a
pre-computed rn, the model loop carries the fused pair across the layer
boundary, skipped before PLE layers where ple(hidden) lands between.

(cherry picked from commit 7636dd76cbe64989cbdec09e5c09a4c67be89542)
(cherry picked from commit 0e847f8cd849e4a75cab781a7846af972afbbc58)
---
 python/freetoken/models/qwen4_exp/hc.py    | 30 +++++++++++++++++-----
 python/freetoken/models/qwen4_exp/model.py | 30 ++++++++++++++++------
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py
index 69aaa58302..2071916db0 100644
--- a/python/freetoken/models/qwen4_exp/hc.py
+++ b/python/freetoken/models/qwen4_exp/hc.py
@@ -17,6 +17,7 @@
 from freetoken.kernel.triton.hc import (
     grouped_gemma_rmsnorm,
     hc_combine,
+    hc_combine_norm,
     hc_gate_mix,
     hc_silu,
 )
@@ -120,14 +121,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]:
         # both slices keep unit inner stride, so the kernels read them without a copy
         return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count]
 
-    def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]:
-        rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count)
+    def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        if rn is None:
+            rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count)
         lora, s = self._down(rn)
         gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count))
         return hc_gate_mix(rn, gate, self.hc_count), s
 
-    def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]:
-        rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count)
+    def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        if rn is None:
+            rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count)
         lora, s = self._down(rn)
         lora = F.silu(lora.float() / self.hc_count)
         gate = self.input_mix_weight_up.forward(lora.to(R.dtype))
@@ -141,9 +144,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t
         out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1)
         return out.flatten(-2).to(R.dtype)
 
-    def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]:
-        """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine)."""
-        return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R)
+    def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).
+
+        ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as
+        this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it."""
+        return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn)
 
     def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
         """Inject the block output ``y [T, hidden]`` back into every stream of ``R``."""
@@ -151,5 +157,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te
             return hc_combine(R, y, s, self.hc_count)
         return self._combine_torch(R, y, s)
 
+    def combine_norm(
+        self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual"
+    ) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None)."""
+        if not R.is_cuda:
+            return self.combine(R, y, s), None
+        return hc_combine_norm(
+            R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count
+        )
+
 
 __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"]
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index fdef4298ba..63fa971745 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -76,17 +76,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
         )
 
     @nvtx_annotate("Layer_{}", layer_id_field="_layer_id")
-    def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor:
+    def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None):
+        """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the
+        next consumer's norm (next layer's attn mix, or the final mixer)."""
         if self.ple is not None:
             hidden = hidden + self.ple.forward(hidden, batch)
-        block_input, inject = self.attn_hyper_connection.mix(hidden)
+        block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn)
         if self._is_linear:
             block_output = self.linear_attn.forward(block_input)
         else:
             block_output = self.self_attn.forward(block_input, batch)
-        hidden = self.attn_hyper_connection.combine(hidden, block_output, inject)
-        block_input, inject = self.mlp_hyper_connection.mix(hidden)
-        return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject)
+        hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection)
+        block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp)
+        return hidden, self.mlp.forward(block_input), inject
 
 
 class Qwen4ExpModel(BaseOP):
@@ -124,13 +126,25 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor:
             meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device)
             for ple in self._ple:  # gather the pinned-host PLE rows while the early layers run
                 ple.start_prefetch(batch, meta)
-        for layer in self.layers.op_list:
-            hidden = layer.forward(hidden, batch)
+        # The last block of each layer combines through the layer boundary, fusing the
+        # combine with the next consumer's hc_norm -- except before a PLE layer, whose
+        # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees.
+        layers = self.layers.op_list
+        rn = None
+        for i, layer in enumerate(layers):
+            hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn)
+            nxt = layers[i + 1] if i + 1 < len(layers) else None
+            if nxt is not None and nxt.ple is not None:
+                hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject)
+                rn = None
+            else:
+                consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer
+                hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer)
         if meta is not None:
             # single writer: the layers only read the context, so a second PLE layer's
             # prefetch sees the un-rolled window
             commit_ngram_context(meta, getattr(batch, "fla_metadata", None))
-        return self.hyper_connection_mixer.mix(hidden)[0]
+        return self.hyper_connection_mixer.mix(hidden, rn=rn)[0]
 
 
 class Qwen4ExpForCausalLM(BaseLLMModel):

From e89e46e6d236b304acdcadbbc80aa2d3e6de9273 Mon Sep 17 00:00:00 2001
From: "Eva (port)" <eva@local>
Date: Sat, 3 Oct 2026 01:21:44 +0300
Subject: [PATCH 210/249] fix(kernel): bit-exact rmsnorm reduction tree in
 hc_combine_norm

the fused kernel summed squares over a two-axis padded tile; the standalone
grouped RMSNorm reduces one flat next_pow2 vector - different fp32 tree,
ulp-level drift caught by the greedy gate. Reshape to the identical tree.

(cherry picked from commit c1c5e76f5d98e08f0ff525fe71415bac028af3bc)
(cherry picked from commit 8f8d7fbd7f0940e6ea4f0fa1b5ddfee4892efe0f)
---
 python/freetoken/kernel/triton/hc.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py
index 45d71bd731..acae534c45 100644
--- a/python/freetoken/kernel/triton/hc.py
+++ b/python/freetoken/kernel/triton/hc.py
@@ -326,9 +326,11 @@ def _hc_combine_norm_kernel(
     tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner)
 
     out = out.to(tl.float32)
-    # Keep the two-axis reduction: flattening the padded tile is ~40% slower
-    # at decode sizes.
-    sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0)
+    # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)]
+    # vector; reshape keeps the identical tree (same element
+    # order, same masked zeros), so the fused statistic is bit-exact vs the split.
+    flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE])
+    sum_sq = tl.sum(flat)
     rrms = tl.rsqrt(sum_sq / HC_DIM + EPS)
 
     if launch_pdl:

From 34de9a96afe6aec310ad2f5943f511dbbcb0b9a9 Mon Sep 17 00:00:00 2001
From: "Eva (port)" <eva@local>
Date: Sat, 3 Oct 2026 01:22:10 +0300
Subject: [PATCH 211/249] test(kernel): bitwise pin for hc_combine_norm against
 the split path

(cherry picked from commit 7cba55ea8ba0d3576523c75ef0fcacb3d3e817bb)
(cherry picked from commit 081534bb5574c7074ff37a2ecb9dca5a645073f1)
---
 tests/kernels/test_hc_combine_norm.py | 48 +++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 tests/kernels/test_hc_combine_norm.py

diff --git a/tests/kernels/test_hc_combine_norm.py b/tests/kernels/test_hc_combine_norm.py
new file mode 100644
index 0000000000..7f46755479
--- /dev/null
+++ b/tests/kernels/test_hc_combine_norm.py
@@ -0,0 +1,48 @@
+"""Pins for the fused hc combine + grouped RMSNorm path (hc_combine_norm).
+
+hc_combine_norm replaces the standalone pair (hc_combine -> grouped_gemma_rmsnorm)
+in the Qwen4-exp decoder loop. The swap is only sound if it is BITWISE identical:
+the hyper-connection streams feed every later layer, and a bf16 flip here compounds
+into different greedy tokens by the end of the stack. We assert torch.equal, not
+assert_close: an approx-tolerance test would not have caught the regression this
+path shipped with at first attempt -- the fused kernel reduced squares over the
+[8, block] tile view (a two-axis tree) while the standalone kernel reduces the
+flat row tree; fp32 ulps that differ between trees flip bf16 rounding. Both paths
+on the same data must return the same bytes: fails before the reduction-tree fix
+in _hc_combine_norm_kernel, passes after.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kernel.triton.hc import hc_combine, grouped_gemma_rmsnorm, hc_combine_norm
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+HC, H = 4, 2560  # model geometry: the reduction tree specializes on the block layout
+DIM = HC * H
+EPS = 1e-6
+
+
+def _inputs(seed: int, T: int):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    bf = lambda shape: torch.randn(shape, generator=g, device="cuda", dtype=torch.float32).to(torch.bfloat16)
+    return (
+        bf((T, DIM)),  # residual streams R
+        bf((T, H)),  # block output y
+        bf((T, HC)),  # inject scores s
+        (torch.randn(DIM, generator=g, device="cuda", dtype=torch.float32) * 0.1).to(torch.bfloat16),  # norm weight
+    )
+
+
+@pytest.mark.parametrize("T", [1, 8])
+@pytest.mark.parametrize("seed", range(6))
+def test_combine_norm_matches_split_bitwise(seed: int, T: int) -> None:
+    R, y, s, w = _inputs(seed, T)
+    r2_split = hc_combine(R, y, s, HC)
+    rn_split = grouped_gemma_rmsnorm(r2_split, w, EPS, HC)
+    r2_fused, rn_fused = hc_combine_norm(R, y, s, w, EPS, HC)
+    assert torch.equal(r2_fused, r2_split)
+    assert torch.equal(rn_fused, rn_split)

From 9f46531d48e4232bf642415dbdea168f658a422c Mon Sep 17 00:00:00 2001
From: Siarhei Berdachuk <berdachuk@berdaflex.com>
Date: Mon, 5 Oct 2026 19:11:42 +0300
Subject: [PATCH 212/249] fix(server): soft-accept OpenAI response_format
 json_object/json_schema

Rejecting response_format with "no constrained decoding" breaks OpenAI-compatible
clients that rely on json_object/json_schema even when the model can emit valid
JSON from a prompt. Accept those wire types and prepend a best-effort system hint;
schema is not token-level enforced. Unknown response_format types still return 400.

(cherry picked from commit ca7d704bb84e255d8677b788f295e9cb6ebb06f0)
(cherry picked from commit db64d8eff1a7778b6bc093634303d31e1324af76)
---
 python/freetoken/server/openai_api.py | 44 +++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 79dff72963..7a345cadfa 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -71,6 +71,43 @@ def _thinking_type(req: Any) -> str | None:
 
 
 
+def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None:
+    """Best-effort JSON mode hint for OpenAI response_format.
+
+    FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients
+    still send response_format=json_object|json_schema. Rejecting them with 400
+    breaks those clients even though the model can emit valid JSON from a
+    prompt. Accept the wire shape and prepend a system reminder so the model is
+    steered toward JSON; the schema is NOT token-level enforced.
+    """
+    if not response_format:
+        return None
+    rtype = response_format.get("type")
+    if rtype == "json_object":
+        return {
+            "role": "system",
+            "content": (
+                "Respond with a single valid JSON object only. "
+                "Do not wrap it in markdown fences or add commentary."
+            ),
+        }
+    if rtype == "json_schema":
+        wrapper = response_format.get("json_schema") or {}
+        name = wrapper.get("name") or "response"
+        schema = wrapper.get("schema", wrapper)
+        return {
+            "role": "system",
+            "content": (
+                f"Respond with a single JSON object that validates against the "
+                f"JSON Schema below (schema name: '{name}'). Put required fields "
+                "at the top level of the object — do not wrap them under an extra "
+                "key named after the schema. Do not wrap the answer in markdown "
+                f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}"
+            ),
+        }
+    return None
+
+
 def chat_request_to_genspec(
     req: ChatCompletionRequest,
     model_sampling: dict[str, Any],
@@ -1063,8 +1100,11 @@ def _usage(
 
 
 def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool:
-    # We have no constrained/guided decoding; only plain text ('text' or unset) is honored.
-    return response_format is not None and response_format.get("type") not in (None, "text")
+    # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding).
+    # Unknown types still 400. See _json_mode_system_message.
+    if response_format is None:
+        return False
+    return response_format.get("type") not in (None, "text", "json_object", "json_schema")
 
 
 def _completion_unsupported_reason(req: CompletionRequest) -> str | None:

From 8248a8f18df437a5b1223dc53fe31b19469ae0db Mon Sep 17 00:00:00 2001
From: Siarhei Berdachuk <berdachuk@berdaflex.com>
Date: Mon, 5 Oct 2026 19:18:06 +0300
Subject: [PATCH 213/249] refactor(server): drop response_format system-hint
 injection

Soft-accepting json_object/json_schema is enough when clients already
request JSON in the prompt. The injected schema wrapper was redundant
and could distort the output shape.

(cherry picked from commit 65907c8bc2c1c11a4af11dff328f7cc293806433)
(cherry picked from commit c82625f6eed239a9edcc6e76e935d3b51359f425)
---
 python/freetoken/server/openai_api.py | 45 ++-------------------------
 1 file changed, 3 insertions(+), 42 deletions(-)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 7a345cadfa..0fc8113bc8 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -68,46 +68,6 @@ def _thinking_type(req: Any) -> str | None:
         if value in ("enabled", "disabled"):
             return value
     return None
-
-
-
-def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None:
-    """Best-effort JSON mode hint for OpenAI response_format.
-
-    FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients
-    still send response_format=json_object|json_schema. Rejecting them with 400
-    breaks those clients even though the model can emit valid JSON from a
-    prompt. Accept the wire shape and prepend a system reminder so the model is
-    steered toward JSON; the schema is NOT token-level enforced.
-    """
-    if not response_format:
-        return None
-    rtype = response_format.get("type")
-    if rtype == "json_object":
-        return {
-            "role": "system",
-            "content": (
-                "Respond with a single valid JSON object only. "
-                "Do not wrap it in markdown fences or add commentary."
-            ),
-        }
-    if rtype == "json_schema":
-        wrapper = response_format.get("json_schema") or {}
-        name = wrapper.get("name") or "response"
-        schema = wrapper.get("schema", wrapper)
-        return {
-            "role": "system",
-            "content": (
-                f"Respond with a single JSON object that validates against the "
-                f"JSON Schema below (schema name: '{name}'). Put required fields "
-                "at the top level of the object — do not wrap them under an extra "
-                "key named after the schema. Do not wrap the answer in markdown "
-                f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}"
-            ),
-        }
-    return None
-
-
 def chat_request_to_genspec(
     req: ChatCompletionRequest,
     model_sampling: dict[str, Any],
@@ -1100,8 +1060,9 @@ def _usage(
 
 
 def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool:
-    # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding).
-    # Unknown types still 400. See _json_mode_system_message.
+    # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained
+    # decoding; clients that already ask for JSON in the prompt still need the
+    # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400.
     if response_format is None:
         return False
     return response_format.get("type") not in (None, "text", "json_object", "json_schema")

From 5f7b3e41fd2ba9fe2a2a380e9ba1b834a77c7805 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 13:14:39 -0400
Subject: [PATCH 214/249] test(engine): drop the pre-volta sampling penalty
 test, whose torch fallback this branch does not carry

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 tests/engine/test_sample.py | 15 ---------------
 1 file changed, 15 deletions(-)

diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py
index c1cc925f00..28e7195b58 100644
--- a/tests/engine/test_sample.py
+++ b/tests/engine/test_sample.py
@@ -67,18 +67,3 @@ def fake_sample_impl(logits, temperatures, top_k, top_p):
     assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0]  # presence: -3 once, not per count
     assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0]
     assert tokens.tolist() == [2, 1]
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch):
-    import freetoken.kernel.backend as backend
-    import freetoken.utils.arch as arch
-
-    monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False)
-    monkeypatch.setattr(arch, "is_sm70_supported", lambda: False)
-    sampler = S.Sampler(torch.device("cuda"), 4)
-    batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0))
-    logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda")
-
-    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1]
-    assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2]

From 675d110c26b314672db3ccf75d2f42f831ca95be Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 13 Oct 2026 13:14:39 -0400
Subject: [PATCH 215/249] fix(moe): align on triton past sgl_kernel's
 1024-expert limit

sgl_kernel's moe_align_block_size (0.4.5 and 0.4.9) reads out of bounds once num_experts + 1
exceeds 1024: an illegal memory access that poisons the CUDA context. A prefill over the
offload slot cache (FlashML-org/FreeToken#676) aligns against the slot count, 2100 here, so
the first short prefill crashed the server; the bf16 and NVFP4 prefill paths both go through
it. The Triton align has no such limit and matches sgl_kernel's output below it.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/moe/fused.py    |  8 +++++++-
 tests/moe/test_nvfp4_backends.py | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py
index 5c34fb33ca..2f8af47c89 100644
--- a/python/freetoken/moe/fused.py
+++ b/python/freetoken/moe/fused.py
@@ -44,6 +44,9 @@ def fused_topk(
     return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded)
 
 
+_SGL_ALIGN_MAX_EXPERTS = 1024
+
+
 def moe_align_block_size(
     topk_ids: torch.Tensor, block_size: int, num_experts: int
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -97,7 +100,10 @@ def moe_align_block_size(
 
             return moe_align_block_size_triton(topk_ids, block_size, num_experts)
 
-    if not is_sgl_kernel_installed():
+    # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of
+    # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the
+    # slot count here (#676), which is thousands; the Triton kernel has no such limit.
+    if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS:
         from freetoken.kernel.triton.moe_align import (
             moe_align_block_size as triton_moe_align_block_size,
         )
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index bb7754206d..ecc401e425 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -710,3 +710,22 @@ def test_b12x_pack_keeps_per_layer_banks_and_flat_alphas():
     assert len(cache.bank_sources["gate_up"]) == L
     assert sum(t.shape[0] for t in cache.bank_sources["gate_up"]) == total
     assert cache.gate_up_alpha.shape == (total,)
+
+
+@cuda
+@pytest.mark.parametrize("num_experts", [1023, 2100])
+def test_moe_align_handles_a_slot_cache_wider_than_1024(num_experts):
+    """A prefill over the offload slot cache aligns against its slot count (#676): sgl_kernel's
+    align reads out of bounds past 1024 counters, so a 2100-slot cache must still align."""
+    from freetoken.moe.fused import moe_align_block_size
+
+    torch.manual_seed(0)
+    ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32)
+    sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts)
+    n = int(padded)
+    flat = ids.flatten()
+    real = sorted_ids[:n][sorted_ids[:n] < flat.numel()]
+    assert sorted(real.tolist()) == list(range(flat.numel()))
+    block_of = torch.arange(n, device="cuda") // 16
+    rows = sorted_ids[:n] < flat.numel()
+    assert torch.equal(flat[sorted_ids[:n][rows].long()], expert_ids[block_of[rows]])

From f4f2d38d3e31fdbe427d59670c0eddba55df4aca Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 08:40:45 -0400
Subject: [PATCH 216/249] fix(moe): refuse the global-slot moe features when
 the cache has several slot pools

Upstream #689 splits the offload expert cache into slot pools whose slot ids are
pool-local. The scored eviction policies, expert prefetch, its verify and slot-audit
debug modes and the #676 on-demand prefill GEMM all address the cache by global slot
over one typed bank, so with several pools they would install or read experts in
another pool's rows. Refuse the combination at construction; one pool (the default,
and this deployment) behaves as before.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/moe/offload_cache.py | 24 ++++++++++++++++++++++++
 tests/moe/test_offload.py             | 23 +++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index fa4b7b823f..eb8547b331 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -389,6 +389,30 @@ def __post_init__(self) -> None:
         self.prefetch = self._init_prefetch()
         self.verify = self._init_verify()
         self.audit = self._init_audit()
+        self._refuse_single_pool_features()
+
+    def _refuse_single_pool_features(self) -> None:
+        """The scored policies, prefetch, its debug checks and the on-demand prefill GEMM address
+        the cache by global slot id over one typed bank, so with several slot pools (#689) they
+        would install and read pool-local slots in the wrong rows: refuse that combination."""
+        if len(self.pool_sizes) == 1:
+            return
+        used = [
+            name
+            for name, on in (
+                (f"--moe-cache-policy {self.cache_policy}", self.cache_policy != "lru"),
+                ("FREETOKEN_MOE_PREFETCH", self.prefetch_mode not in (None, "off")),
+                ("FREETOKEN_MOE_PREFETCH_VERIFY", self.verify_mode not in (None, "off")),
+                ("FREETOKEN_MOE_SLOT_AUDIT", bool(self.slot_audit)),
+                ("--moe-prefill-demand-tokens", bool(self.prefill_demand_tokens)),
+            )
+            if on
+        ]
+        if used:
+            raise ValueError(
+                f"{', '.join(used)} need one MoE slot pool, but this cache has {len(self.pool_sizes)} "
+                "(--moe-cache-scope per-layer, or experts in several quant formats)"
+            )
 
     def _init_audit(self):
         from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 764a8d0083..34add58c13 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -1360,3 +1360,26 @@ def boom(addr, nbytes):
         with hb.PinPipeline() as pins:
             pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)})
     assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value}
+
+
+@pytest.mark.parametrize(
+    "kwargs, flag",
+    [
+        ({"cache_policy": "rule"}, "--moe-cache-policy rule"),
+        ({"prefetch_mode": "on"}, "FREETOKEN_MOE_PREFETCH"),
+        ({"verify_mode": "meta"}, "FREETOKEN_MOE_PREFETCH_VERIFY"),
+        ({"slot_audit": 8}, "FREETOKEN_MOE_SLOT_AUDIT"),
+        ({"prefill_demand_tokens": 64}, "--moe-prefill-demand-tokens"),
+    ],
+)
+def test_offload_cache_refuses_global_slot_features_with_several_pools(kwargs, flag):
+    # these address slots globally over one typed bank: with pool-local slot ids they would
+    # install and read experts in another pool's rows
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    common = dict(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), **kwargs)
+    with pytest.raises(ValueError, match=f"{flag}.*one MoE slot pool"):
+        OffloadMoeCache(pool_of_layer=[0, 1], pool_sizes=[8, 8], **common)
+    OffloadMoeCache(**common)
+    OffloadMoeCache(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), pool_of_layer=[0, 1], pool_sizes=[8, 8])

From 160f3ab2f0b20a2484a8167c6a544982740a71b7 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 08:41:49 -0400
Subject: [PATCH 217/249] test(moe): restore the #676 merge's on-demand test
 updates the replay dropped

The 5080deploy merge of #676 updated these tests inside the merge commit (the
on-demand path runs the prefill GEMM over the slot cache, and the layer hands
ensure_experts unique ids). Re-merging #676's head on the new base does not carry
edits made inside a merge, so they are restored here.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 tests/moe/test_offload.py                 | 9 +++++----
 tests/moe/test_small_prefill_on_demand.py | 4 ++--
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 34add58c13..fbacc8e5fa 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -138,9 +138,10 @@ def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cach
     )
     monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id))
 
-    def ensure(layer_id, expert_ids, *, counts_as_decode=True):
+    def ensure(layer_id, expert_ids, *, counts_as_decode=True, **kwargs):
+        # the layer hands the unique routed ids (decode-width chunks); rewrite each to a slot
         calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode)
-        expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32))
+        expert_ids.add_(3)
 
     monkeypatch.setattr(cache, "ensure_experts", ensure)
     monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True))
@@ -157,11 +158,11 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation
 
     assert out is hidden_states
     assert "materialized" not in calls
-    assert calls["ensured"] == (0, [[2, 1]], False)
+    assert calls["ensured"] == (0, [1, 2], False)
     assert calls["copied"] is True
     # the grouped GEMM reads the whole slot cache at the slots the admission wrote
     assert calls["w1"].shape[0] == cache.cache_size
-    assert calls["topk_ids"].tolist() == [[5, 3]]
+    assert calls["topk_ids"].tolist() == [[5, 4]]
 
     calls.clear()
     layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4))
diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py
index 8b7b9d802c..f7971c67d5 100644
--- a/tests/moe/test_small_prefill_on_demand.py
+++ b/tests/moe/test_small_prefill_on_demand.py
@@ -51,7 +51,7 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill):
     hs = torch.randn(tokens, 16)
     layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32))
     if on_demand:
-        assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False}
+        assert calls == ["ensure", "copy"] and got == {"n": cache.cache_size, "is_prefill": True}
     else:
         assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts
 
@@ -97,7 +97,7 @@ def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_t
     ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32)
     layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone())
     assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)]
-    assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False}
+    assert got == {"ids": (ids + 100).tolist(), "n": cache.cache_size, "is_prefill": True}
 
 
 def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch):

From b282a0acef0796ada279144a32e83a9a22f00ad7 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 08:44:36 -0400
Subject: [PATCH 218/249] fix(moe): keep bank_views for the one-pool slot cache

Upstream #689 replaced OffloadMoeCache.bank_views with per-pool and staging views.
The #676 on-demand prefill GEMM runs over the whole slot cache (n=cache_size) and
still calls bank_views, so the first short prefill raised AttributeError. Restore it
for a cache with one pool and one bank layout, where a slot id is its bank row; the
pool guard already refuses on-demand prefill with several pools.
tests/moe/test_small_prefill_on_demand.py covers the call.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/moe/offload_cache.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index eb8547b331..99d799aabb 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -578,6 +578,16 @@ def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]:
     def pool_slots(self, layer_id: int) -> int:
         return self.pool_sizes[self.pool_of_layer[layer_id]]
 
+    def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]:
+        """Per-bank views of the whole ``[S]`` slot cache, or of its first ``n`` slots, for a cache
+        with one pool and one bank layout, where a slot id is its bank row (the on-demand prefill
+        GEMM over the slot cache and the fork's slot checks)."""
+        assert self.banks, "set_bank_sources must register the banks first"
+        assert len(self.pool_sizes) == 1 and not self._raw_banks, "bank_views needs one slot pool and one bank layout"
+        if n is None:
+            return tuple(cache for _, cache in self.banks)
+        return tuple(cache[:n] for _, cache in self.banks)
+
     def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]:
         """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids."""
         assert self.banks, "set_bank_sources must register the banks first"

From e0023e515b9726af5399063ee1e1abdf23a73d01 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 08:44:37 -0400
Subject: [PATCH 219/249] test: give the rebuild and scored-ensure fakes the
 slot-pool api of #689

#650's engine fake and the decode-forward fake predate upstream #689: the rebuild now
splits the MoE target into pool sizes (decode_target, _moe_pools, _moe_pool_bytes,
validate_rebuild/rebuild take the sizes) and decode reads pool_bank_views.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 tests/engine/test_cache_budget.py                   |  2 +-
 tests/engine/test_rebuild_runtime_cache_ordering.py | 12 ++++++++++--
 tests/moe/test_scored_ensure.py                     |  2 +-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 29603f2a83..1a7d6d66ed 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -428,7 +428,7 @@ def validate_rebuild(config, **kwargs):
         kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild),
         _baseline_free=1 << 30, _weights_bytes=0, num_pages=64,
     )
-    engine._target_moe_and_expert_bytes = lambda size: (0, 0)
+    engine._target_moe_and_expert_bytes = lambda size, sizes=None: (0, 0)
     with pytest.raises(CacheRebuildRejected):
         Engine.rebuild_runtime_cache(engine, num_pages=32)
 
diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py
index aae7d49cf0..ab01f8a18f 100644
--- a/tests/engine/test_rebuild_runtime_cache_ordering.py
+++ b/tests/engine/test_rebuild_runtime_cache_ordering.py
@@ -54,14 +54,15 @@ def __init__(self, ledger: _Ledger, size: int):
         # one uint8 per expert row -> expert_bytes_per_slot == 1 unit
         self.bank_sources = {"w13": [torch.zeros((self.num_experts, 1), dtype=torch.uint8)]}
         self.cache_size = size
+        self.decode_target = "gpu"
         self._held = size
         ledger.alloc("moe", size)
 
-    def validate_rebuild(self, cache_size: int) -> None:
+    def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         if cache_size < self.num_experts:
             raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}")
 
-    def rebuild(self, cache_size: int) -> None:
+    def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         self.validate_rebuild(cache_size)
         self.ledger.free("moe", self._held)
         self._held = 0
@@ -140,6 +141,13 @@ def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: in
         graph_runner=_FakeGraphRunner(),
         _sync_get_memory=lambda: (0, 0),
     )
+    # one MoE slot pool (#689): the rebuild splits the target into its pool sizes first
+    from freetoken.moe.pool_policy import make_pools
+
+    eng._moe_pools = lambda batch: make_pools([[0]], 8, 1, None)
+    eng._moe_expert_counts = None
+    eng._moe_pool_groups, eng._moe_layer_slot_bytes = [[0]], [1]  # one uint8 per expert row
+    eng._moe_pool_bytes = Engine._moe_pool_bytes.__get__(eng)
     eng._target_moe_and_expert_bytes = Engine._target_moe_and_expert_bytes.__get__(eng)
     eng._resize_kv_pool = Engine._resize_kv_pool.__get__(eng)
     eng._refresh_seq_state = Engine._refresh_seq_state.__get__(eng)
diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py
index 6bfa84a5ed..21318c257b 100644
--- a/tests/moe/test_scored_ensure.py
+++ b/tests/moe/test_scored_ensure.py
@@ -570,7 +570,7 @@ def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy):
     )
     monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw))
     monkeypatch.setattr(cache, "copy_missing", lambda: None)
-    monkeypatch.setattr(cache, "bank_views", lambda: ())
+    monkeypatch.setattr(cache, "pool_bank_views", lambda layer_id: ())
     monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None)
     router_logits = torch.randn(2, 8)
     layer.decode_forward(torch.randn(2, 8), router_logits)

From e41664a47324ffbedb0e6617df26e31fb586f981 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 08:58:40 -0400
Subject: [PATCH 220/249] fix(moe): index the per-layout copy row bytes in the
 verify and audit gathers

Upstream #689 keeps one fused-copy row-byte descriptor per bank layout
(_copy_feat_bytes[layout]). The prefetch verify and slot-audit gathers still passed
the whole list, so every enabled check failed at launch (TypeError: expected
DLTensor*, got ffi.Array). Pass the layer's layout entry. The 83 failing
test_skeleton/test_prefetch audit and verify tests pass again (385 passed).

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/moe/slot_audit.py | 2 +-
 python/freetoken/moe/verify.py     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py
index 981239758f..12d9f40ad9 100644
--- a/python/freetoken/moe/slot_audit.py
+++ b/python/freetoken/moe/slot_audit.py
@@ -493,7 +493,7 @@ def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> N
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
 
             fast_index_copy_multi_jit(
-                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, rows, num,
+                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, rows, num,
             )
             return
         from freetoken.kernel import fast_index_copy_jit
diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py
index 1f096c9c3a..cc71e19d9d 100644
--- a/python/freetoken/moe/verify.py
+++ b/python/freetoken/moe/verify.py
@@ -401,7 +401,7 @@ def _gather(self, cache, layer: int) -> None:
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
 
             fast_index_copy_multi_jit(
-                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, self.u_ids, self.u_num,
+                self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, self.u_ids, self.u_num,
             )
             return
         from freetoken.kernel import fast_index_copy_jit

From 8eeced1cd925aa5010fd29d69bb81a58fed8e268 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Fri, 25 Sep 2026 10:56:34 -0400
Subject: [PATCH 221/249] test(kernel): cover the device-sized autotune flush
 buffer

Assisted-by: Claude Opus 5.5 (1M context)
(cherry picked from commit 99c07575e037d41cdc0c16a7a993aee5566f70ba)
---
 tests/kernels/test_autotune_cache.py | 29 ++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 tests/kernels/test_autotune_cache.py

diff --git a/tests/kernels/test_autotune_cache.py b/tests/kernels/test_autotune_cache.py
new file mode 100644
index 0000000000..3416117f79
--- /dev/null
+++ b/tests/kernels/test_autotune_cache.py
@@ -0,0 +1,29 @@
+"""The autotune L2-flush buffer is sized to the device instead of triton's flat 256 MB."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+MB = 1 << 20
+
+
+@pytest.mark.parametrize("l2_bytes, flush_bytes", [
+    (2816 * 1024, 11 * MB),  # GTX 1080 Ti: 4x its 2.75 MB L2
+    (0, 8 * MB),  # no L2 size reported: the floor
+    (128 * MB, 256 * MB),  # a large L2: triton's own 256 MB is the cap
+])
+def test_flush_buffer_is_four_times_the_l2_within_bounds(monkeypatch, l2_bytes, flush_bytes):
+    from triton.backends.nvidia.driver import CudaDriver
+
+    from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer
+
+    monkeypatch.setattr(CudaDriver, "get_empty_cache_for_benchmark", CudaDriver.get_empty_cache_for_benchmark)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(L2_cache_size=l2_bytes))
+    bound_autotune_flush_buffer(0)
+    bound_autotune_flush_buffer(0)  # the engine may call it again; the patch must not stack
+    monkeypatch.setattr(torch, "empty", lambda n, dtype, device: SimpleNamespace(nbytes=n * 4, dtype=dtype, device=device))
+    buf = CudaDriver.get_empty_cache_for_benchmark(None)
+    assert (buf.nbytes, buf.dtype, buf.device) == (flush_bytes, torch.int, "cuda")

From fc1224c088ce207d6be198c9b342c7bdcc5d15b0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Sat, 26 Sep 2026 00:28:53 -0400
Subject: [PATCH 222/249] fix(server): accept top_k <= 0 as "disabled", as vLLM
 and the OpenAI-compatible servers do

Upstream #223 made resolve_sampling reject top_k=0 with HTTP 400 ("top_k must be -1
(disabled) or >= 1"). Clients send top_k=0 to mean "off": Open WebUI with Ollama-style
parameters and some SDKs do, and vLLM and the other OpenAI-compatible servers accept any
top_k <= 0 as "no top-k filter". A checkpoint whose generation_config.json carries
top_k=0 made every request that omits top_k fail the same way under
--sampling-defaults=model.

resolve_sampling is the one place that validates top_k (the chat, legacy completions,
Responses and Anthropic converters and #393's n>1 fan-out all go through it; the
pydantic request models put no bound on it), so it now maps any top_k <= 0, from the
request or from the checkpoint defaults, to the engine's -1. The sampler already read
top_k < 1 as "off". The NaN/inf/negative temperature and top_p (0, 1] checks stay.

#223's test that expected ("top_k", 0) -> 400 now asserts that 0 and -5 are accepted and
reach the engine as -1; tests/server/test_sampling_validation.py covers the request and
checkpoint cases and that the other range checks still raise.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RUamGLmkVPkWmPPBEUqQLP
---
 python/freetoken/server/generation.py    |  7 +++--
 tests/server/test_openai_api.py          | 27 ++++++++++++++++++-
 tests/server/test_sampling_validation.py | 33 ++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 89c898a0e9..00d94abebf 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -476,8 +476,11 @@ def pick(value, key, framework):
     if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1:
         raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}")
     resolved_top_k = pick(top_k, "top_k", -1)
-    if resolved_top_k != -1 and resolved_top_k < 1:
-        raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}")
+    # Any top_k <= 0 means "no top-k filter", as in vLLM and the other OpenAI-compatible
+    # servers: Ollama-style clients (Open WebUI) and some SDKs send top_k=0 for "off", and a
+    # checkpoint's generation_config.json can carry it too. Map it to the engine's -1.
+    if resolved_top_k <= 0:
+        resolved_top_k = -1
     resolved_min_p = float(pick(min_p, "min_p", 0.0))
     resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0))
     if not 0.0 <= resolved_min_p <= 1.0:
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index 86781df84d..fdd2b11159 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -853,7 +853,6 @@ async def collect():
     [
         ("temperature", -0.5),
         ("top_p", 0.0),
-        ("top_k", 0),
     ],
 )
 def test_chat_completion_rejects_invalid_sampling(field, value):
@@ -906,3 +905,29 @@ async def chat_completion(req: ChatCompletionRequest):
 
     assert response.status_code == 200
     assert state.sent is not None
+
+
+@pytest.mark.parametrize("top_k", [0, -5])
+def test_chat_completion_treats_non_positive_top_k_as_disabled(top_k):
+    """Ollama-style clients (Open WebUI) and some SDKs send top_k=0 to mean "off"; vLLM and
+    the other OpenAI-compatible servers accept any top_k <= 0 as disabled, so it must not 400."""
+    app = FastAPI()
+    state = FakeState([])
+
+    @app.post("/v1/chat/completions")
+    async def chat_completion(req: ChatCompletionRequest):
+        return await handle_chat_completion(req, request=None, state=state, model_sampling={})
+
+    response = TestClient(app).post(
+        "/v1/chat/completions",
+        json={
+            "model": "unit-model",
+            "messages": [{"role": "user", "content": "hi"}],
+            "stream": True,
+            "top_k": top_k,
+        },
+    )
+
+    assert response.status_code == 200
+    assert state.sent is not None
+    assert state.sent.sampling_params.top_k == -1
diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py
index 3b2e1d6473..b13db921be 100644
--- a/tests/server/test_sampling_validation.py
+++ b/tests/server/test_sampling_validation.py
@@ -21,3 +21,36 @@ def test_repetition_penalty_must_be_positive_and_finite(value):
 
 def test_a_normal_repetition_penalty_is_kept():
     assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1)
+
+
+def _resolve_with(**overrides):
+    kwargs = dict(
+        temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False,
+        model_sampling={},
+    )
+    kwargs.update(overrides)
+    return resolve_sampling(**kwargs)
+
+
+@pytest.mark.parametrize("value", [0, -1, -5])
+def test_a_non_positive_top_k_means_disabled(value):
+    # vLLM's rule: any top_k <= 0 is "no top-k filter"; clients send 0 for "off"
+    assert _resolve_with(top_k=value).top_k == -1
+
+
+def test_a_checkpoint_default_top_k_of_zero_means_disabled():
+    assert _resolve_with(model_sampling={"top_k": 0}).top_k == -1
+
+
+def test_a_positive_top_k_is_kept():
+    assert _resolve_with(top_k=20).top_k == 20
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [("temperature", float("nan")), ("temperature", float("inf")), ("temperature", -0.5),
+     ("top_p", 0.0), ("top_p", 1.5), ("top_p", float("nan"))],
+)
+def test_the_other_sampling_ranges_are_still_rejected(field, value):
+    with pytest.raises(ValueError, match=field):
+        _resolve_with(**{field: value})

From c750abc3f3fc67d96a971ea1d0b9f48a6eb332fd Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 10:39:11 -0400
Subject: [PATCH 223/249] merge fork: logprobs and usage fixes (token bytes,
 completions offsets and stream entries, host-built logprob row index,
 responses and truncated reasoning_tokens, logprobs null)

---
 python/freetoken/core.py                     |   3 +
 python/freetoken/engine/sample.py            |  80 +++++----
 python/freetoken/message/frontend.py         |   3 +
 python/freetoken/message/tokenizer.py        |   2 +
 python/freetoken/scheduler/scheduler.py      |   1 +
 python/freetoken/server/logprobs.py          |  25 ++-
 python/freetoken/server/openai_api.py        |  56 +++++--
 python/freetoken/server/responses_api.py     |  17 +-
 python/freetoken/tokenizer/detokenize.py     | 128 ++++++++++++--
 python/freetoken/tokenizer/server.py         |  43 ++++-
 tests/engine/test_sample_logprobs.py         |  65 ++++++-
 tests/scheduler/test_cost_accounting_core.py |  22 +++
 tests/server/test_logprobs_api.py            | 168 ++++++++++++++++++-
 tests/server/test_responses_api.py           |  25 ++-
 tests/tokenizer/test_detokenize_extras.py    |  48 +++++-
 tests/tokenizer/test_logprobs_entry.py       |  54 ++++++
 16 files changed, 646 insertions(+), 94 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index a7811ed712..bbba592378 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -54,6 +54,9 @@ class SamplingParams:
     # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives.
     logprobs: bool = False
     top_logprobs: int = 0
+    # Set by the tokenizer worker when the rendered prompt leaves a <think> block open, so
+    # the detokenizer counts an unclosed generation as reasoning tokens.
+    starts_in_reasoning: bool = False
 
     @property
     def is_greedy(self) -> bool:
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 9d12b5f571..62bdd29cc8 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -57,7 +57,10 @@ class BatchSamplingArgs:
     greedy_mask: torch.Tensor | None = None
     penalties: list[RowPenalty] = field(default_factory=list)
     plan: LogitsPlan | None = None
-    logprob_rows: torch.Tensor | None = None
+    # Rows that asked for logprobs, as an int64 device index built on the host (no
+    # nonzero sync). None with logprob_all_rows set when every row asks: no gather then.
+    logprob_row_idx: torch.Tensor | None = None
+    logprob_all_rows: bool = False
     max_top_logprobs: int = 0
 
 
@@ -244,9 +247,12 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             )
         plan = self._plan(batch, params)
         want_logprobs = [p.logprobs for p in params]
-        logprob_rows = (
-            make_device_tensor(want_logprobs, torch.bool, self.device)
-            if any(want_logprobs)
+        logprob_all_rows = bool(want_logprobs) and all(want_logprobs)
+        logprob_row_idx = (
+            make_device_tensor(
+                [i for i, want in enumerate(want_logprobs) if want], torch.int64, self.device
+            )
+            if any(want_logprobs) and not logprob_all_rows
             else None
         )
         max_top_logprobs = min(
@@ -257,7 +263,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
                 temperatures=None,
                 penalties=penalties,
                 plan=plan,
-                logprob_rows=logprob_rows,
+                logprob_row_idx=logprob_row_idx,
+                logprob_all_rows=logprob_all_rows,
                 max_top_logprobs=max_top_logprobs,
             )
 
@@ -289,7 +296,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             greedy_mask=greedy_mask,
             penalties=penalties,
             plan=plan,
-            logprob_rows=logprob_rows,
+            logprob_row_idx=logprob_row_idx,
+            logprob_all_rows=logprob_all_rows,
             max_top_logprobs=max_top_logprobs,
         )
 
@@ -324,46 +332,50 @@ def compute_logprobs(
         sampled_tokens: torch.Tensor,
         args: BatchSamplingArgs,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
-        if args.logprob_rows is None:
-            return None
-
-        requested_rows = torch.nonzero(args.logprob_rows, as_tuple=False).flatten()
-        if requested_rows.numel() == 0:
+        rows = args.logprob_row_idx
+        if not args.logprob_all_rows and (rows is None or rows.numel() == 0):
             return None
 
-        request_logits = logits.index_select(0, requested_rows).float()
+        request_logits = (logits if rows is None else logits.index_select(0, rows)).float()
         # Reported values are raw model logprobs (pre-temperature log_softmax over logits).
         request_logprobs = torch.log_softmax(request_logits, dim=-1)
 
-        request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device).index_select(
-            0, requested_rows
-        )
-        request_row_idx = torch.arange(requested_rows.numel(), device=logits.device)
+        request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device)
+        if rows is not None:
+            request_tokens = request_tokens.index_select(0, rows)
+        request_row_idx = torch.arange(request_logprobs.shape[0], device=logits.device)
         request_chosen_logprobs = request_logprobs[request_row_idx, request_tokens]
 
-        chosen_logprobs = torch.full(
-            (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device
-        )
-        chosen_logprobs.index_copy_(0, requested_rows, request_chosen_logprobs)
+        if rows is None:
+            chosen_logprobs = request_chosen_logprobs
+        else:
+            chosen_logprobs = torch.full(
+                (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device
+            )
+            chosen_logprobs.index_copy_(0, rows, request_chosen_logprobs)
 
         if args.max_top_logprobs > 0:
             request_top_logprobs, request_top_ids = torch.topk(
                 request_logprobs, k=args.max_top_logprobs, dim=-1
             )
-            top_ids = torch.full(
-                (logits.shape[0], args.max_top_logprobs),
-                -1,
-                dtype=torch.int32,
-                device=logits.device,
-            )
-            top_logprobs = torch.full(
-                (logits.shape[0], args.max_top_logprobs),
-                float("-inf"),
-                dtype=torch.float32,
-                device=logits.device,
-            )
-            top_ids[requested_rows] = request_top_ids.to(torch.int32)
-            top_logprobs[requested_rows] = request_top_logprobs
+            if rows is None:
+                top_ids = request_top_ids.to(torch.int32)
+                top_logprobs = request_top_logprobs
+            else:
+                top_ids = torch.full(
+                    (logits.shape[0], args.max_top_logprobs),
+                    -1,
+                    dtype=torch.int32,
+                    device=logits.device,
+                )
+                top_logprobs = torch.full(
+                    (logits.shape[0], args.max_top_logprobs),
+                    float("-inf"),
+                    dtype=torch.float32,
+                    device=logits.device,
+                )
+                top_ids[rows] = request_top_ids.to(torch.int32)
+                top_logprobs[rows] = request_top_logprobs
         else:
             top_ids = torch.empty((logits.shape[0], 0), dtype=torch.int32, device=logits.device)
             top_logprobs = torch.empty((logits.shape[0], 0), dtype=torch.float32, device=logits.device)
diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py
index d70453e910..07634f1652 100644
--- a/python/freetoken/message/frontend.py
+++ b/python/freetoken/message/frontend.py
@@ -65,6 +65,9 @@ class UserReply(BaseFrontendMsg):
     # Neutral sampled-token logprobs entry for this token (see
     # tokenizer.detokenize.build_logprobs_entry); None when the request did not ask.
     logprobs: dict | None = None
+    # Where this token's text starts in the generated text, even while a stop-string
+    # prefix holds it back (completions logprobs.text_offset).
+    text_offset: int = 0
 
 
 @dataclass
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index 1b3b49eeeb..b5c24a6e10 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -40,6 +40,8 @@ class DetokenizeMsg(BaseTokenizerMsg):
     keep_stop_str: bool = False
     # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens).
     skip_special_tokens: bool = False
+    # The prompt left a reasoning block open (SamplingParams.starts_in_reasoning).
+    starts_in_reasoning: bool = False
     # Sampled-token logprobs (None unless the request asked): the chosen token's raw
     # logprob and the top alternatives, already cut to this request's top_logprobs.
     chosen_logprob: float | None = None
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index cb6085d344..358ce31355 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -414,6 +414,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         stop_strs=req.sampling_params.stop_strs or None,
                         keep_stop_str=req.sampling_params.include_stop_str_in_output,
                         skip_special_tokens=req.sampling_params.skip_special_tokens,
+                        starts_in_reasoning=req.sampling_params.starts_in_reasoning,
                         prefill_ms=(
                             0.0 if started_at is None
                             else (time.monotonic() - started_at) * 1000.0
diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py
index 8618569544..0cf6198aef 100644
--- a/python/freetoken/server/logprobs.py
+++ b/python/freetoken/server/logprobs.py
@@ -37,24 +37,31 @@ def chat_content_entry(entry: dict) -> dict:
     }
 
 
-def completions_logprobs(entries: list[dict], start_offset: int = 0) -> dict:
+def completions_logprobs(
+    entries: list[dict], offsets: list[int], text_len: int | None = None
+) -> dict:
+    """``offsets[i]`` is where entry i's token starts in the generated text (the pieces of
+    a split multi-byte character share its start). Given the finished text's ``text_len``,
+    offsets are clipped to it, since a matched stop string trims tokens off the end."""
     tokens: list[str] = []
     token_logprobs: list[float] = []
     top_logprobs: list[dict[str, float]] = []
-    text_offset: list[int] = []
-    offset = start_offset
 
     for entry in entries:
-        token = entry["token"]
-        tokens.append(token)
+        tokens.append(entry["token"])
         token_logprobs.append(entry["logprob"])
-        top_logprobs.append({candidate["token"]: candidate["logprob"] for candidate in entry["top"]})
-        text_offset.append(offset)
-        offset += len(token)
+        top: dict[str, float] = {}
+        for candidate in entry["top"]:
+            # The legacy shape keys alternatives by text, and partial-UTF-8 candidates
+            # all decode to U+FFFD: keep the most likely one per key.
+            key, logprob = candidate["token"], candidate["logprob"]
+            if key not in top or logprob > top[key]:
+                top[key] = logprob
+        top_logprobs.append(top)
 
     return {
         "tokens": tokens,
         "token_logprobs": token_logprobs,
         "top_logprobs": top_logprobs,
-        "text_offset": text_offset,
+        "text_offset": [offset if text_len is None else min(offset, text_len) for offset in offsets],
     }
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 0fc8113bc8..cd7f398c1d 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -399,10 +399,15 @@ async def handle_chat_completion(
             message["reasoning_content"] = result.reasoning
         if result.tool_calls:
             message["tool_calls"] = _tool_calls_to_openai(result.tool_calls)
-        choice: dict[str, Any] = {"index": index, "message": message, "finish_reason": result.finish_reason}
-        if req.logprobs:
-            choice["logprobs"] = {"content": [chat_content_entry(e) for e in result.logprobs]}
-        choices.append(choice)
+        choices.append({
+            "index": index,
+            "message": message,
+            # OpenAI always sends the key, null when logprobs were not requested
+            "logprobs": (
+                {"content": [chat_content_entry(e) for e in result.logprobs]} if req.logprobs else None
+            ),
+            "finish_reason": result.finish_reason,
+        })
 
     first = results[0]
     response: dict[str, Any] = {
@@ -657,6 +662,7 @@ async def handle_completion(
     async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
         text = ""
         entries: list[dict] = []
+        offsets: list[int] = []
         finish_reason = "stop"
         prompt_tokens = completion_tokens = cached_tokens = 0
         async for ack in state.wait_for_ack(uid):
@@ -665,9 +671,10 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
             prompt_tokens += ack.prompt_tokens_delta
             completion_tokens += ack.completion_tokens_delta
             cached_tokens += ack.cached_tokens
-            text += ack.incremental_output
             if req.logprobs is not None and ack.logprobs is not None:
                 entries.append(ack.logprobs)
+                offsets.append(ack.text_offset)
+            text += ack.incremental_output
             if ack.finished:
                 finish_reason = getattr(ack, "finish_reason", None) or "stop"
                 break
@@ -675,7 +682,9 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
             "index": index,
             "text": echo_texts[index // req.n] + text,
             "finish_reason": finish_reason,
-            "logprobs": completions_logprobs(entries) if req.logprobs is not None else None,
+            "logprobs": (
+                completions_logprobs(entries, offsets, len(text)) if req.logprobs is not None else None
+            ),
             "_usage": (prompt_tokens, completion_tokens, cached_tokens),
         }
 
@@ -727,7 +736,11 @@ async def stream_completion_chunks(
     echo_text: str = "",
 ) -> AsyncIterator[bytes]:
     prompt_tokens = 0
-    text_offset = 0
+    sent_len = 0
+    # Logprob entries of acks with no visible text yet (held partial UTF-8 or stop
+    # prefix), carried to the next chunk that has text.
+    pending: list[dict] = []
+    pending_offsets: list[int] = []
     completion_tokens = 0
     cached_tokens = 0
     finish_reason = "stop"
@@ -743,7 +756,11 @@ async def stream_completion_chunks(
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
+        if req.logprobs is not None and ack.logprobs is not None:
+            pending.append(ack.logprobs)
+            pending_offsets.append(ack.text_offset)
         if ack.incremental_output:
+            sent_len += len(ack.incremental_output)
             yield _sse(
                 _completion_chunk(
                     req, uid,
@@ -751,20 +768,37 @@ async def stream_completion_chunks(
                         "text": ack.incremental_output,
                         "index": index,
                         "finish_reason": None,
+                        # Mid-stream, a stamp past sent_len is held-back text still to come.
                         "logprobs": (
-                            completions_logprobs([ack.logprobs], text_offset)
-                            if req.logprobs is not None and ack.logprobs is not None
+                            completions_logprobs(
+                                pending, pending_offsets, sent_len if ack.finished else None
+                            )
+                            if pending
                             else None
                         ),
                     }],
                 )
             )
-            text_offset += len(ack.incremental_output)
+            pending, pending_offsets = [], []
         if ack.finished:
             finish_reason = getattr(ack, "finish_reason", None) or "stop"
             break
 
-    yield _sse(_completion_chunk(req, uid, [{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}]))
+    # Entries whose text never showed (EOS, a stop string, a trailing partial character)
+    # ride the finish chunk, so the joined stream matches the non-stream list.
+    yield _sse(
+        _completion_chunk(
+            req, uid,
+            [{
+                "text": "",
+                "index": index,
+                "finish_reason": finish_reason,
+                "logprobs": (
+                    completions_logprobs(pending, pending_offsets, sent_len) if pending else None
+                ),
+            }],
+        )
+    )
     usage = _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens))
     if not terminal:
         yield _StreamUsage(usage)
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index a84c68cb6d..4240080510 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -467,6 +467,7 @@ def build_responses_response(
         usage=_usage(
             result.prompt_tokens, result.completion_tokens,
             result.cached_tokens if cache_report else 0,
+            reasoning_tokens=result.reasoning_tokens,
         ),
         incomplete_reason="max_output_tokens" if truncated else None,
         metrics=build_metrics(
@@ -504,7 +505,9 @@ def _response_obj(
     return response
 
 
-def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> ResponseUsage:
+def _usage(
+    prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0
+) -> ResponseUsage:
     # input_tokens stays inclusive of the cached prefix (OpenAI semantics);
     # cached_tokens is nonzero only under --enable-cache-report.
     return ResponseUsage(
@@ -512,7 +515,7 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -
         output_tokens=completion_tokens,
         total_tokens=prompt_tokens + completion_tokens,
         input_tokens_details=InputTokensDetails(cached_tokens=cached_tokens, cache_write_tokens=0),
-        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=reasoning_tokens),
     )
 
 
@@ -550,6 +553,7 @@ def snapshot(status, output, usage=None, incomplete_reason=None, request_metrics
     usage_pt = 0
     usage_ct = 0
     usage_cached = 0
+    usage_reasoning = 0
 
     def close_current() -> list[str]:
         nonlocal current, output_index
@@ -735,6 +739,7 @@ def args_delta_frame(fragment: str) -> str:
                 finish_reason = ev.finish_reason
                 usage_pt, usage_ct = ev.prompt_tokens, ev.completion_tokens
                 usage_cached = ev.cached_tokens if cache_report else 0
+                usage_reasoning = ev.reasoning_tokens
                 request_metrics = build_metrics(
                     prompt_tokens=ev.prompt_tokens,
                     completion_tokens=ev.completion_tokens,
@@ -748,7 +753,7 @@ def args_delta_frame(fragment: str) -> str:
                         type="response.incomplete", sequence_number=seq.next(),
                         response=snapshot(
                             "incomplete", output_items,
-                            usage=_usage(usage_pt, usage_ct, usage_cached),
+                            usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning),
                             incomplete_reason="max_output_tokens",
                             request_metrics=request_metrics,
                         ),
@@ -758,7 +763,7 @@ def args_delta_frame(fragment: str) -> str:
                         type="response.completed", sequence_number=seq.next(),
                         response=snapshot(
                             "completed", output_items,
-                            usage=_usage(usage_pt, usage_ct, usage_cached),
+                            usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning),
                             request_metrics=request_metrics,
                         ),
                     ))
@@ -772,7 +777,7 @@ def args_delta_frame(fragment: str) -> str:
             type="response.failed", sequence_number=seq.next(),
             response=_response_obj(
                 response_id, created, req.model, output_items, status="failed",
-                usage=_usage(usage_pt, usage_ct, usage_cached),
+                usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning),
                 # codex reads this code to tell a blown context window from a generic failure.
                 # model_construct because ResponseError.code is a closed Literal in the SDK.
                 error=ResponseError.model_construct(
@@ -787,7 +792,7 @@ def args_delta_frame(fragment: str) -> str:
             type="response.failed", sequence_number=seq.next(),
             response=_response_obj(
                 response_id, created, req.model, output_items, status="failed",
-                usage=_usage(usage_pt, usage_ct, usage_cached),
+                usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning),
                 error=ResponseError(code="server_error", message=str(exc)),
             ),
         ))
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index 2eb3029df0..098082edd3 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -1,5 +1,7 @@
+import json
+import weakref
 from dataclasses import dataclass
-from typing import Dict, FrozenSet, List
+from typing import Any, Dict, FrozenSet, List
 
 from freetoken.message import DetokenizeMsg
 from transformers import PreTrainedTokenizerBase
@@ -73,6 +75,15 @@ class DecodeStatus:
     sent_offset: int  # length of sent out string
     skip_special: bool = False  # decode this request with skip_special_tokens
     think_end_at: int | None = None  # generated tokens up to and incl. the reasoning end tag
+    in_reasoning: bool = False  # the prompt or the output opened a reasoning block
+    generated: int = 0  # sampled tokens, EOS included
+    output_len: int = 0  # decoded text so far, before holdback: where the next token starts
+
+    def reasoning_tokens(self) -> int:
+        if self.think_end_at is not None:
+            return self.think_end_at
+        # stopped inside the block (max_tokens, or EOS before the end tag): all of it
+        return self.generated if self.in_reasoning else 0
 
     def update(self, read_str: str, surr_str: str) -> str:
         new_text = read_str[len(surr_str) :]
@@ -99,6 +110,7 @@ def __init__(
         tokenizer: PreTrainedTokenizerBase,
         eos_token_ids: FrozenSet[int] | None = None,
         think_end_id: int | None = None,
+        think_start_id: int | None = None,
     ) -> None:
         # uid -> DecodeStatus
         self.decode_map: Dict[int, DecodeStatus] = {}
@@ -109,6 +121,7 @@ def __init__(
         # The reasoning end tag's id (e.g. ``</think>``): tokens generated up to and
         # including it are reported as usage.completion_tokens_details.reasoning_tokens.
         self.think_end_id = think_end_id if think_end_id is None else int(think_end_id)
+        self.think_start_id = think_start_id if think_start_id is None else int(think_start_id)
 
     def discard(self, uid: int) -> None:
         """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored
@@ -131,13 +144,21 @@ def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]:
                     out[i] = t
         return [t if t is not None else "" for t in out]
 
-    def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]:
-        """Incremental text per message plus, for a finished message, its reasoning token
-        count (0 otherwise or when the model emitted no reasoning end tag)."""
+    def detokenize_with_meta(
+        self, msgs: List[DetokenizeMsg]
+    ) -> tuple[List[str], List[int], List[int]]:
+        """Incremental text per message; for a finished message, its reasoning token count
+        (0 otherwise or when the generation never was in a reasoning block); and where each
+        message's token starts in the generated text. That offset ignores stop-string
+        holdback, which delays when text is sent but not where it lands."""
         # Each message must advance its request's decode state before the next one.
         if len({msg.uid for msg in msgs}) != len(msgs):
             parts = [self.detokenize_with_meta([msg]) for msg in msgs]
-            return [p[0][0] for p in parts], [p[1][0] for p in parts]
+            return (
+                [p[0][0] for p in parts],
+                [p[1][0] for p in parts],
+                [p[2][0] for p in parts],
+            )
         read_ids: List[List[int]] = []
         surr_ids: List[List[int]] = []
         skip: List[bool] = []
@@ -150,8 +171,10 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li
                     surr_offset=0,
                     sent_offset=0,
                     skip_special=bool(getattr(msg, "skip_special_tokens", False)),
+                    in_reasoning=bool(getattr(msg, "starts_in_reasoning", False)),
                 )
             s = self.decode_map[msg.uid]
+            s.generated += 1
             if not (
                 msg.finished and not msg.matched_stop and msg.next_token in self.eos_token_ids
             ):
@@ -162,6 +185,8 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li
                 and msg.next_token == self.think_end_id
             ):
                 s.think_end_at = len(s.decoded_ids)
+            if s.think_end_at is None and msg.next_token == self.think_start_id:
+                s.in_reasoning = True
             read_ids.append(s.decoded_ids[s.surr_offset :])
             surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset])
             skip.append(s.skip_special)
@@ -171,9 +196,12 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li
 
         incremental_strs: List[str] = []
         reasoning_tokens: List[int] = []
+        text_offsets: List[int] = []
         for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True):
             s = self.decode_map[msg.uid]
+            text_offsets.append(s.output_len)
             output_str = s.update(read_str, surr_str)
+            s.output_len = max(s.output_len, len(output_str))
 
             prev_sent = s.sent_offset
             if msg.finished:
@@ -197,11 +225,77 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li
             incremental_output = output_str[prev_sent:emit_end] if emit_end > prev_sent else ""
             s.sent_offset = max(prev_sent, emit_end)
             incremental_strs.append(incremental_output)
-            reasoning_tokens.append((s.think_end_at or 0) if msg.finished else 0)
+            reasoning_tokens.append(s.reasoning_tokens() if msg.finished else 0)
             if msg.finished:
                 del self.decode_map[msg.uid]
 
-        return incremental_strs, reasoning_tokens
+        return incremental_strs, reasoning_tokens, text_offsets
+
+
+def _gpt2_byte_decoder() -> dict[str, int]:
+    """Inverse of GPT-2's ``bytes_to_unicode``: the printable stand-in char -> its byte."""
+    bs = [*range(ord("!"), ord("~") + 1), *range(ord("¡"), ord("¬") + 1), *range(ord("®"), ord("ÿ") + 1)]
+    cs = bs[:]
+    n = 0
+    for b in range(256):
+        if b not in bs:
+            bs.append(b)
+            cs.append(256 + n)
+            n += 1
+    return {chr(c): b for b, c in zip(bs, cs, strict=True)}
+
+
+class _TokenBytes:
+    """A token id's own bytes. A single-id decode cannot give them for a piece of a
+    multi-byte character (it decodes to U+FFFD), so read the vocab piece instead."""
+
+    def __init__(self, tokenizer: Any) -> None:
+        decoder: dict = {}
+        backend = getattr(tokenizer, "backend_tokenizer", None)
+        try:
+            decoder = json.loads(backend.decoder.__getstate__())
+        except Exception:  # noqa: BLE001 -- no fast backend: fall back to the decoded text
+            pass
+        kinds = {d.get("type") for d in decoder.get("decoders", [decoder])}
+        self.byte_level = _gpt2_byte_decoder() if "ByteLevel" in kinds else None
+        self.byte_fallback = "ByteFallback" in kinds
+        added = getattr(tokenizer, "added_tokens_decoder", None) or {}
+        self.added = {int(tid): str(getattr(tok, "content", tok)) for tid, tok in added.items()}
+        self.convert = getattr(tokenizer, "convert_ids_to_tokens", None)
+
+    def __call__(self, tid: int, text: str) -> bytes:
+        if tid in self.added:
+            # the GPT-2 table cannot invert an added token's raw (e.g. full-width) chars
+            return self.added[tid].encode("utf-8")
+        try:
+            piece = self.convert(tid) if self.convert is not None else None
+        except Exception:  # noqa: BLE001 -- out-of-vocab id
+            piece = None
+        if isinstance(piece, str):
+            if self.byte_level is not None:
+                try:
+                    return bytes(self.byte_level[c] for c in piece)
+                except KeyError:
+                    pass
+            elif self.byte_fallback and len(piece) == 6 and piece.startswith("<0x") and piece.endswith(">"):
+                try:
+                    return bytes([int(piece[3:5], 16)])
+                except ValueError:
+                    pass
+        return text.encode("utf-8")
+
+
+_TOKEN_BYTES: "weakref.WeakKeyDictionary[Any, _TokenBytes]" = weakref.WeakKeyDictionary()
+
+
+def _token_bytes(tokenizer: Any) -> _TokenBytes:
+    try:
+        cached = _TOKEN_BYTES.get(tokenizer)
+        if cached is None:
+            cached = _TOKEN_BYTES[tokenizer] = _TokenBytes(tokenizer)
+    except TypeError:  # not weak-referenceable or unhashable
+        return _TokenBytes(tokenizer)
+    return cached
 
 
 def build_logprobs_entry(
@@ -213,21 +307,27 @@ def build_logprobs_entry(
 ) -> dict:
     """Neutral sampled-token logprobs entry (UserReply.logprobs). Values are the raw
     pre-temperature logprobs the sampler computed; token text comes from a single-id
-    decode, with UTF-8 bytes alongside so clients can reassemble partial-UTF-8 pieces."""
+    decode, and ``bytes`` are the token's own bytes, so clients can reassemble a
+    character split across tokens. Ids past the real vocab (the padded tail of the
+    logits) report an empty token."""
+    token_bytes = _token_bytes(tokenizer)
 
     def _fields(tid: int) -> tuple[str, list[int]]:
-        text = tokenizer.decode([tid])
-        return text, list(text.encode("utf-8"))
+        try:
+            text = tokenizer.decode([tid])
+        except Exception:  # noqa: BLE001 -- a padded id some tokenizers refuse to decode
+            text = ""
+        return text, list(token_bytes(tid, text))
 
-    token, token_bytes = _fields(token_id)
+    token, data = _fields(token_id)
     top = []
     for tid, logprob in zip(top_ids or [], top_logprobs or []):
-        text, data = _fields(int(tid))
-        top.append({"token_id": int(tid), "token": text, "bytes": data, "logprob": float(logprob)})
+        text, cand = _fields(int(tid))
+        top.append({"token_id": int(tid), "token": text, "bytes": cand, "logprob": float(logprob)})
     return {
         "token_id": token_id,
         "token": token,
-        "bytes": token_bytes,
+        "bytes": data,
         "logprob": float(chosen_logprob),
         "top": top,
     }
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index 4c650042b4..3b1068514d 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -84,7 +84,11 @@ def _send_generation_replies(
 def _reasoning_end_id(tokenizer) -> int | None:
     """The id of the reasoning end tag when the vocabulary has one (``</think>`` for the
     Qwen / DeepSeek / GLM families), else None. Only used to count reasoning tokens."""
-    for tag in ("</think>", "<|end_of_thinking|>"):
+    return _vocab_tag_id(tokenizer, ("</think>", "<|end_of_thinking|>"))
+
+
+def _vocab_tag_id(tokenizer, tags: tuple[str, ...]) -> int | None:
+    for tag in tags:
         try:
             tid = tokenizer.convert_tokens_to_ids(tag)
         except Exception:  # noqa: BLE001 -- a tokenizer without the tag answers unk or raises
@@ -94,10 +98,20 @@ def _reasoning_end_id(tokenizer) -> int | None:
     return None
 
 
+def _prompt_opens_reasoning(input_ids: torch.Tensor, start_id: int, end_id: int) -> bool:
+    """The prompt's last reasoning tag is an opening one: the template pre-opened the
+    block (qwen3 with thinking on), so generation starts inside it. A raw prompt that just
+    mentions a literal <think> trips this too; harmless, since only the chat and responses
+    endpoints, whose prompts are always templated, report reasoning_tokens."""
+    tags = torch.nonzero((input_ids == start_id) | (input_ids == end_id)).flatten()
+    return tags.numel() > 0 and int(input_ids[tags[-1]]) == start_id
+
+
 def _tokenize_requests(
     tokenize_manager: Any,
     messages: List[TokenizeMsg],
     logger: Any,
+    reasoning_tags: tuple[int, int] | None = None,
 ) -> tuple[List[UserMsg], List[UserReply]]:
     """Tokenize independently, returning backend work plus terminal frontend errors.
 
@@ -132,6 +146,10 @@ def _tokenize_requests(
                 )
             )
             continue
+        if reasoning_tags is not None:
+            user_msg.sampling_params.starts_in_reasoning = _prompt_opens_reasoning(
+                user_msg.input_ids, *reasoning_tags
+            )
         backend.append(user_msg)
     return backend, errors
 
@@ -163,10 +181,18 @@ def tokenize_worker(
     from .tokenize import TokenizeManager
 
     tokenize_manager = TokenizeManager(tokenizer, get_mm_processor(tokenizer_path, mm))
+    think_start_id = _vocab_tag_id(tokenizer, ("<think>",))
+    think_end_id = _reasoning_end_id(tokenizer)
+    reasoning_tags = (
+        (think_start_id, think_end_id)
+        if think_start_id is not None and think_end_id is not None
+        else None
+    )
     detokenize_manager = DetokenizeManager(
         tokenizer,
         load_eos_token_ids(tokenizer_path, tokenizer),
-        think_end_id=_reasoning_end_id(tokenizer),
+        think_end_id=think_end_id,
+        think_start_id=think_start_id,
     )
 
     if ack_queue is not None:
@@ -224,7 +250,9 @@ def tokenize_worker(
             )
             sampled_replies: List[UserReply] = []
             if len(detokenize_msg) > 0:
-                replies, reasoning_counts = detokenize_manager.detokenize_with_meta(detokenize_msg)
+                replies, reasoning_counts, text_offsets = detokenize_manager.detokenize_with_meta(
+                    detokenize_msg
+                )
                 sampled_replies = [
                     UserReply(
                         uid=msg.uid,
@@ -233,6 +261,7 @@ def tokenize_worker(
                         finish_reason=msg.finish_reason,
                         matched_stop=msg.matched_stop,
                         reasoning_tokens=reasoning_tokens,
+                        text_offset=text_offset,
                         completion_tokens_delta=1,
                         kv_used_pages=msg.kv_used_pages,
                         kv_total_pages=msg.kv_total_pages,
@@ -256,8 +285,8 @@ def tokenize_worker(
                             else None
                         ),
                     )
-                    for msg, reply, reasoning_tokens in zip(
-                        detokenize_msg, replies, reasoning_counts, strict=True
+                    for msg, reply, reasoning_tokens, text_offset in zip(
+                        detokenize_msg, replies, reasoning_counts, text_offsets, strict=True
                     )
                 ]
 
@@ -279,7 +308,9 @@ def tokenize_worker(
                 # Tokenize per-message so a single un-renderable request (e.g. a chat template
                 # that rejects the message layout) becomes a terminal error reply for THAT uid
                 # instead of an uncaught exception that kills the worker and bricks the server.
-                backend, errors = _tokenize_requests(tokenize_manager, tokenize_msg, logger)
+                backend, errors = _tokenize_requests(
+                    tokenize_manager, tokenize_msg, logger, reasoning_tags
+                )
                 if errors:
                     send_frontend.put(
                         errors[0] if len(errors) == 1 else BatchFrontendMsg(data=errors)
diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
index 88fffcc2ed..02dff731ef 100644
--- a/tests/engine/test_sample_logprobs.py
+++ b/tests/engine/test_sample_logprobs.py
@@ -20,7 +20,7 @@ def test_compute_logprobs_matches_raw_log_softmax_and_sorted_top() -> None:
     sampled_tokens = torch.tensor([2, 3, 0], dtype=torch.long)
     args = BatchSamplingArgs(
         temperatures=torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32),
-        logprob_rows=torch.tensor([True, False, True], dtype=torch.bool),
+        logprob_row_idx=torch.tensor([0, 2]),
         max_top_logprobs=3,
     )
 
@@ -58,16 +58,15 @@ def test_compute_logprobs_ignores_temperatures() -> None:
         dtype=torch.float32,
     )
     sampled_tokens = torch.tensor([1, 2], dtype=torch.long)
-    logprob_rows = torch.tensor([True, True], dtype=torch.bool)
 
     cold = BatchSamplingArgs(
         temperatures=torch.full((2,), 0.3, dtype=torch.float32),
-        logprob_rows=logprob_rows,
+        logprob_all_rows=True,
         max_top_logprobs=2,
     )
     hot = BatchSamplingArgs(
         temperatures=torch.full((2,), 2.5, dtype=torch.float32),
-        logprob_rows=logprob_rows,
+        logprob_all_rows=True,
         max_top_logprobs=2,
     )
 
@@ -89,7 +88,7 @@ def test_compute_logprobs_returns_none_without_requested_rows() -> None:
     sampled_tokens = torch.tensor([0, 1], dtype=torch.long)
     args = BatchSamplingArgs(
         temperatures=torch.ones(2),
-        logprob_rows=torch.zeros(2, dtype=torch.bool),
+        logprob_row_idx=torch.empty(0, dtype=torch.int64),
         max_top_logprobs=2,
     )
 
@@ -109,7 +108,7 @@ def test_compute_logprobs_handles_zero_top_logprobs() -> None:
     sampled_tokens = torch.tensor([3, 0, 1], dtype=torch.long)
     args = BatchSamplingArgs(
         temperatures=torch.full((3,), 1.0),
-        logprob_rows=torch.tensor([True, False, True], dtype=torch.bool),
+        logprob_row_idx=torch.tensor([0, 2]),
         max_top_logprobs=0,
     )
 
@@ -166,3 +165,57 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc
     assert top_logprobs.shape == (2, sampler.vocab_size)
     assert torch.equal(top_ids[1], torch.full((sampler.vocab_size,), -1, dtype=torch.int32))
     assert torch.isneginf(top_logprobs[1]).all()
+
+
+def _reference_logprobs(logits, sampled_tokens, want, k):
+    """The pre-index implementation (bool mask + torch.nonzero), kept as the oracle."""
+    mask = torch.tensor(want, dtype=torch.bool)
+    rows = torch.nonzero(mask, as_tuple=False).flatten()
+    if rows.numel() == 0:
+        return None
+    logprobs = torch.log_softmax(logits.index_select(0, rows).float(), dim=-1)
+    chosen = torch.full((logits.shape[0],), float("nan"))
+    chosen.index_copy_(
+        0, rows, logprobs[torch.arange(rows.numel()), sampled_tokens.long().index_select(0, rows)]
+    )
+    top_ids = torch.full((logits.shape[0], k), -1, dtype=torch.int32)
+    top_logprobs = torch.full((logits.shape[0], k), float("-inf"))
+    if k > 0:
+        values, ids = torch.topk(logprobs, k=k, dim=-1)
+        top_ids[rows] = ids.to(torch.int32)
+        top_logprobs[rows] = values
+    return chosen, top_ids, top_logprobs
+
+
+def test_compute_logprobs_builds_row_index_on_host(monkeypatch) -> None:
+    # A nonzero over a CUDA mask is a host-device sync on every logprobs step; the row
+    # index must come from prepare's host list instead, with values unchanged.
+    import freetoken.engine.sample as sample_mod
+
+    monkeypatch.setattr(
+        sample_mod, "make_device_tensor",
+        lambda data, dtype, device: torch.tensor(data, dtype=dtype),
+    )
+    sampler = Sampler(torch.device("cpu"), vocab_size=6)
+    logits = torch.randn(4, 6, generator=torch.Generator().manual_seed(0))
+    sampled_tokens = torch.tensor([5, 0, 3, 1], dtype=torch.int32)
+
+    for want in ([True, False, True, False], [True] * 4, [False] * 4):
+        batch = SimpleNamespace(
+            reqs=[
+                SimpleNamespace(sampling_params=SamplingParams(logprobs=w, top_logprobs=3 if w else 0))
+                for w in want
+            ]
+        )
+        args = sampler.prepare(batch)
+        expected = _reference_logprobs(logits, sampled_tokens, want, 3)
+        with monkeypatch.context() as m:
+            m.setattr(torch, "nonzero", lambda *a, **kw: (_ for _ in ()).throw(AssertionError("sync")))
+            result = sampler.compute_logprobs(logits, sampled_tokens, args)
+        if expected is None:
+            assert result is None
+            continue
+        assert result is not None
+        for got, ref in zip(result, expected, strict=True):
+            assert got.shape == ref.shape and got.dtype == ref.dtype
+            assert torch.equal(got.nan_to_num(-7.0), ref.nan_to_num(-7.0))
diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py
index 512bfa4049..045b9924be 100644
--- a/tests/scheduler/test_cost_accounting_core.py
+++ b/tests/scheduler/test_cost_accounting_core.py
@@ -60,6 +60,28 @@ def tokenize(self, messages):
     assert errors == []  # in particular, no early prompt_tokens_delta UserReply
 
 
+def test_tokenization_flags_a_prompt_that_leaves_reasoning_open():
+    # The detokenizer cannot see the prompt; the tokenizer worker marks a template that
+    # pre-opened <think> (5) so an unclosed generation counts as reasoning (6 = </think>).
+    prompts = {
+        1: [10, 5, 11],  # thinking on: the generation prompt ends inside <think>
+        2: [10, 5, 6, 11],  # thinking off: the template pre-closed the block
+        3: [5, 7, 6, 10, 5, 12],  # an earlier turn's block is closed, the last one open
+        4: [10, 11],  # no reasoning tags
+    }
+
+    class Tokenizer:
+        def tokenize(self, messages):
+            (msg,) = messages
+            ids = torch.tensor(prompts[msg.uid], dtype=torch.int32)
+            return [UserMsg(uid=msg.uid, input_ids=ids, sampling_params=msg.sampling_params)]
+
+    backend, _ = _tokenize_requests(
+        Tokenizer(), [_tokenize_msg(uid) for uid in prompts], _Logger(), (5, 6)
+    )
+    assert [m.sampling_params.starts_in_reasoning for m in backend] == [True, False, True, False]
+
+
 def test_tokenization_failure_and_empty_prompt_are_terminal_without_usage():
     class Tokenizer:
         def tokenize(self, messages):
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 8cf0387c60..69ccf82add 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -2,8 +2,11 @@
 
 import asyncio
 import json
+import os
 from types import SimpleNamespace
 
+import pytest
+
 from freetoken.message import TokenizeMsg, UserReply
 from freetoken.server.openai_api import (
     ChatCompletionRequest,
@@ -43,7 +46,9 @@ async def wait_for_ack(self, uid: int):
             yield reply
 
 
-def reply(text: str, *, finished: bool = False, logprobs: dict | None = None) -> UserReply:
+def reply(
+    text: str, *, finished: bool = False, logprobs: dict | None = None, text_offset: int = 0
+) -> UserReply:
     return UserReply(
         uid=42,
         incremental_output=text,
@@ -52,6 +57,7 @@ def reply(text: str, *, finished: bool = False, logprobs: dict | None = None) ->
         completion_tokens_delta=1 if text else 0,
         cached_tokens=0,
         logprobs=logprobs,
+        text_offset=text_offset,
     )
 
 
@@ -140,7 +146,9 @@ def test_chat_non_stream_logprobs() -> None:
             {},
         )
     )
-    assert without_logprobs["choices"][0].get("logprobs") is None
+    # OpenAI sends the key with null rather than leaving it out
+    assert "logprobs" in without_logprobs["choices"][0]
+    assert without_logprobs["choices"][0]["logprobs"] is None
 
 
 def test_chat_stream_logprobs_follow_content_deltas() -> None:
@@ -175,7 +183,9 @@ def test_completion_logprobs_non_stream_and_stream() -> None:
         handle_completion(
             req,
             None,
-            FakeState([reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second)]),
+            FakeState(
+                [reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second, text_offset=2)]
+            ),
             {},
         )
     )
@@ -215,6 +225,158 @@ def test_completion_logprobs_non_stream_and_stream() -> None:
     }
 
 
+def _crab_replies() -> list[UserReply]:
+    # 🦀 is three byte-level tokens: the first two acks carry no visible text yet.
+    pieces = [
+        reply("", logprobs=entry(9008, "\ufffd", -0.1)),
+        reply("", logprobs=entry(99, "\ufffd", -0.2)),
+        reply("🦀", logprobs=entry(222, "\ufffd", -0.3)),
+        reply(" x", logprobs=entry(1, " x", -0.4), text_offset=1),
+        reply("", finished=True, logprobs=entry(2, "<|im_end|>", -0.5), text_offset=3),
+    ]
+    return pieces
+
+
+def test_completion_text_offset_counts_visible_text() -> None:
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=1, max_tokens=8)
+    result = run(handle_completion(req, None, FakeState(_crab_replies()), {}))
+    logprobs = result["choices"][0]["logprobs"]
+    assert result["choices"][0]["text"] == "🦀 x"
+    # The split character's pieces all point at where it starts, not at
+    # len("\ufffd") steps that drift past the real text.
+    assert logprobs["text_offset"] == [0, 0, 0, 1, 3]
+
+
+def _holdback_replies() -> list[UserReply]:
+    # "Answer: 4\nQuestion 2" under stop=["\n\n"]: the "\n" ack sends nothing (it could
+    # start the stop), and "Question" releases it along with its own text.
+    acks = [("Answer", "Answer", 0), (":", ":", 6), (" ", " ", 7), ("4", "4", 8),
+            ("\n", "", 9), ("Question", "\nQuestion", 10), (" ", " ", 18), ("2", "2", 19)]
+    return [
+        reply(text, finished=i == len(acks) - 1, logprobs=entry(i, token, -0.1), text_offset=offset)
+        for i, (token, text, offset) in enumerate(acks)
+    ]
+
+
+def _joined_stream_logprobs(replies: list[UserReply], **kwargs) -> dict[str, list]:
+    req = CompletionRequest(model="client-model", prompt="hello", max_tokens=8, stream=True, **kwargs)
+    events = parse_sse(run(_collect(stream_completion_chunks(42, req, FakeState(replies)))))
+    joined: dict[str, list] = {"tokens": [], "token_logprobs": [], "top_logprobs": [], "text_offset": []}
+    for event in events:
+        if isinstance(event, dict) and event["choices"] and event["choices"][0]["logprobs"]:
+            for key in joined:
+                joined[key] += event["choices"][0]["logprobs"][key]
+    return joined
+
+
+def test_completion_text_offset_survives_stop_prefix_holdback() -> None:
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["\n\n"])
+    choice = run(handle_completion(req, None, FakeState(_holdback_replies()), {}))["choices"][0]
+    logprobs = choice["logprobs"]
+    assert choice["text"] == "Answer: 4\nQuestion 2"
+    assert logprobs["text_offset"] == [0, 6, 7, 8, 9, 10, 18, 19]
+    for token, offset in zip(logprobs["tokens"], logprobs["text_offset"], strict=True):
+        assert choice["text"][offset:].startswith(token)
+    assert _joined_stream_logprobs(_holdback_replies(), logprobs=0, stop=["\n\n"]) == logprobs
+
+
+def test_completion_text_offset_clipped_to_stop_trimmed_text() -> None:
+    # "a\n\n" under stop=["\n\n"]: the second "\n" completes the stop and both are trimmed,
+    # so their offsets cannot point past the returned text.
+    replies = [
+        reply("a", logprobs=entry(1, "a", -0.1)),
+        reply("", logprobs=entry(2, "\n", -0.2), text_offset=1),
+        reply("", finished=True, logprobs=entry(2, "\n", -0.3), text_offset=2),
+    ]
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["\n\n"])
+    choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
+    assert choice["text"] == "a"
+    assert choice["logprobs"]["text_offset"] == [0, 1, 1]
+    assert _joined_stream_logprobs(replies, logprobs=0, stop=["\n\n"]) == choice["logprobs"]
+
+
+def test_completion_stream_text_offset_ahead_of_sent_text() -> None:
+    # stop=["aab"] over "aaax": two "a"s stay held while the first is sent, so the held
+    # tokens' offsets run past the text sent so far; that is where they land, not an overrun.
+    def replies() -> list[UserReply]:
+        acks = [("a", "", 0), ("a", "", 1), ("a", "a", 2), ("x", "aax", 3)]
+        return [
+            reply(text, finished=i == 3, logprobs=entry(i, token, -0.1), text_offset=offset)
+            for i, (token, text, offset) in enumerate(acks)
+        ]
+
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["aab"])
+    choice = run(handle_completion(req, None, FakeState(replies()), {}))["choices"][0]
+    assert choice["text"] == "aaax"
+    assert choice["logprobs"]["text_offset"] == [0, 1, 2, 3]
+    assert _joined_stream_logprobs(replies(), logprobs=0, stop=["aab"]) == choice["logprobs"]
+
+
+def test_completion_text_offset_real_tokenizer_with_stops() -> None:
+    # The real Qwen3.8 tokenizer through the real detokenizer: every token's offset must
+    # point at its own text, stream and non-stream alike.
+    path = os.path.expanduser("~/models/Qwen3.8-Flash-Next-NVFP4-FTW-fp8")
+    if not os.path.exists(os.path.join(path, "tokenizer.json")):
+        pytest.skip("local Qwen3.8 tokenizer not present")
+    from transformers import AutoTokenizer
+
+    from freetoken.message import DetokenizeMsg
+    from freetoken.tokenizer.detokenize import DetokenizeManager, build_logprobs_entry
+
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    for text, stop in [("Answer: 4\nQuestion 2", "\n\n"), ("Q: hi\nA: Question time 🦀 ok", "Question:")]:
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        manager = DetokenizeManager(tokenizer, frozenset({248046}))
+        replies = []
+        for i, tid in enumerate(ids):
+            last = i == len(ids) - 1
+            msg = DetokenizeMsg(uid=42, next_token=tid, finished=last, stop_strs=[stop])
+            (out,), _, (offset,) = manager.detokenize_with_meta([msg])
+            replies.append(
+                reply(out, finished=last, text_offset=offset,
+                      logprobs=build_logprobs_entry(tokenizer, tid, -0.1, [], []))
+            )
+        req = CompletionRequest(model="client-model", prompt="p", logprobs=0, max_tokens=64, stop=[stop])
+        choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
+        assert choice["text"] == text
+        offsets = choice["logprobs"]["text_offset"]
+        for tid, offset in zip(ids, offsets, strict=True):
+            token = tokenizer.decode([tid])
+            if "\ufffd" not in token:
+                assert text[offset:].startswith(token), (text, token, offset)
+        assert offsets == sorted(offsets)
+        assert _joined_stream_logprobs(replies, logprobs=0, stop=[stop]) == choice["logprobs"]
+
+
+def test_completion_top_logprobs_key_collision_keeps_most_likely() -> None:
+    # Partial-UTF-8 candidates all decode to U+FFFD and collide on the dict key.
+    colliding = entry(9008, "\ufffd", -1.0)
+    colliding["top"] = [
+        {"token_id": 9008, "token": "\ufffd", "bytes": [240, 159], "logprob": -1.0},
+        {"token_id": 9009, "token": "\ufffd", "bytes": [240, 160], "logprob": -2.0},
+        {"token_id": 5, "token": "a", "bytes": [97], "logprob": -3.0},
+    ]
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=3, max_tokens=8)
+    result = run(handle_completion(req, None, FakeState([reply("x", finished=True, logprobs=colliding)]), {}))
+    assert result["choices"][0]["logprobs"]["top_logprobs"] == [{"\ufffd": -1.0, "a": -3.0}]
+
+
+def test_completion_stream_logprobs_match_non_stream() -> None:
+    # Acks with no visible text must not lose their entries: they ride the next chunk
+    # with text (or the finish chunk), so the joined stream equals the non-stream list.
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=1, max_tokens=8)
+    expected = run(handle_completion(req, None, FakeState(_crab_replies()), {}))
+    stream_req = CompletionRequest(model="client-model", prompt="hello", logprobs=1, max_tokens=8, stream=True)
+    events = parse_sse(run(_collect(stream_completion_chunks(42, stream_req, FakeState(_crab_replies())))))
+    joined: dict[str, list] = {"tokens": [], "token_logprobs": [], "top_logprobs": [], "text_offset": []}
+    for event in events:
+        if isinstance(event, dict) and event["choices"] and event["choices"][0]["logprobs"]:
+            for key in joined:
+                joined[key] += event["choices"][0]["logprobs"][key]
+    assert joined == expected["choices"][0]["logprobs"]
+    assert len(joined["tokens"]) == 5
+
+
 def test_logprobs_validation_errors() -> None:
     chat_top_out_of_range = run(
         handle_chat_completion(chat_request(logprobs=True, top_logprobs=25), None, FakeState([]), {})
diff --git a/tests/server/test_responses_api.py b/tests/server/test_responses_api.py
index 6e7721b891..95fbba3ef1 100644
--- a/tests/server/test_responses_api.py
+++ b/tests/server/test_responses_api.py
@@ -338,8 +338,9 @@ def test_stream_tool_call_events():
 # Route smoke tests
 # --------------------------------------------------------------------------- #
 class FakeState:
-    def __init__(self, outputs, finish_reason=None, cached_tokens=0, prefill_ms=0.0):
+    def __init__(self, outputs, finish_reason=None, cached_tokens=0, prefill_ms=0.0, reasoning_tokens=0):
         self._outputs = outputs
+        self._reasoning_tokens = reasoning_tokens  # stamped on the terminal ack
         self._finish_reason = finish_reason  # stamped on the terminal ack
         self._cached_tokens = cached_tokens  # stamped on the first ack (admission reply)
         self._prefill_ms = prefill_ms  # the scheduler stamps it on the first generated token
@@ -366,6 +367,7 @@ async def wait_for_ack(self, uid):
                             finish_reason=self._finish_reason if finished else None,
                             prompt_tokens_delta=pt, completion_tokens_delta=ct,
                             cached_tokens=self._cached_tokens if i == 0 else 0,
+                            reasoning_tokens=self._reasoning_tokens if finished else 0,
                             prefill_ms=self._prefill_ms if i == 0 else 0.0)
 
     async def stream_with_cancellation(self, gen, request, uid):
@@ -780,6 +782,27 @@ def test_stream_length_truncation_emits_incomplete_event():
     assert final["incomplete_details"]["reason"] == "max_output_tokens"
 
 
+def test_route_reports_reasoning_tokens():
+    # Codex shows usage.output_tokens_details.reasoning_tokens; it was hardcoded to 0.
+    def usage(stream):
+        client = _client(FakeState([("Hello", True, 5, 9)], reasoning_tokens=7))
+        r = client.post("/v1/responses", json={"model": "gpt-x", "input": "hi", "stream": stream})
+        assert r.status_code == 200, r.text
+        if not stream:
+            return r.json()["usage"]
+        final = [
+            json.loads(line[len("data:"):])
+            for line in r.text.splitlines()
+            if line.startswith("data:") and "response.completed" in line
+        ]
+        return final[-1]["response"]["usage"]
+
+    for stream in (False, True):
+        u = usage(stream)
+        assert u["output_tokens"] == 9
+        assert u["output_tokens_details"]["reasoning_tokens"] == 7
+
+
 def test_route_length_truncation_reports_incomplete():
     client = _client(FakeState([("partial", True, 5, 8192)], finish_reason="length"))
     r = client.post("/v1/responses", json={"model": "gpt-x", "input": "hi"})
diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py
index 4717617b9a..9767ffb25e 100644
--- a/tests/tokenizer/test_detokenize_extras.py
+++ b/tests/tokenizer/test_detokenize_extras.py
@@ -6,11 +6,12 @@
 from freetoken.message import DetokenizeMsg
 from freetoken.tokenizer.detokenize import DetokenizeManager
 
-EOS, SPECIAL, THINK_END = 0, 99, 7
+EOS, SPECIAL, THINK_END, THINK_START = 0, 99, 7, 8
 
 
 class _Tok:
-    """Token ids decode to letters; 99 is a special marker, 7 the reasoning end tag."""
+    """Token ids decode to letters; 99 is a special marker, 7 / 8 the reasoning end /
+    start tags."""
 
     eos_token_id = EOS
     unk_token_id = -1
@@ -23,6 +24,8 @@ def decode(self, ids, skip_special_tokens=False):
                     out.append("<s>")
             elif t == THINK_END:
                 out.append("</think>")
+            elif t == THINK_START:
+                out.append("<think>")
             else:
                 out.append(chr(ord("a") + t))
         return "".join(out)
@@ -39,7 +42,7 @@ def _run(manager, uid, tokens, *, finished_kwargs=None, **msg_kwargs):
         kwargs = dict(msg_kwargs)
         if last and finished_kwargs:
             kwargs.update(finished_kwargs)
-        strs, counts = manager.detokenize_with_meta(
+        strs, counts, _ = manager.detokenize_with_meta(
             [DetokenizeMsg(uid=uid, next_token=t, finished=last, **kwargs)]
         )
         text += strs[0]
@@ -81,6 +84,27 @@ def test_reasoning_tokens_count_up_to_the_end_tag():
     assert none == 0
 
 
+def test_reasoning_tokens_when_generation_stops_inside_the_block():
+    # The prompt pre-opened <think> (qwen3, thinking on) and the output never closed it:
+    # every completion token was reasoning, whether max_tokens or EOS ended it.
+    manager = DetokenizeManager(
+        _Tok(), frozenset({EOS}), think_end_id=THINK_END, think_start_id=THINK_START
+    )
+    _, by_length = _run(manager, 1, [1, 2, 3], starts_in_reasoning=True)
+    assert by_length == 3
+    _, by_eos = _run(manager, 2, [1, 2, EOS], starts_in_reasoning=True)
+    assert by_eos == 3
+    # the model opened the block itself
+    _, self_opened = _run(manager, 3, [THINK_START, 1, 2])
+    assert self_opened == 3
+    # thinking off: no block, no reasoning
+    _, off = _run(manager, 4, [1, 2, 3])
+    assert off == 0
+    # a closed block still counts up to the end tag
+    _, closed = _run(manager, 5, [1, THINK_END, 2, 3], starts_in_reasoning=True)
+    assert closed == 2
+
+
 def test_a_batch_with_one_uid_twice_advances_its_state_in_order():
     # each message must see the decode state the previous one left, and the reasoning
     # counts must stay aligned with the texts
@@ -91,7 +115,23 @@ def test_a_batch_with_one_uid_twice_advances_its_state_in_order():
         DetokenizeMsg(uid=2, next_token=3, finished=False),
         DetokenizeMsg(uid=1, next_token=2, finished=True),
     ]
-    texts, counts = manager.detokenize_with_meta(msgs)
+    texts, counts, offsets = manager.detokenize_with_meta(msgs)
     assert texts == ["b", "</think>", "d", "c"]
     assert counts == [0, 0, 0, 2]
+    assert offsets == [0, 1, 0, 9]
     assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"]
+
+
+def test_text_offset_is_where_the_token_starts_despite_holdback():
+    # "c" is held back as a possible start of the stop "cx", then sent with "d": its
+    # offset is still where it starts, not how much text had been sent at its ack.
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    sent, offsets = [], []
+    for i, t in enumerate([1, 2, 3, EOS]):
+        strs, _, offs = manager.detokenize_with_meta(
+            [DetokenizeMsg(uid=1, next_token=t, finished=i == 3, stop_strs=["cx"])]
+        )
+        sent.append(strs[0])
+        offsets.append(offs[0])
+    assert sent == ["b", "", "cd", ""]
+    assert offsets == [0, 1, 2, 3]
diff --git a/tests/tokenizer/test_logprobs_entry.py b/tests/tokenizer/test_logprobs_entry.py
index 89cf35cb5e..978824a34d 100644
--- a/tests/tokenizer/test_logprobs_entry.py
+++ b/tests/tokenizer/test_logprobs_entry.py
@@ -1,3 +1,7 @@
+import os
+
+import pytest
+
 from freetoken.tokenizer.detokenize import build_logprobs_entry
 
 
@@ -57,3 +61,53 @@ def test_build_logprobs_entry_empty_top() -> None:
     )
 
     assert entry["top"] == []
+
+
+def _byte_level_tokenizer():
+    """A tiny GPT-2-style ByteLevel BPE tokenizer built in memory (no download): the 256
+    byte stand-ins, one merge ("ðŁ", the first two bytes of most emoji) and one added
+    token whose content the byte table cannot invert."""
+    from tokenizers import AddedToken, Tokenizer, decoders, models, pre_tokenizers
+    from transformers import PreTrainedTokenizerFast
+
+    vocab = {c: i for i, c in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet()))}
+    vocab["ðŁ"] = len(vocab)
+    tk = Tokenizer(models.BPE(vocab=vocab, merges=[("ð", "Ł")]))
+    tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
+    tk.decoder = decoders.ByteLevel()
+    tk.add_tokens([AddedToken("<｜end｜>", special=False)])
+    return PreTrainedTokenizerFast(tokenizer_object=tk)
+
+
+def test_logprobs_bytes_rebuild_a_split_character() -> None:
+    # 🦀 is three byte-level tokens that each decode to U+FFFD; their bytes must still
+    # concatenate back to the character.
+    tokenizer = _byte_level_tokenizer()
+    ids = tokenizer.encode("🦀", add_special_tokens=False)
+    assert len(ids) == 3
+    entries = [build_logprobs_entry(tokenizer, i, -0.5, [i], [-0.5]) for i in ids]
+    assert all(e["token"] == "�" for e in entries)
+    assert bytes(b for e in entries for b in e["bytes"]) == "🦀".encode("utf-8")
+    assert [e["top"][0]["bytes"] for e in entries] == [e["bytes"] for e in entries]
+
+
+def test_logprobs_bytes_added_and_padded_ids() -> None:
+    tokenizer = _byte_level_tokenizer()
+    added = tokenizer.convert_tokens_to_ids("<｜end｜>")
+    padded = len(tokenizer) + 10  # the logits' padded tail: no vocab entry at all
+    entry = build_logprobs_entry(tokenizer, added, -0.1, [added, padded], [-0.1, -9.0])
+    assert entry["token"] == "<｜end｜>"
+    assert entry["bytes"] == list("<｜end｜>".encode("utf-8"))
+    assert entry["top"][1] == {"token_id": padded, "token": "", "bytes": [], "logprob": -9.0}
+
+
+def test_logprobs_bytes_real_qwen_tokenizer() -> None:
+    path = os.path.expanduser("~/models/Qwen3.8-Flash-Next-NVFP4-FTW-fp8")
+    if not os.path.exists(os.path.join(path, "tokenizer.json")):
+        pytest.skip("local Qwen3.8 tokenizer not present")
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    entries = [build_logprobs_entry(tokenizer, i, -0.1, [248300], [-9.0]) for i in (9008, 99, 222)]
+    assert bytes(b for e in entries for b in e["bytes"]) == "🦀".encode("utf-8")
+    assert entries[0]["top"][0]["bytes"] == []

From fadcf830b2558fad7622c9863f17e456785b03cc Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 11:02:00 -0400
Subject: [PATCH 224/249] merge fork: chat logprobs under the qwen3 and
 qwen3_coder parsers, and stream/non-stream parser consistency

With stop strings the detokenizer now holds a logprob entry until its text is
released, so one reply can carry several entries or none, and a per-reply
text_offset can no longer describe them. The detokenizer stamps each entry
with its own token's start instead, and UserReply.text_offset is gone;
completions reads the offset from the entry, keeps the pending carry-forward
for text-less replies, and drops stop_trimmed entries on both paths. The
non-stream chat choice keeps the always-present logprobs key, now with the
segmenter's {"content", "refusal"} object.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/message/frontend.py          |  13 +-
 python/freetoken/server/api_models.py         |   6 +-
 .../freetoken/server/function_call_parser.py  | 226 +++++--
 python/freetoken/server/generation.py         | 135 +++-
 python/freetoken/server/logprobs.py           | 200 +++++-
 python/freetoken/server/openai_api.py         |  74 ++-
 python/freetoken/server/reasoning_parser.py   | 111 +++-
 python/freetoken/tokenizer/detokenize.py      |  64 +-
 python/freetoken/tokenizer/server.py          |  35 +-
 tests/server/test_function_call_parser.py     |  29 +
 tests/server/test_logprobs_api.py             | 586 ++++++++++++++----
 .../test_qwen_semantic_special_tokens.py      |   4 +-
 tests/server/test_reasoning_parser.py         |  15 +
 tests/server/test_streaming_model_matrix.py   | 209 +++++++
 tests/tokenizer/test_detokenize_extras.py     | 109 ++++
 15 files changed, 1542 insertions(+), 274 deletions(-)

diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py
index 07634f1652..35620c7e8e 100644
--- a/python/freetoken/message/frontend.py
+++ b/python/freetoken/message/frontend.py
@@ -62,12 +62,13 @@ class UserReply(BaseFrontendMsg):
     # On the finished reply: tokens generated up to and including the reasoning end tag
     # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens.
     reasoning_tokens: int = 0
-    # Neutral sampled-token logprobs entry for this token (see
-    # tokenizer.detokenize.build_logprobs_entry); None when the request did not ask.
-    logprobs: dict | None = None
-    # Where this token's text starts in the generated text, even while a stop-string
-    # prefix holds it back (completions logprobs.text_offset).
-    text_offset: int = 0
+    # Neutral sampled-token logprobs entries released with this reply (see
+    # tokenizer.detokenize.build_logprobs_entry), normally this token's own; None when
+    # there are none. Each entry's "text_offset" is where its token starts in the generated
+    # text. With stop strings an entry waits for its token's text, and entries
+    # of tokens a matched stop string trimmed arrive marked "stop_trimmed"; the token it
+    # ends inside carries "stop_visible", the length of its text that was released.
+    logprobs: list[dict] | None = None
 
 
 @dataclass
diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 21b1d1f860..0da3d09e77 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -77,9 +77,9 @@ class ChatCompletionRequest(BaseModel):
     presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires
-    # logprobs=true; entries cover generated tokens only (no prompt logprobs).
-    # Fail-closed (400) when a reasoning parser or tool parsing is active: entries
-    # must align 1:1 with visible content tokens (see _chat_logprobs_conflict).
+    # logprobs=true; entries cover message-content tokens only (no prompt, reasoning or
+    # tool-call tokens). 400 for parser families whose hidden tokens the segmenter
+    # cannot tell apart from content (see _chat_logprobs_conflict).
     logprobs: bool = False
     top_logprobs: int | None = None
     chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)
diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py
index 88cc65b367..e45acea77b 100644
--- a/python/freetoken/server/function_call_parser.py
+++ b/python/freetoken/server/function_call_parser.py
@@ -261,6 +261,11 @@ class BaseFormatDetector(ABC):
     # bot_token is a parse trigger, not a uniqueness claim.
     toolcall_opener: str | None = None
 
+    # True when detect_and_parse replays the streaming state machine: its normal_text
+    # is then already every piece of text the stream emits (before, between and after
+    # calls), so parse_non_stream must return it as is.
+    one_shot_replays_stream = False
+
     def __init__(self):
         # Streaming state management
         # Buffer for accumulating incomplete patterns that arrive across multiple streaming chunks
@@ -778,6 +783,12 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
 
 
 
+# A bare invoke opener's name: bounded so a literal "<function=" in prose or code
+# is released after at most this many chars instead of being held.
+_PS_BARE_NAME_RE = re.compile(r"([^\s<>]{1,128})>\s*")
+_PS_BARE_NAME_PARTIAL_RE = re.compile(r"[^\s<>]{0,128}")
+
+
 class InvokeParamStreamMixin:
     """Value-level streaming for invoke/parameter block formats (qwen3_coder,
     minimax — the same shape vLLM's ParserEngine covers): text outside blocks
@@ -793,6 +804,13 @@ class InvokeParamStreamMixin:
       _ps_trim_single  True: at most ONE leading/trailing trim char (qwen3_coder)
       _ps_missing_type schema type assumed for undeclared params ("string" streams,
                        "loose" buffers and loose-parses — per-family legacy typing)
+      _ps_trim_after_block  True: whitespace after a closed block is held, then
+                       dropped if another block or the end follows, kept if
+                       text does -- the same for any chunking (qwen3_coder)
+      _ps_bare_invoke  True: an invoke without the wrapper block is a call when it
+                       names a declared tool and its body starts (qwen3_coder)
+      _ps_literal_open  True: a block opener whose next non-whitespace text cannot
+                       start an invoke or the block close is text (qwen3_coder)
     """
 
     _ps_outer_open: str = ""
@@ -800,9 +818,15 @@ class InvokeParamStreamMixin:
     _ps_trim: str = "\n"
     _ps_trim_single: bool = False
     _ps_missing_type: str = "string"
+    _ps_trim_after_block: bool = False
+    _ps_bare_invoke: bool = False
+    _ps_literal_open: bool = False
 
     def _ps_reset(self) -> None:
         self._ps_mode = "idle"
+        self._ps_after_block = False
+        self._ps_bare = False
+        self._ps_after_bare = False
         self._ps_key = ""
         self._ps_lead = "{"
         self._ps_emitted_any = False
@@ -829,6 +853,58 @@ def _ps_trim_trailing(self, text: str) -> str:
             return text[:-1] if text[-1:] and text[-1:] in self._ps_trim else text
         return text.rstrip(self._ps_trim)
 
+    def _ps_bare_status(self, text: str) -> str:
+        """Classify the bare invoke opener ``text`` starts with: "call", "partial"
+        (may still become one) or "literal". A declared name is required -- not
+        the unknown-tool forwarding a wrapped block gets -- and the body must
+        start, so the format quoted in prose or code stays text."""
+        rest = text[len(self._ps_invoke_open_prefix):]
+        m = _PS_BARE_NAME_RE.match(rest)
+        if m is None:
+            return "partial" if _PS_BARE_NAME_PARTIAL_RE.fullmatch(rest) else "literal"
+        if m.group(1) not in self._tool_indices:
+            return "literal"
+        body = rest[m.end():]
+        for tok in (self._ps_param_open_prefix, self._ps_invoke_close):
+            if body.startswith(tok):
+                return "call"
+            if tok.startswith(body):
+                return "partial"
+        return "literal"
+
+    def _ps_find_bare(self, buf: str, limit: int) -> Tuple[int, str]:
+        """First bare invoke opener before ``limit`` that is not literal."""
+        i = buf.find(self._ps_invoke_open_prefix)
+        while i != -1 and i < limit:
+            status = self._ps_bare_status(buf[i:])
+            if status != "literal":
+                return i, status
+            i = buf.find(self._ps_invoke_open_prefix, i + 1)
+        return -1, ""
+
+    def _ps_opens_block(self, text: str) -> str:
+        """Whether ``text`` starts a block: "yes", "no" or "maybe" (too short to tell)."""
+        if text.startswith(self._ps_outer_open):
+            if not self._ps_literal_open:
+                return "yes"
+            body = text[len(self._ps_outer_open):].lstrip()
+            # A repeated opener is a stutter: the block ignores it, as before.
+            for tok in (self._ps_invoke_open_prefix, self._ps_outer_close, self._ps_outer_open):
+                if body.startswith(tok):
+                    return "yes"
+                if tok.startswith(body):
+                    return "maybe"
+            return "no"
+        if self._ps_outer_open.startswith(text):
+            return "maybe"
+        if self._ps_bare_invoke:
+            prefix = self._ps_invoke_open_prefix
+            if text.startswith(prefix):
+                return {"call": "yes", "partial": "maybe"}.get(self._ps_bare_status(text), "no")
+            if prefix.startswith(text):
+                return "maybe"
+        return "no"
+
     def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult:
         self._buffer += new_text
         if not hasattr(self, "_tool_indices"):
@@ -866,9 +942,42 @@ def _update_prev() -> None:
             if mode == "idle":
                 if calls:
                     break  # text after a call defers to the next step (wire order)
+                if self._ps_after_block:
+                    rest = buf.lstrip()
+                    if self._ps_after_bare:
+                        if rest.startswith(self._ps_outer_close):
+                            # The model dropped the opener but kept the closer.
+                            self._ps_after_bare = False
+                            self._buffer = rest[len(self._ps_outer_close):]
+                            continue
+                        if self._ps_outer_close.startswith(rest):
+                            break
+                        self._ps_after_bare = False
+                    opens = self._ps_opens_block(rest)
+                    if opens == "maybe":
+                        break  # a separator or a text gap: undecided until more arrives
+                    self._ps_after_block = False
+                    if opens == "yes":
+                        buf = self._buffer = rest
                 pos = buf.find(self._ps_outer_open)
+                if self._ps_bare_invoke:
+                    bare, status = self._ps_find_bare(buf, len(buf) if pos == -1 else pos)
+                    if bare > 0:
+                        normal_parts.append(buf[:bare])
+                        self._buffer = buf[bare:]
+                        continue
+                    if bare == 0:
+                        if status == "partial":
+                            break
+                        self._ps_bare = True
+                        self._ps_mode = "block"  # parses the invoke at the buffer head
+                        continue
                 if pos == -1:
                     hold = self._ends_with_partial_token(buf, self._ps_outer_open)
+                    if self._ps_bare_invoke:
+                        hold = max(
+                            hold, self._ends_with_partial_token(buf, self._ps_invoke_open_prefix)
+                        )
                     release = buf[: len(buf) - hold] if hold else buf
                     if release:
                         normal_parts.append(release)
@@ -878,7 +987,13 @@ def _update_prev() -> None:
                     normal_parts.append(buf[:pos])
                     self._buffer = buf[pos:]
                     continue
+                opens = self._ps_opens_block(buf)
+                if opens == "maybe":
+                    break
                 self._buffer = buf[len(self._ps_outer_open):]
+                if opens == "no":
+                    normal_parts.append(self._ps_outer_open)  # the tag named in prose
+                    continue
                 self._ps_mode = "block"
                 continue
 
@@ -888,6 +1003,7 @@ def _update_prev() -> None:
                 if close != -1 and (inv == -1 or close < inv):
                     self._buffer = buf[close + len(self._ps_outer_close):]
                     self._ps_mode = "idle"
+                    self._ps_after_block = self._ps_trim_after_block
                     continue
                 if inv != -1:
                     m = self._ps_invoke_open_re.search(buf, inv)
@@ -939,7 +1055,14 @@ def _update_prev() -> None:
                         while len(self.streamed_args_for_tool) <= self.current_tool_id:
                             self.streamed_args_for_tool.append("")
                     self._buffer = buf[e + len(self._ps_invoke_close):]
-                    self._ps_mode = "block"
+                    if self._ps_bare:
+                        # A bare invoke is its own block: back to text after it.
+                        self._ps_bare = False
+                        self._ps_mode = "idle"
+                        self._ps_after_block = self._ps_trim_after_block
+                        self._ps_after_bare = True
+                    else:
+                        self._ps_mode = "block"
                     continue
                 if p != -1:
                     m = self._ps_param_open_re.search(buf, p)
@@ -1023,6 +1146,8 @@ def _update_prev() -> None:
     def finish_streaming(self) -> str:
         residual, self._buffer = self._buffer, ""
         mode = getattr(self, "_ps_mode", "idle")
+        if getattr(self, "_ps_after_block", False) and not residual.strip():
+            residual = ""
         self._ps_reset()
         if mode != "idle" or (self._ps_outer_open and self._ps_outer_open in residual):
             return ""
@@ -2016,9 +2141,13 @@ def finish_streaming(self) -> str:
 
 class Qwen3CoderDetector(InvokeParamStreamMixin, BaseFormatDetector):
     toolcall_opener = "<tool_call>"
+    one_shot_replays_stream = True
     _ps_trim = "\n"
     _ps_trim_single = True
     _ps_missing_type = "string"
+    _ps_trim_after_block = True
+    _ps_bare_invoke = True
+    _ps_literal_open = True
 
     """
     Detector for Qwen3-Coder XML-style function call format.
@@ -2052,8 +2181,6 @@ def __init__(self):
         self.tool_call_separator = "\n"
 
         # Regex patterns
-        self.tool_call_block_regex = re.compile(r"<tool_call>(.*?)</tool_call>", re.DOTALL)
-        self.function_regex = re.compile(r"<function=(.*?)</function>|<function=(.*)$", re.DOTALL)
         self.parameter_regex = re.compile(
             r"<parameter=(.*?)(?:</parameter>|(?=<parameter=)|(?=</function>)|$)", re.DOTALL
         )
@@ -2063,7 +2190,8 @@ def __init__(self):
         self._ps_outer_open = "<tool_call>"
         self._ps_outer_close = "</tool_call>"
         self._ps_invoke_open_prefix = "<function="
-        self._ps_invoke_open_re = re.compile(r"<function=([^>]*)>")
+        # No "<" in a name: an opener cut off before its ">" must not run into </tool_call>.
+        self._ps_invoke_open_re = re.compile(r"<function=([^<>]*)>")
         self._ps_invoke_close = "</function>"
         self._ps_param_open_prefix = "<parameter="
         self._ps_param_open_re = re.compile(r"<parameter=([^>]*)>")
@@ -2074,44 +2202,6 @@ def has_tool_call(self, text: str) -> bool:
         return "<function=" in text or self.bot_token in text
 
 
-    def _parse_function_call(self, function_str: str, tools: List[Tool]) -> Optional[ToolCallItem]:
-        """Parse a single <function=name>...</function> block into a ToolCallItem."""
-        try:
-            end_index = function_str.index(">")
-        except ValueError:
-            return None
-
-        func_name = function_str[:end_index].strip()
-        tool_indices = self._get_tool_indices(tools)
-        if func_name not in tool_indices and not _should_forward_unknown_tool(func_name):
-            logger.warning(f"Model attempted to call undefined function: {func_name}")
-            return None
-
-        parameters_text = function_str[end_index + 1 :]
-        param_config = self._get_param_config(func_name, tools)
-        param_dict = {}
-
-        for match in self.parameter_regex.findall(parameters_text):
-            try:
-                idx = match.index(">")
-            except ValueError:
-                continue
-            param_name = match[:idx].strip()
-            param_value = match[idx + 1 :]
-            # Strip leading/trailing newlines from value
-            if param_value.startswith("\n"):
-                param_value = param_value[1:]
-            if param_value.endswith("\n"):
-                param_value = param_value[:-1]
-
-            param_dict[param_name] = self._convert_param_value(param_value, param_name, param_config, func_name)
-
-        return ToolCallItem(
-            tool_index=tool_indices.get(func_name, 0),
-            name=func_name,
-            parameters=json.dumps(param_dict, ensure_ascii=False),
-        )
-
     def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools: List[Tool]) -> Optional[str]:
         """Build the current argument JSON from a partial XML tool-call body."""
         param_matches = self.parameter_regex.findall(partial_body)
@@ -2155,28 +2245,38 @@ def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools
         return json.dumps(param_dict, ensure_ascii=False)
 
     def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
-        idx = text.find(self.bot_token)
-        normal_text = text[:idx].strip() if idx != -1 else text
-
-        if "<function=" not in text:
-            return StreamingParseResult(normal_text=normal_text, calls=[])
-
-        # Extract function blocks from tool_call blocks (or raw text as fallback)
-        tool_call_blocks = self.tool_call_block_regex.findall(text)
-        if not tool_call_blocks:
-            tool_call_blocks = [text]
+        """One-shot parse by replaying the streaming state machine on a fresh
+        detector, so text outside blocks (before, between and after calls, in
+        order) and the calls are exactly what streaming emits for any chunking."""
+        shadow = type(self)()
+        normal: List[str] = []
+        opened: List[List] = []  # [detector tool index, name, argument fragments]
+        chunk = text
+        while True:
+            step = shadow.parse_streaming_increment(chunk, tools)
+            chunk = ""
+            normal.append(step.normal_text)
+            for item in step.calls:
+                if item.name is not None:
+                    opened.append([item.tool_index, item.name, ""])
+                elif opened:
+                    opened[-1][2] += item.parameters
+            if not step.normal_text and not step.calls:
+                break
+        normal.append(shadow.finish_streaming())
 
         calls = []
-        for block in tool_call_blocks:
-            func_matches = self.function_regex.findall(block)
-            for match in func_matches:
-                func_str = match[0] if match[0] else match[1]
-                item = self._parse_function_call(func_str, tools)
-                if item:
-                    item.tool_index = len(calls)
-                    calls.append(item)
-
-        return StreamingParseResult(normal_text=normal_text, calls=calls)
+        for index, name, params in opened:
+            try:
+                json.loads(params)
+            except (json.JSONDecodeError, ValueError):
+                # Cut off mid-call: same fallback as the serving layer's stream close.
+                ledger = shadow.prev_tool_call_arr
+                args = ledger[index].get("arguments") if 0 <= index < len(ledger) else None
+                if args is not None:
+                    params = json.dumps(args, ensure_ascii=False)
+            calls.append(ToolCallItem(tool_index=len(calls), name=name, parameters=params or "{}"))
+        return StreamingParseResult(normal_text="".join(normal), calls=calls)
 
 
 class Gemma4Detector(BaseFormatDetector):
@@ -3733,6 +3833,8 @@ def parse_non_stream(self, full_text: str) -> StreamingParseResult:
         if not self.tools:
             return StreamingParseResult(normal_text=full_text, calls=[])
         parsed_result = self.detector.detect_and_parse(full_text, self.tools)
+        if self.detector.one_shot_replays_stream:
+            return parsed_result
         tool_call_list = parsed_result.calls
         if tool_call_list:
             # Keep text after the LAST tool block: the streaming path emits it, so
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 00d94abebf..5a9192fbd7 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -37,8 +37,16 @@
     _TemplateError = ()
 
 from .function_call_parser import FunctionCallParser, TOOLS_TAG_LIST, ToolCallItem
+from .logprobs import (
+    LogprobMarkers,
+    LogprobSegmenter,
+    resolve_logprob_markers,
+    untrimmed,
+)
 from .reasoning_parser import (
     DSV4_SPECIAL_TOKENS,
+    THINK_END_TOKEN,
+    THINK_START_TOKEN,
     ReasoningParser,
     build_reasoning_parser,
     strip_special_tokens,
@@ -386,6 +394,9 @@ class GenDone:
     cached_tokens: int = 0
     reasoning_tokens: int = 0
     timings: GenTimings = field(default_factory=GenTimings)
+    # Content logprob entries no content delta carried (text held to the end, or
+    # dropped as a trailing separator); None when there are none.
+    logprobs: list[dict] | None = None
 
 
 GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone
@@ -403,8 +414,8 @@ class GenResult:
     cached_tokens: int = 0
     reasoning_tokens: int = 0
     timings: GenTimings = field(default_factory=GenTimings)
-    # Neutral sampled-token logprob entries, one per sampled token (empty when the
-    # request did not ask).
+    # Neutral logprob entries of the message-content tokens (empty when the request did
+    # not ask).
     logprobs: list[dict] = field(default_factory=list)
 
 
@@ -740,8 +751,13 @@ async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None:
 
 def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None:
     """Build a reasoning parser for this generation, or None if the server has no
-    reasoning parser configured. ``force_reasoning`` matches the encode-side
-    thinking mode so chat-mode content is never mislabeled as reasoning."""
+    reasoning parser configured."""
+    return build_reasoning_parser(state.config, _force_reasoning(spec, state))
+
+
+def _force_reasoning(spec: GenSpec, state: Any) -> bool:
+    """Whether the generation starts inside reasoning: matches the encode-side thinking
+    mode so chat-mode content is never mislabeled as reasoning."""
     parser_name = getattr(state.config, "reasoning_parser", None)
     if parser_name == "qwen3":
         # The qwen3 chat template opens an implicit <think> (thinking on) unless
@@ -776,7 +792,7 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None:
         force_reasoning = (
             resolve_thinking_mode(spec.chat_template_kwargs, spec.template_tools) == "thinking"
         )
-    return build_reasoning_parser(state.config, force_reasoning)
+    return force_reasoning
 
 
 def _split_reasoning(text: str, spec: GenSpec, state: Any) -> tuple[str, str]:
@@ -817,11 +833,75 @@ def _parse_tool_response(
         return None
     parser = _make_tool_parser(spec, state)
     result = parser.parse_non_stream(text)
-    if not result.calls:
+    # A detector that replays its stream can drop markup without producing a call
+    # (an empty <tool_call> block): the one-shot content must drop it too.
+    if not result.calls and result.normal_text == text:
         return None
     return result.normal_text, result.calls
 
 
+_UNRESOLVED = object()
+
+
+async def chat_logprob_markers(state: Any) -> LogprobMarkers | None:
+    """The token ids chat logprobs are segmented on, resolved once per server from the
+    frontend tokenizer; None when this server has none."""
+    cached = getattr(state, "_logprob_markers", _UNRESOLVED)
+    if cached is not _UNRESOLVED:
+        return cached
+    markers = None
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is not None:
+        try:
+            manager = await asyncio.to_thread(build)
+            hidden = _QWEN_SEMANTIC_HIDDEN_TOKENS if _uses_qwen_semantic_protocol(state) else ()
+            markers = await asyncio.to_thread(
+                resolve_logprob_markers,
+                manager.tokenizer, THINK_START_TOKEN, THINK_END_TOKEN, hidden,
+            )
+        except Exception:  # noqa: BLE001 -- server fault: retried on the next request
+            return None
+    state._logprob_markers = markers
+    return markers
+
+
+def chat_logprobs_supported(state: Any, parse_tools: bool, markers: LogprobMarkers | None) -> bool:
+    """Whether chat logprobs can be confined to message content for this server's parsers:
+    the qwen3 reasoning split needs single-id think markers, and tool markup needs a
+    detector whose one-shot parse replays its stream (the segmenter shadows that stream)."""
+    config = state.config
+    reasoning = getattr(config, "reasoning_parser", None)
+    if reasoning not in (None, "qwen3"):
+        return False
+    if reasoning and (markers is None or markers.think_start is None or markers.think_end is None):
+        return False
+    if parse_tools:
+        name = getattr(config, "tool_call_parser", None)
+        detector = FunctionCallParser.ToolCallParserEnum.get(name)
+        if detector is None or not (detector.one_shot_replays_stream and detector.supports_streaming):
+            return False
+    return True
+
+
+def _wants_logprobs(spec: GenSpec) -> bool:
+    return bool(getattr(spec.sampling_params, "logprobs", False))
+
+
+async def _make_logprob_segmenter(spec: GenSpec, state: Any) -> LogprobSegmenter | None:
+    """The segmenter for a request that asked for logprobs, or None when its entries
+    cannot be confined to message content (they are then all dropped)."""
+    markers = await chat_logprob_markers(state)
+    if not chat_logprobs_supported(state, spec.parse_tools, markers):
+        return None
+    reasoning = getattr(state.config, "reasoning_parser", None) is not None
+    return LogprobSegmenter(
+        markers or LogprobMarkers(None, None, frozenset()),
+        reasoning=reasoning,
+        force_reasoning=reasoning and _force_reasoning(spec, state),
+        tool_parser=_make_tool_parser(spec, state) if spec.parse_tools else None,
+    )
+
+
 def _valid_json(text: str) -> bool:
     if not text:
         return False
@@ -1065,6 +1145,10 @@ def _content_delta(text: str) -> ContentDelta:
         pending_logprobs = []
         return ContentDelta(text, logprobs=logprobs)
 
+    # Content entries ride the next content delta; reasoning and tool-call entries are
+    # never collected, so a hidden token's entry cannot surface on a visible chunk.
+    segmenter = await _make_logprob_segmenter(spec, state) if _wants_logprobs(spec) else None
+
     tool_parser: FunctionCallParser | None = None
     if parse_tools:
         try:
@@ -1074,14 +1158,9 @@ def _content_delta(text: str) -> ContentDelta:
         if candidate is not None and candidate.supports_streaming():
             tool_parser = candidate
     frag_stable = tool_parser.args_fragments_prefix_stable() if tool_parser else True
-
-    # Logprob entries are collected only on the passthrough path: with a reasoning
-    # parser or tool parsing active, text can be hidden, held back, or reclassified,
-    # so an entry could come to describe a token that never reaches visible content.
-    # The API layer fail-closes those request combinations; this guard is the
-    # structural half of the contract — a hidden-token entry must never ride a
-    # later visible delta, whatever the caller.
-    collect_logprobs = reasoning_parser is None and tool_parser is None and not parse_tools
+    # A detector that replays its stream for the one-shot parse already decides which
+    # whitespace after a call is a separator; dropping more here would split the modes.
+    detector_owns_ws = tool_parser is not None and tool_parser.detector.one_shot_replays_stream
 
     # Streaming tool-call assembly: detectors emit fragments (name first, then
     # argument diffs); they accumulate here and the call is emitted complete when
@@ -1105,7 +1184,7 @@ def _semantic_content_events(items: list[Any]) -> list[GenEvent]:
             else:
                 out.append(item)
                 if isinstance(item, ToolCallsDelta):
-                    suppress_ws = True
+                    suppress_ws = not detector_owns_ws
         return out
 
     def _filter_tool_text(piece: str) -> list[GenEvent]:
@@ -1183,8 +1262,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
-        if collect_logprobs and getattr(ack, "logprobs", None) is not None:
-            pending_logprobs.append(ack.logprobs)
+        if segmenter is not None:
+            pending_logprobs.extend(segmenter.feed(untrimmed(getattr(ack, "logprobs", None))))
         content_delta = ack.incremental_output
         if reasoning_parser is not None and content_delta:
             reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta)
@@ -1207,6 +1286,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             engine_matched_stop = getattr(ack, "matched_stop", None)
             break
 
+    if segmenter is not None:
+        pending_logprobs.extend(segmenter.finish())
+
     # Drain residue held in the reasoning parser (a deferred tool block, or a
     # trailing partial token) so it is not silently dropped.
     if reasoning_parser is not None:
@@ -1274,9 +1356,10 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             normal_text = _filter_content(normal_text)
             if normal_text:
                 yield _content_delta(normal_text)
-            yield ToolCallsDelta(tool_calls)
-            if finish_reason != "length":
-                finish_reason = "tool_calls"
+            if tool_calls:
+                yield ToolCallsDelta(tool_calls)
+                if finish_reason != "length":
+                    finish_reason = "tool_calls"
         elif parse_tools and pending:
             stripped_pending = _filter_content(pending)
             if stripped_pending:
@@ -1293,12 +1376,12 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         if trailing_content:
             yield _content_delta(trailing_content)
 
-    # Entries without a content delta are intentionally dropped in streaming mode.
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
         matched_stop=engine_matched_stop, cached_tokens=cached_tokens,
         reasoning_tokens=reasoning_tokens,
         timings=timer.finish(),
+        logprobs=pending_logprobs or None,
     )
 
 
@@ -1322,13 +1405,17 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
         full_content += ack.incremental_output
-        if getattr(ack, "logprobs", None) is not None:
-            logprob_entries.append(ack.logprobs)
+        logprob_entries.extend(untrimmed(getattr(ack, "logprobs", None)))
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
             break
 
+    if _wants_logprobs(spec):
+        segmenter = await _make_logprob_segmenter(spec, state)
+        logprob_entries = (
+            segmenter.feed(logprob_entries) + segmenter.finish() if segmenter is not None else []
+        )
     reasoning_text, content_text = _split_reasoning(full_content, spec, state)
     # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins.
     finish_reason = engine_finish_reason or "stop"
@@ -1336,7 +1423,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
     parsed = _parse_tool_response(content_text, spec, state)
     if parsed is not None:
         content_text, tool_calls = parsed
-        if finish_reason != "length":
+        if tool_calls and finish_reason != "length":
             finish_reason = "tool_calls"
 
     specials = _leaked_special_tokens(state)
diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py
index 0cf6198aef..c3d22141a7 100644
--- a/python/freetoken/server/logprobs.py
+++ b/python/freetoken/server/logprobs.py
@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from collections import deque
+from dataclasses import dataclass
 from typing import Any
 
 
@@ -21,6 +23,11 @@ def completion_logprobs_error(req: Any) -> str | None:
     return None
 
 
+def untrimmed(entries: list[dict] | None) -> list[dict]:
+    """A reply's entries minus those of tokens a matched stop string cut from the text."""
+    return [entry for entry in entries or () if not entry.get("stop_trimmed")]
+
+
 def chat_content_entry(entry: dict) -> dict:
     return {
         "token": entry["token"],
@@ -37,12 +44,11 @@ def chat_content_entry(entry: dict) -> dict:
     }
 
 
-def completions_logprobs(
-    entries: list[dict], offsets: list[int], text_len: int | None = None
-) -> dict:
-    """``offsets[i]`` is where entry i's token starts in the generated text (the pieces of
-    a split multi-byte character share its start). Given the finished text's ``text_len``,
-    offsets are clipped to it, since a matched stop string trims tokens off the end."""
+def completions_logprobs(entries: list[dict], text_len: int | None = None) -> dict:
+    """Each entry's detokenizer-stamped ``text_offset`` is where its token starts in the
+    generated text (the pieces of a split multi-byte character share its start). Given the
+    finished text's ``text_len``, offsets are clipped to it, since a matched stop string
+    trims tokens off the end."""
     tokens: list[str] = []
     token_logprobs: list[float] = []
     top_logprobs: list[dict[str, float]] = []
@@ -63,5 +69,185 @@ def completions_logprobs(
         "tokens": tokens,
         "token_logprobs": token_logprobs,
         "top_logprobs": top_logprobs,
-        "text_offset": [offset if text_len is None else min(offset, text_len) for offset in offsets],
+        "text_offset": [
+            entry["text_offset"] if text_len is None else min(entry["text_offset"], text_len)
+            for entry in entries
+        ],
     }
+
+
+def chat_logprobs_object(entries: list[dict]) -> dict:
+    return {"content": [chat_content_entry(entry) for entry in entries], "refusal": None}
+
+
+@dataclass(frozen=True)
+class LogprobMarkers:
+    """Token ids the chat segmenter keys on, resolved once per server. A think marker is
+    None unless the tokenizer encodes it as exactly one id that decodes back to it."""
+
+    think_start: int | None
+    think_end: int | None
+    # Special and protocol-hidden tokens: never message content.
+    hidden: frozenset[int]
+
+
+def resolve_logprob_markers(
+    tokenizer: Any, think_start: str, think_end: str, hidden_tokens: tuple[str, ...]
+) -> LogprobMarkers:
+    def single_id(marker: str) -> int | None:
+        ids = tokenizer.encode(marker, add_special_tokens=False)
+        if len(ids) == 1 and tokenizer.decode(ids) == marker:
+            return int(ids[0])
+        return None
+
+    hidden = {int(tid) for tid in getattr(tokenizer, "all_special_ids", None) or ()}
+    hidden.update(tid for tid in map(single_id, hidden_tokens) if tid is not None)
+    return LogprobMarkers(single_id(think_start), single_id(think_end), frozenset(hidden))
+
+
+def _visible_text(entry: dict) -> str:
+    """The part of an entry's token that reached the output text."""
+    return entry["token"][: entry.get("stop_visible")]
+
+
+class LogprobSegmenter:
+    """Picks a chat generation's message-content logprob entries out of the sampled ones.
+
+    A pure function of the entry sequence (token ids and single-token strings), so the
+    streaming and non-streaming paths select the same entries for any ack chunking.
+    The reasoning split is keyed on the think marker ids and mirrors the qwen3 parser:
+    a leading ``<think>`` (after whitespace) opens reasoning once, the first ``</think>``
+    ends it for good, and content after it is whitespace-trimmed at both ends. Tool
+    markup is decided by a shadow tool-call parser fed one token string at a time: its
+    held buffer is always a suffix of what it was fed, so every resolved character is
+    either emitted as text or consumed as markup, and a token is content when any of
+    its characters is emitted.
+    """
+
+    def __init__(
+        self,
+        markers: LogprobMarkers,
+        *,
+        reasoning: bool,
+        force_reasoning: bool,
+        tool_parser: Any | None = None,
+    ) -> None:
+        self._markers = markers
+        self._state = "head" if reasoning else "content"
+        self._force = force_reasoning
+        self._trim = False  # content follows a real </think>: trim it at both ends
+        self._content_started = False
+        self._held: list[dict] = []  # whitespace-only tokens awaiting the next text
+        self._tools = tool_parser
+        self._queue: deque[list] = deque()  # [entry, start, end, is_content] being resolved
+        self._fed = 0
+        self._out: list[dict] = []
+
+    def feed(self, entries: list[dict]) -> list[dict]:
+        """Content entries that became final, in generation order."""
+        for entry in entries:
+            # Hidden ids skip the parsers whole: the final EOS adds no text to the output.
+            if entry["token_id"] not in self._markers.hidden:
+                self._split_reasoning(entry)
+        out, self._out = self._out, []
+        return out
+
+    def finish(self) -> list[dict]:
+        held, self._held = self._held, []
+        if self._state == "head" and not self._force:
+            for entry in held:
+                self._content(entry)
+        if self._tools is not None:
+            det = self._tools.detector
+            base = self._fed - len(det._buffer)
+            region = det._buffer
+            if self._tools.recover_truncated_call():
+                self._resolve(base, region, "")
+                base, region = self._fed, ""
+            self._resolve(base, region, self._tools.finish_stream())
+        out, self._out = self._out, []
+        return out
+
+    def _split_reasoning(self, entry: dict) -> None:
+        # A token a stop string cut inside is only its visible text, never a marker.
+        tid = -1 if "stop_visible" in entry else entry["token_id"]
+        text = _visible_text(entry)
+        blank = not text.strip()
+        if self._state == "head":
+            if tid == self._markers.think_start:
+                self._held = []
+                self._state = "reasoning"
+            elif blank:
+                self._held.append(entry)
+            elif self._force:
+                self._held = []
+                self._state = "reasoning"
+                self._split_reasoning(entry)
+            else:
+                self._state = "content"
+                held, self._held = self._held, []
+                for e in held + [entry]:
+                    self._content(e)
+        elif self._state == "reasoning":
+            if tid == self._markers.think_end:
+                self._state = "content"
+                self._trim = True
+        elif not self._trim:
+            self._content(entry)
+        elif blank:
+            if self._content_started:
+                self._held.append(entry)
+        elif not self._content_started:
+            # The reasoning parser strips this leading whitespace before the tool
+            # parser sees it: " <" opening a bare call is markup, not a space.
+            self._content_started = True
+            self._content(entry, text.lstrip())
+        else:
+            held, self._held = self._held, []
+            for e in held + [entry]:
+                self._content(e)
+
+    def _content(self, entry: dict, text: str | None = None) -> None:
+        if self._tools is None:
+            self._out.append(entry)
+            return
+        text = _visible_text(entry) if text is None else text
+        held = self._tools.detector._buffer
+        base = self._fed - len(held)
+        self._queue.append([entry, self._fed, self._fed + len(text), False])
+        self._fed += len(text)
+        normal = "".join(
+            payload for kind, payload in self._tools.parse_stream_events(text) if kind == "text"
+        )
+        resolved = self._fed - len(self._tools.detector._buffer) - base
+        region = (held + text)[:resolved] if normal else ""
+        self._resolve(base, region, normal, resolved)
+
+    def _resolve(self, base: int, region: str, normal: str, length: int | None = None) -> None:
+        """``region`` (``length`` chars from offset ``base``) left the parser's buffer,
+        emitting ``normal``: mark the tokens it emitted from, then release resolved ones.
+        ``region`` itself is only needed when something was emitted."""
+        if normal:
+            if region.startswith(normal):
+                emitted = [(0, len(normal))]
+            elif region.endswith(normal):
+                # Text after markup in one token, as when a bare call's ">" merges with
+                # the next character (">n"): a leftmost match could land inside the markup.
+                emitted = [(len(region) - len(normal), len(region))]
+            else:  # not seen in practice: fall back to the leftmost alignment
+                emitted, i = [], 0
+                for ch in normal:
+                    j = region.find(ch, i)
+                    if j == -1:
+                        break
+                    emitted.append((j, j + 1))
+                    i = j + 1
+            for item in self._queue:
+                start, end = item[1] - base, item[2] - base
+                if any(lo < end and start < hi for lo, hi in emitted):
+                    item[3] = True
+        resolved = base + (len(region) if length is None else length)
+        while self._queue and self._queue[0][2] <= resolved:
+            entry, _, _, is_content = self._queue.popleft()
+            if is_content:
+                self._out.append(entry)
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index cd7f398c1d..4b43074b6e 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -27,10 +27,11 @@
 from .function_call_parser import ToolCallItem
 from .request_logger import log_request
 from .logprobs import (
-    chat_content_entry,
     chat_logprobs_error,
+    chat_logprobs_object,
     completion_logprobs_error,
     completions_logprobs,
+    untrimmed,
 )
 from .generation import (
     DEFAULT_MAX_OUTPUT_TOKENS,
@@ -44,6 +45,8 @@
     ToolCallsDelta,
     ToolCallStart,
     build_metrics,
+    chat_logprob_markers,
+    chat_logprobs_supported,
     generate_events,
     generate_full,
     metrics_enabled,
@@ -261,27 +264,25 @@ def _client_disconnected_response() -> JSONResponse:
     )
 
 
-def _chat_logprobs_conflict(req: ChatCompletionRequest, state: Any) -> str | None:
-    """Fail-close chat logprobs when a semantic parsing layer can hide or reclassify
-    generated tokens. A reasoning parser routes tokens out of visible content (and its
-    stream buffering re-chunks text), and tool parsing consumes tokens into tool_calls,
-    so logprob entries could not be aligned 1:1 with visible content tokens — and hidden
-    reasoning-token strings must never surface through logprobs entries. The raw
-    /v1/completions path has no semantic layer and keeps full logprobs support."""
+async def _chat_logprobs_conflict(req: ChatCompletionRequest, state: Any) -> str | None:
+    """Refuse chat logprobs when the server's parsers can hide or reclassify generated
+    tokens in a way the logprob segmenter cannot follow: entries must cover message
+    content only, never reasoning or tool-call tokens. The raw /v1/completions path has
+    no semantic layer and keeps full logprobs support."""
     if not req.logprobs:
         return None
-    if getattr(state.config, "reasoning_parser", None):
+    markers = await chat_logprob_markers(state)
+    if not chat_logprobs_supported(state, False, markers):
         return (
-            "logprobs on /v1/chat/completions are not supported when the server runs a "
-            "reasoning parser: reasoning tokens are hidden from message content, so "
-            "logprob entries cannot be aligned with it. Use /v1/completions for raw "
-            "token logprobs."
+            "logprobs on /v1/chat/completions are not supported with the "
+            f"{state.config.reasoning_parser!r} reasoning parser: reasoning tokens cannot "
+            "be told apart from message content. Use /v1/completions for raw token logprobs."
         )
-    if _should_parse_tools(req):
+    if not chat_logprobs_supported(state, _should_parse_tools(req), markers):
         return (
-            "logprobs with tool parsing on /v1/chat/completions are not supported: "
-            "tool-call tokens are consumed into tool_calls, so logprob entries cannot "
-            "be aligned with message content. Send tool_choice='none' or use "
+            "logprobs with tool parsing on /v1/chat/completions are not supported for the "
+            f"{state.config.tool_call_parser!r} tool-call parser: tool-call tokens cannot be "
+            "told apart from message content. Send tool_choice='none' or use "
             "/v1/completions for raw token logprobs."
         )
     return None
@@ -298,7 +299,7 @@ async def handle_chat_completion(
     logprobs_error = chat_logprobs_error(req)
     if logprobs_error is not None:
         return create_error_response(logprobs_error, param="top_logprobs")
-    logprobs_conflict = _chat_logprobs_conflict(req, state)
+    logprobs_conflict = await _chat_logprobs_conflict(req, state)
     if logprobs_conflict is not None:
         return create_error_response(logprobs_conflict, param="logprobs")
     if _response_format_unsupported(req.response_format):
@@ -403,9 +404,7 @@ async def handle_chat_completion(
             "index": index,
             "message": message,
             # OpenAI always sends the key, null when logprobs were not requested
-            "logprobs": (
-                {"content": [chat_content_entry(e) for e in result.logprobs]} if req.logprobs else None
-            ),
+            "logprobs": chat_logprobs_object(result.logprobs) if req.logprobs else None,
             "finish_reason": result.finish_reason,
         })
 
@@ -492,7 +491,7 @@ async def stream_chat_completion_chunks(
         elif isinstance(ev, ContentDelta):
             choice: dict[str, Any] = {"delta": {"content": ev.text}, "index": index, "finish_reason": None}
             if ev.logprobs:
-                choice["logprobs"] = {"content": [chat_content_entry(e) for e in ev.logprobs]}
+                choice["logprobs"] = chat_logprobs_object(ev.logprobs)
             yield _sse(_chat_chunk(req, uid, [choice]))
         elif isinstance(ev, ToolCallStart):
             open_tool = {
@@ -573,7 +572,10 @@ async def stream_chat_completion_chunks(
             completion_tokens = ev.completion_tokens
             cached_tokens = ev.cached_tokens
             reasoning_tokens = ev.reasoning_tokens
-            yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": index, "finish_reason": ev.finish_reason}]))
+            done: dict[str, Any] = {"delta": {}, "index": index, "finish_reason": ev.finish_reason}
+            if ev.logprobs:
+                done["logprobs"] = chat_logprobs_object(ev.logprobs)
+            yield _sse(_chat_chunk(req, uid, [done]))
             timings = ev.timings
 
     usage = _usage(
@@ -662,7 +664,6 @@ async def handle_completion(
     async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
         text = ""
         entries: list[dict] = []
-        offsets: list[int] = []
         finish_reason = "stop"
         prompt_tokens = completion_tokens = cached_tokens = 0
         async for ack in state.wait_for_ack(uid):
@@ -671,10 +672,9 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
             prompt_tokens += ack.prompt_tokens_delta
             completion_tokens += ack.completion_tokens_delta
             cached_tokens += ack.cached_tokens
-            if req.logprobs is not None and ack.logprobs is not None:
-                entries.append(ack.logprobs)
-                offsets.append(ack.text_offset)
             text += ack.incremental_output
+            if req.logprobs is not None:
+                entries.extend(untrimmed(ack.logprobs))
             if ack.finished:
                 finish_reason = getattr(ack, "finish_reason", None) or "stop"
                 break
@@ -683,7 +683,7 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
             "text": echo_texts[index // req.n] + text,
             "finish_reason": finish_reason,
             "logprobs": (
-                completions_logprobs(entries, offsets, len(text)) if req.logprobs is not None else None
+                completions_logprobs(entries, len(text)) if req.logprobs is not None else None
             ),
             "_usage": (prompt_tokens, completion_tokens, cached_tokens),
         }
@@ -737,10 +737,9 @@ async def stream_completion_chunks(
 ) -> AsyncIterator[bytes]:
     prompt_tokens = 0
     sent_len = 0
-    # Logprob entries of acks with no visible text yet (held partial UTF-8 or stop
-    # prefix), carried to the next chunk that has text.
+    # Logprob entries of acks with no visible text yet (a held partial UTF-8 character, an
+    # EOS), carried to the next chunk that has text; a reply can bring several at once.
     pending: list[dict] = []
-    pending_offsets: list[int] = []
     completion_tokens = 0
     cached_tokens = 0
     finish_reason = "stop"
@@ -756,9 +755,8 @@ async def stream_completion_chunks(
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
-        if req.logprobs is not None and ack.logprobs is not None:
-            pending.append(ack.logprobs)
-            pending_offsets.append(ack.text_offset)
+        if req.logprobs is not None:
+            pending.extend(untrimmed(ack.logprobs))
         if ack.incremental_output:
             sent_len += len(ack.incremental_output)
             yield _sse(
@@ -770,16 +768,14 @@ async def stream_completion_chunks(
                         "finish_reason": None,
                         # Mid-stream, a stamp past sent_len is held-back text still to come.
                         "logprobs": (
-                            completions_logprobs(
-                                pending, pending_offsets, sent_len if ack.finished else None
-                            )
+                            completions_logprobs(pending, sent_len if ack.finished else None)
                             if pending
                             else None
                         ),
                     }],
                 )
             )
-            pending, pending_offsets = [], []
+            pending = []
         if ack.finished:
             finish_reason = getattr(ack, "finish_reason", None) or "stop"
             break
@@ -794,7 +790,7 @@ async def stream_completion_chunks(
                 "index": index,
                 "finish_reason": finish_reason,
                 "logprobs": (
-                    completions_logprobs(pending, pending_offsets, sent_len) if pending else None
+                    completions_logprobs(pending, sent_len) if pending else None
                 ),
             }],
         )
diff --git a/python/freetoken/server/reasoning_parser.py b/python/freetoken/server/reasoning_parser.py
index 6788675ca9..c571b97a79 100644
--- a/python/freetoken/server/reasoning_parser.py
+++ b/python/freetoken/server/reasoning_parser.py
@@ -393,7 +393,7 @@ def _scan(self, text: str, *, hold_partial: bool) -> tuple[str, str]:
 
 class ThinkReasoningParser(BaseReasoningParser):
     """Generic ``<think>``/``</think>`` reasoning parser for models that wrap their
-    chain-of-thought in think tags (Qwen3/3.5, GLM-4.x, MiniMax-M2). No DSML tool
+    chain-of-thought in think tags (GLM-4.x, MiniMax-M2). No DSML tool
     marker — reasoning ends at ``</think>``.
 
     Assumes the model closes ``</think>`` before any tool call. These families do
@@ -412,6 +412,113 @@ def __init__(self, force_reasoning: bool = False, stream_reasoning: bool = True)
         )
 
 
+class Qwen3ReasoningParser(BaseReasoningParser):
+    """Qwen3 ``<think>`` parser with one positional rule for both paths.
+
+    Reasoning is open from the start when the template pre-opened it
+    (``force_reasoning``) or when the generation itself opens with ``<think>``
+    (stripped once). It ends at the FIRST ``</think>`` and is never re-entered:
+    every later ``<think>``/``</think>`` -- quoted in prose, code or tool-call
+    arguments -- is literal content. Reasoning is whitespace-trimmed at both
+    ends, and so is content that follows a real ``</think>``; streaming holds
+    whitespace-only tails so its deltas join to exactly the one-shot result.
+    The one-shot parse replays the streaming machine, so the two cannot drift.
+    """
+
+    def __init__(self, force_reasoning: bool = False, stream_reasoning: bool = True) -> None:
+        super().__init__(
+            think_start_token=THINK_START_TOKEN,
+            think_end_token=THINK_END_TOKEN,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+        )
+        # The generation start, held while it can still grow into ws* + "<think>".
+        self._head: Optional[str] = ""
+        self._trim_content = False
+        self._reasoning_started = False
+        self._content_started = False
+        self._unstreamed_reasoning = ""
+
+    def detect_and_parse(self, text: str) -> ReasoningParseResult:
+        clone = type(self)(
+            force_reasoning=self.force_reasoning, stream_reasoning=self.stream_reasoning
+        )
+        first = clone.parse_streaming_increment(text)
+        rest = clone.flush()
+        return ReasoningParseResult(
+            reasoning_text=first.reasoning_text + rest.reasoning_text,
+            normal_text=first.normal_text + rest.normal_text,
+        )
+
+    def parse_streaming_increment(self, new_text: str) -> ReasoningParseResult:
+        if self._head is not None:
+            self._head += new_text
+            head = self._head.lstrip()
+            if head.startswith(self.think_start_token):
+                self._head = None
+                self._in_reasoning = True
+                return self._consume(head[len(self.think_start_token):], final=False)
+            if not head or self.think_start_token.startswith(head):
+                return ReasoningParseResult()
+            new_text, self._head = self._head, None
+        return self._consume(new_text, final=False)
+
+    def flush(self) -> ReasoningParseResult:
+        text, self._head = self._head or "", None
+        return self._consume(text, final=True)
+
+    def _reasoning_piece(self, piece: str, closing: bool) -> str:
+        if not self._reasoning_started:
+            piece = piece.lstrip()
+            self._reasoning_started = bool(piece)
+        if self.stream_reasoning:
+            return piece
+        self._unstreamed_reasoning += piece
+        if not closing:
+            return ""
+        piece, self._unstreamed_reasoning = self._unstreamed_reasoning, ""
+        return piece
+
+    def _consume(self, text: str, final: bool) -> ReasoningParseResult:
+        buf = self._buffer + text
+        self._buffer = ""
+        reasoning = ""
+        if self._in_reasoning:
+            end = buf.find(self.think_end_token)
+            if end == -1:
+                if final:
+                    return ReasoningParseResult(
+                        reasoning_text=self._reasoning_piece(buf.rstrip(), closing=True)
+                    )
+                safe, _ = self._split_trailing_partial(buf)
+                body = safe.rstrip()
+                self._buffer = buf[len(body):]
+                return ReasoningParseResult(
+                    reasoning_text=self._reasoning_piece(body, closing=False)
+                )
+            reasoning = self._reasoning_piece(buf[:end].rstrip(), closing=True)
+            self._in_reasoning = False
+            self._trim_content = True
+            buf = buf[end + len(self.think_end_token):]
+        if not self._trim_content:
+            return ReasoningParseResult(reasoning_text=reasoning, normal_text=buf)
+        if not self._content_started:
+            buf = buf.lstrip()
+            self._content_started = bool(buf)
+        body = buf.rstrip()
+        if not final:
+            self._buffer = buf[len(body):]
+        return ReasoningParseResult(reasoning_text=reasoning, normal_text=body)
+
+    def _split_trailing_partial(self, text: str) -> tuple[str, str]:
+        # Only </think> is tracked: a later <think> is literal, never a marker.
+        token = self.think_end_token
+        for k in range(min(len(token) - 1, len(text)), 0, -1):
+            if text.endswith(token[:k]):
+                return text[:-k], text[-k:]
+        return text, ""
+
+
 class MiniMaxM3ReasoningParser(BaseReasoningParser):
     """Reasoning parser for MiniMax-M3's ``<mm:think>...</mm:think>`` protocol.
 
@@ -876,7 +983,7 @@ class ReasoningParser:
     ReasoningParserEnum: Dict[str, Type[BaseReasoningParser]] = {
         "deepseekv32": DeepSeekV32ReasoningParser,
         "gpt_oss": GptOssHarmonyReasoningParser,
-        "qwen3": ThinkReasoningParser,
+        "qwen3": Qwen3ReasoningParser,
         "glm": ThinkReasoningParser,
         "minimax": ThinkReasoningParser,
         "minimax_m3": MiniMaxM3ReasoningParser,
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index 098082edd3..c2e9bee059 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -1,6 +1,6 @@
 import json
 import weakref
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Dict, FrozenSet, List
 
 from freetoken.message import DetokenizeMsg
@@ -77,7 +77,9 @@ class DecodeStatus:
     think_end_at: int | None = None  # generated tokens up to and incl. the reasoning end tag
     in_reasoning: bool = False  # the prompt or the output opened a reasoning block
     generated: int = 0  # sampled tokens, EOS included
-    output_len: int = 0  # decoded text so far, before holdback: where the next token starts
+    text_end: int = 0  # decoded text so far, before holdback: where the next token starts
+    # (text start, entry) of logprob entries whose token text is still held for a stop string
+    held_logprobs: List[tuple] = field(default_factory=list)
 
     def reasoning_tokens(self) -> int:
         if self.think_end_at is not None:
@@ -151,13 +153,34 @@ def detokenize_with_meta(
         (0 otherwise or when the generation never was in a reasoning block); and where each
         message's token starts in the generated text. That offset ignores stop-string
         holdback, which delays when text is sent but not where it lands."""
+        texts, counts, _, offsets = self._detokenize(msgs, [None] * len(msgs))
+        return texts, counts, offsets
+
+    def detokenize_with_logprobs(
+        self, msgs: List[DetokenizeMsg], entries: List[dict | None]
+    ) -> tuple[List[str], List[int], List[List[dict]]]:
+        """``detokenize_with_meta`` plus the logprob entries each reply carries, each stamped
+        with its token's ``text_offset``. With stop strings, an entry waits for the reply that
+        releases all its token's text, so it never precedes text that may still be trimmed;
+        entries of tokens a matched stop string trimmed come on the finished reply marked
+        ``stop_trimmed``, and that of a token it cut inside with ``stop_visible``, the length
+        of the token's text left visible."""
+        texts, counts, released, _ = self._detokenize(msgs, entries)
+        return texts, counts, released
+
+    def _detokenize(
+        self, msgs: List[DetokenizeMsg], entries: List[dict | None]
+    ) -> tuple[List[str], List[int], List[List[dict]], List[int]]:
         # Each message must advance its request's decode state before the next one.
         if len({msg.uid for msg in msgs}) != len(msgs):
-            parts = [self.detokenize_with_meta([msg]) for msg in msgs]
+            parts = [
+                self._detokenize([msg], [entry]) for msg, entry in zip(msgs, entries, strict=True)
+            ]
             return (
                 [p[0][0] for p in parts],
                 [p[1][0] for p in parts],
                 [p[2][0] for p in parts],
+                [p[3][0] for p in parts],
             )
         read_ids: List[List[int]] = []
         surr_ids: List[List[int]] = []
@@ -196,12 +219,16 @@ def detokenize_with_meta(
 
         incremental_strs: List[str] = []
         reasoning_tokens: List[int] = []
+        released: List[List[dict]] = []
         text_offsets: List[int] = []
-        for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True):
+        for msg, entry, read_str, surr_str in zip(
+            msgs, entries, read_texts, surr_texts, strict=True
+        ):
             s = self.decode_map[msg.uid]
-            text_offsets.append(s.output_len)
+            token_start = s.text_end
+            text_offsets.append(token_start)
             output_str = s.update(read_str, surr_str)
-            s.output_len = max(s.output_len, len(output_str))
+            s.text_end = max(s.text_end, len(output_str))
 
             prev_sent = s.sent_offset
             if msg.finished:
@@ -226,10 +253,33 @@ def detokenize_with_meta(
             s.sent_offset = max(prev_sent, emit_end)
             incremental_strs.append(incremental_output)
             reasoning_tokens.append(s.reasoning_tokens() if msg.finished else 0)
+            if entry is not None:
+                s.held_logprobs.append((token_start, {**entry, "text_offset": token_start}))
+            out: List[dict] = []
+            # A held token's text ends where the next one's starts.
+            starts = [start for start, _ in s.held_logprobs]
+            ends = starts[1:] + [len(output_str)] if starts else []
+            if msg.finished or not msg.stop_strs:
+                trimmed = emit_end < len(output_str)
+                for (start, held), end in zip(s.held_logprobs, ends, strict=True):
+                    if trimmed and start >= emit_end:
+                        held = {**held, "stop_trimmed": True}
+                    elif trimmed and emit_end < end:
+                        held = {**held, "stop_visible": emit_end - start}
+                    out.append(held)
+                s.held_logprobs = []
+            elif s.held_logprobs:
+                # Wait for all of a token's text, so a stop match inside it can still mark it.
+                n = 0
+                while n < len(ends) and starts[n] < emit_end and ends[n] <= emit_end:
+                    n += 1
+                out = [held for _, held in s.held_logprobs[:n]]
+                s.held_logprobs = s.held_logprobs[n:]
+            released.append(out)
             if msg.finished:
                 del self.decode_map[msg.uid]
 
-        return incremental_strs, reasoning_tokens, text_offsets
+        return incremental_strs, reasoning_tokens, released, text_offsets
 
 
 def _gpt2_byte_decoder() -> dict[str, int]:
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index 3b1068514d..43c5515761 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -250,8 +250,20 @@ def tokenize_worker(
             )
             sampled_replies: List[UserReply] = []
             if len(detokenize_msg) > 0:
-                replies, reasoning_counts, text_offsets = detokenize_manager.detokenize_with_meta(
-                    detokenize_msg
+                entries = [
+                    build_logprobs_entry(
+                        detokenize_manager.tokenizer,
+                        msg.next_token,
+                        msg.chosen_logprob,
+                        msg.top_ids,
+                        msg.top_logprobs,
+                    )
+                    if msg.chosen_logprob is not None
+                    else None
+                    for msg in detokenize_msg
+                ]
+                replies, reasoning_counts, released = detokenize_manager.detokenize_with_logprobs(
+                    detokenize_msg, entries
                 )
                 sampled_replies = [
                     UserReply(
@@ -261,7 +273,6 @@ def tokenize_worker(
                         finish_reason=msg.finish_reason,
                         matched_stop=msg.matched_stop,
                         reasoning_tokens=reasoning_tokens,
-                        text_offset=text_offset,
                         completion_tokens_delta=1,
                         kv_used_pages=msg.kv_used_pages,
                         kv_total_pages=msg.kv_total_pages,
@@ -271,22 +282,10 @@ def tokenize_worker(
                         swa_total_tokens=msg.swa_total_tokens,
                         gpu_mem_bytes=msg.gpu_mem_bytes,
                         prefill_ms=msg.prefill_ms,
-                        # Stop-string trimming can hide final visible text; keep one logprob
-                        # entry per sampled token.
-                        logprobs=(
-                            build_logprobs_entry(
-                                detokenize_manager.tokenizer,
-                                msg.next_token,
-                                msg.chosen_logprob,
-                                msg.top_ids,
-                                msg.top_logprobs,
-                            )
-                            if msg.chosen_logprob is not None
-                            else None
-                        ),
+                        logprobs=reply_logprobs or None,
                     )
-                    for msg, reply, reasoning_tokens, text_offset in zip(
-                        detokenize_msg, replies, reasoning_counts, text_offsets, strict=True
+                    for msg, reply, reasoning_tokens, reply_logprobs in zip(
+                        detokenize_msg, replies, reasoning_counts, released, strict=True
                     )
                 ]
 
diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py
index 8af5bc88b6..3ba0fb0e9e 100644
--- a/tests/server/test_function_call_parser.py
+++ b/tests/server/test_function_call_parser.py
@@ -365,6 +365,35 @@ def test_dsv41_dsml_tags_lead_with_a_space():
     assert json.loads("".join(c.parameters for c in calls if c.name is None)) == {"pattern": "*.py", "path": "/src"}
 
 
+def _qwen3_coder_read(path):
+    return (
+        "<tool_call>\n<function=read>\n<parameter=filePath>\n"
+        f"{path}\n</parameter>\n</function>\n</tool_call>"
+    )
+
+
+def test_qwen3_coder_one_shot_keeps_text_between_calls():
+    text = "Hi" + _qwen3_coder_read("/a") + "\nThen more.\n" + _qwen3_coder_read("/b") + "\n"
+    result = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="qwen3_coder").parse_non_stream(text)
+    assert result.normal_text == "Hi\nThen more.\n"
+    assert [json.loads(c.parameters) for c in result.calls] == [{"filePath": "/a"}, {"filePath": "/b"}]
+
+    parser = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="qwen3_coder")
+    texts, calls = _feed(parser, text)
+    assert "".join(texts) + parser.finish_stream() == result.normal_text
+    assert [c.name for c in calls if c.name] == ["read", "read"]
+
+
+def test_qwen3_coder_unclosed_function_tag_is_not_a_call():
+    # The name stops at the next tag: never "x\n\n</tool_call" from a cut-off opener.
+    text = "<tool_call>\n<function=x\n\n</tool_call>"
+    result = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="qwen3_coder").parse_non_stream(text)
+    assert result.calls == []
+    parser = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="qwen3_coder")
+    _, calls = _feed(parser, text[:-len("</tool_call>")])
+    assert calls == [] and parser.recover_truncated_call() == []
+
+
 def test_streaming_support_flags():
     # Every registered detector is incremental-safe (buffered fallback remains as
     # the escape hatch for future formats, exercised via monkeypatch in
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 69ccf82add..53ed82f4d2 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -3,6 +3,7 @@
 import asyncio
 import json
 import os
+import random
 from types import SimpleNamespace
 
 import pytest
@@ -23,16 +24,79 @@ def run(coro):
     return asyncio.run(coro)
 
 
+THINK, END_THINK, TOOL_CALL, END_TOOL_CALL = 248068, 248069, 248058, 248059
+IM_END, ENDOFTEXT = 248046, 248044
+_MARKERS = {
+    "<think>": THINK,
+    "</think>": END_THINK,
+    "<tool_call>": TOOL_CALL,
+    "</tool_call>": END_TOOL_CALL,
+    "<|im_end|>": IM_END,
+    "<|endoftext|>": ENDOFTEXT,
+}
+
+
+class QwenTokenizer:
+    """Qwen3.8's marker ids; any other text encodes one id per character."""
+
+    all_special_ids = [ENDOFTEXT, IM_END]
+
+    def __init__(self, single_id_markers: bool = True) -> None:
+        self.single = single_id_markers
+
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        if self.single and text in _MARKERS:
+            return [_MARKERS[text]]
+        return [ord(ch) for ch in text]
+
+    def decode(self, ids: list[int]) -> str:
+        names = {v: k for k, v in _MARKERS.items()}
+        return "".join(names.get(i, chr(i) if i < 0x110000 else "") for i in ids)
+
+
+def token_replies(
+    tokens: list[str], *, finish_reason: str = "stop", texts: list[str] | None = None
+) -> list[UserReply]:
+    """One reply per sampled token, each carrying that token's entry (EOS appended for a
+    "stop" finish). ``texts`` overrides the text each reply releases, the way the
+    detokenizer holds back partial characters and stop prefixes."""
+    tokens = tokens + (["<|im_end|>"] if finish_reason == "stop" else [])
+    if texts is None:
+        texts = ["" if t in ("<|im_end|>", "<|endoftext|>") else t for t in tokens]
+    replies = [UserReply(uid=42, incremental_output="", finished=False, prompt_tokens_delta=3)]
+    for i, (token, text) in enumerate(zip(tokens, texts, strict=True)):
+        last = i == len(tokens) - 1
+        replies.append(
+            UserReply(
+                uid=42,
+                incremental_output=text,
+                finished=last,
+                finish_reason=finish_reason if last else None,
+                completion_tokens_delta=1,
+                logprobs=[entry(_MARKERS.get(token, 1000 + i), token, -0.01 * (i + 1))],
+            )
+        )
+    return replies
+
+
 class FakeState:
-    def __init__(self, replies: list[UserReply], reasoning_parser: str | None = None) -> None:
+    def __init__(
+        self,
+        replies: list[UserReply],
+        reasoning_parser: str | None = None,
+        tool_call_parser: str = "llama3",
+        tokenizer: QwenTokenizer | None = None,
+    ) -> None:
         self.config = SimpleNamespace(
             model_path="/models/unit-model",
             served_model_name="unit-model",
-            tool_call_parser="llama3",
+            tool_call_parser=tool_call_parser,
             reasoning_parser=reasoning_parser,
         )
         self.replies = replies
         self.sent: TokenizeMsg | None = None
+        if tokenizer is not None:
+            self.frontend_tokenizer = lambda: SimpleNamespace(tokenizer=tokenizer)
 
     def new_user(self) -> int:
         return 42
@@ -46,9 +110,7 @@ async def wait_for_ack(self, uid: int):
             yield reply
 
 
-def reply(
-    text: str, *, finished: bool = False, logprobs: dict | None = None, text_offset: int = 0
-) -> UserReply:
+def reply(text: str, *, finished: bool = False, logprobs: list[dict] | None = None) -> UserReply:
     return UserReply(
         uid=42,
         incremental_output=text,
@@ -57,14 +119,14 @@ def reply(
         completion_tokens_delta=1 if text else 0,
         cached_tokens=0,
         logprobs=logprobs,
-        text_offset=text_offset,
     )
 
 
-def entry(token_id: int, token: str, logprob: float) -> dict:
+def entry(token_id: int, token: str, logprob: float, *, text_offset: int = 0) -> dict:
     return {
         "token_id": token_id,
         "token": token,
+        "text_offset": text_offset,
         "bytes": list(token.encode("utf-8")),
         "logprob": logprob,
         "top": [
@@ -111,7 +173,7 @@ def test_chat_non_stream_logprobs() -> None:
         handle_chat_completion(
             chat_request(logprobs=True, top_logprobs=2),
             None,
-            FakeState([reply("Hello", logprobs=first), reply("!", finished=True, logprobs=second)]),
+            FakeState([reply("Hello", logprobs=[first]), reply("!", finished=True, logprobs=[second])]),
             {},
         )
     )
@@ -142,7 +204,7 @@ def test_chat_non_stream_logprobs() -> None:
         handle_chat_completion(
             chat_request(),
             None,
-            FakeState([reply("Hello", finished=True, logprobs=first)]),
+            FakeState([reply("Hello", finished=True, logprobs=[first])]),
             {},
         )
     )
@@ -155,7 +217,7 @@ def test_chat_stream_logprobs_follow_content_deltas() -> None:
     first = entry(1, "Hello", -0.1)
     state = FakeState(
         [
-            reply("Hello", logprobs=first),
+            reply("Hello", logprobs=[first]),
             reply(" world", finished=True),
         ]
     )
@@ -176,16 +238,14 @@ def test_chat_stream_logprobs_follow_content_deltas() -> None:
 
 def test_completion_logprobs_non_stream_and_stream() -> None:
     first = entry(1, "Hi", -0.1)
-    second = entry(2, "!", -0.2)
+    second = entry(2, "!", -0.2, text_offset=2)
     req = CompletionRequest(model="client-model", prompt="hello", logprobs=2, max_tokens=8)
 
     result = run(
         handle_completion(
             req,
             None,
-            FakeState(
-                [reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second, text_offset=2)]
-            ),
+            FakeState([reply("Hi", logprobs=[first]), reply("!", finished=True, logprobs=[second])]),
             {},
         )
     )
@@ -211,7 +271,7 @@ def test_completion_logprobs_non_stream_and_stream() -> None:
                         max_tokens=8,
                         stream=True,
                     ),
-                    FakeState([reply("Hi", finished=True, logprobs=first)]),
+                    FakeState([reply("Hi", finished=True, logprobs=[first])]),
                 )
             )
         )
@@ -228,11 +288,11 @@ def test_completion_logprobs_non_stream_and_stream() -> None:
 def _crab_replies() -> list[UserReply]:
     # 🦀 is three byte-level tokens: the first two acks carry no visible text yet.
     pieces = [
-        reply("", logprobs=entry(9008, "\ufffd", -0.1)),
-        reply("", logprobs=entry(99, "\ufffd", -0.2)),
-        reply("🦀", logprobs=entry(222, "\ufffd", -0.3)),
-        reply(" x", logprobs=entry(1, " x", -0.4), text_offset=1),
-        reply("", finished=True, logprobs=entry(2, "<|im_end|>", -0.5), text_offset=3),
+        reply("", logprobs=[entry(9008, "�", -0.1)]),
+        reply("", logprobs=[entry(99, "�", -0.2)]),
+        reply("🦀", logprobs=[entry(222, "�", -0.3)]),
+        reply(" x", logprobs=[entry(1, " x", -0.4, text_offset=1)]),
+        reply("", finished=True, logprobs=[entry(2, "<|im_end|>", -0.5, text_offset=3)]),
     ]
     return pieces
 
@@ -243,19 +303,34 @@ def test_completion_text_offset_counts_visible_text() -> None:
     logprobs = result["choices"][0]["logprobs"]
     assert result["choices"][0]["text"] == "🦀 x"
     # The split character's pieces all point at where it starts, not at
-    # len("\ufffd") steps that drift past the real text.
+    # len("�") steps that drift past the real text.
     assert logprobs["text_offset"] == [0, 0, 0, 1, 3]
 
 
-def _holdback_replies() -> list[UserReply]:
-    # "Answer: 4\nQuestion 2" under stop=["\n\n"]: the "\n" ack sends nothing (it could
-    # start the stop), and "Question" releases it along with its own text.
-    acks = [("Answer", "Answer", 0), (":", ":", 6), (" ", " ", 7), ("4", "4", 8),
-            ("\n", "", 9), ("Question", "\nQuestion", 10), (" ", " ", 18), ("2", "2", 19)]
-    return [
-        reply(text, finished=i == len(acks) - 1, logprobs=entry(i, token, -0.1), text_offset=offset)
-        for i, (token, text, offset) in enumerate(acks)
-    ]
+def _offset_replies(acks: list[tuple[str, list[tuple[str, int]]]]) -> list[UserReply]:
+    """One reply per (text, [(token, text_offset), ...]) ack, entries in generation order."""
+    replies, n = [], 0
+    for i, (text, tokens) in enumerate(acks):
+        entries = []
+        for token, offset in tokens:
+            entries.append(entry(n, token, -0.1, text_offset=offset))
+            n += 1
+        replies.append(reply(text, finished=i == len(acks) - 1, logprobs=entries or None))
+    return replies
+
+
+# "Answer: 4\nQuestion 2" under stop=["\n\n"]: the "\n" ack sends nothing (it could start
+# the stop), and "Question" releases it along with its own text. "own" gives every ack
+# its own entry; "held" is what the detokenizer sends, holding the "\n" entry until the
+# reply that releases its text, which then carries two entries.
+_HOLDBACK_ACKS = {
+    "own": [("Answer", [("Answer", 0)]), (":", [(":", 6)]), (" ", [(" ", 7)]), ("4", [("4", 8)]),
+            ("", [("\n", 9)]), ("\nQuestion", [("Question", 10)]), (" ", [(" ", 18)]),
+            ("2", [("2", 19)])],
+    "held": [("Answer", [("Answer", 0)]), (":", [(":", 6)]), (" ", [(" ", 7)]), ("4", [("4", 8)]),
+             ("", []), ("\nQuestion", [("\n", 9), ("Question", 10)]), (" ", [(" ", 18)]),
+             ("2", [("2", 19)])],
+}
 
 
 def _joined_stream_logprobs(replies: list[UserReply], **kwargs) -> dict[str, list]:
@@ -269,24 +344,27 @@ def _joined_stream_logprobs(replies: list[UserReply], **kwargs) -> dict[str, lis
     return joined
 
 
-def test_completion_text_offset_survives_stop_prefix_holdback() -> None:
+@pytest.mark.parametrize("shape", ["own", "held"])
+def test_completion_text_offset_survives_stop_prefix_holdback(shape: str) -> None:
     req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["\n\n"])
-    choice = run(handle_completion(req, None, FakeState(_holdback_replies()), {}))["choices"][0]
+    replies = _offset_replies(_HOLDBACK_ACKS[shape])
+    choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
     logprobs = choice["logprobs"]
     assert choice["text"] == "Answer: 4\nQuestion 2"
+    assert logprobs["tokens"] == ["Answer", ":", " ", "4", "\n", "Question", " ", "2"]
     assert logprobs["text_offset"] == [0, 6, 7, 8, 9, 10, 18, 19]
     for token, offset in zip(logprobs["tokens"], logprobs["text_offset"], strict=True):
         assert choice["text"][offset:].startswith(token)
-    assert _joined_stream_logprobs(_holdback_replies(), logprobs=0, stop=["\n\n"]) == logprobs
+    assert _joined_stream_logprobs(replies, logprobs=0, stop=["\n\n"]) == logprobs
 
 
 def test_completion_text_offset_clipped_to_stop_trimmed_text() -> None:
-    # "a\n\n" under stop=["\n\n"]: the second "\n" completes the stop and both are trimmed,
-    # so their offsets cannot point past the returned text.
+    # "a\n\n" under stop=["\n\n"]: the second "\n" completes the stop and both are trimmed.
+    # Unmarked, their offsets are clipped so they cannot point past the returned text.
     replies = [
-        reply("a", logprobs=entry(1, "a", -0.1)),
-        reply("", logprobs=entry(2, "\n", -0.2), text_offset=1),
-        reply("", finished=True, logprobs=entry(2, "\n", -0.3), text_offset=2),
+        reply("a", logprobs=[entry(1, "a", -0.1)]),
+        reply("", logprobs=[entry(2, "\n", -0.2, text_offset=1)]),
+        reply("", finished=True, logprobs=[entry(2, "\n", -0.3, text_offset=2)]),
     ]
     req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["\n\n"])
     choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
@@ -294,27 +372,43 @@ def test_completion_text_offset_clipped_to_stop_trimmed_text() -> None:
     assert choice["logprobs"]["text_offset"] == [0, 1, 1]
     assert _joined_stream_logprobs(replies, logprobs=0, stop=["\n\n"]) == choice["logprobs"]
 
+    # As the detokenizer sends them: held with their text, then marked and dropped.
+    trimmed = [
+        {**entry(2, "\n", -0.2, text_offset=1), "stop_trimmed": True},
+        {**entry(2, "\n", -0.3, text_offset=2), "stop_trimmed": True},
+    ]
+    marked = [replies[0], reply("", logprobs=None), reply("", finished=True, logprobs=trimmed)]
+    choice = run(handle_completion(req, None, FakeState(marked), {}))["choices"][0]
+    assert choice["logprobs"]["tokens"] == ["a"]
+    assert choice["logprobs"]["text_offset"] == [0]
+    assert _joined_stream_logprobs(marked, logprobs=0, stop=["\n\n"]) == choice["logprobs"]
+
+
+# stop=["aab"] over "aaax": two "a"s stay held while the first is sent. "own": the held
+# tokens' entries ride their own acks, so their offsets run past the text sent so far;
+# that is where they land, not an overrun. "held": the detokenizer's shape, each entry
+# waiting for its text, so the finished reply carries three.
+_AHEAD_ACKS = {
+    "own": [("", [("a", 0)]), ("", [("a", 1)]), ("a", [("a", 2)]), ("aax", [("x", 3)])],
+    "held": [("", []), ("", []), ("a", [("a", 0)]), ("aax", [("a", 1), ("a", 2), ("x", 3)])],
+}
 
-def test_completion_stream_text_offset_ahead_of_sent_text() -> None:
-    # stop=["aab"] over "aaax": two "a"s stay held while the first is sent, so the held
-    # tokens' offsets run past the text sent so far; that is where they land, not an overrun.
-    def replies() -> list[UserReply]:
-        acks = [("a", "", 0), ("a", "", 1), ("a", "a", 2), ("x", "aax", 3)]
-        return [
-            reply(text, finished=i == 3, logprobs=entry(i, token, -0.1), text_offset=offset)
-            for i, (token, text, offset) in enumerate(acks)
-        ]
 
+@pytest.mark.parametrize("shape", ["own", "held"])
+def test_completion_stream_text_offset_ahead_of_sent_text(shape: str) -> None:
     req = CompletionRequest(model="client-model", prompt="hello", logprobs=0, max_tokens=8, stop=["aab"])
-    choice = run(handle_completion(req, None, FakeState(replies()), {}))["choices"][0]
+    replies = _offset_replies(_AHEAD_ACKS[shape])
+    choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
     assert choice["text"] == "aaax"
+    assert choice["logprobs"]["tokens"] == ["a", "a", "a", "x"]
     assert choice["logprobs"]["text_offset"] == [0, 1, 2, 3]
-    assert _joined_stream_logprobs(replies(), logprobs=0, stop=["aab"]) == choice["logprobs"]
+    assert _joined_stream_logprobs(replies, logprobs=0, stop=["aab"]) == choice["logprobs"]
 
 
 def test_completion_text_offset_real_tokenizer_with_stops() -> None:
     # The real Qwen3.8 tokenizer through the real detokenizer: every token's offset must
-    # point at its own text, stream and non-stream alike.
+    # point at its own text, stream and non-stream alike, with entries held for their text
+    # and those of stop-trimmed tokens dropped.
     path = os.path.expanduser("~/models/Qwen3.8-Flash-Next-NVFP4-FTW-fp8")
     if not os.path.exists(os.path.join(path, "tokenizer.json")):
         pytest.skip("local Qwen3.8 tokenizer not present")
@@ -324,41 +418,58 @@ def test_completion_text_offset_real_tokenizer_with_stops() -> None:
     from freetoken.tokenizer.detokenize import DetokenizeManager, build_logprobs_entry
 
     tokenizer = AutoTokenizer.from_pretrained(path)
-    for text, stop in [("Answer: 4\nQuestion 2", "\n\n"), ("Q: hi\nA: Question time 🦀 ok", "Question:")]:
+    cases = [
+        ("Answer: 4\nQuestion 2", "\n\n", None),
+        ("Q: hi\nA: Question time 🦀 ok", "Question:", None),
+        # the stop matches on the last token: the "\n\n" token it trimmed loses its entry
+        ("Answer: 4\n\n", "\n\n", "\n\n"),
+        # the stop cuts inside " hello": that entry stays, " world" (trimmed) does not
+        ("Say hello world", "lo w", "lo w"),
+    ]
+    held = False
+    for text, stop, matched in cases:
         ids = tokenizer.encode(text, add_special_tokens=False)
         manager = DetokenizeManager(tokenizer, frozenset({248046}))
         replies = []
         for i, tid in enumerate(ids):
             last = i == len(ids) - 1
-            msg = DetokenizeMsg(uid=42, next_token=tid, finished=last, stop_strs=[stop])
-            (out,), _, (offset,) = manager.detokenize_with_meta([msg])
-            replies.append(
-                reply(out, finished=last, text_offset=offset,
-                      logprobs=build_logprobs_entry(tokenizer, tid, -0.1, [], []))
+            msg = DetokenizeMsg(
+                uid=42, next_token=tid, finished=last, stop_strs=[stop],
+                matched_stop=matched if last else None,
+            )
+            (out,), _, (released,) = manager.detokenize_with_logprobs(
+                [msg], [build_logprobs_entry(tokenizer, tid, -0.1, [], [])]
             )
+            held |= len(released) > 1
+            replies.append(reply(out, finished=last, logprobs=released or None))
+        visible = text[: text.find(matched)] if matched else text
         req = CompletionRequest(model="client-model", prompt="p", logprobs=0, max_tokens=64, stop=[stop])
         choice = run(handle_completion(req, None, FakeState(replies), {}))["choices"][0]
-        assert choice["text"] == text
-        offsets = choice["logprobs"]["text_offset"]
-        for tid, offset in zip(ids, offsets, strict=True):
-            token = tokenizer.decode([tid])
+        assert choice["text"] == visible
+        tokens, offsets = choice["logprobs"]["tokens"], choice["logprobs"]["text_offset"]
+        assert len(tokens) == len(ids) - (1 if matched else 0)
+        for token, offset in zip(tokens, offsets, strict=True):
             if "\ufffd" not in token:
-                assert text[offset:].startswith(token), (text, token, offset)
+                # a token's text sits at its offset; only one the stop cut ends the text early
+                seg = visible[offset : offset + len(token)]
+                assert token.startswith(seg), (text, token, offset)
+                assert seg == token or offset + len(seg) == len(visible), (text, token, offset)
         assert offsets == sorted(offsets)
         assert _joined_stream_logprobs(replies, logprobs=0, stop=[stop]) == choice["logprobs"]
+    assert held, "no reply carried more than one entry"
 
 
 def test_completion_top_logprobs_key_collision_keeps_most_likely() -> None:
     # Partial-UTF-8 candidates all decode to U+FFFD and collide on the dict key.
-    colliding = entry(9008, "\ufffd", -1.0)
+    colliding = entry(9008, "�", -1.0)
     colliding["top"] = [
-        {"token_id": 9008, "token": "\ufffd", "bytes": [240, 159], "logprob": -1.0},
-        {"token_id": 9009, "token": "\ufffd", "bytes": [240, 160], "logprob": -2.0},
+        {"token_id": 9008, "token": "�", "bytes": [240, 159], "logprob": -1.0},
+        {"token_id": 9009, "token": "�", "bytes": [240, 160], "logprob": -2.0},
         {"token_id": 5, "token": "a", "bytes": [97], "logprob": -3.0},
     ]
     req = CompletionRequest(model="client-model", prompt="hello", logprobs=3, max_tokens=8)
-    result = run(handle_completion(req, None, FakeState([reply("x", finished=True, logprobs=colliding)]), {}))
-    assert result["choices"][0]["logprobs"]["top_logprobs"] == [{"\ufffd": -1.0, "a": -3.0}]
+    result = run(handle_completion(req, None, FakeState([reply("x", finished=True, logprobs=[colliding])]), {}))
+    assert result["choices"][0]["logprobs"]["top_logprobs"] == [{"�": -1.0, "a": -3.0}]
 
 
 def test_completion_stream_logprobs_match_non_stream() -> None:
@@ -377,6 +488,52 @@ def test_completion_stream_logprobs_match_non_stream() -> None:
     assert len(joined["tokens"]) == 5
 
 
+def _stop_trimmed_replies() -> list[UserReply]:
+    # "Hi", then a stop string "##" sampled as "#" "#": the detokenizer withheld both
+    # entries with the text and releases them, marked, on the finished reply.
+    trimmed = [{**entry(3, "#", -0.3), "stop_trimmed": True}, {**entry(3, "#", -0.4), "stop_trimmed": True}]
+    return [
+        reply("Hi", logprobs=[entry(1, "Hi", -0.1)]),
+        reply("", logprobs=None),
+        reply("", finished=True, logprobs=trimmed),
+    ]
+
+
+def test_completion_logprobs_drop_stop_trimmed_entries() -> None:
+    req = CompletionRequest(model="client-model", prompt="hello", logprobs=1, max_tokens=8, stop=["##"])
+    result = run(handle_completion(req, None, FakeState(_stop_trimmed_replies()), {}))
+    assert result["choices"][0]["logprobs"]["tokens"] == ["Hi"]
+
+    stream_req = req.model_copy(update={"stream": True})
+    events = parse_sse(run(_collect(stream_completion_chunks(42, stream_req, FakeState(_stop_trimmed_replies())))))
+    tokens = [
+        token
+        for event in events
+        if isinstance(event, dict) and event["choices"] and event["choices"][0]["logprobs"]
+        for token in event["choices"][0]["logprobs"]["tokens"]
+    ]
+    assert tokens == ["Hi"]
+
+
+def test_chat_logprobs_drop_stop_trimmed_entries() -> None:
+    req = chat_request(logprobs=True, stop=["##"])
+    result = run(handle_chat_completion(req, None, FakeState(_stop_trimmed_replies()), {}))
+    assert [e["token"] for e in result["choices"][0]["logprobs"]["content"]] == ["Hi"]
+
+    stream_req = chat_request(logprobs=True, stop=["##"], stream=True)
+    spec = chat_request_to_genspec(stream_req, {})
+    events = parse_sse(
+        run(_collect(stream_chat_completion_chunks(42, stream_req, FakeState(_stop_trimmed_replies()), spec)))
+    )
+    tokens = [
+        e["token"]
+        for event in events
+        if isinstance(event, dict) and event["choices"] and event["choices"][0].get("logprobs")
+        for e in event["choices"][0]["logprobs"]["content"]
+    ]
+    assert tokens == ["Hi"]
+
+
 def test_logprobs_validation_errors() -> None:
     chat_top_out_of_range = run(
         handle_chat_completion(chat_request(logprobs=True, top_logprobs=25), None, FakeState([]), {})
@@ -409,61 +566,280 @@ def test_logprobs_validation_errors() -> None:
     assert completion_echo.status_code == 400
 
 
-def test_chat_logprobs_fail_closed_under_semantic_parsing() -> None:
-    # A server-side reasoning parser hides <think> tokens from message content, so
-    # logprob entries cannot be aligned with it: reject up front, stream and
-    # non-stream alike, before any engine work is submitted.
-    with_parser = FakeState([], reasoning_parser="qwen3")
-    for stream in (False, True):
-        resp = run(
-            handle_chat_completion(chat_request(logprobs=True, stream=stream), None, with_parser, {})
-        )
-        assert resp.status_code == 400
-        assert json.loads(resp.body)["error"]["param"] == "logprobs"
-
-    # Tool parsing consumes tokens into tool_calls -- same conflict.
-    tool = {"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}
-    tools_resp = run(
-        handle_chat_completion(chat_request(logprobs=True, tools=[tool]), None, FakeState([]), {})
+TOOL = {
+    "type": "function",
+    "function": {
+        "name": "read",
+        "parameters": {"type": "object", "properties": {"path": {"type": "string"}}},
+    },
+}
+CALL = [
+    "<tool_call>", "\n", "<", "function", "=read", ">", "\n", "<", "parameter", "=path", ">",
+    "\n", "a", ".py", "\n", "</", "parameter", ">", "\n", "</", "function", ">", "\n",
+    "</tool_call>",
+]
+
+
+def qwen_state(replies: list[UserReply], **kwargs) -> FakeState:
+    """The production stack: qwen3 reasoning parser and qwen3_coder tool parser."""
+    return FakeState(
+        replies,
+        reasoning_parser="qwen3",
+        tool_call_parser="qwen3_coder",
+        tokenizer=kwargs.pop("tokenizer", QwenTokenizer()),
+        **kwargs,
     )
-    assert tools_resp.status_code == 400
-    assert json.loads(tools_resp.body)["error"]["param"] == "logprobs"
 
-    # tool_choice="none" disables parsing, so logprobs stay available.
+
+def _chat_both(replies: list[UserReply], **req_kwargs) -> tuple[dict, list[dict]]:
+    """The non-streaming response and the streamed chunk choices for the same replies."""
+    req = chat_request(logprobs=True, **req_kwargs)
+    result = run(handle_chat_completion(req, None, qwen_state(replies), {}))
+    req = chat_request(logprobs=True, stream=True, **req_kwargs)
+    spec = chat_request_to_genspec(req, {})
+    chunks = stream_chat_completion_chunks(42, req, qwen_state(replies), spec)
+    events = parse_sse(run(_collect(chunks)))
+    choices = [e["choices"][0] for e in events if isinstance(e, dict) and e.get("choices")]
+    return result, choices
+
+
+def _streamed_tokens(choices: list[dict]) -> list[str]:
+    return [e["token"] for c in choices if c.get("logprobs") for e in c["logprobs"]["content"]]
+
+
+def _content_tokens(result: dict) -> list[str]:
+    return [e["token"] for e in result["choices"][0]["logprobs"]["content"]]
+
+
+def test_chat_logprobs_exclude_reasoning() -> None:
+    replies = token_replies(["thought", " more", "</think>", "\n\n", "Hello", " world"])
+    result, choices = _chat_both(replies)
+
+    assert result["choices"][0]["message"]["reasoning_content"] == "thought more"
+    assert result["choices"][0]["logprobs"] == {
+        "content": result["choices"][0]["logprobs"]["content"],
+        "refusal": None,
+    }
+    assert _content_tokens(result) == ["Hello", " world"]
+    assert _streamed_tokens(choices) == ["Hello", " world"]
+    for choice in choices:
+        if choice["delta"].get("reasoning_content"):
+            assert "logprobs" not in choice
+        if choice.get("logprobs"):
+            assert choice["logprobs"]["refusal"] is None
+
+
+def test_chat_logprobs_exclude_tool_call_markup() -> None:
+    replies = token_replies(["plan", "</think>", "\n\n", "Sure", "."] + CALL, finish_reason="stop")
+    result, choices = _chat_both(replies, tools=[TOOL])
+
+    assert result["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "read"
+    assert _content_tokens(result) == ["Sure", "."]
+    assert _streamed_tokens(choices) == ["Sure", "."]
+    for choice in choices:
+        if choice["delta"].get("tool_calls"):
+            assert "logprobs" not in choice
+
+
+def test_chat_logprobs_bare_function_call_is_not_content() -> None:
+    # The qwen3_coder parser takes a declared <function=...> without its <tool_call>
+    # wrapper as a call; its markup tokens are not content either.
+    call = CALL[2:-2]
+    replies = token_replies(["plan", "</think>", "Doing", " it", ".\n"] + call + ["\n", "ok"])
+    result, choices = _chat_both(replies, tools=[TOOL])
+
+    assert result["choices"][0]["message"]["tool_calls"]
+    assert result["choices"][0]["message"]["content"] == "Doing it.\n\nok"
+    assert _content_tokens(result) == ["Doing", " it", ".\n", "\n", "ok"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_text_merged_into_a_bare_call_closer_is_content() -> None:
+    # Qwen3.8 merges the closing ">" with the next character (">n"): the emitted "n"
+    # sits at the end of "</function>n", not on the "n" inside "function".
+    call = CALL[2:-3] + [">n"]
+    replies = token_replies(["plan", "</think>"] + call)
+    result, choices = _chat_both(replies, tools=[TOOL])
+
+    assert result["choices"][0]["message"]["tool_calls"]
+    assert result["choices"][0]["message"]["content"] == "n"
+    assert _content_tokens(result) == [">n"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_leading_space_merged_into_a_bare_call_opener_is_not_content() -> None:
+    # The reasoning parser strips content's leading whitespace before the tool parser
+    # sees it, so the space of " <" is never message text.
+    call = [" <"] + CALL[3:-2]
+    replies = token_replies(["x", "</think>"] + call)
+    result, choices = _chat_both(replies, tools=[TOOL])
+
+    assert result["choices"][0]["message"]["tool_calls"]
+    assert not result["choices"][0]["message"]["content"]
+    assert _content_tokens(result) == []
+    assert _streamed_tokens(choices) == []
+
+
+def _stop_cut_replies(tokens: list[str], visible: int) -> list[UserReply]:
+    """A stop string matched inside the last token: only its first ``visible`` chars
+    were released, and the detokenizer says so on its entry."""
+    texts = tokens[:-1] + [tokens[-1][:visible]]
+    replies = token_replies(tokens, finish_reason="length", texts=texts)
+    replies[-1].finish_reason = "stop"
+    replies[-1].logprobs = [{**replies[-1].logprobs[0], "stop_visible": visible}]
+    return replies
+
+
+def test_chat_logprobs_follow_the_visible_part_of_a_token_a_stop_string_cut() -> None:
+    # Only the trailing space of " !" survives the stop "!", and content trimming drops it.
+    result, choices = _chat_both(_stop_cut_replies(["x", "</think>", "Hi", " !"], 1), stop=["!"])
+    assert result["choices"][0]["message"]["content"] == "Hi"
+    assert _content_tokens(result) == ["Hi"]
+    assert _streamed_tokens(choices) == ["Hi"]
+
+    # A stop "_c" leaves "<tool" of <tool_call> as text, not a call opener.
+    replies = _stop_cut_replies(["x", "</think>", "Use", " ", "<tool_call>"], 5)
+    result, choices = _chat_both(replies, stop=["_c"], tools=[TOOL])
+    assert result["choices"][0]["message"]["content"] == "Use <tool"
+    assert _content_tokens(result) == ["Use", " ", "<tool_call>"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_literal_markers_in_content_stay_content() -> None:
+    # Only a marker that switches state is markup: a later </think> or a <tool_call>
+    # named in prose is message text, and so is its entry.
+    replies = token_replies(
+        ["plan", "</think>", "Use", " </think>", " and", " <tool_call>", " tags", "."]
+    )
+    result, choices = _chat_both(replies, tools=[TOOL])
+
+    content = result["choices"][0]["message"]["content"]
+    assert content == "Use </think> and <tool_call> tags."
+    assert "".join(_content_tokens(result)) == content
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_ride_the_delta_that_releases_a_held_marker_prefix() -> None:
+    # The tool parser holds the "<" of " <" as a possible <tool_call>: the token's
+    # entry waits until all its text is out, then rides that delta, in order.
+    replies = token_replies(["</think>", "a", " <", " b"])
+    _, choices = _chat_both(replies, tools=[TOOL])
+
+    content_chunks = [c for c in choices if c["delta"].get("content")]
+    assert [c["delta"]["content"] for c in content_chunks] == ["a", " ", "< b"]
+    assert [
+        [e["token"] for e in c["logprobs"]["content"]] if c.get("logprobs") else None
+        for c in content_chunks
+    ] == [["a"], None, [" <", " b"]]
+
+
+def test_chat_logprobs_trailing_whitespace_is_dropped() -> None:
+    # Content after </think> is trimmed at both ends, and so are its whitespace entries;
+    # whitespace between words is kept.
+    replies = token_replies(["</think>", "\n\n", "Hi", "\n", "\n", "there", "\n\n"])
+    result, choices = _chat_both(replies)
+
+    assert result["choices"][0]["message"]["content"] == "Hi\n\nthere"
+    assert _content_tokens(result) == ["Hi", "\n", "\n", "there"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+@pytest.mark.parametrize("seed", range(12))
+def test_chat_logprobs_stream_matches_non_stream_for_any_chunking(seed: int) -> None:
+    # The detokenizer can release a token's text on a later reply (a partial character,
+    # a stop prefix), leaving replies with no text at all: the streamed entries, joined
+    # across content chunks and the finish chunk, must still equal the one-shot list.
+    tokens = (
+        ["  ", "<think>", "plan", " it", "</think>", "\n\n", "Hi", " ", "<", "b", ">", "\n\n"]
+        + CALL
+        + ["\n", "Then", " more", "\n"]
+    )
+    full = "".join(tokens)
+    rng = random.Random(seed)
+    cuts, end, prev = [], 0, 0
+    for token in tokens:
+        end += len(token)
+        prev = rng.randint(prev, end)
+        cuts.append(prev)
+    cuts[-1] = len(full)
+    texts = [full[a:b] for a, b in zip([0] + cuts[:-1], cuts, strict=True)] + [""]
+    replies = token_replies(tokens, texts=texts)
+
+    result, choices = _chat_both(replies, tools=[TOOL])
+    assert result["choices"][0]["message"]["content"] == "Hi <b>\n\n\nThen more"
+    assert _content_tokens(result) == ["Hi", " ", "<", "b", ">", "\n\n", "\n", "Then", " more"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_start_in_content_when_thinking_is_off() -> None:
+    replies = token_replies(["  ", "Hi", " </think>", " there"])
+    result, choices = _chat_both(replies, chat_template_kwargs={"enable_thinking": False})
+
+    assert result["choices"][0]["message"]["content"] == "  Hi </think> there"
+    assert _content_tokens(result) == ["  ", "Hi", " </think>", " there"]
+    assert _streamed_tokens(choices) == _content_tokens(result)
+
+
+def test_chat_logprobs_truncated_reasoning_has_no_content_entries() -> None:
+    replies = token_replies(["still", " thinking"], finish_reason="length")
+    result, choices = _chat_both(replies)
+
+    assert result["choices"][0]["finish_reason"] == "length"
+    assert result["choices"][0]["logprobs"] == {"content": [], "refusal": None}
+    assert _streamed_tokens(choices) == []
+
+
+def test_chat_logprobs_drop_the_eos_entry() -> None:
+    # The end-of-turn token is not message text; non-streaming used to list it.
+    replies = token_replies(["Hi", "!"])
+    state = FakeState(replies, tokenizer=QwenTokenizer())
+    result = run(handle_chat_completion(chat_request(logprobs=True), None, state, {}))
+    assert _content_tokens(result) == ["Hi", "!"]
+
+
+def test_chat_logprobs_unsupported_parser_families_still_400() -> None:
+    def status(state: FakeState, **kwargs) -> int | None:
+        resp = run(handle_chat_completion(chat_request(logprobs=True, **kwargs), None, state, {}))
+        if isinstance(resp, dict):
+            return None
+        assert json.loads(resp.body)["error"]["param"] == "logprobs"
+        return resp.status_code
+
+    # A reasoning parser whose split the segmenter does not mirror.
+    assert status(FakeState([], reasoning_parser="glm", tokenizer=QwenTokenizer())) == 400
+    # qwen3 whose think markers are not single tokens, or with no tokenizer to tell.
+    multi = QwenTokenizer(single_id_markers=False)
+    assert status(FakeState([], reasoning_parser="qwen3", tokenizer=multi)) == 400
+    assert status(FakeState([], reasoning_parser="qwen3")) == 400
+    for stream in (False, True):
+        assert status(FakeState([], reasoning_parser="qwen3"), stream=stream) == 400
+    # A tool-call parser the segmenter cannot shadow...
+    assert status(FakeState([], tokenizer=QwenTokenizer()), tools=[TOOL]) == 400
+    # ...unless tool parsing is off for the request.
     ok = run(
         handle_chat_completion(
-            chat_request(logprobs=True, tools=[tool], tool_choice="none"),
+            chat_request(logprobs=True, tools=[TOOL], tool_choice="none"),
             None,
-            FakeState([reply("Hi", finished=True, logprobs=entry(1, "Hi", -0.1))]),
+            FakeState(token_replies(["Hi"]), tokenizer=QwenTokenizer()),
             {},
         )
     )
-    assert ok["choices"][0]["logprobs"]["content"][0]["token"] == "Hi"
+    assert _content_tokens(ok) == ["Hi"]
 
 
-def test_reasoning_logprob_is_not_carried_to_content_delta() -> None:
-    # The generation-layer guard behind the API fail-close: even when a caller
-    # bypasses request validation, a hidden reasoning token's entry is dropped,
-    # never attached to a later visible delta (where its token string would leak).
+def test_chat_logprobs_are_dropped_when_the_segmenter_cannot_build() -> None:
+    # The generation-layer guard behind the API check: a caller that bypasses it still
+    # never gets a reasoning token's entry on a visible delta.
     state = FakeState(
-        [
-            reply("<think>thought</think>", logprobs=entry(1, "thought", -0.1)),
-            reply("answer", finished=True),
-        ],
-        reasoning_parser="qwen3",
+        token_replies(["thought", "</think>", "answer"]), reasoning_parser="qwen3"
     )
     req = chat_request(stream=True, logprobs=True, top_logprobs=2)
     spec = chat_request_to_genspec(req, {})
     events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec))))
 
-    content_choice = next(
-        event["choices"][0]
-        for event in events
-        if isinstance(event, dict)
-        and event["choices"]
-        and event["choices"][0]["delta"].get("content") == "answer"
-    )
-    assert "logprobs" not in content_choice
+    choices = [e["choices"][0] for e in events if isinstance(e, dict) and e.get("choices")]
+    assert any(c["delta"].get("content") == "answer" for c in choices)
+    assert _streamed_tokens(choices) == []
 
 
 async def _collect(iterator):
diff --git a/tests/server/test_qwen_semantic_special_tokens.py b/tests/server/test_qwen_semantic_special_tokens.py
index 11e35ccb3e..da20fdde7a 100644
--- a/tests/server/test_qwen_semantic_special_tokens.py
+++ b/tests/server/test_qwen_semantic_special_tokens.py
@@ -118,8 +118,10 @@ def test_stream_and_nonstream_match_while_preserving_semantic_markers_and_litera
         "<|box_start|>(1,2),(3,4)<|box_end|>"
         "<|quad_start|>(1,2),(3,4),(5,6),(7,8)<|quad_end|>"
     )
+    # The leaked marker sits after the think block: qwen3 opens on <think> only as the
+    # first non-whitespace text, so "<|im_start|><think>" keeps "<think>" as reasoning.
     text = (
-        "<|im_start|><think>plan</think>answer"
+        "<think>plan</think><|im_start|>answer"
         "<|vision_start|><|image_pad|><|vision_end|>"
         f"{grounding} literal \"<|im_start|>\" and `<|image_pad|>`<|im_end|>"
     )
diff --git a/tests/server/test_reasoning_parser.py b/tests/server/test_reasoning_parser.py
index 05c207d0e1..2580af83b9 100644
--- a/tests/server/test_reasoning_parser.py
+++ b/tests/server/test_reasoning_parser.py
@@ -44,10 +44,25 @@ class TestDeepSeekV32ReasoningParser(_ReasoningContract):
 
 
 class TestThinkReasoningParser(_ReasoningContract):
+    name = "glm"
+    opener = "<think>"
+    closer = "</think>"
+
+
+class TestQwen3ReasoningParser(_ReasoningContract):
     name = "qwen3"
     opener = "<think>"
     closer = "</think>"
 
+    def test_markers_after_the_first_close_are_literal(self):
+        text = "plan</think>\n\nSTART = \"<think>\"\nEND = \"</think>\"\n"
+        expected = ("plan", 'START = "<think>"\nEND = "</think>"')
+        parser = ReasoningParser(self.name, force_reasoning=True)
+        assert parser.parse_non_stream(text) == expected
+        parser = ReasoningParser(self.name, force_reasoning=True)
+        parts = [parser.parse_stream_chunk(char) for char in text] + [parser.flush()]
+        assert tuple("".join(pair[i] for pair in parts) for i in (0, 1)) == expected
+
 
 class TestMiniMaxM3ReasoningParser(_ReasoningContract):
     name = "minimax_m3"
diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py
index eb6e3a2ed2..2226fe7fe5 100644
--- a/tests/server/test_streaming_model_matrix.py
+++ b/tests/server/test_streaming_model_matrix.py
@@ -708,3 +708,212 @@ def norm(s):
 
     assert norm(_content(events)) == norm(full.content)
     assert [c.name for c in _calls(events)] == [c.name for c in full.tool_calls] == ["read"]
+
+
+# --------------------------------------------------------------------------- #
+# Axis 3: production Qwen3.8 parsers (qwen3 + qwen3_coder), stream == one-shot
+# --------------------------------------------------------------------------- #
+QWEN_TOOLS = TOOLS + [
+    {
+        "type": "function",
+        "function": {
+            "name": "edit",
+            "parameters": {
+                "type": "object",
+                "properties": {"path": {"type": "string"}, "new": {"type": "string"}},
+            },
+        },
+    },
+]
+
+
+def _qwen_parity(pieces, *, thinking=True, finish_reason=None):
+    """Feed token texts one per ack to the streaming path and joined to the
+    non-streaming path; assert both expose the same content, reasoning and calls
+    (exactly, no strip) and return them."""
+    spec = GenSpec(
+        messages=[], sampling_params=None, parser_tools=QWEN_TOOLS, template_tools=QWEN_TOOLS,
+        chat_template_kwargs={} if thinking else {"enable_thinking": False},
+    )
+
+    def state(chunks):
+        return FakeState(chunks, tool_call_parser="qwen3_coder", reasoning_parser="qwen3",
+                         finish_reason=finish_reason)
+
+    async def stream(chunks):
+        return [ev async for ev in generate_events(42, spec, state(chunks))]
+
+    def calls(items):
+        return [(c.name, json.loads(c.parameters)) for c in items]
+
+    full = asyncio.run(generate_full(42, spec, state(["".join(pieces)])))
+    one_shot = (full.content, full.reasoning, calls(full.tool_calls), full.finish_reason)
+    # One token per delta (what the detokenizer sends), then the worst case: one char.
+    for chunks in (list(pieces), list("".join(pieces))):
+        events = asyncio.run(stream(chunks))
+        streamed = (_content(events), _reasoning(events), calls(_calls(events)),
+                    events[-1].finish_reason)
+        assert streamed == one_shot
+    return one_shot
+
+
+EDIT_WITH_THINK_LITERALS = [
+    "plan", "</think>", "\n\n", "Edit:", "<tool_call>", "\n<function=edit>\n<parameter=new>\n",
+    'START = "', "<think>", '"\nEND = "', "</think>", '"\n', "</parameter>\n</function>\n",
+    "</tool_call>",
+]
+
+
+def test_qwen_think_literals_in_tool_arguments_stay_in_the_call():
+    # Editing a parser's own source: the markers are argument data, and reasoning
+    # ended at the first </think> for good.
+    content, reasoning, calls, finish = _qwen_parity(EDIT_WITH_THINK_LITERALS)
+    assert reasoning == "plan"
+    assert content == "Edit:"
+    assert calls == [("edit", {"new": 'START = "<think>"\nEND = "</think>"'})]
+    assert finish == "tool_calls"
+
+
+@pytest.mark.parametrize("pieces, expected", [
+    (["r", "</think>", "\n\n", "Use ", "<think>", " tags", " and ", "</think>", "."],
+     ("Use <think> tags and </think>.", "r")),
+    (["r", "</think>", "\n\n", "if a ", "<", " b:", " pass"], ("if a < b: pass", "r")),
+    (["r", "\n", "</think>", "\n\nAnswer", " here", "\n"], ("Answer here", "r")),
+    (["\n", "r", "\n\n", "s", "</think>", "\n\n", "A", "\n\n", "B"], ("A\n\nB", "r\n\ns")),
+    (["<think>", "\n", "r", "</think>", "\n\n", "a"], ("a", "r")),
+], ids=["think-literals-in-content", "code-lt", "ws-glued", "ws-separate", "explicit-think"])
+def test_qwen_reasoning_split_parity(pieces, expected):
+    content, reasoning, calls, _ = _qwen_parity(pieces)
+    assert (content, reasoning) == expected
+    assert calls == []
+
+
+@pytest.mark.parametrize("pieces, expected", [
+    (["Use", " <think>", " here", " </think>", " too"], "Use <think> here </think> too"),
+    (["\n", "Plain", " answer", "\n"], "\nPlain answer\n"),
+], ids=["think-literals", "literal-ws"])
+def test_qwen_start_in_content_is_literal(pieces, expected):
+    # enable_thinking=false: the template closed the think block, so nothing the
+    # model writes (except an explicit leading <think>) reopens it.
+    content, reasoning, calls, _ = _qwen_parity(pieces, thinking=False)
+    assert (content, reasoning, calls) == (expected, "", [])
+
+
+def test_qwen_start_in_content_explicit_leading_think():
+    content, reasoning, _, _ = _qwen_parity(["<think>", "r", "</think>", "\n\n", "a"], thinking=False)
+    assert (content, reasoning) == ("a", "r")
+
+
+def test_qwen_truncated_in_reasoning():
+    content, reasoning, calls, finish = _qwen_parity(
+        ["thinking", " more", "\n", "</th"], finish_reason="length")
+    assert (content, reasoning, calls, finish) == ("", "thinking more\n</th", [], "length")
+
+
+def test_qwen_repeated_tool_call_opener_is_one_block():
+    content, _, calls, _ = _qwen_parity(["r", "</think>", "\n\n", "<tool_call>", "\n"] + _qwen_call("/a"))
+    assert (content, calls) == ("", [("read", {"filePath": "/a"})])
+
+
+def test_qwen_tool_call_only_turn():
+    content, reasoning, calls, finish = _qwen_parity([
+        "r", "\n", "</think>", "\n\n", "<tool_call>", "\n", "<function=read>", "\n",
+        "<parameter=filePath>", "\n/tmp/test_calc.py\n", "</parameter>", "\n", "</function>",
+        "\n", "</tool_call>",
+    ])
+    assert (content, reasoning, finish) == ("", "r", "tool_calls")
+    assert calls == [("read", READ_ARGS)]
+
+
+def _qwen_call(path):
+    return ["<tool_call>", "\n", "<function=read>", "\n", "<parameter=filePath>", f"\n{path}\n",
+            "</parameter>", "\n", "</function>", "\n", "</tool_call>"]
+
+
+@pytest.mark.parametrize("between, after, expected", [
+    (["\nThen more.\n"], [], "\nThen more.\n"),
+    (["\n", "Then", " more.", "\n"], ["\n"], "\nThen more.\n"),
+    (["\n\n"], [" Done", "."], " Done."),
+    (["\n", "\n"], [" ", "Done", "."], " Done."),
+    (["\n"], ["\n\n", "<div>", "x"], "\n\n<div>x"),
+    (["\n"], [" ", "<T>", " generic"], " <T> generic"),
+    (["\n"], ["\n", "</tool_call>"], "\n</tool_call>"),
+    (["\n"], ["\n", "<tool_call>", " tags"], "\n<tool_call> tags"),
+], ids=["glued", "separate", "ws-only-glued-tail", "ws-only-separate-tail", "ws-then-tag",
+        "space-then-tag", "stray-close", "literal-open"])
+def test_qwen_text_around_calls_parity(between, after, expected):
+    # Text before, between and after calls is content in both modes, verbatim.
+    # A whitespace-only gap between calls or after the last one is markup
+    # separation and is dropped, however it is chunked.
+    pieces = (["r", "</think>", "\n\n", "Let me look.", "\n\n"] + _qwen_call("/a") + between
+              + _qwen_call("/b") + after)
+    content, reasoning, calls, finish = _qwen_parity(pieces)
+    assert content == "Let me look.\n\n" + expected
+    assert calls == [("read", {"filePath": "/a"}), ("read", {"filePath": "/b"})]
+    assert (reasoning, finish) == ("r", "tool_calls")
+
+
+def test_qwen_empty_tool_block_is_dropped_in_both_modes():
+    content, _, calls, finish = _qwen_parity(
+        ["r", "</think>", "\n\n", "A", "<tool_call>", "\n", "</tool_call>", "\n", "B"])
+    assert (content, calls, finish) == ("A\nB", [], "stop")
+
+
+def test_qwen_truncated_inside_function_tag_is_not_recovered():
+    # Cut off by max_tokens inside "<function=x": no call in either mode.
+    _, _, calls, finish = _qwen_parity(
+        ["</think><tool_cal", "l>\n", ">x<", "div><tool_call><function=x", "\n\n"],
+        thinking=False, finish_reason="length")
+    assert (calls, finish) == ([], "length")
+
+
+@pytest.mark.parametrize("pieces", [
+    ["The parser looks for ", "<tool_call>", " tags", " and splits", " them."],
+    ["Wrap calls in ", "<tool_call>", "...", "</tool_call>", " tags, then answer."],
+    ["Calls:\n\n", "<tool_call>", "\n", '{"name": "read"}', "\n", "</tool_call>"],
+], ids=["opener", "pair", "not-an-invoke"])
+def test_qwen_tool_call_tag_in_prose_stays_content(pieces):
+    # <tool_call> followed by anything but an invoke or its close is the tag
+    # named in text, like a bare <function= lookalike.
+    content, _, calls, finish = _qwen_parity(["r", "</think>", "\n\n"] + pieces)
+    assert (content, calls, finish) == ("".join(pieces), [], "stop")
+
+
+BARE_EDIT = ["<", "function", "=", "edit", ">\n", "<parameter=path>\n", "a.py\n", "</parameter>\n",
+             "</function>"]
+
+
+def test_qwen_bare_function_is_a_call_in_both_modes():
+    # The model skipped <tool_call>: an agent still needs the call executed, and
+    # neither mode may leave the markup in content.
+    content, _, calls, finish = _qwen_parity(
+        ["r", "</think>", "\n\n", "Doing it.\n"] + BARE_EDIT + ["\n", "Done", "."])
+    assert (content, finish) == ("Doing it.\n\nDone.", "tool_calls")
+    assert calls == [("edit", {"path": "a.py"})]
+
+
+def test_qwen_bare_function_after_a_wrapped_call():
+    content, _, calls, _ = _qwen_parity(
+        ["r", "</think>", "\n\n"] + _qwen_call("/a") + ["\n"] + BARE_EDIT)
+    assert content == ""
+    assert calls == [("read", {"filePath": "/a"}), ("edit", {"path": "a.py"})]
+
+
+def test_qwen_bare_function_swallows_a_stray_close():
+    content, _, calls, _ = _qwen_parity(
+        ["r", "</think>", "\n\n", "Doing it.\n"] + BARE_EDIT + ["\n", "</tool_call>", "\n", "Done"])
+    assert content == "Doing it.\n\nDone"
+    assert calls == [("edit", {"path": "a.py"})]
+
+
+@pytest.mark.parametrize("pieces", [
+    ["Write ", "<", "function", "=", "name", ">", " tags."],
+    ["<", "function", "=", "nope", ">\n", "<parameter=x>\n", "1\n", "</parameter>\n", "</function>"],
+    ["x ", "<", "function", "=", "read", ">", " y"],
+    ['p = "', "<", "function", "=", '"', " + name"],
+    ["if a ", "<", "f", "(b):", " pass"],
+    ["tail ", "<", "function"],
+], ids=["undeclared-name", "unknown-tool", "no-body", "code-literal", "lt-call", "truncated-prefix"])
+def test_qwen_bare_function_lookalikes_stay_content(pieces):
+    content, _, calls, finish = _qwen_parity(["r", "</think>", "\n\n"] + pieces)
+    assert (content, calls, finish) == ("".join(pieces), [], "stop")
diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py
index 9767ffb25e..b2573e4ea1 100644
--- a/tests/tokenizer/test_detokenize_extras.py
+++ b/tests/tokenizer/test_detokenize_extras.py
@@ -135,3 +135,112 @@ def test_text_offset_is_where_the_token_starts_despite_holdback():
         offsets.append(offs[0])
     assert sent == ["b", "", "cd", ""]
     assert offsets == [0, 1, 2, 3]
+
+
+PARTIAL, COMPLETES = 50, 51
+
+
+class _ByteTok(_Tok):
+    """50 is the first half of a two-token character: alone it decodes to U+FFFD."""
+
+    def decode(self, ids, skip_special_tokens=False):
+        out = []
+        for i, t in enumerate(ids):
+            if t == PARTIAL:
+                out.append("" if i + 1 < len(ids) and ids[i + 1] == COMPLETES else "�")
+            elif t == COMPLETES:
+                out.append("é" if i and ids[i - 1] == PARTIAL else "�")
+            else:
+                out.append(super().decode([t], skip_special_tokens))
+        return "".join(out)
+
+
+def _run_logprobs(manager, tokens, *, finished_kwargs=None, **msg_kwargs):
+    """(text, entry token ids) per reply, each token's entry carrying its own id."""
+    replies = []
+    for i, t in enumerate(tokens):
+        last = i == len(tokens) - 1
+        kwargs = dict(msg_kwargs)
+        if last and finished_kwargs:
+            kwargs.update(finished_kwargs)
+        texts, _, released = manager.detokenize_with_logprobs(
+            [DetokenizeMsg(uid=1, next_token=t, finished=last, **kwargs)], [{"token_id": t}]
+        )
+        replies.append(
+            (texts[0], [(e["token_id"], bool(e.get("stop_trimmed"))) for e in released[0]])
+        )
+    return replies
+
+
+def test_logprob_entries_of_stop_trimmed_tokens_are_marked():
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    assert _run_logprobs(
+        manager, [1, 2, 3], stop_strs=["cd"], finished_kwargs={"matched_stop": "cd"}
+    ) == [("b", [(1, False)]), ("", []), ("", [(2, True), (3, True)])]
+
+
+def test_logprob_entry_of_a_token_a_stop_string_cuts_inside_carries_its_visible_length():
+    # "b</think>c" with stop "k>c": "</thin" of the </think> token stays visible.
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    final = None
+    for i, t in enumerate([1, THINK_END, 2]):
+        last = i == 2
+        _, _, released = manager.detokenize_with_logprobs(
+            [DetokenizeMsg(uid=1, next_token=t, finished=last, stop_strs=["k>c"],
+                           matched_stop="k>c" if last else None)],
+            [{"token_id": t}],
+        )
+        final = released[0]
+    # each entry also carries where its token starts in the text ("b" is 1 char, "</think>" 8)
+    assert final == [
+        {"token_id": THINK_END, "text_offset": 1, "stop_visible": 6},
+        {"token_id": 2, "text_offset": 9, "stop_trimmed": True},
+    ]
+
+
+def test_logprob_entry_waits_for_its_text_when_a_stop_prefix_is_released():
+    # "b" is held as a prefix of "bc"; the second "b" releases the first, but its own
+    # text is held, so its entry must not ride that reply -- the stop then trims it.
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    assert _run_logprobs(
+        manager, [1, 1, 2], stop_strs=["bc"], finished_kwargs={"matched_stop": "bc"}
+    ) == [("", []), ("b", [(1, False)]), ("", [(1, True), (2, True)])]
+
+
+def test_logprob_entries_kept_with_the_stop_string_or_without_a_match():
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    kept = _run_logprobs(
+        manager, [1, 2, 3], stop_strs=["cd"], keep_stop_str=True,
+        finished_kwargs={"matched_stop": "cd"},
+    )
+    assert kept[-1] == ("cd", [(2, False), (3, False)])
+    # finished by length while "c" was held as a stop prefix: nothing was trimmed
+    by_length = _run_logprobs(manager, [1, 2], stop_strs=["cd"])
+    assert by_length == [("b", [(1, False)]), ("c", [(2, False)])]
+    # no stop strings: every entry rides its own reply, as before
+    plain = _run_logprobs(manager, [1, 2, EOS], finished_kwargs={"finish_reason": "stop"})
+    assert plain == [("b", [(1, False)]), ("c", [(2, False)]), ("", [(EOS, False)])]
+
+
+def test_logprob_entry_of_a_partial_character_rides_the_reply_that_completes_it():
+    manager = DetokenizeManager(_ByteTok(), frozenset({EOS}))
+    assert _run_logprobs(manager, [1, PARTIAL, COMPLETES, 2], stop_strs=["zz"]) == [
+        ("b", [(1, False)]),
+        ("", []),
+        ("é", [(PARTIAL, False), (COMPLETES, False)]),
+        ("c", [(2, False)]),
+    ]
+
+
+def test_logprob_entries_released_together_keep_their_own_text_offsets():
+    # "c" is held as a prefix of "cx" and released with "d": that reply carries both
+    # entries, each stamped with where its own token starts, not one per-reply offset.
+    manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+    replies = []
+    for i, t in enumerate([1, 2, 3, EOS]):
+        texts, _, released = manager.detokenize_with_logprobs(
+            [DetokenizeMsg(uid=1, next_token=t, finished=i == 3, stop_strs=["cx"])],
+            [{"token_id": t}],
+        )
+        replies.append((texts[0], [(e["token_id"], e["text_offset"]) for e in released[0]]))
+    assert replies == [("b", [(1, 0)]), ("", []), ("cd", [(2, 1), (3, 2)]), ("", [(EOS, 3)])]

From bde037162a77355fe7573dad3e442fe111e47b0b Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 11:13:29 -0400
Subject: [PATCH 225/249] perf(tokenizer): skip the logprob hold bookkeeping
 when nothing is held

The stop-string hold block ran for every message, building and replacing
empty lists even for requests that never asked for logprobs, about 0.25 us
per message (6-14% of the detokenizer's per-message cost on the Qwen3.8
tokenizer). Move it into DecodeStatus.release_logprobs and call it only when
the message carries an entry or entries are still held. The new test pins
that held entries are still released by a message without its own entry.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/tokenizer/detokenize.py  | 55 ++++++++++++++---------
 tests/tokenizer/test_detokenize_extras.py | 24 ++++++++++
 2 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index c2e9bee059..44c3149948 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -87,6 +87,34 @@ def reasoning_tokens(self) -> int:
         # stopped inside the block (max_tokens, or EOS before the end tag): all of it
         return self.generated if self.in_reasoning else 0
 
+    def release_logprobs(
+        self, entry: dict | None, token_start: int, emit_end: int, text_len: int, flush: bool
+    ) -> List[dict]:
+        """Hold this token's logprob entry and release those whose text is now sent."""
+        if entry is not None:
+            self.held_logprobs.append((token_start, {**entry, "text_offset": token_start}))
+        out: List[dict] = []
+        # A held token's text ends where the next one's starts.
+        starts = [start for start, _ in self.held_logprobs]
+        ends = starts[1:] + [text_len] if starts else []
+        if flush:
+            trimmed = emit_end < text_len
+            for (start, held), end in zip(self.held_logprobs, ends, strict=True):
+                if trimmed and start >= emit_end:
+                    held = {**held, "stop_trimmed": True}
+                elif trimmed and emit_end < end:
+                    held = {**held, "stop_visible": emit_end - start}
+                out.append(held)
+            self.held_logprobs = []
+        elif self.held_logprobs:
+            # Wait for all of a token's text, so a stop match inside it can still mark it.
+            n = 0
+            while n < len(ends) and starts[n] < emit_end and ends[n] <= emit_end:
+                n += 1
+            out = [held for _, held in self.held_logprobs[:n]]
+            self.held_logprobs = self.held_logprobs[n:]
+        return out
+
     def update(self, read_str: str, surr_str: str) -> str:
         new_text = read_str[len(surr_str) :]
         if len(new_text) > 0 and not new_text.endswith("�"):
@@ -253,28 +281,11 @@ def _detokenize(
             s.sent_offset = max(prev_sent, emit_end)
             incremental_strs.append(incremental_output)
             reasoning_tokens.append(s.reasoning_tokens() if msg.finished else 0)
-            if entry is not None:
-                s.held_logprobs.append((token_start, {**entry, "text_offset": token_start}))
-            out: List[dict] = []
-            # A held token's text ends where the next one's starts.
-            starts = [start for start, _ in s.held_logprobs]
-            ends = starts[1:] + [len(output_str)] if starts else []
-            if msg.finished or not msg.stop_strs:
-                trimmed = emit_end < len(output_str)
-                for (start, held), end in zip(s.held_logprobs, ends, strict=True):
-                    if trimmed and start >= emit_end:
-                        held = {**held, "stop_trimmed": True}
-                    elif trimmed and emit_end < end:
-                        held = {**held, "stop_visible": emit_end - start}
-                    out.append(held)
-                s.held_logprobs = []
-            elif s.held_logprobs:
-                # Wait for all of a token's text, so a stop match inside it can still mark it.
-                n = 0
-                while n < len(ends) and starts[n] < emit_end and ends[n] <= emit_end:
-                    n += 1
-                out = [held for _, held in s.held_logprobs[:n]]
-                s.held_logprobs = s.held_logprobs[n:]
+            if entry is not None or s.held_logprobs:
+                flush = msg.finished or not msg.stop_strs
+                out = s.release_logprobs(entry, token_start, emit_end, len(output_str), flush)
+            else:
+                out = []  # no logprobs asked for: skip the hold bookkeeping
             released.append(out)
             if msg.finished:
                 del self.decode_map[msg.uid]
diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py
index b2573e4ea1..01e02e63ee 100644
--- a/tests/tokenizer/test_detokenize_extras.py
+++ b/tests/tokenizer/test_detokenize_extras.py
@@ -244,3 +244,27 @@ def test_logprob_entries_released_together_keep_their_own_text_offsets():
         )
         replies.append((texts[0], [(e["token_id"], e["text_offset"]) for e in released[0]]))
     assert replies == [("b", [(1, 0)]), ("", []), ("cd", [(2, 1), (3, 2)]), ("", [(EOS, 3)])]
+
+
+def test_held_logprob_entries_release_on_a_message_without_its_own_entry():
+    # "c" is held as a prefix of "cx"; the messages that release it carry no entry
+    # (and skip the hold bookkeeping when nothing is held), yet "c"'s entry still rides them.
+    def run(tokens, entries, last_finishes):
+        manager = DetokenizeManager(_Tok(), frozenset({EOS}))
+        replies = []
+        for i, (t, entry) in enumerate(zip(tokens, entries, strict=True)):
+            finished = last_finishes and i == len(tokens) - 1
+            texts, _, released = manager.detokenize_with_logprobs(
+                [DetokenizeMsg(uid=1, next_token=t, finished=finished, stop_strs=["cx"])], [entry]
+            )
+            replies.append((texts[0], [e["token_id"] for e in released[0]]))
+        assert manager.decode_map.get(1) is None or not manager.decode_map[1].held_logprobs
+        return replies
+
+    e = [{"token_id": t} for t in (1, 2)]
+    # released mid-stream by "d", whose message has no entry
+    assert run([1, 2, 3], [*e, None], False) == [("b", [1]), ("", []), ("cd", [2])]
+    # flushed by the finishing message, which has no entry
+    assert run([1, 2, EOS], [*e, None], True) == [("b", [1]), ("", []), ("c", [2])]
+    # a request without logprobs never carries any
+    assert run([1, 2, 3], [None] * 3, True) == [("b", []), ("", []), ("cd", [])]

From ab70a58004807317463c779fb6a655f66a4eca1b Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Thu, 15 Oct 2026 11:14:21 -0400
Subject: [PATCH 226/249] perf(server): collect non-streaming chat logprob
 entries only when asked

_generate_full_impl ran untrimmed() on every ack even for requests without
logprobs, and its result then kept whatever raw, unsegmented entries the
acks carried. Read the flag once and skip both, so the protocol-neutral
result carries entries only for a request that asked for them, as the
streaming path already does.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt
---
 python/freetoken/server/generation.py |  6 ++++--
 tests/server/test_logprobs_api.py     | 13 +++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 5a9192fbd7..e2d4d343a4 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -1390,6 +1390,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
     tool calls, strip special tokens. The adapters format the GenResult into their wire."""
     full_content = ""
     logprob_entries: list[dict] = []
+    wants_logprobs = _wants_logprobs(spec)
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
@@ -1405,13 +1406,14 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
         full_content += ack.incremental_output
-        logprob_entries.extend(untrimmed(getattr(ack, "logprobs", None)))
+        if wants_logprobs:
+            logprob_entries.extend(untrimmed(getattr(ack, "logprobs", None)))
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
             break
 
-    if _wants_logprobs(spec):
+    if wants_logprobs:
         segmenter = await _make_logprob_segmenter(spec, state)
         logprob_entries = (
             segmenter.feed(logprob_entries) + segmenter.finish() if segmenter is not None else []
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 53ed82f4d2..f983e6d9b5 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -9,6 +9,7 @@
 import pytest
 
 from freetoken.message import TokenizeMsg, UserReply
+from freetoken.server.generation import generate_full
 from freetoken.server.openai_api import (
     ChatCompletionRequest,
     CompletionRequest,
@@ -842,5 +843,17 @@ def test_chat_logprobs_are_dropped_when_the_segmenter_cannot_build() -> None:
     assert _streamed_tokens(choices) == []
 
 
+def test_generate_full_ignores_entries_when_logprobs_were_not_requested() -> None:
+    # The protocol-neutral result carries entries only for a request that asked for them,
+    # so an adapter reading it never sees unsegmented (reasoning) entries.
+    state = FakeState(
+        token_replies(["thought", "</think>", "answer"]), reasoning_parser="qwen3"
+    )
+    spec = chat_request_to_genspec(chat_request(), {})
+    result = run(generate_full(42, spec, state))
+    assert result.content == "answer"
+    assert result.logprobs == []
+
+
 async def _collect(iterator):
     return [chunk async for chunk in iterator]

From 018c4c99ab354a01c4b63cfb15209349d6dc55e5 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:40:02 -0400
Subject: [PATCH 227/249] fix(checkpoint): load an ftw whose nvfp4 global
 scales are fp16

Fork-only shim over #739: FTWReader warns once and the dense, host-bank and
device-bank read paths upcast fp16 NVFP4 globals to the fp32 the kernels read,
so an FTW converted before #739 serves with its old numerics until reconverted.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/checkpoint/ftw.py        | 59 ++++++++++++++++++-----
 tests/checkpoint/test_ftw_fp16_globals.py | 58 ++++++++++++++++++++++
 2 files changed, 104 insertions(+), 13 deletions(-)
 create mode 100644 tests/checkpoint/test_ftw_fp16_globals.py

diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index 8dc53a4cff..23d5278441 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -431,6 +431,10 @@ def _is_nvfp4_global(name: str) -> bool:
     return (m["base"] if m else name).endswith(("gate_up_global", "down_global", ".weight_global"))
 
 
+# fork-only: FTW dirs already warned about fp16 NVFP4 globals (one reader is opened per load stage)
+_FP16_GLOBALS_WARNED: set[str] = set()
+
+
 def is_ftw_checkpoint(path: str) -> bool:
     """True if ``path`` is a directory holding a FreeToken Weight (FTW) index."""
     return os.path.isfile(os.path.join(path, INDEX_NAME))
@@ -578,11 +582,14 @@ def __init__(self, path: str):
         self.dir = path
         self.shards, tensors = _validate_ftw_v1_index(path, self.index)
         self.tensors = {tensor["name"]: tensor for tensor in tensors}
-        fp16_globals = [n for n, t in self.tensors.items() if t["dtype"] == "float16" and _is_nvfp4_global(n)]
-        if fp16_globals:
-            raise ValueError(
-                f"{path!r} stores NVFP4 global scales as fp16 (e.g. {fp16_globals[0]!r}) but FreeToken reads them as fp32; "
-                "reconvert it from the source checkpoint with `ft checkpoint`"
+        # fork-only: an FTW converted before #739 keeps NVFP4 globals in fp16; the read paths upcast them to fp32
+        self.fp16_globals = {n for n, t in self.tensors.items() if t["dtype"] == "float16" and _is_nvfp4_global(n)}
+        if self.fp16_globals and os.path.realpath(path) not in _FP16_GLOBALS_WARNED:
+            _FP16_GLOBALS_WARNED.add(os.path.realpath(path))
+            logger.warning(
+                f"{path!r} stores {len(self.fp16_globals)} NVFP4 global scale entries as fp16 "
+                f"(e.g. {min(self.fp16_globals)!r}); upcasting them to fp32 at load, the fp16 rounding stays. "
+                "Reconvert with `ft checkpoint` for exact globals"
             )
         self._fds: dict[str, int] = {}
         self._maps: dict[str, tuple[mmap.mmap, memoryview]] = {}
@@ -774,6 +781,8 @@ def _producer():
                 else:
                     # torch.frombuffer rejects count=0 even with a non-empty backing map.
                     t = torch.empty(tuple(e["shape"]), dtype=dt)
+                if e["name"] in reader.fp16_globals:
+                    t = t.float()
                 if not _put((e["name"], t, buf, e["nbytes"])):
                     return
         except BaseException as ex:  # surface to consumer
@@ -952,15 +961,22 @@ def _backing(layer_id: int) -> str:
             )
         row_hb[name] = []
         row_view_args[name] = []
+        upcast = name in reader.fp16_globals
         for layer_id in range(num_layers):
             off = e["global_off"] + layer_id * layer_bytes
             win_off = (off // ALIGN) * ALIGN
             win_end = _align_up(off + layer_bytes)
             head_pad = off - win_off
-            bank = HostBank((win_end - win_off,), torch.uint8, backing=_backing(layer_id))
+            if upcast:
+                # read the fp16 window into a transient buffer, keep an fp32 bank
+                bank = HostBank((num_experts, *row_shape), torch.float32, backing=_backing(layer_id))
+                row_view_args[name].append(None)
+            else:
+                bank = HostBank((win_end - win_off,), torch.uint8, backing=_backing(layer_id))
+                row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype))
             row_hb[name].append(bank)
-            row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype))
-            row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
+            row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id,
+                             (head_pad, (num_experts, *row_shape)) if upcast else None))
 
     num_experts = next((e["shape"][0] for by_layer in per_layer_groups.values() for e in by_layer.values()), None)
     for base, by_layer in per_layer_groups.items():
@@ -977,7 +993,9 @@ def _backing(layer_id: int) -> str:
                 raise FTWFormatError(
                     f"FTW bank {base!r} layer {layer_id} is not {ALIGN}-aligned"
                 )
-            bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id))
+            upcast = e["name"] in reader.fp16_globals
+            dtype = torch.float32 if upcast else _dtype_of(e["dtype"])
+            bank = HostBank(tuple(e["shape"]), dtype, backing=_backing(layer_id))
             row_hb[base].append(bank)
             row_view_args[base].append(None)
             layer_jobs.append((base, bank, e, layer_id))
@@ -998,15 +1016,28 @@ def _read_alpha(e):
                 bar.update(e["nbytes"])
 
             def _read_row(job):
-                _name, bank, win_off, win_len, layer_bytes, layer_id = job
-                reader.read_into(bank.memoryview(), {"global_off": win_off, "nbytes": win_len},
-                                 workers=workers, chunk=chunk)
+                _name, bank, win_off, win_len, layer_bytes, layer_id, upcast = job
+                window = {"global_off": win_off, "nbytes": win_len}
+                if upcast is None:
+                    reader.read_into(bank.memoryview(), window, workers=workers, chunk=chunk)
+                else:
+                    head_pad, shape = upcast
+                    buf = _transient_buffer(win_len)
+                    reader.read_into(memoryview(buf), window, workers=workers, chunk=chunk)
+                    raw = torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + layer_bytes)[head_pad:]
+                    bank.tensor.copy_(raw.view(torch.float16).view(shape))
                 pins.submit(bank, residency[layer_id])
                 bar.update(layer_bytes)
 
             def _read_layer(job):
                 _name, bank, entry, layer_id = job
-                reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk)
+                if entry["name"] in reader.fp16_globals:
+                    buf = _transient_buffer(entry["nbytes"])
+                    reader.read_into(memoryview(buf), entry, workers=workers, chunk=chunk)
+                    raw = torch.frombuffer(buf, dtype=torch.float16, count=entry["nbytes"] // 2)
+                    bank.tensor.copy_(raw.view(bank.tensor.shape))
+                else:
+                    reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk)
                 pins.submit(bank, residency[layer_id])
                 bar.update(entry["nbytes"])
 
@@ -1137,6 +1168,8 @@ def _read(job) -> torch.Tensor:
                 ahead = ex.submit(_read, jobs[0])
                 for k, (name, layer_id, *_rest) in enumerate(jobs):
                     host = ahead.result()
+                    if host.dtype == torch.float16 and _is_nvfp4_global(name):
+                        host = host.float()
                     if k + 1 < len(jobs):
                         ahead = ex.submit(_read, jobs[k + 1])
                     if layer_id is None:
diff --git a/tests/checkpoint/test_ftw_fp16_globals.py b/tests/checkpoint/test_ftw_fp16_globals.py
new file mode 100644
index 0000000000..934d7570df
--- /dev/null
+++ b/tests/checkpoint/test_ftw_fp16_globals.py
@@ -0,0 +1,58 @@
+"""An FTW converted before #739 stores NVFP4 global scales as fp16: the reader upcasts them to the fp32 the kernels read."""
+
+import logging
+
+import torch
+
+from freetoken.checkpoint import ftw
+from freetoken.checkpoint.ftw import FTWReader, FTWWriter, iter_ftw_weights, layer_bank_entry_name, load_ftw_banks
+
+E = 4
+
+
+def _globals(rows: int, seed: int) -> torch.Tensor:
+    # the served magnitudes (2e-5 to 4e-4) sit at fp16's normal minimum
+    g = torch.Generator().manual_seed(seed)
+    return (torch.rand(E, rows, generator=g) * 4e-4 + 2e-5).to(torch.float16)
+
+
+def _write(path, *, per_layer: bool) -> dict[str, torch.Tensor]:
+    src = {
+        "gate_up_global": _globals(6, 0),
+        "down_global": _globals(8, 1),
+        "gate_up": torch.arange(E * 6 * 4, dtype=torch.uint8).view(E, 6, 4),
+    }
+    writer = FTWWriter(str(path))
+    writer.add_tensor("model.layers.0.mlp.shared_expert.gate_proj.weight_global", torch.tensor(3e-4, dtype=torch.float16))
+    writer.add_tensor("model.layers.0.input_layernorm.weight", torch.ones(8, dtype=torch.float16))
+    for name, tensor in src.items():
+        writer.add_tensor(layer_bank_entry_name(name, 0) if per_layer else name, tensor, kind="experts_bank")
+    writer.finalize({})
+    return src
+
+
+def test_fp16_globals_load_as_fp32(tmp_path, monkeypatch, caplog):
+    monkeypatch.setenv("FREETOKEN_SKIP_BANK_PIN", "1")
+    monkeypatch.setenv("FREETOKEN_BANK_CUDA_ALLOC", "0")
+    monkeypatch.setattr(ftw, "_FP16_GLOBALS_WARNED", set(), raising=False)
+    for per_layer in (True, False):
+        path = tmp_path / ("per_layer" if per_layer else "flat")
+        src = _write(path, per_layer=per_layer)
+
+        with caplog.at_level(logging.WARNING):
+            caplog.clear()
+            FTWReader(str(path)).close()
+            dense = dict(iter_ftw_weights(str(path)))
+            banks = load_ftw_banks(str(path), num_layers=1)
+        assert sum("fp16" in r.getMessage() for r in caplog.records) == 1
+
+        g = dense["model.layers.0.mlp.shared_expert.gate_proj.weight_global"]
+        assert g.dtype == torch.float32 and g.item() == torch.tensor(3e-4, dtype=torch.float16).item()
+        # only the NVFP4 globals are upcast
+        assert dense["model.layers.0.input_layernorm.weight"].dtype == torch.float16
+
+        for name in ("gate_up_global", "down_global"):
+            bank = banks.sources[name][0]
+            assert bank.dtype == torch.float32
+            assert torch.equal(bank, src[name].float())
+        assert torch.equal(banks.sources["gate_up"][0], src["gate_up"])

From 24de7e68f667c4fea46e0a4ffea64f2ffd56d2d0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:32:06 -0400
Subject: [PATCH 228/249] fix(server): reject unimplemented vllm fields and
 empty prompts on completions

A set value of a result-changing vLLM extra (allowed_token_ids,
return_tokens_as_token_ids, guided_*, ...) is a 400 instead of being dropped
by extra="allow"; null/false/empty still pass. Empty prompts and more than
256 prompts x n are 400s, and an error reply's code (context_length_exceeded)
now reaches the completions body. On c23d2d5d the new deny-list test fails
for all 10 fields with 200 != 400.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/server/api_models.py |  2 +
 python/freetoken/server/openai_api.py | 53 +++++++++++++++++-
 tests/server/test_logprobs_api.py     | 77 +++++++++++++++++++++++++++
 3 files changed, 130 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 0da3d09e77..1bd43b9f49 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -167,6 +167,8 @@ def _sync_max_completion_tokens(self) -> "CompletionRequest":
 
 
 MAX_N = 16  # choices per request; each is a full generation, so keep the fan-out bounded
+# prompts x n per /v1/completions request: every one is queued before any is awaited
+MAX_PROMPTS = 256
 
 
 class TokenizeRequest(BaseModel):
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 4b43074b6e..55d740382a 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -16,6 +16,7 @@
 
 from .api_models import (
     MAX_N,
+    MAX_PROMPTS,
     DetokenizeRequest,
     TokenizeRequest,
     ChatCompletionRequest,
@@ -610,6 +611,11 @@ async def handle_completion(
         return create_error_response(unsupported)
     if not 1 <= req.n <= MAX_N:
         return create_error_response(f"n must be between 1 and {MAX_N}", param="n")
+    num_prompts = 1 if isinstance(req.prompt, str) else len(req.prompt)
+    if num_prompts * req.n > MAX_PROMPTS:
+        return create_error_response(
+            f"prompts x n must be at most {MAX_PROMPTS}, got {num_prompts * req.n}", param="prompt"
+        )
     try:  # surfaces an out-of-range value as a 400 rather than a 500 from the worker
         default_max_tokens = (
             getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS
@@ -668,7 +674,7 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
         prompt_tokens = completion_tokens = cached_tokens = 0
         async for ack in state.wait_for_ack(uid):
             if getattr(ack, "error", None):
-                return create_error_response(ack.error)
+                return create_error_response(ack.error, code=getattr(ack, "error_code", None))
             prompt_tokens += ack.prompt_tokens_delta
             completion_tokens += ack.completion_tokens_delta
             cached_tokens += ack.cached_tokens
@@ -748,7 +754,8 @@ async def stream_completion_chunks(
         yield _sse(_completion_chunk(req, uid, [{"text": echo_text, "index": index, "finish_reason": None, "logprobs": None}]))
     async for ack in state.wait_for_ack(uid):
         if getattr(ack, "error", None):
-            yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}})
+            error = {"message": ack.error, "type": "invalid_request_error", "code": getattr(ack, "error_code", None)}
+            yield _sse({"error": error})
             if terminal:
                 yield b"data: [DONE]\n\n"
             return
@@ -1098,7 +1105,49 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool
     return response_format.get("type") not in (None, "text", "json_object", "json_schema")
 
 
+# vLLM /v1/completions fields that change the result and are not implemented here: a set
+# value is a 400 rather than silently dropped by extra="allow".
+_UNIMPLEMENTED_COMPLETION_FIELDS = (
+    "allowed_token_ids",
+    "return_tokens_as_token_ids",
+    "logprob_token_ids",
+    "prompt_logprobs",
+    "guided_json",
+    "guided_regex",
+    "guided_choice",
+    "guided_grammar",
+    "structured_outputs",
+    "bad_words",
+)
+
+
+def _is_set(value: Any) -> bool:
+    # not truthiness: prompt_logprobs=0 is a real request, and 0 == False
+    if value is None or value is False:
+        return False
+    return not (isinstance(value, (str, list, dict)) and not value)
+
+
+def _empty_prompt(req: CompletionRequest) -> bool:
+    prompt = req.prompt
+    if isinstance(prompt, str):
+        # suffix-only fill-in-the-middle is a real request
+        return not prompt and req.suffix is None
+    if not prompt:
+        return True
+    return any(
+        (isinstance(item, list) and not item)
+        or (isinstance(item, str) and not item and req.suffix is None)
+        for item in prompt
+    )
+
+
 def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
+    for name in _UNIMPLEMENTED_COMPLETION_FIELDS:
+        if _is_set(getattr(req, name, None)):
+            return f"{name} is not supported by this server"
+    if _empty_prompt(req):
+        return "prompt must not be empty"
     if _is_token_prompt(req.prompt):
         return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead"
     logprobs_error = completion_logprobs_error(req)
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index f983e6d9b5..039a89fa1d 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -9,6 +9,7 @@
 import pytest
 
 from freetoken.message import TokenizeMsg, UserReply
+from freetoken.server import openai_api
 from freetoken.server.generation import generate_full
 from freetoken.server.openai_api import (
     ChatCompletionRequest,
@@ -857,3 +858,79 @@ def test_generate_full_ignores_entries_when_logprobs_were_not_requested() -> Non
 
 async def _collect(iterator):
     return [chunk async for chunk in iterator]
+
+
+def _completion(**kwargs) -> CompletionRequest:
+    return CompletionRequest(**{"model": "client-model", "prompt": "hello", "max_tokens": 1, **kwargs})
+
+
+def _one_token_state() -> FakeState:
+    return FakeState([reply("A", finished=True, logprobs=[entry(32, "A", -0.1)])])
+
+
+# Result-changing vLLM fields this server does not implement: a set value must be a 400,
+# never a 200 computed as if the field were absent.
+UNIMPLEMENTED_COMPLETION_FIELDS = [
+    ("allowed_token_ids", [32, 33]),
+    ("return_tokens_as_token_ids", True),
+    ("logprob_token_ids", [32]),
+    ("prompt_logprobs", 0),
+    ("guided_json", {"type": "object"}),
+    ("guided_regex", "[AB]"),
+    ("guided_choice", ["A", "B"]),
+    ("guided_grammar", "root ::= \"A\""),
+    ("structured_outputs", {"choice": ["A"]}),
+    ("bad_words", ["B"]),
+]
+
+
+@pytest.mark.parametrize(("field", "value"), UNIMPLEMENTED_COMPLETION_FIELDS)
+def test_completion_rejects_unimplemented_fields_when_set(field, value) -> None:
+    resp = run(handle_completion(_completion(logprobs=2, **{field: value}), None, _one_token_state(), {}))
+    assert getattr(resp, "status_code", 200) == 400
+    assert field in json.loads(resp.body)["error"]["message"]
+
+
+@pytest.mark.parametrize("value", [None, False, [], {}, ""])
+@pytest.mark.parametrize("field", [f for f, _ in UNIMPLEMENTED_COMPLETION_FIELDS])
+def test_completion_accepts_unimplemented_fields_left_unset(field, value) -> None:
+    resp = run(handle_completion(_completion(**{field: value}), None, _one_token_state(), {}))
+    assert resp["choices"][0]["text"] == "A"
+
+
+@pytest.mark.parametrize("prompt", ["", [], [""], ["ok", ""], [[]]])
+def test_completion_rejects_empty_prompts(prompt) -> None:
+    state = _one_token_state()
+    resp = run(handle_completion(_completion(prompt=prompt), None, state, {}))
+    assert resp.status_code == 400
+    assert state.sent is None  # nothing was queued
+
+
+def test_completion_suffix_only_fim_prompt_is_not_empty() -> None:
+    reason = openai_api._completion_unsupported_reason(_completion(prompt="", suffix="tail"))
+    assert reason is None
+
+
+@pytest.mark.parametrize(("num_prompts", "n", "status"), [(256, 1, 200), (257, 1, 400), (86, 3, 400)])
+def test_completion_caps_prompts_times_n(num_prompts, n, status) -> None:
+    replies = [reply("A", finished=True)]
+    resp = run(handle_completion(_completion(prompt=["p"] * num_prompts, n=n), None, FakeState(replies), {}))
+    assert getattr(resp, "status_code", 200) == status
+
+
+def _too_long() -> UserReply:
+    return UserReply(
+        uid=42, incremental_output="", finished=True,
+        error="prompt is too long: 9 tokens > 8 maximum", error_code="context_length_exceeded",
+    )
+
+
+def test_completion_passes_the_context_length_code_through() -> None:
+    resp = run(handle_completion(_completion(), None, FakeState([_too_long()]), {}))
+    error = json.loads(resp.body)["error"]
+    assert resp.status_code == 400
+    assert error["code"] == "context_length_exceeded"
+    assert error["message"].startswith("prompt is too long: ")
+
+    events = parse_sse(run(_collect(stream_completion_chunks(42, _completion(stream=True), FakeState([_too_long()])))))
+    assert events[0]["error"]["code"] == "context_length_exceeded"

From d4d2606c5802b49bf94b221be71050c550b334a0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:23:00 -0400
Subject: [PATCH 229/249] feat(server): honour allowed_token_ids on completions
 with restricted logprobs

Sampling is masked to the allowed ids at every position (a LogitsPlan mask
like min_tokens'). A row that also asks for logprobs gets them renormalized
over the set from the unprocessed logits, vLLM's processed_logprobs at
temperature 0, so the combinations that mode would fold in are 400s. The
-inf tail and padding are dropped from top entries, non-finite values are
sent as -9999, and completions logprobs go up to 32.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/openai_api.md                       |  8 ++-
 python/freetoken/core.py                 |  4 ++
 python/freetoken/engine/sample.py        | 49 +++++++++++++-
 python/freetoken/scheduler/scheduler.py  | 14 +++-
 python/freetoken/server/api_models.py    |  1 +
 python/freetoken/server/generation.py    | 13 ++++
 python/freetoken/server/logprobs.py      | 21 ++++--
 python/freetoken/server/openai_api.py    | 38 ++++++++++-
 python/freetoken/tokenizer/tokenize.py   |  4 ++
 tests/engine/test_logits_processors.py   | 23 +++++++
 tests/engine/test_sample.py              | 15 +++++
 tests/engine/test_sample_logprobs.py     | 76 +++++++++++++++++++++
 tests/server/test_logprobs_api.py        | 84 +++++++++++++++++++++++-
 tests/server/test_sampling_validation.py | 11 ++++
 tests/tokenizer/test_tokenize.py         | 21 ++++++
 15 files changed, 368 insertions(+), 14 deletions(-)

diff --git a/docs/openai_api.md b/docs/openai_api.md
index 0b1cc9d93c..02ad41832d 100644
--- a/docs/openai_api.md
+++ b/docs/openai_api.md
@@ -62,7 +62,13 @@ parameters as above; plus:
 | `echo` | honoured (the prompt leads the text; not with `logprobs`) | honoured | honoured | no |
 | `suffix` | honoured on models with fill-in-the-middle tokens (`<|fim_prefix|>` ... `<|fim_middle|>`): the prompt becomes prefix-suffix-middle | 400 | 400 | `/infill` |
 | `n` | honoured, prompt-major choice order (`prompt i`, sample `k` -> `index i*n+k`) | honoured | honoured | no |
-| `logprobs` | 400 (see #224 for the sampled-token logprobs) | honoured | honoured | `n_probs` |
+| `logprobs` | honoured, 0..32: raw pre-temperature logprobs of the sampled token and its top alternatives; -inf is sent as -9999 | honoured (`--max-logprobs`, default 20) | honoured | `n_probs` |
+| `allowed_token_ids` | honoured, 1..1024 ids below `len(tokenizer)`: sampling is masked to them at every position. With `logprobs` the values are renormalized over the set (vLLM's `processed_logprobs` at temperature 0), so `temperature > 0`, penalties, `min_tokens` and `logit_bias` are a 400 there; entries past the set are left out | honoured (raw logprobs unless `--logprobs-mode processed_logprobs`) | no | no |
+
+Set values of the vLLM fields this server does not implement (`logprob_token_ids`,
+`prompt_logprobs`, `guided_*`, `structured_outputs`, `bad_words`, ...) are a 400 rather than
+ignored; `null` / `false` / empty pass. Empty prompts and more than 256 prompts x `n` are a 400.
+An over-long prompt's error carries `code: "context_length_exceeded"`.
 
 ## Responses
 
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index bbba592378..ab0141b1e5 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -54,6 +54,9 @@ class SamplingParams:
     # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives.
     logprobs: bool = False
     top_logprobs: int = 0
+    # vLLM allowed_token_ids: sampling is masked to these ids (sorted, unique). A request that
+    # also asks for logprobs gets them renormalized over this set (Sampler.compute_logprobs).
+    allowed_token_ids: list[int] | None = None
     # Set by the tokenizer worker when the rendered prompt leaves a <think> block open, so
     # the detokenizer counts an unclosed generation as reasoning tokens.
     starts_in_reasoning: bool = False
@@ -70,6 +73,7 @@ def needs_logits_processing(self) -> bool:
             self.min_p > 0.0
             or self.logit_bias
             or self.min_tokens > 0
+            or self.allowed_token_ids
         )
 
 
diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index 62bdd29cc8..f07121ba9f 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -33,6 +33,9 @@ class LogitsPlan:
     min_rows: torch.Tensor | None = None  # [m] int64 local rows
     min_ids: torch.Tensor | None = None  # [m, S] int64
     min_p: torch.Tensor | None = None  # [r, 1] float32 (0 = off for that row)
+    # Rows with allowed_token_ids and their ids, padded with vocab_size like min_ids.
+    allow_rows: torch.Tensor | None = None  # [a] int64 local rows
+    allow_ids: torch.Tensor | None = None  # [a, A] int64
 
 
 class RowPenalty(NamedTuple):
@@ -62,6 +65,9 @@ class BatchSamplingArgs:
     logprob_row_idx: torch.Tensor | None = None
     logprob_all_rows: bool = False
     max_top_logprobs: int = 0
+    # Logprob rows restricted to their allowed_token_ids, indexing the logprob rows above.
+    restrict_rows: torch.Tensor | None = None  # [q] int64
+    restrict_ids: torch.Tensor | None = None  # [q, A] int64, padded with vocab_size
 
 
 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
@@ -76,6 +82,13 @@ def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch
     return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device)
 
 
+def disallowed_mask(ids: torch.Tensor, vocab_size: int) -> torch.Tensor:
+    """[a, V] bool, True outside each row's ids (the vocab_size padding column is dropped)."""
+    block = torch.ones((ids.shape[0], vocab_size + 1), dtype=torch.bool, device=ids.device)
+    block.scatter_(1, ids, False)
+    return block[:, :vocab_size]
+
+
 def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor:
     """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order:
     repetition (over prompt + generated ids), then frequency and presence (generated)."""
@@ -98,7 +111,8 @@ def apply_logits_processors(
     logits: torch.Tensor, plan: LogitsPlan, vocab_size: int
 ) -> torch.Tensor:
     """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the
-    min_tokens mask, then min_p. Rows outside the plan are copied unchanged."""
+    allowed_token_ids and min_tokens masks, then min_p. Rows outside the plan are copied
+    unchanged."""
     out = logits.to(torch.float32, copy=True)
     sub = out[plan.rows]  # [r, V] (advanced indexing copies)
     device = sub.device
@@ -107,6 +121,10 @@ def apply_logits_processors(
     if plan.bias_ids is not None:
         sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True)
 
+    if plan.allow_ids is not None:
+        rows = sub[plan.allow_rows]
+        sub[plan.allow_rows] = rows.masked_fill(disallowed_mask(plan.allow_ids, vocab_size), neg_inf)
+
     if plan.min_ids is not None:
         block = torch.zeros((plan.min_ids.shape[0], vocab_size + 1), dtype=torch.bool, device=device)
         block.scatter_(1, plan.min_ids, torch.ones_like(plan.min_ids, dtype=torch.bool))
@@ -199,10 +217,20 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
             plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device)
             plan.min_ids = _padded_rows(min_ids, pad, self.device)
 
+        allow = [(local, self._allowed(p)) for local, (_, p) in enumerate(picked)]
+        allow = [(local, ids) for local, ids in allow if ids]
+        if allow:
+            plan.allow_rows = make_device_tensor([local for local, _ in allow], torch.int64, self.device)
+            plan.allow_ids = _padded_rows([ids for _, ids in allow], pad, self.device)
+
         if any(p.min_p > 0.0 for _, p in picked):
             plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device)
         return plan
 
+    def _allowed(self, p) -> list[int]:
+        # an empty list after the filter leaves the row unmasked rather than all -inf
+        return [t for t in p.allowed_token_ids or () if 0 <= t < self.vocab_size]
+
     def _prompt_mask(self, req) -> torch.Tensor:
         """[V] bool marking the request's prompt ids, uploaded once per request."""
         prompt_len = req.max_device_len - req.output_len
@@ -258,6 +286,12 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
         max_top_logprobs = min(
             max((p.top_logprobs for p in params if p.logprobs), default=0), self.vocab_size
         )
+        restrict = [(i, self._allowed(p)) for i, p in enumerate(p for p in params if p.logprobs)]
+        restrict = [(i, ids) for i, ids in restrict if ids]
+        restrict_rows = restrict_ids = None
+        if restrict:
+            restrict_rows = make_device_tensor([i for i, _ in restrict], torch.int64, self.device)
+            restrict_ids = _padded_rows([ids for _, ids in restrict], self.vocab_size, self.device)
         if all(is_greedy):
             return BatchSamplingArgs(
                 temperatures=None,
@@ -266,6 +300,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
                 logprob_row_idx=logprob_row_idx,
                 logprob_all_rows=logprob_all_rows,
                 max_top_logprobs=max_top_logprobs,
+                restrict_rows=restrict_rows,
+                restrict_ids=restrict_ids,
             )
 
         MIN_P = MIN_T = 1e-6
@@ -299,6 +335,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             logprob_row_idx=logprob_row_idx,
             logprob_all_rows=logprob_all_rows,
             max_top_logprobs=max_top_logprobs,
+            restrict_rows=restrict_rows,
+            restrict_ids=restrict_ids,
         )
 
     @nvtx_annotate("Sampler")
@@ -337,7 +375,14 @@ def compute_logprobs(
             return None
 
         request_logits = (logits if rows is None else logits.index_select(0, rows)).float()
-        # Reported values are raw model logprobs (pre-temperature log_softmax over logits).
+        if args.restrict_rows is not None:
+            # not in place: on an fp32 batch request_logits can be the engine's logits view
+            masked = request_logits[args.restrict_rows].masked_fill(
+                disallowed_mask(args.restrict_ids, self.vocab_size), float("-inf")
+            )
+            request_logits = request_logits.index_put((args.restrict_rows,), masked)
+        # Reported values are raw model logprobs (pre-temperature log_softmax over logits),
+        # renormalized over allowed_token_ids on the restricted rows.
         request_logprobs = torch.log_softmax(request_logits, dim=-1)
 
         request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device)
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 358ce31355..659a8bab3f 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 import os
 import time
 
@@ -63,6 +64,13 @@ class ForwardInput(NamedTuple):
 ForwardData: TypeAlias = "Tuple[ForwardInput, ForwardOutput]"
 
 
+def logprob_candidates(ids: list[int], logprobs: list[float]) -> tuple[list[int], list[float]]:
+    """A row's top-logprob slice minus what is not a candidate: padding (-1 ids) and the
+    -inf tail of a row restricted to fewer allowed_token_ids than it asked for."""
+    keep = [k for k, (t, lp) in enumerate(zip(ids, logprobs)) if t >= 0 and math.isfinite(lp)]
+    return [int(ids[k]) for k in keep], [logprobs[k] for k in keep]
+
+
 class Scheduler(SchedulerIOMixin):
     def __init__(self, config: SchedulerConfig):
         from freetoken.engine import Engine
@@ -371,8 +379,10 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                     row_chosen_logprob = float(outputs.chosen_logprobs_cpu[i].item())
                     requested_top = req.sampling_params.top_logprobs
                     if requested_top > 0 and outputs.top_ids_cpu is not None:
-                        row_top_ids = [int(t) for t in outputs.top_ids_cpu[i, :requested_top].tolist()]
-                        row_top_logprobs = outputs.top_logprobs_cpu[i, :requested_top].tolist()
+                        row_top_ids, row_top_logprobs = logprob_candidates(
+                            outputs.top_ids_cpu[i, :requested_top].tolist(),
+                            outputs.top_logprobs_cpu[i, :requested_top].tolist(),
+                        )
                     else:
                         row_top_ids = []
                         row_top_logprobs = []
diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 1bd43b9f49..4c4ec53432 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -152,6 +152,7 @@ class CompletionRequest(BaseModel):
     stop_token_ids: list[int] | None = None
     include_stop_str_in_output: bool = False
     skip_special_tokens: bool | None = None
+    allowed_token_ids: list[int] | None = None
     # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no
     # per-request generator (seed); user / metadata / store are opaque to the server.
     seed: int | None = None
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index e2d4d343a4..95ae304fe0 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -444,6 +444,8 @@ def parse_tools(self) -> bool:
 # --max-output-tokens (the Responses adapter passes that through); clamped to the remaining
 # context by the scheduler regardless.
 DEFAULT_MAX_OUTPUT_TOKENS = 32768
+# vLLM's MAX_NUM_ALLOWED_TOKEN_IDS: bounds the per-row mask upload
+MAX_ALLOWED_TOKEN_IDS = 1024
 
 
 def resolve_sampling(
@@ -465,6 +467,7 @@ def resolve_sampling(
     stop_token_ids: list[int] | None = None,
     include_stop_str_in_output: bool = False,
     skip_special_tokens: bool | None = None,
+    allowed_token_ids: list[int] | None = None,
 ) -> SamplingParams:
     """Map a protocol's sampling fields onto the engine's neutral SamplingParams,
     filling unspecified fields from the checkpoint's recommended defaults. Range errors
@@ -518,6 +521,15 @@ def pick(value, key, framework):
     ids = [int(t) for t in (stop_token_ids or [])]
     if any(t < 0 for t in ids):
         raise ValueError("stop_token_ids must be non-negative token ids")
+    allowed = None
+    if allowed_token_ids is not None:
+        if not 1 <= len(allowed_token_ids) <= MAX_ALLOWED_TOKEN_IDS:
+            raise ValueError(
+                f"allowed_token_ids must hold 1 to {MAX_ALLOWED_TOKEN_IDS} ids, got {len(allowed_token_ids)}"
+            )
+        allowed = sorted({int(t) for t in allowed_token_ids})
+        if allowed[0] < 0:
+            raise ValueError("allowed_token_ids must be non-negative token ids")
     params = SamplingParams(
         ignore_eos=ignore_eos,
         max_tokens=default_max_tokens if max_tokens is None else max_tokens,
@@ -533,6 +545,7 @@ def pick(value, key, framework):
         min_tokens=int(min_tokens),
         stop_token_ids=ids,
         include_stop_str_in_output=bool(include_stop_str_in_output),
+        allowed_token_ids=allowed,
     )
     if skip_special_tokens is not None:
         params.skip_special_tokens = bool(skip_special_tokens)
diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py
index c3d22141a7..ffc666816d 100644
--- a/python/freetoken/server/logprobs.py
+++ b/python/freetoken/server/logprobs.py
@@ -6,6 +6,15 @@
 from dataclasses import dataclass
 from typing import Any
 
+# AnyJev reads up to 26 letter labels; its vLLM recipe runs --max-logprobs 32
+MAX_COMPLETION_LOGPROBS = 32
+# vLLM's stand-in for -inf: Starlette's JSON encoder refuses non-finite floats
+LOGPROB_FLOOR = -9999.0
+
+
+def _clamped(logprob: float) -> float:
+    return logprob if logprob >= LOGPROB_FLOOR else LOGPROB_FLOOR  # NaN too
+
 
 def chat_logprobs_error(req: Any) -> str | None:
     if req.top_logprobs is not None and not req.logprobs:
@@ -16,8 +25,8 @@ def chat_logprobs_error(req: Any) -> str | None:
 
 
 def completion_logprobs_error(req: Any) -> str | None:
-    if req.logprobs is not None and not 0 <= req.logprobs <= 5:
-        return "logprobs must be between 0 and 5"
+    if req.logprobs is not None and not 0 <= req.logprobs <= MAX_COMPLETION_LOGPROBS:
+        return f"logprobs must be between 0 and {MAX_COMPLETION_LOGPROBS}"
     if req.echo and req.logprobs is not None:
         return "echo with logprobs is not supported"
     return None
@@ -31,12 +40,12 @@ def untrimmed(entries: list[dict] | None) -> list[dict]:
 def chat_content_entry(entry: dict) -> dict:
     return {
         "token": entry["token"],
-        "logprob": entry["logprob"],
+        "logprob": _clamped(entry["logprob"]),
         "bytes": entry["bytes"],
         "top_logprobs": [
             {
                 "token": candidate["token"],
-                "logprob": candidate["logprob"],
+                "logprob": _clamped(candidate["logprob"]),
                 "bytes": candidate["bytes"],
             }
             for candidate in entry["top"]
@@ -55,12 +64,12 @@ def completions_logprobs(entries: list[dict], text_len: int | None = None) -> di
 
     for entry in entries:
         tokens.append(entry["token"])
-        token_logprobs.append(entry["logprob"])
+        token_logprobs.append(_clamped(entry["logprob"]))
         top: dict[str, float] = {}
         for candidate in entry["top"]:
             # The legacy shape keys alternatives by text, and partial-UTF-8 candidates
             # all decode to U+FFFD: keep the most likely one per key.
-            key, logprob = candidate["token"], candidate["logprob"]
+            key, logprob = candidate["token"], _clamped(candidate["logprob"])
             if key not in top or logprob > top[key]:
                 top[key] = logprob
         top_logprobs.append(top)
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 55d740382a..df1ede267b 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -623,6 +623,14 @@ async def handle_completion(
         sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens)
     except ValueError as exc:
         return create_error_response(str(exc))
+    if sampling.allowed_token_ids:
+        # the tokenizer worker checks again; this one answers before anything is queued
+        vocab = await _vocab_len(state)
+        if vocab is not None and sampling.allowed_token_ids[-1] >= vocab:
+            return create_error_response(
+                f"allowed_token_ids must be below the vocabulary size {vocab}",
+                param="allowed_token_ids",
+            )
 
     prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
     assert isinstance(prompts, list)
@@ -871,6 +879,18 @@ async def pump(slot: int, stream: AsyncIterator[Any]) -> None:
 _FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")
 
 
+async def _vocab_len(state: Any) -> int | None:
+    """``len(tokenizer)`` (added tokens included), or None without a frontend tokenizer."""
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is None:
+        return None
+    try:
+        manager = await asyncio.to_thread(build)
+        return len(manager.tokenizer)
+    except Exception:  # noqa: BLE001 -- the tokenizer worker still range-checks
+        return None
+
+
 async def _fim_markers(state: Any) -> tuple[str, str, str] | None:
     """The model's fill-in-the-middle marker strings when its vocabulary has them
     (Qwen / DeepSeek-Coder style), else None."""
@@ -985,13 +1005,30 @@ def _resolve_sampling(
         stop_token_ids=req.stop_token_ids,
         include_stop_str_in_output=req.include_stop_str_in_output,
         skip_special_tokens=req.skip_special_tokens,
+        allowed_token_ids=req.allowed_token_ids if isinstance(req, CompletionRequest) else None,
     )
     if isinstance(req, CompletionRequest) and req.logprobs is not None:
         sampling_params.logprobs = True
         sampling_params.top_logprobs = req.logprobs
+        if sampling_params.allowed_token_ids:
+            _check_restricted_logprobs(sampling_params)
     return sampling_params
 
 
+def _check_restricted_logprobs(sp: SamplingParams) -> None:
+    """Logprobs restricted to allowed_token_ids are vLLM's processed_logprobs at temperature 0:
+    log_softmax over the allowed ids of the unprocessed logits. Anything vLLM would also fold
+    into them is refused rather than silently left out."""
+    if sp.temperature > 0:
+        raise ValueError("allowed_token_ids with logprobs needs temperature=0")
+    if sp.presence_penalty or sp.frequency_penalty or sp.repetition_penalty != 1.0:
+        raise ValueError("allowed_token_ids with logprobs does not support penalties")
+    if sp.min_tokens > 0:
+        raise ValueError("allowed_token_ids with logprobs does not support min_tokens")
+    if sp.logit_bias:
+        raise ValueError("allowed_token_ids with logprobs does not support logit_bias")
+
+
 def _response_id(prefix: str, req: Any, fallback: Any) -> str:
     """``request_id`` from the client when given (vLLM request_id / SGLang rid), else the
     server's own id."""
@@ -1108,7 +1145,6 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool
 # vLLM /v1/completions fields that change the result and are not implemented here: a set
 # value is a 400 rather than silently dropped by extra="allow".
 _UNIMPLEMENTED_COMPLETION_FIELDS = (
-    "allowed_token_ids",
     "return_tokens_as_token_ids",
     "logprob_token_ids",
     "prompt_logprobs",
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index 31b74bdac2..bf9b63d2fc 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -103,6 +103,10 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
         results: List[UserMsg] = []
         # TODO: batch tokenization
         for msg in msgs:
+            allowed = msg.sampling_params.allowed_token_ids
+            # past the real vocabulary an id indexes the padded logits tail or beyond them
+            if allowed and not (min(allowed) >= 0 and max(allowed) < len(self.tokenizer)):
+                raise ValueError(f"allowed_token_ids must be in [0, {len(self.tokenizer)})")
             prompt = self.render_prompt(msg)
             # A jinja chat template owns every special token (HF's apply_chat_template
             # tokenizes with add_special_tokens=False for the same reason): tokenizers
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index 1829f35f32..769c610c62 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -232,3 +232,26 @@ def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_devic
     assert req.prompt_token_mask is mask
     out = apply_penalties(logits, args.penalties)
     assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0]
+
+
+def test_allowed_token_ids_mask_everything_else_in_the_row_only():
+    logits = torch.arange(2 * V, dtype=torch.float32).view(2, V)
+    # duplicates and the vocab_size padding column are harmless
+    plan = _plan([0, 1], allow_rows=torch.tensor([1]), allow_ids=_ids([[2, 2, 5]]))
+    out = apply_logits_processors(logits, plan, V)
+    assert torch.isfinite(out[1]).tolist() == [i in (2, 5) for i in range(V)]
+    assert out[1, 2].item() == logits[1, 2].item() and out[1, 5].item() == logits[1, 5].item()
+    assert torch.equal(out[0], logits[0])
+
+
+def test_prepare_masks_sampling_to_the_allowed_ids_at_every_step():
+    sampler = Sampler(CPU, V)
+    req = _req([1], [], 4, allowed_token_ids=[2, 4, 99])  # 99: past the logits, dropped
+    other = _req([1], [], 4)
+    logits = torch.zeros((2, V))
+    logits[:, 7] = 10.0
+    logits[:, 4] = 1.0
+    for _ in range(3):
+        args = sampler.prepare(SimpleNamespace(reqs=[req, other]))
+        assert args.plan.allow_rows.tolist() == [0] and args.plan.allow_ids.tolist() == [[2, 4]]
+        assert sampler.sample(logits, args).tolist() == [4, 7]
diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py
index 28e7195b58..03099ab88d 100644
--- a/tests/engine/test_sample.py
+++ b/tests/engine/test_sample.py
@@ -67,3 +67,18 @@ def fake_sample_impl(logits, temperatures, top_k, top_p):
     assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0]  # presence: -3 once, not per count
     assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0]
     assert tokens.tolist() == [2, 1]
+
+
+def test_sampled_rows_draw_only_allowed_ids(monkeypatch):
+    def fake_sample_impl(logits, temperatures, top_k, top_p):
+        probs = torch.softmax(logits / temperatures.unsqueeze(-1), dim=-1)
+        return torch.multinomial(probs, 1, generator=torch.Generator().manual_seed(0)).flatten()
+
+    monkeypatch.setattr(S, "sample_impl", fake_sample_impl)
+    sampler = S.Sampler(torch.device("cpu"), 6)
+    allowed = [1, 4]
+    batch = _batch(SamplingParams(temperature=1.5, allowed_token_ids=allowed), SamplingParams(temperature=0.0))
+    logits = torch.tensor([[3.0, 0.0, 3.0, 3.0, 0.0, 3.0], [0.0, 0.0, 0.0, 5.0, 0.0, 0.0]])
+    for _ in range(50):
+        tokens = sampler.sample(logits, sampler.prepare(batch))
+        assert tokens[0].item() in allowed and tokens[1].item() == 3
diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py
index 02dff731ef..8f0915ed10 100644
--- a/tests/engine/test_sample_logprobs.py
+++ b/tests/engine/test_sample_logprobs.py
@@ -1,9 +1,12 @@
+import math
 from types import SimpleNamespace
 
+import pytest
 import torch
 
 from freetoken.core import SamplingParams
 from freetoken.engine.sample import BatchSamplingArgs, Sampler
+from freetoken.scheduler.scheduler import logprob_candidates
 
 
 def test_compute_logprobs_matches_raw_log_softmax_and_sorted_top() -> None:
@@ -219,3 +222,76 @@ def test_compute_logprobs_builds_row_index_on_host(monkeypatch) -> None:
         for got, ref in zip(result, expected, strict=True):
             assert got.shape == ref.shape and got.dtype == ref.dtype
             assert torch.equal(got.nan_to_num(-7.0), ref.nan_to_num(-7.0))
+
+
+def _restricted_batch(*params: SamplingParams) -> SimpleNamespace:
+    return SimpleNamespace(reqs=[SimpleNamespace(sampling_params=p) for p in params])
+
+
+def test_restricted_logprobs_renormalize_over_the_allowed_ids() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=6)
+    logits = torch.tensor([[0.5, 2.0, -1.0, 3.0, 0.0, 1.0]])
+    allowed = [1, 2, 4]
+    args = sampler.prepare(
+        _restricted_batch(SamplingParams(logprobs=True, top_logprobs=3, allowed_token_ids=allowed))
+    )
+    chosen, top_ids, top_logprobs = sampler.compute_logprobs(logits, torch.tensor([1]), args)
+
+    expected = logits[0, allowed] - torch.logsumexp(logits[0, allowed], dim=-1)
+    assert torch.isclose(chosen[0], expected[0])
+    assert top_ids[0].tolist() == [1, 4, 2]
+    assert torch.allclose(top_logprobs[0], expected[[0, 2, 1]])
+    assert torch.exp(top_logprobs[0]).sum().item() == pytest.approx(1.0)
+
+
+def test_restricted_and_raw_rows_share_a_batch() -> None:
+    # K=20 over 4 labels next to a raw K=20 row: topk still runs over the whole row, and
+    # the restricted row's tail past its labels is -inf (dropped by the scheduler).
+    vocab = 64
+    sampler = Sampler(torch.device("cpu"), vocab_size=vocab)
+    logits = torch.randn(3, vocab, generator=torch.Generator().manual_seed(1))
+    labels = [32, 33, 34, 35]
+    raw = SamplingParams(logprobs=True, top_logprobs=20)
+    restricted = SamplingParams(logprobs=True, top_logprobs=20, allowed_token_ids=labels)
+    unasked = SamplingParams()
+    sampled = torch.tensor([5, 33, 7])
+
+    snapshot = logits.clone()
+    mixed = sampler.compute_logprobs(
+        logits, sampled, sampler.prepare(_restricted_batch(raw, restricted, unasked))
+    )
+    plain = sampler.compute_logprobs(
+        logits, sampled, sampler.prepare(_restricted_batch(raw, raw, unasked))
+    )
+    assert torch.equal(logits, snapshot)  # the engine's logits are not masked in place
+
+    chosen, top_ids, top_logprobs = mixed
+    assert torch.equal(chosen[0], plain[0][0])
+    assert torch.equal(top_ids[0], plain[1][0]) and torch.equal(top_logprobs[0], plain[2][0])
+
+    finite = torch.isfinite(top_logprobs[1])
+    assert finite.sum().item() == 4
+    assert sorted(top_ids[1][finite].tolist()) == labels
+    expected = torch.log_softmax(logits[1, labels], dim=-1)
+    assert torch.isclose(chosen[1], expected[1])
+
+    ids, values = logprob_candidates(top_ids[1].tolist(), top_logprobs[1].tolist())
+    assert sorted(ids) == labels and all(map(math.isfinite, values))
+
+
+def test_restricted_logprobs_on_an_all_rows_fp32_batch_leave_the_logits_alone() -> None:
+    sampler = Sampler(torch.device("cpu"), vocab_size=5)
+    logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    args = sampler.prepare(
+        _restricted_batch(SamplingParams(logprobs=True, top_logprobs=2, allowed_token_ids=[0, 0, 1]))
+    )
+    assert args.logprob_all_rows
+    chosen, top_ids, _ = sampler.compute_logprobs(logits, torch.tensor([1]), args)
+    assert top_ids[0].tolist() == [1, 0]
+    assert torch.isclose(chosen[0], torch.log_softmax(torch.tensor([1.0, 2.0]), dim=-1)[1])
+    assert logits.tolist() == [[1.0, 2.0, 3.0, 4.0, 5.0]]
+
+
+def test_logprob_candidates_drop_padding_and_non_finite_entries() -> None:
+    inf = float("inf")
+    assert logprob_candidates([3, -1, 4, 5], [-0.1, -inf, -inf, float("nan")]) == ([3], [-0.1])
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 039a89fa1d..bc45f21e79 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -549,7 +549,7 @@ def test_logprobs_validation_errors() -> None:
 
     completion_out_of_range = run(
         handle_completion(
-            CompletionRequest(model="client-model", prompt="hello", logprobs=7),
+            CompletionRequest(model="client-model", prompt="hello", logprobs=33),
             None,
             FakeState([]),
             {},
@@ -871,7 +871,6 @@ def _one_token_state() -> FakeState:
 # Result-changing vLLM fields this server does not implement: a set value must be a 400,
 # never a 200 computed as if the field were absent.
 UNIMPLEMENTED_COMPLETION_FIELDS = [
-    ("allowed_token_ids", [32, 33]),
     ("return_tokens_as_token_ids", True),
     ("logprob_token_ids", [32]),
     ("prompt_logprobs", 0),
@@ -934,3 +933,84 @@ def test_completion_passes_the_context_length_code_through() -> None:
 
     events = parse_sse(run(_collect(stream_completion_chunks(42, _completion(stream=True), FakeState([_too_long()])))))
     assert events[0]["error"]["code"] == "context_length_exceeded"
+
+
+class VocabTokenizer(QwenTokenizer):
+    def __len__(self) -> int:
+        return 248077
+
+
+def _status(resp) -> int:
+    return getattr(resp, "status_code", 200)
+
+
+def test_completion_logprobs_cap_is_32() -> None:
+    assert _status(run(handle_completion(_completion(logprobs=32), None, _one_token_state(), {}))) == 200
+    assert _status(run(handle_completion(_completion(logprobs=33), None, _one_token_state(), {}))) == 400
+
+
+def test_completion_allowed_token_ids_reach_the_engine_sorted_and_unique() -> None:
+    state = _one_token_state()
+    req = _completion(allowed_token_ids=[33, 32, 33], logprobs=2, temperature=0)
+    assert _status(run(handle_completion(req, None, state, {"temperature": 0.6}))) == 200
+    assert state.sent.sampling_params.allowed_token_ids == [32, 33]
+
+
+@pytest.mark.parametrize(
+    ("ids", "fragment"),
+    [([], "1 to 1024"), (list(range(1025)), "1 to 1024"), ([-1, 3], "non-negative"), ([248077], "vocabulary")],
+)
+def test_completion_allowed_token_ids_out_of_range(ids, fragment) -> None:
+    state = FakeState([], tokenizer=VocabTokenizer())
+    resp = run(handle_completion(_completion(allowed_token_ids=ids), None, state, {}))
+    assert resp.status_code == 400 and fragment in json.loads(resp.body)["error"]["message"]
+    assert state.sent is None
+
+
+def test_completion_allowed_token_ids_at_the_vocab_edge_pass() -> None:
+    state = FakeState([reply("A", finished=True)], tokenizer=VocabTokenizer())
+    assert _status(run(handle_completion(_completion(allowed_token_ids=[0, 248076]), None, state, {}))) == 200
+
+
+@pytest.mark.parametrize(
+    ("extra", "model_sampling"),
+    [
+        ({"temperature": 0.7}, {}),
+        ({}, {"temperature": 0.6}),  # the checkpoint default counts, not just the request
+        ({"temperature": 0, "presence_penalty": 0.5}, {}),
+        ({"temperature": 0, "frequency_penalty": 0.5}, {}),
+        ({"temperature": 0, "repetition_penalty": 1.1}, {}),
+        ({"temperature": 0}, {"repetition_penalty": 1.05}),
+        ({"temperature": 0, "min_tokens": 1}, {}),
+        ({"temperature": 0, "logit_bias": {"32": 1.0}}, {}),
+    ],
+)
+def test_restricted_logprobs_refuse_what_they_would_not_fold_in(extra, model_sampling) -> None:
+    req = _completion(allowed_token_ids=[32, 33], logprobs=2, **extra)
+    assert _status(run(handle_completion(req, None, _one_token_state(), model_sampling))) == 400
+
+
+def test_allowed_token_ids_without_logprobs_may_sample() -> None:
+    req = _completion(allowed_token_ids=[32, 33], temperature=0.7, presence_penalty=0.5)
+    assert _status(run(handle_completion(req, None, _one_token_state(), {}))) == 200
+
+
+def test_non_finite_logprobs_are_clamped_to_the_vllm_floor() -> None:
+    inf = float("inf")
+    sampled = entry(32, "A", -inf)
+    sampled["top"][1]["logprob"] = float("nan")
+    state = FakeState([reply("A", finished=True, logprobs=[sampled])])
+    result = run(handle_completion(_completion(logprobs=2), None, state, {}))
+    lp = result["choices"][0]["logprobs"]
+    assert lp["token_logprobs"] == [-9999.0] and lp["top_logprobs"] == [{"A": -9999.0, "x": -9999.0}]
+    json.dumps(result, allow_nan=False)
+
+    chat = run(
+        handle_chat_completion(
+            chat_request(logprobs=True, top_logprobs=2), None,
+            FakeState([reply("A", finished=True, logprobs=[sampled])]), {},
+        )
+    )
+    content = chat["choices"][0]["logprobs"]["content"][0]
+    assert content["logprob"] == -9999.0
+    json.dumps(chat, allow_nan=False)
diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py
index b13db921be..61f521c7c0 100644
--- a/tests/server/test_sampling_validation.py
+++ b/tests/server/test_sampling_validation.py
@@ -54,3 +54,14 @@ def test_a_positive_top_k_is_kept():
 def test_the_other_sampling_ranges_are_still_rejected(field, value):
     with pytest.raises(ValueError, match=field):
         _resolve_with(**{field: value})
+
+
+def test_allowed_token_ids_are_sorted_and_deduplicated():
+    assert _resolve_with(allowed_token_ids=[9, 3, 9, 0]).allowed_token_ids == [0, 3, 9]
+    assert _resolve_with().allowed_token_ids is None
+
+
+@pytest.mark.parametrize("ids", [[], list(range(1025)), [4, -2]])
+def test_allowed_token_ids_size_and_sign_are_checked(ids):
+    with pytest.raises(ValueError, match="allowed_token_ids"):
+        _resolve_with(allowed_token_ids=ids)
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 9e441886b2..545d92eb4f 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -848,3 +848,24 @@ def test_mapped_developer_messages_are_merged_into_one_leading_system_message():
     image = [{"type": "text", "text": "B"}, {"type": "image", "freetoken_ref": {"kind": "url", "data": "x"}}]
     with pytest.raises(ValueError, match="System message cannot contain images"):
         _map_developer_role([{"role": "developer", "content": "A"}, {"role": "system", "content": image}], template)
+
+
+class RawTokenizer:
+    """A raw-prompt tokenizer with a 10-id vocabulary."""
+
+    def __len__(self) -> int:
+        return 10
+
+    def encode(self, prompt, return_tensors=None, add_special_tokens=True):
+        return torch.tensor([[1, 2, 3]], dtype=torch.long)
+
+
+@pytest.mark.parametrize(("allowed", "ok"), [([0, 9], True), ([3, 10], False)])
+def test_tokenize_range_checks_allowed_token_ids(allowed, ok):
+    manager = TokenizeManager(RawTokenizer())
+    msg = TokenizeMsg(uid=1, text="hi", sampling_params=SamplingParams(allowed_token_ids=allowed))
+    if ok:
+        assert manager.tokenize([msg])[0].input_ids.tolist() == [1, 2, 3]
+    else:
+        with pytest.raises(ValueError, match="allowed_token_ids"):
+            manager.tokenize([msg])

From c501d830ea93a25e0a3849857a117d113752e25e Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:23:36 -0400
Subject: [PATCH 230/249] feat(server): honour return_tokens_as_token_ids on
 completions

tokens and top_logprobs keys become "token_id:N", vLLM's id-keyed shape, which
also keeps candidates that decode to the same text from collapsing.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/openai_api.md                    |  2 ++
 python/freetoken/server/api_models.py |  1 +
 python/freetoken/server/logprobs.py   | 15 +++++++++++----
 python/freetoken/server/openai_api.py | 15 +++++++++++----
 tests/server/test_logprobs_api.py     | 23 ++++++++++++++++++++++-
 5 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/docs/openai_api.md b/docs/openai_api.md
index 02ad41832d..210c67c548 100644
--- a/docs/openai_api.md
+++ b/docs/openai_api.md
@@ -65,6 +65,8 @@ parameters as above; plus:
 | `logprobs` | honoured, 0..32: raw pre-temperature logprobs of the sampled token and its top alternatives; -inf is sent as -9999 | honoured (`--max-logprobs`, default 20) | honoured | `n_probs` |
 | `allowed_token_ids` | honoured, 1..1024 ids below `len(tokenizer)`: sampling is masked to them at every position. With `logprobs` the values are renormalized over the set (vLLM's `processed_logprobs` at temperature 0), so `temperature > 0`, penalties, `min_tokens` and `logit_bias` are a 400 there; entries past the set are left out | honoured (raw logprobs unless `--logprobs-mode processed_logprobs`) | no | no |
 
+| `return_tokens_as_token_ids` | honoured: `tokens` and the `top_logprobs` keys are `token_id:N`, so ids that decode to the same text stay apart | honoured | n/a | no |
+
 Set values of the vLLM fields this server does not implement (`logprob_token_ids`,
 `prompt_logprobs`, `guided_*`, `structured_outputs`, `bad_words`, ...) are a 400 rather than
 ignored; `null` / `false` / empty pass. Empty prompts and more than 256 prompts x `n` are a 400.
diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py
index 4c4ec53432..344b274829 100644
--- a/python/freetoken/server/api_models.py
+++ b/python/freetoken/server/api_models.py
@@ -153,6 +153,7 @@ class CompletionRequest(BaseModel):
     include_stop_str_in_output: bool = False
     skip_special_tokens: bool | None = None
     allowed_token_ids: list[int] | None = None
+    return_tokens_as_token_ids: bool = False
     # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no
     # per-request generator (seed); user / metadata / store are opaque to the server.
     seed: int | None = None
diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py
index ffc666816d..1dce7cc413 100644
--- a/python/freetoken/server/logprobs.py
+++ b/python/freetoken/server/logprobs.py
@@ -53,23 +53,30 @@ def chat_content_entry(entry: dict) -> dict:
     }
 
 
-def completions_logprobs(entries: list[dict], text_len: int | None = None) -> dict:
+def completions_logprobs(
+    entries: list[dict], text_len: int | None = None, as_token_ids: bool = False
+) -> dict:
     """Each entry's detokenizer-stamped ``text_offset`` is where its token starts in the
     generated text (the pieces of a split multi-byte character share its start). Given the
     finished text's ``text_len``, offsets are clipped to it, since a matched stop string
-    trims tokens off the end."""
+    trims tokens off the end. ``as_token_ids`` (vLLM return_tokens_as_token_ids) names
+    every token ``token_id:N`` instead of by its text."""
+
+    def name(item: dict) -> str:
+        return f"token_id:{item['token_id']}" if as_token_ids else item["token"]
+
     tokens: list[str] = []
     token_logprobs: list[float] = []
     top_logprobs: list[dict[str, float]] = []
 
     for entry in entries:
-        tokens.append(entry["token"])
+        tokens.append(name(entry))
         token_logprobs.append(_clamped(entry["logprob"]))
         top: dict[str, float] = {}
         for candidate in entry["top"]:
             # The legacy shape keys alternatives by text, and partial-UTF-8 candidates
             # all decode to U+FFFD: keep the most likely one per key.
-            key, logprob = candidate["token"], _clamped(candidate["logprob"])
+            key, logprob = name(candidate), _clamped(candidate["logprob"])
             if key not in top or logprob > top[key]:
                 top[key] = logprob
         top_logprobs.append(top)
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index df1ede267b..b1e52ebc8f 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -697,7 +697,9 @@ async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse:
             "text": echo_texts[index // req.n] + text,
             "finish_reason": finish_reason,
             "logprobs": (
-                completions_logprobs(entries, len(text)) if req.logprobs is not None else None
+                completions_logprobs(entries, len(text), req.return_tokens_as_token_ids)
+                if req.logprobs is not None
+                else None
             ),
             "_usage": (prompt_tokens, completion_tokens, cached_tokens),
         }
@@ -783,7 +785,11 @@ async def stream_completion_chunks(
                         "finish_reason": None,
                         # Mid-stream, a stamp past sent_len is held-back text still to come.
                         "logprobs": (
-                            completions_logprobs(pending, sent_len if ack.finished else None)
+                            completions_logprobs(
+                                pending,
+                                sent_len if ack.finished else None,
+                                req.return_tokens_as_token_ids,
+                            )
                             if pending
                             else None
                         ),
@@ -805,7 +811,9 @@ async def stream_completion_chunks(
                 "index": index,
                 "finish_reason": finish_reason,
                 "logprobs": (
-                    completions_logprobs(pending, sent_len) if pending else None
+                    completions_logprobs(pending, sent_len, req.return_tokens_as_token_ids)
+                    if pending
+                    else None
                 ),
             }],
         )
@@ -1145,7 +1153,6 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool
 # vLLM /v1/completions fields that change the result and are not implemented here: a set
 # value is a 400 rather than silently dropped by extra="allow".
 _UNIMPLEMENTED_COMPLETION_FIELDS = (
-    "return_tokens_as_token_ids",
     "logprob_token_ids",
     "prompt_logprobs",
     "guided_json",
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index bc45f21e79..5e6f299a0e 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -871,7 +871,6 @@ def _one_token_state() -> FakeState:
 # Result-changing vLLM fields this server does not implement: a set value must be a 400,
 # never a 200 computed as if the field were absent.
 UNIMPLEMENTED_COMPLETION_FIELDS = [
-    ("return_tokens_as_token_ids", True),
     ("logprob_token_ids", [32]),
     ("prompt_logprobs", 0),
     ("guided_json", {"type": "object"}),
@@ -1014,3 +1013,25 @@ def test_non_finite_logprobs_are_clamped_to_the_vllm_floor() -> None:
     content = chat["choices"][0]["logprobs"]["content"][0]
     assert content["logprob"] == -9999.0
     json.dumps(chat, allow_nan=False)
+
+
+def test_completion_return_tokens_as_token_ids_keys_by_id() -> None:
+    # two ids that decode to the same text stay two entries
+    a = entry(32, "\ufffd", -0.1)
+    a["top"][1].update(token_id=900, token="\ufffd")
+    req = _completion(logprobs=2, return_tokens_as_token_ids=True)
+    result = run(handle_completion(req, None, FakeState([reply("A", finished=True, logprobs=[a])]), {}))
+    expected = {
+        "tokens": ["token_id:32"],
+        "token_logprobs": [-0.1],
+        "top_logprobs": [{"token_id:32": -0.1, "token_id:900": -1.1}],
+        "text_offset": [0],
+    }
+    assert result["choices"][0]["logprobs"] == expected
+
+    stream_req = _completion(logprobs=2, return_tokens_as_token_ids=True, stream=True)
+    events = parse_sse(
+        run(_collect(stream_completion_chunks(42, stream_req, FakeState([reply("A", finished=True, logprobs=[a])]))))
+    )
+    streamed = [e["choices"][0]["logprobs"] for e in events if isinstance(e, dict) and e["choices"][0]["logprobs"]]
+    assert streamed == [expected]

From 7ade750018c7989e7c8097d27fffeabe3e69c64f Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:25:20 -0400
Subject: [PATCH 231/249] feat(server): accept token-id prompts on completions

prompt as list[int] or list[list[int]] goes to the engine as is, one request
per prompt x n in prompt-major order. Ids are range-checked at the API when
a frontend tokenizer is up and always in tokenize(), so an id past the
vocabulary never reaches the embedding. echo or suffix with one is a 400.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/openai_api.md                     |  3 +-
 python/freetoken/message/tokenizer.py  |  2 +
 python/freetoken/server/openai_api.py  | 44 +++++++++++++++++-----
 python/freetoken/tokenizer/tokenize.py | 13 ++++++-
 tests/server/test_logprobs_api.py      | 51 ++++++++++++++++++++++++++
 tests/server/test_message_wire.py      |  8 ++++
 tests/server/test_openai_api.py        |  6 +--
 tests/tokenizer/test_tokenize.py       | 12 ++++++
 8 files changed, 124 insertions(+), 15 deletions(-)

diff --git a/docs/openai_api.md b/docs/openai_api.md
index 210c67c548..1079bf8156 100644
--- a/docs/openai_api.md
+++ b/docs/openai_api.md
@@ -54,7 +54,8 @@ effect; "400" means the server refuses the request with a clear message.
 
 ## Legacy completions: request
 
-`prompt` as a string or a list of strings (token-id prompts are refused); the same sampling
+`prompt` as a string, a list of strings, a list of token ids or a list of token-id lists
+(each id in `[0, len(tokenizer))`; not with `echo` or `suffix`); the same sampling
 parameters as above; plus:
 
 | parameter | FreeToken | vLLM | SGLang | llama.cpp |
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index b5c24a6e10..30efdac0d7 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -90,6 +90,8 @@ class TokenizeMsg(BaseTokenizerMsg):
     tools: List[Dict[str, Any]] | None = None
     images: List[bytes] | None = None
     inline_system_policy: str | None = None
+    # A token-id prompt (OpenAI / vLLM completions): used as is, ``text`` is ignored.
+    input_ids: List[int] | None = None
 
 
 @dataclass
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b1e52ebc8f..5c148b0b36 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -611,7 +611,8 @@ async def handle_completion(
         return create_error_response(unsupported)
     if not 1 <= req.n <= MAX_N:
         return create_error_response(f"n must be between 1 and {MAX_N}", param="n")
-    num_prompts = 1 if isinstance(req.prompt, str) else len(req.prompt)
+    prompts = _prompt_list(req.prompt)
+    num_prompts = len(prompts)
     if num_prompts * req.n > MAX_PROMPTS:
         return create_error_response(
             f"prompts x n must be at most {MAX_PROMPTS}, got {num_prompts * req.n}", param="prompt"
@@ -623,17 +624,21 @@ async def handle_completion(
         sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens)
     except ValueError as exc:
         return create_error_response(str(exc))
-    if sampling.allowed_token_ids:
-        # the tokenizer worker checks again; this one answers before anything is queued
-        vocab = await _vocab_len(state)
-        if vocab is not None and sampling.allowed_token_ids[-1] >= vocab:
+    token_prompts = _is_token_prompt(req.prompt)
+    if sampling.allowed_token_ids or token_prompts:
+        # The tokenizer worker checks again against its own tokenizer; this one answers
+        # before anything is queued (and keeps ids msgpack can carry without one).
+        bound = await _vocab_len(state) or _MAX_WIRE_TOKEN_ID
+        if sampling.allowed_token_ids and sampling.allowed_token_ids[-1] >= bound:
             return create_error_response(
-                f"allowed_token_ids must be below the vocabulary size {vocab}",
+                f"allowed_token_ids must be below the vocabulary size {bound}",
                 param="allowed_token_ids",
             )
+        if token_prompts and not all(0 <= t < bound for ids in prompts for t in ids):
+            return create_error_response(
+                f"prompt token ids must be in [0, {bound})", param="prompt"
+            )
 
-    prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
-    assert isinstance(prompts, list)
     if req.suffix is not None:
         markers = await _fim_markers(state)
         if markers is None:
@@ -653,7 +658,12 @@ async def handle_completion(
     for prompt in prompts:
         for _ in range(req.n):
             uid = state.new_user()
-            await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling))
+            msg = (
+                TokenizeMsg(uid=uid, text="", input_ids=prompt, sampling_params=sampling)
+                if token_prompts
+                else TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling)
+            )
+            await state.send_one(msg)
             uids.append(uid)
 
     if req.stream:
@@ -887,6 +897,10 @@ async def pump(slot: int, stream: AsyncIterator[Any]) -> None:
 _FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")
 
 
+# Without a frontend tokenizer: the int32 the engine holds ids in (msgpack also needs a bound)
+_MAX_WIRE_TOKEN_ID = 2**31 - 1
+
+
 async def _vocab_len(state: Any) -> int | None:
     """``len(tokenizer)`` (added tokens included), or None without a frontend tokenizer."""
     build = getattr(state, "frontend_tokenizer", None)
@@ -1192,7 +1206,10 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
     if _empty_prompt(req):
         return "prompt must not be empty"
     if _is_token_prompt(req.prompt):
-        return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead"
+        if req.echo:
+            return "echo is not supported with token-id prompts"
+        if req.suffix is not None:
+            return "suffix is not supported with token-id prompts"
     logprobs_error = completion_logprobs_error(req)
     if logprobs_error is not None:
         return logprobs_error
@@ -1205,6 +1222,13 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
     return None
 
 
+def _prompt_list(prompt: Any) -> list:
+    """One entry per prompt: a string or a token-id list."""
+    if isinstance(prompt, str) or (prompt and all(isinstance(t, int) for t in prompt)):
+        return [prompt]
+    return list(prompt)
+
+
 def _is_token_prompt(prompt: Any) -> bool:
     return (
         isinstance(prompt, list)
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index bf9b63d2fc..fdb65d8c53 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -87,6 +87,10 @@ def _map_developer_role(
     return [{"role": "system", "content": join_system_contents(system)}] + rest
 
 
+def _ids_below(ids: List[int], bound: int) -> bool:
+    return min(ids) >= 0 and max(ids) < bound
+
+
 class TokenizeManager:
     def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor | None = None) -> None:
         self.tokenizer = tokenizer
@@ -105,8 +109,15 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
         for msg in msgs:
             allowed = msg.sampling_params.allowed_token_ids
             # past the real vocabulary an id indexes the padded logits tail or beyond them
-            if allowed and not (min(allowed) >= 0 and max(allowed) < len(self.tokenizer)):
+            if allowed and not _ids_below(allowed, len(self.tokenizer)):
                 raise ValueError(f"allowed_token_ids must be in [0, {len(self.tokenizer)})")
+            if msg.input_ids is not None:
+                # an id past the embedding rows would fault the engine, not just this request
+                if msg.input_ids and not _ids_below(msg.input_ids, len(self.tokenizer)):
+                    raise ValueError(f"prompt token ids must be in [0, {len(self.tokenizer)})")
+                ids = torch.tensor(msg.input_ids, dtype=torch.int32)
+                results.append(UserMsg(uid=msg.uid, input_ids=ids, sampling_params=msg.sampling_params))
+                continue
             prompt = self.render_prompt(msg)
             # A jinja chat template owns every special token (HF's apply_chat_template
             # tokenizes with add_special_tokens=False for the same reason): tokenizers
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 5e6f299a0e..e0f9fd2081 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -1035,3 +1035,54 @@ def test_completion_return_tokens_as_token_ids_keys_by_id() -> None:
     )
     streamed = [e["choices"][0]["logprobs"] for e in events if isinstance(e, dict) and e["choices"][0]["logprobs"]]
     assert streamed == [expected]
+
+
+class ManyUsersState(FakeState):
+    """Hands out a fresh uid per queued request and replies with the uid as the text."""
+
+    def __init__(self, tokenizer=None) -> None:
+        super().__init__([], tokenizer=tokenizer)
+        self.msgs = []
+
+    def new_user(self) -> int:
+        return 100 + len(self.msgs)
+
+    async def send_one(self, msg):
+        self.msgs.append(msg)
+
+    async def wait_for_ack(self, uid: int):
+        yield reply(str(uid), finished=True)
+
+
+def test_completion_token_id_prompts_fan_out_prompt_major() -> None:
+    state = ManyUsersState(tokenizer=VocabTokenizer())
+    req = _completion(prompt=[[1, 2], [3, 4, 5]], n=2)
+    result = run(handle_completion(req, None, state, {}))
+    assert [m.input_ids for m in state.msgs] == [[1, 2], [1, 2], [3, 4, 5], [3, 4, 5]]
+    assert all(m.text == "" for m in state.msgs)
+    choices = result["choices"]
+    assert [c["index"] for c in choices] == [0, 1, 2, 3]  # prompt i, sample k -> i*n + k
+    assert [c["text"] for c in choices] == ["100", "101", "102", "103"]
+
+
+def test_completion_flat_token_id_prompt_is_one_prompt() -> None:
+    state = ManyUsersState()
+    result = run(handle_completion(_completion(prompt=list(range(1, 300))), None, state, {}))
+    assert len(result["choices"]) == 1 and state.msgs[0].input_ids == list(range(1, 300))
+
+
+@pytest.mark.parametrize(
+    ("prompt", "tokenizer"),
+    [([[1], [248077]], VocabTokenizer()), ([5, -1], None), ([[2**40]], None)],
+)
+def test_completion_token_id_prompts_are_range_checked_before_queueing(prompt, tokenizer) -> None:
+    state = ManyUsersState(tokenizer=tokenizer)
+    resp = run(handle_completion(_completion(prompt=prompt), None, state, {}))
+    assert resp.status_code == 400 and "prompt token ids" in json.loads(resp.body)["error"]["message"]
+    assert state.msgs == []
+
+
+@pytest.mark.parametrize("extra", [{"echo": True}, {"suffix": "tail"}])
+def test_completion_token_id_prompts_refuse_echo_and_suffix(extra) -> None:
+    resp = run(handle_completion(_completion(prompt=[[1, 2]], **extra), None, ManyUsersState(), {}))
+    assert resp.status_code == 400
diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py
index fdc095921d..25b8eb1e67 100644
--- a/tests/server/test_message_wire.py
+++ b/tests/server/test_message_wire.py
@@ -40,6 +40,14 @@ def test_inline_system_policy_tokenizer_roundtrip():
     assert BaseTokenizerMsg.decoder(legacy).inline_system_policy is None
 
 
+def test_token_id_prompt_and_allowed_ids_survive_the_wire():
+    msg = TokenizeMsg(uid=1, text="", input_ids=[248045, 32, 0],
+                      sampling_params=SamplingParams(allowed_token_ids=[32, 33]))
+    decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg))
+    assert decoded.input_ids == [248045, 32, 0]
+    assert decoded.sampling_params.allowed_token_ids == [32, 33]
+
+
 def test_cache_rebuild_msg_roundtrip():
     msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle")
     out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg))
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index fdd2b11159..ff84745955 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -359,11 +359,11 @@ def test_tool_choice_none_keeps_tool_tags_as_content():
     assert choice["message"] == {"role": "assistant", "content": output}
 
 
-def test_completion_rejects_token_id_prompts():
+def test_completion_rejects_echo_with_token_id_prompts():
     state = FakeState([])
     response = run(
         handle_completion(
-            CompletionRequest(model="client-model", prompt=[1, 2, 3]),
+            CompletionRequest(model="client-model", prompt=[1, 2, 3], echo=True),
             request=None,
             state=state,
             model_sampling={},
@@ -373,7 +373,7 @@ def test_completion_rejects_token_id_prompts():
     assert response.status_code == 400
     body = json.loads(response.body)
     assert body["error"]["type"] == "invalid_request_error"
-    assert "token-id prompt" in body["error"]["message"]
+    assert "token-id prompts" in body["error"]["message"]
 
 
 def test_completion_accepts_text_prompt():
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 545d92eb4f..6fade1641b 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -869,3 +869,15 @@ def test_tokenize_range_checks_allowed_token_ids(allowed, ok):
     else:
         with pytest.raises(ValueError, match="allowed_token_ids"):
             manager.tokenize([msg])
+
+
+@pytest.mark.parametrize(("ids", "ok"), [([0, 9, 4], True), ([3, 10], False), ([-1], False)])
+def test_tokenize_takes_token_id_prompts_as_is_within_the_vocabulary(ids, ok):
+    manager = TokenizeManager(RawTokenizer())
+    msg = TokenizeMsg(uid=1, text="", input_ids=ids, sampling_params=SamplingParams())
+    if ok:
+        out = manager.tokenize([msg])[0].input_ids
+        assert out.tolist() == ids and out.dtype == torch.int32
+    else:
+        with pytest.raises(ValueError, match="prompt token ids"):
+            manager.tokenize([msg])

From abc6dec1a6e9b5c7de007a83cd15c327dfab748c Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:30:40 -0400
Subject: [PATCH 232/249] docs(server): drop the stale logprobs 400 note from
 the chat table

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/openai_api.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/openai_api.md b/docs/openai_api.md
index 1079bf8156..82b9daf885 100644
--- a/docs/openai_api.md
+++ b/docs/openai_api.md
@@ -40,7 +40,7 @@ effect; "400" means the server refuses the request with a clear message.
 | `skip_special_tokens` | honoured; default off here, the reasoning and tool parsers read the tags | honoured (default on) | honoured (default on) | n/a |
 | `n` | honoured, 1..16, one generation per choice (the prefix cache shares the prompt) | honoured | honoured | `n_cmpl` |
 | `stream`, `stream_options.include_usage` | honoured | honoured | honoured | honoured |
-| `logprobs`, `top_logprobs` | 400 on this route (see #224); `/v1/completions` `logprobs` 400 | honoured | honoured | `n_probs` |
+| `logprobs`, `top_logprobs` | honoured (`top_logprobs` 0..20; message content only, under the supported reasoning and tool parsers) | honoured | honoured | `n_probs` |
 | `tools`, `tool_choice` (`auto` / `none` / `required` / named), `parallel_tool_calls` | honoured | honoured | honoured | honoured |
 | `reasoning_effort`, `thinking`, `chat_template_kwargs` | honoured | `chat_template_kwargs` | honoured | honoured |
 | `continue_final_message` | honoured (the last assistant message is continued, no generation prompt) | honoured | honoured | no |

From 2bfdae17911e63cc55fe60c5ba30e3663f1f75ba Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:33:01 -0400
Subject: [PATCH 233/249] fix(server): reject unimplemented vllm fields on chat
 completions

The completions deny-list now covers /v1/chat/completions too, plus
allowed_token_ids and return_tokens_as_token_ids, which only completions
implements: a set value is a 400 instead of a 200 that ignored it. Before
this, the new chat test fails for all 10 fields with 200 != 400.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/openai_api.md                    |  3 ++-
 python/freetoken/server/openai_api.py | 24 +++++++++++++++++++++---
 tests/server/test_logprobs_api.py     | 24 ++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/docs/openai_api.md b/docs/openai_api.md
index 82b9daf885..ec795b0823 100644
--- a/docs/openai_api.md
+++ b/docs/openai_api.md
@@ -49,7 +49,8 @@ effect; "400" means the server refuses the request with a clear message.
 | `user`, `metadata`, `store`, `service_tier` | accepted, ignored | accepted, ignored | accepted, ignored | accepted, ignored |
 | `request_id` | honoured: becomes the response id | honoured | `rid` | `id_slot` (different meaning) |
 | `function_call` (legacy) | 400 | 400 | 400 | no |
-| `best_of`, `use_beam_search`, `length_penalty`, `prompt_logprobs`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly |
+| `allowed_token_ids`, `return_tokens_as_token_ids`, `logprob_token_ids`, `prompt_logprobs`, `guided_*`, `structured_outputs`, `bad_words` | 400 when set (`null` / `false` / empty pass) | honoured | partly | partly |
+| `best_of`, `use_beam_search`, `length_penalty`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly |
 | `ignore_eos` | honoured (benchmarking) | honoured | honoured | honoured |
 
 ## Legacy completions: request
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 5c148b0b36..816fdfd255 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -297,6 +297,9 @@ async def handle_chat_completion(
 ):
     if req.function_call is not None:
         return create_error_response("function_call is not supported; use tools/tool_choice instead")
+    unimplemented = _unimplemented_field(req, _UNIMPLEMENTED_CHAT_FIELDS)
+    if unimplemented is not None:
+        return create_error_response(unimplemented)
     logprobs_error = chat_logprobs_error(req)
     if logprobs_error is not None:
         return create_error_response(logprobs_error, param="top_logprobs")
@@ -1178,6 +1181,21 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool
 )
 
 
+# chat has neither of the completions-only ones
+_UNIMPLEMENTED_CHAT_FIELDS = (
+    "allowed_token_ids",
+    "return_tokens_as_token_ids",
+    *_UNIMPLEMENTED_COMPLETION_FIELDS,
+)
+
+
+def _unimplemented_field(req: Any, names: tuple[str, ...]) -> str | None:
+    for name in names:
+        if _is_set(getattr(req, name, None)):
+            return f"{name} is not supported by this server"
+    return None
+
+
 def _is_set(value: Any) -> bool:
     # not truthiness: prompt_logprobs=0 is a real request, and 0 == False
     if value is None or value is False:
@@ -1200,9 +1218,9 @@ def _empty_prompt(req: CompletionRequest) -> bool:
 
 
 def _completion_unsupported_reason(req: CompletionRequest) -> str | None:
-    for name in _UNIMPLEMENTED_COMPLETION_FIELDS:
-        if _is_set(getattr(req, name, None)):
-            return f"{name} is not supported by this server"
+    unimplemented = _unimplemented_field(req, _UNIMPLEMENTED_COMPLETION_FIELDS)
+    if unimplemented is not None:
+        return unimplemented
     if _empty_prompt(req):
         return "prompt must not be empty"
     if _is_token_prompt(req.prompt):
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index e0f9fd2081..38a1e135dc 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -1086,3 +1086,27 @@ def test_completion_token_id_prompts_are_range_checked_before_queueing(prompt, t
 def test_completion_token_id_prompts_refuse_echo_and_suffix(extra) -> None:
     resp = run(handle_completion(_completion(prompt=[[1, 2]], **extra), None, ManyUsersState(), {}))
     assert resp.status_code == 400
+
+
+# Chat implements none of them, allowed_token_ids and return_tokens_as_token_ids included.
+CHAT_UNIMPLEMENTED_FIELDS = [
+    ("allowed_token_ids", [32, 33]),
+    ("return_tokens_as_token_ids", True),
+    *UNIMPLEMENTED_COMPLETION_FIELDS,
+]
+
+
+@pytest.mark.parametrize(("field", "value"), CHAT_UNIMPLEMENTED_FIELDS)
+def test_chat_rejects_unimplemented_fields_when_set(field, value) -> None:
+    state = _one_token_state()
+    resp = run(handle_chat_completion(chat_request(**{field: value}), None, state, {}))
+    assert getattr(resp, "status_code", 200) == 400
+    assert field in json.loads(resp.body)["error"]["message"]
+    assert state.sent is None
+
+
+@pytest.mark.parametrize("value", [None, False, [], {}, ""])
+@pytest.mark.parametrize("field", [f for f, _ in CHAT_UNIMPLEMENTED_FIELDS])
+def test_chat_accepts_unimplemented_fields_left_unset(field, value) -> None:
+    resp = run(handle_chat_completion(chat_request(**{field: value}), None, _one_token_state(), {}))
+    assert resp["choices"][0]["message"]["content"] == "A"

From 605b1ae0b5e48e551cbc3b26656ada5dde14f556 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:35:05 -0400
Subject: [PATCH 234/249] fix(server): bound logit_bias and stop_token_ids ids
 to the vocabulary

An id at or past len(tokenizer) (2**31 without a frontend tokenizer) in
logit_bias or stop_token_ids is a 400 on chat and completions before
anything is queued, and the tokenizer worker checks them again. Before, a
huge id overflowed msgpack at send_one (a 500) and one past the vocabulary
was silently ignored by the sampler.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/core.py                 | 10 +++++++
 python/freetoken/server/generation.py    | 15 ++++++----
 python/freetoken/server/openai_api.py    | 35 ++++++++++++++----------
 python/freetoken/tokenizer/tokenize.py   |  6 ++--
 tests/server/test_logprobs_api.py        | 35 +++++++++++++++++++++++-
 tests/server/test_sampling_validation.py | 11 ++++++++
 tests/tokenizer/test_tokenize.py         | 15 ++++++++++
 7 files changed, 102 insertions(+), 25 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index ab0141b1e5..541c1de761 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -61,6 +61,16 @@ class SamplingParams:
     # the detokenizer counts an unclosed generation as reasoning tokens.
     starts_in_reasoning: bool = False
 
+    def client_token_ids(self) -> list[tuple[str, list[int]]]:
+        """The token ids a client supplied, by request field (empty ones left out): each
+        must be checked against the tokenizer before the request reaches the engine."""
+        fields = [
+            ("allowed_token_ids", self.allowed_token_ids or []),
+            ("logit_bias", [int(tid) for tid, _ in self.logit_bias or ()]),
+            ("stop_token_ids", self.stop_token_ids),
+        ]
+        return [(name, ids) for name, ids in fields if ids]
+
     @property
     def is_greedy(self) -> bool:
         return self.temperature <= 0.0 or self.top_k == 1
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 95ae304fe0..615be79852 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -446,6 +446,9 @@ def parse_tools(self) -> bool:
 DEFAULT_MAX_OUTPUT_TOKENS = 32768
 # vLLM's MAX_NUM_ALLOWED_TOKEN_IDS: bounds the per-row mask upload
 MAX_ALLOWED_TOKEN_IDS = 1024
+# The int32 the engine holds ids in; also keeps them within what msgpack carries. The API
+# narrows this to len(tokenizer) when it has a frontend tokenizer, the tokenizer worker always.
+MAX_TOKEN_ID = 2**31 - 1
 
 
 def resolve_sampling(
@@ -515,12 +518,12 @@ def pick(value, key, framework):
                 tid = int(key)
             except (TypeError, ValueError):
                 raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None
-            if tid < 0:
-                raise ValueError(f"logit_bias token id must be >= 0, got {tid}")
+            if not 0 <= tid <= MAX_TOKEN_ID:
+                raise ValueError(f"logit_bias token id must be in [0, {MAX_TOKEN_ID}], got {tid}")
             bias.append([tid, max(-100.0, min(100.0, float(value)))])  # the OpenAI range
     ids = [int(t) for t in (stop_token_ids or [])]
-    if any(t < 0 for t in ids):
-        raise ValueError("stop_token_ids must be non-negative token ids")
+    if not all(0 <= t <= MAX_TOKEN_ID for t in ids):
+        raise ValueError(f"stop_token_ids must be token ids in [0, {MAX_TOKEN_ID}]")
     allowed = None
     if allowed_token_ids is not None:
         if not 1 <= len(allowed_token_ids) <= MAX_ALLOWED_TOKEN_IDS:
@@ -528,8 +531,8 @@ def pick(value, key, framework):
                 f"allowed_token_ids must hold 1 to {MAX_ALLOWED_TOKEN_IDS} ids, got {len(allowed_token_ids)}"
             )
         allowed = sorted({int(t) for t in allowed_token_ids})
-        if allowed[0] < 0:
-            raise ValueError("allowed_token_ids must be non-negative token ids")
+        if not 0 <= allowed[0] <= allowed[-1] <= MAX_TOKEN_ID:
+            raise ValueError(f"allowed_token_ids must be token ids in [0, {MAX_TOKEN_ID}]")
     params = SamplingParams(
         ignore_eos=ignore_eos,
         max_tokens=default_max_tokens if max_tokens is None else max_tokens,
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index 816fdfd255..d3cc50a23b 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -36,6 +36,7 @@
 )
 from .generation import (
     DEFAULT_MAX_OUTPUT_TOKENS,
+    MAX_TOKEN_ID,
     ContentDelta,
     GenDone,
     GenerationError,
@@ -344,6 +345,8 @@ async def handle_chat_completion(
         )
     except ValueError as exc:
         return create_error_response(str(exc))
+    if (bad := await _token_id_error(state, spec.sampling_params.client_token_ids())) is not None:
+        return bad
 
     if req.stream:
         # Non-stream requests already surface render failures as a clean 400
@@ -628,19 +631,11 @@ async def handle_completion(
     except ValueError as exc:
         return create_error_response(str(exc))
     token_prompts = _is_token_prompt(req.prompt)
-    if sampling.allowed_token_ids or token_prompts:
-        # The tokenizer worker checks again against its own tokenizer; this one answers
-        # before anything is queued (and keeps ids msgpack can carry without one).
-        bound = await _vocab_len(state) or _MAX_WIRE_TOKEN_ID
-        if sampling.allowed_token_ids and sampling.allowed_token_ids[-1] >= bound:
-            return create_error_response(
-                f"allowed_token_ids must be below the vocabulary size {bound}",
-                param="allowed_token_ids",
-            )
-        if token_prompts and not all(0 <= t < bound for ids in prompts for t in ids):
-            return create_error_response(
-                f"prompt token ids must be in [0, {bound})", param="prompt"
-            )
+    id_fields = sampling.client_token_ids()
+    if token_prompts:
+        id_fields += [("prompt", ids) for ids in prompts]
+    if (bad := await _token_id_error(state, id_fields)) is not None:
+        return bad
 
     if req.suffix is not None:
         markers = await _fim_markers(state)
@@ -900,8 +895,18 @@ async def pump(slot: int, stream: AsyncIterator[Any]) -> None:
 _FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")
 
 
-# Without a frontend tokenizer: the int32 the engine holds ids in (msgpack also needs a bound)
-_MAX_WIRE_TOKEN_ID = 2**31 - 1
+async def _token_id_error(state: Any, fields: list[tuple[str, list[int]]]) -> JSONResponse | None:
+    """A 400 for the first client token id outside the vocabulary, before anything is
+    queued. Without a frontend tokenizer the bound is the engine's int32; the tokenizer
+    worker checks against its own tokenizer either way."""
+    if not fields:
+        return None
+    bound = await _vocab_len(state) or MAX_TOKEN_ID + 1
+    for name, ids in fields:
+        if not all(0 <= t < bound for t in ids):
+            what = "prompt token ids" if name == "prompt" else name
+            return create_error_response(f"{what} must be in [0, {bound})", param=name)
+    return None
 
 
 async def _vocab_len(state: Any) -> int | None:
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index fdb65d8c53..86960299fc 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -107,10 +107,10 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
         results: List[UserMsg] = []
         # TODO: batch tokenization
         for msg in msgs:
-            allowed = msg.sampling_params.allowed_token_ids
             # past the real vocabulary an id indexes the padded logits tail or beyond them
-            if allowed and not _ids_below(allowed, len(self.tokenizer)):
-                raise ValueError(f"allowed_token_ids must be in [0, {len(self.tokenizer)})")
+            for name, ids in msg.sampling_params.client_token_ids():
+                if not _ids_below(ids, len(self.tokenizer)):
+                    raise ValueError(f"{name} must be in [0, {len(self.tokenizer)})")
             if msg.input_ids is not None:
                 # an id past the embedding rows would fault the engine, not just this request
                 if msg.input_ids and not _ids_below(msg.input_ids, len(self.tokenizer)):
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 38a1e135dc..91aaa874ff 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -957,7 +957,7 @@ def test_completion_allowed_token_ids_reach_the_engine_sorted_and_unique() -> No
 
 @pytest.mark.parametrize(
     ("ids", "fragment"),
-    [([], "1 to 1024"), (list(range(1025)), "1 to 1024"), ([-1, 3], "non-negative"), ([248077], "vocabulary")],
+    [([], "1 to 1024"), (list(range(1025)), "1 to 1024"), ([-1, 3], "[0, 2147483647]"), ([248077], "[0, 248077)")],
 )
 def test_completion_allowed_token_ids_out_of_range(ids, fragment) -> None:
     state = FakeState([], tokenizer=VocabTokenizer())
@@ -1110,3 +1110,36 @@ def test_chat_rejects_unimplemented_fields_when_set(field, value) -> None:
 def test_chat_accepts_unimplemented_fields_left_unset(field, value) -> None:
     resp = run(handle_chat_completion(chat_request(**{field: value}), None, _one_token_state(), {}))
     assert resp["choices"][0]["message"]["content"] == "A"
+
+
+ID_FIELD_CASES = [
+    ({"logit_bias": {"248077": 1.0}}, "logit_bias", True),
+    ({"stop_token_ids": [5, 248077]}, "stop_token_ids", True),
+    ({"logit_bias": {str(2**40): 1.0}}, "logit_bias", False),  # would overflow msgpack
+    ({"stop_token_ids": [2**40]}, "stop_token_ids", False),
+]
+
+
+@pytest.mark.parametrize(("extra", "field", "with_tokenizer"), ID_FIELD_CASES)
+@pytest.mark.parametrize("route", ["completions", "chat"])
+def test_token_id_fields_outside_the_vocabulary_are_400s(extra, field, with_tokenizer, route) -> None:
+    state = FakeState([reply("A", finished=True)], tokenizer=VocabTokenizer() if with_tokenizer else None)
+    if route == "completions":
+        resp = run(handle_completion(_completion(**extra), None, state, {}))
+    else:
+        resp = run(handle_chat_completion(chat_request(**extra), None, state, {}))
+    assert getattr(resp, "status_code", 200) == 400
+    assert field in json.loads(resp.body)["error"]["message"]
+    assert state.sent is None
+
+
+@pytest.mark.parametrize("route", ["completions", "chat"])
+def test_token_id_fields_at_the_vocab_edge_pass(route) -> None:
+    extra = {"logit_bias": {"248076": 1.0}, "stop_token_ids": [0, 248076]}
+    state = FakeState([reply("A", finished=True)], tokenizer=VocabTokenizer())
+    if route == "completions":
+        resp = run(handle_completion(_completion(**extra), None, state, {}))
+    else:
+        resp = run(handle_chat_completion(chat_request(**extra), None, state, {}))
+    assert getattr(resp, "status_code", 200) == 200
+    assert state.sent.sampling_params.stop_token_ids == [0, 248076]
diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py
index 61f521c7c0..590d6bc419 100644
--- a/tests/server/test_sampling_validation.py
+++ b/tests/server/test_sampling_validation.py
@@ -65,3 +65,14 @@ def test_allowed_token_ids_are_sorted_and_deduplicated():
 def test_allowed_token_ids_size_and_sign_are_checked(ids):
     with pytest.raises(ValueError, match="allowed_token_ids"):
         _resolve_with(allowed_token_ids=ids)
+
+
+@pytest.mark.parametrize(
+    "extra",
+    [{"logit_bias": {str(2**31): 1.0}}, {"stop_token_ids": [2**31]}, {"allowed_token_ids": [2**31]}],
+)
+def test_token_ids_are_bounded_by_the_engine_int32(extra):
+    with pytest.raises(ValueError, match=next(iter(extra))):
+        _resolve_with(**extra)
+    edge = {k: ({str(2**31 - 1): 1.0} if k == "logit_bias" else [2**31 - 1]) for k in extra}
+    _resolve_with(**edge)
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index 6fade1641b..3dfadd3a23 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -881,3 +881,18 @@ def test_tokenize_takes_token_id_prompts_as_is_within_the_vocabulary(ids, ok):
     else:
         with pytest.raises(ValueError, match="prompt token ids"):
             manager.tokenize([msg])
+
+
+@pytest.mark.parametrize(
+    ("params", "field"),
+    [
+        (SamplingParams(logit_bias=[[10, 1.0]]), "logit_bias"),
+        (SamplingParams(stop_token_ids=[3, 10]), "stop_token_ids"),
+    ],
+)
+def test_tokenize_range_checks_logit_bias_and_stop_token_ids(params, field):
+    manager = TokenizeManager(RawTokenizer())
+    with pytest.raises(ValueError, match=field):
+        manager.tokenize([TokenizeMsg(uid=1, text="hi", sampling_params=params)])
+    ok = SamplingParams(logit_bias=[[9, 1.0]], stop_token_ids=[0, 9])
+    assert manager.tokenize([TokenizeMsg(uid=1, text="hi", sampling_params=ok)])[0].input_ids.tolist() == [1, 2, 3]

From 10ea0821bf0b24db559d9711ef3671efaa08e141 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:37:48 -0400
Subject: [PATCH 235/249] feat(scheduler): limit the gdn snapshots a logprob
 read donates

--logprob-read-snapshot-donate {all,boundary,none} (default all, unchanged)
applies to finished requests with max_tokens=1 and logprobs on hybrid GDN
models: boundary donates only the frozen x64 prefill snapshot, none donates
neither it nor the live state and frees the pages, since a node without a
snapshot can never be resumed. A burst of label reads then cannot evict a
long conversation's snapshots.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/cli.md                                  |  1 +
 python/freetoken/scheduler/cache.py          | 25 +++++-
 python/freetoken/scheduler/config.py         |  2 +
 python/freetoken/scheduler/scheduler.py      |  1 +
 python/freetoken/server/args.py              | 11 +++
 tests/scheduler/test_hybrid_cache_manager.py | 90 ++++++++++++++++++++
 6 files changed, 127 insertions(+), 3 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index ebec8b405f..b5af7aa2d3 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -74,6 +74,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) |
 | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 |
 | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` |
+| `--logprob-read-snapshot-donate` | all | hybrid GDN models: snapshots a finished `max_tokens=1` + `logprobs` request donates to the prefix cache: `all` (x64 prefill boundary and final state), `boundary` (only the former) or `none` |
 | `--kv-cache-dtype` | bf16 | `bf16`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |
 
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index ff4467328e..c29d578b7b 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -29,9 +29,13 @@ def _swa_eviction_interval() -> int:
 _SWA_RETAIN_GAP = 16
 
 
+LOGPROB_READ_SNAPSHOT_DONATE = ("all", "boundary", "none")
+
+
 class CacheManager:
     def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str,
-                 linear_state_pool=None, swa_pool=None, sliding_window_size=None):
+                 linear_state_pool=None, swa_pool=None, sliding_window_size=None,
+                 logprob_read_snapshot_donate: str = "all"):
         # The `_free_slots` follows a page-aligned manner. For example, if page_size = 2,
         # the `_free_slots` may look like [0, 2, 4, 6, ...], and each slot represents a page.
         device = page_table.device
@@ -44,6 +48,12 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ
         self.swa_pool = swa_pool
         self.sliding_window_size = sliding_window_size
         self.is_hybrid = type == "hybrid_radix"
+        # GDN snapshots a finished logprob read (max_tokens=1 + logprobs) donates: "all" (the
+        # x64 prefill boundary and the live final state), "boundary" (only the former), "none".
+        # A burst of such reads otherwise evicts long conversations' snapshots.
+        if logprob_read_snapshot_donate not in LOGPROB_READ_SNAPSHOT_DONATE:
+            raise ValueError(f"unknown snapshot donate policy {logprob_read_snapshot_donate!r}")
+        self.logprob_read_snapshot_donate = logprob_read_snapshot_donate
         self.is_swa = type == "swa_radix"
         # swa_paged: this SWA model drives the global-paged swa pool -- true for BOTH the naive
         # (NaivePrefixCache, no reuse) and radix (SWARadixCache) paths. Gates the swa slot
@@ -384,10 +394,17 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             # is this request's dup to free. The frozen slot is consumed either way (taken by
             # the tree or freed here) and both ping-pong refs are dropped before
             # _free_req_slots so nothing double-frees.
+            sp = req.sampling_params
+            donate = (
+                self.logprob_read_snapshot_donate
+                if sp.logprobs and sp.max_tokens == 1
+                else "all"
+            )
             free_upto = old_handle.cached_len
             L = req.mamba_last_track_seqlen
             if (
-                L is not None
+                donate != "none"
+                and L is not None
                 and 0 < L <= req.cached_len
                 and align_down(L, self.page_size) == L
                 and req.mamba_ping_pong is not None
@@ -407,7 +424,9 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             # remain as reuse points).
             insert_len = align_down(req.cached_len, self.page_size)
             keep_live = False
-            if insert_len == req.cached_len and insert_len > 0:
+            # Without a donated snapshot a node could never be resumed (match_prefix needs
+            # one), so a non-donating read frees its pages instead of inserting them.
+            if donate == "all" and insert_len == req.cached_len and insert_len > 0:
                 prefix_len, mamba_exist = self.prefix_cache.insert(
                     req.input_ids[:insert_len], page_indices[:insert_len], req.linear_slot_idx)
                 self.unlock(old_handle)
diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py
index b2bbe0aefe..75502210ca 100644
--- a/python/freetoken/scheduler/config.py
+++ b/python/freetoken/scheduler/config.py
@@ -18,6 +18,8 @@ class SchedulerConfig(EngineConfig):
     offline_mode: bool = False
     decode_log_interval: int = 40
     special_token_ckpt: bool = False
+    # GDN snapshots a finished logprob read donates (CacheManager): all | boundary | none
+    logprob_read_snapshot_donate: str = "all"
 
     # networking config
     _unique_suffix: str = field(default_factory=_get_pid_suffix)
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 659a8bab3f..78d2edf25e 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -99,6 +99,7 @@ def __init__(self, config: SchedulerConfig):
                 (g.sliding_window for g in config.model_config.kv_cache_group_specs() if g.is_swa),
                 None,
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
+            logprob_read_snapshot_donate=getattr(config, "logprob_read_snapshot_donate", "all"),
         )
         self.decode_manager = DecodeManager(config.page_size)
         # Bound how long a run of prefill steps may starve in-flight decodes. A long
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 90904353c4..4fe0a7a8e5 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -545,6 +545,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "as a GDN-aware radix (cross-request GDN-state prefix reuse); pass 'naive' to opt out.",
     )
 
+    parser.add_argument(
+        "--logprob-read-snapshot-donate",
+        type=str,
+        default=ServerArgs.logprob_read_snapshot_donate,
+        choices=["all", "boundary", "none"],
+        help="Hybrid (GDN) models: which state snapshots a finished logprob read (max_tokens=1 "
+        "with logprobs, e.g. AnyJev label scoring) donates to the prefix cache. 'all' (default): "
+        "the x64 prefill boundary and the final state; 'boundary': only the former; 'none': "
+        "neither, so a burst of reads cannot evict other conversations' snapshots.",
+    )
+
     parser.add_argument(
         "--text-model-only",
         action="store_true",
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index 165689ab3b..6b4f164b52 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -254,3 +254,93 @@ def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch):
 
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))
+
+
+def _finish_read(policy: str, sp: SamplingParams, *, warm_prefix: bool = False):
+    """A 130-token prompt finishing on its prefill forward, with a frozen x64 snapshot at 128
+    (the deepest boundary strictly inside it) and its live state at 130. ``warm_prefix`` first
+    donates a 64-token prefix so the read resumes from a locked prefix-cache hit."""
+    num_pages = 256
+    pool = _pool()
+    pt = torch.zeros(4, 256, dtype=torch.int32)
+    cm = CacheManager(num_pages, 1, pt, "hybrid_radix", linear_state_pool=pool,
+                      logprob_read_snapshot_donate=policy)
+    prompt = list(range(1, 131))
+    if warm_prefix:
+        warm = Req(input_ids=torch.tensor(prompt[:64] + [999], dtype=torch.int32), table_idx=1,
+                   cached_len=64, output_len=1, uid=9, sampling_params=SamplingParams(),
+                   cache_handle=cm.match_req(_pend(prompt[:64])).cuda_handle)
+        pt[1, :64] = cm._allocate(64)
+        warm.linear_slot_idx, warm.mamba_ping_pong = pool.alloc(1)[0], tuple(pool.alloc(2))
+        cm.lock(warm.cache_handle)
+        cm.cache_req(warm, finished=True)
+    mr = cm.match_req(_pend(prompt))
+    hit = mr.cuda_handle.cached_len
+    pt[0, :hit] = mr.cuda_handle.get_matched_indices()
+    pt[0, hit:130] = cm._allocate(130 - hit)
+    req = Req(input_ids=torch.tensor(prompt + [7], dtype=torch.int32), table_idx=0, cached_len=130,
+              output_len=1, uid=0, sampling_params=sp, cache_handle=mr.cuda_handle)
+    req.linear_slot_idx, req.mamba_ping_pong = pool.alloc(1)[0], tuple(pool.alloc(2))
+    req.mamba_next_track_idx, req.mamba_last_track_seqlen = 1, 128
+    cm.lock(mr.cuda_handle)
+    cm.cache_req(req, finished=True)
+
+    pc = cm.prefix_cache
+    tree_pages = pc.full_evictable + pc.full_protected
+    assert len(cm.free_slots) + tree_pages == num_pages  # every page is free or tree-owned
+    assert pc.full_protected == 0 and pc.mamba_protected == 0  # nothing left locked
+    assert pool.num_free_slots + pc.mamba_evictable_size == pool.num_slots - 1  # no GDN slot leak
+    cm.check_integrity()
+    return cm, prompt
+
+
+def _resume_len(cm, prompt, n):
+    return cm.match_req(_pend(prompt[:n] + [12345])).cuda_handle.cached_len
+
+
+READ = SamplingParams(max_tokens=1, logprobs=True, top_logprobs=4, allowed_token_ids=[32, 33])
+
+
+@pytest.mark.parametrize("warm_prefix", [False, True])
+@pytest.mark.parametrize(
+    ("policy", "snapshots", "resume_at_128", "resume_at_130"),
+    [("all", 2, 128, 130), ("boundary", 1, 128, 128), ("none", 0, 0, 0)],
+)
+def test_logprob_read_snapshot_donate_policy(policy, snapshots, resume_at_128, resume_at_130, warm_prefix):
+    cm, prompt = _finish_read(policy, READ, warm_prefix=warm_prefix)
+    warm = 1 if warm_prefix else 0  # the warm-up request's own 64-token snapshot
+    assert cm.prefix_cache.mamba_evictable_size == snapshots + warm
+    assert _resume_len(cm, prompt, 128) == max(resume_at_128, 64 * warm)
+    assert _resume_len(cm, prompt, 130) == max(resume_at_130, 64 * warm)
+
+
+@pytest.mark.parametrize("policy", ["boundary", "none"])
+@pytest.mark.parametrize(
+    "sp",
+    [
+        SamplingParams(max_tokens=1),  # no logprobs
+        SamplingParams(max_tokens=2, logprobs=True, top_logprobs=4),  # not a single-token read
+    ],
+)
+def test_snapshot_donate_policy_leaves_other_requests_alone(policy, sp):
+    cm, prompt = _finish_read(policy, sp)
+    assert cm.prefix_cache.mamba_evictable_size == 2
+    assert _resume_len(cm, prompt, 130) == 130
+
+
+def test_snapshot_donate_policy_is_validated():
+    with pytest.raises(ValueError, match="snapshot donate"):
+        CacheManager(8, 1, torch.zeros(1, 8, dtype=torch.int32), "hybrid_radix",
+                     linear_state_pool=_pool(), logprob_read_snapshot_donate="some")
+
+
+def test_snapshot_donate_policy_is_a_server_arg():
+    from unittest.mock import patch
+
+    from freetoken.server.args import parse_args
+
+    config = SimpleNamespace(to_dict=lambda: {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"})
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):
+        assert parse_args(["--model", "/models/anon"])[0].logprob_read_snapshot_donate == "all"
+        args, _ = parse_args(["--model", "/models/anon", "--logprob-read-snapshot-donate", "none"])
+    assert args.logprob_read_snapshot_donate == "none"

From d75874ba1943f048ca06375ddde5f479308663f8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 11:00:54 -0400
Subject: [PATCH 236/249] test(moe): drive #334's admission tests through
 #700's method loader

#700 renamed _legacy_expert_banks to _method_expert_banks and replaced load_expert_banks' dtype argument with the quant method, so the tests patched a missing name and passed an unknown keyword.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 tests/moe/test_expert_bank_memory_admission.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/moe/test_expert_bank_memory_admission.py b/tests/moe/test_expert_bank_memory_admission.py
index 221814212d..c3e4fc368d 100644
--- a/tests/moe/test_expert_bank_memory_admission.py
+++ b/tests/moe/test_expert_bank_memory_admission.py
@@ -44,14 +44,14 @@ def effective_fixture():
     chosen = []
 
     def build(*args, **kwargs):
-        chosen.append(args[5])  # _legacy_expert_banks(..., dummy, parallel, ...)
+        chosen.append(args[5])  # _method_expert_banks(..., dummy, parallel, ...)
         return expert_banks.ExpertBanks("bf16", {})
 
-    monkeypatch.setattr(expert_banks, "_legacy_expert_banks", build)
+    monkeypatch.setattr(expert_banks, "_method_expert_banks", build)
     config = SimpleNamespace(num_moe_layers=1, expert_quant="none")
 
     expert_banks.load_expert_banks(
-        str(tmp_path), config, device=torch.device("cpu"), dtype=torch.bfloat16
+        str(tmp_path), config, device=torch.device("cpu"), method=SimpleNamespace(kind=None)
     )
 
     assert chosen == [False]
@@ -76,14 +76,14 @@ def build(*args, **kwargs):
         chosen.append(args[5])
         return expert_banks.ExpertBanks("bf16", {})
 
-    monkeypatch.setattr(expert_banks, "_legacy_expert_banks", build)
+    monkeypatch.setattr(expert_banks, "_method_expert_banks", build)
     config = SimpleNamespace(num_moe_layers=1, expert_quant="none")
 
     expert_banks.load_expert_banks(
         str(tmp_path),
         config,
         device=torch.device("cpu"),
-        dtype=torch.bfloat16,
+        method=SimpleNamespace(kind=None),
         parallel=True,
     )
 

From cbf345b78215ce4011a7bc023ec9f6f783def872 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 11:00:54 -0400
Subject: [PATCH 237/249] test(moe): stub #697's capture-time graph upload in
 the warm offload capture test

The fake graph has no CUDA graph exec for cuGraphUpload; the test fails the same way on upstream cc56b6bd.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 tests/moe/test_offload.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index fbacc8e5fa..780cdc4e3a 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -710,9 +710,13 @@ def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch):
     from freetoken.core import Context, Req, get_global_ctx
     from freetoken.engine.graph import GraphRunner
 
+    import freetoken.engine.graph as graph_mod
+
     events = []
     _init_tp()
     monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
+    # the fake graph has no CUDA graph exec for #697's capture-time cuGraphUpload
+    monkeypatch.setattr(graph_mod, "_upload_graph", lambda graph, stream: None)
 
     class FakeGraph:
         def pool(self):

From 840aa6db7790dc5a41bd10774e2e6b341c8d97c3 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 11:00:54 -0400
Subject: [PATCH 238/249] test(minimax_m3): keep the swigluoai kernel choice
 independent of the test host's gpu

W4A4 asks the real device whether it is worth it, so on an sm_120 host auto picks w4a4 and the triton assertion fails (also on upstream cc56b6bd).

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 tests/models/test_minimax_m3.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/models/test_minimax_m3.py b/tests/models/test_minimax_m3.py
index dc0ab90df5..e8a8f2919a 100644
--- a/tests/models/test_minimax_m3.py
+++ b/tests/models/test_minimax_m3.py
@@ -275,10 +275,13 @@ def test_nvfp4_experts_restricted_to_triton_for_swigluoai(monkeypatch):
     from freetoken.kernel import backend
     from freetoken.layers.quantization import KernelSelectionError, ModelOptConfig, select_kernel
     from freetoken.layers.quantization.moe import Nvfp4MoEMethod
+    from freetoken.layers.quantization.moe.nvfp4 import W4a4Nvfp4MoEKernel
 
     monkeypatch.setattr(backend, "device_capability", lambda: (12, 0))
     monkeypatch.setattr(backend, "is_vllm_installed", lambda: True)
     monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: True)
+    # W4A4 asks the real device whether it is worth it; keep the selection independent of the test host
+    monkeypatch.setattr(W4a4Nvfp4MoEKernel, "worth_it", lambda self, cfg: False)
     cfg = MoEConfig(num_experts=256, hidden=6144, intermediate=3072, top_k=4, scheme=ModelOptConfig.SCHEMES["NVFP4"], strategy="offload", activation="swigluoai", alpha=1.702, limit=7.0)
     assert select_kernel(Nvfp4MoEMethod.candidates, "auto", cfg).name == "triton"
     with pytest.raises(KernelSelectionError):

From 76f6f1e71033cc241d98e797d03caba44f131670 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:42:48 -0400
Subject: [PATCH 239/249] fix(server): report the admitted prompt usage in the
 anthropic message_start

Streamed /v1/messages opened with input_tokens 0 and no cache fields (#718).
generate_events now yields PromptAdmitted on the scheduler's admission ack and
message_start waits for it, so its input/cache split matches message_delta. A
keepalive or error before admission still sends message_start first, with zeros.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/server/anthropic_api.py    | 39 +++++++++---
 python/freetoken/server/generation.py       | 16 ++++-
 tests/server/test_anthropic_api.py          | 69 +++++++++++++++++++++
 tests/server/test_streaming_model_matrix.py |  4 +-
 4 files changed, 118 insertions(+), 10 deletions(-)

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 35d902b9d6..167db5a995 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -40,6 +40,7 @@
     GenerationError,
     GenResult,
     GenSpec,
+    PromptAdmitted,
     ReasoningDelta,
     ToolCallArgsDelta,
     ToolCallsDelta,
@@ -502,14 +503,21 @@ def _stop_block() -> list[str]:
         block_index += 1
         return frames
 
-    # message_start (input_tokens filled in by the terminal message_delta).
-    yield _event(AnthropicStreamEvent(
-        type="message_start",
-        message=AnthropicMessagesResponse(
-            id=f"msg_{uid}", content=[], model=model,
-            usage=AnthropicUsage(input_tokens=0, output_tokens=0),
-        ),
-    ))
+    started = False
+
+    def _message_start(usage: AnthropicUsage) -> str:
+        nonlocal started
+        started = True
+        return _event(AnthropicStreamEvent(
+            type="message_start",
+            message=AnthropicMessagesResponse(
+                id=f"msg_{uid}", content=[], model=model, usage=usage,
+            ),
+        ))
+
+    # message_start waits for the prompt's admission so it can carry the real input/cache
+    # usage; anything that must go out first (a keepalive, an error) sends it with zeros.
+    zero_usage = AnthropicUsage(input_tokens=0, output_tokens=0)
 
     def _open_tool(name: str | None, ordinal: int | None, stable: bool = True) -> list[str]:
         nonlocal block_open, tool_args_sent, tool_ordinal, tool_stable
@@ -543,6 +551,15 @@ def _tool_args_delta(fragment: str) -> str:
     events = with_keepalive(events, KEEPALIVE_INTERVAL_S)
     try:
         async for ev in events:
+            if isinstance(ev, PromptAdmitted):
+                if not started:
+                    yield _message_start(
+                        _anthropic_usage(ev.prompt_tokens, 0, ev.cached_tokens, cache_report)
+                    )
+                continue
+            if not started:
+                yield _message_start(zero_usage)
+
             if ev is KEEPALIVE:
                 yield _event(AnthropicStreamEvent(type="ping"))
 
@@ -625,10 +642,14 @@ def _tool_args_delta(fragment: str) -> str:
                 yield _event(AnthropicStreamEvent(type="message_stop"))
                 # Anthropic streams terminate on message_stop — no OpenAI-style
                 # `data: [DONE]` sentinel (it would be an unknown event to strict clients).
+        if not started:
+            yield _message_start(zero_usage)
     except GenerationError as exc:
         # Request-side failure (template rejection / over-length prompt). Match the
         # non-streaming path's classification so Claude Code treats it as a client error
         # rather than a server fault to retry.
+        if not started:
+            yield _message_start(zero_usage)
         if block_open:
             for f in _stop_block():
                 yield f
@@ -636,6 +657,8 @@ def _tool_args_delta(fragment: str) -> str:
             type="error", error=AnthropicError(type="invalid_request_error", message=str(exc)),
         ))
     except Exception as exc:  # noqa: BLE001 — surface as an Anthropic error event
+        if not started:
+            yield _message_start(zero_usage)
         if block_open:
             for f in _stop_block():
                 yield f
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 615be79852..f10a73b61f 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -319,6 +319,16 @@ def __init__(self, message: str, code: str | None = None):
 # so generation/parsing logic lives in exactly one place and no adapter has to
 # re-parse a serialized OpenAI stream.
 # --------------------------------------------------------------------------- #
+@dataclass
+class PromptAdmitted:
+    """The scheduler admitted the prompt: its token count and prefix-cache hit are known
+    before any output. Sent once, ahead of the first delta, so a protocol that reports
+    input usage up front (Anthropic's message_start) need not wait for GenDone."""
+
+    prompt_tokens: int
+    cached_tokens: int = 0
+
+
 @dataclass
 class ReasoningDelta:
     text: str
@@ -399,7 +409,7 @@ class GenDone:
     logprobs: list[dict] | None = None
 
 
-GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone
+GenEvent = PromptAdmitted | ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone
 
 
 @dataclass
@@ -1270,6 +1280,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
 
     engine_finish_reason: str | None = None
     engine_matched_stop: str | None = None
+    admitted = False
     async for ack in state.wait_for_ack(uid):
         if getattr(ack, "error", None):
             raise GenerationError(ack.error, getattr(ack, "error_code", None))
@@ -1278,6 +1289,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
+        if ack.prompt_tokens_delta and not admitted:
+            admitted = True
+            yield PromptAdmitted(prompt_tokens, cached_tokens)
         if segmenter is not None:
             pending_logprobs.extend(segmenter.feed(untrimmed(getattr(ack, "logprobs", None))))
         content_delta = ack.incremental_output
diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py
index 890339522c..c9faf5a6c2 100644
--- a/tests/server/test_anthropic_api.py
+++ b/tests/server/test_anthropic_api.py
@@ -475,6 +475,71 @@ def test_stream_empty_turn_terminates():
     assert types == ["message_start", "message_delta", "message_stop"]
 
 
+@pytest.mark.parametrize("cache_report", [False, True])
+def test_stream_message_start_reports_the_admitted_prompt_usage(cache_report):
+    # #718: message_start carries the prompt's input/cache split, the same one the final
+    # message_delta reports, instead of hard-coded zeros.
+    from freetoken.server.generation import PromptAdmitted
+
+    collected = _collect_events(
+        [PromptAdmitted(1200, 1024), ContentDelta("hi"),
+         GenDone("stop", 1200, 5, cached_tokens=1024)],
+        cache_report=cache_report,
+    )
+    types = [e[0] for e in collected]
+    assert types.count("message_start") == 1 and types[0] == "message_start"
+    start = collected[0][1]["message"]["usage"]
+    delta = next(d for t, d in collected if t == "message_delta")["usage"]
+    assert start == {**delta, "output_tokens": 0}
+    expected = {"input_tokens": 176, "cache_read_input_tokens": 1024} if cache_report \
+        else {"input_tokens": 1200}
+    assert start == {**expected, "output_tokens": 0}
+
+
+def test_stream_keepalive_before_admission_sends_message_start_first(monkeypatch):
+    # A long queue/prefill must not open with a ping: message_start goes out first (zeros,
+    # the counts are not known yet), and the later admission does not send a second one.
+    from freetoken.server.generation import PromptAdmitted
+
+    monkeypatch.setattr(A, "KEEPALIVE_INTERVAL_S", 0.01)
+
+    async def slow():
+        await asyncio.sleep(0.1)
+        yield PromptAdmitted(10, 0)
+        yield ContentDelta("hi")
+        yield GenDone("stop", 10, 1)
+
+    async def run():
+        out = []
+        async for frame in A.anthropic_event_stream(slow(), "claude-x", 1):
+            out.append(frame.split("\n")[0][len("event: "):])
+        return out
+
+    types = asyncio.run(run())
+    assert types[0] == "message_start" and types[1] == "ping"
+    assert types.count("message_start") == 1
+    assert types[-2:] == ["message_delta", "message_stop"]
+
+
+def test_generate_events_announces_prompt_admission():
+    # The admission ack's counts leave the primitive as an event before any content.
+    from freetoken.server.generation import PromptAdmitted, generate_events
+
+    fake = FakeState([("", False, 1200, 0), ("hi", True, 0, 1)])
+    fake.cached = 1024
+    spec = A.convert_anthropic_to_genspec(
+        AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "x"}]),
+        {},
+    )
+
+    async def run():
+        return [ev async for ev in generate_events(1, spec, fake)]
+
+    events = asyncio.run(run())
+    assert events[0] == PromptAdmitted(1200, 1024)
+    assert sum(isinstance(ev, PromptAdmitted) for ev in events) == 1
+
+
 # --------------------------------------------------------------------------- #
 # Route smoke tests (real app + shared primitive + fake engine)
 # --------------------------------------------------------------------------- #
@@ -492,6 +557,7 @@ def __init__(self, outputs):
         )
         self._uid = 0
         self._count_manager = None
+        self.cached = 0
 
     def frontend_tokenizer(self):
         return self._count_manager
@@ -511,6 +577,7 @@ async def wait_for_ack(self, uid):
                 finished=finished,
                 prompt_tokens_delta=pt,
                 completion_tokens_delta=ct,
+                cached_tokens=self.cached if pt else 0,
             )
 
     async def stream_with_cancellation(self, gen, request, uid):
@@ -574,6 +641,8 @@ def test_route_stream_text():
             etypes.append(etype)
         if etype == "content_block_delta" and data and data["delta"].get("type") == "text_delta":
             text += data["delta"]["text"]
+        if etype == "message_start":
+            assert data["message"]["usage"]["input_tokens"] == 5
     assert etypes[0] == "message_start"
     assert "content_block_start" in etypes
     assert etypes[-1] == "message_stop"
diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py
index 2226fe7fe5..bdd71b2e22 100644
--- a/tests/server/test_streaming_model_matrix.py
+++ b/tests/server/test_streaming_model_matrix.py
@@ -523,7 +523,9 @@ def test_tool_call_events_shape_and_order():
     block = '<tool_call>\n{"name": "read", "arguments": {"filePath": "/a"}}\n</tool_call>'
     events = _stream_events(["Checking. ", block + "\n", "Done"], "qwen25")
     kinds = [type(ev).__name__ for ev in events]
-    # text, call open + streamed args + close, trailing text, done — in wire order
+    # admission, text, call open + streamed args + close, trailing text, done — in wire order
+    assert kinds.pop(0) == "PromptAdmitted"
+    events = events[1:]
     assert kinds[0] == "ContentDelta" and kinds[1] == "ToolCallStart"
     assert kinds[-3:] == ["ToolCallsDelta", "ContentDelta", "GenDone"]
     assert all(k == "ToolCallArgsDelta" for k in kinds[2:-3])

From 27cd6f153e1a06f4ad6ead59e349ca7828670e05 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:58:56 -0400
Subject: [PATCH 240/249] fix(server): reject an over-long prompt with a 400
 before the stream starts

The context check ran only in the scheduler, after the worker tokenized, so a
streamed request whose prompt did not fit got HTTP 200 and an in-stream error
(#749). Every endpoint now renders and tokenizes the prompt frontend-side before
any response and rejects a prompt that leaves no room for one output token with
the scheduler's "prompt is too long" text and code context_length_exceeded,
against the /v1/models context. A max_tokens that does not fit is still clamped
by the scheduler: agents always send a large one. The worker serves the counted
ids instead of tokenizing again.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/message/tokenizer.py  |   6 +-
 python/freetoken/server/generation.py  | 151 +++++++++++++------
 python/freetoken/server/openai_api.py  |  62 +++-----
 python/freetoken/tokenizer/tokenize.py |  34 +++--
 tests/server/test_anthropic_api.py     |   9 ++
 tests/server/test_logprobs_api.py      |   8 +-
 tests/server/test_openai_api.py        | 191 +++++++++++++++++++++++++
 tests/server/test_openai_extras.py     |   2 +
 8 files changed, 360 insertions(+), 103 deletions(-)

diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index 30efdac0d7..fcd599681f 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -90,8 +90,10 @@ class TokenizeMsg(BaseTokenizerMsg):
     tools: List[Dict[str, Any]] | None = None
     images: List[bytes] | None = None
     inline_system_policy: str | None = None
-    # A token-id prompt (OpenAI / vLLM completions): used as is, ``text`` is ignored.
-    input_ids: List[int] | None = None
+    # Prompt ids served as they are (``text`` is then ignored for tokenization): a client
+    # token-id prompt (OpenAI / vLLM completions; a list, range-checked by the worker), or
+    # the int32 tensor the frontend already rendered and counted (no second tokenization).
+    input_ids: Any = None
 
 
 @dataclass
diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index f10a73b61f..4fe7d13580 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import json
 import math
 import time
@@ -441,6 +442,9 @@ class GenSpec:
     template_tools: list[dict[str, Any]] | None = None   # tools the model sees (TokenizeMsg.tools)
     parser_tools: list[dict[str, Any]] | None = None     # tools for FunctionCallParser; None disables parsing
     inline_system_policy: str | None = None
+    # Prompt ids counted by the frontend pre-check, handed to the worker so it does not
+    # render and tokenize again (also reused by every choice of an n > 1 request).
+    prompt_ids: Any = field(default=None, repr=False, compare=False)
 
     @property
     def parse_tools(self) -> bool:
@@ -685,22 +689,112 @@ def split_tool_lists(
 # --------------------------------------------------------------------------- #
 async def submit_generation(spec: GenSpec, state: Any) -> int:
     """Enqueue one generation from a GenSpec; return its uid. Every protocol adapter
-    calls this — it takes the neutral spec, not a wire request type."""
+    calls this — it takes the neutral spec, not a wire request type. Raises
+    GenerationError (a 400 in every adapter) when the prompt cannot be rendered or does
+    not fit the context, before anything is committed to the client."""
     refs = collect_image_refs(spec.messages)
     images = await _resolve_images(refs, state) if refs else None
-    uid = state.new_user()
-    await state.send_one(
-        TokenizeMsg(
-            uid=uid,
-            text=spec.messages,
-            sampling_params=spec.sampling_params,
-            chat_template_kwargs=spec.chat_template_kwargs,
-            tools=spec.template_tools,
-            images=images,
-            inline_system_policy=spec.inline_system_policy,
-        )
+    msg = TokenizeMsg(
+        uid=0,
+        text=spec.messages,
+        sampling_params=spec.sampling_params,
+        chat_template_kwargs=spec.chat_template_kwargs,
+        tools=spec.template_tools,
+        images=images,
+        inline_system_policy=spec.inline_system_policy,
     )
-    return uid
+    if spec.prompt_ids is None:
+        spec.prompt_ids = await frontend_prompt_ids(msg, state)
+    msg.uid = state.new_user()
+    msg.input_ids = spec.prompt_ids
+    await state.send_one(msg)
+    return msg.uid
+
+
+def model_context_length(state: Any) -> int | None:
+    """The usable context: ``min(model ceiling, allocated KV tokens)``.
+
+    This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__,
+    and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler
+    admits against. The frontend process holds an unclamped ServerArgs copy, so publishing
+    ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144
+    against a 178176-token pool, so the engine and this route disagreed about one quantity.
+
+    `ft launch` reads this to size each agent's context window (opencode's ``limit.context``,
+    codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops
+    those agents compacting before the KV pool runs out.
+    """
+    try:  # never 500 a metadata route: max_seq_len walks into the HF config on some builds
+        value = int(state.config.max_seq_len)
+    except Exception:  # noqa: BLE001
+        return None
+    if value <= 0:
+        return None
+    try:
+        # Local import: api_server imports the protocol modules, which import this one.
+        from .api_server import kv_pool_geometry
+
+        num_pages, page_size = kv_pool_geometry(state)
+        kv_tokens = num_pages * page_size
+    except Exception:  # noqa: BLE001
+        kv_tokens = 0
+    # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then.
+    return min(value, kv_tokens) if kv_tokens > 0 else value
+
+
+def check_context_fits(prompt_tokens: int, state: Any) -> None:
+    """Raise the scheduler's context_length_exceeded error up front when the prompt leaves
+    no room for one output token. A max_tokens that does not fit is not an error: the
+    scheduler clamps it to the room left, as agents that always send a large one expect."""
+    limit = model_context_length(state)
+    if limit is not None and prompt_tokens >= limit:
+        # "prompt is too long: N tokens > M" is the phrasing Claude Code and OpenClaw match
+        # on (the scheduler sends the same text).
+        raise GenerationError(
+            f"prompt is too long: {prompt_tokens} tokens > {limit} maximum (prompt + generation); "
+            f"shorten the prompt or increase the KV cache budget",
+            "context_length_exceeded",
+        )
+
+
+async def frontend_prompt_ids(msg: TokenizeMsg, state: Any) -> Any | None:
+    """Render and tokenize ``msg`` frontend-side and check it against the context, so a
+    request that cannot run is a 400 *before* an adapter commits an SSE stream (once headers
+    go out the failure can only ride in-stream, where some agents show nothing but "empty
+    response", and clients that compact on "prompt is too long" never see it).
+
+    Returns the prompt ids for the worker to serve as they are (no second tokenization), or
+    None when there is nothing to hand over: no frontend tokenizer (minimal embeddings), one
+    that fails to *initialize* (a server fault, not this request's), or an image prompt (the
+    worker's multimodal processor expands it; only the render is checked here and the
+    scheduler keeps the length check). Raises GenerationError for a render or encode failure
+    (the worker's own classification) and for a prompt that does not fit."""
+    build = getattr(state, "frontend_tokenizer", None)
+    if build is None:
+        return None
+    try:
+        manager = await asyncio.to_thread(build)
+    except Exception:  # noqa: BLE001 -- server fault, not this request's problem
+        return None
+    if manager is None:
+        return None
+
+    if msg.images:
+        def work() -> Any:
+            manager.render_prompt(dataclasses.replace(msg, images=None))
+            return None
+    else:
+        def work() -> Any:
+            return manager.encode_prompt(msg, manager.render_prompt(msg))
+
+    try:
+        ids = await asyncio.to_thread(work)
+    except Exception as exc:  # noqa: BLE001 -- mirror the worker's classification
+        raise GenerationError(f"could not encode request: {exc}") from exc
+    if ids is None:
+        return None
+    check_context_fits(int(ids.numel()), state)
+    return ids
 
 
 async def count_prompt_tokens(
@@ -744,37 +838,6 @@ class the generation path maps to a 400. A tokenizer *initialization* failure (m
     return int(user_msg.input_ids.numel())
 
 
-async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None:
-    """Render ``spec``'s prompt frontend-side, returning the failure a streaming
-    adapter should surface as an HTTP 400 *before* committing an SSE stream —
-    once headers go out, a template rejection can only ride in-stream, where
-    some agents show nothing but "empty response". Render only; the worker
-    still renders and encodes authoritatively. Best-effort: a state without a
-    frontend tokenizer, or one that fails to *initialize*, skips validation
-    rather than blocking the generation path.
-    """
-    build = getattr(state, "frontend_tokenizer", None)
-    if build is None:
-        return None
-    msg = TokenizeMsg(
-        uid=0,
-        text=spec.messages,
-        sampling_params=SamplingParams(),
-        chat_template_kwargs=spec.chat_template_kwargs,
-        tools=spec.template_tools,
-        inline_system_policy=spec.inline_system_policy,
-    )
-    try:
-        manager = await asyncio.to_thread(build)
-    except Exception:  # noqa: BLE001 -- server fault, not this request's problem
-        return None
-    try:
-        await asyncio.to_thread(manager.render_prompt, msg)
-    except Exception as exc:  # noqa: BLE001 -- mirror the worker's classification
-        return GenerationError(f"could not encode request: {exc}")
-    return None
-
-
 def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None:
     """Build a reasoning parser for this generation, or None if the server has no
     reasoning parser configured."""
diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index d3cc50a23b..d447f2a3b2 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -49,10 +49,12 @@
     build_metrics,
     chat_logprob_markers,
     chat_logprobs_supported,
+    check_context_fits,
+    frontend_prompt_ids,
     generate_events,
     generate_full,
     metrics_enabled,
-    prerender_error,
+    model_context_length,
     render_messages,
     resolve_sampling,
     submit_generation,
@@ -160,7 +162,7 @@ async def detokenize(req: DetokenizeRequest):
     async def v1_models():
         state = get_state()
         model_id = _served_model_name(state)
-        ctx = _model_context_length(state)
+        ctx = model_context_length(state)
         efforts, default_effort = await _effort_fields(state)
         return ModelList(data=[ModelCard(
             id=model_id,
@@ -193,7 +195,7 @@ async def v1_model(model_id: str):
                 code="model_not_found",
             )
 
-        ctx = _model_context_length(state)
+        ctx = model_context_length(state)
         efforts, default_effort = await _effort_fields(state)
         return ModelCard(
             id=served_model_id,
@@ -348,13 +350,6 @@ async def handle_chat_completion(
     if (bad := await _token_id_error(state, spec.sampling_params.client_token_ids())) is not None:
         return bad
 
-    if req.stream:
-        # Non-stream requests already surface render failures as a clean 400
-        # through GenerationError; only the stream path needs the pre-check.
-        err = await prerender_error(spec, state)
-        if err is not None:
-            return create_error_response(str(err), code=err.code)
-
     # n > 1: one generation per choice, submitted together (the prefix cache serves the
     # shared prompt after the first prefill).
     uids: list[int] = []
@@ -652,14 +647,26 @@ async def handle_completion(
         echo_texts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt)
 
     # choices are ordered prompt-major: prompt i, sample k -> index i * n + k (OpenAI)
+    prompt_ids = []
+    try:
+        for prompt in prompts:
+            if token_prompts:
+                check_context_fits(len(prompt), state)
+                prompt_ids.append(None)
+                continue
+            prompt_ids.append(await frontend_prompt_ids(
+                TokenizeMsg(uid=0, text=prompt, sampling_params=sampling), state
+            ))
+    except GenerationError as exc:
+        return create_error_response(str(exc), code=exc.code)
     uids: list[int] = []
-    for prompt in prompts:
+    for prompt, ids in zip(prompts, prompt_ids):
         for _ in range(req.n):
             uid = state.new_user()
             msg = (
                 TokenizeMsg(uid=uid, text="", input_ids=prompt, sampling_params=sampling)
                 if token_prompts
-                else TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling)
+                else TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling, input_ids=ids)
             )
             await state.send_one(msg)
             uids.append(uid)
@@ -973,7 +980,7 @@ async def handle_tokenize(req: TokenizeRequest, state: Any):
         return create_error_response(str(exc))
     body: dict[str, Any] = {
         "count": len(ids),
-        "max_model_len": _model_context_length(state),
+        "max_model_len": model_context_length(state),
         "tokens": [int(t) for t in ids],
     }
     if req.return_token_strs:
@@ -1288,32 +1295,3 @@ def _served_model_name(state: Any) -> str:
     return getattr(state.config, "served_model_name", None) or state.config.model_path
 
 
-def _model_context_length(state: Any) -> int | None:
-    """The usable context: ``min(model ceiling, allocated KV tokens)``.
-
-    This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__,
-    and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler
-    admits against. The frontend process holds an unclamped ServerArgs copy, so publishing
-    ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144
-    against a 178176-token pool, so the engine and this route disagreed about one quantity.
-
-    `ft launch` reads this to size each agent's context window (opencode's ``limit.context``,
-    codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops
-    those agents compacting before the KV pool runs out.
-    """
-    try:  # never 500 a metadata route: max_seq_len walks into the HF config on some builds
-        value = int(state.config.max_seq_len)
-    except Exception:  # noqa: BLE001
-        return None
-    if value <= 0:
-        return None
-    try:
-        # Local import: api_server imports this module, so a module-level one would cycle.
-        from .api_server import kv_pool_geometry
-
-        num_pages, page_size = kv_pool_geometry(state)
-        kv_tokens = num_pages * page_size
-    except Exception:  # noqa: BLE001
-        kv_tokens = 0
-    # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then.
-    return min(value, kv_tokens) if kv_tokens > 0 else value
diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index 86960299fc..dc7144d574 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -111,26 +111,18 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
             for name, ids in msg.sampling_params.client_token_ids():
                 if not _ids_below(ids, len(self.tokenizer)):
                     raise ValueError(f"{name} must be in [0, {len(self.tokenizer)})")
-            if msg.input_ids is not None:
+            if isinstance(msg.input_ids, torch.Tensor) and not msg.images:
+                # the frontend's own encoding of this prompt (frontend_prompt_ids)
+                input_ids = msg.input_ids.view(-1).to(torch.int32)
+            elif msg.input_ids is not None:
                 # an id past the embedding rows would fault the engine, not just this request
                 if msg.input_ids and not _ids_below(msg.input_ids, len(self.tokenizer)):
                     raise ValueError(f"prompt token ids must be in [0, {len(self.tokenizer)})")
                 ids = torch.tensor(msg.input_ids, dtype=torch.int32)
                 results.append(UserMsg(uid=msg.uid, input_ids=ids, sampling_params=msg.sampling_params))
                 continue
-            prompt = self.render_prompt(msg)
-            # A jinja chat template owns every special token (HF's apply_chat_template
-            # tokenizes with add_special_tokens=False for the same reason): tokenizers
-            # that auto-add bos (muse-glimmer's, llama's) would otherwise double it --
-            # the template already rendered one. Raw-string prompts and the dsv4
-            # encoder path keep the default.
-            templated = isinstance(msg.text, list) and self._dsv4_encoder is None
-            input_ids: torch.Tensor = (  # type: ignore
-                self.tokenizer.encode(
-                    prompt, return_tensors="pt", add_special_tokens=not templated
-                )
-            )
-            input_ids = input_ids.view(-1).to(torch.int32)
+            else:
+                input_ids = self.encode_prompt(msg, self.render_prompt(msg))
             if msg.images:
                 if self.mm_processor is None:
                     raise ValueError("image input is not supported for this model")
@@ -151,6 +143,20 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]:
                 )
         return results
 
+    def encode_prompt(self, msg: TokenizeMsg, prompt: str) -> torch.Tensor:
+        """The tokenizer half of ``tokenize``: ``prompt`` (from ``render_prompt``) as 1-D
+        int32 ids, exposed so the frontend can count a request before committing a stream."""
+        # A jinja chat template owns every special token (HF's apply_chat_template
+        # tokenizes with add_special_tokens=False for the same reason): tokenizers
+        # that auto-add bos (muse-glimmer's, llama's) would otherwise double it --
+        # the template already rendered one. Raw-string prompts and the dsv4
+        # encoder path keep the default.
+        templated = isinstance(msg.text, list) and self._dsv4_encoder is None
+        input_ids: torch.Tensor = self.tokenizer.encode(  # type: ignore
+            prompt, return_tensors="pt", add_special_tokens=not templated
+        )
+        return input_ids.view(-1).to(torch.int32)
+
     def render_prompt(self, msg: TokenizeMsg) -> str:
         """The template/encoder half of ``tokenize``, exposed so the frontend can
         validate a request before committing an SSE stream. Sanitizes
diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py
index c9faf5a6c2..362c0ca1b4 100644
--- a/tests/server/test_anthropic_api.py
+++ b/tests/server/test_anthropic_api.py
@@ -753,6 +753,15 @@ def tokenize(self, msgs):
             for m in msgs
         ]
 
+    # The generation path's frontend pre-check renders and encodes separately.
+    def render_prompt(self, msg):
+        return str(msg.text)
+
+    def encode_prompt(self, msg, prompt):
+        import torch
+
+        return torch.zeros(len(prompt), dtype=torch.int32)
+
 
 def _count_client(manager=None, maintenance_state="serving"):
     fake = FakeState([])
diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py
index 91aaa874ff..26a815c8b2 100644
--- a/tests/server/test_logprobs_api.py
+++ b/tests/server/test_logprobs_api.py
@@ -7,6 +7,7 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 from freetoken.message import TokenizeMsg, UserReply
 from freetoken.server import openai_api
@@ -98,7 +99,12 @@ def __init__(
         self.replies = replies
         self.sent: TokenizeMsg | None = None
         if tokenizer is not None:
-            self.frontend_tokenizer = lambda: SimpleNamespace(tokenizer=tokenizer)
+            manager = SimpleNamespace(
+                tokenizer=tokenizer,
+                render_prompt=lambda msg: json.dumps(msg.text),
+                encode_prompt=lambda msg, prompt: torch.zeros(len(prompt), dtype=torch.int32),
+            )
+            self.frontend_tokenizer = lambda: manager
 
     def new_user(self) -> int:
         return 42
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index ff84745955..7ff7e96a99 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -931,3 +931,194 @@ async def chat_completion(req: ChatCompletionRequest):
     assert response.status_code == 200
     assert state.sent is not None
     assert state.sent.sampling_params.top_k == -1
+
+
+# --------------------------------------------------------------------------- #
+# Context-length pre-check (#749): every endpoint answers an over-long request with a
+# 400 before any StreamingResponse, counting with the frontend tokenizer.
+# --------------------------------------------------------------------------- #
+CONTEXT = 64
+
+
+class _CharManager:
+    """Frontend tokenizer stand-in: the prompt is the messages' text, one id per char."""
+
+    def render_prompt(self, msg):
+        if isinstance(msg.text, str):
+            return msg.text
+        return "".join(str(m.get("content") or "") for m in msg.text)
+
+    def encode_prompt(self, msg, prompt):
+        import torch
+
+        return torch.arange(len(prompt), dtype=torch.int32)
+
+
+def _ctx_state() -> FakeState:
+    state = FakeState([
+        UserReply(uid=42, incremental_output="ok", finished=True,
+                  prompt_tokens_delta=2, completion_tokens_delta=1),
+    ])
+    state.config.max_seq_len = CONTEXT
+    state.config.max_output_tokens = None
+    state.config.enable_cache_report = False
+    state.maintenance_state = "serving"
+    manager = _CharManager()
+    state.frontend_tokenizer = lambda: manager
+
+    async def stream_with_cancellation(gen, request, uid):
+        async for chunk in gen:
+            yield chunk
+
+    state.stream_with_cancellation = stream_with_cancellation
+    return state
+
+
+def _ctx_call(endpoint: str, text: str, stream: bool, max_tokens: int | None, state):
+    from freetoken.server.anthropic_api import handle_anthropic_messages
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+    from freetoken.server.responses_api import ResponsesRequest, handle_responses
+
+    limit = {} if max_tokens is None else {"max_tokens": max_tokens}
+    if endpoint == "chat":
+        req = ChatCompletionRequest(
+            model="m", messages=[{"role": "user", "content": text}], stream=stream, **limit
+        )
+        return run(handle_chat_completion(req, None, state, {}))
+    if endpoint == "completions":
+        req = CompletionRequest(model="m", prompt=text, stream=stream, **limit)
+        return run(handle_completion(req, None, state, {}))
+    if endpoint == "responses":
+        extra = {} if max_tokens is None else {"max_output_tokens": max_tokens}
+        req = ResponsesRequest(model="m", input=text, stream=stream, **extra)
+        return run(handle_responses(req, None, state, {}))
+    req = AnthropicMessagesRequest(
+        model="m", max_tokens=max_tokens or 8, stream=stream,
+        messages=[{"role": "user", "content": text}],
+    )
+    return run(handle_anthropic_messages(req, None, state, {}))
+
+
+def _error_message(endpoint: str, body: dict) -> tuple[str, str | None]:
+    err = body["error"]
+    if endpoint == "anthropic":
+        assert err["type"] == "invalid_request_error"
+        return err["message"], None
+    return err["message"], err.get("code")
+
+
+import pytest  # noqa: E402
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("endpoint", ["chat", "completions", "responses", "anthropic"])
+def test_an_overlong_prompt_is_a_400_before_any_stream(endpoint, stream):
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call(endpoint, "a" * (CONTEXT + 10), stream, None, state)
+    assert isinstance(response, JSONResponse), type(response)
+    assert response.status_code == 400
+    message, code = _error_message(endpoint, json.loads(response.body))
+    # Claude Code and OpenClaw match on this exact prefix.
+    assert message.startswith(f"prompt is too long: {CONTEXT + 10} tokens > {CONTEXT} maximum")
+    if endpoint != "anthropic":
+        assert code == "context_length_exceeded"
+    assert state.sent is None  # never reached the worker
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("endpoint", ["chat", "completions", "responses", "anthropic"])
+def test_an_explicit_max_tokens_that_does_not_fit_is_clamped_not_rejected(endpoint, stream):
+    # Agents always send a large max_tokens and run near a full context: only a prompt
+    # that does not fit by itself is a 400; the scheduler clamps the output budget.
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call(endpoint, "a" * 40, stream, CONTEXT * 4, state)
+    assert not isinstance(response, JSONResponse) or response.status_code == 200
+    assert state.sent is not None and state.sent.sampling_params.max_tokens == CONTEXT * 4
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("endpoint", ["chat", "completions", "responses", "anthropic"])
+def test_a_prompt_one_token_short_of_the_context_still_runs(endpoint, stream):
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call(endpoint, "a" * (CONTEXT - 1), stream, None, state)
+    assert not isinstance(response, JSONResponse) or response.status_code == 200
+    assert state.sent is not None
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("endpoint", ["chat", "completions", "responses", "anthropic"])
+def test_a_request_that_fits_carries_its_counted_ids_to_the_worker(endpoint, stream):
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call(endpoint, "a" * 40, stream, None, state)
+    assert not (isinstance(response, JSONResponse) and response.status_code != 200)
+    # The worker serves exactly the ids the frontend counted instead of tokenizing again.
+    assert state.sent is not None and state.sent.input_ids is not None
+    assert state.sent.input_ids.numel() == 40
+
+
+@pytest.mark.parametrize("endpoint", ["chat", "completions", "responses"])
+def test_an_omitted_max_tokens_is_still_clamped_not_rejected(endpoint):
+    # The server's default output budget also passes through; the scheduler clamps it.
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call(endpoint, "a" * 40, False, None, state)
+    assert not isinstance(response, JSONResponse) or response.status_code == 200
+    assert state.sent is not None and state.sent.sampling_params.max_tokens > CONTEXT - 40
+
+
+def test_the_tokenizer_worker_serves_prepared_ids_without_retokenizing():
+    import torch
+
+    from freetoken.core import SamplingParams
+    from freetoken.tokenizer.tokenize import TokenizeManager
+
+    def boom(*args, **kwargs):
+        raise AssertionError("re-tokenized a prepared prompt")
+
+    manager = TokenizeManager(SimpleNamespace(chat_template="t", encode=boom, apply_chat_template=boom))
+    ids = torch.tensor([5, 6, 7], dtype=torch.int32)
+    msg = TokenizeMsg(uid=1, text=[{"role": "user", "content": "x"}],
+                      sampling_params=SamplingParams(), input_ids=ids)
+    (user_msg,) = manager.tokenize([msg])
+    assert user_msg.input_ids.tolist() == [5, 6, 7]
+
+
+def test_a_prompt_that_fills_the_context_exactly_is_a_400():
+    # Same condition as the scheduler: no room for one output token.
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    response = _ctx_call("chat", "a" * CONTEXT, True, None, state)
+    assert isinstance(response, JSONResponse) and response.status_code == 400
+    assert json.loads(response.body)["error"]["code"] == "context_length_exceeded"
+    assert state.sent is None
+
+
+@pytest.mark.parametrize("stream", [False, True])
+def test_an_overlong_token_id_prompt_is_a_400_counted_by_its_length(stream):
+    from fastapi.responses import JSONResponse
+
+    state = _ctx_state()
+    req = CompletionRequest(model="m", prompt=[1] * (CONTEXT + 3), stream=stream)
+    response = run(handle_completion(req, None, state, {}))
+    assert isinstance(response, JSONResponse) and response.status_code == 400
+    body = json.loads(response.body)["error"]
+    assert body["message"].startswith(f"prompt is too long: {CONTEXT + 3} tokens > {CONTEXT}")
+    assert body["code"] == "context_length_exceeded"
+    assert state.sent is None
+
+
+def test_a_token_id_prompt_that_fits_reaches_the_worker_as_its_ids():
+    state = _ctx_state()
+    req = CompletionRequest(model="m", prompt=[5, 6, 7])
+    run(handle_completion(req, None, state, {}))
+    assert state.sent is not None and state.sent.input_ids == [5, 6, 7]
diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py
index 19140d8b72..0335b8d089 100644
--- a/tests/server/test_openai_extras.py
+++ b/tests/server/test_openai_extras.py
@@ -9,6 +9,7 @@
 from types import SimpleNamespace
 
 import pydantic
+import torch
 import pytest
 from freetoken.message import TokenizeMsg, UserReply
 from freetoken.server.api_models import (
@@ -52,6 +53,7 @@ def __init__(self, text=("hel", "lo"), reasoning_tokens=0, tokenizer=None):
             manager = SimpleNamespace(
                 tokenizer=tokenizer,
                 render_prompt=lambda msg: "RENDERED:" + json.dumps(msg.text),
+                encode_prompt=lambda msg, prompt: torch.zeros(len(prompt), dtype=torch.int32),
             )
             self.frontend_tokenizer = lambda: manager
 

From a36322e5dc63a0ac2cecfa361770df05564dafd1 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 10:51:53 -0400
Subject: [PATCH 241/249] fix(server): close a qwen3_coder call cut inside its
 parameter closer

A call cut by max_tokens inside </parameter> streamed {"city":"Lima and then
fell back to {} at close, while non-streaming returned {} (#710). The
qwen3_coder detector now finalizes an open call at end of stream, dropping the
partial closer and closing the JSON, and the one-shot parse replays that
finalize, so both modes return the same valid arguments with the value kept.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 .../freetoken/server/function_call_parser.py  | 48 +++++++++++++++++++
 tests/server/test_function_call_parser.py     | 37 ++++++++++++++
 tests/server/test_streaming_model_matrix.py   | 20 ++++++++
 3 files changed, 105 insertions(+)

diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py
index e45acea77b..0d2906c84c 100644
--- a/python/freetoken/server/function_call_parser.py
+++ b/python/freetoken/server/function_call_parser.py
@@ -1143,6 +1143,47 @@ def _update_prev() -> None:
 
         return StreamingParseResult(normal_text="".join(normal_parts), calls=calls)
 
+    def _ps_finalize_open_call(self) -> List[ToolCallItem]:
+        """The generation ended (max_tokens) inside an open call: close its arguments the
+        way the missing closers would have, so the streamed fragments end as valid JSON and
+        the one-shot replay returns the same arguments. A held suffix that could only grow
+        into the parameter closer is markup and is dropped, the value before it is kept.
+        Outside an open call (a cut header, a skipped unknown tool) nothing changes."""
+        mode = getattr(self, "_ps_mode", "idle")
+        if mode not in ("invoke", "pstr", "pbuf"):
+            return []
+        buf = self._buffer
+        ledger = self.streamed_args_for_tool[self.current_tool_id]
+        held = self._ends_with_partial_token(buf, self._ps_param_close) if mode != "invoke" else 0
+        value = buf[: len(buf) - held]
+        if mode == "pstr":
+            if not self._ps_lead_trimmed:
+                value = self._ps_trim_leading(value)
+            frag = self._json_escape_chunk(self._ps_trim_trailing(value)) + '"}'
+        elif mode == "pbuf" and held:
+            # Only the closer was cut: the typed value itself is complete.
+            raw = self._ps_trim_trailing(self._ps_trim_leading(value))
+            converted = self._ps_convert_value(self._ps_key, raw)
+            frag = (
+                self._ps_lead + json.dumps(self._ps_key, ensure_ascii=False) + ":"
+                + json.dumps(converted, ensure_ascii=False) + "}"
+            )
+        else:
+            frag = "}" if ledger else "{}"
+        index = self.current_tool_id
+        self.streamed_args_for_tool[index] += frag
+        try:
+            parsed = json.loads(self.streamed_args_for_tool[index])
+        except (json.JSONDecodeError, ValueError):
+            parsed = {}
+        self.prev_tool_call_arr[index]["arguments"] = parsed if isinstance(parsed, dict) else {}
+        self.current_tool_id += 1
+        while len(self.streamed_args_for_tool) <= self.current_tool_id:
+            self.streamed_args_for_tool.append("")
+        self._buffer = ""
+        self._ps_reset()
+        return [ToolCallItem(tool_index=index, name=None, parameters=frag)]
+
     def finish_streaming(self) -> str:
         residual, self._buffer = self._buffer, ""
         mode = getattr(self, "_ps_mode", "idle")
@@ -2201,6 +2242,10 @@ def __init__(self):
     def has_tool_call(self, text: str) -> bool:
         return "<function=" in text or self.bot_token in text
 
+    def finalize_streaming(self) -> List[ToolCallItem]:
+        """End-of-stream hook (FunctionCallParser.finalize_stream): close a call that
+        max_tokens cut off mid-arguments."""
+        return self._ps_finalize_open_call()
 
     def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools: List[Tool]) -> Optional[str]:
         """Build the current argument JSON from a partial XML tool-call body."""
@@ -2263,6 +2308,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
                     opened[-1][2] += item.parameters
             if not step.normal_text and not step.calls:
                 break
+        for item in shadow.finalize_streaming():
+            if opened:
+                opened[-1][2] += item.parameters
         normal.append(shadow.finish_streaming())
 
         calls = []
diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py
index 3ba0fb0e9e..2726a7dcf0 100644
--- a/tests/server/test_function_call_parser.py
+++ b/tests/server/test_function_call_parser.py
@@ -394,6 +394,43 @@ def test_qwen3_coder_unclosed_function_tag_is_not_a_call():
     assert calls == [] and parser.recover_truncated_call() == []
 
 
+_CUT_TOOLS = [{"type": "function", "function": {
+    "name": "w", "parameters": {"type": "object", "properties": {"c": {"type": "string"}}}}}]
+_CUT_CALL = "<tool_call>\n<function=w>\n<parameter=c>\nLima\n</parameter>\n</function>\n</tool_call>"
+
+
+def _cut_both_ways(text):
+    one_shot = FunctionCallParser(_CUT_TOOLS, tool_call_parser="qwen3_coder").parse_non_stream(text)
+    parser = FunctionCallParser(_CUT_TOOLS, tool_call_parser="qwen3_coder")
+    texts, calls = _feed(parser, list(text))
+    calls += parser.finalize_stream()
+    streamed = "".join(c.parameters for c in calls if c.name is None)
+    texts.append(parser.finish_stream())
+    return one_shot, streamed, "".join(texts)
+
+
+@pytest.mark.parametrize("cut", range(len("Lima"), len("Lima\n</parameter>")))
+def test_qwen3_coder_call_cut_inside_the_parameter_close_keeps_its_value(cut):
+    # #710: max_tokens ends the call between the value and the ">" of </parameter>.
+    # Both modes must close the arguments as the same valid JSON, value kept and no
+    # markup in it or in the content.
+    start = _CUT_CALL.index("Lima")
+    one_shot, streamed, content = _cut_both_ways(_CUT_CALL[: start + cut])
+    assert [c.name for c in one_shot.calls] == ["w"]
+    assert json.loads(streamed) == {"c": "Lima"}
+    assert one_shot.calls[0].parameters == streamed
+    assert one_shot.normal_text == content == ""
+
+
+@pytest.mark.parametrize("cut", range(1, len("Lima")))
+def test_qwen3_coder_call_cut_inside_a_string_value_closes_it(cut):
+    # The streamed fragments already showed the partial value: one-shot agrees.
+    start = _CUT_CALL.index("Lima")
+    one_shot, streamed, _ = _cut_both_ways(_CUT_CALL[: start + cut])
+    assert json.loads(streamed) == {"c": "Lima"[:cut]}
+    assert one_shot.calls[0].parameters == streamed
+
+
 def test_streaming_support_flags():
     # Every registered detector is incremental-safe (buffered fallback remains as
     # the escape hatch for future formats, exercised via monkeypatch in
diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py
index bdd71b2e22..aeca9968f7 100644
--- a/tests/server/test_streaming_model_matrix.py
+++ b/tests/server/test_streaming_model_matrix.py
@@ -869,6 +869,26 @@ def test_qwen_truncated_inside_function_tag_is_not_recovered():
     assert (calls, finish) == ([], "length")
 
 
+@pytest.mark.parametrize("tail", ["\n", "\n<", "\n</", "\n</param", "\n</parameter"])
+def test_qwen_truncated_inside_parameter_close_keeps_the_value(tail):
+    # #710: a call cut by max_tokens after its value, inside </parameter>: the same valid
+    # arguments in both modes, and the streamed fragments add up to them.
+    pieces = ["<tool_call>", "\n<function=get_weather>\n<parameter=city>\n", "Lima", tail]
+    _, _, calls, finish = _qwen_parity(pieces, thinking=False, finish_reason="length")
+    assert (calls, finish) == ([("get_weather", {"city": "Lima"})], "length")
+    spec = GenSpec(messages=[], sampling_params=None, parser_tools=QWEN_TOOLS,
+                   template_tools=QWEN_TOOLS, chat_template_kwargs={"enable_thinking": False})
+    state = FakeState(pieces, tool_call_parser="qwen3_coder", reasoning_parser="qwen3",
+                      finish_reason="length")
+
+    async def stream():
+        return [ev async for ev in generate_events(42, spec, state)]
+
+    events = asyncio.run(stream())
+    fragments = "".join(ev.fragment for ev in events if isinstance(ev, ToolCallArgsDelta))
+    assert fragments == _calls(events)[0].parameters
+
+
 @pytest.mark.parametrize("pieces", [
     ["The parser looks for ", "<tool_call>", " tags", " and splits", " them."],
     ["Wrap calls in ", "<tool_call>", "...", "</tool_call>", " tags, then answer."],

From e67bd429d9daf3d420b82d249374b1c8697e39b8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 11:12:37 -0400
Subject: [PATCH 242/249] fix(checkpoint): upcast fp16 nvfp4 globals under
 inference mode on the bank reader threads

The engine allocates the host banks under torch.inference_mode(), which the reader pool's threads do not inherit, so the in-place fp32 copy raised and the server refused to start. The test now loads the banks under inference mode, as the engine does.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/checkpoint/ftw.py        | 6 ++++--
 tests/checkpoint/test_ftw_fp16_globals.py | 4 +++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index 23d5278441..c7068a587c 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -1025,7 +1025,8 @@ def _read_row(job):
                     buf = _transient_buffer(win_len)
                     reader.read_into(memoryview(buf), window, workers=workers, chunk=chunk)
                     raw = torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + layer_bytes)[head_pad:]
-                    bank.tensor.copy_(raw.view(torch.float16).view(shape))
+                    with torch.inference_mode():  # the banks are inference tensors; this thread does not inherit the mode
+                        bank.tensor.copy_(raw.view(torch.float16).view(shape))
                 pins.submit(bank, residency[layer_id])
                 bar.update(layer_bytes)
 
@@ -1035,7 +1036,8 @@ def _read_layer(job):
                     buf = _transient_buffer(entry["nbytes"])
                     reader.read_into(memoryview(buf), entry, workers=workers, chunk=chunk)
                     raw = torch.frombuffer(buf, dtype=torch.float16, count=entry["nbytes"] // 2)
-                    bank.tensor.copy_(raw.view(bank.tensor.shape))
+                    with torch.inference_mode():  # as in _read_row
+                        bank.tensor.copy_(raw.view(bank.tensor.shape))
                 else:
                     reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk)
                 pins.submit(bank, residency[layer_id])
diff --git a/tests/checkpoint/test_ftw_fp16_globals.py b/tests/checkpoint/test_ftw_fp16_globals.py
index 934d7570df..71aafac074 100644
--- a/tests/checkpoint/test_ftw_fp16_globals.py
+++ b/tests/checkpoint/test_ftw_fp16_globals.py
@@ -43,7 +43,9 @@ def test_fp16_globals_load_as_fp32(tmp_path, monkeypatch, caplog):
             caplog.clear()
             FTWReader(str(path)).close()
             dense = dict(iter_ftw_weights(str(path)))
-            banks = load_ftw_banks(str(path), num_layers=1)
+            # the engine loads the banks under inference mode, which the bank reader threads do not inherit
+            with torch.inference_mode():
+                banks = load_ftw_banks(str(path), num_layers=1)
         assert sum("fp16" in r.getMessage() for r in caplog.records) == 1
 
         g = dense["model.layers.0.mlp.shared_expert.gate_proj.weight_global"]

From fc30523843d79f9cf9c3d7ce497c31d33fd4810f Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 15:54:11 -0400
Subject: [PATCH 243/249] feat(scheduler): add the --logprob-read-isolate flag
 and the is_logprob_read predicate

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 docs/cli.md                                  |  1 +
 python/freetoken/core.py                     |  5 ++++
 python/freetoken/scheduler/cache.py          |  6 ++---
 python/freetoken/scheduler/config.py         |  2 ++
 python/freetoken/scheduler/prefill.py        |  5 ++++
 python/freetoken/scheduler/scheduler.py      |  2 ++
 python/freetoken/server/args.py              | 10 ++++++++
 tests/scheduler/test_hybrid_cache_manager.py | 25 ++++++++++++++++++++
 8 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index b5af7aa2d3..d9ec62a0d0 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -75,6 +75,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 |
 | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` |
 | `--logprob-read-snapshot-donate` | all | hybrid GDN models: snapshots a finished `max_tokens=1` + `logprobs` request donates to the prefix cache: `all` (x64 prefill boundary and final state), `boundary` (only the former) or `none` |
+| `--logprob-read-isolate` | off | a `max_tokens=1` + `logprobs` request prefills in a pass of its own and neither resumes from nor stores prefix-cache entries, so its logprobs do not depend on other traffic; costs read throughput under concurrency and prefix reuse for reads |
 | `--kv-cache-dtype` | bf16 | `bf16`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |
 
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 541c1de761..a375dc9dcf 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -71,6 +71,11 @@ def client_token_ids(self) -> list[tuple[str, list[int]]]:
         ]
         return [(name, ids) for name, ids in fields if ids]
 
+    @property
+    def is_logprob_read(self) -> bool:
+        """A scoring read (e.g. AnyJev label logprobs): one token whose logprobs are the answer."""
+        return bool(self.logprobs) and self.max_tokens == 1
+
     @property
     def is_greedy(self) -> bool:
         return self.temperature <= 0.0 or self.top_k == 1
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index c29d578b7b..384e2e9d87 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -35,7 +35,7 @@ def _swa_eviction_interval() -> int:
 class CacheManager:
     def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str,
                  linear_state_pool=None, swa_pool=None, sliding_window_size=None,
-                 logprob_read_snapshot_donate: str = "all"):
+                 logprob_read_snapshot_donate: str = "all", isolate_logprob_reads: bool = False):
         # The `_free_slots` follows a page-aligned manner. For example, if page_size = 2,
         # the `_free_slots` may look like [0, 2, 4, 6, ...], and each slot represents a page.
         device = page_table.device
@@ -54,6 +54,7 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ
         if logprob_read_snapshot_donate not in LOGPROB_READ_SNAPSHOT_DONATE:
             raise ValueError(f"unknown snapshot donate policy {logprob_read_snapshot_donate!r}")
         self.logprob_read_snapshot_donate = logprob_read_snapshot_donate
+        self.isolate_logprob_reads = isolate_logprob_reads
         self.is_swa = type == "swa_radix"
         # swa_paged: this SWA model drives the global-paged swa pool -- true for BOTH the naive
         # (NaivePrefixCache, no reuse) and radix (SWARadixCache) paths. Gates the swa slot
@@ -394,10 +395,9 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             # is this request's dup to free. The frozen slot is consumed either way (taken by
             # the tree or freed here) and both ping-pong refs are dropped before
             # _free_req_slots so nothing double-frees.
-            sp = req.sampling_params
             donate = (
                 self.logprob_read_snapshot_donate
-                if sp.logprobs and sp.max_tokens == 1
+                if req.sampling_params.is_logprob_read
                 else "all"
             )
             free_upto = old_handle.cached_len
diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py
index 75502210ca..18981135e2 100644
--- a/python/freetoken/scheduler/config.py
+++ b/python/freetoken/scheduler/config.py
@@ -20,6 +20,8 @@ class SchedulerConfig(EngineConfig):
     special_token_ckpt: bool = False
     # GDN snapshots a finished logprob read donates (CacheManager): all | boundary | none
     logprob_read_snapshot_donate: str = "all"
+    # logprob reads prefill alone and bypass the prefix cache, so their logprobs do not depend on other traffic
+    logprob_read_isolate: bool = False
 
     # networking config
     _unique_suffix: str = field(default_factory=_get_pid_suffix)
diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index 9019604707..25437060e3 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -49,6 +49,8 @@ class PrefillAdder:
     encoder_cache: EncoderCache | None = None
     # end a chunk before an image it would cut; only models whose image spans attend in both directions need it
     keep_images_whole: bool = False
+    # a logprob read (max_tokens=1 + logprobs) neither resumes from nor feeds the prefix cache
+    isolate_logprob_reads: bool = False
     # the whole budget of this pass; token_budget shrinks as requests are admitted
     pass_budget: int = 0
     # SWA-pool tokens charged to reqs admitted so far this pass. Mirrors reserved_size: swa is
@@ -282,6 +284,8 @@ class PrefillManager:
     decode_manager: DecodeManager
     encoder_cache: EncoderCache | None = None
     keep_images_whole: bool = False
+    # a logprob read prefills in a pass of its own (see PrefillAdder for the cache side)
+    isolate_logprob_reads: bool = False
     pending_list: List[PendingReq] = field(default_factory=list)
 
     def add_one_req(self, req: UserMsg) -> None:
@@ -308,6 +312,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
             table_manager=self.table_manager,
             encoder_cache=self.encoder_cache,
             keep_images_whole=self.keep_images_whole,
+            isolate_logprob_reads=self.isolate_logprob_reads,
         )
         reqs: List[Req] = []
         chunked_list: List[PendingReq] = []
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index 78d2edf25e..d2fb4c3c8f 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -100,6 +100,7 @@ def __init__(self, config: SchedulerConfig):
                 None,
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
             logprob_read_snapshot_donate=getattr(config, "logprob_read_snapshot_donate", "all"),
+            isolate_logprob_reads=getattr(config, "logprob_read_isolate", False),
         )
         self.decode_manager = DecodeManager(config.page_size)
         # Bound how long a run of prefill steps may starve in-flight decodes. A long
@@ -117,6 +118,7 @@ def __init__(self, config: SchedulerConfig):
             self.decode_manager,
             encoder_cache=self.engine.encoder_cache,
             keep_images_whole=self._bidirectional_mm,
+            isolate_logprob_reads=getattr(config, "logprob_read_isolate", False),
         )
 
         # some alias for easy access
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 4fe0a7a8e5..6ed30ff073 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -556,6 +556,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "neither, so a burst of reads cannot evict other conversations' snapshots.",
     )
 
+    parser.add_argument(
+        "--logprob-read-isolate",
+        action="store_true",
+        default=ServerArgs.logprob_read_isolate,
+        help="Make a logprob read's (max_tokens=1 with logprobs, e.g. AnyJev label scoring) "
+        "logprobs independent of other traffic: it prefills in a pass of its own and neither "
+        "resumes from nor stores prefix-cache entries. Costs read throughput under concurrency "
+        "and all prefix reuse for reads; other requests are unchanged.",
+    )
+
     parser.add_argument(
         "--text-model-only",
         action="store_true",
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index 6b4f164b52..c5373959bc 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -344,3 +344,28 @@ def test_snapshot_donate_policy_is_a_server_arg():
         assert parse_args(["--model", "/models/anon"])[0].logprob_read_snapshot_donate == "all"
         args, _ = parse_args(["--model", "/models/anon", "--logprob-read-snapshot-donate", "none"])
     assert args.logprob_read_snapshot_donate == "none"
+
+
+@pytest.mark.parametrize(
+    ("sp", "read"),
+    [
+        (READ, True),
+        (SamplingParams(max_tokens=1, logprobs=True), True),
+        (SamplingParams(max_tokens=1), False),
+        (SamplingParams(max_tokens=2, logprobs=True, top_logprobs=4), False),
+    ],
+)
+def test_is_logprob_read(sp, read):
+    assert sp.is_logprob_read is read
+
+
+def test_logprob_read_isolate_is_a_server_arg():
+    from unittest.mock import patch
+
+    from freetoken.server.args import parse_args
+
+    config = SimpleNamespace(to_dict=lambda: {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"})
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):
+        assert parse_args(["--model", "/models/anon"])[0].logprob_read_isolate is False
+        args, _ = parse_args(["--model", "/models/anon", "--logprob-read-isolate"])
+    assert args.logprob_read_isolate is True

From fed01568cffe56f5594a61e388d4abe7f6bed896 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 15:55:07 -0400
Subject: [PATCH 244/249] feat(scheduler): prefill an isolated logprob read in
 a pass of its own

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/scheduler/prefill.py         |  6 ++
 .../test_scheduler_chunked_prefill.py         | 59 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index 25437060e3..e367d1e935 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -324,6 +324,10 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
         log_cached_tokens = 0
         for pending_req in self.pending_list:
             is_continuation = pending_req.chunked_req is not None
+            # break, not continue: the slicing below assumes the admitted requests are a prefix
+            read = self.isolate_logprob_reads and pending_req.sampling_params.is_logprob_read
+            if read and reqs:
+                break  # a read never shares a pass
             if req := adder.try_add_one(pending_req):
                 pending_req.chunked_req = None
                 if isinstance(req, ChunkedReq):
@@ -346,6 +350,8 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
                 log_new_tokens += req.extend_len
                 if not is_continuation:
                     log_cached_tokens += req.cache_handle.cached_len
+                if read:
+                    break  # nobody joins a read's pass
             else:
                 break  # We cannot add more requests
         if len(reqs) == 0:
diff --git a/tests/scheduler/test_scheduler_chunked_prefill.py b/tests/scheduler/test_scheduler_chunked_prefill.py
index a4c7e0659a..f0fe3a7a1b 100644
--- a/tests/scheduler/test_scheduler_chunked_prefill.py
+++ b/tests/scheduler/test_scheduler_chunked_prefill.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 
 CHUNK = 8
@@ -237,3 +238,61 @@ def test_the_sliding_window_pool_cap_cannot_recut_an_image():
         for r in batch.reqs:
             r.complete_one()
     assert ends[0] == 4900 and ends[-1] == 6000
+
+
+def _passes(pm, cm, prompts, budget, isolate=True):
+    """Drain ``prompts`` ([(uid, prompt_len, read)]) and return each pass's (uid, extend_len) list."""
+    from freetoken.core import SamplingParams
+    from freetoken.scheduler.utils import PendingReq
+
+    pm.isolate_logprob_reads = isolate
+    pm.pending_list = [
+        PendingReq(uid, torch.arange(n, dtype=torch.int32),
+                   SamplingParams(max_tokens=1, logprobs=True) if read else SamplingParams(max_tokens=4))
+        for uid, n, read in prompts
+    ]
+    passes = []
+    while pm.runnable:
+        batch = pm.schedule_next_batch(budget)
+        passes.append([(r.uid, r.extend_len) for r in batch.reqs])
+        cm.allocate_paged(batch.reqs)
+        for r in batch.reqs:
+            r.complete_one()
+    return passes
+
+
+READ, CHAT = True, False
+
+
+@pytest.mark.parametrize(
+    ("prompts", "expected"),
+    [
+        ([(1, 5, READ), (2, 6, READ)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, CHAT), (2, 6, READ)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, READ), (2, 6, CHAT)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, CHAT), (2, 6, CHAT)], [[(1, 5), (2, 6)]]),
+        ([(1, 5, READ), (2, 6, READ), (3, 4, CHAT), (4, 3, CHAT)], [[(1, 5)], [(2, 6)], [(3, 4), (4, 3)]]),
+    ],
+)
+def test_isolated_logprob_read_prefills_alone(prompts, expected):
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    assert _passes(pm, cm, prompts, budget=32) == expected
+
+
+def test_isolation_off_keeps_reads_batched():
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, 5, READ), (2, 6, READ), (3, 4, CHAT)]
+    assert _passes(pm, cm, prompts, budget=32, isolate=False) == [[(1, 5), (2, 6), (3, 4)]]
+
+
+def test_isolated_chunked_read_continuations_stay_alone():
+    """Every chunk of a read runs alone and on the same cuts as a solo run (budget multiples)."""
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, 2 * CHUNK + 3, READ), (2, 4, CHAT)]
+    assert _passes(pm, cm, prompts, budget=CHUNK) == [[(1, CHUNK)], [(1, CHUNK)], [(1, 3)], [(2, 4)]]
+
+
+def test_isolated_read_waits_behind_a_chunked_chat_prompt():
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, CHUNK + 2, CHAT), (2, 4, READ)]
+    assert _passes(pm, cm, prompts, budget=CHUNK) == [[(1, CHUNK)], [(1, 2)], [(2, 4)]]

From cc13fecce538bc7f38efa459b0f10cc313dfa588 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 15:56:42 -0400
Subject: [PATCH 245/249] feat(scheduler): keep isolated logprob reads out of
 the prefix cache

An isolated read matches with max_len 0 (no host-tier lookup) and donates nothing.
The prefill-commit branch also skips the snapshot insert for a non-donating read;
the scheduler only reaches it for requests that keep decoding, never a read.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/scheduler/cache.py          | 17 +++--
 python/freetoken/scheduler/prefill.py        |  6 +-
 tests/scheduler/test_hybrid_cache_manager.py | 68 +++++++++++++++++++-
 3 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 384e2e9d87..a9ef9dca11 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -146,7 +146,7 @@ def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult:
             # (KV bytes + GDN snapshot). If the tree match fell short, rehydrate the deeper
             # host entry into fresh pages, re-insert it into the tree and re-match.
             tier = self.host_tier
-            if tier is not None and tier.has_kv:
+            if tier is not None and tier.has_kv and len(ids) > 0:
                 hit = tier.lookup(ids, self.page_size)
                 if (hit is not None and hit.kv_pages is not None and hit.gdn_slot is not None
                         and hit.host_len > m.cached_len):
@@ -375,6 +375,11 @@ def cache_req(self, req: Req, *, finished: bool) -> None:
             req.cache_handle = new_handle
             self.lock(new_handle)
 
+    def _donate_policy(self, req: Req) -> str:
+        if not req.sampling_params.is_logprob_read:
+            return "all"
+        return "none" if self.isolate_logprob_reads else self.logprob_read_snapshot_donate
+
     def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
         """Hybrid (GDN) cache_req: commit KV like radix AND manage the GDN state snapshot.
         Prefill chunk commit: DONATE the frozen ping-pong slot (the snapshot the forward wrote
@@ -395,11 +400,7 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             # is this request's dup to free. The frozen slot is consumed either way (taken by
             # the tree or freed here) and both ping-pong refs are dropped before
             # _free_req_slots so nothing double-frees.
-            donate = (
-                self.logprob_read_snapshot_donate
-                if req.sampling_params.is_logprob_read
-                else "all"
-            )
+            donate = self._donate_policy(req)
             free_upto = old_handle.cached_len
             L = req.mamba_last_track_seqlen
             if (
@@ -439,6 +440,10 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             return
 
         # Prefill chunk commit: donate the frozen snapshot at the tracked ×64 boundary.
+        # A non-donating read keeps its own pages and frozen slot until the finish frees them.
+        if self._donate_policy(req) == "none":
+            req.mamba_last_track_seqlen = None
+            return
         L = req.mamba_last_track_seqlen
         if L is None:
             return  # no ×64 boundary crossed this chunk; req keeps its pages (committed later)
diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index e367d1e935..7536e7a786 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -75,7 +75,11 @@ def _try_allocate_one(self, req: PendingReq):
 
         # TODO: consider host cache match case (it needs the same replay cap)
         replay = self.cache_manager.prefix_replay_tokens
-        mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None)
+        if self.isolate_logprob_reads and req.sampling_params.is_logprob_read:
+            max_len = 0  # a resumed prefill differs numerically from a cold one
+        else:
+            max_len = max(0, req.input_len - replay) if replay else None
+        mr = self.cache_manager.match_req(req, max_len=max_len)
         handle = mr.cuda_handle
         cached_len = handle.cached_len
         # TODO: better estimate policy
diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py
index c5373959bc..d24de93d5c 100644
--- a/tests/scheduler/test_hybrid_cache_manager.py
+++ b/tests/scheduler/test_hybrid_cache_manager.py
@@ -256,11 +256,10 @@ def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch):
     raise SystemExit(pytest.main([__file__]))
 
 
-def _finish_read(policy: str, sp: SamplingParams, *, warm_prefix: bool = False):
+def _finish_read(policy: str, sp: SamplingParams, *, warm_prefix: bool = False, num_pages: int = 256):
     """A 130-token prompt finishing on its prefill forward, with a frozen x64 snapshot at 128
     (the deepest boundary strictly inside it) and its live state at 130. ``warm_prefix`` first
     donates a 64-token prefix so the read resumes from a locked prefix-cache hit."""
-    num_pages = 256
     pool = _pool()
     pt = torch.zeros(4, 256, dtype=torch.int32)
     cm = CacheManager(num_pages, 1, pt, "hybrid_radix", linear_state_pool=pool,
@@ -369,3 +368,68 @@ def test_logprob_read_isolate_is_a_server_arg():
         assert parse_args(["--model", "/models/anon"])[0].logprob_read_isolate is False
         args, _ = parse_args(["--model", "/models/anon", "--logprob-read-isolate"])
     assert args.logprob_read_isolate is True
+
+
+def _assert_no_leak(cm, num_pages=256):
+    pc, pool = cm.prefix_cache, cm.linear_state_pool
+    assert len(cm.free_slots) + pc.full_evictable + pc.full_protected == num_pages
+    assert pc.full_protected == 0 and pc.mamba_protected == 0
+    assert pool.num_free_slots + pc.mamba_evictable_size == pool.num_slots - 1
+    cm.check_integrity()
+
+
+@pytest.mark.parametrize("isolate", [False, True])
+def test_isolated_logprob_read_ignores_and_keeps_out_of_the_prefix_cache(isolate):
+    """An identical earlier read donated under "all"; an isolated read still prefills cold and stores nothing."""
+    from freetoken.scheduler.decode import DecodeManager
+    from freetoken.scheduler.prefill import PrefillManager
+    from freetoken.scheduler.table import TableManager
+    from freetoken.scheduler.utils import PendingReq
+
+    cm, prompt = _finish_read("all", READ, num_pages=512)  # room for both reads: no eviction
+    cm.isolate_logprob_reads = isolate
+    tree = (cm.prefix_cache.full_evictable, cm.prefix_cache.mamba_evictable_size)
+    pm = PrefillManager(cm, TableManager(max_running_reqs=2, page_table=cm.page_table),
+                        DecodeManager(page_size=1), isolate_logprob_reads=isolate)
+    pm.pending_list = [PendingReq(1, torch.tensor(prompt, dtype=torch.int32), READ)]
+    if isolate:  # nor may the host tier rehydrate a prefix for it
+        cm.host_tier = SimpleNamespace(has_kv=True, lookup=lambda *_: pytest.fail("host tier lookup"))
+    batch = pm.schedule_next_batch(4096)
+    (req,) = batch.reqs
+    assert req.cache_handle.cached_len == batch.log_cached_tokens == (0 if isolate else 128)
+
+    cm.allocate_paged(batch.reqs)
+    req.complete_one()
+    if isolate:  # a cold 130-token extend tracks the x64 boundary at 128 (build_fla_metadata)
+        req.mamba_next_track_idx, req.mamba_last_track_seqlen = 1, 128
+    cm.cache_req(req, finished=True)
+    if isolate:
+        assert (cm.prefix_cache.full_evictable, cm.prefix_cache.mamba_evictable_size) == tree
+    _assert_no_leak(cm, num_pages=512)
+
+
+@pytest.mark.parametrize(("policy", "isolate"), [("none", False), ("all", True), ("boundary", True)])
+def test_non_donating_read_chunk_commit_leaves_the_tree_alone(policy, isolate):
+    """The prefill-commit branch donated the x64 snapshot whatever the read policy said."""
+    num_pages, pool = 256, _pool()
+    pt = torch.zeros(4, 256, dtype=torch.int32)
+    cm = CacheManager(num_pages, 1, pt, "hybrid_radix", linear_state_pool=pool,
+                      logprob_read_snapshot_donate=policy, isolate_logprob_reads=isolate)
+    prompt = list(range(1, 131))
+    mr = cm.match_req(_pend(prompt))
+    pt[0, :130] = cm._allocate(130)
+    row = pt[0, :130].clone()
+    req = Req(input_ids=torch.tensor(prompt + [7], dtype=torch.int32), table_idx=0, cached_len=130,
+              output_len=1, uid=0, sampling_params=READ, cache_handle=mr.cuda_handle)
+    req.linear_slot_idx, req.mamba_ping_pong = pool.alloc(1)[0], tuple(pool.alloc(2))
+    req.mamba_next_track_idx, req.mamba_last_track_seqlen = 1, 128
+    cm.lock(mr.cuda_handle)
+
+    cm.cache_req(req, finished=False)
+    assert cm.prefix_cache.full_evictable + cm.prefix_cache.full_protected == 0
+    assert cm.prefix_cache.mamba_evictable_size + cm.prefix_cache.mamba_protected == 0
+    assert torch.equal(pt[0, :130], row)  # the request still owns its pages
+
+    cm.cache_req(req, finished=True)
+    assert len(cm.free_slots) == num_pages
+    _assert_no_leak(cm)

From 77874d22da600c5fb2220707b1333b63bc98fa92 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 20:02:50 -0400
Subject: [PATCH 246/249] perf(engine): make the capture-time cuda graph upload
 switchable with FREETOKEN_GRAPH_UPLOAD

On Qwen3.8-Flash / RTX 5080 a decode graph uploaded with cuGraphUpload at capture (#697) replays
with ~0.5 us instead of ~0.2 us between kernels: +0.19 ms/step (-1% decode), +0.43 ms with
--moe-collect-stats. FREETOKEN_GRAPH_UPLOAD=0 keeps the driver's first-launch upload; default on.

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/engine/graph.py  |  3 ++-
 python/freetoken/env.py           |  3 +++
 tests/engine/test_graph_upload.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 tests/engine/test_graph_upload.py

diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index a995316ddc..0d79bad93f 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -10,6 +10,7 @@
 import torch
 from freetoken.core import Batch, Req, get_global_ctx
 from freetoken.distributed import get_tp_info
+from freetoken.env import ENV
 from freetoken.utils import init_logger, mem_GB
 from freetoken.utils.progress import emit_progress
 from tqdm import tqdm
@@ -120,7 +121,7 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non
 
     Otherwise the first replay pushes every node behind the graph's own stream WAIT, and a large graph then blocks cudaGraphLaunch until the WAIT is released. PLE wait-sync releases it only after the launch returns, so that deadlocks.
     """
-    upload = _cu_graph_upload()
+    upload = _cu_graph_upload() if ENV.GRAPH_UPLOAD else None
     if upload is None:
         return
     rc = upload(graph.raw_cuda_graph_exec(), stream.cuda_stream)
diff --git a/python/freetoken/env.py b/python/freetoken/env.py
index 9a06b1c029..03e9388024 100644
--- a/python/freetoken/env.py
+++ b/python/freetoken/env.py
@@ -77,6 +77,9 @@ class EnvClassSingleton:
     MAMBA_SSM_DTYPE = EnvStr("float32")
     # cap the longest prefill warmup forward; 0 = the engine's own chunk cap
     WARMUP_MAX_LEN = EnvInt(0)
+    # cuGraphUpload each decode graph at capture (#697); off keeps the first-launch upload, which
+    # replays ~0.2 ms/step faster on Qwen3.8-Flash / RTX 5080 but can block the first launch on a stream WAIT
+    GRAPH_UPLOAD = EnvBool(True)
     # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert
     MOE_COPY_OVERLAP = EnvBool(False)
     # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy)
diff --git a/tests/engine/test_graph_upload.py b/tests/engine/test_graph_upload.py
new file mode 100644
index 0000000000..de61152162
--- /dev/null
+++ b/tests/engine/test_graph_upload.py
@@ -0,0 +1,28 @@
+import pytest
+
+import freetoken.engine.graph as graph_mod
+from freetoken.env import ENV
+
+
+class _Graph:
+    def raw_cuda_graph_exec(self):
+        return 1234
+
+
+class _Stream:
+    cuda_stream = 5678
+    synced = False
+
+    def synchronize(self):
+        self.synced = True
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+def test_graph_upload_follows_env(monkeypatch, enabled):
+    calls = []
+    monkeypatch.setattr(graph_mod, "_cu_graph_upload", lambda: lambda exec_, stream: calls.append((exec_, stream)) or 0)
+    monkeypatch.setattr(ENV.GRAPH_UPLOAD, "value", enabled)
+    stream = _Stream()
+    graph_mod._upload_graph(_Graph(), stream)
+    assert calls == ([(1234, 5678)] if enabled else [])
+    assert stream.synced is enabled

From aab6bdd83e6705577d5c3e3e7260c0f461d7b8f5 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 21 Oct 2026 12:12:49 -0400
Subject: [PATCH 247/249] perf(qsa): score prefill rows in tiles of 16 query
 rows per program

The prefill indexer score ran one query row per program and padded the 4
index heads to a 16-wide dot, so every row re-read all visible keys and 3/4
of the tensor-core work was padding. qsa_mqa_paged_rows gives each program
16 rows of one request (a host-built tile map cut at the cu_seqlens request
boundaries) and does one [BLOCK_N, 16] dot per head, summing the heads in
the one-row kernel's order, so every score is bitwise the old one. Decode
(graph-captured) keeps the one-row kernel.

Microbench, one 4096-row chunk at 249K context: 18.5 -> 3.2 ms per layer,
no spills (134 regs).

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/attention/qsa_sparse.py      |  53 ++++-
 .../freetoken/kernel/triton/qsa/__init__.py   |   4 +-
 python/freetoken/kernel/triton/qsa/score.py   | 221 ++++++++++++++++++
 tests/kernels/test_qsa_score_rows.py          | 136 +++++++++++
 4 files changed, 402 insertions(+), 12 deletions(-)
 create mode 100644 tests/kernels/test_qsa_score_rows.py

diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index 7d20092efc..9d2a3bc8b0 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -50,6 +50,7 @@
 # Block-score transient budget (vLLM's number): the fp32 [rows, n_blocks] logits tile is
 # 256 KB per row at a 1M-token context, so a long prefill must be scored in row chunks.
 _LOGITS_WORKSPACE_BYTES = 128 << 20
+_SCORE_ROWS = 16  # prefill query rows per score program (qsa_mqa_paged_rows)
 
 
 TORCH_TOPK_ENV = "FREETOKEN_QSA_TORCH_TOPK"
@@ -96,6 +97,8 @@ class QSASparseMetadata(BaseAttnMetadata):
     # fetch query so no mid-forward eviction can take them) and the physical out_loc.
     write_pages:      torch.Tensor | None = None
     out_loc_gpu:      torch.Tensor | None = None
+    # Prefill scoring: row-tile maps per (start, end) row chunk, built by the first QSA layer.
+    score_tiles:      dict | None = None
     # fmt: on
 
     def get_last_indices(self, bs: int) -> torch.Tensor:
@@ -464,6 +467,8 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor:
             expand_qsa_block_indices,
             qsa_index_norm_rope,
             qsa_mqa_paged,
+            qsa_mqa_paged_rows,
+            qsa_score_tile_map,
         )
 
         rows = index.q.shape[0]
@@ -488,22 +493,48 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor:
         columns = md.block_table.shape[1] * self.cmp_page_size
         indices = self._scratch("indices", rows, self.select_width, dtype=torch.int32)
         rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1))
+        row_tiles = not md.is_decode and self.index_heads == 4
         for start in range(0, rows, rows_per_chunk):
             end = min(start + rows_per_chunk, rows)
             chunk = slice(start, end)
             logits = self._scratch("logits", end - start, columns, dtype=torch.float32)
             visible = self._scratch("visible", end - start, dtype=torch.int32)
-            qsa_mqa_paged(
-                q_index[chunk],
-                cmp_pages,
-                md.block_table,
-                md.token_to_req[chunk],
-                positions[chunk],
-                md.seq_lens,
-                self.ratio,
-                logits,
-                visible,
-            )
+            if row_tiles:
+                # Prefill: several query rows of one request per program (decode, which a
+                # graph may capture, keeps the one-row kernel). Same scores bit for bit.
+                if md.score_tiles is None:
+                    md.score_tiles = {}
+                tile_map = md.score_tiles.get((start, end))
+                if tile_map is None:
+                    tile_map = qsa_score_tile_map(
+                        md.qo_indptr_cpu, start, end, _SCORE_ROWS, self.device
+                    )
+                    md.score_tiles[(start, end)] = tile_map
+                qsa_mqa_paged_rows(
+                    q_index[chunk],
+                    cmp_pages,
+                    md.block_table,
+                    tile_map,
+                    md.token_to_req[chunk],
+                    positions[chunk],
+                    md.seq_lens,
+                    self.ratio,
+                    logits,
+                    visible,
+                    rows_per_tile=_SCORE_ROWS,
+                )
+            else:
+                qsa_mqa_paged(
+                    q_index[chunk],
+                    cmp_pages,
+                    md.block_table,
+                    md.token_to_req[chunk],
+                    positions[chunk],
+                    md.seq_lens,
+                    self.ratio,
+                    logits,
+                    visible,
+                )
             blocks = self._scratch("blocks", end - start, self.block_topk, dtype=torch.int32)
             self._top_blocks(logits, visible, blocks)
             expand_qsa_block_indices(
diff --git a/python/freetoken/kernel/triton/qsa/__init__.py b/python/freetoken/kernel/triton/qsa/__init__.py
index 91753b7922..b04d27f050 100644
--- a/python/freetoken/kernel/triton/qsa/__init__.py
+++ b/python/freetoken/kernel/triton/qsa/__init__.py
@@ -3,7 +3,7 @@
 from .attend import qsa_sparse_paged_attention
 from .compress import qsa_compress_groups, qsa_index_norm_rope, qsa_store_rows
 from .expand import expand_qsa_block_indices
-from .score import qsa_mqa_paged
+from .score import qsa_mqa_paged, qsa_mqa_paged_rows, qsa_score_tile_map
 from .topk import qsa_block_topk, qsa_block_topk_scratch_width
 
 __all__ = [
@@ -13,6 +13,8 @@
     "qsa_compress_groups",
     "qsa_index_norm_rope",
     "qsa_mqa_paged",
+    "qsa_mqa_paged_rows",
+    "qsa_score_tile_map",
     "qsa_sparse_paged_attention",
     "qsa_store_rows",
 ]
diff --git a/python/freetoken/kernel/triton/qsa/score.py b/python/freetoken/kernel/triton/qsa/score.py
index 7b65da8d6f..24526d0e86 100644
--- a/python/freetoken/kernel/triton/qsa/score.py
+++ b/python/freetoken/kernel/triton/qsa/score.py
@@ -115,6 +115,227 @@ def _qsa_mqa_paged_kernel(
         )
 
 
+@triton.jit
+def _qsa_mqa_paged_rows_kernel(
+    q_ptr,
+    k_cache_ptr,
+    page_table_ptr,
+    tile_map_ptr,
+    token_to_req_ptr,
+    query_positions_ptr,
+    sequence_lengths_ptr,
+    visible_blocks_ptr,
+    logits_ptr,
+    stride_q_row,
+    stride_q_head,
+    stride_q_dim,
+    stride_cache_block,
+    stride_cache_token,
+    stride_cache_dim,
+    stride_table_req,
+    stride_table_page,
+    stride_logits_row,
+    num_tiles,
+    num_columns,
+    num_pages,
+    num_requests,
+    score_divisor,
+    PAGE_SIZE: tl.constexpr,
+    PAGE_TABLE_WIDTH: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    ROWS: tl.constexpr,
+    TILES_PER_PROG: tl.constexpr,
+    STAGES: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+) -> None:
+    """Prefill scoring, ROWS query rows x 4 heads per program.
+
+    The rows of a tile all belong to one request (``tile_map`` holds each tile's first row
+    and row count, cut at request boundaries), so one key tile serves every row. Each head
+    is its own ``[BLOCK_N, ROWS]`` dot and the head sum keeps the one-row kernel's order
+    ``(h0 + h1) + (h2 + h3)``: every score is bitwise the one-row kernel's.
+    """
+    tile = tl.program_id(0)
+    row0 = tl.load(tile_map_ptr + tile)
+    nrows = tl.load(tile_map_ptr + num_tiles + tile)
+    dims = tl.arange(0, BLOCK_D)
+    lanes = tl.arange(0, ROWS)
+    row_live = lanes < nrows
+    rows = row0 + lanes
+    request = tl.load(token_to_req_ptr + row0)
+    safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
+    query_position = tl.load(query_positions_ptr + rows, mask=row_live, other=0)
+    sequence_length = tl.load(
+        sequence_lengths_ptr + safe_request,
+        mask=(request >= 0) & (request < num_requests),
+        other=0,
+    )
+    visible = tl.minimum(
+        (query_position + 1) // COMPRESS_RATIO,
+        sequence_length // COMPRESS_RATIO,
+    )
+    visible = tl.where(row_live, visible, 0)
+    if tl.program_id(1) == 0:
+        tl.store(visible_blocks_ptr + rows, visible, mask=row_live)
+    visible_max = tl.max(visible, axis=0)
+    tile_start = tl.program_id(1) * TILES_PER_PROG
+    # Top-k is bounded by visible_blocks, so columns beyond it need no value.
+    if tile_start * BLOCK_N >= visible_max:
+        return
+    tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible_max, BLOCK_N))
+    tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N))
+
+    q_rows = q_ptr + rows[None, :].to(tl.int64) * stride_q_row + dims[:, None] * stride_q_dim
+    q_mask = row_live[None, :] & (dims[:, None] < HEAD_DIM)
+    query0 = tl.load(q_rows + 0 * stride_q_head, mask=q_mask, other=0.0)
+    query1 = tl.load(q_rows + 1 * stride_q_head, mask=q_mask, other=0.0)
+    query2 = tl.load(q_rows + 2 * stride_q_head, mask=q_mask, other=0.0)
+    query3 = tl.load(q_rows + 3 * stride_q_head, mask=q_mask, other=0.0)
+    column_offsets = tl.arange(0, BLOCK_N)
+    out_rows = logits_ptr + rows[None, :].to(tl.int64) * stride_logits_row
+    for col_tile in tl.range(tile_start, tile_end, num_stages=STAGES):
+        columns = col_tile * BLOCK_N + column_offsets
+        live = columns < visible_max
+        logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)
+        page_offset = columns % PAGE_SIZE
+        physical_page = tl.load(
+            page_table_ptr
+            + safe_request * stride_table_req
+            + logical_page * stride_table_page,
+            mask=live,
+            other=-1,
+        )
+        page_valid = live & (physical_page >= 0) & (physical_page < num_pages)
+        # physical_page * block stride can overflow int32 for large caches.
+        safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64)
+        keys = tl.load(
+            k_cache_ptr
+            + safe_physical_page[:, None] * stride_cache_block
+            + page_offset[:, None] * stride_cache_token
+            + dims[None, :] * stride_cache_dim,
+            mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM),
+            other=0.0,
+            eviction_policy="evict_first",
+        )
+        s0 = tl.maximum(tl.dot(keys, query0, out_dtype=tl.float32), 0.0)
+        s1 = tl.maximum(tl.dot(keys, query1, out_dtype=tl.float32), 0.0)
+        s2 = tl.maximum(tl.dot(keys, query2, out_dtype=tl.float32), 0.0)
+        s3 = tl.maximum(tl.dot(keys, query3, out_dtype=tl.float32), 0.0)
+        score = ((s0 + s1) + (s2 + s3)) / score_divisor
+        tl.store(
+            out_rows + columns[:, None],
+            tl.where(page_valid[:, None], score, -float("inf")),
+            mask=(columns[:, None] < visible[None, :]) & (columns[:, None] < num_columns),
+        )
+
+
+def qsa_score_tile_map(
+    qo_indptr: torch.Tensor, start: int, end: int, rows_per_tile: int, device: torch.device
+) -> torch.Tensor:
+    """``[2, tiles]`` int32 (first row, row count) of rows ``[start, end)``, relative to
+    ``start``, in tiles of at most ``rows_per_tile`` rows that never cross a request boundary
+    of the host-side ``qo_indptr`` (cu_seqlens)."""
+
+    bounds = [int(x) for x in qo_indptr.tolist()]
+    firsts: list[int] = []
+    counts: list[int] = []
+    for lo, hi in zip(bounds, bounds[1:]):
+        lo, hi = max(lo, start), min(hi, end)
+        for row in range(lo, hi, rows_per_tile):
+            firsts.append(row - start)
+            counts.append(min(rows_per_tile, hi - row))
+    host = torch.tensor([firsts, counts], dtype=torch.int32)
+    if device.type == "cuda":
+        host = host.pin_memory()
+    return host.to(device, non_blocking=True)
+
+
+SCORE_ROWS = 16
+
+
+def qsa_mqa_paged_rows(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    tile_map: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    compress_ratio: int,
+    logits: torch.Tensor,
+    visible_blocks: torch.Tensor,
+    score_scale: float | None = None,
+    rows_per_tile: int = SCORE_ROWS,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Prefill form of :func:`qsa_mqa_paged`: the same scores bit for bit, with
+    ``rows_per_tile`` query rows of one request per program instead of one row padded to 16
+    heads. ``tile_map`` comes from :func:`qsa_score_tile_map` over the same rows and must use
+    the same ``rows_per_tile``. Needs exactly 4 index heads."""
+
+    if q.ndim != 3 or q.shape[1] != 4 or q.shape[2] <= 0:
+        raise ValueError("QSA row-tiled scoring needs a [rows, 4, head_dim] query")
+    if k_cache.ndim != 4 or k_cache.shape[2] != 1 or k_cache.shape[3] != q.shape[2]:
+        raise ValueError("QSA cache must be [pages, page_size, 1, head_dim] matching the query")
+    if token_to_req.shape != (q.shape[0],) or query_positions.shape != (q.shape[0],):
+        raise ValueError("QSA request mapping and positions must match query rows")
+    if sequence_lengths.shape != (page_table.shape[0],):
+        raise ValueError("QSA sequence lengths must match page-table requests")
+    if q.dtype not in (torch.bfloat16, torch.float16) or k_cache.dtype is not q.dtype:
+        raise ValueError(
+            f"QSA scoring is 16-bit only, got query={q.dtype} keys={k_cache.dtype}"
+        )
+    if tile_map.ndim != 2 or tile_map.shape[0] != 2 or tile_map.dtype != torch.int32:
+        raise ValueError("QSA tile map must be [2, tiles] int32")
+    if rows_per_tile < 16 or rows_per_tile & (rows_per_tile - 1):
+        raise ValueError("rows_per_tile must be a power of two >= 16 (the dot's N)")
+    score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale
+    columns = logits.shape[1]
+    tiles = tile_map.shape[1]
+    if not q.shape[0] or not columns or not tiles:
+        return logits, visible_blocks
+    BLOCK_N = 64
+    BLOCK_D = max(16, triton.next_power_of_2(q.shape[2]))
+    tiles_per_program = 8
+    _qsa_mqa_paged_rows_kernel[(tiles, triton.cdiv(columns, BLOCK_N * tiles_per_program))](
+        q,
+        k_cache,
+        page_table,
+        tile_map,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+        visible_blocks,
+        logits,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k_cache.stride(0),
+        k_cache.stride(1),
+        k_cache.stride(3),
+        page_table.stride(0),
+        page_table.stride(1),
+        logits.stride(0),
+        tiles,
+        columns,
+        k_cache.shape[0],
+        page_table.shape[0],
+        float(score_divisor),
+        PAGE_SIZE=k_cache.shape[1],
+        PAGE_TABLE_WIDTH=page_table.shape[1],
+        HEAD_DIM=q.shape[2],
+        BLOCK_N=BLOCK_N,
+        BLOCK_D=BLOCK_D,
+        ROWS=rows_per_tile,
+        TILES_PER_PROG=tiles_per_program,
+        STAGES=2,
+        COMPRESS_RATIO=compress_ratio,
+        num_warps=4,
+    )
+    return logits, visible_blocks
+
+
 def qsa_mqa_paged(
     q: torch.Tensor,
     k_cache: torch.Tensor,
diff --git a/tests/kernels/test_qsa_score_rows.py b/tests/kernels/test_qsa_score_rows.py
new file mode 100644
index 0000000000..fec0898679
--- /dev/null
+++ b/tests/kernels/test_qsa_score_rows.py
@@ -0,0 +1,136 @@
+"""The prefill row-tiled QSA score kernel (``qsa_mqa_paged_rows``) against the one-row kernel.
+
+Both must give the same visible counts, bitwise-equal logits on every visible column, and so the
+same top-k block selection, for ragged multi-request prefill chunks whose request starts are not
+aligned to the row tile, and for a row slice in the middle of a forward (the ``_select`` chunking).
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+CMP_PAGE = 16
+RATIO = 4
+HEADS = 4
+DIM = 128
+TOPK = 512
+
+
+def _case(lengths, rows_per_req, seed, with_invalid=False):
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(seed)
+    max_blocks = max(lengths) // RATIO
+    width = -(-max_blocks // CMP_PAGE)
+    bs = len(lengths)
+    total = bs * width + 3
+    block_table = torch.randperm(total, device=device, generator=gen)[: bs * width]
+    block_table = block_table.reshape(bs, width).to(torch.int32)
+    if with_invalid:
+        block_table[0, 1] = -1
+        block_table[-1, 0] = total + 5
+    cache = torch.randn(total, CMP_PAGE, 1, DIM, device=device, dtype=torch.bfloat16, generator=gen)
+    rows = sum(rows_per_req)
+    q = torch.randn(rows, HEADS, DIM, device=device, dtype=torch.bfloat16, generator=gen)
+    token_to_req = torch.repeat_interleave(
+        torch.arange(bs, dtype=torch.int32), torch.tensor(rows_per_req)
+    ).to(device)
+    positions = torch.cat(
+        [
+            torch.arange(n - r, n, dtype=torch.int32, device=device)
+            for n, r in zip(lengths, rows_per_req)
+        ]
+    )
+    seq_lens = torch.tensor(lengths, dtype=torch.int32, device=device)
+    qo_indptr = torch.tensor([0] + list(rows_per_req), dtype=torch.int32).cumsum(0).to(torch.int32)
+    return q, cache, block_table, token_to_req, positions, seq_lens, qo_indptr, width * CMP_PAGE
+
+
+def _run_both(case, start, end, rows_per_tile):
+    from freetoken.kernel.triton.qsa import (
+        qsa_block_topk,
+        qsa_mqa_paged,
+        qsa_mqa_paged_rows,
+        qsa_score_tile_map,
+    )
+
+    q, cache, table, t2r, pos, seq, qo, columns = case
+    sl = slice(start, end)
+    n = end - start
+    out = []
+    for rows_kernel in (False, True):
+        logits = torch.full((n, columns), float("nan"), dtype=torch.float32, device=q.device)
+        visible = torch.full((n,), -7, dtype=torch.int32, device=q.device)
+        if rows_kernel:
+            tiles = qsa_score_tile_map(qo, start, end, rows_per_tile, q.device)
+            qsa_mqa_paged_rows(
+                q[sl], cache, table, tiles, t2r[sl], pos[sl], seq, RATIO, logits, visible,
+                rows_per_tile=rows_per_tile,
+            )
+        else:
+            qsa_mqa_paged(q[sl], cache, table, t2r[sl], pos[sl], seq, RATIO, logits, visible)
+        blocks = torch.empty((n, TOPK), dtype=torch.int32, device=q.device)
+        qsa_block_topk(logits.clone(), visible, blocks)
+        out.append((logits, visible, blocks))
+    return out
+
+
+def _assert_same(old, new):
+    (lo, vo, bo), (ln, vn, bn) = old, new
+    assert torch.equal(vo, vn)
+    col = torch.arange(lo.shape[1], device=lo.device)
+    live = col[None, :] < vo[:, None].to(torch.int64)
+    assert live.any()
+    a, b = lo[live], ln[live]
+    assert not torch.isnan(a).any() and not torch.isnan(b).any()
+    # bitwise, including the sign of zero and -inf on invalid pages
+    assert torch.equal(a.view(torch.int32), b.view(torch.int32)), (
+        f"{(a.view(torch.int32) != b.view(torch.int32)).sum().item()} of {a.numel()} scores differ"
+    )
+    assert torch.equal(bo, bn)
+
+
+@requires_cuda
+@pytest.mark.parametrize("rows_per_tile", [16, 32])
+@pytest.mark.parametrize(
+    "lengths, rows_per_req",
+    [
+        ([9000], [517]),
+        ([20000, 3001], [37, 300]),
+        ([700, 12345, 64, 5000], [5, 210, 33, 17]),
+        ([4096], [4096]),
+    ],
+)
+def test_rows_kernel_is_bitwise_the_one_row_kernel(lengths, rows_per_req, rows_per_tile):
+    case = _case(lengths, rows_per_req, seed=sum(lengths) + rows_per_tile)
+    rows = sum(rows_per_req)
+    old, new = _run_both(case, 0, rows, rows_per_tile)
+    _assert_same(old, new)
+
+
+@requires_cuda
+@pytest.mark.parametrize("start, end", [(100, 612), (0, 512), (517, 818)])
+def test_rows_kernel_on_a_row_slice(start, end):
+    case = _case([30000, 2000, 9000], [600, 101, 117], seed=start + end, with_invalid=True)
+    old, new = _run_both(case, start, end, 16)
+    _assert_same(old, new)
+
+
+@requires_cuda
+def test_rows_kernel_at_long_context():
+    """A late 249K-context chunk: ~62K visible blocks per row, 512-row slice."""
+    case = _case([249000], [512], seed=3)
+    old, new = _run_both(case, 0, 512, 16)
+    _assert_same(old, new)
+
+
+def test_tile_map_never_crosses_a_request():
+    from freetoken.kernel.triton.qsa import qsa_score_tile_map
+
+    qo = torch.tensor([0, 5, 40, 41, 100], dtype=torch.int32)
+    tiles = qsa_score_tile_map(qo, 3, 90, 16, torch.device("cpu"))
+    firsts, counts = tiles.tolist()
+    assert firsts == [0, 2, 18, 34, 37, 38, 54, 70, 86]
+    assert counts == [2, 16, 16, 3, 1, 16, 16, 16, 1]

From fa5e0ca70e008b0b8e6a59e6b4d04039077f97dd Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 14:08:31 -0400
Subject: [PATCH 248/249] test(fp8): pin a dynamic fp8 linear row against the
 other rows of its batch

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 tests/models/qwen4_exp/test_fp8_dense.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py
index 74da973341..c7ba264f8e 100644
--- a/tests/models/qwen4_exp/test_fp8_dense.py
+++ b/tests/models/qwen4_exp/test_fp8_dense.py
@@ -113,6 +113,22 @@ def test_fp8_linear_matches_bf16_on_the_dequantized_weight(rows: int):
     )
 
 
+@requires_cuda
+def test_fp8_linear_row_does_not_depend_on_the_other_rows():
+    """Rows of one varlen prefill batch belong to different sequences: a row's output must not
+    move when another sequence's activations do (same M, so the GEMM config is unchanged)."""
+    from freetoken.layers.fp8_dynamic import fp8_dynamic_linear
+
+    torch.manual_seed(0)
+    w8, scale = _quantize_per_tensor((torch.randn(256, 128) * 0.02).to(torch.bfloat16))
+    w8, scale = w8.cuda(), scale.cuda()
+    mine = torch.randn(61, 128, device="cuda", dtype=torch.bfloat16)
+    other = torch.randn(53, 128, device="cuda", dtype=torch.bfloat16)
+    quiet = fp8_dynamic_linear(torch.cat([mine, other]), w8, scale)[:61]
+    loud = fp8_dynamic_linear(torch.cat([mine, other * 8]), w8, scale)[:61]
+    assert torch.equal(quiet, loud)
+
+
 @requires_cuda
 def test_quant_per_tensor_zero_input_is_finite():
     from freetoken.layers.fp8_dynamic import quant_per_tensor

From 6c381c366e8017d644ff6390f71b7e0b1db1b3be Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Tue, 20 Oct 2026 14:09:32 -0400
Subject: [PATCH 249/249] fix(fp8): scale dynamic fp8 activations per row so
 batched sequences stay independent

Assisted-by: Claude Opus 5.5
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
---
 python/freetoken/layers/fp8_dynamic.py | 49 +++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py
index caeebb3221..1d8713cc8e 100644
--- a/python/freetoken/layers/fp8_dynamic.py
+++ b/python/freetoken/layers/fp8_dynamic.py
@@ -2,7 +2,8 @@
 
 The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``),
 produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding).
-The activation is quantized per call with a dynamic per-tensor scale: one Triton launch (a
+The activation is quantized per call with a dynamic per-row scale (``quant_per_row``, applied
+after the GEMM). ``quant_per_tensor`` is the older per-tensor cast: one Triton launch (a
 single program walking the tensor) for tiny inputs, two above ``_SPLIT_MIN_ELEMENTS`` (a
 block-parallel amax, then a reduce+cast). Both compute the same arithmetic, so the split is
 a pure speed choice. No host sync anywhere (the scales stay on the device), so the decode
@@ -116,17 +117,56 @@ def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     return out, scale
 
 
+@triton.jit
+def _quant_rows_kernel(x_ptr, out_ptr, scale_ptr, k, BLOCK: tl.constexpr):
+    """One program per row: that row's max|x|, then its cast."""
+    row = tl.program_id(0).to(tl.int64)
+    x_row = x_ptr + row * k
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for start in range(0, k, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        acc = tl.maximum(acc, tl.abs(tl.load(x_row + offs, mask=offs < k, other=0.0).to(tl.float32)))
+    amax = tl.maximum(tl.max(acc, axis=0), 1e-12)
+    tl.store(scale_ptr + row, amax / 448.0)
+    inv = 448.0 / amax
+    for start in range(0, k, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        mask = offs < k
+        v = tl.load(x_row + offs, mask=mask, other=0.0).to(tl.float32) * inv
+        v = tl.minimum(tl.maximum(v, -448.0), 448.0)
+        tl.store(out_ptr + row * k + offs, v.to(tl.float8e4nv), mask=mask)
+
+
+def quant_per_row(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous ``[M, K]``, ``scale`` fp32 ``[M, 1]``."""
+    m, k = x.shape
+    out = torch.empty_like(x, dtype=FP8)
+    scale = torch.empty((m, 1), dtype=torch.float32, device=x.device)
+    if m:
+        _quant_rows_kernel[(m,)](x, out, scale, k, BLOCK=1024, num_warps=4)
+    return out, scale
+
+
 def fp8_dynamic_linear(
     x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor
 ) -> torch.Tensor:
     """``x @ (weight * weight_scale)^T`` in W8A8; ``weight`` [N, K] e4m3 row-major, whose ``.t()``
-    is the column-major operand cuBLASLt wants (a stride change, never a copy)."""
+    is the column-major operand cuBLASLt wants (a stride change, never a copy).
+
+    The activation scale is per row, so a row never depends on the other sequences of its
+    batch. It is applied after the GEMM: cuBLASLt's rowwise-scaled mode is 1.4-2.9x slower
+    than the tensorwise one on sm_120."""
     *lead, k = x.shape
     x2 = x.reshape(-1, k).contiguous()
-    x8, scale = quant_per_tensor(x2)
+    if x2.shape[0] == 1:  # one row: per-tensor is per-row, and the scale rides in the GEMM epilogue
+        x8, scale = quant_per_tensor(x2)
+        y = torch._scaled_mm(x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype)
+        return y.reshape(*lead, weight.shape[0])
+    x8, row_scale = quant_per_row(x2)
     y = torch._scaled_mm(
-        x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype
+        x8, weight.t(), scale_a=weight_scale.new_ones(()), scale_b=weight_scale, out_dtype=x.dtype
     )
+    y.mul_(row_scale)
     return y.reshape(*lead, weight.shape[0])
 
 
@@ -207,5 +247,6 @@ def _logits(self, x: torch.Tensor) -> torch.Tensor:
     "Fp8DynamicLinear",
     "Fp8DynamicRowParallel",
     "fp8_dynamic_linear",
+    "quant_per_row",
     "quant_per_tensor",
 ]